Making ModeData a collection of ModeAmpsData and ModeIndexData

tidy3d/components/data.py

@@ -147,7 +147,7 @@ def data(self) -> Tidy3dDataArray:
         # make DataArray
         data_dict = self.dict()
         coords = {dim: data_dict[dim] for dim in self._dims}
-        data_array = Tidy3dDataArray(self.values, coords=coords)
+        data_array = Tidy3dDataArray(self.values, coords=coords, dims=self._dims)
 
         # assign attrs for xarray
         if self.data_attrs:
@@ -585,7 +585,13 @@ class FluxTimeData(AbstractFluxData, TimeData):
     _dims = ("t",)
 
 
-class ModeData(PlanarData, FreqData):
+class AbstractModeData(PlanarData, FreqData, ABC):
+    """Abstract class for mode data as a function of frequency and mode index."""
+
+    mode_index: Array[int]
+
+
+class ModeAmpsData(AbstractModeData):
     """Stores modal amplitdudes from a :class:`ModeMonitor`.
 
     Parameters
@@ -606,14 +612,13 @@ class ModeData(PlanarData, FreqData):
 
     >>> f = np.linspace(2e14, 3e14, 1001)
     >>> values = (1+1j) * np.random.random((1, 2, len(f)))
-    >>> data = ModeData(values=values, direction=['+'], mode_index=np.arange(1, 3), f=f)
+    >>> data = ModeAmpsData(values=values, direction=['+'], mode_index=np.arange(1, 3), f=f)
     """
 
     direction: List[Direction] = ["+", "-"]
-    mode_index: Array[int]
     values: Array[complex]
     data_attrs: Dict[str, str] = {"units": "sqrt(W)", "long_name": "mode amplitudes"}
-    type: Literal["ModeData"] = "ModeData"
+    type: Literal["ModeAmpsData"] = "ModeAmpsData"
 
     _dims = ("direction", "mode_index", "f")
 
@@ -622,6 +627,94 @@ def normalize(self, source_freq_amps: Array[complex]) -> None:
         self.values /= 1j * source_freq_amps  # pylint: disable=no-member
 
 
+class ModeIndexData(AbstractModeData):
+    """Stores modal amplitdudes from a :class:`ModeMonitor`.
+
+    Parameters
+    ----------
+    mode_index : numpy.ndarray
+        Array of integer indices into the original monitor's :attr:`ModeMonitor.modes`.
+    f : numpy.ndarray
+        Frequency coordinates (Hz).
+    values : numpy.ndarray
+        Complex-valued array of mode amplitude values
+        with shape ``values.shape=(len(direction), len(mode_index), len(f))``.
+
+    Example
+    -------
+
+    >>> f = np.linspace(2e14, 3e14, 1001)
+    >>> values = (1+1j) * np.random.random((2, len(f)))
+    >>> data = ModeIndexData(values=values, mode_index=np.arange(1, 3), f=f)
+    """
+
+    values: Array[complex]
+    data_attrs: Dict[str, str] = {"units": "None", "long_name": "complex effective index"}
+    type: Literal["ModeIndexData"] = "ModeIndexData"
+
+    _dims = ("mode_index", "f")
+
+    def normalize(self, source_freq_amps: Array[complex]) -> None:
+        """normalize the values by the amplitude of the source."""
+        return
+
+
+class ModeData(CollectionData):
+    """Stores a collection of mode decomposition amplitudes and mode effective indexes for all
+    modes in a :class:`.ModeMonitor`.
+
+    Parameters
+    ----------
+    data_dict : Dict[str, :class:`ScalarFieldData`]
+        Mapping of field name (eg. 'Ex') to its scalar field data.
+
+    Example
+    -------
+    >>> f = np.linspace(1e14, 2e14, 1001)
+    >>> x = np.linspace(-1, 1, 10)
+    >>> y = np.linspace(-2, 2, 20)
+    >>> z = np.linspace(0, 0, 1)
+    >>> values = (1+1j) * np.random.random((len(x), len(y), len(z), len(f)))
+    >>> field = ScalarFieldData(values=values, x=x, y=y, z=z, f=f)
+    >>> data = FieldData(data_dict={'Ex': field, 'Ey': field})
+    """
+
+    data_dict: Dict[str, AbstractModeData]
+    type: Literal["ModeData"] = "ModeData"
+
+    @property
+    def amps(self):
+        """Get mode amplitudes."""
+        scalar_data = self.data_dict.get("amps")
+        if scalar_data:
+            return scalar_data.data
+        return None
+
+    @property
+    def n_complex(self):
+        """Get complex effective indexes."""
+        scalar_data = self.data_dict.get("n_complex")
+        if scalar_data:
+            return scalar_data.data
+        return None
+
+    @property
+    def n_eff(self):
+        """Get real part of effective index."""
+        scalar_data = self.data_dict.get("n_complex")
+        if scalar_data:
+            return scalar_data.data.real
+        return None
+
+    @property
+    def k_eff(self):
+        """Get imaginary part of effective index."""
+        scalar_data = self.data_dict.get("n_complex")
+        if scalar_data:
+            return scalar_data.data.imag
+        return None
+
+
 # maps MonitorData.type string to the actual type, for MonitorData.from_file()
 DATA_TYPE_MAP = {
     "ScalarFieldData": ScalarFieldData,
@@ -630,6 +723,8 @@ def normalize(self, source_freq_amps: Array[complex]) -> None:
     "FieldTimeData": FieldTimeData,
     "FluxData": FluxData,
     "FluxTimeData": FluxTimeData,
+    "ModeAmpsData": ModeAmpsData,
+    "ModeIndexData": ModeIndexData,
     "ModeData": ModeData,
 }
 
