feat: scalar_objective decorator to ensure autograd-compatibility of objective functions

Author: yaugenst-flex

CHANGELOG.md

@@ -11,6 +11,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Autograd support for local field projections using `FieldProjectionKSpaceMonitor`.
 - Function `components.geometry.utils.flatten_groups` now also flattens transformed groups when requested.
 - Differentiable `smooth_min`, `smooth_max`, and `least_squares` functions in `tidy3d.plugins.autograd`.
+- Differential operators `grad` and `value_and_grad` in `tidy3d.plugins.autograd` that behave similarly to the autograd operators but support auxiliary data via `aux_data=True` as well as differentiation w.r.t. `DataArray`.
+- `@scalar_objective` decorator in `tidy3d.plugins.autograd` that wraps objective functions to ensure they return a scalar value and performs additional checks to ensure compatibility of objective functions with autograd. Used by default in `tidy3d.plugins.autograd.value_and_grad` as well as `tidy3d.plugins.autograd.grad`.
+
 
 ### Changed
 - `CustomMedium` design regions require far less data when performing inverse design by reducing adjoint field monitor size for dims with one pixel.

tests/test_plugins/autograd/test_differential_operators.py

@@ -1,23 +1,66 @@
 import autograd.numpy as np
+import pytest
+from autograd import grad as grad_ag
 from autograd import value_and_grad as value_and_grad_ag
 from numpy.testing import assert_allclose
-from tidy3d.plugins.autograd.differential_operators import value_and_grad
+from tidy3d.components.data.data_array import DataArray
+from tidy3d.plugins.autograd import grad, value_and_grad
 
 
-def test_value_and_grad(rng):
+@pytest.mark.parametrize("argnum", [0, 1])
+@pytest.mark.parametrize("has_aux", [True, False])
+def test_grad(rng, argnum, has_aux):
     """Test the custom value_and_grad function against autograd's implementation"""
     x = rng.random(10)
+    y = rng.random(10)
     aux_val = "aux"
 
-    vg_fun = value_and_grad(lambda x: (np.linalg.norm(x), aux_val), has_aux=True)
-    vg_fun_ag = value_and_grad_ag(lambda x: np.linalg.norm(x))
+    def f(x, y):
+        ret = DataArray(x * y).sum()  # still DataArray
+        if has_aux:
+            return ret, aux_val
+        return ret
 
-    (v, g), aux = vg_fun(x)
-    v_ag, g_ag = vg_fun_ag(x)
+    grad_fun = grad(f, argnum=argnum, has_aux=has_aux)
+    grad_fun_ag = grad_ag(
+        lambda x, y: f(x, y)[0].item() if has_aux else f(x, y).item(), argnum=argnum
+    )
+
+    if has_aux:
+        g, aux = grad_fun(x, y)
+        assert aux == aux_val
+    else:
+        g = grad_fun(x, y)
+    g_ag = grad_fun_ag(x, y)
 
-    # assert that values and gradients match
-    assert_allclose(v, v_ag)
     assert_allclose(g, g_ag)
 
-    # check that auxiliary output is correctly returned
-    assert aux == aux_val
+
+@pytest.mark.parametrize("argnum", [0, 1])
+@pytest.mark.parametrize("has_aux", [True, False])
+def test_value_and_grad(rng, argnum, has_aux):
+    """Test the custom value_and_grad function against autograd's implementation"""
+    x = rng.random(10)
+    y = rng.random(10)
+    aux_val = "aux"
+
+    def f(x, y):
+        ret = DataArray(np.linalg.norm(x * y)).sum()  # still DataArray
+        if has_aux:
+            return ret, aux_val
+        return ret
+
+    vg_fun = value_and_grad(f, argnum=argnum, has_aux=has_aux)
+    vg_fun_ag = value_and_grad_ag(
+        lambda x, y: f(x, y)[0].item() if has_aux else f(x, y).item(), argnum=argnum
+    )
+
+    if has_aux:
+        (v, g), aux = vg_fun(x, y)
+        assert aux == aux_val
+    else:
+        v, g = vg_fun(x, y)
+    v_ag, g_ag = vg_fun_ag(x, y)
+
+    assert_allclose(v, v_ag)
+    assert_allclose(g, g_ag)

