PR comments

Author: dbochkov-flexcompute

tests/test_components/test_mode_interp.py

@@ -3,14 +3,14 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Optional, Union, get_args
+from typing import Any, Callable, Literal, Optional, Union, get_args
 
 import numpy as np
 import pydantic.v1 as pd
 import xarray as xr
 
 from tidy3d.components.base import Tidy3dBaseModel
-from tidy3d.components.types import Axis, xyz
+from tidy3d.components.types import Axis, FreqArray, xyz
 from tidy3d.constants import C_0, PICOSECOND_PER_NANOMETER_PER_KILOMETER, UnitScaling
 from tidy3d.exceptions import DataError
 from tidy3d.log import log
@@ -50,6 +50,139 @@ def data_arrs(self) -> dict:
         return data_arrs
 
 
+class FreqDataset(Dataset, ABC):
+    """Abstract base class for objects that store collections of `:class:`.DataArray`s."""
+
+    def _interp_in_freq_update_dict(
+        self,
+        freqs: FreqArray,
+        method: Literal["linear", "cubic", "cheb"] = "linear",
+    ) -> dict[str, DataArray]:
+        """Interpolate mode data to new frequency points.
+
+        Interpolates all stored mode data (effective indices, field components, group indices,
+        and dispersion) from the current frequency grid to a new set of frequencies. This is
+        useful for obtaining mode data at many frequencies from computations at fewer frequencies,
+        when modes vary smoothly with frequency.
+
+        Parameters
+        ----------
+        freqs : FreqArray
+            New frequency points to interpolate to. Should generally span a similar range
+            as the original frequencies to avoid extrapolation.
+        method : Literal["linear", "cubic", "cheb"]
+            Interpolation method. ``"linear"`` for linear interpolation (requires 2+ source
+            frequencies), ``"cubic"`` for cubic spline interpolation (requires 4+ source
+            frequencies), ``"cheb"`` for Chebyshev polynomial interpolation using barycentric
+            formula (requires 3+ source frequencies at Chebyshev nodes).
+            For complex-valued data, real and imaginary parts are interpolated independently.
+
+        Returns
+        -------
+        ModeSolverData
+            New :class:`ModeSolverData` object with data interpolated to the requested frequencies.
+
+        Raises
+        ------
+        DataError
+            If interpolation parameters are invalid (e.g., too few source frequencies for the
+            chosen method, or source frequencies not at Chebyshev nodes for 'cheb' method).
+
+        Note
+        ----
+            Interpolation assumes modes vary smoothly with frequency. Results may be inaccurate
+            near mode crossings or regions of rapid mode variation. Use frequency tracking
+            (``mode_spec.sort_spec.track_freq``) to help maintain mode ordering consistency.
+
+            For Chebyshev interpolation, source frequencies must be at Chebyshev nodes of the
+            second kind within the frequency range.
+
+        Example
+        -------
+        >>> # Compute modes at 5 frequencies
+        >>> import numpy as np
+        >>> freqs_sparse = np.linspace(1e14, 2e14, 5)
+        >>> # ... create mode_solver and compute modes ...
+        >>> # mode_data = mode_solver.solve()
+        >>> # Interpolate to 50 frequencies
+        >>> freqs_dense = np.linspace(1e14, 2e14, 50)
+        >>> # mode_data_interp = mode_data.interp(freqs=freqs_dense, method='linear')
+        """
+        freqs = np.array(freqs)
+
+        modify_data = {}
+        for key, data in self.data_arrs.items():
+            modify_data[key] = self._interp_dataarray_in_freq(data, freqs, method)
+
+        return modify_data
+
+    @staticmethod
+    def _interp_dataarray_in_freq(
+        data: DataArray,
+        freqs: FreqArray,
+        method: Literal["linear", "cubic", "cheb", "nearest"],
+    ) -> DataArray:
+        """Interpolate a DataArray along the frequency coordinate.
+
+        Parameters
+        ----------
+        data : DataArray
+            Data array to interpolate. Must have a frequency coordinate ``"f"``.
+        freqs : FreqArray
+            New frequency points.
+        method : Literal["linear", "cubic", "cheb", "nearest"]
+            Interpolation method (``"linear"``, ``"cubic"``, ``"cheb"``, or ``"nearest"``).
+            For ``"cheb"``, uses barycentric formula for Chebyshev interpolation.
+
+        Returns
+        -------
+        DataArray
+            Interpolated data array with the same structure but new frequency points.
+        """
+        # Map 'cheb' to xarray's 'barycentric' method
+        xr_method = "barycentric" if method == "cheb" else method
+
+        # Use xarray's built-in interpolation
+        # For complex data, this automatically interpolates real and imaginary parts
+        interp_kwargs = {"method": xr_method}
+
+        if method == "nearest":
+            return data.sel(f=freqs, method="nearest")
+        else:
+            if method != "cheb":
+                interp_kwargs["kwargs"] = {"fill_value": "extrapolate"}
+            return data.interp(f=freqs, **interp_kwargs)
+
+
+class ModeFreqDataset(FreqDataset, ABC):
+    """Abstract base class for objects that store collections of `:class:`.DataArray`s."""
+
+    def _apply_mode_reorder(self, sort_inds_2d):
+        """Apply a mode reordering along mode_index for all frequency indices.
+
+        Parameters
+        ----------
+        sort_inds_2d : np.ndarray
+            Array of shape (num_freqs, num_modes) where each row is the
+            permutation to apply to the mode_index for that frequency.
+        """
+        num_freqs, num_modes = sort_inds_2d.shape
+        modify_data = {}
+        for key, data in self.data_arrs.items():
+            if "mode_index" not in data.dims or "f" not in data.dims:
+                continue
+            dims_orig = data.dims
+            f_coord = data.coords["f"]
+            slices = []
+            for ifreq in range(num_freqs):
+                sl = data.isel(f=ifreq, mode_index=sort_inds_2d[ifreq])
+                slices.append(sl.assign_coords(mode_index=np.arange(num_modes)))
+            # Concatenate along the 'f' dimension name and then restore original frequency coordinates
+            data = xr.concat(slices, dim="f").assign_coords(f=f_coord).transpose(*dims_orig)
+            modify_data[key] = data
+        return self.updated_copy(**modify_data)
+
+
 class AbstractFieldDataset(Dataset, ABC):
     """Collection of scalar fields with some symmetry properties."""
 
