Nearly all tests

Author: daquinteroflex

tests/test_plugins/smatrix/test_terminal_component_modeler.py

@@ -40,8 +40,8 @@ def run_component_modeler(
     sim_dict = modeler.sim_dict
     batch_data = {task_name: run_emulated(sim) for task_name, sim in sim_dict.items()}
     port_data = PortSimulationData(
-        ports=[port for port in batch_data.keys()],
-        data=[sim_data for sim_data in batch_data.values()],
+        ports=list(batch_data.keys()),
+        data=list(batch_data.values()),
     )
     modeler_data = TerminalComponentModelerData(modeler=modeler, data=port_data)
     monkeypatch.setattr(AbstractComponentModeler, "inv", lambda matrix: np.eye(len(modeler.ports)))
@@ -727,11 +727,11 @@ def test_port_impedance_check():
     Z_numpy = np.ones((50, 3))
     Z_numpy[:, 1] = -1.0
     # All ok if same sign for every frequency
-    TerminalComponentModeler.check_port_impedance_sign(Z_numpy)
+    TerminalComponentModelerData.check_port_impedance_sign(Z_numpy)
     Z_numpy[25, 1] = 1.0
     # Change of sign is unexpected
     with pytest.raises(Tidy3dError):
-        TerminalComponentModeler.check_port_impedance_sign(Z_numpy)
+        TerminalComponentModelerData.check_port_impedance_sign(Z_numpy)
 
 
 def test_antenna_helpers(monkeypatch, tmp_path):
@@ -756,6 +756,7 @@ def test_antenna_helpers(monkeypatch, tmp_path):
     # Run simulation to get data
     modeler_data = run_component_modeler(monkeypatch, modeler)
     rad_mon_data = modeler_data.data[radiation_monitor.name]
+    sim_data = modeler_data.data[modeler_data.modeler.get_task_name(modeler.ports[0])]
 
     # Test monitor helper
     found_mon = modeler.get_radiation_monitor_by_name(radiation_monitor.name)

tidy3d/plugins/smatrix/analysis/antenna.py

@@ -67,12 +67,14 @@ def get_antenna_metrics_data(
     # Retrieve associated simulation data
     combined_directivity_data = None
     for port, amplitude in port_dict.items():
-        sim_data_port = terminal_component_modeler_data.data.data[port]
+        sim_data_port = terminal_component_modeler_data.data[
+            terminal_component_modeler_data.modeler.get_task_name(port)
+        ]
         radiation_data = sim_data_port[rad_mon.name]
 
         a, b = compute_power_wave_amplitudes_at_each_port(
             modeler=terminal_component_modeler_data.modeler,
-            port_reference_impedances=terminal_component_modeler_data.modeler.port_reference_impedances,
+            port_reference_impedances=terminal_component_modeler_data.port_reference_impedances,
             sim_data=sim_data_port,
         )
         # Select a possible subset of frequencies

tidy3d/plugins/smatrix/analysis/terminal.py

@@ -54,7 +54,7 @@ def terminal_construct_smatrix(modeler_data: TerminalComponentModelerData) -> Te
 
     # loop through source ports
     for port_in in modeler_data.modeler.ports:
-        sim_data = modeler_data.data[port_in]
+        sim_data = modeler_data.data[modeler_data.modeler.get_task_name(port=port_in)]
         a, b = compute_power_wave_amplitudes_at_each_port(
             modeler_data.modeler, port_impedances, sim_data
         )
@@ -95,7 +95,7 @@ def port_reference_impedances(modeler_data: TerminalComponentModelerData) -> Por
     for port in modeler_data.modeler.ports:
         if isinstance(port, WavePort):
             # Mode solver data for each wave port is stored in its associated SimulationData
-            sim_data_port = modeler_data.data[port]
+            sim_data_port = modeler_data.data[modeler_data.modeler.get_task_name(port=port)]
             # WavePorts have a port impedance calculated from its associated modal field distribution
             # and is frequency dependent.
             impedances = port.compute_port_impedance(sim_data_port).values

tidy3d/plugins/smatrix/data/modal.py

@@ -15,6 +15,24 @@ class PortSimulationData(Tidy3dBaseModel):
     ports: tuple[PortReferenceType, ...]
     data: tuple[SimulationData, ...]
 
+    def __getitem__(self, port_name: str) -> SimulationData:
+        """
+        Allows retrieving simulation data by the port name.
+
+        Args:
+            port_name: The string name of the port to look up.
+
+        Returns:
+            The SimulationData object corresponding to the given port name.
+
+        Raises:
+            KeyError: If no port with the given name is found.
+        """
+        for i, port_i in enumerate(self.ports):
+            if port_i == port_name:
+                return self.data[i]
+        raise KeyError(f"Port '{port_name}' not found.")
+
 
 class ComponentModelerData(Tidy3dBaseModel):
     modeler: ComponentModeler = pd.Field(

tidy3d/plugins/smatrix/data/terminal.py

@@ -110,7 +110,7 @@ def _monitor_data_at_port_amplitude(
         incident from the port into the system.
         """
         a_raw, _ = self.compute_power_wave_amplitudes_at_each_port(
-            self.port_reference_impedances, sim_data
+            self.port_reference_impedances, sim_data=sim_data
         )
         a_raw_port = a_raw.sel(port=port.name)
         if not isinstance(a_port, FreqDataArray):
@@ -173,7 +173,7 @@ def port_reference_impedances(self) -> PortDataArray:
         """
         from tidy3d.plugins.smatrix.analysis.terminal import port_reference_impedances
 
-        return port_reference_impedances(self.modeler)
+        return port_reference_impedances(self)
 
     def compute_power_wave_amplitudes_at_each_port(
         self, port_reference_impedances: PortDataArray, sim_data: SimulationData

tidy3d/plugins/smatrix/data/types.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from typing import Union
+
+from tidy3d.plugins.smatrix.data.modal import ComponentModelerData
+from tidy3d.plugins.smatrix.data.terminal import TerminalComponentModelerData
+
+ComponentModelerDataType = Union[TerminalComponentModelerData, ComponentModelerData]

tidy3d/plugins/smatrix/run.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import os
+
 from tidy3d.plugins.smatrix.component_modelers.modal import ComponentModeler
 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 ComponentModelerData, PortSimulationMap
+from tidy3d.plugins.smatrix.data.modal import ComponentModelerData, PortSimulationData
 from tidy3d.plugins.smatrix.data.terminal import TerminalComponentModelerData
 from tidy3d.plugins.smatrix.data.types import ComponentModelerDataType
 from tidy3d.web import Batch, BatchData
@@ -19,47 +21,96 @@ def create_batch(
     file_name: str = "batch.hdf5",
     **kwargs,
 ) -> Batch:
-    batch = Batch(simulations=modeler.sim_dict, **kwargs)
-    batch.to_file(file_name)
+    """Creates a simulation Batch from a component modeler and saves it to a file.
+
+    Args:
+        modeler: The component modeler that defines the set of simulations.
+        path_dir: Directory where the batch file will be saved.
+        file_name: Name for the HDF5 file where the batch is stored.
+        **kwargs: Additional keyword arguments passed to the `Batch` constructor.
 
+    Returns:
+        The configured `Batch` object ready for execution.
+    """
+    filepath = os.path.join(path_dir, file_name)
+    batch = Batch(simulations=modeler.sim_dict, **kwargs)
+    batch.to_file(filepath)
     return batch
 
 
 def compose_terminal_modeler_data(
     modeler: TerminalComponentModeler,
-    batch_data: BatchData = None,
+    batch_data: BatchData,
 ) -> TerminalComponentModelerData:
-    port_to_sim_data_map = {
-        port_i: batch_data[modeler.get_task_name(port=port_i)] for port_i in modeler.ports
-    }
-    port_simulation_data = PortSimulationMap(data=port_to_sim_data_map)
+    """Assembles `TerminalComponentModelerData` from simulation results.
+
+    This function maps the simulation data from a completed batch run back to the
+    ports of the terminal component modeler.
+
+    Args:
+        modeler: The `TerminalComponentModeler` used to generate the simulations.
+        batch_data: The results obtained from running the simulation `Batch`.
+
+    Returns:
+        A `TerminalComponentModelerData` object containing the results mapped to
+        their respective ports.
+    """
+    ports = [modeler.get_task_name(port=port_i) for port_i in modeler.ports]
+    data = [batch_data[modeler.get_task_name(port=port_i)] for port_i in modeler.ports]
+    port_simulation_data = PortSimulationData(ports=ports, data=data)
     return TerminalComponentModelerData(modeler=modeler, data=port_simulation_data)
 
 
 def compose_component_modeler_data(
     modeler: ComponentModeler,
-    batch_data: BatchData = None,
+    batch_data: BatchData,
 ) -> ComponentModelerData:
-    port_to_sim_data_map = {
-        port_i: batch_data[modeler.get_task_name(port=port_i)] for port_i in modeler.ports
-    }
-    port_simulation_data = PortSimulationMap(data=port_to_sim_data_map)
+    """Assembles `ComponentModelerData` from simulation results.
+
+    This function maps the simulation data from a completed batch run back to the
+    ports of the component modeler.
+
+    Args:
+        modeler: The `ComponentModeler` used to generate the simulations.
+        batch_data: The results obtained from running the simulation `Batch`.
+
+    Returns:
+        A `ComponentModelerData` object containing the results mapped to
+        their respective ports.
+    """
+    ports = [modeler.get_task_name(port=port_i) for port_i in modeler.ports]
+    data = [batch_data[modeler.get_task_name(port=port_i)] for port_i in modeler.ports]
+    port_simulation_data = PortSimulationData(ports=ports, data=data)
     return ComponentModelerData(modeler=modeler, data=port_simulation_data)
 
 
 def compose_modeler_data(
     modeler: ComponentModelerType,
-    batch_data: BatchData = None,
+    batch_data: BatchData,
 ) -> ComponentModelerDataType:
-    """
-    This method internally determines which functions to run to compose the corresponding ComponentModelerDataType.
+    """Selects the correct composer based on the modeler type and creates the data object.
+
+    This method acts as a dispatcher, inspecting the type of `modeler` to determine
+    which composer function (`compose_component_modeler_data` or
+    `compose_terminal_modeler_data`) to invoke.
+
+    Args:
+        modeler: The component modeler, which can be either a `ComponentModeler` or
+            a `TerminalComponentModeler`.
+        batch_data: The results obtained from running the simulation `Batch`.
+
+    Returns:
+        The appropriate `ComponentModelerDataType` object containing the simulation results.
+
+    Raises:
+        TypeError: If the provided `modeler` is not a recognized type.
     """
     if isinstance(modeler, ComponentModeler):
         modeler_data = compose_component_modeler_data(modeler=modeler, batch_data=batch_data)
     elif isinstance(modeler, TerminalComponentModeler):
         modeler_data = compose_terminal_modeler_data(modeler=modeler, batch_data=batch_data)
     else:
-        raise Exception("sasas")
+        raise TypeError(f"Unsupported modeler type: {type(modeler).__name__}")
 
     return modeler_data
 
@@ -68,8 +119,20 @@ def run(
     modeler: ComponentModelerType,
     path_dir: str = DEFAULT_DATA_DIR,
 ) -> ComponentModelerDataType:
-    """
-    This method internally determines which functions to run to compose the corresponding ComponentModelerDataType.
+    """Executes the full simulation workflow for a given component modeler.
+
+    This function orchestrates the end-to-end process:
+    1. Creates a `Batch` of simulations from the `modeler`.
+    2. Submits the `Batch` for execution and waits for results.
+    3. Composes the results into a structured `ComponentModelerDataType` object.
+
+    Args:
+        modeler: The component modeler defining the simulations to be run.
+        path_dir: The directory where the batch file will be saved.
+
+    Returns:
+        A `ComponentModelerDataType` object containing the processed simulation data,
+        ready for S-parameter extraction and analysis.
     """
     batch = create_batch(modeler=modeler, path_dir=path_dir)
     batch_data = batch.run()