fix[adjoint]: fix magnitude of gradient calculation in CustomMedium w…

…hen permittivity and medium defined over different number of dimensions

CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Support for quasi-uniform grid specifications via `QuasiUniformGrid` that subclasses from `GridSpec1d`. The grids are almost uniform, but can adjust locally to the edge of structure bounding boxes, and snapping points.
 - New field `min_steps_per_sim_size` in `AutoGrid` that sets minimal number of grid steps per longest edge length of simulation domain.
 - Validation error in inverse design plugin when simulation has no sources by adding a source existence check before validating pixel size.
+- Fixed magnitude of gradient computation in `CustomMedium` by accounting properly for full volume element when permittivity data is defined over less dimensions than the medium.
 
 ## [2.8.0rc1] - 2024-12-17
 

tidy3d/components/medium.py

@@ -2648,7 +2648,17 @@ def _derivative_field_cmp(
 
         coords_interp = {key: eps_data.coords[key] for key in "xyz"}
         coords_interp = {key: val for key, val in coords_interp.items() if len(val) > 1}
-        dims_sum = [dim for dim in "xyz" if dim not in coords_interp]
+
+        E_der_dim_interp = E_der_map[f"E{dim}"]
+
+        for dim_ in "xyz":
+            if dim_ not in coords_interp:
+                bound_max = np.max(E_der_dim_interp.coords[dim_])
+                bound_min = np.min(E_der_dim_interp.coords[dim_])
+                dimension_size = bound_max - bound_min
+
+                if dimension_size > 0.0:
+                    E_der_dim_interp = E_der_dim_interp.integrate(dim_)
 
         # compute sizes along each of the interpolation dimensions
         sizes_list = []
@@ -2676,27 +2686,24 @@ def _derivative_field_cmp(
             d_vol = np.array(1.0)
 
         # TODO: probably this could be more robust. eg if the DataArray has weird edge cases
-        E_der_dim = E_der_map[f"E{dim}"]
         E_der_dim_interp = (
-            E_der_dim.interp(**coords_interp, assume_sorted=True).fillna(0.0).sum(dims_sum).real
+            E_der_dim_interp.interp(**coords_interp, assume_sorted=True).fillna(0.0).real.sum("f")
         )
-        E_der_dim_interp = E_der_dim_interp.sum("f")
-
-        vjp_array = np.array(E_der_dim_interp.values, dtype=float)
-        vjp_array = vjp_array.reshape(eps_data.shape)
 
-        # multiply by volume elements (if possible, being defensive here..)
         try:
-            vjp_array *= d_vol.reshape(vjp_array.shape)
+            E_der_dim_interp = E_der_dim_interp * d_vol.reshape(E_der_dim_interp.shape)
         except ValueError:
             log.warning(
                 "Skipping volume element normalization of 'CustomMedium' gradients. "
                 f"Could not reshape the volume elements of shape {d_vol.shape} "
-                f"to the shape of the gradient {vjp_array.shape}. "
+                f"to the shape of the fields {E_der_dim_interp.shape}. "
                 "If you encounter this warning, gradient direction will be accurate but the norm "
                 "will be inaccurate. Please raise an issue on the tidy3d front end with this "
                 "message and some information about your simulation setup and we will investigate. "
             )
+        vjp_array = E_der_dim_interp.values
+        vjp_array = vjp_array.reshape(eps_data.shape)
+
         return vjp_array