maybe need damian's help

Author: daquinteroflex

tidy3d/plugins/smatrix/component_modelers/terminal.py

@@ -438,7 +438,7 @@ def _construct_smatrix(self) -> TerminalPortDataArray:
 
         return _construct_smatrix(self)
 
-    def _internal_construct_smatrix(self, batch_data: typing.Any) -> TerminalPortDataArray:
+    def _internal_construct_smatrix(self, batch_data: typing.Any = None) -> TerminalPortDataArray:
         from tidy3d.plugins.smatrix.run import _internal_construct_smatrix
 
-        return _internal_construct_smatrix(self)
+        return _internal_construct_smatrix(self, batch_data=batch_data)

tidy3d/plugins/smatrix/run.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import typing
+
 import numpy as np
 
 from tidy3d.components.data.sim_data import SimulationData
@@ -47,39 +49,6 @@ def _port_reference_impedances(simulation) -> PortDataArray:
     return port_impedances
 
 
-def _internal_construct_smatrix(simulation) -> TerminalPortDataArray:
-    """Post process :class:`.BatchData` to generate scattering matrix, for internal use only."""
-    from tidy3d.plugins.smatrix.utils import ab_to_s
-
-    port_names = [port.name for port in simulation.ports]
-
-    values = np.zeros(
-        (len(simulation.freqs), len(port_names), len(port_names)),
-        dtype=complex,
-    )
-    coords = {
-        "f": np.array(simulation.freqs),
-        "port_out": port_names,
-        "port_in": port_names,
-    }
-    a_matrix = TerminalPortDataArray(values, coords=coords)
-    b_matrix = a_matrix.copy(deep=True)
-
-    # Tabulate the reference impedances at each port and frequency
-    port_impedances = _port_reference_impedances(simulation=simulation)
-
-    # loop through source ports
-    for port_in in simulation.ports:
-        sim_data = simulation.batch_data[simulation._task_name(port=port_in)]
-        a, b = simulation.compute_power_wave_amplitudes_at_each_port(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
-
-    s_matrix = ab_to_s(a_matrix, b_matrix)
-    return s_matrix
-
-
 def compute_power_wave_amplitudes_at_each_port(
     simulation, port_reference_impedances: PortDataArray, sim_data: SimulationData
 ) -> tuple[PortDataArray, PortDataArray]:
@@ -138,3 +107,41 @@ def compute_power_wave_amplitudes_at_each_port(
     b.values = F_numpy * (V_numpy - np.conj(Z_numpy) * I_numpy)
 
     return a, b
+
+
+def _internal_construct_smatrix(simulation, batch_data: typing.Any = None) -> TerminalPortDataArray:
+    """Post process :class:`.BatchData` to generate scattering matrix, for internal use only."""
+    if batch_data:
+        pass
+    else:
+        batch_data = simulation.batch_data
+
+    from tidy3d.plugins.smatrix.utils import ab_to_s
+
+    port_names = [port.name for port in simulation.ports]
+
+    values = np.zeros(
+        (len(simulation.freqs), len(port_names), len(port_names)),
+        dtype=complex,
+    )
+    coords = {
+        "f": np.array(simulation.freqs),
+        "port_out": port_names,
+        "port_in": port_names,
+    }
+    a_matrix = TerminalPortDataArray(values, coords=coords)
+    b_matrix = a_matrix.copy(deep=True)
+
+    # Tabulate the reference impedances at each port and frequency
+    port_impedances = _port_reference_impedances(simulation=simulation)
+
+    # loop through source ports
+    for port_in in simulation.ports:
+        sim_data = batch_data[simulation._task_name(port=port_in)]
+        a, b = compute_power_wave_amplitudes_at_each_port(simulation, 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
+
+    s_matrix = ab_to_s(a_matrix, b_matrix)
+    return s_matrix