tests/test_plugins/autograd/test_utilities.py

@@ -1,7 +1,15 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
-from tidy3d.plugins.autograd.utilities import chain, get_kernel_size_px, make_kernel
+import xarray as xr
+from tidy3d.exceptions import Tidy3dError
+from tidy3d.plugins.autograd import (
+    chain,
+    get_kernel_size_px,
+    make_kernel,
+    scalar_objective,
+    value_and_grad,
+)
 
 
 @pytest.mark.parametrize("size", [(3, 3), (4, 4), (5, 5)])
@@ -104,3 +112,55 @@ def add_one(x):
         funcs = [add_one, "not_a_function"]
         with pytest.raises(TypeError, match="All elements in funcs must be callable"):
             chain(funcs)
+
+
+class TestScalarObjective:
+    def test_scalar_objective_no_aux(self):
+        """Test scalar_objective decorator without auxiliary data."""
+
+        @scalar_objective
+        def objective(x):
+            da = xr.DataArray(x)
+            return da.sum()
+
+        x = np.array([1.0, 2.0, 3.0])
+        result, grad = value_and_grad(objective)(x)
+        assert np.allclose(grad, np.ones_like(grad))
+        assert np.isclose(result, 6.0)
+
+    def test_scalar_objective_with_aux(self):
+        """Test scalar_objective decorator with auxiliary data."""
+
+        @scalar_objective(has_aux=True)
+        def objective(x):
+            da = xr.DataArray(x)
+            return da.sum(), "auxiliary_data"
+
+        x = np.array([1.0, 2.0, 3.0])
+        (result, grad), aux_data = value_and_grad(objective, has_aux=True)(x)
+        assert np.allclose(grad, np.ones_like(grad))
+        assert np.isclose(result, 6.0)
+        assert aux_data == "auxiliary_data"
+
+    def test_scalar_objective_invalid_return(self):
+        """Test scalar_objective decorator with invalid return value."""
+
+        @scalar_objective
+        def objective(x):
+            da = xr.DataArray(x)
+            return da  # Returning the array directly, not a scalar
+
+        x = np.array([1, 2, 3])
+        with pytest.raises(Tidy3dError, match="must be a scalar"):
+            objective(x)
+
+    def test_scalar_objective_float(self):
+        """Test scalar_objective decorator with a Python float return value."""
+
+        @scalar_objective
+        def objective(x):
+            return x**2
+
+        result, grad = value_and_grad(objective)(3.0)
+        assert np.isclose(grad, 6.0)
+        assert np.isclose(result, 9.0)

tidy3d/plugins/autograd/__init__.py

@@ -1,16 +1,22 @@
-from .differential_operators import value_and_grad
+from .differential_operators import grad, value_and_grad
 from .functions import (
+    add_at,
     convolve,
     grey_closing,
     grey_dilation,
     grey_erosion,
     grey_opening,
+    interpn,
+    least_squares,
     morphological_gradient,
     morphological_gradient_external,
     morphological_gradient_internal,
     pad,
     rescale,
+    smooth_max,
+    smooth_min,
     threshold,
+    trapz,
 )
 from .invdes import (
     CircularFilter,
@@ -28,7 +34,7 @@
     tanh_projection,
 )
 from .primitives import gaussian_filter
-from .utilities import chain, get_kernel_size_px, make_kernel
+from .utilities import chain, get_kernel_size_px, make_kernel, scalar_objective
 
 __all__ = [
     "CircularFilter",
@@ -60,4 +66,12 @@
     "rescale",
     "threshold",
     "value_and_grad",
+    "smooth_min",
+    "smooth_max",
+    "add_at",
+    "interpn",
+    "least_squares",
+    "grad",
+    "scalar_objective",
+    "trapz",
 ]

tidy3d/plugins/autograd/differential_operators.py

@@ -1,65 +1,72 @@
-from typing import Any, Callable
+from typing import Callable
 
-from autograd import value_and_grad as value_and_grad_ag
 from autograd.builtins import tuple as atuple
 from autograd.core import make_vjp
 from autograd.extend import vspace
 from autograd.wrap_util import unary_to_nary
 from numpy.typing import ArrayLike
 
+from .utilities import scalar_objective
+
 __all__ = [
     "value_and_grad",
+    "grad",
 ]
 
 
 @unary_to_nary
