add ability to project fields through a manually specified material

shashwat-sh · shashwat-sh · commit e9fc81b70a34 · 2023-12-07T22:58:58.000-05:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Warning if nonlinear mediums are used in an `adjoint` simulation. In this case, the gradients will not be accurate, but may be approximately correct if the nonlinearity is weak.
 - Validator for surface field projection monitors that warns if projecting backwards relative to the monitor's normal direction.
 - Validator for field projection monitors when far field approximation is enabled but the projection distance is small relative to the near field domain.
+- Ability to manually specify a medium through which to project fields, when using field projection monitors.
 
 ### Changed
 - Credit cost for remote mode solver has been modified to be defined in advance based on the mode solver details. Previously, the cost was based on elapsed runtime. On average, there should be little difference in the cost.
diff --git a/tidy3d/components/field_projection.py b/tidy3d/components/field_projection.py
@@ -341,6 +341,7 @@ def _far_fields_for_surface(
         phi: ArrayLikeN2F,
         surface: FieldProjectionSurface,
         currents: xr.Dataset,
+        medium: MediumType,
     ):
         """Compute far fields at an angle in spherical coordinates
         for a given set of surface currents and observation angles.
@@ -358,13 +359,14 @@ def _far_fields_for_surface(
             :class:`FieldProjectionSurface` object to use as source of near field.
         currents : xarray.Dataset
             xarray Dataset containing surface currents associated with the surface monitor.
+        medium : :class:`.MediumType`
+            Background medium through which to project fields.
 
         Returns
         -------
         tuple(numpy.ndarray[float], ...)
             ``Er``, ``Etheta``, ``Ephi``, ``Hr``, ``Htheta``, ``Hphi`` for the given surface.
         """
-
         pts = [currents[name].values for name in ["x", "y", "z"]]
 
         try:
