Propagate latest changes to autograd CM

tests/test_plugins/smatrix/test_component_modeler_autograd.py

@@ -8,7 +8,7 @@
 import tidy3d as td
 import tidy3d.web as web
 from tidy3d.plugins.smatrix.analysis import terminal as terminal_analysis
-from tidy3d.plugins.smatrix.component_modelers.modal import ComponentModeler
+from tidy3d.plugins.smatrix.component_modelers.modal import ModalComponentModelerData
 from tidy3d.plugins.smatrix.component_modelers.terminal import TerminalComponentModeler
 from tidy3d.plugins.smatrix.data.data_array import TerminalPortDataArray
 from tidy3d.plugins.smatrix.ports.modal import Port as ModalPort
@@ -162,7 +162,7 @@ def _build_base_sim(scale: float) -> td.Simulation:
     )
 
 
-def build_modal_modeler(scale: float) -> ComponentModeler:
+def build_modal_modeler(scale: float) -> ModalComponentModelerData:
     sim = _build_base_sim(scale)
 
     # two modal ports on +/- z sides
@@ -183,7 +183,7 @@ def build_modal_modeler(scale: float) -> ComponentModeler:
     )
 
     freqs = [2.0e14]
-    return ComponentModeler(simulation=sim, ports=(p1, p2), freqs=freqs)
+    return ModalComponentModelerData(simulation=sim, ports=(p1, p2), freqs=freqs)
 
 
 def build_terminal_modeler(scale: float) -> TerminalComponentModeler:

tidy3d/plugins/smatrix/component_modelers/base.py

@@ -2,9 +2,8 @@
 
 from __future__ import annotations
 
-import os
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union, get_args
+from typing import Generic, Optional, TypeVar, Union, get_args
 
 import numpy as np
 import pydantic.v1 as pd
@@ -27,6 +26,8 @@
 from tidy3d.plugins.smatrix.ports.modal import Port
 from tidy3d.plugins.smatrix.ports.types import TerminalPortType
 from tidy3d.plugins.smatrix.ports.wave import WavePort
+
+# DO NOT import from web if it can be avoided, to avoid circular imports
 from tidy3d.web.core.types import PayType
 
 # fwidth of gaussian pulse in units of central frequency
@@ -37,73 +38,6 @@
 IndexType = TypeVar("IndexType")
 ElementType = TypeVar("ElementType")
 
-if TYPE_CHECKING:
-    from tidy3d.plugins.smatrix.data.types import ComponentModelerDataType
-
-    from .types import ComponentModelerType
-
-
-def _run_component_modeler(
-    modeler: ComponentModelerType,
-    task_name: str,
-    folder_name: str,
-    path: str,
-    callback_url: Optional[str],
-    verbose: bool,
-    solver_version: Optional[str],
-    local_gradient: bool,
-    max_num_adjoint_per_fwd: int,
-    pay_type: Union[PayType, str],
-) -> ComponentModelerDataType:
-    """Run a Component Modeler via autograd by batching its underlying simulations."""
-
-    from tidy3d.web.api.autograd.autograd import DEFAULT_DATA_DIR, _run_async
-
-    path_dir = os.dirname(path) if path else DEFAULT_DATA_DIR
-    if not path_dir:
-        path_dir = DEFAULT_DATA_DIR
-
-    sims = modeler.sim_dict
-
-    sim_data_map = _run_async(
-        simulations=sims,
-        folder_name=folder_name,
-        path_dir=path_dir,
-        callback_url=callback_url,
-        verbose=verbose,
-        simulation_type="tidy3d_autograd_async",
-        solver_version=solver_version,
-        parent_tasks=None,
-        local_gradient=local_gradient,
-        max_num_adjoint_per_fwd=max_num_adjoint_per_fwd,
-        pay_type=pay_type,
-    )
-
-    return _compose_modeler_data_from_sim_map(modeler=modeler, sim_data_map=sim_data_map)
-
-
-def _compose_modeler_data_from_sim_map(
-    modeler: ComponentModelerType, sim_data_map: dict
-) -> ComponentModelerDataType:
-    """Create ComponentModelerDataType from a dict of SimulationData keyed by task name."""
-
-    # local imports to avoid cycles through tidy3d.web
-    from tidy3d.components.data.index import IndexSimulationData
-    from tidy3d.plugins.smatrix.component_modelers.modal import ComponentModeler
-    from tidy3d.plugins.smatrix.component_modelers.terminal import TerminalComponentModeler
-    from tidy3d.plugins.smatrix.data.modal import ComponentModelerData
-    from tidy3d.plugins.smatrix.data.terminal import TerminalComponentModelerData
-
-    # preserve mapping order
-    index = tuple(sim_data_map.keys())
-    data = tuple(sim_data_map.values())
-    indexed = IndexSimulationData(index=index, data=data)
-
-    if isinstance(modeler, ComponentModeler):
-        return ComponentModelerData(modeler=modeler, data=indexed)
-    if isinstance(modeler, TerminalComponentModeler):
-        return TerminalComponentModelerData(modeler=modeler, data=indexed)
-
 
 class AbstractComponentModeler(ABC, Generic[IndexType, ElementType], Tidy3dBaseModel):
     """Tool for modeling devices and computing port parameters."""
