From 0865c1d6ade2b2bcc1ea00a5e10128025ceef3b2 Mon Sep 17 00:00:00 2001
From: Matthew Filipovich <42307495+MatthewFilipovich@users.noreply.github.com>
Date: Sun, 8 Dec 2024 19:10:05 -0800
Subject: [PATCH] Add bessel beam

---
 tests/test_fields.py             |  2 +-
 tests/test_profiles.py           | 23 ++++++++++++
 torchoptics/profiles/__init__.py |  1 +
 torchoptics/profiles/bessel.py   | 61 ++++++++++++++++++++++++++++++++
 4 files changed, 86 insertions(+), 1 deletion(-)
 create mode 100644 torchoptics/profiles/bessel.py

diff --git a/tests/test_fields.py b/tests/test_fields.py
index 2830320..4c81aee 100644
--- a/tests/test_fields.py
+++ b/tests/test_fields.py
@@ -275,7 +275,7 @@ def test_visualize(self):
         field = Field(data, wavelength=1, spacing=1)
 
         # Test visualization of the field data
-        fig = field.visualize(intensity=True, show=False, return_fig=True)
+        fig = field.visualize(show=False, return_fig=True)
         self.assertIsInstance(fig, plt.Figure)
 
 
diff --git a/tests/test_profiles.py b/tests/test_profiles.py
index 2f4ee72..2b84b5b 100644
--- a/tests/test_profiles.py
+++ b/tests/test_profiles.py
@@ -354,5 +354,28 @@ def test_incoherent(self):
         self.assertTrue(torch.all(incoherent_data == 0))  # off-diagonal elements should be zero
 
 
+class TestBesselProfile(unittest.TestCase):
+    def setUp(self):
+        self.shape = (100, 100)
+        self.cone_angle = torch.pi / 4
+        self.wavelength = 0.5
+        self.spacing = (0.1, 0.1)
+        self.offset = (0.0, 0.0)
+
+        self.profile = bessel(
+            shape=self.shape,
+            cone_angle=self.cone_angle,
+            wavelength=self.wavelength,
+            spacing=self.spacing,
+            offset=self.offset,
+        )
+
+    def test_bessel_shape(self):
+        self.assertEqual(self.profile.shape, self.shape)
+
+    def test_bessel_values(self):
+        self.assertTrue(torch.all(self.profile.abs() <= 1))
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/torchoptics/profiles/__init__.py b/torchoptics/profiles/__init__.py
index b754a68..a8b711b 100644
--- a/torchoptics/profiles/__init__.py
+++ b/torchoptics/profiles/__init__.py
@@ -1,5 +1,6 @@
 """This module contains functions to generate different types of profiles."""
 
+from .bessel import bessel
 from .gratings import blazed_grating, sinusoidal_amplitude_grating, sinusoidal_phase_grating
 from .hermite_gaussian import gaussian, hermite_gaussian
 from .laguerre_gaussian import laguerre_gaussian
diff --git a/torchoptics/profiles/bessel.py b/torchoptics/profiles/bessel.py
new file mode 100644
index 0000000..c389c72
--- /dev/null
+++ b/torchoptics/profiles/bessel.py
@@ -0,0 +1,61 @@
+"""This module defines a function to generate the bessel beam profile."""
+
+import math
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from ..config import get_default_wavelength
+from ..planar_geometry import PlanarGeometry
+from ..type_defs import Scalar, Vector2
+
+__all__ = ["bessel"]
+
+
+def bessel(
+    shape: Vector2,
+    cone_angle: float,
+    wavelength: Optional[Scalar] = None,
+    spacing: Optional[Vector2] = None,
+    offset: Optional[Vector2] = None,
+) -> Tensor:
+    r"""
+    Generates a zeroth-order Bessel beam.
+
+    The zeroth-order Bessel beam is defined by the following equation:
+
+    .. math::
+        \psi(r) = J_0(k \, r \sin(\theta)),
+
+    where:
+
+    - :math:`J_0` is the zeroth-order Bessel function of the first kind,
+    - :math:`\theta` is the cone angle of the Bessel beam, and
+    - :math:`k` is the wave number, :math:`k = 2\pi / \lambda`.
+
+    Args:
+        shape (Vector2): Number of grid points along the planar dimensions.
+        cone_angle (float): The cone angle in radians.
+        wavelength (Scalar, optional): The wavelength of the beam. Default: if `None`, uses a global default
+            (see :meth:`torchoptics.set_default_wavelength()`).
+        spacing (Optional[Vector2]): Distance between grid points along planar dimensions. Default: if
+            `None`, uses a global default (see :meth:`torchoptics.set_default_spacing()`).
+        offset (Optional[Vector2]): Center coordinates of the beam. Default: `(0, 0)`.
+
+    Returns:
+        Tensor: The generated zeroth-order Bessel beam profile.
+
+    """
+    wavelength = wavelength if wavelength is not None else get_default_wavelength()
+
+    # Calculate the wave number k and its radial component
+    k = 2 * torch.pi / wavelength  # type: ignore
+    k_r = k * math.sin(cone_angle)
+
+    # Generate the planar grid
+    x, y = PlanarGeometry(shape, spacing=spacing, offset=offset).meshgrid()
+    r = torch.sqrt(x**2 + y**2)
+
+    # Calculate the zeroth-order Bessel beam
+    return torch.special.bessel_j0(k_r * r)  # pylint: disable=not-callable