Differentiable local field projection

CHANGELOG.md

@@ -11,11 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Users can pass the `shading` argument in the `SimulationData.plot_field` to `Xarray.plot` method. When `shading='gouraud'`, the image is interpolated, producing a smoother visualization.
 - Users can manually specify the background medium for a structure to be used for geometry gradient calculations by supplying `Structure.background_permittivity`. This is useful when there are overlapping structures or structures embedded in other mediums.
 - Autograd functions can now be called directly on `DataArray` (e.g., `np.sum(data_array)`) in objective functions.
+- Automatic differentiation support for local field projections with `FieldProjectionAngleMonitor` and `FieldProjectionCartesianMonitor` using `FieldProjector.project_fields(far_field_monitor)`.
 
 ### Changed
 - Improved autograd tracer handling in `DataArray`, resulting in significant speedups for differentiation involving large monitors.
-
-### Changed
 - Triangulation of `PolySlab` polygons now supports polygons with collinear vertices.
 
 ### Fixed

tests/ruff.toml

@@ -10,4 +10,5 @@ extend-ignore = [
     "E731", # lambda assignment
     "F841", # unused local variable
     "S101", # asserts allowed in tests
+    "NPY201", # numpy 2.* compatibility check
 ]

tests/test_components/test_autograd.py

@@ -1002,6 +1002,109 @@ def objective(args):
     assert np.any(grads > 0)
 
 
+@pytest.mark.parametrize("far_field_approx", [True, False])
+@pytest.mark.parametrize("projection_type", ["angular", "cartesian"])
+@pytest.mark.parametrize("sim_2d", [True, False])
+class TestFieldProjection:
+    @staticmethod
+    def setup(far_field_approx, projection_type, sim_2d):
+        if sim_2d and not far_field_approx:
+            pytest.skip("Exact field projection not implemented for 2d simulations")
+
+        r_proj = 50 * WVL
+        monitor = td.FieldMonitor(
+            center=(0, SIM_BASE.size[1] / 2 - 0.1, 0),
+            size=(td.inf, 0, td.inf),
+            freqs=[FREQ0],
+            name="near_field",
+            colocate=False,
+        )
+
+        if projection_type == "angular":
+            theta_proj = np.linspace(np.pi / 10, np.pi - np.pi / 10, 2)
+            phi_proj = np.linspace(np.pi / 10, np.pi - np.pi / 10, 3)
+            monitor_far = td.FieldProjectionAngleMonitor(
+                center=monitor.center,
+                size=monitor.size,
+                freqs=monitor.freqs,
+                phi=tuple(phi_proj),
+                theta=tuple(theta_proj),
+                proj_distance=r_proj,
+                far_field_approx=far_field_approx,
+                name="far_field",
+            )
+        elif projection_type == "cartesian":
+            x_proj = np.linspace(-10, 10, 2)
+            y_proj = np.linspace(-10, 10, 3)
+            monitor_far = td.FieldProjectionCartesianMonitor(
+                center=monitor.center,
+                size=monitor.size,
+                freqs=monitor.freqs,
+                x=x_proj,
+                y=y_proj,
+                proj_axis=1,
+                proj_distance=r_proj,
+                far_field_approx=far_field_approx,
+                name="far_field",
+            )
+
+        sim = SIM_BASE.updated_copy(monitors=[monitor])
+
+        if sim_2d and IS_3D:
+            sim = sim.updated_copy(size=(0, *sim.size[1:]))
+
+        return sim, monitor_far
+
+    @staticmethod
+    def objective(sim_data, monitor_far):
+        projector = td.FieldProjector.from_near_field_monitors(
+            sim_data=sim_data,
+            near_monitors=[sim_data.simulation.monitors[0]],
+            normal_dirs=["+"],
+        )
+
+        projected_fields = projector.project_fields(monitor_far)
+
+        return projected_fields.power.sum().item()
+
+    def test_field_projection_grad_prop(
+        self, use_emulated_run, far_field_approx, projection_type, sim_2d
+    ):
+        """Tests whether field projection gradients are propagated through simulation.
+        x0 <-> structures <-> sim <-> run <-> fields <-> projection <-> objval
+        Does _not_ test gradient accuracy!
+        """
+        sim_base, monitor_far = self.setup(far_field_approx, projection_type, sim_2d)
+
+        def objective(args):
+            structures_traced_dict = make_structures(args)
+            structures = list(SIM_BASE.structures)
+            for structure_key in structure_keys_:
+                structures.append(structures_traced_dict[structure_key])
+
+            sim = sim_base.updated_copy(structures=structures)
+            sim_data = run(sim, task_name="field_projection_test")
+
+            return self.objective(sim_data, monitor_far)
+
+        grads = ag.grad(objective)(params0)
+        assert np.linalg.norm(grads) > 0
+
+    def test_field_projection_grads(
+        self, use_emulated_run, far_field_approx, projection_type, sim_2d
+    ):
+        """Tests projection gradient accuracy w.r.t. fields.
+        fields <-> projection <-> objval
+        """
+        sim_base, monitor_far = self.setup(far_field_approx, projection_type, sim_2d)
+
+        def objective(x0):
+            sim_data = run_emulated(sim_base, task_name="field_projection_test", x0=x0)
+            return self.objective(sim_data, monitor_far)
+
+        check_grads(objective, modes=["rev"], order=1)(1.0)
+
+
 def test_autograd_deepcopy():
     """make sure deepcopy works as expected in autograd."""
 