@@ -863,8 +958,9 @@ def normalize(self, normalize_index: int = 0):
         try:
             source = self.simulation.sources[normalize_index]
             source_time = source.source_time
-        except Exception as e:
-            raise DataError(f"Could not locate source at normalize_index={normalize_index}.") from e
+        except Exception:  # pylint:disable=broad-except
+            logging.warning(f"Could not locate source at normalize_index={normalize_index}.")
+            return self
 
         source_time = source.source_time
         sim_data_norm = self.copy(deep=True)
@@ -881,9 +977,9 @@ def normalize_data(monitor_data):
 
             if isinstance(monitor_data, (FieldData, FluxData, ModeData)):
 
-                if isinstance(monitor_data, FieldData):
-                    for scalar_field_data in monitor_data.data_dict.values():
-                        normalize_data(scalar_field_data)
+                if isinstance(monitor_data, CollectionData):
+                    for attr_data in monitor_data.data_dict.values():
+                        normalize_data(attr_data)
 
                 else:
                     normalize_data(monitor_data)

tidy3d/components/simulation.py

@@ -1213,7 +1213,7 @@ def make_eps_data(coords: Coords):
                 eps_structure = get_eps(structure.medium, freq)
                 is_inside = structure.geometry.inside(x, y, z)
                 eps_array[np.where(is_inside)] = eps_structure
-            return xr.DataArray(eps_array, coords={"x": xs, "y": ys, "z": zs})
+            return xr.DataArray(eps_array, coords={"x": xs, "y": ys, "z": zs}, dims=("x", "y", "z"))
 
         # combine all data into dictionary
         coords = sub_grid[coord_key]

tidy3d/convert.py

@@ -64,9 +64,6 @@ def old_json_parameters(sim: Simulation) -> Dict:
         "courant": sim.courant,
         "shutoff": sim.shutoff,
         "subpixel": sim.subpixel,
-        "time_steps": 100,
-        "nodes": 100,
-        "compute_weight": 1.0,
     }
 
     """ TODO: Support nonuniform coordinates """

tidy3d/plugins/mode/mode_solver.py

@@ -1,20 +1,22 @@
 """Turn Mode Specifications into Mode profiles 
 """
 
-from typing import List
+from typing import List, Dict
 from dataclasses import dataclass
 
 import numpy as np
 import xarray as xr
 
+from ...components.base import Tidy3dBaseModel
 from ...components import Box
 from ...components import Simulation
 from ...components import ModeSpec
 from ...components import ModeMonitor
 from ...components import ModeSource, GaussianPulse
 from ...components.types import Direction
