feature[frontend]: data array method for extrapolating to endpoints outside of coordinate array that is autograd compatible

Author: Gregory Roberts

CHANGELOG.md

@@ -37,6 +37,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed `ElectromagneticFieldData.to_zbf()` to support single frequency monitors and apply the correct flattening order.
 - Bug in `TerminalComponentModeler.get_antenna_metrics_data` when port amplitudes are set to zero.
 - Added missing `solver_version` keyword argument to `run_async`.
+- Fixed `interpn` data array method to be compatible with extrapolation outside of data array coordinates. 
 
 ## [2.9.0] - 2025-08-04
 

tests/test_data/test_data_arrays.py

@@ -6,6 +6,7 @@
 
 import autograd as ag
 import autograd.numpy as np
+import numpy
 import pytest
 import xarray.testing as xrt
 from autograd.test_util import check_grads
@@ -516,3 +517,66 @@ def test_with_updated_data_shape():
 
     with pytest.raises(ValueError):
         arr2 = arr._with_updated_data(data=data, coords=coords)
+
+
+@pytest.mark.parametrize("method", ["nearest", "linear"])
+def test_interpn_with_extrapolation(rng, method):
+    """Checks that the extrapolation in `interpn` works as expected and that
+    it is autograd compatible."""
+    arr = td.SpatialDataArray(
+        rng.random((1, 3, 4, 5), dtype=np.float64),
+        coords={"x": [1], "y": [1, 2, 3], "z": [2, 3, 4, 5], "f": [0, 1, 2, 3, 4]},
+    )
+
+    for coord in arr.dims:
+        endpoints = [
+            arr.coords[coord].values[0] - 1.0,
+            arr.coords[coord].values[-1] + 1.0,
+        ]
+
+        method_coord = method if (len(arr.coords[coord]) > 1) else "nearest"
+
+        offset_interp_coords = arr.coords[coord].values + 0.5
+        coords_interp = {coord: [endpoints[0], *offset_interp_coords, endpoints[1]]}
+
+        extrapolate = arr._ag_interp(
+            coords_interp, method=method_coord, kwargs={"fill_value": "extrapolate"}
+        )
+
+        compare = arr.interp(
+            coords_interp, method=method_coord, kwargs={"fill_value": "extrapolate"}
+        )
+
+        numpy.testing.assert_allclose(
+            extrapolate.data, compare.data, err_msg="Expected data to be close!"
+        )
+
+    def f(params):
+        arr = td.SpatialDataArray(
+            params.reshape((1, 3, 4, 5)),
+            coords={"x": [1], "y": [1, 2, 3], "z": [2, 3, 4, 5], "f": [0, 1, 2, 3, 4]},
+        )
+
+        result = 0.0
+
+        for coord in arr.dims:
+            endpoints = [
+                arr.coords[coord].values[0] - 1.0,
+                arr.coords[coord].values[-1] + 1.0,
+            ]
+
+            method_coord = method if (len(arr.coords[coord]) > 1) else "nearest"
+
+            offset_interp_coords = arr.coords[coord].values + 0.5
+            coords_interp = {coord: [endpoints[0], *offset_interp_coords, endpoints[1]]}
+
+            interp_data = arr.interp(
+                coords_interp, method=method_coord, kwargs={"fill_value": "extrapolate"}
+            )
+
+            result += np.sum(interp_data.data)
+
+        return result
+
+    data = rng.random((1, 3, 4, 5), dtype=np.float64)
+    check_grads(f, order=1, modes=["rev"])(data)

tidy3d/components/autograd/functions.py

@@ -98,6 +98,7 @@ def interpn(
     xi: tuple[NDArray[np.float64], ...],
     *,
     method: InterpolationType = "linear",
+    **kwargs,
 ) -> NDArray[np.float64]:
     """Interpolate over a rectilinear grid in arbitrary dimensions.
 
@@ -137,7 +138,12 @@ def interpn(
     else:
         raise ValueError(f"Unsupported interpolation method: {method}")
 
-    itrp = RegularGridInterpolator(points, values, method=method)
+    if kwargs.get("fill_value") == "extrapolate":
+        itrp = RegularGridInterpolator(
+            points, values, method=method, fill_value=None, bounds_error=False
+        )
+    else:
+        itrp = RegularGridInterpolator(points, values, method=method)
 
     # Prepare the grid for interpolation
     # This step reshapes the grid, checks for NaNs and out-of-bounds values

tidy3d/components/data/data_array.py

@@ -467,12 +467,7 @@ def _ag_interp_func(var, indexes_coords, method, **kwargs):
             data = anp.transpose(var.data, combined_permutation)
             xi = anp.stack([anp.ravel(new_xi.data) for new_xi in new_x], axis=-1)
 
-            result = interpn(
-                [xn.data for xn in x],
-                data,
-                xi,
-                method=method,
-            )
+            result = interpn([xn.data for xn in x], data, xi, method=method, **kwargs)
 
             result = anp.moveaxis(result, 0, -1)
             result = anp.reshape(result, result.shape[:-1] + new_x[0].shape)

tidy3d/plugins/smatrix/component_modelers/terminal.py

@@ -4,7 +4,7 @@
 
 from typing import Literal, Optional, Union
 
-import numpy as np
+import autograd.numpy as np
 import pydantic.v1 as pd
 
 from tidy3d.components.base import cached_property