feat[adjoint]: allow traceable kwargs in primitives interpolate_spline and add_at

Author: bzhangflex

CHANGELOG.md

@@ -29,6 +29,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Adjoint source frequency width is adjusted to decay sufficiently before zero frequency when possible to improve accuracy of simulation normalization when using custom current sources.
 - Change `VisualizationSpec` validator for checking validity of user specified colors to only issue a warning if matplotlib is not installed instead of an error.
 
+### Changed
+- `tidy3d.plugins.autograd.interpolate_spline()` and `tidy3d.plugins.autograd.add_at()` can now be called with keyword arguments during tracing.
+
 ## [2.8.4] - 2025-05-15
 
 ### Added

tests/test_plugins/autograd/primitives/test_interpolate.py

@@ -1,6 +1,7 @@
 import autograd.numpy as np
 import numpy.testing as npt
 import pytest
+from autograd import grad
 from autograd.test_util import check_grads
 from tidy3d.plugins.autograd import interpolate_spline
 
@@ -36,6 +37,22 @@ def test_interpolate_spline_grads(rng, order, num_points, endpoint_derivs, x_dis
     )
 
 
+def test_interpolate_spline_grads_kwargs(rng):
+    """Test interpolate_spline function can be called with kwargs."""
+    x = np.linspace(0, 1, 10)
+    y = rng.random(x.size)
+    # this should not error
+    grad(
+        lambda y_: interpolate_spline(
+            x_points=x,
+            y_points=y_,
+            num_points=10,
+            order=3,
+            endpoint_derivatives=(None, None),
+        )[1][0]
+    )(y)
+
+
 @pytest.mark.parametrize("order", [1, 2, 3])
 @pytest.mark.parametrize("x", [np.linspace(0, 1, 10), np.linspace(1, 0, 10)])
 def test_interpolate_spline_vals(rng, order, x):

tests/test_plugins/autograd/test_functions.py

@@ -3,6 +3,7 @@
 import pytest
 import scipy.interpolate
 import scipy.ndimage
