feat: Adding gaussian monitors

Author: momchil-flex

tests/test_components/test_monitor.py

@@ -344,6 +344,14 @@ def test_diffraction_validators():
 FREQS = np.array([1, 2, 3]) * 1e12
 
 
+def test_gaussian_overlap_monitors_basic():
+    g = td.GaussianOverlapMonitor(size=(1, 1, 0), name="g", freqs=FREQS)
+    a = td.AstigmaticGaussianOverlapMonitor(size=(1, 1, 0), name="a", freqs=FREQS)
+    for m in (g, a):
+        s = m.storage_size(num_cells=10, tmesh=[0.0, 1.0])
+        assert isinstance(s, int) and s > 0
+
+
 def test_monitor():
     size = (1, 2, 3)
     center = (1, 2, 3)
@@ -381,10 +389,14 @@ def test_monitor():
     m10 = td.PermittivityMonitor(size=size, center=center, freqs=FREQS, name="perm")
     m11 = td.AuxFieldTimeMonitor(size=size, center=center, name="aux_field_time", fields=("Nfx",))
     m12 = td.MediumMonitor(size=size, center=center, freqs=FREQS, name="mat")
+    m13 = td.GaussianOverlapMonitor(size=(1, 1, 0), center=center, freqs=FREQS, name="gauss")
+    m14 = td.AstigmaticGaussianOverlapMonitor(
+        size=(1, 1, 0), center=center, freqs=FREQS, name="astigauss"
+    )
 
     tmesh = np.linspace(0, 1, 10)
 
-    for m in [m1, m2, m3, m4, m5, m6, m7, m8, m9, m10, m12]:
+    for m in [m1, m2, m3, m4, m5, m6, m7, m8, m9, m10, m11, m12, m13, m14]:
         m.storage_size(num_cells=100, tmesh=tmesh)
 
     for m in [m2, m4]:

tidy3d/components/monitor.py

@@ -330,7 +330,79 @@ def normal_axis(self) -> Axis:
         return self.size.index(0.0)
 
 
