fix: Updates to speed up some mode data operations

Author: momchil-flex

tidy3d/components/data/monitor_data.py

@@ -149,20 +149,6 @@ def amplitude_fn(freq: list[float]) -> complex:
 
         return self.normalize(amplitude_fn)
 
-    def _updated(self, update: dict) -> MonitorData:
-        """Similar to ``updated_copy``, but does not actually copy components, for speed.
-
-        Note
-        ----
-            This does **not** produce a copy of mutable objects, so e.g. if some of the data arrays
-            are not updated, they will point to the values in the original data. This method should
-            thus be used carefully.
-
-        """
-        data_dict = self.dict()
-        data_dict.update(update)
-        return type(self).parse_obj(data_dict)
-
     def _make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[Source]:
         """Generate adjoint sources for this ``MonitorData`` instance."""
 
@@ -261,7 +247,7 @@ def symmetry_expanded(self):
         if all(sym == 0 for sym in self.symmetry):
             return self
 
-        return self._updated(self._symmetry_update_dict)
+        return self.updated_copy(**self._symmetry_update_dict, deep=False, validate=False)
 
     @property
     def symmetry_expanded_copy(self) -> AbstractFieldData:
@@ -780,21 +766,38 @@ def dot(
             fields_self = {key: field.conj() for key, field in fields_self.items()}
 
         fields_other = field_data._interpolated_tangential_fields(self._plane_grid_boundaries)
+        dim1, dim2 = self._tangential_dims
+        d_area = self._diff_area
 
-        # Drop size-1 dimensions in the other data
-        fields_other = {key: field.squeeze(drop=True) for key, field in fields_other.items()}
+        try:
+            for field_self, field_other in zip(fields_self.values(), fields_other.values()):
+                for key, val in field_self.coords:
+                    if not np.all(val.values == field_other.coords[key].values):
+                        raise ValueError("Coordinates do not match.")
+
+            # Coordinates match, so we can use .values for speed
+            e_self_x_h_other = fields_self["E" + dim1].values * fields_other["H" + dim2].values
+            e_self_x_h_other -= fields_self["E" + dim2].values * fields_other["H" + dim1].values
+            h_self_x_e_other = fields_self["H" + dim1].values * fields_other["E" + dim2].values
+            h_self_x_e_other -= fields_self["H" + dim2].values * fields_other["E" + dim1].values
+            integrand = xr.DataArray(
+                e_self_x_h_other - h_self_x_e_other, coords=fields_self["E" + dim1].coords
+            )
+            integrand *= d_area
+        except:  # noqa: E722
+            # Catching a broad exception here in case anything went wrong in the check.
 
-        # Cross products of fields
-        dim1, dim2 = self._tangential_dims
-        e_self_x_h_other = fields_self["E" + dim1] * fields_other["H" + dim2]
-        e_self_x_h_other -= fields_self["E" + dim2] * fields_other["H" + dim1]
-        h_self_x_e_other = fields_self["H" + dim1] * fields_other["E" + dim2]
-        h_self_x_e_other -= fields_self["H" + dim2] * fields_other["E" + dim1]
+            # Drop size-1 dimensions in the other data
+            fields_other = {key: field.squeeze(drop=True) for key, field in fields_other.items()}
 
-        # Integrate over plane
-        d_area = self._diff_area
-        integrand = (e_self_x_h_other - h_self_x_e_other) * d_area
+            # Cross products of fields
+            e_self_x_h_other = fields_self["E" + dim1] * fields_other["H" + dim2]
+            e_self_x_h_other -= fields_self["E" + dim2] * fields_other["H" + dim1]
+            h_self_x_e_other = fields_self["H" + dim1] * fields_other["E" + dim2]
+            h_self_x_e_other -= fields_self["H" + dim2] * fields_other["E" + dim1]
+            integrand = (e_self_x_h_other - h_self_x_e_other) * d_area
 
+        # Integrate over plane
         return ModeAmpsDataArray(0.25 * integrand.sum(dim=d_area.dims))
 
     def _interpolated_tangential_fields(self, coords: ArrayFloat2D) -> dict[str, DataArray]:
@@ -811,6 +814,19 @@ def _interpolated_tangential_fields(self, coords: ArrayFloat2D) -> dict[str, Dat
         """
         fields = self._tangential_fields
 
+        try:
+            # If coords already match, just return the tangential fields directly
+            for field in fields.values():
+                for idim, dim in enumerate(self._tangential_dims):
+                    if field.coords[dim].values.size != coords[idim].size or not np.all(
+                        field.coords[dim].values == coords[idim]
+                    ):
+                        raise ValueError("Coordinates do not match.")
+            return fields
+        except:  # noqa: E722
+            # Catching a broad exception here in case anything went wrong in the check.
+            pass
+
         # Interpolate if data has more than one coordinate along a dimension
         interp_dict = {"assume_sorted": True}
         # If single coordinate, just sel "nearest", i.e. just propagate the same data everywhere
@@ -1711,10 +1727,12 @@ def overlap_sort(
 
         # Normalizing the flux to 1, does not guarantee self terms of overlap integrals
         # are also normalized to 1 when the non-conjugated product is used.
-        if self.monitor.conjugated_dot_product:
+        data_expanded = self.symmetry_expanded
+        if data_expanded.monitor.conjugated_dot_product:
             self_overlap = np.ones((num_freqs, num_modes))
         else:
-            self_overlap = np.abs(self.dot(self, self.monitor.conjugated_dot_product).values)
+            self_overlap = data_expanded.dot(data_expanded, self.monitor.conjugated_dot_product)
+            self_overlap = np.abs(self_overlap.values)
             threshold_array = overlap_thresh * self_overlap
 
         # Compute sorting order and overlaps with neighboring frequencies
@@ -1727,20 +1745,19 @@ def overlap_sort(
         # Sort in two directions from the base frequency
         for step, last_ind in zip([-1, 1], [-1, num_freqs]):
             # Start with the base frequency
-            data_template = self._isel(f=[f0_ind])
+            data_template = data_expanded._isel(f=[f0_ind])
 
             # March to lower/higher frequencies
             for freq_id in range(f0_ind + step, last_ind, step):
                 # Calculate threshold array for this frequency
-                if not self.monitor.conjugated_dot_product:
+                if not data_expanded.monitor.conjugated_dot_product:
                     overlap_thresh = threshold_array[freq_id, :]
                 # Get next frequency to sort
-                data_to_sort = self._isel(f=[freq_id])
+                data_to_sort = data_expanded._isel(f=[freq_id])
                 # Assign to the base frequency so that outer_dot will compare them
                 data_to_sort = data_to_sort._assign_coords(f=[self.monitor.freqs[f0_ind]])
 
                 # Compute "sorting w.r.t. to neighbor" and overlap values
-
                 sorting_one_mode, amps_one_mode = data_template._find_ordering_one_freq(
                     data_to_sort, overlap_thresh
                 )
@@ -1756,8 +1773,8 @@ def overlap_sort(
                 for mode_ind in list(np.nonzero(overlap[freq_id, :] < overlap_thresh)[0]):
                     log.warning(
                         f"Mode '{mode_ind}' appears to undergo a discontinuous change "
-                        f"between frequencies '{self.monitor.freqs[freq_id]}' "
-                        f"and '{self.monitor.freqs[freq_id - step]}' "
+                        f"between frequencies '{data_expanded.monitor.freqs[freq_id]}' "
+                        f"and '{data_expanded.monitor.freqs[freq_id - step]}' "
                         f"(overlap: '{overlap[freq_id, mode_ind]:.2f}')."
                     )
 
@@ -1796,7 +1813,7 @@ def _isel(self, **isel_kwargs):
             for key, field in update_dict.items()
             if isinstance(field, DataArray)
         }
-        return self._updated(update=update_dict)
+        return self.updated_copy(**update_dict, deep=False, validate=False)
 
     def _assign_coords(self, **assign_coords_kwargs):
         """Wraps ``xarray.DataArray.assign_coords`` for all data fields that are defined over frequency and
@@ -1808,7 +1825,7 @@ def _assign_coords(self, **assign_coords_kwargs):
         update_dict = {
             key: field.assign_coords(**assign_coords_kwargs) for key, field in update_dict.items()
         }
-        return self._updated(update=update_dict)
+        return self.updated_copy(**update_dict, deep=False, validate=False)
 
     def _find_ordering_one_freq(
         self,
@@ -2214,20 +2231,58 @@ def _apply_mode_reorder(self, sort_inds_2d):
             Array of shape (num_freqs, num_modes) where each row is the
             permutation to apply to the mode_index for that frequency.
         """
+        sort_inds_2d = np.asarray(sort_inds_2d, dtype=int)
         num_freqs, num_modes = sort_inds_2d.shape
+
+        # Fast no-op
+        identity = np.arange(num_modes)
+        if np.all(sort_inds_2d == identity[None, :]):
+            return self
+
         modify_data = {}
+        new_mode_index_coord = identity
+
         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
+
+            dims_orig = tuple(data.dims)
+            # Preserve coords (as numpy)
+            coords_out = {
+                k: (v.values if hasattr(v, "values") else np.asarray(v))
+                for k, v in data.coords.items()
+            }
+            f_axis = data.get_axis_num("f")
+            m_axis = data.get_axis_num("mode_index")
+
+            # Move axes so array is (..., f, mode)
+            move_order = [ax for ax in range(data.ndim) if ax not in (f_axis, m_axis)] + [
+                f_axis,
+                m_axis,
+            ]
+            arr = np.moveaxis(data.data, move_order, range(data.ndim))
+            lead_shape = arr.shape[:-2]
+            nf, nm = arr.shape[-2], arr.shape[-1]
+            if nf != num_freqs or nm != num_modes:
+                raise DataError(
+                    "sort_inds_2d shape does not match array shape in _apply_mode_reorder."
+                )
+
+            # Vectorized gather: reshape to (nf, Nlead, nm), gather along last axis
+            arr3 = arr.reshape((-1, nf, nm)).transpose(1, 0, 2)  # (nf, Nlead, nm)
+            inds = sort_inds_2d[:, None, :]  # (nf, 1, nm)
+            arr3_sorted = np.take_along_axis(arr3, inds, axis=2)
+            arr_sorted = arr3_sorted.transpose(1, 0, 2).reshape(*lead_shape, nf, nm)
+
+            # Move axes back to original order
+            arr_sorted = np.moveaxis(arr_sorted, range(data.ndim), move_order)
+
+            # Update coords: keep f, reset mode_index to 0..num_modes-1
+            coords_out["mode_index"] = new_mode_index_coord
+            coords_out["f"] = data.coords["f"].values
+
+            modify_data[key] = DataArray(arr_sorted, coords=coords_out, dims=dims_orig)
+
         return self.updated_copy(**modify_data)
 
     def sort_modes(

tidy3d/components/mode/mode_solver.py

@@ -1334,8 +1334,7 @@ def _colocate_data(self, mode_solver_data: ModeSolverData) -> ModeSolverData:
         mode_solver_monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME)
         grid_expanded = self.simulation.discretize_monitor(mode_solver_monitor)
         data_dict_colocated.update({"monitor": mode_solver_monitor, "grid_expanded": grid_expanded})
-        mode_solver_data = mode_solver_data._updated(update=data_dict_colocated)
-
+        mode_solver_data = mode_solver_data.updated_copy(**data_dict_colocated, deep=False)
         return mode_solver_data
 
     def _normalize_modes(self, mode_solver_data: ModeSolverData):