-from ...components.data import ScalarFieldData, FieldData
+from ...components.data import ScalarFieldData, FieldData, Tidy3dData
 from ...log import SetupError
+from ...constants import C_0
 
 from .solver import compute_modes
 
@@ -50,8 +52,7 @@
 """
 
 
-@dataclass
-class ModeInfo:
+class ModeInfo(Tidy3dBaseModel):
     """stores information about a (solved) mode.
     Attributes
     ----------
@@ -110,15 +111,13 @@ def solve(self, mode_spec: ModeSpec) -> List[ModeInfo]:
 
         normal_axis = self.plane.size.index(0.0)
 
-        # note discretizing, need to make consistent
-        eps_xx = self.simulation.epsilon(self.plane, "Ex", self.freq)
-        eps_yy = self.simulation.epsilon(self.plane, "Ey", self.freq)
-        eps_zz = self.simulation.epsilon(self.plane, "Ez", self.freq)
+        # Get diagonal epsilon components in the plane
+        (eps_xx, eps_yy, eps_zz) = self.get_epsilon()
 
-        # make numpy array and get rid of normal axis
-        eps_xx = np.squeeze(eps_xx.values, axis=normal_axis)
-        eps_yy = np.squeeze(eps_yy.values, axis=normal_axis)
-        eps_zz = np.squeeze(eps_zz.values, axis=normal_axis)
+        # get rid of normal axis
+        eps_xx = np.squeeze(eps_xx, axis=normal_axis)
+        eps_yy = np.squeeze(eps_yy, axis=normal_axis)
+        eps_zz = np.squeeze(eps_zz, axis=normal_axis)
 
         # swap axes to waveguide coordinates (propagating in z)
         eps_wg_zz, (eps_wg_xx, eps_wg_yy) = self.plane.pop_axis(
@@ -162,6 +161,12 @@ def rotate_field_coords(e_field, h_field):
             # Get E and H fields at the current mode_index
             E, H = mode_fields[..., mode_index]
 
+            # Set gauge to highest-amplitude in-plane E being real and positive
+            ind_max = np.argmax(np.abs(E[:2]))
+            phi = np.angle(E[:2].ravel()[ind_max])
+            E *= np.exp(-1j * phi)
+            H *= np.exp(-1j * phi)
+
             # # Handle symmetries
             # if mode.symmetries[0] != 0:
             #     E_half = E[:, 1:, ...]
@@ -202,7 +207,7 @@ def rotate_field_coords(e_field, h_field):
                 )
 
             mode_info = ModeInfo(
-                field_data=FieldData(data_dict=data_dict).data,
+                field_data=FieldData(data_dict=data_dict),
                 mode_spec=mode_spec,
                 mode_index=mode_index,
                 n_eff=n_eff_complex[mode_index].real,
@@ -213,6 +218,15 @@ def rotate_field_coords(e_field, h_field):
 
         return modes
 
+    def get_epsilon(self):
+        """Compute the diagonal components of the epsilon tensor in the plane."""
+
+        eps_xx = self.simulation.epsilon(self.plane, "Ex", self.freq)
+        eps_yy = self.simulation.epsilon(self.plane, "Ey", self.freq)
+        eps_zz = self.simulation.epsilon(self.plane, "Ez", self.freq)
+
+        return np.stack((eps_xx, eps_yy, eps_zz), axis=0)
+
     # def make_source(self, mode_spec: ModeSpec, fwidth: float, direction: Direction) -> ModeSource:
     #     """Creates ``ModeSource`` from a Mode + additional specifications.
 

tidy3d/web/webapi.py

@@ -459,7 +459,7 @@ def _upload_task(  # pylint:disable=too-many-locals,too-many-arguments
     """upload with all kwargs exposed"""
 
     if solver_version[:6] == "revamp":
-        json_string = simulation.json()
+        json_string = simulation._json_string()  # pylint:disable=protected-access
     else:
         # convert to old json and get string version
         sim_dict = export_old_json(simulation)