@@ -492,7 +625,7 @@ class AuxFieldTimeDataset(AuxFieldDataset):
     )
 
 
-class ModeSolverDataset(ElectromagneticFieldDataset):
+class ModeSolverDataset(ElectromagneticFieldDataset, ModeFreqDataset):
     """Dataset storing scalar components of E and H fields as a function of freq. and mode_index.
 
     Example

tidy3d/components/data/dataset.py

@@ -2393,6 +2393,12 @@ def normalize(self, source_spectrum_fn: Callable[[float], complex]) -> ModeSolve
         """Return copy of self after normalization is applied using source spectrum function."""
         return self.copy()
 
+    def _normalize_modes(self):
+        """Normalize modes. Note: this modifies ``self`` in-place."""
+        scaling = np.sqrt(np.abs(self.flux))
+        for field in self.field_components.values():
+            field /= scaling
+
     @staticmethod
     def _validate_cheb_nodes(freqs: np.ndarray) -> None:
         """Validate that frequencies are approximately at Chebyshev nodes.
@@ -2415,22 +2421,21 @@ def _validate_cheb_nodes(freqs: np.ndarray) -> None:
         freqs_sorted = np.sort(freqs)
         expected_sorted = np.sort(expected_freqs)
 
-        # Check relative error
-        freq_range = np.abs(expected_freqs[-1] - expected_freqs[0])
-        max_error = np.max(np.abs(freqs_sorted - expected_sorted)) / freq_range
-
-        if max_error > CHEB_NODES_TOLERANCE:
+        # Check if frequencies are close to Chebyshev nodes
+        if not np.allclose(
+            freqs_sorted, expected_sorted, atol=CHEB_NODES_TOLERANCE, rtol=CHEB_NODES_TOLERANCE
+        ):
             raise DataError(
-                f"For Chebyshev interpolation ('cheb'), source frequencies must be at "
-                f"Chebyshev nodes of the second kind. Maximum relative error: {max_error:.2e}, "
-                f"tolerance: {CHEB_NODES_TOLERANCE:.2e}. Use ModeInterpSpec.sampling_points() to generate "
-                f"appropriate frequencies."
+                "For Chebyshev interpolation ('cheb'), source frequencies must be at "
+                "Chebyshev nodes of the second kind. Use 'ModeInterpSpec' to generate "
+                "appropriate frequencies."
             )
 
