improve performance of antenna metrics calculation

Author: dmarek-flex

CHANGELOG.md

@@ -6,6 +6,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [Unreleased]
+
+### Changed
+- Improved performance of antenna metrics calculation by utilizing cached wave amplitude calculations instead of recomputing wave amplitudes for each port excitation in the `TerminalComponentModelerData`.
+
 ## [v2.10.0rc2] - 2025-10-01
 
 ### Added

tests/test_plugins/smatrix/test_terminal_component_modeler.py

@@ -1058,11 +1058,12 @@ def test_antenna_helpers(monkeypatch, tmp_path):
 
     # Test monitor data normalization with different amplitude types
     a_array = FreqDataArray(np.ones(len(modeler.freqs)), {"f": modeler.freqs})
+    a_array_raw = 2.0 * a_array
     normalized_data_array = modeler_data._monitor_data_at_port_amplitude(
-        modeler.ports[0], radiation_monitor.name, a_array
+        modeler.ports[0], radiation_monitor.name, a_array, a_array_raw
     )
     normalized_data_const = modeler_data._monitor_data_at_port_amplitude(
-        modeler.ports[0], radiation_monitor.name, 1.0
+        modeler.ports[0], radiation_monitor.name, 1.0, a_array_raw
     )
     assert isinstance(normalized_data_array, td.DirectivityData)
     assert isinstance(normalized_data_const, td.DirectivityData)
@@ -1198,7 +1199,7 @@ def test_run_only_and_element_mappings(monkeypatch, tmp_path):
     xy_grid = td.UniformGrid(dl=0.1 * 1e3)
     grid_spec = td.GridSpec(grid_x=xy_grid, grid_y=xy_grid, grid_z=z_grid)
     modeler = make_coaxial_component_modeler(
-        port_types=(CoaxialLumpedPort, WavePort), grid_spec=grid_spec
+        port_types=(CoaxialLumpedPort, CoaxialLumpedPort), grid_spec=grid_spec
     )
     port0_idx = modeler.network_index(modeler.ports[0])
     port1_idx = modeler.network_index(modeler.ports[1])

tidy3d/plugins/smatrix/analysis/antenna.py

@@ -5,9 +5,6 @@
 import numpy as np
 
 from tidy3d.components.microwave.data.monitor_data import AntennaMetricsData
-from tidy3d.plugins.smatrix.analysis.terminal import (
-    compute_wave_amplitudes_at_each_port,
-)
 from tidy3d.plugins.smatrix.data.data_array import PortDataArray
 from tidy3d.plugins.smatrix.data.terminal import TerminalComponentModelerData
 
@@ -69,26 +66,27 @@ def get_antenna_metrics_data(
         coords=coords,
     )
     b_sum = a_sum.copy()
+    a_matrix, b_matrix = terminal_component_modeler_data.port_power_wave_matrices
     # Retrieve associated simulation data
     combined_directivity_data = None
     for port, amplitude in port_dict.items():
+        port_in_index = terminal_component_modeler_data.modeler.network_index(port)
         if amplitude is not None:
             if np.isclose(amplitude, 0.0):
                 continue
         sim_data_port = terminal_component_modeler_data.data[
             terminal_component_modeler_data.modeler.get_task_name(port)
         ]
 
