WIP: local changes before rebase

Author: George-Guryev-flxcmp

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- Added S-parameter de-embedding to `TerminalComponentModelerData`, enabling recalculation with shifted reference planes.
 - Added support for `tidy3d-extras`, an optional plugin that enables more accurate local mode solving via subpixel averaging.
 
 ### Changed
@@ -21,8 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [v2.10.0rc2] - 2025-10-01
 
-### Added
-- Added S-parameter de-embedding to `TerminalComponentModelerData`, enabling recalculation with shifted reference planes.  
+### Added  
 - New `MediumMonitor` that returns both permittivity and permeability profiles.
 - Task names are now optional when using `run(sim)` or `Job`. When running multiple jobs (via `run_async` or `Batch`), you can also provide simulations as a list without specifying task names. The previous dictionary-based format with explicit task names is still supported.
 - Enabled lazy loading of data via `web.load(..., lazy=True)`. When used, this returns a lightweight proxy object holding a reference to the data. On first access to any field or method, the proxy transparently loads the full object (same as with the default lazy=False).

tests/test_plugins/smatrix/test_terminal_component_modeler.py

@@ -29,6 +29,7 @@
     TerminalPortDataArray,
     WavePort,
 )
+from tidy3d.plugins.smatrix.data.data_array import PortNameDataArray
 from tidy3d.plugins.smatrix.ports.base_lumped import AbstractLumpedPort
 from tidy3d.plugins.smatrix.utils import s_to_z, validate_square_matrix
 
@@ -1421,29 +1422,57 @@ def test_S_parameter_deembedding(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=(WavePort, WavePort), grid_spec=grid_spec)
-    port0_idx = modeler.network_index(modeler.ports[0])
-    port1_idx = modeler.network_index(modeler.ports[1])
-    modeler_run1 = modeler.updated_copy(run_only=(port0_idx,))
 
     # Make sure the smatrix and impedance calculations work for reduced simulations
     modeler_data = run_component_modeler(monkeypatch, modeler)
     s_matrix = modeler_data.smatrix()
 
-    # test for invalid dimensions in port shifts
-    port_shifts = np.array([0])
-    with pytest.raises(ValueError):
-        S_dmb = modeler_data.change_port_reference_planes(smatrix=s_matrix, port_shifts=port_shifts)
-
-    port_shifts = np.zeros((3, 1))
-    with pytest.raises(ValueError):
-        S_dmb = modeler_data.change_port_reference_planes(smatrix=s_matrix, port_shifts=port_shifts)
+    # set up port shifts
+    port_names = [port.name for port in modeler.ports]
+    coords = {"port": port_names}
+    shift_vec = [0, 0]
+    port_shifts = PortNameDataArray(data=shift_vec, coords=coords)
 
-    # ensure matrices are identical if there is no shift in reference planes
-    port_shifts = np.array([0, 0])
+    # make sure that de-embedded S-matrices are identical to the original one if reference planes are not shifted
     S_dmb = modeler_data.change_port_reference_planes(smatrix=s_matrix, port_shifts=port_shifts)
+    S_dmb_shortcut = modeler_data.smatrix_deembedded(port_shifts=port_shifts)
     assert np.allclose(S_dmb.data.values, s_matrix.data.values)
+    assert np.allclose(S_dmb_shortcut.data.values, s_matrix.data.values)
 
     # make sure S-parameters are different if reference planes are moved
-    port_shifts = np.array([-2, 10])
+    port_shifts = PortNameDataArray(data=[-100, 200], coords=coords)
     S_dmb = modeler_data.change_port_reference_planes(smatrix=s_matrix, port_shifts=port_shifts)
+    S_dmb_shortcut = modeler_data.smatrix_deembedded(port_shifts=port_shifts)
     assert not np.allclose(S_dmb.data.values, s_matrix.data.values)
