First structural refactor complete

Author: daquinteroflex

Untitled.ipynb

@@ -1,4 +1,5 @@
-import tidy3d as td
+from __future__ import annotations
+
 import tidy3d.web as web
 
 task = web.load("he-172a2b2e-f576-4a3e-a90b-06a64194ed51")

test_task.py

@@ -765,7 +765,7 @@ def test_antenna_helpers(monkeypatch, tmp_path):
     # Run simulation to get data
     _ = run_component_modeler(monkeypatch, modeler)
     batch_data = modeler.batch_data
-    sim_data = batch_data[modeler._task_name(modeler.ports[0])]
+    sim_data = batch_data[modeler.get_task_name(modeler.ports[0])]
     rad_mon_data = sim_data[radiation_monitor.name]
 
     # Test monitor helper

tests/test_plugins/smatrix/test_terminal_component_modeler.py

@@ -0,0 +1,89 @@
+"""
+Tool for generating an S matrix automatically from a Tidy3d simulation and lumped port definitions.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from tidy3d.plugins.smatrix.data.data_array import ModalPortDataArray
+from tidy3d.plugins.smatrix.data.modal import ComponentModelerData
+
+
+def modal_construct_smatrix(modeler_data: ComponentModelerData) -> ModalPortDataArray:
+    """Post process :class:`.BatchData` to generate scattering matrix, for internal use only."""
+
+    max_mode_index_out, max_mode_index_in = modeler_data.modeler.max_mode_index
+    num_modes_out = max_mode_index_out + 1
+    num_modes_in = max_mode_index_in + 1
+    port_names_out, port_names_in = modeler_data.modeler.port_names
+
+    values = np.zeros(
+        (
+            len(port_names_out),
+            len(port_names_in),
+            num_modes_out,
+            num_modes_in,
+            len(modeler_data.modeler.freqs),
+        ),
+        dtype=complex,
+    )
+    coords = {
+        "port_out": port_names_out,
+        "port_in": port_names_in,
+        "mode_index_out": range(num_modes_out),
+        "mode_index_in": range(num_modes_in),
+        "f": np.array(modeler_data.modeler.freqs),
+    }
+    s_matrix = ModalPortDataArray(values, coords=coords)
+
+    # loop through source ports
+    for col_index in modeler_data.modeler.matrix_indices_run_sim:
+        port_name_in, mode_index_in = col_index
+        port_in = modeler_data.modeler.get_port_by_name(port_name=port_name_in)
+
+        sim_data = modeler_data.data[
+            modeler_data.modeler.get_task_name(port=port_in, mode_index=mode_index_in)
+        ]
+
+        for row_index in modeler_data.modeler.matrix_indices_monitor:
+            port_name_out, mode_index_out = row_index
+            port_out = modeler_data.modeler.get_port_by_name(port_name=port_name_out)
+
+            # directly compute the element
+            mode_amps_data = sim_data[port_out.name].copy().amps
+            dir_out = "-" if port_out.direction == "+" else "+"
+            amp = mode_amps_data.sel(f=coords["f"], direction=dir_out, mode_index=mode_index_out)
+            source_norm = modeler_data.modeler._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
+
+    # element can be determined by user-defined mapping
+    for (row_in, col_in), (row_out, col_out), mult_by in modeler_data.modeler.element_mappings:
+        port_out_from, mode_index_out_from = row_in
+        port_in_from, mode_index_in_from = col_in
+        coords_from = {
+            "port_in": port_in_from,
+            "mode_index_in": mode_index_in_from,
+            "port_out": port_out_from,
+            "mode_index_out": mode_index_out_from,
+        }
+
+        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
+
+    return s_matrix

tidy3d/plugins/smatrix/run/__init__.py renamed to tidy3d/plugins/smatrix/analysis/__init__.py

@@ -3,8 +3,10 @@
 import numpy as np
 
 from tidy3d.components.data.sim_data import SimulationData
-from tidy3d.plugins.smatrix.data.terminal import TerminalComponentModelerData
+from tidy3d.plugins.smatrix.component_modelers.terminal import TerminalComponentModeler
 from tidy3d.plugins.smatrix.data.data_array import PortDataArray, TerminalPortDataArray
+from tidy3d.plugins.smatrix.data.terminal import TerminalComponentModelerData
+from tidy3d.plugins.smatrix.ports.wave import WavePort
 from tidy3d.plugins.smatrix.utils import (
     ab_to_s,
     check_port_impedance_sign,
@@ -13,9 +15,7 @@
 )
 
 
-def terminal_construct_smatrix(
-    modeler_data: TerminalComponentModelerData
-) -> TerminalPortDataArray:
+def terminal_construct_smatrix(modeler_data: TerminalComponentModelerData) -> TerminalPortDataArray:
     """
     Constructs the scattering matrix (S-matrix) from raw simulation data.
 
