Documentation and Greptile comments

Author: daquinteroflex

tidy3d/plugins/smatrix/local_run.py

@@ -1,12 +1,5 @@
 """
 Tool for generating an S matrix automatically from a Tidy3d simulation and lumped port definitions.
-
-One major problem that prevents a stronger use of typing in these functions is the backwards compatibility of
-the imports. Ideally, we would do
-
-run(simulation: MySimulation) -> MySimulationData
-
-but because these
 """
 
 from __future__ import annotations
@@ -34,13 +27,37 @@ def run(
     batch_data: BatchData = None,
     path_dir: str = DEFAULT_DATA_DIR,
 ) -> TerminalComponentModelerData:
+    """
+    Runs the S-matrix calculation for a given terminal component modeler simulation.
+
+    This function orchestrates the entire process of obtaining the S-matrix. It takes the
+    simulation definition, runs the batch simulations if data is not already provided,
+    constructs the S-matrix from the results, and packages all relevant data into a
+    single `TerminalComponentModelerData` object for analysis and export.
+
+    Parameters
+    ----------
+    simulation : TerminalComponentModeler
+        The simulation setup defining the component, ports, and frequencies.
+    batch_data : BatchData, optional
+        Pre-computed simulation results from a Tidy3D batch run. If None, the function
+        will execute the batch run defined within the `simulation` object.
+    path_dir : str, optional
+        Directory path to store simulation data. Defaults to `DEFAULT_DATA_DIR`.
+
+    Returns
+    -------
+    TerminalComponentModelerData
+        A data container holding the original simulation definition, the calculated
+        S-matrix data, and the raw simulation data for each port excitation.
+    """
     _ = simulation.get_path_dir(path_dir)
 
     if batch_data is None:
         # TODO Remove once fully decoupled simulation from batch_data
         batch_data = simulation.batch_data
 
-    terminal_port_data = construct_smatrix(simulation=simulation)
+    terminal_port_data = construct_smatrix(simulation=simulation, batch_data=batch_data)
 
     port_to_sim_data_map = {
         port_i: batch_data[simulation._task_name(port=port_i)] for port_i in simulation.ports
@@ -55,7 +72,28 @@ def run(
 def construct_smatrix(
     simulation: TerminalComponentModeler, batch_data: BatchData = None
 ) -> TerminalPortDataArray:
-    """Post process :class:`.BatchData` to generate scattering matrix."""
+    """ "
+    Constructs the scattering matrix (S-matrix) from raw simulation data.
+
+    This function iterates through each port, treating it as an input source once.
+    For each input source, it calculates the resulting incident (a) and reflected (b)
+    power wave amplitudes at all ports. These amplitudes are compiled into matrices,
+    which are then used to compute the final S-matrix.
+
+    Parameters
+    ----------
+    simulation : TerminalComponentModeler
+        The simulation setup defining the component, ports, and frequencies.
+    batch_data : BatchData, optional
+        Pre-computed simulation results. If None, the function will execute the
+        batch run defined within the `simulation` object.
+
+    Returns
+    -------
+    TerminalPortDataArray
+        The computed S-matrix as a data array with dimensions for frequency,
+        output port, and input port.
+    """
     if batch_data is None:
         batch_data = simulation.batch_data
 
@@ -93,8 +131,28 @@ def construct_smatrix(
 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`.
+    """
+    Calculates the reference impedance for each port across all frequencies.
+
+    This function determines the characteristic impedance for every port in the simulation.
+    It handles two types of ports differently:
+    - `WavePort`: The impedance is frequency-dependent and is computed from the modal
+      properties stored in the simulation data.
+    - Other ports (e.g., `LumpedPort`): The impedance is a constant value defined by the
+      user.
+
+    Parameters
+    ----------
+    simulation : TerminalComponentModeler
+        The simulation setup defining the ports.
+    batch_data : BatchData, optional
+        Pre-computed simulation results, required for `WavePort` impedance calculation.
+        If None, the batch data from the simulation object is used.
+
+    Returns
+    -------
+    PortDataArray
+        A data array containing the complex impedance for each port at each frequency.
     """
     if batch_data is None:
         # TODO Remove once fully decoupled simulation from batch_data
@@ -131,20 +189,30 @@ def compute_power_wave_amplitudes_at_each_port(
     port_reference_impedances: PortDataArray,
     sim_data: SimulationData,
 ) -> tuple[PortDataArray, PortDataArray]:
-    """Compute the incident and reflected power wave amplitudes at each port.
-    The computed amplitudes have not been normalized.
+    """
+    Computes the incident (a) and reflected (b) power wave amplitudes at all ports
+    from a single simulation run where one port was excited.
+
+    This function converts the raw voltage (V) and current (I) data into power wave
+    amplitudes using standard microwave engineering formulas. It also performs a
+    sanity check to ensure the real part of the reference impedance is positive,
+    flipping signs of V and Z if necessary to maintain physical consistency.
 
     Parameters
     ----------
-    port_reference_impedances : :class:`.PortDataArray`
-        Reference impedance at each port.
-    sim_data : :class:`.SimulationData`
-        Results from the simulation.
+    simulation : TerminalComponentModeler
+        The simulation setup defining the ports.
+    port_reference_impedances : PortDataArray
+        The characteristic impedance for each port at each frequency.
+    sim_data : SimulationData
+        The raw results (fields, currents, etc.) from a single simulation run.
 
     Returns
     -------
-    tuple[:class:`.PortDataArray`, :class:`.PortDataArray`]
-        Incident (a) and reflected (b) power wave amplitudes at each port.
+    tuple[PortDataArray, PortDataArray]
+        A tuple containing two PortDataArrays:
+        - a: The incident power wave amplitudes at each port.
+        - b: The reflected power wave amplitudes at each port.
     """
     port_names = [port.name for port in simulation.ports]
     values = np.zeros(