fix: autograd regressions

[yaugenst-flex]

===================== 8 passed in 1581.62s (0:26:21) =================== 🎉

fixes regressions introduced in 8bdf76e for changes that weren't carried over from c038f17

Greptile Overview

Greptile Summary

Fixes autograd gradient computation regressions by refactoring the derivative calculation approach for Box and PolySlab geometries.

Key changes:

• Box gradient refactoring: Replaced the xarray-based field interpolation and integration approach (~340 lines) with a simpler point-based grid sampling method (~150 lines) that evaluates gradients at discrete points on geometry faces
• PolySlab 2D fix: Fixed gradient calculation for 2D geometries by clamping the center point to simulation bounds (prevents using out-of-bounds coordinates)
• PolySlab intersection fix: Added explicit polygon intersection with simulation bounds box to ensure proper gradient computation when geometry extends beyond simulation domain
• Removed unused imports: Cleaned up xarray, EPSILON_0, MU_0, and field data utility imports that are no longer needed

The refactoring simplifies the gradient computation logic significantly while fixing numerical issues that were causing test failures.

Confidence Score: 4/5

• This PR is safe to merge with minimal risk after addressing the debug comment
• Score reflects a significant refactoring that simplifies complex autograd logic (340 lines reduced to 150), passes all 8 autograd tests, and fixes known gradient regressions. The only issue found is a leftover debug comment. The changes are well-tested and the simpler implementation reduces future maintenance burden.
• Minor style cleanup needed in tidy3d/components/geometry/base.py (debug comment on line 2616)

Important Files Changed

File Analysis

Filename	Score	Overview
tidy3d/components/geometry/base.py	4/5	Refactored Box._derivative_face method from xarray-based integration to point-based grid sampling approach, removing 300+ lines of complex field interpolation logic. Contains minor leftover debug comment.
tidy3d/components/geometry/polyslab.py	5/5	Fixed 2D geometry gradient calculation by clamping center point to simulation bounds and added polygon intersection with simulation box to ensure proper gradient computation.

Sequence Diagram

sequenceDiagram
    participant Caller as Autograd System
    participant BoxFaces as Box._derivative_faces
    participant BoxFace as Box._derivative_face
    participant GridBuilder as Grid Builder
    participant Evaluator as evaluate_gradient_at_points
    
    Caller->>BoxFaces: compute gradients for box
    BoxFaces->>BoxFaces: determine axes_to_compute from paths
    
    loop For each face (6 faces)
        BoxFaces->>BoxFace: compute gradient for face
        BoxFace->>BoxFace: check if face outside simulation
        alt Face outside bounds
            BoxFace-->>BoxFaces: return 0.0
        else Face inside bounds
            BoxFace->>BoxFace: compute intersection bounds
            BoxFace->>BoxFace: detect if 2D geometry
            alt is 2D
                BoxFace->>GridBuilder: build 1D grid along edge
                GridBuilder-->>BoxFace: grid_points (1D)
            else is 3D
                BoxFace->>GridBuilder: build 2D grid on face
                GridBuilder-->>BoxFace: grid_points (2D mesh)
            end
            BoxFace->>Evaluator: evaluate_gradient_at_points
            Evaluator-->>BoxFace: gradient values at points
            BoxFace->>BoxFace: integrate (sum * weight)
            BoxFace-->>BoxFaces: vjp_value
        end
    end
    
    BoxFaces-->>Caller: vjp_faces array

Loading

[greptile-apps]

_{2 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[github-actions]

Diff Coverage

Diff: origin/develop...HEAD, staged and unstaged changes

• tidy3d/components/geometry/base.py (94.4%): Missing lines 2555,2584,2604,2627
• tidy3d/components/geometry/polyslab.py (100%)

Summary

• Total: 75 lines
• Missing: 4 lines
• Coverage: 94%

tidy3d/components/geometry/base.py

  2551                 continue
  2552             is_2d_map.append(np.isclose(extents[axis_idx], 0.0))
  2553 
  2554         if np.all(is_2d_map):
! 2555             return 0.0
  2556 
  2557         is_2d = np.any(is_2d_map)
  2558 
  2559         sim_bounds_normal, sim_bounds_perp = self.pop_axis(

---

  2580 
  2581         def compute_integration_weight(grid_points: NDArray[float]) -> float:
  2582             grid_spacing = grid_points[1] - grid_points[0]
  2583             if grid_spacing == 0.0:
! 2584                 integration_weight = 1.0 / len(grid_points)
  2585             else:
  2586                 integration_weight = grid_points[1] - grid_points[0]
  2587 
  2588             return integration_weight

---

  2600                 intersect_max_perp[nonzero_dim],
  2601             )
  2602 
  2603             if not verify_integration_interval(integration_bounds_perp):
! 2604                 return 0.0
  2605 
  2606             grid_points_linear = spacing_to_grid_points(
  2607                 adaptive_spacing, integration_bounds_perp[0], integration_bounds_perp[1]
  2608             )

---

  2623                 (intersect_min_perp[1], intersect_max_perp[1]),
  2624             )
  2625 
  2626             if not np.all([verify_integration_interval(b) for b in integration_bounds_perp]):
! 2627                 return 0.0
  2628 
  2629             grid_points_perp_1 = spacing_to_grid_points(
  2630                 adaptive_spacing, integration_bounds_perp[0][0], integration_bounds_perp[0][1]
  2631             )

---

[groberts-flex]

looks good to me, thanks for tracking down these changes!