add windowing feature to field projection monitors

- revisions based on PR 1275 review

- update docstring and comment

Author: shashwat-sh

CHANGELOG.md

@@ -29,6 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Ability to downsample recorded near fields to speed up server-side far field projections.
 - `FieldData.apply_phase(phase)` to multiply field data by a phase.
 - Optional `phase` argument to `SimulationData.plot_field` that applies a phase to complex-valued fields.
+- Ability to window near fields for spatial filtering of far fields for both client- and server-side field projections.
 
 ### Changed
 - Indent for the json string of Tidy3D models has been changed to `None` when used internally; kept as `indent=4` for writing to `json` and `yaml` files.

tests/sims/simulation_2_5_0rc3.json

@@ -1706,6 +1706,10 @@
                 3.0
             ],
             "far_field_approx": true,
+            "window_size": [
+                0.0,
+                0.0
+            ],
             "proj_distance": 1000000.0,
             "theta": [
                 -1.5707963267948966,
@@ -1851,6 +1855,10 @@
                 3.0
             ],
             "far_field_approx": true,
+            "window_size": [
+                0.0,
+                0.0
+            ],
             "proj_axis": 2,
             "proj_distance": 5.0,
             "x": [
@@ -1903,6 +1911,10 @@
                 3.0
             ],
             "far_field_approx": true,
+            "window_size": [
+                0.0,
+                0.0
+            ],
             "proj_axis": 2,
             "proj_distance": 1000000.0,
             "ux": [
@@ -1952,6 +1964,10 @@
                 3.0
             ],
             "far_field_approx": true,
