add selective simulation capabilities to TerminalComponentModeler

refactored common code into AbstractComponentModeler

add toggle for controlling how s matrix is calculated

added validation of run_only and element_mappings for modelers

Author: dmarek-flex

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Access field decay values in `SimulationData` via `sim_data.field_decay` as `TimeDataArray`.
+- 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.
 
 ### Changed
 - By default, batch downloads will skip files that already exist locally. To force re-downloading and replace existing files, pass the `replace_existing=True` argument to `Batch.load()`, `Batch.download()`, or `BatchData.load()`.

tests/test_plugins/smatrix/terminal_component_modeler_def.py

@@ -276,7 +276,7 @@ def make_port(center, direction, type, name) -> Union[CoaxialLumpedPort, WavePor
                 inner_diameter=2 * Rinner,
                 normal_axis=2,
                 direction=direction,
-                name=name,
+                name="coax" + name,
                 num_grid_cells=port_cells,
                 impedance=reference_impedance,
             )
@@ -311,7 +311,7 @@ def make_port(center, direction, type, name) -> Union[CoaxialLumpedPort, WavePor
                 center=center,
                 size=[2 * Router, 2 * Router, 0],
                 direction=direction,
-                name=name,
+                name="wave" + name,
                 mode_spec=td.ModeSpec(num_modes=1),
                 mode_index=0,
                 voltage_integral=voltage_integral,
@@ -321,9 +321,9 @@ def make_port(center, direction, type, name) -> Union[CoaxialLumpedPort, WavePor
         return port
 
     center_src1 = [0, 0, -length / 2]
-    port_1 = make_port(center_src1, direction="+", type=port_types[0], name="coax_port_1")
+    port_1 = make_port(center_src1, direction="+", type=port_types[0], name="_1")
     center_src2 = [0, 0, length / 2]
-    port_2 = make_port(center_src2, direction="-", type=port_types[1], name="coax_port_2")
+    port_2 = make_port(center_src2, direction="-", type=port_types[1], name="_2")
     ports = [port_1, port_2]
     freqs = np.linspace(freq_start, freq_stop, 100)
 

tests/test_plugins/smatrix/test_component_modeler.py

@@ -374,6 +374,35 @@ def test_mapping_exclusion(monkeypatch):
     _test_mappings(element_mappings, s_matrix)
 
 
+def test_mapping_with_run_only():
+    """Make sure that the Modeler is correctly validated when both run_only and
+    element_mappings are provided."""
+    ports = make_ports()
+
+    EXCLUDE_INDEX = ("right_bot", 0)
+    element_mappings = []
+    run_only = []
+    # add a mapping to each element in the row of EXCLUDE_INDEX
+    for port in ports:
+        for mode_index in range(port.mode_spec.num_modes):
+            row_index = (port.name, mode_index)
+            run_only.append(row_index)
+            if row_index != EXCLUDE_INDEX:
+                mapping = ((row_index, row_index), (row_index, EXCLUDE_INDEX), +1)
+                element_mappings.append(mapping)
+
+    # add the self-self coupling element to complete row
+    mapping = ((("right_bot", 1), ("right_bot", 1)), (EXCLUDE_INDEX, EXCLUDE_INDEX), +1)
+    element_mappings.append(mapping)
+
+    # Will pass, since run_only covers all source indices in element_mapping
+    _ = make_component_modeler(element_mappings=element_mappings, run_only=run_only)
+
+    run_only.remove(EXCLUDE_INDEX)
+    with pytest.raises(pydantic.ValidationError):
+        _ = make_component_modeler(element_mappings=element_mappings, run_only=run_only)
+
+
 def test_batch_filename(tmp_path):
     modeler = make_component_modeler()
     path = modeler._batch_path

tests/test_plugins/smatrix/test_terminal_component_modeler.py

@@ -900,3 +900,44 @@ def test_get_combined_antenna_parameters_data(monkeypatch, tmp_path):
     assert not np.allclose(
         antenna_params.radiation_efficiency, single_port_params.radiation_efficiency
     )
+
+
+def test_run_only_and_element_mappings(monkeypatch, tmp_path):
+    """Checks the terminal component modeler works when running with a subset of excitations."""
+    z_grid = td.UniformGrid(dl=1 * 1e3)
+    xy_grid = td.UniformGrid(dl=0.1 * 1e3)
+    grid_spec = td.GridSpec(grid_x=xy_grid, grid_y=xy_grid, grid_z=z_grid)
+    modeler = make_coaxial_component_modeler(
+        path_dir=str(tmp_path), port_types=(CoaxialLumpedPort, WavePort), grid_spec=grid_spec
+    )
+    port0_idx = modeler.network_index(modeler.ports[0])
+    port1_idx = modeler.network_index(modeler.ports[1])
+    modeler_run1 = modeler.updated_copy(run_only=(port0_idx,))
+
+    # Make sure the smatrix and impedance calculations work for reduced simulations
+    s_matrix = run_component_modeler(monkeypatch, modeler_run1)
+    with pytest.raises(ValueError):
+        TerminalComponentModeler._validate_square_matrix(s_matrix, "test_method")
+    _ = modeler_run1.port_reference_impedances
+
+    assert len(modeler_run1.sim_dict) == 1
+    S11 = (port0_idx, port0_idx)
+    S21 = (port1_idx, port0_idx)
+    S12 = (port0_idx, port1_idx)
+    S22 = (port1_idx, port1_idx)
+    element_mappings = ((S11, S22, 1),)
+    modeler_with_mappings = modeler.updated_copy(element_mappings=element_mappings)
+    assert len(modeler_with_mappings.sim_dict) == 2
+
+    # Column 1 is mapped to column 2, resulting in one simulation
+    element_mappings = ((S11, S22, 1), (S21, S12, 1))
+    modeler_with_mappings = modeler.updated_copy(element_mappings=element_mappings)
+    s_matrix = run_component_modeler(monkeypatch, modeler_with_mappings)
+    assert np.all(s_matrix.values[:, 0, 0] == s_matrix.values[:, 1, 1])
+    assert np.all(s_matrix.values[:, 0, 1] == s_matrix.values[:, 1, 0])
+    assert len(modeler_with_mappings.sim_dict) == 1
+
+    # Mapping is incomplete, so two simulations are run
+    element_mappings = ((S11, S22, 1), (S12, S21, 1))
+    modeler_with_mappings = modeler.updated_copy(element_mappings=element_mappings)
+    assert len(modeler_with_mappings.sim_dict) == 2

tidy3d/plugins/smatrix/component_modelers/base.py

@@ -4,7 +4,7 @@
 
 import os
 from abc import ABC, abstractmethod
-from typing import Optional, Union, get_args
+from typing import Generic, Optional, TypeVar, Union, get_args
 
 import numpy as np
 import pydantic.v1 as pd
@@ -13,7 +13,8 @@
 from tidy3d.components.data.data_array import DataArray
 from tidy3d.components.data.sim_data import SimulationData
 from tidy3d.components.simulation import Simulation
-from tidy3d.components.types import FreqArray
+from tidy3d.components.types import Complex, FreqArray
+from tidy3d.components.validators import assert_unique_names
 from tidy3d.config import config
 from tidy3d.constants import HERTZ
 from tidy3d.exceptions import SetupError, Tidy3dKeyError
@@ -31,8 +32,12 @@
 LumpedPortType = Union[LumpedPort, CoaxialLumpedPort]
 TerminalPortType = Union[LumpedPortType, WavePort]
 
+# Generic type variables for matrix indices and elements
+IndexType = TypeVar("IndexType")
+ElementType = TypeVar("ElementType")
 
-class AbstractComponentModeler(ABC, Tidy3dBaseModel):
+
+class AbstractComponentModeler(ABC, Generic[IndexType, ElementType], Tidy3dBaseModel):
     """Tool for modeling devices and computing port parameters."""
 
     simulation: Simulation = pd.Field(
@@ -108,6 +113,24 @@ class AbstractComponentModeler(ABC, Tidy3dBaseModel):
         "fields that were not used to create the task will cause errors.",
     )
 
+    run_only: Optional[tuple[IndexType, ...]] = pd.Field(
+        None,
+        title="Run Only",
+        description="Set of matrix indices that define the simulations to run. "
+        "If ``None``, simulations will be run for all indices in the scattering matrix. "
+        "If a tuple is given, simulations will be run only for the given matrix indices.",
+    )
+
+    element_mappings: tuple[tuple[ElementType, ElementType, Complex], ...] = pd.Field(
+        (),
+        title="Element Mappings",
+        description="Tuple of S matrix element mappings, each described by a tuple of "
+        "(input_element, output_element, coefficient), where the coefficient is the "
+        "element_mapping coefficient describing the relationship between the input and output "
+        "matrix element. If all elements of a given column of the scattering matrix are defined "
+        "by ``element_mappings``, the simulation corresponding to this column is skipped automatically.",
+    )
+
     @pd.validator("simulation", always=True)
     def _sim_has_no_sources(cls, val):
         """Make sure simulation has no sources as they interfere with tool."""
@@ -131,6 +154,30 @@ def _warn_rf_license(cls, val):
             )
         return val
 
+    @pd.validator("element_mappings", always=True)
+    def _validate_element_mappings(cls, element_mappings, values):
+        """
+        Validate that each source index referenced in element_mappings is included in run_only.
+        """
+        run_only = values.get("run_only")
+        if run_only is None:
+            return element_mappings
+
+        valid_set = set(run_only)
+        invalid_indices = set()
+        for mapping in element_mappings:
+            input_element = mapping[0]
+            output_element = mapping[1]
+            for source_index in [input_element[1], output_element[1]]:
+                if source_index not in valid_set:
+                    invalid_indices.add(source_index)
+        if invalid_indices:
+            raise SetupError(
+                f"'element_mappings' references source index(es) {invalid_indices} "
+                f"that are not present in run_only: {run_only}."
+            )
+        return element_mappings
+
     @staticmethod
     def _task_name(port: Port, mode_index: Optional[int] = None) -> str:
         """The name of a task, determined by the port of the source and mode index, if given."""
@@ -230,6 +277,44 @@ def get_port_by_name(self, port_name: str) -> Port:
             raise Tidy3dKeyError(f'Port "{port_name}" not found.')
         return ports[0]
 
+    @property
+    @abstractmethod
+    def matrix_indices_monitor(self) -> tuple[IndexType, ...]:
+        """Abstract property for all matrix indices that will be used to collect data."""
+
+    @cached_property
+    def matrix_indices_source(self) -> tuple[IndexType, ...]:
+        """Tuple of all the source matrix indices, which may be less than the total number of ports."""
+        if self.run_only is not None:
+            return self.run_only
+        return self.matrix_indices_monitor
+
+    @cached_property
+    def matrix_indices_run_sim(self) -> tuple[IndexType, ...]:
+        """Tuple of all the matrix indices that will be used to run simulations."""
+
+        if not self.element_mappings:
+            return self.matrix_indices_source
+
+        # all the (i, j) pairs in `S_ij` that are tagged as covered by `element_mappings`
+        elements_determined_by_map = [element_out for (_, element_out, _) in self.element_mappings]
+
+        # loop through rows of the full s matrix and record rows that still need running.
+        source_indices_needed = []
+        for col_index in self.matrix_indices_source:
+            # loop through columns and keep track of whether each element is covered by mapping.
+            matrix_elements_covered = []
+            for row_index in self.matrix_indices_monitor:
+                element = (row_index, col_index)
+                element_covered_by_map = element in elements_determined_by_map
+                matrix_elements_covered.append(element_covered_by_map)
+
+            # if any matrix elements in row still not covered by map, a source is needed for row.
+            if not all(matrix_elements_covered):
+                source_indices_needed.append(col_index)
+
+        return source_indices_needed
+
     @abstractmethod
     def _construct_smatrix(self, batch_data: BatchData) -> DataArray:
         """Post process :class:`.BatchData` to generate scattering matrix."""
@@ -308,3 +393,5 @@ def sim_data_by_task_name(self, task_name: str) -> SimulationData:
         sim_data = self.batch_data[task_name]
         config.logging_level = log_level_cache
         return sim_data
+
+    _unique_port_names = assert_unique_names("ports")

tidy3d/plugins/smatrix/component_modelers/modal.py

@@ -15,9 +15,8 @@
 from tidy3d.components.simulation import Simulation
 from tidy3d.components.source.field import ModeSource
 from tidy3d.components.source.time import GaussianPulse
-from tidy3d.components.types import Ax, Complex
+from tidy3d.components.types import Ax
 from tidy3d.components.viz import add_ax_if_none, equal_aspect
-from tidy3d.exceptions import SetupError
 from tidy3d.plugins.smatrix.ports.modal import ModalPortDataArray, Port
 from tidy3d.web.api.container import BatchData
 
@@ -27,7 +26,7 @@
 Element = tuple[MatrixIndex, MatrixIndex]  # the 'ij' in S_ij
 
 
-class ComponentModeler(AbstractComponentModeler):
+class ComponentModeler(AbstractComponentModeler[MatrixIndex, Element]):
     """
     Tool for modeling devices and computing scattering matrix elements.
 
@@ -47,52 +46,6 @@ class ComponentModeler(AbstractComponentModeler):
         "For each input mode, one simulation will be run with a modal source.",
     )
 
-    element_mappings: tuple[tuple[Element, Element, Complex], ...] = pd.Field(
-        (),
-        title="Element Mappings",
-        description="Mapping between elements of the scattering matrix, "
-        "as specified by pairs of ``(port name, mode index)`` matrix indices, where the "
-        "first element of the pair is the output and the second element of the pair is the input."
-        "Each item of ``element_mappings`` is a tuple of ``(element1, element2, c)``, where "
-        "the scattering matrix ``Smatrix[element2]`` is set equal to ``c * Smatrix[element1]``."
-        "If all elements of a given column of the scattering matrix are defined by "
-        " ``element_mappings``, the simulation corresponding to this column "
-        "is skipped automatically.",
-    )
-
-    run_only: Optional[tuple[MatrixIndex, ...]] = pd.Field(
-        None,
-        title="Run Only",
-        description="If given, a tuple of matrix indices, specified by (:class:`.Port`, ``int``),"
-        " to run only, excluding the other rows from the scattering matrix. "
-        "If this option is used, "
-        "the data corresponding to other inputs will be missing in the resulting matrix.",
-    )
-    """Finally, to exclude some rows of the scattering matrix, one can supply a ``run_only`` parameter to the
-    :class:`ComponentModeler`. ``run_only`` contains the scattering matrix indices that the user wants to run as a
-    source. If any indices are excluded, they will not be run."""
-
-    verbose: bool = pd.Field(
-        False,
-        title="Verbosity",
-        description="Whether the :class:`.ComponentModeler` should print status and progressbars.",
-    )
-
-    callback_url: str = pd.Field(
-        None,
-        title="Callback URL",
-        description="Http PUT url to receive simulation finish event. "
-        "The body content is a json file with fields "
-        "``{'id', 'status', 'name', 'workUnit', 'solverVersion'}``.",
-    )
-
-    @pd.validator("simulation", always=True)
-    def _sim_has_no_sources(cls, val):
-        """Make sure simulation has no sources as they interfere with tool."""
-        if len(val.sources) > 0:
-            raise SetupError("'ComponentModeler.simulation' must not have any sources.")
-        return val
-
     @cached_property
     def sim_dict(self) -> dict[str, Simulation]:
         """Generate all the :class:`.Simulation` objects for the S matrix calculation."""
@@ -121,39 +74,6 @@ def matrix_indices_monitor(self) -> tuple[MatrixIndex, ...]:
                 matrix_indices.append((port.name, mode_index))
         return tuple(matrix_indices)
 
-    @cached_property
-    def matrix_indices_source(self) -> tuple[MatrixIndex, ...]:
-        """Tuple of all the source matrix indices (port, mode_index) in the Component Modeler."""
-        if self.run_only is not None:
-            return self.run_only
-        return self.matrix_indices_monitor
-
-    @cached_property
-    def matrix_indices_run_sim(self) -> tuple[MatrixIndex, ...]:
-        """Tuple of all the source matrix indices (port, mode_index) in the Component Modeler."""
-
-        if self.element_mappings is None or self.element_mappings == {}:
-            return self.matrix_indices_source
-
-        # all the (i, j) pairs in `S_ij` that are tagged as covered by `element_mappings`
-        elements_determined_by_map = [element_out for (_, element_out, _) in self.element_mappings]
-
-        # loop through rows of the full s matrix and record rows that still need running.
-        source_indices_needed = []
-        for col_index in self.matrix_indices_source:
-            # loop through columns and keep track of whether each element is covered by mapping.
-            matrix_elements_covered = []
-            for row_index in self.matrix_indices_monitor:
-                element = (row_index, col_index)
-                element_covered_by_map = element in elements_determined_by_map
-                matrix_elements_covered.append(element_covered_by_map)
-
-            # if any matrix elements in row still not covered by map, a source is needed for row.
-            if not all(matrix_elements_covered):
-                source_indices_needed.append(col_index)
-
-        return source_indices_needed
-
     @cached_property
     def port_names(self) -> tuple[list[str], list[str]]:
         """List of port names for inputs and outputs, respectively."""