better solution to with_updated_copy

Author: tylerflex

tidy3d/components/data/data_array.py

@@ -489,24 +489,37 @@ def _ag_interp_func(var, indexes_coords, method, **kwargs):
                 result = result.transpose(*out_dims)
         return result
 
-    def with_updated_data(self, data, coords: dict[str, Any]) -> DataArray:
+    def with_updated_data(self, data: xr, coords: dict[str, Any]) -> DataArray:
         """Make copy of ``DataArray`` with ``data`` at specified ``coords``, autograd-approved."""
 
+        old_values = self.values
+        shape_in = data.shape
+        shape_in_new = []
+        i = 0
+        for dim in self.dims:
+            if dim in coords:
+                shape_in_new.append(1)
+            else:
+                shape_in_new.append(shape_in[i])
+                i += 1
+
+        data = data.reshape(shape_in_new)
+        data = data + np.zeros_like(old_values)
+
         mask = xr.zeros_like(self, dtype=bool)
         mask.loc[coords] = True
 
-        old_values = self.values
-        replacement_data = np.zeros_like(old_values)
-        axes_to_fill = [self.get_axis_num(key) for key in coords]
+        modified_values = np.where(mask, data, old_values)
 
-        reshape_guide = [1] * len(old_values.shape)
-        for idx, dim in enumerate(axes_to_fill):
-            reshape_guide[dim] = data.shape[idx]
+        return self.copy(deep=True, data=modified_values)
 
-        replacement_data = replacement_data + data.reshape(reshape_guide)
-        modified_values = np.where(mask, replacement_data, old_values)
+    def _with_updated_data(self, data: DataArray, coords: dict[str, Any]) -> DataArray:
+        """Make copy of ``DataArray`` with ``data`` at specified ``coords``, autograd-approved."""
 
-        return self.copy(deep=True, data=modified_values)
+        mask = xr.zeros_like(self, dtype=bool)
+        mask.loc[coords] = True
+
+        return xr.where(mask, data, self)
 
 
 class FreqDataArray(DataArray):

tidy3d/plugins/smatrix/component_modelers/modal.py

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

tidy3d/plugins/smatrix/component_modelers/terminal.py

@@ -224,12 +224,9 @@ def _internal_construct_smatrix(self, batch_data: BatchData) -> TerminalPortData
         for _port_in_idx, port_in in enumerate(self.ports):
             sim_data = batch_data[self._task_name(port=port_in)]
             a, b = self.compute_power_wave_amplitudes_at_each_port(port_impedances, sim_data)
-            indexer = {"f": a.f, "port_in": port_in.name, "port_out": a.port}
-
-            a_data = np.expand_dims(a.data, axis=1)
-            b_data = np.expand_dims(b.data, axis=1)
-            a_matrix = a_matrix.with_updated_data(data=a_data, coords=indexer)
-            b_matrix = b_matrix.with_updated_data(data=b_data, coords=indexer)
+            indexer = {"port_in": port_in.name}
+            a_matrix = a_matrix.with_updated_data(data=a.data, coords=indexer)
+            b_matrix = b_matrix.with_updated_data(data=b.data, coords=indexer)
 
         s_matrix = self.ab_to_s(a_matrix, b_matrix)
 
@@ -322,12 +319,9 @@ def compute_power_wave_amplitudes_at_each_port(
 
         for port_out in self.ports:
             V_out, I_out = self.compute_port_VI(port_out, sim_data)
-            indexer = {"f": V_out.f, "port": port_out.name}
-
-            V_out_data = np.expand_dims(V_out.data, axis=1)
-            I_out_data = np.expand_dims(I_out.data, axis=1)
-            V_matrix = V_matrix.with_updated_data(data=V_out_data, coords=indexer)
-            I_matrix = V_matrix.with_updated_data(data=I_out_data, coords=indexer)
+            indexer = {"port": port_out.name}
+            V_matrix = V_matrix.with_updated_data(data=V_out.data, coords=indexer)
+            I_matrix = V_matrix.with_updated_data(data=I_out.data, coords=indexer)
 
         V_numpy = V_matrix.values
         I_numpy = I_matrix.values