reduce computation time for adaptive_vjp_spacing via caching for GeometryGroup

[groberts-flex]

adapative_vjp_spacing takes longer to run now that we use the full eps_data inside of DerivativeInfo. For geometry groups this was running for every geometry and adding a significant overhead. Just caching the property didn't work because the derivative_info variable is copied and updated with new derivative paths in GeometryGroup. Instead, we allow the variable to be cached inside of DerivativeInfo.

---

Note

Low Risk
Performance-only refactor with scoped caching and test coverage; main risk is subtle cache leakage or incorrect reuse, which is explicitly checked and cleared.

Overview
Speeds up autograd shape-gradient evaluation for GeometryGroup by caching the permittivity-derived spacing used by DerivativeInfo.adaptive_vjp_spacing, avoiding recomputation for each child geometry.

This refactors the spacing calculation into a new DerivativeInfo.min_spacing_from_permittivity property plus a scoped cache_min_spacing_from_permittivity() context manager, and wraps GeometryGroup._compute_derivatives() so all child derivative computations share the cached value. Tests are updated/added to ensure caching doesn’t change VJP results and that cached state is cleared after use; the changelog notes the performance improvement.

^{Written by Cursor Bugbot for commit c2191cb. This will update automatically on new commits. Configure here.}

Greptile Overview

Greptile Summary

This PR optimizes the performance of adaptive_vjp_spacing for GeometryGroup by implementing a caching mechanism for the permittivity-based spacing calculation.

Key Changes:

• Extracted the core permittivity-based spacing computation into a new @cached_property method min_spacing_from_permittivity in DerivativeInfo
• Added a cached_min_spacing_from_permittivity field to enable explicit caching that survives dataclasses.replace() operations
• Added helper methods cache_min_spacing_from_permittivity() and uncache_min_spacing_from_permittivity() to control cache lifecycle
• Modified GeometryGroup._compute_derivatives() to cache the value once before the loop and uncache after, so all child geometries share the same cached value

Implementation Details:
The caching approach cleverly handles the fact that GeometryGroup creates copies of DerivativeInfo for each child geometry using updated_copy() (which uses dataclasses.replace()). Since @cached_property stores its cache in the instance's __dict__, it doesn't survive across replace() operations. The solution is to store the computed value in an explicit field (cached_min_spacing_from_permittivity) which DOES get copied during replace(). The @cached_property decorator checks this field first, so copied instances use the pre-computed value without recalculation.

Performance Impact:
Previously, adaptive_vjp_spacing would compute the permittivity-based spacing for every geometry in a GeometryGroup, even though all geometries share the same eps_data. With this change, it's computed once and reused across all geometries in the group.

Confidence Score: 4/5

• This PR is safe to merge with low risk - it's a performance optimization with a well-thought-out caching mechanism
• The implementation is sound and the caching logic correctly handles the interaction between @cached_property and dataclasses.replace(). The confidence is 4 (not 5) because this is a subtle performance optimization involving cached properties and dataclass copies, which could benefit from additional testing to verify the performance improvement and ensure no edge cases were missed.
• No files require special attention - all changes are straightforward

Important Files Changed

Filename	Overview
CHANGELOG.md	Added performance improvement entry for adaptive_vjp_spacing caching in GeometryGroup
tidy3d/components/autograd/derivative_utils.py	Refactored adaptive_vjp_spacing to cache permittivity-based spacing via new cached_min_spacing_from_permittivity field and @cached_property decorator
tidy3d/components/geometry/base.py	Added caching calls around geometry loop in GeometryGroup._compute_derivatives to enable cache sharing across geometries

Sequence Diagram

sequenceDiagram
    participant GG as GeometryGroup
    participant DI as DerivativeInfo (original)
    participant DI_copy as DerivativeInfo (copy)
    participant G1 as Geometry 1
    participant G2 as Geometry 2
    
    Note over GG: _compute_derivatives called
    GG->>DI: cache_min_spacing_from_permittivity()
    activate DI
    DI->>DI: Access min_spacing_from_permittivity (@cached_property)
    DI->>DI: Compute spacing from eps_data
    DI->>DI: Store in cached_min_spacing_from_permittivity field
    deactivate DI
    
    loop For each geometry in GeometryGroup
        GG->>DI: updated_copy(paths, bounds, ...)
        DI-->>DI_copy: Create shallow copy via dataclasses.replace()
        Note over DI_copy: cached_min_spacing_from_permittivity<br/>is copied to new instance
        
        GG->>G1: _compute_derivatives(DI_copy)
        activate G1
        G1->>DI_copy: adaptive_vjp_spacing()
        DI_copy->>DI_copy: Access min_spacing_from_permittivity
        Note over DI_copy: Check cached_min_spacing_from_permittivity<br/>Return cached value immediately
        DI_copy-->>G1: Return cached spacing
        G1-->>GG: Return gradient
        deactivate G1
    end
    
    GG->>DI: uncache_min_spacing_from_permittivity()
    Note over DI: Set cached_min_spacing_from_permittivity = None
    
    Note over GG: All geometries used cached value,<br/>avoiding redundant computation

Loading

^{Cursor Bugbot found 1 potential issue for commit c2191cb}

[greptile-apps]

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[marcorudolphflex]

Good to have that! How many impact do we observe from that caching?

[yaugenst-flex]

Thanks @groberts-flex LGTM! Agree with @marcorudolphflex's comment that this might be nicer with a context manager, alternatively a try/finally to ensure things are cleaned up.
Only nit would be that it might be good to add a unit test to ensure that cached vs uncached adaptive_vjp_spacing results are identical. Helps catch future regressions.

[groberts-flex]

Good to have that! How many impact do we observe from that caching?

It is really beneficial in geometry groups with a lot of geometries. For example, in the performance profiling test, I have a 35x35 array of geometries in a geometry group. Without caching, the postprocess_adj function takes 602sec. With caching it takes 327sec.

[groberts-flex]

Thanks @groberts-flex LGTM! Agree with @marcorudolphflex's comment that this might be nicer with a context manager, alternatively a try/finally to ensure things are cleaned up. Only nit would be that it might be good to add a unit test to ensure that cached vs uncached adaptive_vjp_spacing results are identical. Helps catch future regressions.

great point, I'll definitely add a test!

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.}

[cursor]

Test mock doesn't propagate cache in updated_copy

Low Severity

The SimpleDerivativeInfo.updated_copy method does not include cached_min_spacing_from_permittivity in the data dictionary being copied. The real DerivativeInfo.updated_copy() uses dataclasses.replace() which copies all fields automatically. This means the caching mechanism added in this PR won't work correctly for this test mock - when GeometryGroup._compute_derivatives creates copies via updated_copy() inside the caching context, the cached value won't propagate to child instances.

 

[github-actions]

Diff Coverage

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

• tidy3d/components/autograd/derivative_utils.py (100%)
• tidy3d/components/geometry/base.py (88.9%): Missing lines 3607

Summary

• Total: 42 lines
• Missing: 1 line
• Coverage: 97%

tidy3d/components/geometry/base.py

Lines 3603-3611

  3603 
  3604                 vjp_dict_geo = geo._compute_derivatives(geo_info)
  3605 
  3606                 if len(vjp_dict_geo) != 1:
! 3607                     raise AssertionError("Got multiple gradients for single geometry field.")
  3608 
  3609                 grad_vjps[field_path] = vjp_dict_geo.popitem()[1]
  3610 
  3611         return grad_vjps

---