reorganized code

Author: rahul-flex

tidy3d/components/geometry/base.py

@@ -1640,6 +1640,125 @@ def __invert__(self):
             operation="difference", geometry_a=Box(size=(inf, inf, inf)), geometry_b=self
         )
 
+    def build_box_face_mesh(
+        self,
+        center: np.ndarray,
+        size: np.ndarray,
+        axis_normal: int,  # 0,1,2  → x,y,z faces
+        min_max_index: int,  # 0 = − side, 1 = + side
+        rotation_matrix: np.ndarray | None = None,
+    ) -> tuple[DerivativeSurfaceMesh, np.ndarray, np.ndarray]:
+        """Build a mesh for the face of a box, given the center, size, axis normal, and min/max index.
+        The mesh is built in the local coordinate system of the box, and then transformed to the global
+        coordinate system using the rotation matrix."""
+
+        if axis_normal == 0:
+            canonical_normal = np.array([1.0, 0.0, 0.0])
+        elif axis_normal == 1:
+            canonical_normal = np.array([0.0, 1.0, 0.0])
+        elif axis_normal == 2:
+            canonical_normal = np.array([0.0, 0.0, 1.0])
+        else:
+            raise ValueError("Invalid axis_normal")
+
+        if min_max_index == 0:
+            canonical_normal *= -1.0
+
+        if rotation_matrix is None:
+            rotation_matrix = np.eye(3)
+
+        n_local = rotation_matrix @ canonical_normal
+        n_local = n_local / np.linalg.norm(n_local)
+
+        def compute_tangential_vectors(
+            normal: np.ndarray, eps: float = 1e-8
+        ) -> tuple[np.ndarray, np.ndarray]:
+            """Compute any two perpendicular tangential vectors t1, t2, given a normal."""
+            if abs(normal[0]) > abs(normal[2]):
+                t1 = np.array([-normal[1], normal[0], 0.0])
+            else:
+                t1 = np.array([0.0, -normal[2], normal[1]])
+            t1_norm = np.linalg.norm(t1)
+            if t1_norm < eps:
+                raise ValueError("Degenerate normal vector.")
+            t1 = t1 / t1_norm
+            t2 = np.cross(normal, t1)
+            t2 /= np.linalg.norm(t2)
+            return t1, t2
+
+        t1_local, t2_local = compute_tangential_vectors(n_local)
+
+        min_bound = np.array(center) - np.array(size) / 2.0
+        max_bound = np.array(center) + np.array(size) / 2.0
+        bounds_old = np.column_stack((min_bound, max_bound))
+
+        corners = np.array(
+            [
+                [bounds_old[0, i], bounds_old[1, j], bounds_old[2, k]]
+                for i in (0, 1)
+                for j in (0, 1)
+                for k in (0, 1)
+            ]
+        )
+
+        connectivity = {
+            0: {  # Faces perpendicular to x-axis
+                0: [0, 1, 3, 2],
+                1: [4, 5, 7, 6],
+            },
+            1: {  # Faces perpendicular to y-axis
+                0: [0, 1, 5, 4],
+                1: [2, 3, 7, 6],
+            },
+            2: {  # Faces perpendicular to z-axis
+                0: [0, 4, 6, 2],
+                1: [1, 5, 7, 3],
+            },
+        }
+
+        face_indices = connectivity[axis_normal][min_max_index]
+        face_corners = corners[face_indices, :]
+
+        rotated_corners = (rotation_matrix @ face_corners.T).T
+        p1, p2, p3, p4 = rotated_corners
+
+        num_s = _NUM_PTS_DIM_BOX_FACE
+        num_t = _NUM_PTS_DIM_BOX_FACE
+        s_vals = np.linspace(0, 1, 2 * num_s + 1)[1::2]
+        t_vals = np.linspace(0, 1, 2 * num_t + 1)[1::2]
+        S, T = np.meshgrid(s_vals, t_vals, indexing="ij")
+
+        X = (1 - S) * (1 - T) * p1[0] + S * (1 - T) * p2[0] + S * T * p3[0] + (1 - S) * T * p4[0]
+        Y = (1 - S) * (1 - T) * p1[1] + S * (1 - T) * p2[1] + S * T * p3[1] + (1 - S) * T * p4[1]
+        Z = (1 - S) * (1 - T) * p1[2] + S * (1 - T) * p2[2] + S * T * p3[2] + (1 - S) * T * p4[2]
+
+        centers = np.column_stack([X.ravel(), Y.ravel(), Z.ravel()])
+
+        tri1_area = 0.5 * np.linalg.norm(np.cross((p2 - p1), (p3 - p1)))
+        tri2_area = 0.5 * np.linalg.norm(np.cross((p4 - p1), (p3 - p1)))
+        face_area = tri1_area + tri2_area
+
+        num_cells = (num_s) * (num_t)
+        if num_cells > 0:
+            cell_area = face_area / num_cells
+        else:
+            cell_area = face_area
+
+        areas = cell_area * np.ones(centers.shape[0])
+
+        normals = np.tile(n_local, (centers.shape[0], 1))
+        perps1 = np.tile(t1_local, (centers.shape[0], 1))
+        perps2 = np.tile(t2_local, (centers.shape[0], 1))
+
+        surface_mesh = DerivativeSurfaceMesh(
+            centers=centers,
+            areas=areas,
+            normals=normals,
+            perps1=perps1,
+            perps2=perps2,
+        )
+        return surface_mesh, n_local
+
 
 """ Abstract subclasses """
 