@@ -298,6 +232,7 @@ def run(
         deprecation_warning: bool = True,
     ):
         """Run component modeler locally, with autograd support."""
+        from tidy3d.plugins.smatrix.run import _run_component_modeler
 
         if deprecation_warning:
             log.warning(

tidy3d/plugins/smatrix/run.py

@@ -2,32 +2,40 @@
 
 import json
 import os
-from typing import Optional
+from typing import Optional, Union
 
 from tidy3d.components.base import Tidy3dBaseModel
+from tidy3d.components.data.index import SimulationDataMap
 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,
 )
-from tidy3d.plugins.smatrix.data.modal import ModalComponentModelerData, SimulationDataMap
+from tidy3d.plugins.smatrix.data.modal import ModalComponentModelerData
 from tidy3d.plugins.smatrix.data.terminal import TerminalComponentModelerData
 from tidy3d.plugins.smatrix.data.types import ComponentModelerDataType
 from tidy3d.web import Batch, BatchData
+from tidy3d.web.api.autograd.autograd import DEFAULT_DATA_DIR, _run_async
+from tidy3d.web.core.types import PayType
 
-DEFAULT_DATA_DIR = "."
 
+def compose_simulation_data_map(sim_data_map: dict) -> SimulationDataMap:
+    # preserve mapping order
+    index = tuple(sim_data_map.keys())
+    data = tuple(sim_data_map.values())
+    indexed = SimulationDataMap(keys=index, values=data)
+    return indexed
 
-def compose_simulation_data_index(port_task_map: dict[str, str]) -> SimulationDataMap:
-    port_data_dict = {}
-    for _, _ in port_task_map.items():
-        pass
-        # FIXME: get simulationdata for each port
-        # port_data_dict[port] = sim_data_i
 
-    return SimulationDataMap(
-        keys=tuple(port_data_dict.keys()), values=tuple(port_data_dict.values())
-    )
+def _compose_modeler_data_from_sim_map(
+    modeler: ComponentModelerType, sim_data_map: dict
+) -> ComponentModelerDataType:
+    """Create ComponentModelerDataType from a dict of SimulationData keyed by task name."""
+    indexed = compose_simulation_data_map(sim_data_map)
+    if isinstance(modeler, ModalComponentModeler):
+        return ModalComponentModelerData(modeler=modeler, data=indexed)
+    if isinstance(modeler, TerminalComponentModeler):
+        return TerminalComponentModelerData(modeler=modeler, data=indexed)
 
 
 def compose_terminal_modeler_data(
@@ -45,7 +53,7 @@ def compose_terminal_modeler_data(
         A `TerminalComponentModelerData` object containing the results mapped to
         their respective ports.
     """
-    port_simulation_data = compose_simulation_data_index(port_task_map)
+    port_simulation_data = compose_simulation_data_map(port_task_map)
     return TerminalComponentModelerData(modeler=modeler, data=port_simulation_data)
 
 
@@ -65,7 +73,7 @@ def compose_component_modeler_data(
         A `ModalComponentModelerData` object containing the results mapped to
         their respective ports.
     """
-    port_simulation_data = compose_simulation_data_index(port_task_map)
+    port_simulation_data = compose_simulation_data_map(port_task_map)
     return ModalComponentModelerData(modeler=modeler, data=port_simulation_data)
 
 
@@ -98,7 +106,7 @@ def compose_modeler(
     elif modeler_type == "TerminalComponentModeler":
         modeler = TerminalComponentModeler.from_file(modeler_file)
     else:
-        raise TypeError(f"Unsupported modeler type: {type(modeler).__name__}")
+        raise TypeError(f"Unsupported modeler type: {modeler_type}")
     return modeler
 
 
@@ -273,3 +281,39 @@ def run(
     batch_data = batch.run()
     modeler_data = compose_modeler_data_from_batch_data(modeler=modeler, batch_data=batch_data)
     return modeler_data
+
+
+def _run_component_modeler(
+    modeler: ComponentModelerType,
+    task_name: str,
+    folder_name: str,
+    path: str,
+    callback_url: Optional[str],
+    verbose: bool,
+    solver_version: Optional[str],
+    local_gradient: bool,
+    max_num_adjoint_per_fwd: int,
+    pay_type: Union[PayType, str],
+) -> ComponentModelerDataType:
+    """Run a Component Modeler via autograd by batching its underlying simulations."""
+    path_dir = os.dirname(path) if path else DEFAULT_DATA_DIR
+    if not path_dir:
+        path_dir = DEFAULT_DATA_DIR
+
+    sims = modeler.sim_dict
+
+    sim_data_map = _run_async(
+        simulations=sims,
+        folder_name=folder_name,
+        path_dir=path_dir,
+        callback_url=callback_url,
+        verbose=verbose,
+        simulation_type="tidy3d_autograd_async",
+        solver_version=solver_version,
+        parent_tasks=None,
+        local_gradient=local_gradient,
+        max_num_adjoint_per_fwd=max_num_adjoint_per_fwd,
+        pay_type=pay_type,
+    )
+
+    return _compose_modeler_data_from_sim_map(modeler=modeler, sim_data_map=sim_data_map)