reduce computation time for adaptive_vjp_spacing via caching for GeometryGroup

CHANGELOG.md

@@ -33,6 +33,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added deprecation warning for `TemperatureMonitor` and `SteadyPotentialMonitor` when `unstructured` parameter is not explicitly set. The default value of `unstructured` will change from `False` to `True` in the next release.
 - Added deprecation warning for ``TemperatureMonitor`` and ``SteadyPotentialMonitor`` when ``unstructured`` parameter is not explicitly set. The default value of ``unstructured`` will change from ``False`` to ``True`` after the 2.11 release.
 - Added validation to `GaussianDoping` to ensure `ref_con < concentration`, validate `source` face identifier, and warn the user when the box size is not sufficient for the specified transition width.
+- Reduced computation time of `adaptive_vjp_spacing` for `GeometryGroup` by allowing permittivity based spacing value to be cached.
 
 ### Fixed
 - Fixed intermittent "API key not found" errors in parallel job launches by making configuration directory detection race-safe.

tests/test_components/autograd/test_autograd.py

@@ -5,6 +5,7 @@
 import cProfile
 import typing
 import warnings
+from contextlib import nullcontext
 from dataclasses import dataclass
 from importlib import reload
 from os.path import join
@@ -1967,6 +1968,96 @@ def test_adaptive_spacing(eps_real):
         assert np.isclose(expected_vjp_spacing, vjp_spacing), "Unexpected adaptive vjp spacing!"
 
 
+def test_adaptive_spacing_cache(rng, redirect_stdout_to_stderr, monkeypatch):
+    """Ensure the cache is not affecting `GeometryGroup` results or persisting after derivative computation."""
+    x_geom = [-0.5, 0.5]
+    y_geom = [-0.5, 0.5]
+
+    radius = 0.3
+
+    geometries = []
+    for x_center in x_geom:
+        for y_center in y_geom:
+            geometries.append(
+                td.Cylinder(center=(x_center, y_center, 0.0), length=0.5, radius=radius)
+            )
+
+    field_paths = []
+    for idx in range(len(geometries)):
+        field_paths.append(("geometries", idx, "radius"))
+
+    geometry = td.GeometryGroup(geometries=geometries)
+
+    eps_keys = ["eps_xx", "eps_yy", "eps_zz"]
+
+    N = 20
+    xcoord = np.linspace(-1, 1, N)
+    ycoord = np.linspace(-1, 1, N)
+    zcoord = np.linspace(-1, 1, N)
+
+    eps_out_data = rng.uniform(1, 2, (N, N, N, 1))
+    eps_in_data = rng.uniform(1, 2, (N, N, N, 1))
+
+    freq = 1.94e14
+
+    def random_scalar_data_array():
+        return td.ScalarFieldDataArray(
+            rng.uniform(1, 2, (N, N, N, 1)),
+            coords={"x": xcoord, "y": ycoord, "z": zcoord, "f": [freq]},
+        )
+
+    E_fwd = {key: random_scalar_data_array() for key in ["Ex", "Ey", "Ez"]}
+    E_adj = {key: random_scalar_data_array() for key in ["Ex", "Ey", "Ez"]}
+    D_fwd = {key: random_scalar_data_array() for key in ["Ex", "Ey", "Ez"]}
+    D_adj = {key: random_scalar_data_array() for key in ["Ex", "Ey", "Ez"]}
+    E_der_map = {key: random_scalar_data_array() for key in ["Ex", "Ey", "Ez"]}
+    D_der_map = {key: random_scalar_data_array() for key in ["Ex", "Ey", "Ez"]}
+
+    derivative_info = DerivativeInfo(
+        paths=tuple(field_paths),
+        E_der_map=E_der_map,
+        D_der_map=D_der_map,
+        E_fwd=E_fwd,
+        D_fwd=D_fwd,
+        E_adj=E_adj,
+        D_adj=D_adj,
+        eps_data={
+            key: td.ScalarFieldDataArray(
+                rng.uniform(1, 2, (N, N, N, 1)),
+                coords={"x": xcoord, "y": ycoord, "z": zcoord, "f": [freq]},
+            )
+            for key in eps_keys
+        },
+        frequencies=[freq],
+        bounds=((-1, -1, -1), (1, 1, 1)),
+        eps_out=td.ScalarFieldDataArray(
+            eps_out_data, coords={"x": xcoord, "y": ycoord, "z": zcoord, "f": [freq]}
+        ),
+        eps_in=td.ScalarFieldDataArray(
+            eps_in_data, coords={"x": xcoord, "y": ycoord, "z": zcoord, "f": [freq]}
+        ),
+        bounds_intersect=((-1, -1, -1), (1, 1, 1)),
+        simulation_bounds=((-2, -2, -2), (2, 2, 2)),
+    )
+
+    vjp_with_cache = geometry._compute_derivatives(derivative_info)
+
+    assert derivative_info.cached_min_spacing_from_permittivity is None, (
+        "Unexpected cached variable persistence."
+    )
+
+    monkeypatch.setattr(
+        derivative_info, "cache_min_spacing_from_permittivity", lambda: nullcontext()
+    )
+
+    vjp_without_cache = geometry._compute_derivatives(derivative_info)
+
+    for k, v in vjp_with_cache.items():
+        assert v == vjp_without_cache[k], (
+            "Geometry group computation changed when running with cache."
+        )
+
+
 @pytest.mark.parametrize("eps_real", [1e6, -1e8])
 def test_cylinder_discretization(eps_real):
     freq = 5e9
