feat: add EMESimulationData.coeffs (FXC-4275)

Author: caseyflex

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Added `symmetrize_mirror`, `symmetrize_rotation`, `symmetrize_diagonal` functions to the autograd plugin. They can be used for enforcing symmetries in topology optimization.
+- Added `EMESimulationData.coeffs` to store coefficients from the EME solver, including interface S matrices and effective propagation indices.
 
 ### Changed
 - Removed validator that would warn if `PerturbationMedium` values could become numerically unstable, since an error will anyway be raised if this actually happens when the medium is converted using actual perturbation data.

tests/test_components/test_eme.py

@@ -311,16 +311,16 @@ def test_eme_simulation():
 
     # test warning for not providing wavelength in autogrid
     grid_spec = td.GridSpec.auto(min_steps_per_wvl=20)
+    sim = sim.updated_copy(grid_spec=grid_spec)
     with AssertLogLevel("INFO", contains_str="wavelength"):
-        sim = sim.updated_copy(grid_spec=grid_spec)
+        _ = sim.updated_copy(monitors=[])
     # multiple freqs are ok, but not for autogrid
     _ = sim.updated_copy(
         grid_spec=td.GridSpec.uniform(dl=0.2), freqs=[10000000000.0, *list(sim.freqs)]
     )
     with AssertLogLevel("INFO", contains_str="wavelength"):
         _ = sim.updated_copy(
-            freqs=[*list(sim.freqs), 10000000000.0],
-            grid_spec=grid_spec,
+            freqs=[*list(sim.freqs), 10000000000.0], grid_spec=grid_spec, monitors=[]
         )
 
     # test port offsets
@@ -558,12 +558,14 @@ def test_eme_simulation():
             constraint="passive",
             eme_grid_spec=td.EMEUniformGrid(num_cells=1, mode_spec=td.EMEModeSpec(num_modes=40)),
             grid_spec=sim.grid_spec.updated_copy(wavelength=1),
+            monitors=[],
         )
         sim_good.validate_pre_upload()
         sim_good = sim.updated_copy(
             constraint=None,
             eme_grid_spec=td.EMEUniformGrid(num_cells=1, mode_spec=td.EMEModeSpec(num_modes=60)),
             grid_spec=sim.grid_spec.updated_copy(wavelength=1),
+            monitors=[],
         )
         sim_good.validate_pre_upload()
     # warn about num modes with constraint
@@ -1305,7 +1307,7 @@ def test_eme_periodicity():
 
     # remove the field monitor, now it passes
     desired_cell_index_pairs = set([(i, i + 1) for i in range(6)] + [(5, 1)])
-    with AssertLogLevel(None):
+    with AssertLogLevel("WARNING", contains_str="deprecated"):
         sim = sim.updated_copy(
             monitors=[m for m in sim.monitors if not isinstance(m, td.EMEFieldMonitor)]
         )

tidy3d/__init__.py

@@ -185,6 +185,8 @@
     ChargeDataArray,
     DiffractionDataArray,
     EMECoefficientDataArray,
+    EMEFluxDataArray,
+    EMEInterfaceSMatrixDataArray,
     EMEModeIndexDataArray,
     EMEScalarFieldDataArray,
     EMEScalarModeFieldDataArray,
@@ -605,8 +607,10 @@ def set_logging_level(level: str) -> None:
     "EMEFieldData",
     "EMEFieldDataset",
     "EMEFieldMonitor",
+    "EMEFluxDataArray",
     "EMEFreqSweep",
     "EMEGrid",
+    "EMEInterfaceSMatrixDataArray",
     "EMELengthSweep",
     "EMEModeIndexDataArray",
     "EMEModeSolverData",

tidy3d/components/data/data_array.py

@@ -1188,6 +1188,31 @@ class EMESMatrixDataArray(DataArray):
     _data_attrs = {"long_name": "scattering matrix element"}
 
 
+class EMEInterfaceSMatrixDataArray(DataArray):
+    """Scattering matrix elements at a single cell interface for a fixed pair of ports,
+    possibly with an extra sweep index.
+    Example
+    -------
+    >>> mode_index_in = [0, 1]
+    >>> mode_index_out = [0, 1, 2]
+    >>> eme_cell_index = [2, 4]
+    >>> f = [2e14]
+    >>> sweep_index = np.arange(10)
+    >>> coords = dict(
+    ...     f=f,
+    ...     sweep_index=sweep_index,
+    ...     eme_cell_index=eme_cell_index,
+    ...     mode_index_out=mode_index_out,
+    ...     mode_index_in=mode_index_in,
+    ... )
+    >>> fd = EMEInterfaceSMatrixDataArray((1 + 1j) * np.random.random((1, 10, 2, 3, 2)), coords=coords)
+    """
+
+    __slots__ = ()
+    _dims = ("f", "sweep_index", "eme_cell_index", "mode_index_out", "mode_index_in")
+    _data_attrs = {"long_name": "scattering matrix element"}
+
+
 class EMEModeIndexDataArray(DataArray):
     """Complex-valued effective propagation index of an EME mode,
     also indexed by EME cell.
@@ -1206,6 +1231,23 @@ class EMEModeIndexDataArray(DataArray):
     _data_attrs = {"long_name": "Propagation index"}
 
 
+class EMEFluxDataArray(DataArray):
+    """Power flux of an EME mode, also indexed by EME cell.
+
+    Example
+    -------
+    >>> f = [2e14, 3e14]
+    >>> mode_index = np.arange(4)
+    >>> eme_cell_index = np.arange(5)
+    >>> coords = dict(f=f, mode_index=mode_index, eme_cell_index=eme_cell_index)
+    >>> data = EMEModeIndexDataArray((1+1j) * np.random.random((2,4,5)), coords=coords)
+    """
+
+    __slots__ = ()
+    _dims = ("f", "sweep_index", "eme_cell_index", "mode_index")
+    _data_attrs = {"units": WATT, "long_name": "flux"}
+
+
 class ChargeDataArray(DataArray):
     """Charge data array.
 
