Add translated_copy function to ElectromagneticFieldData

Author: caseyflex

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Advanced option `dist_type` that allows `LinearLumpedElement` to be distributed across grid cells in different ways. For example, restricting the network portion to a single cell and using PEC wires to connect to the desired location of the terminals. 
 - New `layer_refinement_specs` field in `GridSpec` that takes a list of `LayerRefinementSpec` for automatic mesh refinement and snapping in layered structures. Structure corners on the cross section perpendicular to layer thickness direction can be automatically identified. Mesh is automatically snapped and refined around those corners.
 - New field `drop_outside_sim` in `MeshOverrideStructure` to specify whether to drop an override structure if it is outside the simulation domain, but it overlaps with the simulation domain when projected to an axis.
+- Function `translated_copy` in `ElectromagneticFieldData` to facilitate computing overlaps between field data at different locations.
 
 ### Changed
 - The coordinate of snapping points in `GridSpec` can take value `None`, so that mesh can be selectively snapped only along certain dimensions.

tests/test_data/test_data_arrays.py

@@ -120,20 +120,24 @@
 def get_xyz(
     monitor: td.components.monitor.MonitorType, grid_key: str, symmetry: bool
 ) -> Tuple[List[float], List[float], List[float]]:
-    if symmetry:
-        grid = SIM_SYM.discretize_monitor(monitor)
+    sim = SIM_SYM if symmetry else SIM
+    grid = sim.discretize_monitor(monitor)
+    if monitor.colocate:
+        x, y, z = grid.boundaries.to_list
+    else:
         x, y, z = grid[grid_key].to_list
+    if symmetry:
         x = [_x for _x in x if _x >= 0]
         y = [_y for _y in y if _y >= 0]
         z = [_z for _z in z if _z >= 0]
-    else:
-        grid = SIM.discretize_monitor(monitor)
-        x, y, z = grid[grid_key].to_list
     return x, y, z
 
 
-def make_scalar_field_data_array(grid_key: str, symmetry=True):
-    XS, YS, ZS = get_xyz(FIELD_MONITOR, grid_key, symmetry)
+def make_scalar_field_data_array(grid_key: str, symmetry=True, colocate: bool = None):
+    monitor = FIELD_MONITOR
+    if colocate is not None:
+        monitor = monitor.updated_copy(colocate=colocate)
+    XS, YS, ZS = get_xyz(monitor, grid_key, symmetry)
     values = (1 + 1j) * np.random.random((len(XS), len(YS), len(ZS), len(FS)))
     return td.ScalarFieldDataArray(values, coords=dict(x=XS, y=YS, z=ZS, f=FS))
 
@@ -144,8 +148,11 @@ def make_scalar_field_time_data_array(grid_key: str, symmetry=True):
     return td.ScalarFieldTimeDataArray(values, coords=dict(x=XS, y=YS, z=ZS, t=TS))
 
 