+            "window_size": [
+                0.0,
+                0.0
+            ],
             "proj_distance": 1000000.0,
             "theta": [
                 -1.5707963267948966,

tests/test_components/test_monitor.py

@@ -197,6 +197,26 @@ def test_fieldproj_local_origin():
     M.local_origin
 
 
+def test_fieldproj_window():
+    M = td.FieldProjectionAngleMonitor(
+        size=(2, 0, 2), theta=[1, 2], phi=[0], name="f", freqs=[2e12], window_size=(0.2, 1)
+    )
+    window_size, window_minus, window_plus = M.window_parameters()
+    window_size, window_minus, window_plus = M.window_parameters(M.bounds)
+    points = np.linspace(0, 10, 100)
+    _ = M.window_function(points, window_size, window_minus, window_plus, 2)
+    # do not allow a window size larger than 1
+    with pytest.raises(pydantic.ValidationError):
+        _ = td.FieldProjectionAngleMonitor(
+            size=(2, 0, 2), theta=[1, 2], phi=[0], name="f", freqs=[2e12], window_size=(0.2, 1.1)
+        )
+    # do not allow non-zero windows for volume monitors
+    with pytest.raises(pydantic.ValidationError):
+        _ = td.FieldProjectionAngleMonitor(
+            size=(2, 1, 2), theta=[1, 2], phi=[0], name="f", freqs=[2e12], window_size=(0.2, 0)
+        )
+
+
 PROJ_MNTS = [
     td.FieldProjectionAngleMonitor(size=(2, 0, 2), theta=[1, 2], phi=[0], name="f", freqs=[2e12]),
     td.FieldProjectionCartesianMonitor(

tidy3d/components/field_projection.py

@@ -459,6 +459,45 @@ def integrate_for_one_theta(i_th: int):
 
         return Er, Etheta, Ephi, Hr, Htheta, Hphi
 
+    @staticmethod
+    def apply_window_to_currents(
+        proj_monitor: AbstractFieldProjectionMonitor, currents: xr.Dataset
+    ) -> xr.Dataset:
+        """Apply windowing function to the surface currents."""
+        if proj_monitor.size.count(0.0) == 0:
+            return currents
+        if proj_monitor.window_size == (0, 0):
+            return currents
+
+        pts = [currents[name].values for name in ["x", "y", "z"]]
+
+        custom_bounds = [
+            [pts[i][0] for i in range(3)],
+            [pts[i][-1] for i in range(3)],
+        ]
+
+        window_size, window_minus, window_plus = proj_monitor.window_parameters(
+            custom_bounds=custom_bounds
+        )
+
+        new_currents = currents.copy(deep=True)
+        for dim, (dim_name, points) in enumerate(zip("xyz", pts)):
+            window_fn = proj_monitor.window_function(
+                points=points,
+                window_size=window_size,
+                window_minus=window_minus,
+                window_plus=window_plus,
+                dim=dim,
+            )
+            window_data = xr.DataArray(
+                window_fn,
+                dims=[dim_name],
+                coords=[points],
+            )
+            new_currents *= window_data
+
+        return new_currents
+
     def project_fields(
         self, proj_monitor: AbstractFieldProjectionMonitor
     ) -> AbstractFieldProjectionData:
@@ -514,10 +553,17 @@ def _project_fields_angular(
 
         for surface in self.surfaces:
 
+            # apply windowing to currents
+            currents = self.apply_window_to_currents(monitor, self.currents[surface.monitor.name])
+
             if monitor.far_field_approx:
                 for idx_f, frequency in enumerate(freqs):
                     _fields = self._far_fields_for_surface(
-                        frequency, theta, phi, surface, self.currents[surface.monitor.name]
+                        frequency=frequency,
+                        theta=theta,
+                        phi=phi,
+                        surface=surface,
+                        currents=currents,
                     )
                     for field, _field in zip(fields, _fields):
                         field[..., idx_f] += _field * phase[idx_f]
@@ -534,7 +580,7 @@ def _project_fields_angular(
                 ):
                     _x, _y, _z = monitor.sph_2_car(monitor.proj_distance, _theta, _phi)
                     _fields = self._fields_for_surface_exact(
-                        _x, _y, _z, surface, self.currents[surface.monitor.name]
+                        x=_x, y=_y, z=_z, surface=surface, currents=currents
                     )
                     for field, _field in zip(fields, _fields):
                         field[0, i, j, :] += _field
@@ -596,16 +642,25 @@ def _project_fields_cartesian(
 
             for surface in self.surfaces:
 
+                # apply windowing to currents
+                currents = self.apply_window_to_currents(
+                    monitor, self.currents[surface.monitor.name]
+                )
+
                 if monitor.far_field_approx:
                     for idx_f, frequency in enumerate(freqs):
                         _fields = self._far_fields_for_surface(
-                            frequency, theta, phi, surface, self.currents[surface.monitor.name]
+                            frequency=frequency,
+                            theta=theta,
+                            phi=phi,
+                            surface=surface,
+                            currents=currents,
                         )
                         for field, _field in zip(fields, _fields):
                             field[i, j, k, idx_f] += _field * phase[idx_f]
                 else:
                     _fields = self._fields_for_surface_exact(
-                        _x, _y, _z, surface, self.currents[surface.monitor.name]
+                        x=_x, y=_y, z=_z, surface=surface, currents=currents
                     )
                     for field, _field in zip(fields, _fields):
                         field[i, j, k, :] += _field
@@ -658,18 +713,27 @@ def _project_fields_kspace(
 
             for surface in self.surfaces:
 
+                # apply windowing to currents
+                currents = self.apply_window_to_currents(
+                    monitor, self.currents[surface.monitor.name]
+                )
+
                 if monitor.far_field_approx:
                     for idx_f, frequency in enumerate(freqs):
                         _fields = self._far_fields_for_surface(
-                            frequency, theta, phi, surface, self.currents[surface.monitor.name]
+                            frequency=frequency,
+                            theta=theta,
+                            phi=phi,
+                            surface=surface,
+                            currents=currents,
                         )
                         for field, _field in zip(fields, _fields):
                             field[i, j, 0, idx_f] += _field * phase[idx_f]
 
                 else:
                     _x, _y, _z = monitor.sph_2_car(monitor.proj_distance, theta, phi)
                     _fields = self._fields_for_surface_exact(
-                        _x, _y, _z, surface, self.currents[surface.monitor.name]
+                        x=_x, y=_y, z=_z, surface=surface, currents=currents
                     )
                     for field, _field in zip(fields, _fields):
                         field[i, j, 0, :] += _field

tidy3d/components/monitor.py

@@ -5,7 +5,7 @@
 import pydantic.v1 as pydantic
 import numpy as np
 
-from .types import Ax, EMField, ArrayFloat1D, FreqArray, FreqBound
+from .types import Ax, EMField, ArrayFloat1D, FreqArray, FreqBound, Bound, Size
 from .types import Literal, Direction, Coordinate, Axis, ObsGridArray, BoxSurface
 from .validators import assert_plane
 from .base import cached_property, Tidy3dBaseModel
@@ -24,6 +24,12 @@
 WARN_NUM_FREQS = 2000
 WARN_NUM_MODES = 100
 
+# Field projection windowing factor that determines field decay at the edges of surface field
+# projection monitors. A value of 15 leads to a decay of < 1e-3x in field amplitude.
+# This number relates directly to the standard deviation of the Gaussian function which is used
+# for windowing the monitor.
+WINDOW_FACTOR = 15
+
 
 class Monitor(AbstractMonitor):
     """Abstract base class for monitors."""
@@ -666,6 +672,48 @@ class AbstractFieldProjectionMonitor(SurfaceIntegrationMonitor, FreqMonitor):
         "projecting the recorded near fields to the far field.",
     )
 
+    window_size: Tuple[pydantic.NonNegativeFloat, pydantic.NonNegativeFloat] = pydantic.Field(
+        (0, 0),
+        title="Spatial filtering window size",
+        description="Size of the transition region of the windowing function used to ensure that "
+        "the recorded near fields decay to zero near the edges of the monitor. "
+        "The two components refer to the two tangential directions associated with each surface. "
+        "For surfaces with the normal along ``x``, the two components are (``y``, ``z``). "
+        "For surfaces with the normal along ``y``, the two components are (``x``, ``z``). "
+        "For surfaces with the normal along ``z``, the two components are (``x``, ``y``). "
+        "Each value must be between 0 and 1, inclusive, and denotes the size of the transition "
+        "region over which fields are scaled to less than a thousandth of the original amplitude, "
+        "relative to half the size of the monitor in that direction. A value of 0 turns windowing "
+        "off in that direction, while a value of 1 indicates that the window will be applied to "
+        "the entire monitor in that direction. This field is applicable for surface monitors only, "
+        "and otherwise must remain (0, 0).",
+    )
+
+    @pydantic.validator("window_size", always=True)
+    def window_size_for_surface(cls, val, values):
+        """Ensures that windowing is applied for surface monitors only."""
+        size = values.get("size")
+        name = values.get("name")
+
+        if size.count(0.0) != 1:
+            if val != (0, 0):
+                raise ValidationError(
+                    f"A non-zero 'window_size' cannot be used for projection monitor '{name}'. "
+                    "Windowing can be applied only for surface projection monitors."
+                )
+        return val
+
+    @pydantic.validator("window_size", always=True)
+    def window_size_leq_one(cls, val, values):
+        """Ensures that each component of the window size is less than or equal to 1."""
+        name = values.get("name")
+        if val[0] > 1 or val[1] > 1:
+            raise ValidationError(
+                f"Each component of 'window_size' for monitor '{name}' "
+                "must be less than or equal to 1."
+            )
+        return val
+
     @property
     def projection_surfaces(self) -> Tuple[FieldProjectionSurface, ...]:
         """Surfaces of the monitor where near fields will be recorded for subsequent projection."""
@@ -691,6 +739,62 @@ def local_origin(self) -> Coordinate:
             return self.center
         return self.custom_origin
 
+    def window_parameters(self, custom_bounds: Bound = None) -> Tuple[Size, Coordinate, Coordinate]:
+        """Return the physical size of the window transition region based on the monitor's size
+        and optional custom bounds (useful in case the monitor has infinite dimensions). The window
+        size is returned in 3D. Also returns the coordinate where the transition region beings on
+        the minus and plus side of the monitor."""
+
+        window_size = [0, 0, 0]
+        window_minus = [0, 0, 0]
+        window_plus = [0, 0, 0]
+
+        # windowing is for surface monitors only
+        if self.size.count(0.0) != 1:
+            return window_size, window_minus, window_plus
+
+        _, plane_inds = self.pop_axis([0, 1, 2], axis=self.size.index(0.0))
+
+        for i, ind in enumerate(plane_inds):
+            if custom_bounds:
+                size = min(self.size[ind], custom_bounds[1][ind] - custom_bounds[0][ind])
+                bound_min = max(self.bounds[0][ind], custom_bounds[0][ind])
+                bound_max = min(self.bounds[1][ind], custom_bounds[1][ind])
+            else:
+                size = self.size[ind]
+                bound_min = self.bounds[0][ind]
+                bound_max = self.bounds[1][ind]
+
+            window_size[ind] = self.window_size[i] * size / 2
+            window_minus[ind] = bound_min + window_size[ind]
+            window_plus[ind] = bound_max - window_size[ind]
+
+        return window_size, window_minus, window_plus
+
+    @staticmethod
+    def window_function(
+        points: ArrayFloat1D,
+        window_size: Size,
+        window_minus: Coordinate,
+        window_plus: Coordinate,
+        dim: int,
+    ) -> ArrayFloat1D:
+        """Get the windowing function along a given direction for a given set of points."""
+        rising_window = np.exp(
+            -0.5
+            * WINDOW_FACTOR
+            * ((points[points < window_minus[dim]] - window_minus[dim]) / window_size[dim]) ** 2
+        )
+        falling_window = np.exp(
+            -0.5
+            * WINDOW_FACTOR
+            * ((points[points > window_plus[dim]] - window_plus[dim]) / window_size[dim]) ** 2
+        )
+        window_fn = np.ones_like(points)
+        window_fn[points < window_minus[dim]] = rising_window
+        window_fn[points > window_plus[dim]] = falling_window
+        return window_fn
+
 
 class FieldProjectionAngleMonitor(AbstractFieldProjectionMonitor):
     """:class:`Monitor` that samples electromagnetic near fields in the frequency domain