differentiable s-matrix calculation

Author: tylerflex

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Selective simulation capabilities to `TerminalComponentModeler` via `run_only` and `element_mappings` fields, allowing users to run fewer simulations and extract only needed scattering matrix elements.
 - Added KLayout plugin, with DRC functionality for running design rule checks in `plugins.klayout.drc`. Supports running DRC on GDS files as well as `Geometry`, `Structure`, and `Simulation` objects.
 - Added "mil" and "in" (inch) units to `plot_length_units`.
+- Objective functions that involve running `tidy3d.plugins.smatrix.ComponentModeler` can be differentiated with autograd.
 
 ### Changed
 - Validate mode solver object for large number of grid points on the modal plane.

docs/notebooks

@@ -4,9 +4,11 @@
 
 from typing import Optional
 
-import numpy as np
+import autograd as ag
+import autograd.numpy as np
 import pytest
 import xarray.testing as xrt
+from autograd.test_util import check_grads
 
 import tidy3d as td
 from tidy3d.exceptions import DataError
@@ -468,3 +470,49 @@ def test_interp(method, scalar_index):
     xr_interp = data.interp(f=f)
     ag_interp = data._ag_interp(f=f)
     xrt.assert_allclose(xr_interp, ag_interp)
+
+
+def test_with_updated_data_grad():
+    """Check the ``DataArray.with_updated_data()`` method."""
+
+    arr = td.SpatialDataArray(
+        np.ones((2, 3, 4, 5), dtype=np.complex128),
+        coords={"x": [0, 1], "y": [1, 2, 3], "z": [2, 3, 4, 5], "w": [0, 1, 2, 3, 4]},
+    )
+
+    data = np.zeros((1, 1, 1, 5))
+
+    coords = {"x": 0, "y": 2, "z": 3}
+
+    arr2 = arr._with_updated_data(data=data, coords=coords)
+
+    data_expected = np.ones(arr.shape) + 0j
+    data_expected[0, 1, 1, :] = 0.0 + 0j
+    assert np.all(arr2.data == data_expected), "DataArray.with_updated_copy() failed"
+
+    def f(x):
+        arr2 = arr._with_updated_data(data=x, coords=coords)
+        return np.abs(np.sum(arr2.data))
+
+    # grad should just be all 1s because of sum, so check that this is true
+    g = ag.grad(f)(data)
+    assert np.all(g == np.ones_like(data))
+
+    check_grads(f, order=1, modes=["rev"])(data)
+
+
+def test_with_updated_data_shape():
+    """Check the ``DataArray.with_updated_data()`` method."""
+
+    arr = td.SpatialDataArray(
+        np.ones((2, 3, 4, 5), dtype=np.complex128),
+        coords={"x": [0, 1], "y": [1, 2, 3], "z": [2, 3, 4, 5], "w": [0, 1, 2, 3, 4]},
+    )
+
+    # wrong shape
+    data = np.zeros((1, 1, 1, 3))
+
+    coords = {"x": 0, "y": 2, "z": 3}
+
+    with pytest.raises(ValueError):
+        arr2 = arr._with_updated_data(data=data, coords=coords)

tests/test_data/test_data_arrays.py

@@ -489,6 +489,40 @@ def _ag_interp_func(var, indexes_coords, method, **kwargs):
                 result = result.transpose(*out_dims)
         return result
 
