feat: add EMESimulationData.coeffs (FXC-4275)

Author: caseyflex

CHANGELOG.md

@@ -10,6 +10,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.
 - Klayout plugin automatically finds klayout installation path at common locations.
+- Added `EMESimulationData.coeffs` to store coefficients from the EME solver, including mode overlaps, 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.

schemas/EMESimulation.json

@@ -12879,6 +12879,10 @@
       },
       "type": "array"
     },
+    "store_coeffs": {
+      "default": true,
+      "type": "boolean"
+    },
     "store_port_modes": {
       "default": true,
       "type": "boolean"

tests/sims/full_fdtd.h5

@@ -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
@@ -731,6 +733,47 @@ def _get_eme_smatrix_data_array(num_modes_in=2, num_modes_out=3, num_freqs=2, nu
     return smatrix_entry
 
 
+def _get_eme_interface_smatrix_data_array(
+    num_modes_in=2, num_modes_out=3, num_freqs=2, num_sweep=0
+):
+    if num_modes_in != 0:
+        mode_index_in = np.arange(num_modes_in)
+    else:
+        mode_index_in = [0]
+    if num_modes_out != 0:
+        mode_index_out = np.arange(num_modes_out)
+    else:
+        mode_index_out = [0]
+    if num_sweep != 0:
+        sweep_index = np.arange(num_sweep)
+    else:
+        sweep_index = [0]
+    eme_cell_index = np.arange(3)
+
+    f = td.C_0 * np.linspace(1, 2, num_freqs)
+
+    data = (1 + 1j) * np.random.random(
+        (len(f), len(sweep_index), len(eme_cell_index), len(mode_index_out), len(mode_index_in))
+    )
+    coords = {
+        "f": f,
+        "sweep_index": sweep_index,
+        "eme_cell_index": eme_cell_index,
+        "mode_index_out": mode_index_out,
+        "mode_index_in": mode_index_in,
+    }
+    smatrix_entry = td.EMEInterfaceSMatrixDataArray(data, coords=coords)
+
+    if num_modes_in == 0:
+        smatrix_entry = smatrix_entry.drop_vars("mode_index_in")
+    if num_modes_out == 0:
+        smatrix_entry = smatrix_entry.drop_vars("mode_index_out")
+    if num_sweep == 0:
+        smatrix_entry = smatrix_entry.drop_vars("sweep_index")
+
+    return smatrix_entry
+
+
 def _get_eme_smatrix_dataset(num_modes_1=3, num_modes_2=4, num_sweep=0):
     S11 = _get_eme_smatrix_data_array(
         num_modes_in=num_modes_1, num_modes_out=num_modes_1, num_sweep=num_sweep
@@ -747,6 +790,35 @@ def _get_eme_smatrix_dataset(num_modes_1=3, num_modes_2=4, num_sweep=0):
     return td.EMESMatrixDataset(S11=S11, S12=S12, S21=S21, S22=S22)
 
 
+def _get_eme_interface_smatrix_dataset(num_modes_1=3, num_modes_2=4, num_sweep=0):
+    S11 = _get_eme_interface_smatrix_data_array(
+        num_modes_in=num_modes_1, num_modes_out=num_modes_1, num_sweep=num_sweep
+    )
+    S12 = _get_eme_interface_smatrix_data_array(
+        num_modes_in=num_modes_2, num_modes_out=num_modes_1, num_sweep=num_sweep
+    )
+    S21 = _get_eme_interface_smatrix_data_array(
+        num_modes_in=num_modes_1, num_modes_out=num_modes_2, num_sweep=num_sweep
+    )
+    S22 = _get_eme_interface_smatrix_data_array(
+        num_modes_in=num_modes_2, num_modes_out=num_modes_2, num_sweep=num_sweep
+    )
+    return td.EMEInterfaceSMatrixDataset(S11=S11, S12=S12, S21=S21, S22=S22)
+
+
+def _get_eme_overlaps_dataset(num_modes_1=3, num_modes_2=4, num_sweep=0):
+    O11 = _get_eme_interface_smatrix_data_array(
+        num_modes_in=num_modes_1, num_modes_out=num_modes_1, num_sweep=num_sweep
+    )
+    O12 = _get_eme_interface_smatrix_data_array(
+        num_modes_in=num_modes_2, num_modes_out=num_modes_1, num_sweep=num_sweep
+    )
+    O21 = _get_eme_interface_smatrix_data_array(
+        num_modes_in=num_modes_1, num_modes_out=num_modes_2, num_sweep=num_sweep
+    )
+    return td.EMEOverlapDataset(O11=O11, O12=O12, O21=O21)
+
+
 def _get_eme_coeff_data_array(num_sweep=0):
     f = [2e14]
     mode_index_out = [0, 1]
@@ -787,7 +859,26 @@ def _get_eme_coeff_data_array(num_sweep=0):
 def _get_eme_coeff_dataset(num_sweep=0):
     A = _get_eme_coeff_data_array(num_sweep=num_sweep)
     B = _get_eme_coeff_data_array(num_sweep=num_sweep)
-    return td.EMECoefficientDataset(A=A, B=B)
+    flux = _get_eme_flux_data_array(num_sweep=num_sweep)
+    n_complex = _get_eme_mode_index_data_array(num_sweep=num_sweep)
+    interface_smatrices = _get_eme_interface_smatrix_dataset(num_sweep=num_sweep)
+    overlaps = _get_eme_overlaps_dataset(num_sweep=num_sweep)
+    return td.EMECoefficientDataset(
+        A=A,
+        B=B,
+        flux=flux,
+        n_complex=n_complex,
+        interface_smatrices=interface_smatrices,
+        overlaps=overlaps,
+    )
+
+
+def test_eme_normalize_coeff_dataset():
+    coeffs = _get_eme_coeff_dataset()
+    coeffs_normalized = coeffs.normalized_copy
+    assert coeffs_normalized.flux is None
+    with pytest.raises(ValidationError):
+        _ = coeffs_normalized.normalized_copy
 
 
 def test_eme_coeff_data_array():
@@ -819,6 +910,29 @@ def _get_eme_mode_index_data_array(num_sweep=0):
     return data
 
 
+def _get_eme_flux_data_array(num_sweep=0):
+    f = [td.C_0, 3e14]
+    mode_index = np.arange(10)
+    eme_cell_index = np.arange(7)
+    if num_sweep != 0:
+        sweep_index = np.arange(num_sweep)
+    else:
+        sweep_index = [0]
+    coords = {
+        "f": f,
+        "sweep_index": sweep_index,
+        "eme_cell_index": eme_cell_index,
+        "mode_index": mode_index,
+    }
+    data = td.EMEFluxDataArray(
+        np.random.random((len(f), len(sweep_index), len(eme_cell_index), len(mode_index))),
+        coords=coords,
+    )
+    if num_sweep == 0:
+        data = data.drop_vars("sweep_index")
+    return data
+
+
 def test_eme_mode_index_data_array():
     _ = _get_eme_mode_index_data_array()
 
@@ -1305,7 +1419,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)]
         )