-    def interp(
+    def interp_in_freq(
         self,
         freqs: FreqArray,
         method: Literal["linear", "cubic", "cheb"] = "linear",
+        renormalize: Optional[bool] = False,
     ) -> ModeSolverData:
         """Interpolate mode data to new frequency points.
 
@@ -2450,6 +2455,8 @@ def interp(
             frequencies), ``"cheb"`` for Chebyshev polynomial interpolation using barycentric
             formula (requires 3+ source frequencies at Chebyshev nodes).
             For complex-valued data, real and imaginary parts are interpolated independently.
+        renormalize : Optional[bool] = False
+            Whether to renormalize the mode profiles to unity power after interpolation.
 
         Returns
         -------
@@ -2507,36 +2514,27 @@ def interp(
                 f"Invalid interpolation method '{method}'. Use 'linear', 'cubic', or 'cheb'."
             )
 
-        # Build update dictionary
-        update_dict = {}
-
-        # Interpolate n_complex (required field)
-        update_dict["n_complex"] = self._interp_dataarray(self.n_complex, freqs, method)
-
-        # Interpolate field components if present
-        for field_name, field_data in self.field_components.items():
-            if field_data is not None:
-                update_dict[field_name] = self._interp_dataarray(field_data, freqs, method)
-
-        # Interpolate n_group_raw if present
-        if self.n_group_raw is not None:
-            update_dict["n_group_raw"] = self._interp_dataarray(self.n_group_raw, freqs, method)
+        # Check if we're extrapolating significantly and warn
+        freq_min, freq_max = np.min(source_freqs), np.max(source_freqs)
+        new_freq_min, new_freq_max = np.min(freqs), np.max(freqs)
 
-        # Interpolate dispersion_raw if present
-        if self.dispersion_raw is not None:
-            update_dict["dispersion_raw"] = self._interp_dataarray(
-                self.dispersion_raw, freqs, method
+        if new_freq_min < freq_min * (
+            1 - MODE_INTERP_EXTRAPOLATION_TOLERANCE
+        ) or new_freq_max > freq_max * (1 + MODE_INTERP_EXTRAPOLATION_TOLERANCE):
+            log.warning(
+                f"Interpolating to frequencies outside original range "
+                f"[{freq_min:.3e}, {freq_max:.3e}] Hz. New range: "
+                f"[{new_freq_min:.3e}, {new_freq_max:.3e}] Hz. "
+                "Results may be inaccurate due to extrapolation."
             )
 
-        # Interpolate grid correction data if present
-        for key, data in self._grid_correction_dict.items():
-            if isinstance(data, DataArray) and "f" in data.coords:
-                update_dict[key] = self._interp_dataarray(data, freqs, method)
+        # Build update dictionary
+        update_dict = self._interp_in_freq_update_dict(freqs, method)
 
         # Handle eps_spec if present - use nearest neighbor interpolation
         if self.eps_spec is not None:
             update_dict["eps_spec"] = list(
-                self._interp_dataarray(
+                self._interp_dataarray_in_freq(
                     FreqDataArray(self.eps_spec, coords={"f": self.monitor.freqs}),
                     freqs,
                     "nearest",
@@ -2546,57 +2544,13 @@ def interp(
         # Update monitor with new frequencies
         update_dict["monitor"] = self.monitor.updated_copy(freqs=list(freqs))
 
-        return self.copy(update=update_dict)
+        updated_data = self.updated_copy(**update_dict)
+        # print(updated_data.poynting)
+        # print(updated_data._diff_area)
+        if renormalize:
+            updated_data._normalize_modes()
 
-    @staticmethod
-    def _interp_dataarray(
-        data: DataArray,
-        freqs: FreqArray,
-        method: str,
-    ) -> DataArray:
-        """Interpolate a DataArray along the frequency coordinate.
-
-        Parameters
-        ----------
-        data : DataArray
-            Data array to interpolate. Must have a frequency coordinate ``"f"``.
-        freqs : FreqArray
-            New frequency points.
-        method : str
-            Interpolation method (``"linear"``, ``"cubic"``, or ``"cheb"``).
-            For ``"cheb"``, uses barycentric formula for Chebyshev interpolation.
-
-        Returns
-        -------
-        DataArray
-            Interpolated data array with the same structure but new frequency points.
-        """
-        # Map 'cheb' to xarray's 'barycentric' method
-        xr_method = "barycentric" if method == "cheb" else method
-
-        # Use xarray's built-in interpolation
-        # For complex data, this automatically interpolates real and imaginary parts
-        interp_kwargs = {"method": xr_method}
-
-        # Check if we're extrapolating significantly and warn
-        freq_min, freq_max = float(data.coords["f"].min()), float(data.coords["f"].max())
-        new_freq_min, new_freq_max = float(freqs.min()), float(freqs.max())
-
-        if new_freq_min < freq_min * (
-            1 - MODE_INTERP_EXTRAPOLATION_TOLERANCE
-        ) or new_freq_max > freq_max * (1 + MODE_INTERP_EXTRAPOLATION_TOLERANCE):
-            log.warning(
-                f"Interpolating to frequencies outside original range "
-                f"[{freq_min:.3e}, {freq_max:.3e}] Hz. New range: "
-                f"[{new_freq_min:.3e}, {new_freq_max:.3e}] Hz. "
-                "Results may be inaccurate due to extrapolation."
-            )
-            interp_kwargs["kwargs"] = {"fill_value": "extrapolate"}
-
-        if method == "nearest":
-            return data.sel(f=freqs, method="nearest")
-        else:
-            return data.interp(f=freqs, **interp_kwargs)
+        return updated_data
 
     @property
     def time_reversed_copy(self) -> FieldData:

tidy3d/components/data/monitor_data.py

@@ -9,10 +9,10 @@
     ImpedanceFreqModeDataArray,
     VoltageFreqModeDataArray,
 )
-from tidy3d.components.data.dataset import Dataset
+from tidy3d.components.data.dataset import ModeFreqDataset
 
 
-class TransmissionLineDataset(Dataset):
+class TransmissionLineDataset(ModeFreqDataset):
     """Holds mode data that is specific to transmission lines in microwave and RF applications,
     like characteristic impedance.