-def value_and_grad(
-    fun: Callable, x: ArrayLike, *, has_aux: bool = False
-) -> tuple[tuple[float, ArrayLike], Any]:
-    """Returns a function that returns both value and gradient.
-
-    This function wraps and extends autograd's 'value_and_grad' function by adding
-    support for auxiliary data.
+def grad(fun: Callable, x: ArrayLike, *, has_aux: bool = False) -> Callable:
+    """Returns a function that computes the gradient of `fun` with respect to `x`.
 
     Parameters
     ----------
     fun : Callable
-        The function to differentiate. Should take a single argument and return
-        a scalar value, or a tuple where the first element is a scalar value if has_aux is True.
+        The function to differentiate. Should return a scalar value, or a tuple of
+        (scalar_value, auxiliary_data) if `has_aux` is True.
     x : ArrayLike
-        The point at which to evaluate the function and its gradient.
+        The point at which to evaluate the gradient.
     has_aux : bool = False
-        If True, the function returns auxiliary data as the second element of a tuple.
+        If True, `fun` returns auxiliary data as the second element of a tuple.
 
     Returns
     -------
-    tuple[tuple[float, ArrayLike], Any]
-        A tuple containing:
-        - A tuple with the function value (float) and its gradient (ArrayLike)
-        - The auxiliary data returned by the function (if has_aux is True)
+    Callable
+        A function that takes the same arguments as `fun` and returns its gradient at `x`.
+    """
+    wrapped_fun = scalar_objective(fun, has_aux=has_aux)
+    vjp, result = make_vjp(lambda x: atuple(wrapped_fun(x)) if has_aux else wrapped_fun(x), x)
 
-    Raises
-    ------
-    TypeError
-        If the function does not return a scalar value.
+    if has_aux:
+        ans, aux = result
+        return vjp((vspace(ans).ones(), None)), aux
+    ans = result
+    return vjp(vspace(ans).ones())
 
-    Notes
-    -----
-    This function uses autograd for automatic differentiation. If the function
-    does not return auxiliary data (has_aux is False), it delegates to autograd's
-    value_and_grad function. The main extension is the support for auxiliary data
-    when has_aux is True.
-    """
-    if not has_aux:
-        return value_and_grad_ag(fun)(x)
 
-    vjp, (ans, aux) = make_vjp(lambda x: atuple(fun(x)), x)
+@unary_to_nary
+def value_and_grad(fun: Callable, x: ArrayLike, *, has_aux: bool = False) -> Callable:
+    """Returns a function that computes both the value and gradient of `fun` with respect to `x`.
+
+    Parameters
+    ----------
+    fun : Callable
+        The function to differentiate. Should return a scalar value, or a tuple of
+        (scalar_value, auxiliary_data) if `has_aux` is True.
+    x : ArrayLike
+        The point at which to evaluate the function and its gradient.
+    has_aux : bool = False
+        If True, `fun` returns auxiliary data as the second element of a tuple.
 
-    if not vspace(ans).size == 1:
-        raise TypeError(
-            "value_and_grad only applies to real scalar-output "
-            "functions. Try jacobian, elementwise_grad or "
-            "holomorphic_grad."
-        )
+    Returns
+    -------
+    Callable
+        A function that takes the same arguments as `fun` and returns its value and gradient at `x`.
+    """
+    wrapped_fun = scalar_objective(fun, has_aux=has_aux)
+    vjp, result = make_vjp(lambda x: atuple(wrapped_fun(x)) if has_aux else wrapped_fun(x), x)
 
-    return (ans, vjp((vspace(ans).ones(), None))), aux
+    if has_aux:
+        ans, aux = result
+        return (ans, vjp((vspace(ans).ones(), None))), aux
+    ans = result
+    return ans, vjp(vspace(ans).ones())

tidy3d/plugins/autograd/functions.py

@@ -670,7 +670,7 @@ def least_squares(
     >>> initial_guess = (0.0, 0.0)
     >>> params = least_squares(linear_model, x_data, y_data, initial_guess)
     >>> print(params)
-    array([2.0, -3.0])
+    [ 2. -3.]
     """
     params = np.array(initial_guess, dtype="f8")
     jac = jacobian(lambda params: func(x, *params))