@@ -393,7 +395,7 @@ def _far_fields_for_surface(
 
         phase = [None] * 3
         propagation_factor = -1j * AbstractFieldProjectionData.wavenumber(
-            medium=self.medium, frequency=frequency
+            medium=medium, frequency=frequency
         )
 
         def integrate_for_one_theta(i_th: int):
@@ -448,7 +450,7 @@ def integrate_for_one_theta(i_th: int):
         # Lphi  (8.34b)
         Lphi = -M[0] * sin_phi[None, :] + M[1] * cos_phi[None, :]
 
-        eta = ETA_0 / np.sqrt(self.medium.eps_model(frequency))
+        eta = ETA_0 / np.sqrt(medium.eps_model(frequency))
 
         Etheta = -(Lphi + eta * Ntheta)
         Ephi = Ltheta - eta * Nphi
@@ -546,7 +548,8 @@ def _project_fields_angular(
             np.zeros((1, len(theta), len(phi), len(freqs)), dtype=complex) for _ in field_names
         ]
 
-        k = AbstractFieldProjectionData.wavenumber(medium=self.medium, frequency=freqs)
+        medium = monitor.medium if monitor.medium else self.medium
+        k = AbstractFieldProjectionData.wavenumber(medium=medium, frequency=freqs)
         phase = np.atleast_1d(
             AbstractFieldProjectionData.propagation_phase(dist=monitor.proj_distance, k=k)
         )
@@ -564,6 +567,7 @@ def _project_fields_angular(
                         phi=phi,
                         surface=surface,
                         currents=currents,
+                        medium=medium,
                     )
                     for field, _field in zip(fields, _fields):
                         field[..., idx_f] += _field * phase[idx_f]
@@ -580,7 +584,7 @@ def _project_fields_angular(
                 ):
                     _x, _y, _z = monitor.sph_2_car(monitor.proj_distance, _theta, _phi)
                     _fields = self._fields_for_surface_exact(
-                        x=_x, y=_y, z=_z, surface=surface, currents=currents
+                        x=_x, y=_y, z=_z, surface=surface, currents=currents, medium=medium
                     )
                     for field, _field in zip(fields, _fields):
                         field[0, i, j, :] += _field
@@ -591,7 +595,7 @@ def _project_fields_angular(
             for name, field in zip(field_names, fields)
         }
         return FieldProjectionAngleData(
-            monitor=monitor, projection_surfaces=self.surfaces, medium=self.medium, **fields
+            monitor=monitor, projection_surfaces=self.surfaces, medium=medium, **fields
         )
 
     def _project_fields_cartesian(
@@ -622,7 +626,8 @@ def _project_fields_cartesian(
             np.zeros((len(x), len(y), len(z), len(freqs)), dtype=complex) for _ in field_names
         ]
 
-        wavenumber = AbstractFieldProjectionData.wavenumber(medium=self.medium, frequency=freqs)
+        medium = monitor.medium if monitor.medium else self.medium
+        wavenumber = AbstractFieldProjectionData.wavenumber(medium=medium, frequency=freqs)
 
         # Zip together all combinations of observation points for better progress tracking
         iter_coords = [
@@ -655,12 +660,13 @@ def _project_fields_cartesian(
                             phi=phi,
                             surface=surface,
                             currents=currents,
+                            medium=medium,
                         )
                         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=_x, y=_y, z=_z, surface=surface, currents=currents
+                        x=_x, y=_y, z=_z, surface=surface, currents=currents, medium=medium
                     )
                     for field, _field in zip(fields, _fields):
                         field[i, j, k, :] += _field
@@ -671,7 +677,7 @@ def _project_fields_cartesian(
             for name, field in zip(field_names, fields)
         }
         return FieldProjectionCartesianData(
-            monitor=monitor, projection_surfaces=self.surfaces, medium=self.medium, **fields
+            monitor=monitor, projection_surfaces=self.surfaces, medium=medium, **fields
         )
 
     def _project_fields_kspace(
@@ -698,7 +704,8 @@ def _project_fields_kspace(
         field_names = ("Er", "Etheta", "Ephi", "Hr", "Htheta", "Hphi")
         fields = [np.zeros((len(ux), len(uy), 1, len(freqs)), dtype=complex) for _ in field_names]
 
-        k = AbstractFieldProjectionData.wavenumber(medium=self.medium, frequency=freqs)
+        medium = monitor.medium if monitor.medium else self.medium
+        k = AbstractFieldProjectionData.wavenumber(medium=medium, frequency=freqs)
         phase = np.atleast_1d(
             AbstractFieldProjectionData.propagation_phase(dist=monitor.proj_distance, k=k)
         )
@@ -726,14 +733,15 @@ def _project_fields_kspace(
                             phi=phi,
                             surface=surface,
                             currents=currents,
+                            medium=medium,
                         )
                         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=_x, y=_y, z=_z, surface=surface, currents=currents
+                        x=_x, y=_y, z=_z, surface=surface, currents=currents, medium=medium
                     )
                     for field, _field in zip(fields, _fields):
                         field[i, j, 0, :] += _field
@@ -749,7 +757,7 @@ def _project_fields_kspace(
             for name, field in zip(field_names, fields)
         }
         return FieldProjectionKSpaceData(
-            monitor=monitor, projection_surfaces=self.surfaces, medium=self.medium, **fields
+            monitor=monitor, projection_surfaces=self.surfaces, medium=medium, **fields
         )
 
     """Exact projections"""
@@ -761,6 +769,7 @@ def _fields_for_surface_exact(
         z: float,
         surface: FieldProjectionSurface,
         currents: xr.Dataset,
+        medium: MediumType,
     ):
         """Compute projected fields in spherical coordinates at a given projection point on a
         Cartesian grid for a given set of surface currents using the exact homogeneous medium
@@ -778,20 +787,21 @@ def _fields_for_surface_exact(
             :class:`FieldProjectionSurface` object to use as source of near field.
         currents : xarray.Dataset
             xarray Dataset containing surface currents associated with the surface monitor.
+        medium : :class:`.MediumType`
+            Background medium through which to project fields.
 
         Returns
         -------
         tuple(np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray)
             ``Er``, ``Etheta``, ``Ephi``, ``Hr``, ``Htheta``, ``Hphi`` projected fields for
             each frequency.
         """
-
         freqs = np.array(self.frequencies)
         i_omega = 1j * 2.0 * np.pi * freqs[None, None, None, :]
-        wavenumber = AbstractFieldProjectionData.wavenumber(frequency=freqs, medium=self.medium)
+        wavenumber = AbstractFieldProjectionData.wavenumber(frequency=freqs, medium=medium)
         wavenumber = wavenumber[None, None, None, :]  # add space dimensions
 
-        eps_complex = self.medium.eps_model(frequency=freqs)
+        eps_complex = medium.eps_model(frequency=freqs)
         epsilon = EPSILON_0 * eps_complex[None, None, None, :]
 
         # source points
diff --git a/tidy3d/components/monitor.py b/tidy3d/components/monitor.py
@@ -11,6 +11,7 @@
 from .base import cached_property, Tidy3dBaseModel
 from .mode import ModeSpec
 from .apodization import ApodizationSpec
+from .medium import MediumType
 from .viz import ARROW_COLOR_MONITOR, ARROW_ALPHA
 from ..constants import HERTZ, SECOND, MICROMETER, RADIAN, inf
 from ..exceptions import SetupError, ValidationError
@@ -685,6 +686,16 @@ class AbstractFieldProjectionMonitor(SurfaceIntegrationMonitor, FreqMonitor):
         "and otherwise must remain (0, 0).",
     )
 
+    medium: MediumType = pydantic.Field(
+        None,
+        title="Projection medium",
+        description="Medium through which to project fields. Generally, the fields should be "
+        "projected through the same medium as the one in which this monitor is placed, and "
+        "this is the default behavior when ``medium=None``. A custom ``medium`` can be useful "
+        "in some situations for advanced users, but we recommend trying to avoid using a "
+        "non-default ``medium``.",
+    )
+
     @pydantic.validator("window_size", always=True)
     def window_size_for_surface(cls, val, values):
         """Ensures that windowing is applied for surface monitors only."""
