grab inside pixel at index=3 or 1 depending on cells

Author: tylerflex

CHANGELOG.md

@@ -12,7 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 - Ensure same `Grid` is generated in forward and adjoint simulations by setting `GridSpec.wavelength` manually in adjoint.
-
+- Proper handling of `JaxBox` derivatives both for multi-cell and single cell thickness.
 
 ## [2.4.1] - 2023-9-20
 

tidy3d/plugins/adjoint/components/geometry.py

@@ -237,9 +237,20 @@ def store_vjp(
                     eps_data = grad_data_eps.field_components[eps_field_name].isel(f=0)
 
                     # get the permittivity values just inside and outside the edge
-                    n_cells_in = 2
+
+                    num_cells_normal_dim = len(eps_data.coords[dim_normal])
+
+                    # if the eps_data is <=3 grid cells thick, pick the middle cell
+                    if num_cells_normal_dim <= 3:
+                        isel_ins = num_cells_normal_dim // 2
+
+                    # otherwise, pick the cell 4 pixels from outside cell
+                    else:
+                        n_cells_in = 3
+                        isel_ins = n_cells_in if min_max_index == 0 else -n_cells_in - 1
+
                     isel_out = 0 if min_max_index == 0 else -1
-                    isel_ins = n_cells_in if min_max_index == 0 else -n_cells_in - 1
+
                     eps2 = eps_data.isel(**{dim_normal: isel_out})
                     eps1 = eps_data.isel(**{dim_normal: isel_ins})