nicer

Author: tylerflex

tidy3d/components/data/data_array.py

@@ -489,6 +489,25 @@ def _ag_interp_func(var, indexes_coords, method, **kwargs):
                 result = result.transpose(*out_dims)
         return result
 
+    def with_updated_data(
+        self,
+        data: float,
+        coords: dict[str, Any],
+    ) -> DataArray:
+        """Make copy of ``DataArray`` with ``data`` at specified ``coords``, autograd-approved."""
+        old_values = self.values
+        slice_indices = [slice(None)] * len(self.dims)  # Start with full slices for all dims
+        for dim_idx, dim_name in enumerate(self.dims):
+            if dim_name in coords:
+                coord_value = coords[dim_name]
+                pos = int(self.get_index(dim_name).get_loc(coord_value))
+                slice_indices[dim_idx] = pos
+        mask_slice = tuple(slice_indices)
+        mask = np.zeros_like(old_values, dtype=bool)
+        mask[mask_slice] = True
+        modified_values = np.where(mask, data, old_values)
+        return self.copy(deep=True, data=modified_values)
+
 
 class FreqDataArray(DataArray):
     """Frequency-domain array.

tidy3d/plugins/smatrix/component_modelers/modal.py

@@ -298,23 +298,6 @@ def _internal_construct_smatrix(self, batch_data: BatchData) -> ModalPortDataArr
         }
         s_matrix = ModalPortDataArray(values, coords=coords)
 
-        def set_new_values(
-            values: np.ndarray,
-            new_values: np.ndarray,
-            port_name_in: str,
-            mode_index_in: int,
-            port_name_out: str,
-            mode_index_out: int,
-        ) -> np.ndarray:
-            """Replace ``values`` with ``new_values`` at indices given by dims in ``s_matrix```"""
-            port_in_pos = int(s_matrix.get_index("port_in").get_loc(port_name_in))
-            mode_in_pos = int(s_matrix.get_index("mode_index_in").get_loc(mode_index_in))
-            port_out_pos = int(s_matrix.get_index("port_out").get_loc(port_name_out))
-            mode_out_pos = int(s_matrix.get_index("mode_index_out").get_loc(mode_index_out))
-            mask = np.zeros_like(values)
-            mask[port_out_pos, port_in_pos, mode_out_pos, mode_in_pos, :] = 1
-            return np.where(mask, new_values, values)
-
         # loop through source ports
         for col_index in self.matrix_indices_run_sim:
             port_name_in, mode_index_in = col_index
@@ -335,16 +318,14 @@ def set_new_values(
                 source_norm = self._normalization_factor(port_in, sim_data)
                 s_matrix_elements = np.array(amp.data) / np.array(source_norm)
 
-                values = set_new_values(
-                    values=values,
-                    new_values=s_matrix_elements,
-                    port_name_in=port_name_in,
-                    mode_index_in=mode_index_in,
-                    port_name_out=port_name_out,
-                    mode_index_out=mode_index_out,
-                )
+                coords_set = {
+                    "port_in": port_name_in,
+                    "mode_index_in": mode_index_in,
+                    "port_out": port_name_out,
+                    "mode_index_out": mode_index_out,
+                }
 
-        s_matrix = ModalPortDataArray(values, coords=coords)
+                s_matrix = s_matrix.with_updated_data(data=s_matrix_elements, coords=coords_set)
 
         # element can be determined by user-defined mapping
         for (row_in, col_in), (row_out, col_out), mult_by in self.element_mappings:
@@ -361,16 +342,13 @@ def set_new_values(
             port_in_to, mode_index_in_to = col_out
 
             elements_from = mult_by * s_matrix.loc[coords_from].values
+            coords_to = {
+                "port_in": port_in_to,
+                "mode_index_in": mode_index_in_to,
+                "port_out": port_out_to,
+                "mode_index_out": mode_index_out_to,
+            }
 
-            values = set_new_values(
-                values=values,
-                new_values=elements_from,
-                port_name_in=port_in_to,
-                mode_index_in=mode_index_in_to,
-                port_name_out=port_out_to,
-                mode_index_out=mode_index_out_to,
-            )
-
-            s_matrix = ModalPortDataArray(values, coords=coords)
+            s_matrix = s_matrix.with_updated_data(data=elements_from, coords=coords_to)
 
         return s_matrix