+    def _with_updated_data(self, data: np.ndarray, coords: dict[str, Any]) -> DataArray:
+        """Make copy of ``DataArray`` with ``data`` at specified ``coords``, autograd compatible
+
+        Constraints / Edge cases:
+            - `coords` must map to a specific value eg {x: '1'}, does not broadcast to arrays
+            - `data` will be reshaped to try to match `self.shape` except where `coords` present
+        """
+
+        # make mask
+        mask = xr.zeros_like(self, dtype=bool)
+        mask.loc[coords] = True
+
+        # reshape `data` to line up with `self.dims`, with shape of 1 along the selected axis
+        old_data = self.data
+        new_shape = list(old_data.shape)
+        for i, dim in enumerate(self.dims):
+            if dim in coords:
+                new_shape[i] = 1
+        try:
+            new_data = data.reshape(new_shape)
+        except ValueError as e:
+            raise ValueError(
+                "Couldn't reshape the supplied 'data' to update 'DataArray'. The provided data was "
+                f"of shape {data.shape} and tried to reshape to {new_shape}. If you encounter this "
+                "error please raise an issue on the tidy3d github repository with the context."
+            ) from e
+
+        # broadcast data to repeat data along the selected dimensions to match mask
+        new_data = new_data + np.zeros_like(old_data)
+
+        new_data = np.where(mask, new_data, old_data)
+
+        return self.copy(deep=True, data=new_data)
+
 
 class FreqDataArray(DataArray):
     """Frequency-domain array.

tidy3d/components/data/data_array.py

@@ -217,6 +217,7 @@ We also support the following high-level features:
 - We automatically determine the number of adjoint simulations to run from a given forward simulation to maintain gradient accuracy.
   Adjoint sources are automatically grouped by either frequency or spatial port (whichever yields fewer adjoint simulations), and all adjoint simulations are run in a single batch (applies to both `run` and `run_async`).
   The parameter `max_num_adjoint_per_fwd` (default `10`) prevents launching unexpectedly large numbers of adjoint simulations automatically.
+- Differentiation of objective functions involving the scattering matrix produced by `tidy3d.plugins.smatrix.ComponentModeler.run()` and `tidy3d.plugins.smatrix.TerminalComponentModeler.run()`.
 
 We currently have the following restrictions:
 

tidy3d/plugins/autograd/README.md

@@ -6,7 +6,7 @@
 from abc import ABC, abstractmethod
 from typing import Generic, Optional, TypeVar, Union, get_args
 
-import numpy as np
+import autograd.numpy as np
 import pydantic.v1 as pd
 
 from tidy3d.components.base import Tidy3dBaseModel, cached_property
@@ -23,12 +23,16 @@
 from tidy3d.plugins.smatrix.ports.modal import Port
 from tidy3d.plugins.smatrix.ports.rectangular_lumped import LumpedPort
 from tidy3d.plugins.smatrix.ports.wave import WavePort
+from tidy3d.web import run_async
 from tidy3d.web.api.container import Batch, BatchData
 
 # fwidth of gaussian pulse in units of central frequency
 FWIDTH_FRAC = 1.0 / 10
 DEFAULT_DATA_DIR = "."
 
+# whether to run gradient calculation for component modeler locally
+LOCAL_GRADIENT = False
+
 LumpedPortType = Union[LumpedPort, CoaxialLumpedPort]
 TerminalPortType = Union[LumpedPortType, WavePort]
 
@@ -251,7 +255,26 @@ def batch_path(self) -> str:
     @cached_property
     def batch_data(self) -> BatchData:
         """The :class:`.BatchData` associated with the simulations run for this component modeler."""
-        return self.batch.run(path_dir=self.path_dir)
+
+        # NOTE: uses run_async because Batch is not differentiable.
+        batch = self.batch
+        run_async_kwargs = batch.dict(
+            exclude={
+                "type",
+                "path_dir",
+                "attrs",
+                "solver_version",
+                "jobs_cached",
+                "num_workers",
+                "simulations",
+            }
+        )
+        return run_async(
+            batch.simulations,
+            **run_async_kwargs,
+            local_gradient=LOCAL_GRADIENT,
+            path_dir=self.path_dir,
+        )
 
     def get_path_dir(self, path_dir: str) -> None:
         """Check whether the supplied 'path_dir' matches the internal field value."""

tidy3d/plugins/smatrix/component_modelers/base.py

@@ -6,7 +6,7 @@
 
 from typing import Optional
 
-import numpy as np
+import autograd.numpy as np
 import pydantic.v1 as pd
 
 from tidy3d.components.base import cached_property
@@ -237,14 +237,15 @@ def _internal_construct_smatrix(self, batch_data: BatchData) -> ModalPortDataArr
                 )
                 source_norm = self._normalization_factor(port_in, sim_data)
                 s_matrix_elements = np.array(amp.data) / np.array(source_norm)
-                s_matrix.loc[
-                    {
-                        "port_in": port_name_in,
-                        "mode_index_in": mode_index_in,
-                        "port_out": port_name_out,
-                        "mode_index_out": mode_index_out,
-                    }
-                ] = s_matrix_elements
+
+                coords_set = {
+                    "port_in": port_name_in,
+                    "mode_index_in": mode_index_in,
+                    "port_out": port_name_out,
+                    "mode_index_out": mode_index_out,
+                }
+
+                s_matrix = s_matrix._with_updated_data(data=s_matrix_elements, coords=coords_set)
 
         # element can be determined by user-defined mapping
         for (row_in, col_in), (row_out, col_out), mult_by in self.element_mappings:
@@ -259,12 +260,14 @@ def _internal_construct_smatrix(self, batch_data: BatchData) -> ModalPortDataArr
 
             port_out_to, mode_index_out_to = row_out
             port_in_to, mode_index_in_to = col_out
+
+            elements_from = mult_by * s_matrix.loc[coords_from].values
             coords_to = {
                 "port_in": port_in_to,
                 "mode_index_in": mode_index_in_to,
                 "port_out": port_out_to,
                 "mode_index_out": mode_index_out_to,
             }
-            s_matrix.loc[coords_to] = mult_by * s_matrix.loc[coords_from].values
+            s_matrix = s_matrix._with_updated_data(data=elements_from, coords=coords_to)
 
         return s_matrix

tidy3d/plugins/smatrix/component_modelers/modal.py

@@ -291,9 +291,10 @@ def _internal_construct_smatrix(
             port, mode_index = self.network_dict[source_index]
             sim_data = batch_data[self._task_name(port=port, mode_index=mode_index)]
             a, b = self.compute_power_wave_amplitudes_at_each_port(port_impedances, sim_data)
-            indexer = {"f": a.f, "port_out": a.port, "port_in": source_index}
-            a_matrix.loc[indexer] = a
-            b_matrix.loc[indexer] = b
+
+            indexer = {"port_in": source_index}
+            a_matrix = a_matrix._with_updated_data(data=a.data, coords=indexer)
+            b_matrix = b_matrix._with_updated_data(data=b.data, coords=indexer)
 
         # If excitation is assumed ideal, a_matrix is assumed to be diagonal
         # and the explicit inverse can be avoided. When only a subset of excitations
@@ -315,7 +316,8 @@ def _internal_construct_smatrix(
                 "port_in": col_out,
                 "port_out": row_out,
             }
-            s_matrix.loc[coords_to] = mult_by * s_matrix.loc[coords_from].values
+            data = mult_by * s_matrix.loc[coords_from].data
+            s_matrix = s_matrix._with_updated_data(data=data, coords=coords_to)
 
         return s_matrix
 
@@ -408,8 +410,8 @@ def compute_power_wave_amplitudes_at_each_port(
             port, mode_index = self.network_dict[network_index]
             V_out, I_out = self.compute_port_VI(port, sim_data)
             indexer = {"port": network_index}
-            V_matrix.loc[indexer] = V_out
-            I_matrix.loc[indexer] = I_out
+            V_matrix = V_matrix._with_updated_data(data=V_out.data, coords=indexer)
+            I_matrix = I_matrix._with_updated_data(data=I_out.data, coords=indexer)
 
         V_numpy = V_matrix.values
         I_numpy = I_matrix.values
@@ -577,16 +579,16 @@ def _port_reference_impedances(self, batch_data: BatchData) -> PortDataArray:
             )
         for network_index in self.matrix_indices_monitor:
             port, mode_index = self.network_dict[network_index]
+            indexer = {"port": network_index}
             if isinstance(port, WavePort):
                 # WavePorts have a port impedance calculated from its associated modal field distribution
                 # and is frequency dependent.
-                impedances = port.compute_port_impedance(sim_data).values
-                port_impedances.loc[{"port": network_index}] = impedances.squeeze()
+                data = port.compute_port_impedance(sim_data).data
+                port_impedances = port_impedances._with_updated_data(data=data, coords=indexer)
             else:
                 # LumpedPorts have a constant reference impedance
-                port_impedances.loc[{"port": network_index}] = np.full(
-                    len(self.freqs), port.impedance
-                )
+                data = np.full(len(self.freqs), port.impedance)
+                port_impedances = port_impedances._with_updated_data(data=data, coords=indexer)
 
         port_impedances = TerminalComponentModeler._set_port_data_array_attributes(port_impedances)
         return port_impedances