fix(adjoint): disallow autograd with element_mappings for a ComponentModeler

Author: Gregory Roberts

tests/test_plugins/smatrix/test_component_modeler_autograd.py

@@ -232,6 +232,35 @@ def objective(scale: float) -> float:
     assert not np.isclose(g, 0.0)
 
 
+def test_component_modeler_autograd_error_with_element_mappings(
+    patch_web_autograd_emulator, tmp_path
+):
+    """Verify that we get the expected error when running a component modeler with `element_mappings` in autograd."""
+    td.config.logging_level = "ERROR"
+    td.config.log_suppression = True
+
+    def objective(scale: float) -> float:
+        modeler = build_modal_modeler(scale)
+        # set up symmetry (port names only; directions are handled internally)
+        left = ("p1", 0)
+        right = ("p2", 0)
+
+        element_mappings = [
+            ((left, right), (right, left), 1.0),
+        ]
+
+        modeler_with_mappings = modeler.updated_copy(element_mappings=element_mappings)
+        s = modeler_with_mappings.run(
+            path_dir=str(tmp_path),
+            verbose=False,
+            local_gradient=True,
+        )
+        return anp.real(anp.sum(s.data))
+
+    g = ag.grad(objective)(1.0)
+    assert np.isfinite(g)
+
+
 def test_component_modeler_autograd_tracing_modeler_run(patch_web_autograd_emulator, tmp_path):
     td.config.logging_level = "ERROR"
     td.config.log_suppression = True

tidy3d/plugins/smatrix/component_modelers/base.py

@@ -227,7 +227,8 @@ def run(
     ):
         log.warning(
             "'ComponentModeler.run()' is deprecated and will be removed in a future release. "
-            "Use 'web.run(modeler)' instead.",
+            "Use web.run(modeler) instead. 'web.run' returns a 'ComponentModelerData' object; "
+            "get the scattering matrix via 'data.smatrix()'.",
             log_once=True,
         )
         from tidy3d.plugins.smatrix.run import _run_local

tidy3d/plugins/smatrix/run.py

@@ -4,6 +4,7 @@
 
 from tidy3d.components.base import Tidy3dBaseModel
 from tidy3d.components.data.index import SimulationDataMap
+from tidy3d.log import log
 from tidy3d.plugins.smatrix.component_modelers.modal import ModalComponentModeler
 from tidy3d.plugins.smatrix.component_modelers.terminal import TerminalComponentModeler
 from tidy3d.plugins.smatrix.component_modelers.types import ComponentModelerType
@@ -179,15 +180,24 @@ def _run_local(
     # autograd path if any sim is valid for autograd
     from tidy3d.web.api.autograd import autograd as web_ag
 
-    sims = getattr(modeler, "sim_dict", None) or {}
+    sims = modeler.sim_dict
     if any(web_ag.is_valid_for_autograd(sim) for sim in sims.values()):
+        if len(modeler.element_mappings) > 0:
+            log.warning(
+                "Element mappings are used to populate S-matrix values, but autograd gradients "
+                "are computed only for simulated elements. Gradients for mapped elements are not "
+                "included. For optimization with autograd, prefer enforcing symmetry in geometry/"
+                "objective functions and use 'run_only' to select unique sources.",
+                log_once=True,
+            )
+
         from tidy3d.web.api.autograd.autograd import _run_async
 
         kwargs.setdefault("folder_name", "default")
         kwargs.setdefault("simulation_type", "tidy3d_autograd_async")
         kwargs.setdefault("path_dir", path_dir)
 
-        sim_data_map = _run_async(simulations=modeler.sim_dict, **kwargs)
+        sim_data_map = _run_async(simulations=sims, **kwargs)
 
         return compose_modeler_data_from_batch_data(modeler=modeler, batch_data=sim_data_map)
 

tidy3d/web/api/autograd/autograd.py

@@ -23,6 +23,7 @@
 from tidy3d.components.autograd.derivative_utils import DerivativeInfo
 from tidy3d.components.data.data_array import DataArray
 from tidy3d.components.grid.grid_spec import GridSpec
+from tidy3d.components.types.workflow import WorkflowDataType, WorkflowType
 from tidy3d.exceptions import AdjointError
 from tidy3d.web.api.asynchronous import DEFAULT_DATA_DIR
 from tidy3d.web.api.asynchronous import run_async as run_async_webapi
@@ -94,7 +95,7 @@ def is_valid_for_autograd_async(simulations: dict[str, td.Simulation]) -> bool:
 
 
 def run(
-    simulation: typing.Any,
+    simulation: WorkflowType,
     task_name: str,
     folder_name: str = "default",
     path: str = "simulation_data.hdf5",
@@ -111,14 +112,14 @@ def run(
     reduce_simulation: typing.Literal["auto", True, False] = "auto",
     pay_type: typing.Union[PayType, str] = PayType.AUTO,
     priority: typing.Optional[int] = None,
-) -> td.SimulationData:
+) -> WorkflowDataType:
     """
     Submits a :class:`.Simulation` to server, starts running, monitors progress, downloads,
     and loads results as a :class:`.WorkflowDataType` object.
 
     Parameters
     ----------
-    simulation : Union[:class:`.Simulation`, :class:`.HeatSimulation`, :class:`.EMESimulation`]
+    simulation : Union[:class:`.Simulation`, :class:`.HeatSimulation`, :class:`.EMESimulation`, :class:`.ModalComponentModeler`, :class:`.TerminalComponentModeler`]
         Simulation to upload to server.
     task_name : str
         Name of task.
@@ -154,8 +155,8 @@ def run(
         Task priority for vGPU queue (1=lowest, 10=highest).
     Returns
     -------
-    Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`]
-        Object containing solver results for the supplied simulation.
+    Union[:class:`.SimulationData`, :class:`.HeatSimulationData`, :class:`.EMESimulationData`, :class:`.ModalComponentModelerData`, :class:`.TerminalComponentModelerData`]
+        Object containing solver results for the supplied input.
 
     Notes
     -----
@@ -203,10 +204,7 @@ def run(
     from tidy3d.plugins.smatrix.component_modelers.types import ComponentModelerType
 
     if isinstance(simulation, typing.get_args(ComponentModelerType)):
-        sims = getattr(simulation, "sim_dict", None)
-        if local_gradient or (
-            isinstance(sims, dict) and any(is_valid_for_autograd(s) for s in sims.values())
-        ):
+        if any(is_valid_for_autograd(s) for s in simulation.sim_dict.values()):
             from tidy3d.plugins.smatrix import run as smatrix_run
 
             path_dir = dirname(path) or "."