More autograd helpers

CHANGELOG.md

@@ -10,9 +10,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - 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.
+- Calling `.values` on `DataArray` no longer raises a `DeprecationWarning` during automatic differentiation
 
 ### Fixed
 - Regression in local field projection leading to incorrect results for `far_field_approx=True`.

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_functions.py

@@ -5,19 +5,22 @@
 import scipy.ndimage
 from autograd.test_util import check_grads
 from scipy.signal import convolve as convolve_sp
-from tidy3d.plugins.autograd.functions import (
+from tidy3d.plugins.autograd 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,
 )
@@ -55,7 +58,7 @@ def test_pad_val(self, rng, mode, size, pad_width, axis):
     def test_pad_grad(self, rng, mode, size, pad_width, axis):
         """Test gradients of padding function for various modes, sizes, pad widths, and axes."""
         x = rng.random(size)
-        check_grads(pad, modes=["fwd", "rev"], order=1)(x, pad_width, mode=mode, axis=axis)
+        check_grads(pad, modes=["fwd", "rev"], order=2)(x, pad_width, mode=mode, axis=axis)
 
 
 class TestPadExceptions:
@@ -136,7 +139,7 @@ def test_convolve_grad(self, rng, mode, padding, ary_size, kernel_size, square_k
             )
 
         x, k = self._ary_and_kernel(rng, ary_size, kernel_size, square_kernel)
-        check_grads(convolve, modes=["rev"], order=1)(x, k, padding=padding, mode=mode)
+        check_grads(convolve, modes=["rev"], order=2)(x, k, padding=padding, mode=mode)
 
 
 class TestConvolveExceptions:
@@ -194,7 +197,7 @@ def test_morphology_val_size(self, rng, op, sp_op, mode, ary_size, kernel_size):
     def test_morphology_val_grad(self, rng, op, sp_op, mode, ary_size, kernel_size):
         """Test gradients of morphological operations for various modes, array sizes, and kernel sizes."""
         x = rng.random(ary_size)