+    assert np.allclose(S_dmb.data.values, S_dmb_shortcut.data.values)
+
+    # test if `.smatrix_deembedded()` raises a `ValueError` when at least one port to be shifted is not defined in TCM
+    port_shifts_wrong = PortNameDataArray(data=[10, -10], coords={"port": ["wave_1", "LP_wave_2"]})
+
+    with pytest.raises(ValueError):
+        S_dmb = modeler_data.smatrix_deembedded(port_shifts=port_shifts_wrong)
+
+    # set up a new TCM with a mixture of `WavePort` and `CoaxialLumpedPort`
+    modeler_LP = make_coaxial_component_modeler(
+        port_types=(WavePort, CoaxialLumpedPort), grid_spec=grid_spec
+    )
+    modeler_data_LP = run_component_modeler(monkeypatch, modeler_LP)
+
+    # test if `.smatrix_deembedded()` raises a `ValueError` when one tries to de-embed a lumped port
+    port_shifts_LP = PortNameDataArray(data=[10, -10], coords={"port": ["wave_1", "coax_2"]})
+    with pytest.raises(ValueError):
+        S_dmb = modeler_data_LP.smatrix_deembedded(port_shifts=port_shifts_LP)
+
+    # update port shifts so that a reference plane is shifted only for `WavePort` port
+    port_shifts_LP = PortNameDataArray(data=[100], coords={"port": ["wave_1"]})
+
+    # get a new S-matrix
+    s_matrix_LP = modeler_data_LP.smatrix()
+
+    # de-embed S-matrix
+    S_dmb = modeler_data_LP.change_port_reference_planes(
+        smatrix=s_matrix_LP, port_shifts=port_shifts_LP
+    )
+    S_dmb_shortcut = modeler_data_LP.smatrix_deembedded(port_shifts=port_shifts_LP)
+    assert not np.allclose(S_dmb.data.values, s_matrix_LP.data.values)
+    assert np.allclose(S_dmb_shortcut.data.values, S_dmb.data.values)

tidy3d/plugins/smatrix/data/data_array.py

@@ -84,3 +84,19 @@ def _warn_rf_license(cls, values):
             log_once=True,
         )
         return values
+
+
+class PortNameDataArray(DataArray):
+    """Port shifts Data Array.
+
+    Example
+    -------
+    >>> import numpy as np
+    >>> port_names = ["port1", "port2"]
+    >>> coords = dict(port_name=port_names)
+    >>> data = (1 + 1j) * np.random.random((2,))
+    >>> port_data = PortNameDataArray(data, coords=coords)
+    """
+
+    __slots__ = ()
+    _dims = "port_name"

tidy3d/plugins/smatrix/data/terminal.py

@@ -7,7 +7,6 @@
 import numpy as np
 import pydantic.v1 as pd
 
-from tidy3d import ModeData
 from tidy3d.components.base import Tidy3dBaseModel, cached_property
 from tidy3d.components.data.data_array import FreqDataArray
 from tidy3d.components.data.monitor_data import MonitorData
@@ -17,8 +16,12 @@
 from tidy3d.log import log
 from tidy3d.plugins.smatrix.component_modelers.terminal import TerminalComponentModeler
 from tidy3d.plugins.smatrix.data.base import AbstractComponentModelerData
