feature[adjoint]: pec gradient support for Box

Author: Gregory Roberts

tests/test_components/test_autograd_rf_box.py

@@ -22,6 +22,7 @@
 from tidy3d.components.autograd.constants import GRADIENT_DTYPE_FLOAT
 from tidy3d.components.autograd.derivative_utils import DerivativeInfo, integrate_within_bounds
 from tidy3d.components.base import Tidy3dBaseModel, cached_property
+from tidy3d.components.geometry.bound_ops import bounds_intersection, bounds_union
 from tidy3d.components.transformation import ReflectionFromPlane, RotationAroundAxis
 from tidy3d.components.types import (
     ArrayFloat2D,
@@ -51,7 +52,7 @@
     polygon_patch,
     set_default_labels_and_title,
 )
-from tidy3d.constants import LARGE_NUMBER, MICROMETER, RADIAN, fp_eps, inf
+from tidy3d.constants import EPSILON_0, LARGE_NUMBER, MICROMETER, MU_0, RADIAN, fp_eps, inf
 from tidy3d.exceptions import (
     SetupError,
     Tidy3dError,
@@ -62,8 +63,6 @@
 from tidy3d.log import log
 from tidy3d.packaging import verify_packages_import
 
-from .bound_ops import bounds_intersection, bounds_union
-
 POLY_GRID_SIZE = 1e-12
 POLY_TOLERANCE_RATIO = 1e-12
 
@@ -2439,11 +2438,12 @@ def _derivative_face(
 
         # normal and tangential dims
         dim_normal, dims_perp = self.pop_axis("xyz", axis=axis_normal)
-        fld_normal, flds_perp = self.pop_axis(("Ex", "Ey", "Ez"), axis=axis_normal)
+        fld_E_normal, flds_E_perp = self.pop_axis(("Ex", "Ey", "Ez"), axis=axis_normal)
+        fld_H_normal, flds_H_perp = self.pop_axis(("Hx", "Hy", "Hz"), axis=axis_normal)
 
         # fields and bounds
-        D_normal = derivative_info.D_der_map[fld_normal]
-        Es_perp = tuple(derivative_info.E_der_map[key] for key in flds_perp)
+        D_normal = derivative_info.D_der_map[fld_E_normal]
+        Es_perp = tuple(derivative_info.E_der_map[key] for key in flds_E_perp)
         bounds_normal, bounds_perp = self.pop_axis(
             np.array(derivative_info.bounds).T, axis=axis_normal
         )
@@ -2497,33 +2497,159 @@ def _derivative_face(
         delta_eps_perps = [eps_in - eps_out for eps_in, eps_out in zip(eps_in_perps, eps_out_perps)]
         delta_eps_inv_normal = 1.0 / eps_in_normal - 1.0 / eps_out_normal
 
-        def integrate_face(arr: xr.DataArray) -> complex:
+        def integrate_face(
+            arr: xr.DataArray,
+            snap_coords_outside: bool = False,
+            singularity_correction=None,
+            integration_dims=dims_perp,
+            integration_bounds=bounds_perp,
+        ) -> complex:
             """Interpolate and integrate a scalar field data over the face using bounds."""
 
-            arr_at_face = arr.interp(**{dim_normal: float(coord_normal_face)}, assume_sorted=True)
+            edge_correction = 1.0
+
+            interpolation_point = coord_normal_face
+            if snap_coords_outside:
+                snap_coords_values = arr.coords[dim_normal].values
+
+                index_face = np.argmin(np.abs(snap_coords_values - coord_normal_face))
+
+                if min_max_index == 0:
+                    min_boundary_mapping = np.where(snap_coords_values > coord_normal_face)[0]
+                    index_face = (
+                        0
+                        if (len(min_boundary_mapping) == 0)
+                        else np.maximum(0, min_boundary_mapping[0] - 1)
+                    )
+
+                    if snap_coords_values[index_face] > coord_normal_face:
+                        log.warning("Was not able to snap coordinates outside of min face.")
+                else:
+                    max_boundary_mapping = np.where(snap_coords_values < coord_normal_face)[0]
+                    index_face = (
+                        0
+                        if (len(max_boundary_mapping) == 0)
+                        else np.minimum(len(snap_coords_values) - 1, max_boundary_mapping[-1] + 1)
+                    )
+
+                    if snap_coords_values[index_face] < coord_normal_face:
+                        log.warning("Was not able to snap coordinates outside of max face.")
+
+                interpolation_point = snap_coords_values[index_face]
+
+                if singularity_correction:
+                    edge_correction = singularity_correction(
+                        np.abs(interpolation_point - coord_normal_face)
+                    )
+
+            arr_at_face = (
+                arr
+                if (len(arr.coords[dim_normal]) == 1)
+                else arr.interp(**{dim_normal: float(interpolation_point)}, assume_sorted=True)
+            )
 
             integral_result = integrate_within_bounds(
                 arr=arr_at_face,
-                dims=dims_perp,
-                bounds=bounds_perp,
+                dims=integration_dims,
+                bounds=integration_bounds,
             )
 
-            return complex(integral_result.sum("f"))
+            return complex(edge_correction * integral_result.sum("f"))
 
         # compute vjp from field integrals
         vjp_value = 0.0
 
-        # perform D-normal integral
-        integrand_D = -delta_eps_inv_normal * D_normal
-        integral_D = integrate_face(integrand_D)
-        vjp_value += integral_D
+        if derivative_info.is_medium_pec:
+            Hs_perp = tuple(derivative_info.H_der_map[key] for key in flds_H_perp)
+
+            # detect whether the box is 2-dimensional
+            zero_size_map = [s == 0.0 for s in self.size]
+
+            dimension = 3 - np.sum(zero_size_map)
+            if dimension < 2:
+                log.error(
+                    "Derivative of PEC material with less than 2 dimensions is unsupported. "
+                    "Specified PEC box is {dimension}-dimesional"
+                )
+
+            do_singularity_correction = False
+            is_2d = np.any(zero_size_map)
+            if is_2d:
+                zero_dim_idx = np.where(zero_size_map)[0][0]
+                zero_dimension = "xyz"[zero_dim_idx]
+
+                # for singularity correction, need 2-dimensional box and that
+                # the face we are integrating over is the 1-dimensional (i.e. -
+                # the normal for the face is not the same as the flat dimension
+                do_singularity_correction = not (axis_normal == zero_dim_idx)
+
+            # apply singularity correction only when integrating along the line
+            singularity_correction = (
+                (lambda d: 0.5 * np.pi * d) if do_singularity_correction else (lambda d: 1.0)
+            )
+
+            integration_dims = dims_perp
+            integration_bounds = bounds_perp
+            # if we are correcting for singularity and integrating over the line, then adjust
+            # integration dimensions and bounds to exclude the flat dimension
+            if do_singularity_correction:
+                integration_dims = [dim for dim in dims_perp if (not (dim == zero_dimension))]
+
+                integration_bounds = []
+                zero_dim_idx = 0
+                # find the index into the bounds that corresponds to the flat dimension
+                for dim in dims_perp:
+                    if dim == zero_dimension:
+                        break
+                    zero_dim_idx += 1
+
+                # trim the zero dimension from the integration bounds
+                for bound in bounds_perp:
+                    new_bound = [b for b_idx, b in enumerate(bound) if b_idx != zero_dim_idx]
+
+                    integration_bounds.append(new_bound)
+
+            integrand_E = D_normal / np.real(eps_out_normal)
+
+            integral_E = integrate_face(
+                integrand_E,
+                snap_coords_outside=True,
+                singularity_correction=singularity_correction,
+                integration_dims=integration_dims,
+                integration_bounds=integration_bounds,
+            )
 
-        # perform E-perpendicular integrals
-        for E_perp, delta_eps_perp in zip(Es_perp, delta_eps_perps):
-            integrand_E = E_perp * delta_eps_perp
-            integral_E = integrate_face(integrand_E)
             vjp_value += integral_E
 
+            for H_perp_idx, H_perp in enumerate(Hs_perp):
+                integrand_H = MU_0 * H_perp / EPSILON_0
+
+                if (is_2d and not (flds_H_perp[H_perp_idx] == f"H{zero_dimension}")) and (
+                    dim_normal != zero_dimension
+                ):
+                    continue
+
+                integral_H = integrate_face(
+                    integrand_H,
+                    snap_coords_outside=True,
+                    singularity_correction=singularity_correction,
+                    integration_dims=integration_dims,
+                    integration_bounds=integration_bounds,
+                )
+
+                vjp_value += integral_H
+
+        else:
+            integrand_D = -delta_eps_inv_normal * D_normal
+            integral_D = integrate_face(integrand_D)
+            vjp_value += integral_D
+
+            # perform E-perpendicular integrals
+            for E_perp, delta_eps_perp in zip(Es_perp, delta_eps_perps):
+                integrand_E = E_perp * delta_eps_perp
+                integral_E = integrate_face(integrand_E)
+                vjp_value += integral_E
+
         return np.real(vjp_value)
 
 

tests/test_components/test_autograd_rf_polyslab.py

@@ -1447,13 +1447,15 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField
         """
         vjps: AutogradFieldMap = {}
 
-        sim_min, sim_max = map(np.asarray, derivative_info.bounds_intersect)
-        extents = sim_max - sim_min
+        intersect_min, intersect_max = map(np.asarray, derivative_info.bounds_intersect)
+        sim_min, sim_max = map(np.asarray, derivative_info.simulation_bounds)
+
+        extents = intersect_max - intersect_min
         is_2d = np.isclose(extents[self.axis], 0.0)
 
         # early return if polyslab is not in simulation domain
         slab_min, slab_max = self.slab_bounds
-        if (slab_max <= sim_min[self.axis]) or (slab_min >= sim_max[self.axis]):
+        if (slab_max < sim_min[self.axis]) or (slab_min > sim_max[self.axis]):
             log.warning(
                 "'PolySlab' lies completely outside the simulation domain.",
                 log_once=True,
@@ -1473,6 +1475,7 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField
                 vjps[path] = self._compute_derivative_vertices(
                     derivative_info, sim_min, sim_max, is_2d, interpolators
                 )
+
             elif path[0] == "slab_bounds":
                 idx = path[1]
                 face_coord = self.slab_bounds[idx]
@@ -1508,6 +1511,12 @@ def _compute_derivative_slab_bounds(
         is split equally between the two vertices that bound the edge segment.
         """
         # rmin/rmax over the geometry and simulation box
+        if np.isclose(self.slab_bounds[1] - self.slab_bounds[0], 0.0):
+            log.warning(
+                "Computing slab face derivatives for flat structures is not fully supported and "
+                "may give zero for the derivative. Try using a structure with a small, but nonzero"
+                "thickness for slab bound derivatives."
+            )
         rmin, rmax = derivative_info.bounds_intersect
         _, (r1_min, r2_min) = self.pop_axis(rmin, axis=self.axis)
         _, (r1_max, r2_max) = self.pop_axis(rmax, axis=self.axis)

tidy3d/components/autograd/derivative_utils.py

@@ -321,11 +321,15 @@ def _make_adjoint_monitors(
         else:
             interval_space = AUTOGRAD_MONITOR_INTERVAL_SPACE_POLY
 
+        field_components_for_adjoint = [f"E{dim}" for dim in "xyz"]
+        if self.medium.is_pec:
+            field_components_for_adjoint += [f"H{dim}" for dim in "xyz"]
+
         mnt_fld = FieldMonitor(
             size=size,
             center=center,
             freqs=freqs,
-            fields=("Ex", "Ey", "Ez"),
+            fields=field_components_for_adjoint,
             name=self._get_monitor_name(index=index, data_type="fld"),
             interval_space=interval_space,
             colocate=False,