@@ -2564,119 +2683,17 @@ def derivative_face(
     ) -> float:
         """
         Compute the derivative (VJP) with respect to shifting a face of a rotated box,
-        using full integration over that face. This version uses bilinear interpolation
-        of the four corners to sample interior points.
-        """
-        if axis_normal == 0:
-            canonical_normal = np.array([1.0, 0.0, 0.0])
-        elif axis_normal == 1:
-            canonical_normal = np.array([0.0, 1.0, 0.0])
-        elif axis_normal == 2:
-            canonical_normal = np.array([0.0, 0.0, 1.0])
-        else:
-            raise ValueError("Invalid axis_normal")
-
-        if min_max_index == 0:
-            canonical_normal *= -1.0
-
-        if rotation_matrix is None:
-            rotation_matrix = np.eye(3)
-
-        n_local = rotation_matrix @ canonical_normal
-        n_local = n_local / np.linalg.norm(n_local)
-
-        def compute_tangential_vectors(
-            normal: np.ndarray, eps: float = 1e-8
-        ) -> tuple[np.ndarray, np.ndarray]:
-            """Compute any two perpendicular tangential vectors t1, t2, given a normal."""
-            if abs(normal[0]) > abs(normal[2]):
-                t1 = np.array([-normal[1], normal[0], 0.0])
-            else:
-                t1 = np.array([0.0, -normal[2], normal[1]])
-            t1_norm = np.linalg.norm(t1)
-            if t1_norm < eps:
-                raise ValueError("Degenerate normal vector.")
-            t1 = t1 / t1_norm
-            t2 = np.cross(normal, t1)
-            t2 /= np.linalg.norm(t2)
-            return t1, t2
-
-        t1_local, t2_local = compute_tangential_vectors(n_local)
-
-        min_bound = np.array(self.center) - np.array(self.size) / 2.0
-        max_bound = np.array(self.center) + np.array(self.size) / 2.0
-        bounds_old = np.column_stack((min_bound, max_bound))
-
-        corners = np.array(
-            [
-                [bounds_old[0, i], bounds_old[1, j], bounds_old[2, k]]
-                for i in (0, 1)
-                for j in (0, 1)
-                for k in (0, 1)
-            ]
+        using full integration over that face.
+        """
+        mesh, _ = self.build_box_face_mesh(
+            center=np.asarray(self.center, float),
+            size=np.asarray(self.size, float),
+            axis_normal=axis_normal,
+            min_max_index=min_max_index,
+            rotation_matrix=rotation_matrix,
         )
-
-        connectivity = {
-            0: {  # Faces perpendicular to x-axis
-                0: [0, 1, 3, 2],
-                1: [4, 5, 7, 6],
-            },
-            1: {  # Faces perpendicular to y-axis
-                0: [0, 1, 5, 4],
-                1: [2, 3, 7, 6],
-            },
-            2: {  # Faces perpendicular to z-axis
-                0: [0, 4, 6, 2],
-                1: [1, 5, 7, 3],
-            },
-        }
-
-        face_indices = connectivity[axis_normal][min_max_index]
-        face_corners = corners[face_indices, :]
-
-        rotated_corners = (rotation_matrix @ face_corners.T).T
-        p1, p2, p3, p4 = rotated_corners
-
-        num_s = _NUM_PTS_DIM_BOX_FACE
-        num_t = _NUM_PTS_DIM_BOX_FACE
-        s_vals = np.linspace(0, 1, 2 * num_s + 1)[1::2]
-        t_vals = np.linspace(0, 1, 2 * num_t + 1)[1::2]
-        S, T = np.meshgrid(s_vals, t_vals, indexing="ij")
-
-        X = (1 - S) * (1 - T) * p1[0] + S * (1 - T) * p2[0] + S * T * p3[0] + (1 - S) * T * p4[0]
-        Y = (1 - S) * (1 - T) * p1[1] + S * (1 - T) * p2[1] + S * T * p3[1] + (1 - S) * T * p4[1]
-        Z = (1 - S) * (1 - T) * p1[2] + S * (1 - T) * p2[2] + S * T * p3[2] + (1 - S) * T * p4[2]
-
-        centers = np.column_stack([X.ravel(), Y.ravel(), Z.ravel()])
-
-        tri1_area = 0.5 * np.linalg.norm(np.cross((p2 - p1), (p3 - p1)))
-        tri2_area = 0.5 * np.linalg.norm(np.cross((p4 - p1), (p3 - p1)))
-        face_area = tri1_area + tri2_area
-
-        num_cells = (num_s) * (num_t)
-        if num_cells > 0:
-            cell_area = face_area / num_cells
-        else:
-            cell_area = face_area
-
-        areas = cell_area * np.ones(centers.shape[0])
-
-        normals = np.tile(n_local, (centers.shape[0], 1))
-        perps1 = np.tile(t1_local, (centers.shape[0], 1))
-        perps2 = np.tile(t2_local, (centers.shape[0], 1))
-
-        surface_mesh = DerivativeSurfaceMesh(
-            centers=centers,
-            areas=areas,
-            normals=normals,
-            perps1=perps1,
-            perps2=perps2,
-        )
-
-        grads = derivative_info.grad_surfaces(surface_mesh=surface_mesh)
-        vjp_value = np.real(np.sum(grads).item())
-
-        return vjp_value
+        vjp = derivative_info.grad_surfaces(surface_mesh=mesh)
+        return float(np.real(np.sum(vjp)))
 
 
 """Compound subclasses"""

tidy3d/web/api/autograd/autograd.py

@@ -1013,77 +1013,56 @@ def postprocess_adj(
                 eps_background = None
 
             # manually override simulation medium as the background structure
-            if not isinstance(structure.geometry, td.Box):
-                # auto permittivity detection
-                sim_orig = sim_data_orig.simulation
-                plane_eps = eps_fwd.monitor.geometry
-
-                # get permittivity without this structure
-                structs_no_struct = list(sim_orig.structures)
-                structs_no_struct.pop(structure_index)
-                sim_no_structure = sim_orig.updated_copy(structures=structs_no_struct)
-                eps_no_structure = sim_no_structure.epsilon(
-                    box=plane_eps, coord_key="centers", freq=freq_adj
-                )
-
-        eps_in = np.mean(structure.medium.eps_model(freq_adj))
-        eps_out = np.mean(sim_data_orig.simulation.medium.eps_model(freq_adj))
-        if structure.background_medium:
-            eps_background = structure.background_medium.eps_model(freq_adj)
-        else:
-            eps_background = None
-
-        # manually override simulation medium as the background structure
-        # auto permittivity detection
-        sim_orig = sim_data_orig.simulation
-        plane_eps = eps_fwd.monitor.geometry
-
-        # get permittivity without this structure
-        structs_no_struct = list(sim_orig.structures)
-        structs_no_struct.pop(structure_index)
-        sim_no_structure = sim_orig.updated_copy(structures=structs_no_struct)
-        eps_no_structure = sim_no_structure.epsilon(
-            box=plane_eps, coord_key="centers", freq=freq_adj
-        )
+            # auto permittivity detection
+            sim_orig = sim_data_orig.simulation
+            plane_eps = eps_fwd.monitor.geometry
+
+            # get permittivity without this structure
+            structs_no_struct = list(sim_orig.structures)
+            structs_no_struct.pop(structure_index)
+            sim_no_structure = sim_orig.updated_copy(structures=structs_no_struct)
+            eps_no_structure = sim_no_structure.epsilon(
+                box=plane_eps, coord_key="centers", freq=freq_adj
+            )
 
-        # get permittivity with structures on top of an infinite version of this structure
-        structs_inf_struct = list(sim_orig.structures)[structure_index + 1 :]
-        sim_inf_structure = sim_orig.updated_copy(
-            structures=structs_inf_struct,
-            medium=structure.medium,
-            monitors=[],
-        )
-        eps_inf_structure = sim_inf_structure.epsilon(
-            box=plane_eps, coord_key="centers", freq=freq_adj
-        )
+            # get permittivity with structures on top of an infinite version of this structure
+            structs_inf_struct = list(sim_orig.structures)[structure_index + 1 :]
+            sim_inf_structure = sim_orig.updated_copy(
+                structures=structs_inf_struct,
+                medium=structure.medium,
+                monitors=[],
+            )
+            eps_inf_structure = sim_inf_structure.epsilon(
+                box=plane_eps, coord_key="centers", freq=freq_adj
+            )
 
-        # get minimum intersection of bounds with structure and sim
-        struct_bounds = rmin_struct, rmax_struct = structure.geometry.bounds
-        rmin_sim, rmax_sim = sim_data_orig.simulation.bounds
-        rmin_intersect = tuple([max(a, b) for a, b in zip(rmin_sim, rmin_struct)])
-        rmax_intersect = tuple([min(a, b) for a, b in zip(rmax_sim, rmax_struct)])
-        bounds_intersect = (rmin_intersect, rmax_intersect)
-
-        derivative_info = DerivativeInfo(
-            paths=structure_paths,
-            E_der_map=E_der_map.field_components,
-            D_der_map=D_der_map.field_components,
-            E_fwd=E_fwd.field_components,
-            E_adj=E_adj.field_components,
-            D_fwd=D_fwd.field_components,
-            D_adj=D_adj.field_components,
-            eps_data=eps_fwd.field_components,
-            eps_in=eps_in,
-            eps_out=eps_out,
-            eps_background=eps_background,
-            frequency=freq_adj,
-            eps_no_structure=eps_no_structure,
-            eps_inf_structure=eps_inf_structure,
-            bounds=struct_bounds,
-            bounds_intersect=bounds_intersect,
-        )
+            # get minimum intersection of bounds with structure and sim
+            struct_bounds = rmin_struct, rmax_struct = structure.geometry.bounds
+            rmin_sim, rmax_sim = sim_data_orig.simulation.bounds
+            rmin_intersect = tuple([max(a, b) for a, b in zip(rmin_sim, rmin_struct)])
+            rmax_intersect = tuple([min(a, b) for a, b in zip(rmax_sim, rmax_struct)])
+            bounds_intersect = (rmin_intersect, rmax_intersect)
+
+            derivative_info = DerivativeInfo(
+                paths=structure_paths,
+                E_der_map=E_der_map.field_components,
+                D_der_map=D_der_map.field_components,
+                E_fwd=E_fwd.field_components,
+                E_adj=E_adj.field_components,
+                D_fwd=D_fwd.field_components,
+                D_adj=D_adj.field_components,
+                eps_data=eps_fwd.field_components,
+                eps_in=eps_in,
+                eps_out=eps_out,
+                eps_background=eps_background,
+                frequency=freq_adj,
+                eps_no_structure=eps_no_structure,
+                eps_inf_structure=eps_inf_structure,
+                bounds=struct_bounds,
+                bounds_intersect=bounds_intersect,
+            )
 
-        vjp_value_map = structure.compute_derivatives(derivative_info)
+            vjp_value_map = structure.compute_derivatives(derivative_info)
 
             # extract VJPs and put back into sim_fields_vjp AutogradFieldMap
             for structure_path, vjp_value in vjp_value_map.items():