@@ -3279,6 +3370,18 @@ class SimpleDerivativeInfo:
         bounds_intersect: tuple
         simulation_bounds: tuple
         interpolators: dict | None = None
+        cached_min_spacing_from_permittivity: float | None = None
+        eps_data = {
+            key: td.ScalarFieldDataArray(
+                [[[[2.0]]]], coords={"x": [0], "y": [0], "z": [0], "f": [200e12]}
+            )
+            for key in ["eps_xx", "eps_yy", "eps_zz"]
+        }
+        frequencies = [200e12]
+
+        cache_min_spacing_from_permittivity = DerivativeInfo.cache_min_spacing_from_permittivity
+        min_spacing_from_permittivity = DerivativeInfo.min_spacing_from_permittivity
+        wavelength_min = property(lambda self: DerivativeInfo.wavelength_min.fget(self))
 
         def create_interpolators(self, dtype: float = float):
             return self.interpolators or {}
@@ -3290,6 +3393,7 @@ def updated_copy(self, **kwargs):
                 "bounds_intersect": self.bounds_intersect,
                 "simulation_bounds": self.simulation_bounds,
                 "interpolators": self.interpolators,
+                "cached_min_spacing_from_permittivity": self.cached_min_spacing_from_permittivity,
             }
             data.update({k: v for k, v in kwargs.items() if k in data})
             return SimpleDerivativeInfo(**data)

tests/test_components/autograd/test_autograd_polyslab.py

@@ -132,6 +132,9 @@ def __init__(
             else ((-2e3, -2e3, -2e3), (2e3, 2e3, 2e3))
         )
 
+        self.cached_min_spacing_from_permittivity = None
+
+    min_spacing_from_permittivity = DerivativeInfo.min_spacing_from_permittivity
     adaptive_vjp_spacing = DerivativeInfo.adaptive_vjp_spacing
     wavelength_min = property(lambda self: DerivativeInfo.wavelength_min.fget(self))
     wavelength_max = property(lambda self: DerivativeInfo.wavelength_max.fget(self))

tidy3d/components/autograd/derivative_utils.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Iterator
+from contextlib import contextmanager
 from dataclasses import dataclass, field, replace
 from functools import reduce
 from typing import Any, Callable, Optional