-        check_grads(op, modes=["rev"], order=1)(x, size=kernel_size, mode=mode)
+        check_grads(op, modes=["rev"], order=2)(x, size=kernel_size, mode=mode)
 
     @pytest.mark.parametrize(
         "full",
@@ -238,7 +241,7 @@ def test_morphology_val_structure_grad(
         ):
             """Test gradients of morphological operations for various kernel structures."""
             x, k = self._ary_and_kernel(rng, ary_size, kernel_size, full, square, flat)
-            check_grads(op, modes=["rev"], order=1)(x, size=kernel_size, mode=mode)
+            check_grads(op, modes=["rev"], order=2)(x, size=kernel_size, mode=mode)
 
 
 @pytest.mark.parametrize(
@@ -317,11 +320,9 @@ def test_interpn_val(self, rng, dim, method):
         result_scipy = scipy.interpolate.interpn(points, values, tuple(xi_grid), method=method)
         npt.assert_allclose(result_custom, result_scipy)
 
-    @pytest.mark.parametrize("order", [1, 2])
-    @pytest.mark.parametrize("mode", ["fwd", "rev"])
-    def test_interpn_values_grad(self, rng, dim, method, order, mode):
+    def test_interpn_values_grad(self, rng, dim, method):
         points, values, xi = self.generate_points_values_xi(rng, dim)
-        check_grads(lambda v: interpn(points, v, xi, method=method), modes=[mode], order=order)(
+        check_grads(lambda v: interpn(points, v, xi, method=method), modes=["fwd", "rev"], order=2)(
             values
         )
 
@@ -356,12 +357,10 @@ def test_trapz_val(self, rng, shape, axis, use_x):
         result_numpy = np.trapz(y, x=x, dx=dx, axis=axis)
         npt.assert_allclose(result_custom, result_numpy)
 
-    @pytest.mark.parametrize("order", [1, 2])
-    @pytest.mark.parametrize("mode", ["fwd", "rev"])
-    def test_trapz_grad(self, rng, shape, axis, use_x, order, mode):
+    def test_trapz_grad(self, rng, shape, axis, use_x):
         """Test gradients of trapz function for different array dimensions and integration axes."""
         y, x, dx = self.generate_y_x_dx(rng, shape, use_x)
-        check_grads(lambda y: trapz(y, x=x, dx=dx, axis=axis), modes=[mode], order=order)(y)
+        check_grads(lambda y: trapz(y, x=x, dx=dx, axis=axis), modes=["fwd", "rev"], order=2)(y)
 
 
 @pytest.mark.parametrize("shape", [(10,), (10, 10)])
@@ -381,10 +380,137 @@ def test_add_at_val(self, rng, shape, indices):
         result_numpy[indices] += y
         npt.assert_allclose(result_custom, result_numpy)
 
-    @pytest.mark.parametrize("order", [1, 2])
-    @pytest.mark.parametrize("mode", ["fwd", "rev"])
-    def test_add_at_grad(self, rng, shape, indices, order, mode):
+    def test_add_at_grad(self, rng, shape, indices):
         """Test gradients of add_at function for different array dimensions and indices."""
         x, y = self.generate_x_y(rng, shape, indices)
-        check_grads(lambda x: add_at(x, indices, y), modes=[mode], order=order)(x)
-        check_grads(lambda y: add_at(x, indices, y), modes=[mode], order=order)(y)
+        check_grads(lambda x: add_at(x, indices, y), modes=["fwd", "rev"], order=2)(x)
+        check_grads(lambda y: add_at(x, indices, y), modes=["fwd", "rev"], order=2)(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])
+class TestSmoothMax:
+    def test_smooth_max_values(self, rng, shape, tau, axis):
+        """Test `smooth_max` values for various shapes, tau, and axes."""
+
+        if axis == 1 and len(shape) == 1:
+            pytest.skip()
+
+        x = rng.uniform(-10, 10, size=shape)
+        result = smooth_max(x, tau=tau, axis=axis)
+
+        expected = np.max(x, axis=axis)
+        npt.assert_allclose(result, expected, atol=10 * tau)
+
+    def test_smooth_max_grad(self, rng, shape, tau, axis):
+        """Test gradients of `smooth_max` for various parameters."""
+
+        if axis == 1 and len(shape) == 1:
+            pytest.skip()
+
+        x = rng.uniform(-1, 1, size=shape)
+        func = lambda x: smooth_max(x, tau=tau, axis=axis)
+        check_grads(func, modes=["fwd", "rev"], order=2)(x)
+
+
+@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])
+class TestSmoothMin:
+    def test_smooth_min_values(self, rng, shape, tau, axis):
+        """Test `smooth_min` values for various shapes, tau, and axes."""
+
+        if axis == 1 and len(shape) == 1:
+            pytest.skip()
+
+        x = rng.uniform(-10, 10, size=shape)
+        result = smooth_min(x, tau=tau, axis=axis)
+
+        expected = np.min(x, axis=axis)
+        npt.assert_allclose(result, expected, atol=10 * tau)
+
+    def test_smooth_min_grad(self, rng, shape, tau, axis):
+        """Test gradients of `smooth_min` for various parameters."""
+
+        if axis == 1 and len(shape) == 1:
+            pytest.skip()
+
+        x = rng.uniform(-1, 1, size=shape)
+        func = lambda x: smooth_min(x, tau=tau, axis=axis)
+        check_grads(func, modes=["fwd", "rev"], order=2)(x)
+
+
+class TestLeastSquares:
+    @pytest.mark.parametrize(
+        "model, params_true, initial_guess, x, y",
+        [
+            (
+                lambda x, a, b: a * x + b,
+                np.array([2.0, -3.0]),
+                (0.0, 0.0),
+                np.linspace(0, 10, 50),
+                2.0 * np.linspace(0, 10, 50) - 3.0,
+            ),
+            (
+                lambda x, a, b, c: a * x**2 + b * x + c,
+                np.array([1.0, -2.0, 1.0]),
+                (0.0, 0.0, 0.0),
+                np.linspace(-5, 5, 100),
+                1.0 * np.linspace(-5, 5, 100) ** 2 - 2.0 * np.linspace(-5, 5, 100) + 1.0,
+            ),
+            (
+                lambda x, a, b: a * np.exp(b * x),
+                np.array([1.5, 0.5]),
+                (1.0, 0.0),
+                np.linspace(0, 2, 50),
+                1.5 * np.exp(0.5 * np.linspace(0, 2, 50)),
+            ),
+        ],
+    )
+    def test_least_squares(self, model, params_true, initial_guess, x, y):
+        """Test least_squares function with different models."""
+        params_estimated = least_squares(model, x, y, initial_guess)
+        npt.assert_allclose(params_estimated, params_true, rtol=1e-5)
+
+    def test_least_squares_with_noise(self, rng):
+        """Test least_squares function with noisy data."""
+
+        model = lambda x, a, b: a * x + b
+        a_true, b_true = -1.0, 4.0
+        params_true = np.array([a_true, b_true])
+        x = np.linspace(0, 10, 100)
+        noise = rng.normal(scale=0.1, size=x.shape)
+        y = a_true * x + b_true + noise
+        initial_guess = (0.0, 0.0)
+
+        params_estimated = least_squares(model, x, y, initial_guess)
+
+        npt.assert_allclose(params_estimated, params_true, rtol=1e-1)
+
+    def test_least_squares_no_convergence(self):
+        """Test that least_squares function raises an error when not converging."""
+
+        def constant_model(x, a):
+            return a
+
+        x = np.linspace(0, 10, 50)
+        y = 2.0 * x - 3.0  # Linear data
+        initial_guess = (0.0,)
+
+        with pytest.raises(np.linalg.LinAlgError):
+            least_squares(constant_model, x, y, initial_guess, max_iterations=10, tol=1e-12)
+
+    def test_least_squares_gradient(self):
+        """Test gradients of least_squares function with respect to parameters."""
+
+        def linear_model(x, a, b):
+            return a * x + b
+
+        x = np.linspace(0, 10, 50)
+        y = 2.0 * x - 3.0
+        initial_guess = (1.0, 0.0)
+
+        check_grads(
+            lambda params: least_squares(linear_model, x, y, params), modes=["fwd", "rev"], order=2
+        )(initial_guess)

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)