@@ -50,12 +50,14 @@ def terminal_construct_smatrix(
     b_matrix = a_matrix.copy(deep=True)
 
     # Tabulate the reference impedances at each port and frequency
-    port_impedances = port_reference_impedances(modeler=modeler_data.modeler)
+    port_impedances = port_reference_impedances(modeler_data=modeler_data)
 
     # loop through source ports
-    for port_in in modeler.ports:
-        sim_data = batch_data[modeler._task_name(port=port_in)]
-        a, b = compute_power_wave_amplitudes_at_each_port(modeler, port_impedances, sim_data)
+    for port_in in modeler_data.modeler.ports:
+        sim_data = modeler_data.data[port_in]
+        a, b = compute_power_wave_amplitudes_at_each_port(
+            modeler_data.modeler, port_impedances, sim_data
+        )
         indexer = {"f": a.f, "port_in": port_in.name, "port_out": a.port}
         a_matrix.loc[indexer] = a
         b_matrix.loc[indexer] = b
@@ -64,9 +66,7 @@ def terminal_construct_smatrix(
     return s_matrix
 
 
-def port_reference_impedances(
-    modeler_data: TerminalComponentModelerData
-) -> PortDataArray:
+def port_reference_impedances(modeler_data: TerminalComponentModelerData) -> PortDataArray:
     """
     Calculates the reference impedance for each port across all frequencies.
 
@@ -98,17 +98,19 @@ def port_reference_impedances(
     )
     coords = {"f": np.array(modeler_data.modeler.freqs), "port": port_names}
     port_impedances = PortDataArray(values, coords=coords)
-    for port in modeler.ports:
+    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 = batch_data[modeler._task_name(port=port)]
+            sim_data_port = modeler_data.data[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
             port_impedances.loc[{"port": port.name}] = impedances.squeeze()
         else:
             # LumpedPorts have a constant reference impedance
-            port_impedances.loc[{"port": port.name}] = np.full(len(modeler.freqs), port.impedance)
+            port_impedances.loc[{"port": port.name}] = np.full(
+                len(modeler_data.modeler.freqs), port.impedance
+            )
 
     port_impedances = modeler_data.modeler._set_port_data_array_attributes(port_impedances)
     return port_impedances

tidy3d/plugins/smatrix/analysis/modal.py

@@ -16,6 +16,7 @@
 from tidy3d.exceptions import SetupError, Tidy3dKeyError
 from tidy3d.log import log
 from tidy3d.plugins.smatrix.ports.modal import Port
+from tidy3d.plugins.smatrix.ports.types import TerminalPortType
 from tidy3d.plugins.smatrix.ports.wave import WavePort
 
 # fwidth of gaussian pulse in units of central frequency
@@ -86,7 +87,7 @@ def _warn_rf_license(cls, val):
         return val
 
     @staticmethod
-    def _task_name(port: Port, mode_index: Optional[int] = None) -> str:
+    def get_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."""
         if mode_index is not None:
             return f"smatrix_{port.name}_{mode_index}"

tidy3d/plugins/smatrix/run/terminal.py renamed to tidy3d/plugins/smatrix/analysis/terminal.py

@@ -93,7 +93,7 @@ def sim_dict(self) -> dict[str, Simulation]:
 
             new_mnts = list(self.simulation.monitors) + mode_monitors
             sim_copy = self.simulation.copy(update={"sources": [mode_source], "monitors": new_mnts})
-            task_name = self._task_name(port=port, mode_index=mode_index)
+            task_name = self.get_task_name(port=port, mode_index=mode_index)
             sim_dict[task_name] = sim_copy
         return sim_dict