tests/test_components/test_eme.py

@@ -185,6 +185,8 @@
     ChargeDataArray,
     DiffractionDataArray,
     EMECoefficientDataArray,
+    EMEFluxDataArray,
+    EMEInterfaceSMatrixDataArray,
     EMEModeIndexDataArray,
     EMEScalarFieldDataArray,
     EMEScalarModeFieldDataArray,
@@ -243,7 +245,9 @@
 from .components.eme.data.dataset import (
     EMECoefficientDataset,
     EMEFieldDataset,
+    EMEInterfaceSMatrixDataset,
     EMEModeSolverDataset,
+    EMEOverlapDataset,
     EMESMatrixDataset,
 )
 from .components.eme.data.monitor_data import EMECoefficientData, EMEFieldData, EMEModeSolverData
@@ -605,8 +609,11 @@ def set_logging_level(level: str) -> None:
     "EMEFieldData",
     "EMEFieldDataset",
     "EMEFieldMonitor",
+    "EMEFluxDataArray",
     "EMEFreqSweep",
     "EMEGrid",
+    "EMEInterfaceSMatrixDataArray",
+    "EMEInterfaceSMatrixDataset",
     "EMELengthSweep",
     "EMEModeIndexDataArray",
     "EMEModeSolverData",
@@ -615,6 +622,7 @@ def set_logging_level(level: str) -> None:
     "EMEModeSpec",
     "EMEModeSweep",
     "EMEMonitor",
+    "EMEOverlapDataset",
     "EMEPeriodicitySweep",
     "EMESMatrixDataArray",
     "EMESMatrixDataset",

tidy3d/__init__.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,24 @@ 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]
+    >>> sweep_index = np.arange(2)
+    >>> eme_cell_index = np.arange(5)
+    >>> mode_index = np.arange(4)
+    >>> coords = dict(f=f, sweep_index=sweep_index, eme_cell_index=eme_cell_index, mode_index=mode_index)
+    >>> data = EMEFluxDataArray(np.random.random((2,2,5,4)), 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 +1638,10 @@ def _make_impedance_data_array(result: DataArray) -> ImpedanceResultType:
     EMEScalarFieldDataArray,
     EMEScalarModeFieldDataArray,
     EMESMatrixDataArray,
+    EMEInterfaceSMatrixDataArray,
     EMECoefficientDataArray,
     EMEModeIndexDataArray,
+    EMEFluxDataArray,
     EMEFreqModeDataArray,
     ChargeDataArray,
     SteadyVoltageDataArray,