@@ -152,6 +154,10 @@ class DerivativeInfo:
     sharing the same field data. This significantly improves performance for
     GeometryGroup processing."""
 
+    cached_min_spacing_from_permittivity: Optional[float] = None
+    """Cached `min_spacing_from_permittivity` to be used for objects like GeometryGroup
+    to avoid recomputing this value multiple times in `adaptive_vjp_spacing`."""
+
     # private cache for interpolators
     _interpolators_cache: dict = field(default_factory=dict, init=False, repr=False)
 
@@ -915,6 +921,54 @@ def project_der_map_to_axis(
                 projected[key] = field_map[key]
         return projected
 
+    @property
+    def min_spacing_from_permittivity(self):
+        if self.cached_min_spacing_from_permittivity is not None:
+            return self.cached_min_spacing_from_permittivity
+
+        def spacing_by_permittivity(eps_array):
+            eps_real = np.asarray(eps_array.values, dtype=np.complex128).real
+
+            dx_candidates = []
+            max_frequency = np.max(self.frequencies)
+
+            # wavelength-based sampling for dielectrics
+            if np.any(eps_real > 0):
+                eps_max = eps_real[eps_real > 0].max()
+                lambda_min = self.wavelength_min / np.sqrt(eps_max)
+                dx_candidates.append(lambda_min)
+
+            # skin depth sampling for metals
+            if np.any(eps_real <= 0):
+                omega = 2 * np.pi * max_frequency
+                eps_neg = eps_real[eps_real <= 0]
+                delta_min = C_0 / (omega * np.sqrt(np.abs(eps_neg).max()))
+                dx_candidates.append(delta_min)
+
+            computed_spacing = min(dx_candidates)
+
+            return computed_spacing
+
+        eps_spacings = [
+            spacing_by_permittivity(eps_array) for _, eps_array in self.eps_data.items()
+        ]
+        min_spacing = np.min(eps_spacings)
+
+        return min_spacing
+
+    @contextmanager
+    def cache_min_spacing_from_permittivity(self) -> Iterator[None]:
+        """
+        Cache min_spacing_from_permittivity for the duration of the block. Cache
+        is always cleared on exit.
+        """
+
+        self.cached_min_spacing_from_permittivity = self.min_spacing_from_permittivity
+        try:
+            yield
+        finally:
+            self.cached_min_spacing_from_permittivity = None
+
     def adaptive_vjp_spacing(
         self,
         wl_fraction: Optional[float] = None,
@@ -948,33 +1002,7 @@ def adaptive_vjp_spacing(
             if min_allowed_spacing_fraction is None:
                 min_allowed_spacing_fraction = config.adjoint.minimum_spacing_fraction
 
-        def spacing_by_permittivity(eps_array):
-            eps_real = np.asarray(eps_array.values, dtype=np.complex128).real
-
-            dx_candidates = []
-            max_frequency = np.max(self.frequencies)
-
-            # wavelength-based sampling for dielectrics
-            if np.any(eps_real > 0):
-                eps_max = eps_real[eps_real > 0].max()
-                lambda_min = self.wavelength_min / np.sqrt(eps_max)
-                dx_candidates.append(wl_fraction * lambda_min)
-
-            # skin depth sampling for metals
-            if np.any(eps_real <= 0):
-                omega = 2 * np.pi * max_frequency
-                eps_neg = eps_real[eps_real <= 0]
-                delta_min = C_0 / (omega * np.sqrt(np.abs(eps_neg).max()))
-                dx_candidates.append(wl_fraction * delta_min)
-
-            computed_spacing = min(dx_candidates)
-
-            return computed_spacing
-
-        eps_spacings = [
-            spacing_by_permittivity(eps_array) for _, eps_array in self.eps_data.items()
-        ]
-        computed_spacing = np.min(eps_spacings)
+        computed_spacing = wl_fraction * self.min_spacing_from_permittivity
 
         min_allowed_spacing = self.wavelength_min * min_allowed_spacing_fraction
 

tidy3d/components/geometry/base.py

@@ -3585,26 +3585,28 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField
         # create interpolators once for all geometries to avoid redundant field data conversions
         interpolators = derivative_info.interpolators or derivative_info.create_interpolators()
 
-        for field_path in derivative_info.paths:
-            _, index, *geo_path = field_path
-            geo = self.geometries[index]
-            # pass pre-computed interpolators if available
-            geo_info = derivative_info.updated_copy(
-                paths=[tuple(geo_path)],
-                bounds=geo.bounds,
-                bounds_intersect=self.bounds_intersection(
-                    geo.bounds, derivative_info.simulation_bounds
-                ),
-                deep=False,
-                interpolators=interpolators,
-            )
+        with derivative_info.cache_min_spacing_from_permittivity():
+            for field_path in derivative_info.paths:
+                _, index, *geo_path = field_path
+
+                geo = self.geometries[index]
+                # pass pre-computed interpolators if available
+                geo_info = derivative_info.updated_copy(
+                    paths=[tuple(geo_path)],
+                    bounds=geo.bounds,
+                    bounds_intersect=self.bounds_intersection(
+                        geo.bounds, derivative_info.simulation_bounds
+                    ),
+                    deep=False,
+                    interpolators=interpolators,
+                )
 
-            vjp_dict_geo = geo._compute_derivatives(geo_info)
+                vjp_dict_geo = geo._compute_derivatives(geo_info)
 
-            if len(vjp_dict_geo) != 1:
-                raise AssertionError("Got multiple gradients for single geometry field.")
+                if len(vjp_dict_geo) != 1:
+                    raise AssertionError("Got multiple gradients for single geometry field.")
 
-            grad_vjps[field_path] = vjp_dict_geo.popitem()[1]
+                grad_vjps[field_path] = vjp_dict_geo.popitem()[1]
 
         return grad_vjps