Ok ready to review

Author: daquinteroflex

tidy3d/plugins/smatrix/component_modelers/terminal.py

@@ -288,14 +288,14 @@ def get_radiation_monitor_by_name(self, monitor_name: str) -> DirectivityMonitor
                 return monitor
         raise Tidy3dKeyError(f"No radiation monitor named '{monitor_name}'.")
 
+    ##### Backwards compatibility methods, should be external functions really#####
     def run(self, path_dir: str = DEFAULT_DATA_DIR):
         """Solves for the scattering matrix of the system."""
         from tidy3d.plugins.smatrix.local_run import run
 
         terminal_component_modeler_data = run(simulation=self, path_dir=path_dir)
         return terminal_component_modeler_data
 
-    ##### Backwards compatibility methods, should be used within TerminalComponentModelerData really#####
     @staticmethod
     def _check_port_impedance_sign(Z_numpy: np.ndarray):
         """Sanity check for consistent sign of real part of Z for each port across all frequencies."""
@@ -329,6 +329,10 @@ def _monitor_data_at_port_amplitude(
         monitor_data: MonitorData,
         a_port: Union[FreqDataArray, complex],
     ) -> MonitorData:
+        """Normalize the monitor data to a desired complex amplitude of a port,
+        represented by ``a_port``, where :math:`\\frac{1}{2}|a|^2` is the power
+        incident from the port into the system.
+        """
         from tidy3d.plugins.smatrix.local_run import run
 
         terminal_component_modeler_data = run(simulation=self)
@@ -476,6 +480,28 @@ def get_antenna_metrics_data(
         port_amplitudes: Optional[dict[str, complex]] = None,
         monitor_name: Optional[str] = None,
     ) -> AntennaMetricsData:
+        """Calculate antenna parameters using superposition of fields from multiple port excitations.
+
+        The method computes the radiated far fields and port excitation power wave amplitudes
+        for a superposition of port excitations, which can be used to analyze antenna radiation
+        characteristics.
+
+        Parameters
+        ----------
+        port_amplitudes : dict[str, complex] = None
+            Dictionary mapping port names to their desired excitation amplitudes. For each port,
+            :math:`\\frac{1}{2}|a|^2` represents the incident power from that port into the system.
+            If None, uses only the first port without any scaling of the raw simulation data.
+        monitor_name : str = None
+            Name of the :class:`.DirectivityMonitor` to use for calculating far fields.
+            If None, uses the first monitor in `radiation_monitors`.
+
+        Returns
+        -------
+        :class:`.AntennaMetricsData`
+            Container with antenna parameters including directivity, gain, and radiation efficiency,
+            computed from the superposition of fields from all excited ports.
+        """
         from tidy3d.plugins.smatrix.local_run import run
 
         terminal_component_modeler_data = run(simulation=self)

tidy3d/plugins/smatrix/local_run.py

@@ -11,8 +11,6 @@
 
 from __future__ import annotations
 
-import typing
-
 import numpy as np
 
 from tidy3d.components.data.sim_data import SimulationData
@@ -28,17 +26,18 @@
 from tidy3d.plugins.smatrix.data.data_array import PortDataArray, TerminalPortDataArray
 from tidy3d.plugins.smatrix.ports.wave import WavePort
 from tidy3d.plugins.smatrix.utils import check_port_impedance_sign, compute_F, compute_port_VI
+from tidy3d.web import BatchData
 
 
 def run(
     simulation: TerminalComponentModeler,
-    batch_data: typing.Any = None,
+    batch_data: BatchData = None,
     path_dir: str = DEFAULT_DATA_DIR,
 ) -> TerminalComponentModelerData:
     _ = simulation.get_path_dir(path_dir)
 
     if batch_data is None:
-        # Remove once fully decoupled simulation from batch_data
+        # TODO Remove once fully decoupled simulation from batch_data
         batch_data = simulation.batch_data
 
     terminal_port_data = construct_smatrix(simulation=simulation)
@@ -54,7 +53,7 @@ def run(
 
 
 def construct_smatrix(
-    simulation: TerminalComponentModeler, batch_data: typing.Any = None
+    simulation: TerminalComponentModeler, batch_data: BatchData = None
 ) -> TerminalPortDataArray:
     """Post process :class:`.BatchData` to generate scattering matrix."""
     if batch_data is None:
@@ -91,11 +90,15 @@ def construct_smatrix(
     return s_matrix
 
 
-def port_reference_impedances(simulation: TerminalComponentModeler) -> PortDataArray:
+def port_reference_impedances(
+    simulation: TerminalComponentModeler, batch_data: BatchData = None
+) -> PortDataArray:
     """Tabulates the reference impedance of each port at each frequency using the
     supplied :class:`.BatchData`.
     """
-    # TODO properly refactor, plugins data types should not have web methods.
+    if batch_data is None:
+        # TODO Remove once fully decoupled simulation from batch_data
+        batch_data = simulation.batch_data
 
     port_names = [port.name for port in simulation.ports]
 
@@ -108,7 +111,7 @@ def port_reference_impedances(simulation: TerminalComponentModeler) -> PortDataA
     for port in simulation.ports:
         if isinstance(port, WavePort):
             # Mode solver data for each wave port is stored in its associated SimulationData
-            sim_data_port = simulation.batch_data[simulation._task_name(port=port)]
+            sim_data_port = batch_data[simulation._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