@@ -1595,8 +1637,10 @@ def _make_impedance_data_array(result: DataArray) -> ImpedanceResultType:
     EMEScalarFieldDataArray,
     EMEScalarModeFieldDataArray,
     EMESMatrixDataArray,
+    EMEInterfaceSMatrixDataArray,
     EMECoefficientDataArray,
     EMEModeIndexDataArray,
+    EMEFluxDataArray,
     EMEFreqModeDataArray,
     ChargeDataArray,
     SteadyVoltageDataArray,

tidy3d/components/eme/data/dataset.py

@@ -2,10 +2,14 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import pydantic.v1 as pd
 
 from tidy3d.components.data.data_array import (
     EMECoefficientDataArray,
+    EMEFluxDataArray,
+    EMEInterfaceSMatrixDataArray,
     EMEModeIndexDataArray,
     EMEScalarFieldDataArray,
     EMEScalarModeFieldDataArray,
@@ -40,21 +44,72 @@ class EMESMatrixDataset(Dataset):
 
 
 class EMECoefficientDataset(Dataset):
-    """Dataset storing expansion coefficients for the modes in a cell.
+    """Dataset storing various coefficients related to the EME simulation.
+    These coefficients can be used for debugging or optimization.
+
+    The ``A`` and ``B`` fields store the expansion coefficients for the modes in a cell.
     These are defined at the cell centers.
+
+    The ``n_complex`` and ``flux`` fields respectively store the complex-valued effective
+    propagation index and the power flux associated with the modes used in the
+    EME calculation.
+
+    The ``interface_Sij`` fields store the S matrices associated with the interfaces
+    between EME cells.
     """
 
-    A: EMECoefficientDataArray = pd.Field(
-        ...,
+    A: Optional[EMECoefficientDataArray] = pd.Field(
+        None,
         title="A coefficient",
         description="Coefficient for forward mode in this cell.",
     )
-    B: EMECoefficientDataArray = pd.Field(
-        ...,
+
+    B: Optional[EMECoefficientDataArray] = pd.Field(
+        None,
         title="B coefficient",
         description="Coefficient for backward mode in this cell.",
     )
 
+    n_complex: Optional[EMEModeIndexDataArray] = pd.Field(
+        None,
+        title="Propagation Index",
+        description="Complex-valued effective propagation indices associated with the EME modes.",
+    )
+
+    flux: Optional[EMEFluxDataArray] = pd.Field(
+        None,
+        title="Flux",
+        description="Power flux of the EME modes.",
+    )
+
+    interface_S11: Optional[EMEInterfaceSMatrixDataArray] = pd.Field(
+        None,
+        title="Interface S11",
+        description="S matrix relating output modes at port 1 to input modes at port 1 "
+        "at a given interface between EME cells.",
+    )
+
+    interface_S12: Optional[EMEInterfaceSMatrixDataArray] = pd.Field(
+        None,
+        title="Interface S12",
+        description="S matrix relating output modes at port 1 to input modes at port 2 "
+        "at a given interface between EME cells.",
+    )
+
+    interface_S21: Optional[EMEInterfaceSMatrixDataArray] = pd.Field(
+        None,
+        title="Interface S21",
+        description="S matrix relating output modes at port 2 to input modes at port 1 "
+        "at a given interface between EME cells.",
+    )
+
+    interface_S22: Optional[EMEInterfaceSMatrixDataArray] = pd.Field(
+        None,
+        title="Interface S22",
+        description="S matrix relating output modes at port 2 to input modes at port 2 "
+        "at a given interface between EME cells.",
+    )
+
 
 class EMEFieldDataset(ElectromagneticFieldDataset):
     """Dataset storing scalar components of E and H fields as a function of freq, mode_index, and port_index."""

tidy3d/components/eme/data/sim_data.py

@@ -17,7 +17,7 @@
 from tidy3d.exceptions import SetupError
 from tidy3d.log import log
 
-from .dataset import EMESMatrixDataset
+from .dataset import EMECoefficientDataset, EMESMatrixDataset
 from .monitor_data import EMEFieldData, EMEModeSolverData, EMEMonitorDataType
 
 
@@ -39,6 +39,12 @@ class EMESimulationData(AbstractYeeGridSimulationData):
         None, title="S Matrix", description="Scattering matrix of the EME simulation."
     )
 
+    coeffs: Optional[EMECoefficientDataset] = pd.Field(
+        None,
+        title="Coefficients",
+        description="Coefficients from the EME simulation. Useful for debugging and optimization.",
+    )
+
     port_modes_raw: Optional[EMEModeSolverData] = pd.Field(
         None,
         title="Port Modes",

tidy3d/components/eme/simulation.py

@@ -752,6 +752,12 @@ def _validate_sweep_spec(self) -> None:
     def _validate_monitor_setup(self) -> None:
         """Check monitor setup."""
         for i, monitor in enumerate(self.monitors):
+            if isinstance(monitor, EMECoefficientMonitor):
+                log.warning(
+                    "'EMECoefficientMonitor' is deprecated. "
+                    "The full coefficient data is stored in "
+                    "'EMESimulationData.coeffs'."
+                )
             if isinstance(monitor, EMEMonitor):
                 _ = self._monitor_eme_cell_indices(monitor=monitor)
             if (