assume sorted

Author: dbochkov-flexcompute

tests/test_components/test_mode_interp.py

@@ -657,6 +657,84 @@ def test_mode_solver_data_interp_extrapolation_warning():
         mode_data.interp_in_freq(freqs=freqs_extrap, method="linear")
 
 
+# ============================================================================
+# assume_sorted Tests (for source frequencies)
+# ============================================================================
+
+
+@pytest.mark.parametrize("method", ["linear", "cubic", "cheb"])
+def test_interp_assume_sorted_source_frequencies(method):
+    """Test that assume_sorted correctly handles sorted vs unsorted source frequencies.
+
+    This test verifies that:
+    - assume_sorted=True works correctly when source frequencies are sorted
+    - assume_sorted=False correctly handles unsorted source frequencies
+    - Both approaches produce identical interpolation results
+
+    Parameters
+    ----------
+    method : str
+        Interpolation method to test: "linear", "cubic", or "cheb"
+    """
+    from tidy3d.components.data.data_array import ModeIndexDataArray
+    from tidy3d.components.data.dataset import FreqDataset
+
+    # Define number of source frequencies based on method requirements
+    num_source_points = {"linear": 5, "cubic": 5, "cheb": 5}[method]
+
+    # Generate sorted source frequencies using ModeInterpSpec.sampling_points
+    interp_spec = td.ModeInterpSpec(num_points=num_source_points, method=method)
+    freqs_full_range = np.linspace(1e14, 2e14, 100)  # Full frequency range
+    freqs_source_sorted = interp_spec.sampling_points(freqs_full_range)
+
+    # Create unsorted version by shuffling
+    rng = np.random.RandomState(42)  # Fixed seed for reproducibility
+    shuffle_indices = np.arange(len(freqs_source_sorted))
+    rng.shuffle(shuffle_indices)
+    freqs_source_unsorted = freqs_source_sorted[shuffle_indices]
+
+    # Create test data values
+    mode_indices = np.arange(2)
+    data_values = (1.5 + 0.1j) * np.random.random((num_source_points, 2))
+
+    # Create dataset with SORTED source frequencies
+    data_sorted = ModeIndexDataArray(
+        data_values.copy(), coords={"f": freqs_source_sorted, "mode_index": mode_indices}
+    )
+
+    # Create dataset with UNSORTED source frequencies
+    # Need to reorder data values to match the shuffled frequencies
+    data_unsorted = ModeIndexDataArray(
+        data_values[shuffle_indices],
+        coords={"f": freqs_source_unsorted, "mode_index": mode_indices},
+    )
+
+    # Define target frequencies for interpolation
+    freqs_target = np.linspace(1e14, 2e14, 50)
+
+    # Interpolate with sorted source frequencies and assume_sorted=True
+    result_sorted = FreqDataset._interp_dataarray_in_freq(
+        data_sorted, freqs_target, method=method, assume_sorted=True
+    )
+
+    # Interpolate with unsorted source frequencies and assume_sorted=False
+    # This should internally sort the data before interpolating
+    result_unsorted = FreqDataset._interp_dataarray_in_freq(
+        data_unsorted, freqs_target, method=method, assume_sorted=False
+    )
+
+    # Both approaches should produce identical results
+    assert np.allclose(result_sorted.values, result_unsorted.values, rtol=1e-10), (
+        f"{method} interpolation: sorted and unsorted source frequencies "
+        "should produce identical results"
+    )
+
+    # Verify the interpolation actually happened correctly
+    assert result_sorted.shape == (50, 2)
+    assert len(result_sorted.coords["f"]) == 50
+    assert np.allclose(result_sorted.coords["f"].values, freqs_target)
+
+
 # ============================================================================
 # ModeSolver Integration Tests (Phase 5)
 # ============================================================================

tidy3d/components/data/dataset.py

@@ -57,6 +57,7 @@ def _interp_in_freq_update_dict(
         self,
         freqs: FreqArray,
         method: Literal["linear", "cubic", "cheb"] = "linear",
+        assume_sorted: bool = False,
     ) -> dict[str, DataArray]:
         """Interpolate mode data to new frequency points.
 
@@ -76,6 +77,8 @@ def _interp_in_freq_update_dict(
             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.
+        assume_sorted: bool = False,
+            Whether to assume the frequency points are sorted.
 
         Returns
         -------
@@ -112,7 +115,7 @@ def _interp_in_freq_update_dict(
 
         modify_data = {}
         for key, data in self.data_arrs.items():
-            modify_data[key] = self._interp_dataarray_in_freq(data, freqs, method)
+            modify_data[key] = self._interp_dataarray_in_freq(data, freqs, method, assume_sorted)
 
         return modify_data
 
@@ -121,6 +124,7 @@ def _interp_dataarray_in_freq(
         data: DataArray,
         freqs: FreqArray,
         method: Literal["linear", "cubic", "cheb", "nearest"],
+        assume_sorted: bool = False,
     ) -> DataArray:
         """Interpolate a DataArray along the frequency coordinate.
 
@@ -133,6 +137,8 @@ def _interp_dataarray_in_freq(
         method : Literal["linear", "cubic", "cheb", "nearest"]
             Interpolation method (``"linear"``, ``"cubic"``, ``"cheb"``, or ``"nearest"``).
             For ``"cheb"``, uses barycentric formula for Chebyshev interpolation.
+        assume_sorted: bool = False,
+            Whether to assume the frequency points are sorted.
 
         Returns
         -------
@@ -151,7 +157,7 @@ def _interp_dataarray_in_freq(
         else:
             if method != "cheb":
                 interp_kwargs["kwargs"] = {"fill_value": "extrapolate"}
-            return data.interp(f=freqs, **interp_kwargs)
+            return data.interp(f=freqs, assume_sorted=assume_sorted, **interp_kwargs)
 
 
 class ModeFreqDataset(FreqDataset, ABC):

tidy3d/components/data/monitor_data.py

@@ -2608,6 +2608,7 @@ def interp_in_freq(
         method: Literal["linear", "cubic", "cheb"] = "linear",
         renormalize: bool = True,
         recalculate_grid_correction: bool = True,
+        assume_sorted: bool = False,
     ) -> ModeSolverData:
         """Interpolate mode data to new frequency points.
 
@@ -2632,6 +2633,8 @@ def interp_in_freq(
         recalculate_grid_correction : bool = True
             Whether to recalculate the grid correction factors after interpolation or use interpolated
             grid corrections.
+        assume_sorted: bool = False,
+            Whether to assume the frequency points are sorted.
 
         Returns
         -------
@@ -2705,7 +2708,7 @@ def interp_in_freq(
             )
 
         # Build update dictionary
-        update_dict = self._interp_in_freq_update_dict(freqs, method)
+        update_dict = self._interp_in_freq_update_dict(freqs, method, assume_sorted)
 
         # Handle eps_spec if present - use nearest neighbor interpolation
         if self.eps_spec is not None:
@@ -2753,6 +2756,7 @@ def interpolated_copy(self) -> ModeSolverData:
             method=self.monitor.mode_spec.interp_spec.method,
             renormalize=True,
             recalculate_grid_correction=True,
+            assume_sorted=True,
         )
         return interpolated_data