add exception handling for reshaping with DataArray._with_updated_copy

Author: tylerflex

tests/test_data/test_data_arrays.py

@@ -472,7 +472,7 @@ def test_interp(method, scalar_index):
     xrt.assert_allclose(xr_interp, ag_interp)
 
 
-def test_with_updated_data():
+def test_with_updated_data_grad():
     """Check the ``DataArray.with_updated_data()`` method."""
 
     arr = td.SpatialDataArray(
@@ -499,3 +499,20 @@ def f(x):
     assert np.all(g == np.ones_like(data))
 
     check_grads(f, order=1, modes=["rev"])(data)
+
+
+def test_with_updated_data_shape():
+    """Check the ``DataArray.with_updated_data()`` method."""
+
+    arr = td.SpatialDataArray(
+        np.ones((2, 3, 4, 5), dtype=np.complex128),
+        coords={"x": [0, 1], "y": [1, 2, 3], "z": [2, 3, 4, 5], "w": [0, 1, 2, 3, 4]},
+    )
+
+    # wrong shape
+    data = np.zeros((1, 1, 1, 3))
+
+    coords = {"x": 0, "y": 2, "z": 3}
+
+    with pytest.raises(ValueError):
+        arr2 = arr._with_updated_data(data=data, coords=coords)

tidy3d/components/data/data_array.py

@@ -507,7 +507,14 @@ def _with_updated_data(self, data: np.ndarray, coords: dict[str, Any]) -> DataAr
         for i, dim in enumerate(self.dims):
             if dim in coords:
                 new_shape[i] = 1
-        new_data = data.reshape(new_shape)
+        try:
+            new_data = data.reshape(new_shape)
+        except ValueError as e:
+            raise ValueError(
+                "Couldn't reshape the supplied 'data' to update 'DataArray'. The provided data was "
+                f"of shape {data.shape} and tried to reshape to {new_shape}. If you encounter this "
+                "error please raise an issue on the tidy3d github repository with the context."
+            ) from e
 
         # broadcast data to repeat data along the selected dimensions to match mask
         new_data = new_data + np.zeros_like(old_data)