gregs refactor

Author: tylerflex

tidy3d/components/data/data_array.py

@@ -489,23 +489,23 @@ 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:
+    def with_updated_data(self, data, coords: dict[str, Any]) -> DataArray:
         """Make copy of ``DataArray`` with ``data`` at specified ``coords``, autograd-approved."""
+
+        mask = xr.zeros_like(self, dtype=bool)
+        mask.loc[coords] = True
+
         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)
+        replacement_data = np.zeros_like(old_values)
+        axes_to_fill = [self.get_axis_num(key) for key in coords]
+
+        reshape_guide = [1] * len(old_values.shape)
+        for idx, dim in enumerate(axes_to_fill):
+            reshape_guide[dim] = data.shape[idx]
+
+        replacement_data = replacement_data + data.reshape(reshape_guide)
+        modified_values = np.where(mask, replacement_data, old_values)
+
         return self.copy(deep=True, data=modified_values)
 
 

tidy3d/plugins/smatrix/component_modelers/modal.py

@@ -325,7 +325,8 @@ def _internal_construct_smatrix(self, batch_data: BatchData) -> ModalPortDataArr
                     "mode_index_out": mode_index_out,
                 }
 
-                s_matrix = s_matrix.with_updated_data(data=s_matrix_elements, coords=coords_set)
+                reshaped_smatrix_data = np.reshape(s_matrix_elements, (1, 1, 1, -1))
+                s_matrix = s_matrix.with_updated_data(data=reshaped_smatrix_data, 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:
@@ -349,6 +350,7 @@ def _internal_construct_smatrix(self, batch_data: BatchData) -> ModalPortDataArr
                 "mode_index_out": mode_index_out_to,
             }
 
-            s_matrix = s_matrix.with_updated_data(data=elements_from, coords=coords_to)
+            reshaped_elements_from = np.reshape(elements_from, (1, 1, 1, -1))
+            s_matrix = s_matrix.with_updated_data(data=reshaped_elements_from, coords=coords_to)
 
         return s_matrix