+from autograd import grad
 from autograd.test_util import check_grads
 from scipy.signal import convolve as convolve_sp
 from tidy3d.plugins.autograd import (
@@ -387,6 +388,15 @@ def test_add_at_grad(self, rng, shape, indices):
         check_grads(lambda y: add_at(x, indices, y), modes=["fwd", "rev"], order=2)(y)
 
 
+def test_add_at_grad_kwargs(rng):
+    """Test add_at function for different array dimensions and indices, with kwargs."""
+    indices = (0,)
+    x = rng.uniform(-1, 1, (10,))
+    y = rng.uniform(-1, 1, x[tuple(indices)].shape)
+    # this should not error
+    grad(lambda y_: add_at(x=x, y=y_, indices_x=indices)[0])(y)
+
+
 @pytest.mark.parametrize("shape", [(5,), (5, 5), (5, 5, 5)])
 @pytest.mark.parametrize("tau", [1e-3, 1.0])
 @pytest.mark.parametrize("axis", [None, 0, 1, -1])

tidy3d/components/autograd/functions.py

@@ -198,6 +198,34 @@ def trapz(y: NDArray, x: NDArray = None, dx: float = 1.0, axis: int = -1) -> flo
 
 
 @primitive
+def _add_at(x: NDArray, indices_x: tuple, y: NDArray) -> NDArray:
+    """
+    Add values to specified indices of an array.
+
+    Autograd requires that arguments to primitives are passed in positionally.
+    ``add_at`` is the public-facing wrapper for this function,
+    which allows keyword arguments in case users pass in kwargs.
+    """
+    out = np.copy(x)  # Copy to preserve 'x' for gradient computation
+    out[tuple(indices_x)] += y
+    return out
+
+
+defvjp(
+    _add_at,
+    lambda ans, x, indices_x, y: unbroadcast_f(x, lambda g: g),
+    lambda ans, x, indices_x, y: lambda g: g[tuple(indices_x)],
+    argnums=(0, 2),
+)
+
+defjvp(
+    _add_at,
+    lambda g, ans, x, indices_x, y: broadcast(g, ans),
+    lambda g, ans, x, indices_x, y: _add_at(anp.zeros_like(ans), indices_x, g),
+    argnums=(0, 2),
+)
+
+
 def add_at(x: NDArray, indices_x: tuple, y: NDArray) -> NDArray:
     """
     Add values to specified indices of an array.
@@ -219,24 +247,7 @@ def add_at(x: NDArray, indices_x: tuple, y: NDArray) -> NDArray:
     np.ndarray
         The modified array with values added at the specified indices.
     """
-    out = np.copy(x)  # Copy to preserve 'x' for gradient computation
-    out[tuple(indices_x)] += y
-    return out
-
-
-defvjp(
-    add_at,
-    lambda ans, x, indices_x, y: unbroadcast_f(x, lambda g: g),
-    lambda ans, x, indices_x, y: lambda g: g[tuple(indices_x)],
-    argnums=(0, 2),
-)
-
-defjvp(
-    add_at,
-    lambda g, ans, x, indices_x, y: broadcast(g, ans),
-    lambda g, ans, x, indices_x, y: add_at(anp.zeros_like(ans), indices_x, g),
-    argnums=(0, 2),
-)
+    return _add_at(x, indices_x, y)
 
 
 __all__ = [

tidy3d/plugins/autograd/primitives/interpolate.py

@@ -682,7 +682,7 @@ def get_spline_derivatives_wrt_y(
 
 
 @primitive
-def interpolate_spline(
+def _interpolate_spline(
     x_points: NDArray,
     y_points: NDArray,
     num_points: int,
@@ -692,46 +692,9 @@ def interpolate_spline(
     """Primitive function to perform spline interpolation of a given order
     with optional endpoint derivatives.
 
-    Parameters
-    ----------
-    x_points : np.ndarray
-        X coordinates of the data points (must be strictly monotonic)
-    y_points : np.ndarray
-        Y coordinates of the data points
-    num_points : int
-        Number of points in the output interpolation
-    order : int
-        Order of the spline (1=linear, 2=quadratic, 3=cubic)
-    endpoint_derivatives : tuple[float, float] = (None, None)
-        Derivatives at the endpoints (left, right)
-        Note: For order=1 (linear), all endpoint derivatives are ignored.
-              For order=2 (quadratic), only the left endpoint derivative is used.
-              For order=3 (cubic), both endpoint derivatives are used if provided.
-
-    Returns
-    -------
-    tuple[np.ndarray, np.ndarray]
-        Tuple of (x_interpolated, y_interpolated) values
-
-    Examples
-    --------
-    >>> import numpy as np
-    >>> x = np.array([0, 1, 2])
-    >>> y = np.array([0, 1, 0])
-    >>> # Linear interpolation
-    >>> x_interp, y_interp = interpolate_spline(x, y, num_points=5, order=1)
-    >>> print(y_interp)
-    [0.   0.5  1.   0.5  0. ]
-
-    >>> # Quadratic interpolation with left endpoint derivative
-    >>> x_interp, y_interp = interpolate_spline(x, y, num_points=5, endpoint_derivatives=(0, None), order=2)
-    >>> print(np.round(y_interp, 3))
-    [0.    0.75  1.    0.5   0.  ]
-
-    >>> # Cubic interpolation with both endpoint derivatives
-    >>> x_interp, y_interp = interpolate_spline(x, y, num_points=5, endpoint_derivatives=(0, 0), order=3)
-    >>> print(np.round(y_interp, 3))
-    [0.    0.75  1.    0.75  0.  ]
+    Autograd requires that arguments to primitives are passed in positionally.
+    ``interpolate_spline`` is the public-facing wrapper for this function,
+    which allows keyword arguments in case users pass in kwargs.
     """
     if order not in (1, 2, 3):
         raise NotImplementedError(f"Spline order '{order}' not implemented.")
@@ -810,4 +773,64 @@ def vjp(g):
     return vjp
 
 
-defvjp(interpolate_spline, None, interpolate_spline_y_vjp)
+defvjp(_interpolate_spline, None, interpolate_spline_y_vjp)
+
+
+def interpolate_spline(
+    x_points: NDArray,
+    y_points: NDArray,
+    num_points: int,
+    order: int,
+    endpoint_derivatives: tuple[Optional[float], Optional[float]] = (None, None),
+) -> tuple[NDArray, NDArray]:
+    """Differentiable spline interpolation of a given order
+    with optional endpoint derivatives.
+
+    Parameters
+    ----------
+    x_points : np.ndarray
+        X coordinates of the data points (must be strictly monotonic)
+    y_points : np.ndarray
+        Y coordinates of the data points
+    num_points : int
+        Number of points in the output interpolation
+    order : int
+        Order of the spline (1=linear, 2=quadratic, 3=cubic)
+    endpoint_derivatives : tuple[float, float] = (None, None)
+        Derivatives at the endpoints (left, right)
+        Note: For order=1 (linear), all endpoint derivatives are ignored.
+              For order=2 (quadratic), only the left endpoint derivative is used.
+              For order=3 (cubic), both endpoint derivatives are used if provided.
+
+    Returns
+    -------
+    tuple[np.ndarray, np.ndarray]
+        Tuple of (x_interpolated, y_interpolated) values
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> x = np.array([0, 1, 2])
+    >>> y = np.array([0, 1, 0])
+    >>> # Linear interpolation
+    >>> x_interp, y_interp = interpolate_spline(x, y, num_points=5, order=1)
+    >>> print(y_interp)
+    [0.   0.5  1.   0.5  0. ]
+
+    >>> # Quadratic interpolation with left endpoint derivative
+    >>> x_interp, y_interp = interpolate_spline(x, y, num_points=5, endpoint_derivatives=(0, None), order=2)
+    >>> print(np.round(y_interp, 3))
+    [0.    0.75  1.    0.5   0.  ]
+
+    >>> # Cubic interpolation with both endpoint derivatives
+    >>> x_interp, y_interp = interpolate_spline(x, y, num_points=5, endpoint_derivatives=(0, 0), order=3)
+    >>> print(np.round(y_interp, 3))
+    [0.    0.75  1.    0.75  0.  ]
+    """
+    return _interpolate_spline(
+        x_points,
+        y_points,
+        num_points,
+        order,
+        endpoint_derivatives,
+    )