-class AbstractModeMonitor(PlanarMonitor, FreqMonitor):
+class AbstractOverlapMonitor(PlanarMonitor, FreqMonitor):
+    """:class:`Monitor` that projects fields onto a specified basis and stores overlap amplitudes.
+
+    This base is shared by ModeMonitor and Gaussian-overlap monitors.
+    """
+
+    store_fields_direction: Direction = pydantic.Field(
+        None,
+        title="Store Fields",
+        description="Propagation direction for the field profiles stored from overlap calculation.",
+    )
+
+    colocate: bool = pydantic.Field(
+        True,
+        title="Colocate Fields",
+        description="Toggle whether fields should be colocated to grid cell boundaries (i.e. primal grid nodes).",
+    )
+
+    conjugated_dot_product: bool = pydantic.Field(
+        True,
+        title="Conjugated Dot Product",
+        description="Use conjugated or non-conjugated dot product for overlap/decomposition.",
+    )
+
+    def plot(
+        self,
+        x: Optional[float] = None,
+        y: Optional[float] = None,
+        z: Optional[float] = None,
+        ax: Ax = None,
+        **patch_kwargs,
+    ) -> Ax:
+        """Plot this overlap monitor with an arrow indicating propagation direction."""
+        ax = super().plot(x=x, y=y, z=z, ax=ax, **patch_kwargs)
+
+        kwargs_alpha = patch_kwargs.get("alpha")
+        arrow_alpha = ARROW_ALPHA if kwargs_alpha is None else kwargs_alpha
+
+        ax = self._plot_arrow(
+            x=x,
+            y=y,
+            z=z,
+            ax=ax,
+            direction=self._dir_arrow,
+            bend_radius=None,
+            bend_axis=None,
+            color=ARROW_COLOR_MONITOR,
+            alpha=arrow_alpha,
+            both_dirs=True,
+        )
+        return ax
+
+    @cached_property
+    def _dir_arrow(self) -> tuple[float, float, float]:
+        """Direction vector from angles used for overlap (children must define angles)."""
+        theta, phi = self._angles
+        dx = np.cos(phi) * np.sin(theta)
+        dy = np.sin(phi) * np.sin(theta)
+        dz = np.cos(theta)
+        return self.unpop_axis(dz, (dx, dy), axis=self.normal_axis)
+
+    @property
+    def _angles(self) -> tuple[float, float]:
+        """Angle tuple (theta, phi) in radians. Children override to supply values."""
+        return (0.0, 0.0)
+
+    def _storage_size_solver(self, num_cells: int, tmesh: ArrayFloat1D) -> int:
+        """Default size of intermediate data; store all fields on plane for overlap."""
+        num_sample = len(getattr(self, "freqs", [0]))
+        return BYTES_COMPLEX * num_cells * num_sample * 6
+
+
+class AbstractModeMonitor(AbstractOverlapMonitor):
     """:class:`Monitor` that records mode-related data."""
 
     mode_spec: ModeSpec = pydantic.Field(
@@ -393,6 +465,10 @@ def plot(
         )
         return ax
 
+    @cached_property
+    def _angles(self) -> tuple[float, float]:
+        return (self.mode_spec.angle_theta, self.mode_spec.angle_phi)
+
     @cached_property
     def _dir_arrow(self) -> tuple[float, float, float]:
         """Source direction normal vector in cartesian coordinates."""
@@ -433,6 +509,138 @@ def _storage_size_solver(self, num_cells: int, tmesh: ArrayFloat1D) -> int:
         return bytes_single
 
 
+class AbstractGaussianOverlapMonitor(AbstractOverlapMonitor):
+    """:class:`Monitor` that records amplitudes from decomposition onto a Gaussian-like beam.
+
+    Common fields and behavior shared by GaussianOverlapMonitor and
+    AstigmaticGaussianOverlapMonitor.
+    """
+
+    angle_theta: float = pydantic.Field(
+        0.0,
+        title="Polar Angle",
+        description="Polar angle of propagation direction.",
+        units=RADIAN,
+    )
+
+    angle_phi: float = pydantic.Field(
+        0.0,
+        title="Azimuth Angle",
+        description="Azimuth angle of propagation direction.",
+        units=RADIAN,
+    )
+
+    pol_angle: float = pydantic.Field(
+        0,
+        title="Polarization Angle",
+        description="Specifies the angle between the electric field polarization of the "
+        "source and the plane defined by the injection axis and the propagation axis (rad). "
+        "``pol_angle=0`` (default) specifies P polarization, "
+        "while ``pol_angle=np.pi/2`` specifies S polarization. "
+        "At normal incidence when S and P are undefined, ``pol_angle=0`` defines: "
+        "- ``Ey`` polarization for propagation along ``x``."
+        "- ``Ex`` polarization for propagation along ``y``."
+        "- ``Ex`` polarization for propagation along ``z``.",
+        units=RADIAN,
+    )
+
+    @property
+    def _angles(self) -> tuple[float, float]:
+        return (self.angle_theta, self.angle_phi)
+
+    def storage_size(self, num_cells: int, tmesh: ArrayFloat1D) -> int:
+        """Size of monitor storage given the number of points after discretization."""
+        # store complex amplitudes for +/- directions
+        num_dirs = 2
+        return BYTES_COMPLEX * len(self.freqs) * num_dirs
+
+
+class GaussianOverlapMonitor(AbstractGaussianOverlapMonitor):
+    """:class:`Monitor` that records amplitudes from decomposition onto a Gaussian beam.
+
+    Example
+    -------
+    >>> gauss = GaussianOverlapMonitor(
+    ...     size=(0,3,3),
+    ...     source_time=pulse,
+    ...     pol_angle=np.pi / 2,
+    ...     waist_radius=1.0)
+
+    Notes
+    --------
+    If one wants the focus 'in front' of the source, a negative value of ``waist_distance`` is needed.
+
+    .. image:: ../../_static/img/beam_waist.png
+        :width: 30%
+        :align: center
+    """
+
+    waist_radius: pydantic.PositiveFloat = pydantic.Field(
+        1.0,
+        title="Waist Radius",
+        description="Radius of the beam at the waist.",
+        units=MICROMETER,
+    )
+
+    waist_distance: float = pydantic.Field(
+        0.0,
+        title="Waist Distance",
+        description="Distance from the beam waist along the propagation direction. "
+        "A positive value means the waist is positioned behind the source, considering the propagation direction. "
+        "For example, for a beam propagating in the ``+`` direction, a positive value of ``beam_distance`` "
+        "means the beam waist is positioned in the ``-`` direction (behind the source). "
+        "A negative value means the beam waist is in the ``+`` direction (in front of the source). "
+        "For an angled source, the distance is defined along the rotated propagation direction.",
+        units=MICROMETER,
+    )
+
+
+class AstigmaticGaussianOverlapMonitor(AbstractGaussianOverlapMonitor):
+    """:class:`Monitor` that records amplitudes from decomposition onto an astigmatic Gaussian beam.
+
+    The simple astigmatic Gaussian distribution allows
+    both an elliptical intensity profile and different waist locations for the two principal axes
+    of the ellipse. When equal waist sizes and equal waist distances are specified in the two
+    directions, this source becomes equivalent to :class:`GaussianBeam`.
+
+    Notes
+    -----
+
+        This class implements the simple astigmatic Gaussian beam described in _`[1]`.
+
+        **References**:
+
+        .. [1] Kochkina et al., Applied Optics, vol. 52, issue 24, 2013.
+
+    Example
+    -------
+    >>> gauss = AstigmaticGaussianOverlapMonitor(
+    ...     size=(0,3,3),
+    ...     pol_angle=np.pi / 2,
+    ...     waist_sizes=(1.0, 2.0),
+    ...     waist_distances = (3.0, 4.0))
+    """
+
+    waist_sizes: tuple[pydantic.PositiveFloat, pydantic.PositiveFloat] = pydantic.Field(
+        (1.0, 1.0),
+        title="Waist sizes",
+        description="Size of the beam at the waist in the local x and y directions.",
+        units=MICROMETER,
+    )
+
+    waist_distances: tuple[float, float] = pydantic.Field(
+        (0.0, 0.0),
+        title="Waist distances",
+        description="Distance to the beam waist along the propagation direction "
+        "for the waist sizes in the local x and y directions. "
+        "When ``direction`` is ``+`` and ``waist_distances`` are positive, the waist "
+        "is on the ``-`` side (behind) the source plane. When ``direction`` is ``+`` and "
+        "``waist_distances`` are negative, the waist is on the ``+`` side (in front) of "
+        "the source plane.",
+        units=MICROMETER,
+    )
+
+
 class FieldMonitor(AbstractFieldMonitor, FreqMonitor):
     """:class:`Monitor` that records electromagnetic fields in the frequency domain.
 

tidy3d/components/types/monitor.py

@@ -20,6 +20,8 @@
     ModeMonitor,
     ModeSolverMonitor,
     PermittivityMonitor,
+    GaussianOverlapMonitor,
+    AstigmaticGaussianOverlapMonitor,
 )
 
 # types of monitors that are accepted by simulation
@@ -40,4 +42,6 @@
     DirectivityMonitor,
     MicrowaveModeMonitor,
     MicrowaveModeSolverMonitor,
+    GaussianOverlapMonitor,
+    AstigmaticGaussianOverlapMonitor,
 ]