tests/test_components/test_field_projection.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pytest
 import tidy3d as td
+from tidy3d.components.field_projection import FieldProjector
 from tidy3d.exceptions import DataError
 
 MEDIUM = td.Medium(permittivity=3)
@@ -594,3 +595,25 @@ def test_2d_proj_clientside():
 
     for plane in planes:
         make_2d_proj(plane)
+
+
+@pytest.mark.parametrize(
+    "array, pts, axes, expected",
+    [
+        # 1D array, integrate over axis 0
+        (np.array([1, 2, 3]), np.array([0, 1, 2]), 0, 4.0),
+        # 2D array, integrate over axis 0
+        (np.array([[1, 2, 3], [4, 5, 6]]), np.array([0, 1]), 0, np.array([2.5, 3.5, 4.5])),
+        # 2D array, integrate over axis 1
+        (np.array([[1, 2], [3, 4], [5, 6]]), np.array([0, 1]), 1, np.array([1.5, 3.5, 5.5])),
+        # 3D array, integrate over axes 0 and 1
+        (np.ones((2, 2, 2)), [np.array([0, 1]), np.array([0, 1])], [0, 1], np.array([1.0, 1.0])),
+        # one element along integration axis but two points in pts
+        (np.array([[1, 1], [2, 2], [3, 3]]), np.array([0, 1]), 1, np.array([1.0, 2.0, 3.0])),
+        # 2D array of shape (1, 3), integrate over both axes
+        (np.array([[1, 2, 3]]), [np.array([0]), np.array([0, 1, 2])], [0, 1], 4.0),
+    ],
+)
+def test_trapezoid(array, pts, axes, expected):
+    result = FieldProjector.trapezoid(array, pts, axes)
+    assert np.allclose(result, expected)

tests/test_plugins/autograd/test_functions.py

@@ -6,6 +6,7 @@
 from autograd.test_util import check_grads
 from scipy.signal import convolve as convolve_sp
 from tidy3d.plugins.autograd.functions import (
+    add_at,
     convolve,
     grey_closing,
     grey_dilation,
@@ -18,6 +19,7 @@
     pad,
     rescale,
     threshold,
+    trapz,
 )
 from tidy3d.plugins.autograd.types import PaddingType
 
@@ -330,3 +332,59 @@ def test_invalid_method(self, rng):
         points, values, xi = TestInterpn.generate_points_values_xi(rng, 2)
         with pytest.raises(ValueError, match="interpolation method"):
             interpn(points, values, xi, method="invalid_method")
+
+
+@pytest.mark.parametrize("axis", [0, -1])
+@pytest.mark.parametrize("shape", [(10,), (10, 10)])
+@pytest.mark.parametrize("use_x", [True, False])
+class TestTrapz:
+    @staticmethod
+    def generate_y_x_dx(rng, shape, use_x):
+        y = rng.uniform(-1, 1, shape)
+        if use_x:
+            x = rng.random(shape)
+            dx = 1.0  # dx is not used when x is provided
+        else:
+            x = None
+            dx = rng.random() + 0.1  # ensure dx is not zero
+        return y, x, dx
+
+    def test_trapz_val(self, rng, shape, axis, use_x):
+        """Test trapz values against NumPy for different array dimensions and integration axes."""
+        y, x, dx = self.generate_y_x_dx(rng, shape, use_x)
+        result_custom = trapz(y, x=x, dx=dx, axis=axis)
+        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):
+        """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)
+
+
+@pytest.mark.parametrize("shape", [(10,), (10, 10)])
+@pytest.mark.parametrize("indices", [(0,), (slice(3, 8),)])
+class TestAddAt:
+    @staticmethod
+    def generate_x_y(rng, shape, indices):
+        x = rng.uniform(-1, 1, shape)
+        y = rng.uniform(-1, 1, x[tuple(indices)].shape)
+        return x, y
+
+    def test_add_at_val(self, rng, shape, indices):
+        """Test add_at values against NumPy for different array dimensions and indices."""
+        x, y = self.generate_x_y(rng, shape, indices)
+        result_custom = add_at(x, indices, y)
+        result_numpy = np.array(x)
+        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):
+        """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)

tests/utils.py

@@ -881,6 +881,8 @@ def run_emulated(simulation: td.Simulation, path=None, **kwargs) -> td.Simulatio
     """Emulates a simulation run."""
     from scipy.ndimage.filters import gaussian_filter
 
+    x = kwargs.get("x0", 1.0)
+
     def make_data(
         coords: dict, data_array_type: type, is_complex: bool = False
     ) -> td.components.data.data_array.DataArray:
@@ -891,7 +893,7 @@ def make_data(
 
         data = (1 + 0.5j) * data if is_complex else data
         data = gaussian_filter(data, sigma=1.0)  # smooth out the data a little so it isnt random
-        data_array = data_array_type(data, coords=coords)
+        data_array = data_array_type(x * data, coords=coords)
         return data_array
 
     def make_field_data(monitor: td.FieldMonitor) -> td.FieldData:

tidy3d/components/autograd/__init__.py

@@ -29,4 +29,6 @@
     "interpn",
     "split_list",
     "is_tidy_box",
+    "trapz",
+    "add_at",
 ]