differentiable s-matrix calculation

Author: tylerflex

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Objective functions that involve running `tidy3d.plugins.smatrix.ComponentModeler` can be differentiated with autograd.
+
+
 ## [2.9.0rc1] - 2025-06-10
 
 ### Added

tests/test_plugins/smatrix/test_component_modeler.py

@@ -322,12 +322,13 @@ def _test_mappings(element_mappings, s_matrix):
             "mode_index_out": mode_index_out_from,
         }
 
-        coords_to = {
-            "port_in": port_in_to,
-            "port_out": port_out_to,
-            "mode_index_in": mode_index_in_to,
-            "mode_index_out": mode_index_out_to,
-        }
+
+        coords_to = dict(
+            port_in=port_in_to,
+            port_out=port_out_to,
+            mode_index_in=mode_index_in_to,
+            mode_index_out=mode_index_out_to,
+        )
 
         assert np.all(
             s_matrix.sel(**coords_to).values == mult_by * s_matrix.sel(**coords_from).values

tidy3d/plugins/autograd/README.md

@@ -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`.
 
 We currently have the following restrictions:
 

tidy3d/plugins/smatrix/component_modelers/base.py

@@ -6,7 +6,7 @@
 from abc import ABC, abstractmethod
 from typing import Optional, 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
@@ -24,6 +24,7 @@
 from tidy3d.plugins.smatrix.ports.wave import WavePort
 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 = "."
@@ -196,7 +197,21 @@ 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, 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/modal.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
@@ -298,6 +298,23 @@ def _internal_construct_smatrix(self, batch_data: BatchData) -> ModalPortDataArr
         }
         s_matrix = ModalPortDataArray(values, coords=coords)
 
+        def set_new_values(
+            values: np.ndarray,
+            new_values: np.ndarray,
+            port_name_in: str,
+            mode_index_in: int,
+            port_name_out: str,
+            mode_index_out: int,
+        ) -> np.ndarray:
+            """Replace ``values`` with ``new_values`` at indices given by dims in ``s_matrix```"""
+            port_in_pos = int(s_matrix.get_index("port_in").get_loc(port_name_in))
+            mode_in_pos = int(s_matrix.get_index("mode_index_in").get_loc(mode_index_in))
+            port_out_pos = int(s_matrix.get_index("port_out").get_loc(port_name_out))
+            mode_out_pos = int(s_matrix.get_index("mode_index_out").get_loc(mode_index_out))
+            mask = np.zeros_like(values)
+            mask[port_out_pos, port_in_pos, mode_out_pos, mode_in_pos, :] = 1
+            return np.where(mask, new_values, values)
+
         # loop through source ports
         for col_index in self.matrix_indices_run_sim:
             port_name_in, mode_index_in = col_index
@@ -317,14 +334,17 @@ 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
+
+                values = set_new_values(
+                    values=values,
+                    new_values=s_matrix_elements,
+                    port_name_in=port_name_in,
+                    mode_index_in=mode_index_in,
+                    port_name_out=port_name_out,
+                    mode_index_out=mode_index_out,
+                )
+
+        s_matrix = ModalPortDataArray(values, coords=coords)
 
         # element can be determined by user-defined mapping
         for (row_in, col_in), (row_out, col_out), mult_by in self.element_mappings:
@@ -339,12 +359,18 @@ 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
-            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
+
+            elements_from = mult_by * s_matrix.loc[coords_from].values
+
+            values = set_new_values(
+                values=values,
+                new_values=elements_from,
+                port_name_in=port_in_to,
+                mode_index_in=mode_index_in_to,
+                port_name_out=port_out_to,
+                mode_index_out=mode_index_out_to,
+            )
+
+            s_matrix = ModalPortDataArray(values, coords=coords)
 
         return s_matrix