-def make_scalar_mode_field_data_array(grid_key: str, symmetry=True):
-    XS, YS, ZS = get_xyz(MODE_MONITOR_WITH_FIELDS, grid_key, symmetry)
+def make_scalar_mode_field_data_array(grid_key: str, symmetry=True, colocate: bool = None):
+    monitor = MODE_MONITOR_WITH_FIELDS
+    if colocate is not None:
+        monitor = monitor.updated_copy(colocate=colocate)
+    XS, YS, ZS = get_xyz(monitor, grid_key, symmetry)
     values = (1 + 0.1j) * np.random.random((len(XS), 1, len(ZS), len(FS), len(MODE_INDICES)))
 
     return td.ScalarModeFieldDataArray(

tests/test_data/test_monitor_data.py

@@ -168,9 +168,9 @@ def make_permittivity_data(symmetry: bool = True):
     sim = SIM_SYM if symmetry else SIM
     return PermittivityData(
         monitor=PERMITTIVITY_MONITOR,
-        eps_xx=make_scalar_field_data_array("Ex", symmetry),
-        eps_yy=make_scalar_field_data_array("Ey", symmetry),
-        eps_zz=make_scalar_field_data_array("Ez", symmetry),
+        eps_xx=make_scalar_field_data_array("Ex", symmetry, colocate=False),
+        eps_yy=make_scalar_field_data_array("Ey", symmetry, colocate=False),
+        eps_zz=make_scalar_field_data_array("Ez", symmetry, colocate=False),
         symmetry=sim.symmetry,
         symmetry_center=sim.center,
         grid_expanded=sim.discretize_monitor(PERMITTIVITY_MONITOR),
@@ -234,7 +234,7 @@ def test_field_data():
     # Compute flux as dot product with itself
     flux2 = np.abs(data_2d.dot(data_2d))
     # Assert result is the same
-    assert np.all(flux1 == flux2)
+    assert np.allclose(flux1, flux2)
 
 
 def test_field_data_to_source():
@@ -267,7 +267,7 @@ def test_mode_solver_data():
     # Compute flux as dot product with itself
     flux2 = np.abs(data.dot(data))
     # Assert result is the same
-    assert np.all(flux1 == flux2)
+    assert np.allclose(flux1, flux2)
     # Compute dot product with a field data
     field_data = make_field_data_2d()
     dot = data.dot(field_data)
@@ -611,6 +611,91 @@ def isel(data, freqs):
     assert len(dot.f) == 2
 
 
+def test_translated_copy():
+    mode_data = make_mode_solver_data()
+    field_data = make_field_data_2d()
+
+    vector = (1, 0, 0)
+    mode_data_translated = mode_data.translated_copy(vector=vector)
+    field_data_translated = field_data.translated_copy(vector=vector)
+
+    field1 = mode_data.symmetry_expanded_copy.Ex.isel(mode_index=0, f=0)
+    field2 = mode_data_translated.symmetry_expanded_copy.Ex.isel(mode_index=0, f=0)
+
+    atol = 1e-10
+
+    assert np.allclose(field1.data, field2.data)
+
+    assert np.allclose(
+        mode_data.dot(mode_data), mode_data_translated.dot(mode_data_translated), atol=atol
+    )
+    assert np.allclose(
+        mode_data.outer_dot(mode_data),
+        mode_data_translated.outer_dot(mode_data_translated),
+        atol=atol,
+    )
+    assert np.allclose(
+        mode_data.dot(field_data), mode_data_translated.dot(field_data_translated), atol=atol
+    )
+    assert np.allclose(
+        mode_data.outer_dot(field_data),
+        mode_data_translated.outer_dot(field_data_translated),
+        atol=atol,
+    )
+    assert np.allclose(
+        field_data.dot(mode_data), field_data_translated.dot(mode_data_translated), atol=atol
+    )
+    assert np.allclose(
+        field_data.outer_dot(mode_data),
+        field_data_translated.outer_dot(mode_data_translated),
+        atol=atol,
+    )
+    assert np.allclose(
+        field_data.dot(field_data), field_data_translated.dot(field_data_translated), atol=atol
+    )
+    assert np.allclose(
+        field_data.outer_dot(field_data),
+        field_data_translated.outer_dot(field_data_translated),
+        atol=atol,
+    )
+
+    assert np.allclose(
+        mode_data.dot(mode_data),
+        mode_data_translated.translated_copy(vector=[-v for v in vector]).dot(mode_data),
+        atol=atol,
+    )
+
+    assert np.allclose(
+        mode_data.outer_dot(mode_data),
+        mode_data_translated.translated_copy(vector=[-v for v in vector]).outer_dot(mode_data),
+        atol=atol,
+    )
+
+    # test warning for mismatch between monitor and field colocation
+    # monitor colocated, data colocated
+    with AssertLogLevel(None):
+        _ = mode_data.symmetry_expanded_copy
+    monitor = mode_data.monitor.updated_copy(colocate=False)
+    grid_expanded = SIM_SYM.discretize_monitor(monitor)
+    mode_data_warn1 = mode_data.updated_copy(monitor=monitor, grid_expanded=grid_expanded)
+    # monitor not colocated, data colocated
+    with AssertLogLevel("WARNING", contains_str="Interpolating"):
+        _ = mode_data_warn1.symmetry_expanded_copy
+    field_kwargs = {}
+    for key in mode_data.field_components.keys():
+        field_kwargs[key] = make_scalar_mode_field_data_array(key, colocate=False)
+    mode_data_warn2 = mode_data.updated_copy(**field_kwargs)
+    # monitor colocated, data not colocated
+    with AssertLogLevel("WARNING", contains_str="Interpolating"):
+        _ = mode_data_warn2.symmetry_expanded_copy
+    # neither colocated
+    mode_data_uncolocated = mode_data_warn2.updated_copy(
+        monitor=monitor, grid_expanded=grid_expanded
+    )
+    with AssertLogLevel(None):
+        _ = mode_data_uncolocated.symmetry_expanded_copy
+
+
 @pytest.mark.parametrize("phase_shift", np.linspace(0, 2 * np.pi, 10))
 def test_field_data_phase(phase_shift):
     def get_combined_phase(data):

tests/test_plugins/test_mode_solver.py

@@ -1202,3 +1202,45 @@ def test_gauge_robustness():
 
     array[1, -1] = np.nan
     assert ModeSolver._weighted_coord_max(array, ij, ij) == (0, 0)
+
+
+def test_translated_dot():
+    sim_size = (5, 5, 5)
+    lambda0 = 1.55
+    freq0 = td.C_0 / lambda0
+    si = td.material_library["cSi"]["Li1993_293K"]
+    sio2 = td.material_library["SiO2"]["Horiba"]
+    wg = td.Structure(geometry=td.Box(size=(0.22, 0.5, td.inf)), medium=si)
+    mode_spec = td.ModeSpec(num_modes=3)
+    grid_spec = td.GridSpec.auto(wavelength=lambda0, min_steps_per_wvl=20)
+
+    sim = td.Simulation(
+        size=sim_size, medium=sio2, structures=[wg], grid_spec=grid_spec, run_time=1e-30
+    )
+    mode_plane = td.Box(size=(3, 3, 0))
+    mode_solver = ModeSolver(simulation=sim, plane=mode_plane, mode_spec=mode_spec, freqs=[freq0])
+
+    data = mode_solver.data_raw
+
+    # now create a translated copy
+    vector = (0.5, 0, 0)
+    mode_solver2 = mode_solver.updated_copy(center=vector, path="simulation/structures/0/geometry")
+
+    data2 = mode_solver2.data_raw
+
+    # self-overlaps are close to 1, others are close to 0
+    atol = 1e-2
+
+    # just make sure the mode overlaps in the translated waveguide are the same
+    assert np.allclose(data.dot(data), data2.dot(data2), atol=atol)
+    assert np.allclose(data.outer_dot(data), data2.outer_dot(data2), atol=atol)
+
+    # now translate the data, and check that its overlaps with the modes
+    # of the translated waveguide agree with the self-overlaps of those modes
+    data_translated = data.translated_copy(vector)
+
+    assert np.allclose(data2.dot(data_translated), data2.dot(data2), atol=atol)
+    assert np.allclose(data_translated.dot(data2), data2.dot(data2), atol=atol)
+
+    assert np.allclose(data2.outer_dot(data_translated), data2.outer_dot(data2), atol=atol)
+    assert np.allclose(data_translated.outer_dot(data2), data2.outer_dot(data2), atol=atol)

tidy3d/components/data/monitor_data.py

@@ -246,6 +246,7 @@ def _symmetry_update_dict(self) -> Dict:
         """Dictionary of data fields to create data with expanded symmetry."""
 
         update_dict = {}
+        warn_interp = False
         for field_name, scalar_data in self.field_components.items():
             eigenval_fn = self.symmetry_eigenvalues[field_name]
 
@@ -273,8 +274,44 @@ def _symmetry_update_dict(self) -> Dict:
                 # Interpolate. There generally shouldn't be values out of bounds except potentially
                 # when handling modes, in which case they should be at the boundary and close to 0.
 
-                scalar_data = scalar_data.sel(**{dim_name: coords_interp}, method="nearest")
-                scalar_data = scalar_data.assign_coords({dim_name: coords})
+                # using sel vs interp is faster, and should always be fine
+                # if the data is set up correctly such that its colocation
+                # matches the monitor colocation settings. If these do not match,
+                # then we need to interpolate, which is slower.
+                use_sel = (
+                    len(scalar_data.coords[dim_name]) == 1
+                    or coords[-1] in scalar_data.coords[dim_name]
+                )
+                if use_sel:
+                    scalar_data = scalar_data.sel(**{dim_name: coords_interp}, method="nearest")
+                    scalar_data = scalar_data.assign_coords({dim_name: coords})
+                else:
+                    warn_interp = True
+                    no_flip_inds = np.where(coords >= sym_loc)[0]
+                    scalar_data_arrays = []
+                    if len(scalar_data.coords[dim_name]) == 1:
+                        scalar_data = scalar_data.sel(**{dim_name: coords_interp}, method="nearest")
+                    else:
+                        if len(flip_inds) > 0:
+                            scalar_data_flip = scalar_data.interp(
+                                **{dim_name: coords_interp[flip_inds][::-1]},
+                                method="linear",
+                                kwargs={"fill_value": "extrapolate"},
+                                assume_sorted=True,
+                            ).isel({dim_name: slice(None, None, -1)})
+                            scalar_data_flip = scalar_data_flip.assign_coords(
+                                {dim_name: coords[flip_inds]}
+                            )
+                            scalar_data_arrays.append(scalar_data_flip)
+                        if len(no_flip_inds) > 0:
+                            scalar_data_no_flip = scalar_data.interp(
+                                **{dim_name: coords_interp[no_flip_inds]},
+                                method="linear",
+                                kwargs={"fill_value": "extrapolate"},
+                                assume_sorted=True,
+                            )
+                            scalar_data_arrays.append(scalar_data_no_flip)
+                        scalar_data = xr.concat(scalar_data_arrays, dim=dim_name)
 
                 # apply the symmetry eigenvalue (if defined) to the flipped values
                 if eigenval_fn is not None:
@@ -288,6 +325,13 @@ def _symmetry_update_dict(self) -> Dict:
 
         update_dict.update({"symmetry": (0, 0, 0), "symmetry_center": None})
 
+        if warn_interp:
+            log.warning(
+                "Interpolating 'ElectromagneticFieldData'. This may be due to "
+                "mismatch between monitor colocation and data colocation, "
+                "and can lead to performance issues."
+            )
+
         return update_dict
 
     def at_coords(self, coords: Coords) -> xr.Dataset:
@@ -686,10 +730,11 @@ def dot(
         # Tangential fields for current and other field data
         fields_self = self._colocated_tangential_fields
 
-        fields_other = field_data._colocated_tangential_fields
         if conjugate:
             fields_self = {key: field.conj() for key, field in fields_self.items()}
 
+        fields_other = field_data._interpolated_tangential_fields(self._plane_grid_boundaries)
+
         # Drop size-1 dimensions in the other data
         fields_other = {key: field.squeeze(drop=True) for key, field in fields_other.items()}
 
@@ -703,6 +748,7 @@ def dot(
         # Integrate over plane
         d_area = self._diff_area
         integrand = (e_self_x_h_other - h_self_x_e_other) * d_area
+
         return ModeAmpsDataArray(0.25 * integrand.sum(dim=d_area.dims))
 
     def _interpolated_tangential_fields(self, coords: ArrayFloat2D) -> Dict[str, DataArray]:
@@ -941,6 +987,46 @@ def _check_fields_stored(self, components: list[str]):
                 "the 'fields' argument of a field monitor to select which components are stored."
             )
 
+    def translated_copy(self, vector: Coordinate) -> ElectromagneticFieldData:
+        """Create a copy of the :class:`.ElectromagneticFieldData` with fields translated
+        by the provided vector. Can be used together with ``dot`` or ``outer_dot``
+        to compute overlaps between field data at different locations.
+
+        Parameters
+        ----------
+        vector: :class:`.Coordinate`
+            Translation vector to apply to the field data.
+
+        Returns
+        -------
+        :class:`ElectromagneticFieldData`
+            A data object with the translated fields.
+        """
+        field_kwargs = {}
+        for key, val in self.field_components.items():
+            coords = dict(val.coords)
+            coords["x"] = coords["x"] + vector[0]
+            coords["y"] = coords["y"] + vector[1]
+            coords["z"] = coords["z"] + vector[2]
+            field_kwargs[key] = val.assign_coords(coords)
+
+        symmetry_center = self.symmetry_center
+        if symmetry_center is not None:
+            symmetry_center = tuple([x + y for (x, y) in zip(symmetry_center, vector)])
+        grid_expanded = self.grid_expanded._translated_copy(vector=vector)
+
+        monitor_center = tuple([x + y for (x, y) in zip(self.monitor.center, vector)])
+        monitor = self.monitor.updated_copy(center=monitor_center)
+
+        return self.updated_copy(
+            monitor=monitor,
+            symmetry=self.symmetry,
+            symmetry_center=symmetry_center,
+            grid_expanded=grid_expanded,
+            **self._grid_correction_dict,
+            **field_kwargs,
+        )
+
 
 class FieldData(FieldDataset, ElectromagneticFieldData):
     """

tidy3d/components/grid/grid.py

@@ -12,7 +12,7 @@
 from ..data.data_array import DataArray, ScalarFieldDataArray, SpatialDataArray
 from ..data.utils import UnstructuredGridDataset, UnstructuredGridDatasetType
 from ..geometry.base import Box
-from ..types import ArrayFloat1D, Axis, InterpMethod, Literal
+from ..types import ArrayFloat1D, Axis, Coordinate, InterpMethod, Literal
 
 # data type of one dimensional coordinate array.
 Coords1D = ArrayFloat1D
@@ -625,3 +625,13 @@ def snap_to_box_zero_dim(self, box: Box):
                     raise ValueError("Cannot snap grid to box center outside of grid domain.")
                 boundary_dict[dim] = np.array([center, center])
         return self.updated_copy(boundaries=Coords(**boundary_dict))
+
+    def _translated_copy(self, vector: Coordinate) -> Grid:
+        """Translate the grid by a vector. Not officially supported as resulting
+        grid may not be aligned with original Yee grid."""
+        boundaries = Coords(
+            x=self.boundaries.x + vector[0],
+            y=self.boundaries.y + vector[1],
+            z=self.boundaries.z + vector[2],
+        )
+        return self.updated_copy(boundaries=boundaries)