fix adjoint plugin with complex-valued permittivity inputs

Author: tylerflex

CHANGELOG.md

@@ -7,7 +7,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Support for differentiating with respect to `JaxMedium.conductivity`.
-
 - Validating that every surface (unless excluded in ``exclude_surfaces``) of a 3D ``SurfaceIntegrationMonitor`` (flux monitor or field projection monitor) is not completely outside the simulation domain.
 
 ### Changed
@@ -26,6 +25,7 @@ that the fields match exactly except for a ``pi`` phase shift. This interpretati
 - Cleaner display of `ArrayLike` in docs.
 - `ArrayLike` validation properly fails with `None` or `nan` contents.
 - Apply finite grid correction to the fields when calculating the Poynting vector from 2D monitors.
+- `JaxCustomMedium` properly handles complex-valued permittivity.
 
 ## [2.3.0] - 2023-6-30
 

tests/test_plugins/test_adjoint.py

@@ -227,6 +227,11 @@ def make_sim(
 
     jax_box_custom = JaxBox(size=size, center=(1, 0, 2))
     values = base_eps_val + np.random.random((Nx, Ny, Nz, 1))
+
+    # adding this line breaks things without enforcing that the vjp for custom medium is complex
+    values = (1 + 1j) * values
+    values = values + (1 + 1j) * values / 0.5
+
     eps_ii = JaxDataArray(values=values, coords=coords)
     field_components = {f"eps_{dim}{dim}": eps_ii for dim in "xyz"}
     jax_eps_dataset = JaxPermittivityDataset(**field_components)

tidy3d/plugins/adjoint/components/medium.py

@@ -523,7 +523,14 @@ def store_vjp(
 
             # reshape values to the expected vjp shape to be more safe
             vjp_shape = tuple(len(coord) for _, coord in coords.items())
-            vjp_values = e_dotted.real.values.reshape(vjp_shape)
+
+            # make sure this has the same dtype as the original
+            dtype_orig = np.array(orig_data_array.values).dtype
+
+            vjp_values = e_dotted.values.reshape(vjp_shape)
+            if dtype_orig.kind == "f":
+                vjp_values = vjp_values.real
+            vjp_values = vjp_values.astype(dtype_orig)
 
             # construct a DataArray storing the vjp
             vjp_data_array = JaxDataArray(values=vjp_values, coords=coords)