-from tidy3d.plugins.smatrix.data.data_array import PortDataArray, TerminalPortDataArray
-from tidy3d.plugins.smatrix.ports.types import TerminalPortType
+from tidy3d.plugins.smatrix.data.data_array import (
+    PortDataArray,
+    PortNameDataArray,
+    TerminalPortDataArray,
+)
+from tidy3d.plugins.smatrix.ports.types import LumpedPortType, TerminalPortType
 from tidy3d.plugins.smatrix.types import SParamDef
 from tidy3d.plugins.smatrix.utils import (
     ab_to_s,
@@ -120,17 +123,20 @@ def smatrix(
         return smatrix_data
 
     def change_port_reference_planes(
-        self, smatrix: MicrowaveSMatrixData, port_shifts: np.ndarray
+        self, smatrix: MicrowaveSMatrixData, port_shifts: PortNameDataArray = None
     ) -> MicrowaveSMatrixData:
         """
-        Performs S-parameter de-embedding by shifting reference planes `port_shifts` um.
+        Performs S-parameter de-embedding by shifting reference planes ``port_shifts`` um.
 
         Parameters
         ----------
         smatrix : :class:`.MicrowaveSMatrixData`
             S-parameters before reference planes are shifted.
-        port_shifts : np.ndarray
+        port_shifts : :class:`.PortNameDataArray`
             Numpy array of shifts of wave ports' reference planes.
+            The sign of a port shift reflects direction with respect to the axis normal to a ``WavePort`` plane:
+            E.g.: ``port_shift[0]==-a`` defines a shift in the first ``WavePort`` by
+            ``a`` um in direction opposite to the positive axis direction (the axis normal to the port plane).
 
         Returns
         -------
@@ -143,41 +149,96 @@ def change_port_reference_planes(
         S_new = np.zeros_like(S_matrix, dtype=complex)
         N_freq, N_ports, _ = S_matrix.shape
 
-        if len(port_shifts) != N_ports:
-            raise ValueError(
-                "A vector of WavePort reference plane shifts has to match a total number of `WavePort`'s in a simulation."
-                f"The expected length was {N_ports}, while a vector of {len(port_shifts)} was provided."
-            )
-
-        # extract `ModeSource` directions to ensure correct sign is used
-        port_shifts = np.ravel(port_shifts)
-        directions = np.array([1 if port.direction == "+" else -1 for port in self.modeler.ports])
-        directions = np.ravel(directions)
-
         # pre-allocate memory for effective propagation constants
-        kvecs = np.zeros((N_ports, N_freq), dtype=complex)
+        kvecs = np.zeros((N_freq, N_ports), dtype=complex)
+        shifts_vec = np.zeros(N_ports)
+        directions_vec = np.ones(N_ports)
+
+        modes_data = []
+        mode_indices = []
+        port_idxs = []
 
-        # extract mode data
+        # extract raw data
         key = self.data.keys_tuple[0]
         data = self.data[key].data
-        modes_data = tuple(mode_data for mode_data in data if isinstance(mode_data, ModeData))
+        ports = self.modeler.ports
+
+        # get port names and names of ports to be shifted
+        port_names = [port.name for port in ports]
+        shift_names = port_shifts.coords["port"].values
+
+        # Build a mapping for quick lookup from monitor name to monitor data
+        mode_map = {mode_data.monitor.name: mode_data for mode_data in data}
+
+        # form a numpy vector of port shifts
+        for shift_name in shift_names:
+            # ensure that port shifts were defined for valid ports
+            if shift_name not in port_names:
+                raise ValueError(
+                    "The specified port could not be found in the simulation! "
+                    f"Please, make sure the port name is from the following list {port_names}"
+                )
+
+            # get index of a shifted port in port_names list
+            idx = port_names.index(shift_name)
+            port = ports[idx]
+
+            # if de-embedding is requested for lumped port
+            if isinstance(port, LumpedPortType):
+                raise ValueError(
+                    "De-embedding currently supports only 'WavePort' instances. "
+                    f"Received type: '{type(port).__name__}'."
+                )
+                # alternatively we can send a warning and set `shifts_vector[index]` to 0.
+                # shifts_vector[index] = 0.0
+            else:
+                shifts_vec[idx] = port_shifts.sel(port=shift_name).values
+                directions_vec[idx] = -1 if port.direction == "-" else 1
+                mode_indices.append(port.mode_index)
+                port_idxs.append(idx)
+
+                # Collect corresponding mode_data
+                modes_data.append(mode_map[port._mode_monitor_name])
+
+        # Convert modes_data to tuple (same as original)
+        modes_data = tuple(modes_data)
+
+        # flatten port shift vector
+        shifts_vec = np.ravel(shifts_vec)
+        directions_vec = np.ravel(directions_vec)
+
+        # Collect relevant propagation constants and frequencies
+        n_complex_new, freqs = zip(
+            *(
+                (
+                    np.squeeze(mode_data.n_complex.sel(mode_index=mode_indices[i]).data),
+                    np.squeeze(mode_data.n_complex.f),
+                )
+                for i, mode_data in enumerate(modes_data)
+            )
+        )
 
-        # infer propagation constants from modal data
-        for i, mode_data in enumerate(modes_data):
-            n_complex = mode_data.n_complex
-            kvecs[i, :] = (2 * np.pi * n_complex.f * n_complex / C_0).squeeze()
+        # Convert to stacked arrays
+        n_complex_new = np.stack(n_complex_new, axis=1)
+        freqs = np.stack(freqs, axis=1)
 
-        # updated/de-embed S-parameters with respect to shifted reference planes
-        for i in range(N_freq):
-            phase = kvecs[:, i] * port_shifts * directions
-            P_inv = np.diag(np.exp(-1j * phase))
-            S_new[i, :, :] = P_inv @ S_matrix[i, :, :] @ P_inv
+        # construct transformation matrix P_inv
+        kvecs[:, port_idxs] = 2 * np.pi * freqs * n_complex_new / C_0
+        phase = -kvecs * shifts_vec * directions_vec
+        P_inv = np.exp(1j * phase)
+
+        # de-embed S-parameters: S_new = P_inv @ S_matrix @ P_inv
+        S_new = S_matrix * P_inv[:, :, np.newaxis] * P_inv[:, np.newaxis, :]
 
         # create a new Port Data Array
         smat_data = TerminalPortDataArray(S_new, coords=smatrix.data.coords)
 
         return smatrix.updated_copy(data=smat_data)
 
+    def smatrix_deembedded(self, port_shifts: np.ndarray = None) -> MicrowaveSMatrixData:
+        """Interface function returns  de-embedded S-parameter matrix."""
+        return self.change_port_reference_planes(self.smatrix(), port_shifts=port_shifts)
+
     @pd.root_validator(pre=False)
     def _warn_rf_license(cls, values):
         log.warning(