Move CustomMedium._interp to Coords.spatial_interp

CHANGELOG.md

@@ -8,13 +8,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Specification of spatial permittivity distribution of dispersive material using user-supplied data through `CustomPoleResidue`, `CustomSellmeier`, `CustomLorentz`, `CustomDebye`, and `CustomDrude` components.
 - `CustomAnisotropicMedium` where each component can take user-supplied data to define spatial permittivity distribution of non-dispersive or dispersive material.
+- `Coords.spatial_interp` to interpolate spatial data avoiding pitfalls of `xarray.interp` with linear interpolation in directions where there a single data point.
 
 ### Changed
 - Add `Medium2D` to full simulation in tests.
 - `DispersionFitter` and `StableDispersionFitter` unified in a single `DispersionFitter` interface.
 - `StableDispersionFitter` deprecated, with stable fitter now being run instead through `plugins.dispersion.web.run(DispersionFitter)`.
 - Removed validator from `CustomFieldSource` that ensured data spanned the source geometry. Now the current values are extrapolated outside of the supplied data ranges.
 - `CustomMedium` now take fields `permittivity` and `conductivity`. `eps_dataset` will be deprecated in v3.0.
+- Moved `CustomMedium._interp` to `Coords.spatial_interp` to be used by custom current sources.
 
 ### Fixed
 - Plotting 2D materials in `SimulationData.plot_field` and other circumstances.

tests/test_components/test_custom.py

@@ -239,8 +239,8 @@ def test_medium_interp():
 
     # more than one entries per each axis
     orig_data = make_scalar_data()
-    data_fit_nearest = CustomMedium._interp(orig_data, coord_interp, "nearest")
-    data_fit_linear = CustomMedium._interp(orig_data, coord_interp, "linear")
+    data_fit_nearest = coord_interp.spatial_interp(orig_data, "nearest")
+    data_fit_linear = coord_interp.spatial_interp(orig_data, "linear")
     assert np.allclose(data_fit_nearest.shape[:3], [len(f) for f in coord_interp.to_list])
     assert np.allclose(data_fit_linear.shape[:3], [len(f) for f in coord_interp.to_list])
     # maximal or minimal values shouldn't exceed that in the supplied data
@@ -254,8 +254,8 @@ def test_medium_interp():
     X = [1.1]
     data = np.random.random((Nx, Ny, Nz, 1))
     orig_data = ScalarFieldDataArray(data, coords=dict(x=X, y=Y, z=Z, f=freqs))
-    data_fit_nearest = CustomMedium._interp(orig_data, coord_interp, "nearest")
-    data_fit_linear = CustomMedium._interp(orig_data, coord_interp, "linear")
+    data_fit_nearest = coord_interp.spatial_interp(orig_data, "nearest")
+    data_fit_linear = coord_interp.spatial_interp(orig_data, "linear")
     assert np.allclose(data_fit_nearest.shape[:3], [len(f) for f in coord_interp.to_list])
     assert np.allclose(data_fit_linear.shape[:3], [len(f) for f in coord_interp.to_list])
     # maximal or minimal values shouldn't exceed that in the supplied data

tidy3d/components/grid/grid.py

@@ -1,12 +1,13 @@
 """Defines the FDTD grid."""
-
-from typing import Tuple, List
+from __future__ import annotations
+from typing import Tuple, List, Union
 
 import numpy as np
 import pydantic as pd
 
 from ..base import Tidy3dBaseModel, cached_property
-from ..types import ArrayFloat1D, Axis, TYPE_TAG_STR
+from ..data.data_array import SpatialDataArray, ScalarFieldDataArray
+from ..types import ArrayFloat1D, Axis, TYPE_TAG_STR, InterpMethod
 from ..geometry import Box
 
 from ...exceptions import SetupError
@@ -48,6 +49,73 @@ def to_list(self):
         """Return a list of the three Coord1D objects as numpy arrays."""
         return list(self.to_dict.values())
 
+    def spatial_interp(
+        self,
+        spatial_dataarray: Union[SpatialDataArray, ScalarFieldDataArray],
+        interp_method: InterpMethod,
+    ) -> Union[SpatialDataArray, ScalarFieldDataArray]:
+        """
+        Enhance xarray's ``.interp`` in two ways:
+            1) Check if the coordinate of the supplied data are in monotically increasing order.
+            If they are, apply the faster ``assume_sorted=True``.
+
+            2) For axes of single entry, instead of error, apply ``isel()`` along the axis.
+
+            3) When linear interp is applied, in the extrapolated region, filter values smaller
+            or larger than the original data's min(max) will be replaced with the original min(max).
+
+        Parameters
+        ----------
+        spatial_dataarray: Union[:class:`.SpatialDataArray`, :class:`.ScalarFieldDataArray`]
+            Supplied scalar dataset
+        interp_method : :class:`.InterpMethod`
+            Interpolation method.
+
+        Returns
+        -------
+        Union[:class:`.SpatialDataArray`, :class:`.ScalarFieldDataArray`]
+            The interpolated spatial dataset.
+        """
+
+        all_coords = "xyz"
+        is_single_entry = [spatial_dataarray.sizes[ax] == 1 for ax in all_coords]
+        interp_ax = [
+            ax for (ax, single_entry) in zip(all_coords, is_single_entry) if not single_entry
+        ]
+        isel_ax = [ax for ax in all_coords if ax not in interp_ax]
+
+        # apply isel for the axis containing single entry
+        if len(isel_ax) > 0:
+            spatial_dataarray = spatial_dataarray.isel(
+                {ax: [0] * len(self.to_dict[ax]) for ax in isel_ax}
+            )
+            spatial_dataarray = spatial_dataarray.assign_coords(
+                {ax: self.to_dict[ax] for ax in isel_ax}
+            )
+            if len(interp_ax) == 0:
+                return spatial_dataarray
+
+        # Apply interp for the rest
+        #   first check if it's sorted
+        is_sorted = all((np.all(np.diff(spatial_dataarray.coords[f]) > 0) for f in interp_ax))
+        interp_param = dict(
+            kwargs={"fill_value": "extrapolate"},
+            assume_sorted=is_sorted,
+            method=interp_method,
+        )
+        #   interpolation
+        interp_dataarray = spatial_dataarray.interp(
+            {ax: self.to_dict[ax] for ax in interp_ax},
+            **interp_param,
+        )
+
+        # filter any values larger/smaller than the original data's max/min.
+        max_val = max(spatial_dataarray.values.ravel())
+        min_val = min(spatial_dataarray.values.ravel())
+        interp_dataarray = interp_dataarray.where(interp_dataarray >= min_val, min_val)
+        interp_dataarray = interp_dataarray.where(interp_dataarray <= max_val, max_val)
+        return interp_dataarray
+
 
 class FieldGrid(Tidy3dBaseModel):
     """Holds the grid data for a single field.