-        a, b = compute_wave_amplitudes_at_each_port(
-            modeler=terminal_component_modeler_data.modeler,
-            port_reference_impedances=terminal_component_modeler_data.port_reference_impedances,
-            sim_data=sim_data_port,
-            s_param_def="power",
+        a, b = (
+            a_matrix.sel(port_in=port_in_index, drop=True),
+            b_matrix.sel(port_in=port_in_index, drop=True),
         )
+
         # Select a possible subset of frequencies
         a = a.sel(f=f)
         b = b.sel(f=f)
-        a_raw = a.sel(port=terminal_component_modeler_data.modeler.network_index(port))
+        a_raw = a.sel(port_out=port_in_index)
 
         if amplitude is None:
             # No scaling performed when amplitude is None
@@ -97,7 +95,7 @@ def get_antenna_metrics_data(
         else:
             scaled_directivity_data = (
                 terminal_component_modeler_data._monitor_data_at_port_amplitude(
-                    port, rad_mon.name, amplitude
+                    port, rad_mon.name, amplitude, a_raw
                 )
             )
             scale_factor = amplitude / a_raw
@@ -109,8 +107,8 @@ def get_antenna_metrics_data(
             combined_directivity_data = scaled_directivity_data
         else:
             combined_directivity_data = combined_directivity_data + scaled_directivity_data
-        a_sum += a
-        b_sum += b
+        a_sum += a.rename({"port_out": "port"})
+        b_sum += b.rename({"port_out": "port"})
 
     # Compute and add power measures to results
     power_incident = np.real(0.5 * a_sum * np.conj(a_sum)).sum(dim="port")

tidy3d/plugins/smatrix/analysis/terminal.py

@@ -71,33 +71,10 @@ def terminal_construct_smatrix(
     source_indices = list(modeler_data.modeler.matrix_indices_source)
     run_source_indices = list(modeler_data.modeler.matrix_indices_run_sim)
 
-    values = np.zeros(
-        (len(modeler_data.modeler.freqs), len(monitor_indices), len(source_indices)),
-        dtype=complex,
-    )
-    coords = {
-        "f": np.array(modeler_data.modeler.freqs),
-        "port_out": monitor_indices,
-        "port_in": source_indices,
-    }
-    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(modeler_data=modeler_data)
-
-    for source_index in run_source_indices:
-        port, mode_index = modeler_data.modeler.network_dict[source_index]
-        sim_data = modeler_data.data[
-            modeler_data.modeler.get_task_name(port=port, mode_index=mode_index)
-        ]
-        a, b = modeler_data.compute_wave_amplitudes_at_each_port(
-            port_reference_impedances=port_impedances, sim_data=sim_data, s_param_def=s_param_def
-        )
-
-        indexer = {"port_in": source_index}
-        a_matrix = a_matrix._with_updated_data(data=a.data, coords=indexer)
-        b_matrix = b_matrix._with_updated_data(data=b.data, coords=indexer)
+    if s_param_def == "pseudo":
+        a_matrix, b_matrix = modeler_data.port_pseudo_wave_matrices
+    else:
+        a_matrix, b_matrix = modeler_data.port_power_wave_matrices
 
     # If excitation is assumed ideal, a_matrix is assumed to be diagonal
     # and the explicit inverse can be avoided. When only a subset of excitations
@@ -109,6 +86,8 @@ def terminal_construct_smatrix(
         # Scale each column by the corresponding diagonal entry
         s_matrix = b_matrix / a_diag[:, np.newaxis, :]
 
+    # Expand the smatrix using user defined mappings
+    s_matrix_expanded = s_matrix.reindex(port_in=source_indices, fill_value=0.0)
     # 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:
         coords_from = {
@@ -119,9 +98,9 @@ def terminal_construct_smatrix(
             "port_in": col_out,
             "port_out": row_out,
         }
-        data = mult_by * s_matrix.loc[coords_from].data
-        s_matrix = s_matrix._with_updated_data(data=data, coords=coords_to)
-    return s_matrix
+        data = mult_by * s_matrix_expanded.loc[coords_from].data
+        s_matrix_expanded = s_matrix_expanded._with_updated_data(data=data, coords=coords_to)
+    return s_matrix_expanded
 
 
 def port_reference_impedances(modeler_data: TerminalComponentModelerData) -> PortDataArray:
@@ -179,32 +158,29 @@ def port_reference_impedances(modeler_data: TerminalComponentModelerData) -> Por
     return port_impedances
 
 
-def compute_wave_amplitudes_at_each_port(
+def _compute_port_voltages_currents(
     modeler: TerminalComponentModeler,
-    port_reference_impedances: PortDataArray,
     sim_data: SimulationData,
-    s_param_def: SParamDef = "pseudo",
 ) -> tuple[PortDataArray, PortDataArray]:
-    """Compute the incident and reflected amplitudes at each port.
+    """Compute voltage and current values at all ports for a single simulation.
 
-    The computed amplitudes have not been normalized.
+    This function calculates the voltage and current at each monitor port from the
+    electromagnetic field data in a single simulation result. The voltages and currents
+    are computed according to the specific port type (e.g., lumped, wave) and are used
+    as inputs for subsequent wave amplitude calculations.
 
     Parameters
     ----------
     modeler : :class:`.TerminalComponentModeler`
-        The component modeler defining the ports and simulation settings.
-    port_reference_impedances : :class:`.PortDataArray`
-        Reference impedance at each port.
+        The component modeler containing port definitions and network mapping.
     sim_data : :class:`.SimulationData`
-        Results from a single simulation run.
-    s_param_def : SParamDef
-        The type of waves computed, either pseudo waves defined by Equation 53 and
-        Equation 54 in [1], or power waves defined by Equation 4.67 in [2].
+        Simulation results containing the electromagnetic field data.
 
     Returns
     -------
     tuple[:class:`.PortDataArray`, :class:`.PortDataArray`]
-        Incident (a) and reflected (b) wave amplitudes at each port.
+        A tuple containing the voltage and current arrays with dimensions (f, port),
+        where voltages and currents are computed for each frequency and monitor port.
     """
     network_indices = list(modeler.matrix_indices_monitor)
     values = np.zeros(
@@ -218,18 +194,56 @@ def compute_wave_amplitudes_at_each_port(
 
     V_matrix = PortDataArray(values, coords=coords)
     I_matrix = V_matrix.copy(deep=True)
-    a = V_matrix.copy(deep=True)
-    b = V_matrix.copy(deep=True)
 
     for network_index in network_indices:
         port, mode_index = modeler.network_dict[network_index]
         V_out, I_out = compute_port_VI(port, sim_data)
         indexer = {"port": network_index}
         V_matrix = V_matrix._with_updated_data(data=V_out.data, coords=indexer)
         I_matrix = I_matrix._with_updated_data(data=I_out.data, coords=indexer)
+    return (V_matrix, I_matrix)
+
+
+def _compute_wave_amplitudes_from_VI(
+    port_reference_impedances: PortDataArray,
+    port_voltages: PortDataArray,
+    port_currents: PortDataArray,
+    s_param_def: SParamDef = "pseudo",
+) -> tuple[PortDataArray, PortDataArray]:
+    """Convert port voltages and currents to incident and reflected wave amplitudes.
+
+    This function transforms voltage and current data at each port into forward-traveling
+    (incident, 'a') and backward-traveling (reflected, 'b') wave amplitudes using the
+    specified wave definition. The conversion handles impedance sign consistency and
+    applies the appropriate normalization based on the chosen S-parameter definition.
+
+    The wave amplitudes are computed using:
+    - Pseudo waves: Equations 53-54 from Marks and Williams [1]
+    - Power waves: Equation 4.67 from Pozar [2]
 
-    V_numpy = V_matrix.values
-    I_numpy = I_matrix.values
+    Parameters
+    ----------
+    port_reference_impedances : :class:`.PortDataArray`
+        Reference impedance values for each port with dimensions (f, port).
+    port_voltages : :class:`.PortDataArray`
+        Voltage values at each port with dimensions (f, port).
+    port_currents : :class:`.PortDataArray`
+        Current values at each port with dimensions (f, port).
+    s_param_def : SParamDef, optional
+        Wave definition type: "pseudo" for pseudo waves or "power" for power waves.
+        Defaults to "pseudo".
+
+    Returns
+    -------
+    tuple[:class:`.PortDataArray`, :class:`.PortDataArray`]
+        A tuple containing the incident (a) and reflected (b) wave amplitude arrays,
+        each with dimensions (f, port) representing the wave amplitudes at each
+        frequency and port.
+    """
+    a = port_voltages.copy(deep=True)
+    b = port_currents.copy(deep=True)
+    V_numpy = port_voltages.values
+    I_numpy = port_currents.values
     Z_numpy = port_reference_impedances.values
 
     # Check to make sure sign is consistent for all impedance values
@@ -254,6 +268,41 @@ def compute_wave_amplitudes_at_each_port(
     return a, b
 
 
+def compute_wave_amplitudes_at_each_port(
+    modeler: TerminalComponentModeler,
+    port_reference_impedances: PortDataArray,
+    sim_data: SimulationData,
+    s_param_def: SParamDef = "pseudo",
+) -> tuple[PortDataArray, PortDataArray]:
+    """Compute the incident and reflected amplitudes at each port.
+
+    The computed amplitudes have not been normalized.
+
+    Parameters
+    ----------
+    modeler : :class:`.TerminalComponentModeler`
+        The component modeler defining the ports and simulation settings.
+    port_reference_impedances : :class:`.PortDataArray`
+        Reference impedance at each port.
+    sim_data : :class:`.SimulationData`
+        Results from a single simulation run.
+    s_param_def : SParamDef
+        The type of waves computed, either pseudo waves defined by Equation 53 and
+        Equation 54 in [1], or power waves defined by Equation 4.67 in [2].
+
+    Returns
+    -------
+    tuple[:class:`.PortDataArray`, :class:`.PortDataArray`]
+        Incident (a) and reflected (b) wave amplitudes at each port.
+    """
+
+    port_voltages, port_currents = _compute_port_voltages_currents(modeler, sim_data)
+
+    return _compute_wave_amplitudes_from_VI(
+        port_reference_impedances, port_voltages, port_currents, s_param_def=s_param_def
+    )
+
+
 def compute_power_wave_amplitudes_at_each_port(
     modeler: TerminalComponentModeler,
     port_reference_impedances: PortDataArray,