Improve speed of Box derivative computation when in GeometryGroup

[groberts-flex]

Connect box face derivatives to evaluate_gradient_at_points in DerivativeInfo so that interpolators can be shared when using inside a GeometryGroup.

Greptile Overview

Updated On: 2025-10-06 19:50:07 UTC

Summary

This PR optimizes the performance of Box geometry derivative computation when used inside a GeometryGroup by refactoring the implementation to leverage the existing `evaluate_gradient_at_points` infrastructure in `DerivativeInfo`. The key changes include:

• Major refactoring of Box derivative computation: The _derivative_face method in tidy3d/components/geometry/base.py was completely rewritten, removing 380+ lines of complex field integration code and replacing it with a simpler grid-based approach that builds spatial grids and evaluates gradients at specific points.

• Performance optimization through interpolator sharing: By connecting to the evaluate_gradient_at_points method, the new implementation allows interpolators to be shared when multiple Box geometries are processed together in a GeometryGroup, significantly improving computational efficiency.

• Code simplification: The refactoring eliminates the need for PEC-specific handling, complex epsilon data management, and manual interpolation logic, making the codebase more maintainable while leveraging existing gradient evaluation frameworks.

• Comprehensive testing: A new test file was added to validate that the optimized Box gradient implementation produces mathematically equivalent results to the established PolySlab approach across different dimensional configurations (2D and 3D cases).

• Test cleanup: Obsolete Box-specific test functions were removed since the functionality is now covered by the broader gradient evaluation system.

This change fits into the broader Tidy3D architecture by unifying Box derivative computation with the existing DerivativeInfo framework, enabling better resource sharing and performance optimization when working with multiple geometries in electromagnetic simulations.

Changed Files

Filename	Score	Overview
CHANGELOG.md	5/5	Added changelog entry documenting the Box derivative performance improvement
tidy3d/components/geometry/base.py	4/5	Major refactoring of Box _derivative_face method to use grid-based approach with evaluate_gradient_at_points
tests/test_components/autograd/numerical/test_autograd_box_polyslab_numerical.py	5/5	New comprehensive test validating Box vs PolySlab gradient equivalence across 2D/3D configurations
tests/test_components/test_geometry.py	4/5	Removed obsolete Box-specific test functions that are no longer needed after refactoring

Confidence score: 4/5

• This PR appears safe to merge with good performance benefits and comprehensive testing coverage
• Score reflects the substantial refactoring of core gradient computation logic which requires careful validation
• Pay close attention to the geometry base module changes and ensure the new grid-based approach handles all edge cases correctly

Sequence Diagram

sequenceDiagram
    participant User
    participant Geometry
    participant DerivativeInfo
    participant Interpolators
    participant GradientEvaluator
    
    User->>Geometry: "_compute_derivatives(derivative_info)"
    Geometry->>DerivativeInfo: "create_interpolators(dtype=GRADIENT_DTYPE_FLOAT)"
    DerivativeInfo->>Interpolators: "create field data interpolators"
    Interpolators-->>DerivativeInfo: "interpolators"
    DerivativeInfo-->>Geometry: "interpolators"
    
    loop For each derivative path
        Geometry->>Geometry: "_derivative_faces(derivative_info)"
        
        loop For min/max faces and axes
            Geometry->>Geometry: "_derivative_face(min_max_index, axis_normal, derivative_info)"
            Geometry->>DerivativeInfo: "evaluate_gradient_at_points(spatial_coords, normals, perps1, perps2, interpolators)"
            DerivativeInfo->>GradientEvaluator: "interpolate fields at points"
            GradientEvaluator->>Interpolators: "get field values"
            Interpolators-->>GradientEvaluator: "field values"
            GradientEvaluator-->>DerivativeInfo: "gradient values"
            DerivativeInfo-->>Geometry: "gradient_at_points"
            Geometry->>Geometry: "sum(integration_weight * gradient_at_points)"
            Geometry-->>Geometry: "vjp_face"
        end
        
        Geometry->>Geometry: "_derivatives_center_size(vjps_faces)"
        Geometry-->>Geometry: "vjps_center_size"
    end
    
    Geometry-->>User: "derivative_map"

Loading

[greptile-apps]

_{4 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.2%): Missing lines 2498,2525,2545,2574

Summary

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

tidy3d/components/geometry/base.py

Lines 2494-2502

  2494 
  2495             is_2d_map.append(np.isclose(extents[axis_idx], 0.0))
  2496 
  2497         if np.all(is_2d_map):
! 2498             return 0.0
  2499 
  2500         is_2d = np.any(is_2d_map)
  2501 
  2502         sim_bounds_normal, sim_bounds_perp = self.pop_axis(

---

Lines 2521-2529

  2521 
  2522         def compute_integration_weight(grid_points):
  2523             grid_spacing = grid_points[1] - grid_points[0]
  2524             if grid_spacing == 0.0:
! 2525                 integration_weight = 1.0 / len(grid_points)
  2526             else:
  2527                 integration_weight = grid_points[1] - grid_points[0]
  2528 
  2529             return integration_weight

---

Lines 2541-2549

  2541                 np.minimum(bounds_perp[nonzero_dim][1], sim_bounds_perp[nonzero_dim][1]),
  2542             )
  2543 
  2544             if not verify_integration_interval(integration_bounds_perp):
! 2545                 return 0.0
  2546 
  2547             grid_points_linear = spacing_to_grid_points(
  2548                 adaptive_spacing, integration_bounds_perp[0], integration_bounds_perp[1]
  2549             )

---

Lines 2570-2578

  2570                 ),
  2571             )
  2572 
  2573             if not np.all([verify_integration_interval(b) for b in integration_bounds_perp]):
! 2574                 return 0.0
  2575 
  2576             grid_points_perp_1 = spacing_to_grid_points(
  2577                 adaptive_spacing, integration_bounds_perp[0][0], integration_bounds_perp[0][1]
  2578             )

---

[yaugenst-flex]

+99 / -346 love to see it 😄
Caught some minor issues but overall looks great, thanks @groberts-flex

[yaugenst-flex]

One last thing I spotted but then it's ready to go I think.

[yaugenst-flex]

Just noticed this, I think here since you always take the first perpendicular axis, this would break if the zero-extent axis is the second entry, right? I think the fix would just be grid_points[:, axis_perp[zero_dim]] = bounds_perp[zero_dim][0]