autograd config section

Author: yaugenst-flex

tests/config/test_manager.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import numpy as np
 import pytest
 
 
@@ -39,3 +40,38 @@ def test_builtin_profiles(profile, config_manager):
     config_manager.switch_profile(profile)
     web = config_manager.get_section("web")
     assert web.s3_region is not None
+
+
+def test_autograd_defaults(config_manager):
+    autograd = config_manager.get_section("autograd")
+    assert autograd.min_wvl_fraction == pytest.approx(5e-2)
+    assert autograd.points_per_wavelength == 10
+    assert autograd.monitor_interval_poly == (1, 1, 1)
+    assert autograd.quadrature_sample_fraction == pytest.approx(0.4)
+    assert autograd.gradient_precision == "single"
+    assert autograd.max_traced_structures == 500
+    assert autograd.max_adjoint_per_fwd == 10
+
+
+def test_autograd_update_section(config_manager):
+    config_manager.update_section(
+        "autograd",
+        min_wvl_fraction=0.08,
+        points_per_wavelength=12,
+        solver_freq_chunk_size=3,
+        gradient_precision="double",
+        max_traced_structures=600,
+        max_adjoint_per_fwd=7,
+    )
+    autograd = config_manager.get_section("autograd")
+    assert autograd.min_wvl_fraction == pytest.approx(0.08)
+    assert autograd.points_per_wavelength == 12
+    assert autograd.solver_freq_chunk_size == 3
+    assert autograd.gradient_precision == "double"
+    assert autograd.max_traced_structures == 600
+    assert autograd.max_adjoint_per_fwd == 7
+
+    from tidy3d.components.autograd import derivative_utils as autograd_utils
+
+    assert autograd_utils.GRADIENT_PRECISION == "double"
+    assert autograd_utils.GRADIENT_DTYPE_FLOAT is np.float64

tests/test_components/autograd/test_autograd.py

@@ -19,14 +19,13 @@
 
 import tidy3d as td
 import tidy3d.web as web
-from tidy3d.components.autograd.constants import (
-    MAX_NUM_TRACED_STRUCTURES,
-    MIN_WVL_FRACTION_CYLINDER_DISCRETIZE,
+from tidy3d.components.autograd.derivative_utils import (
     MINIMUM_SPACING_FRACTION,
+    DerivativeInfo,
 )
-from tidy3d.components.autograd.derivative_utils import DerivativeInfo
 from tidy3d.components.autograd.utils import is_tidy_box
 from tidy3d.components.data.data_array import DataArray
+from tidy3d.config import config
 from tidy3d.exceptions import AdjointError
 from tidy3d.plugins.polyslab import ComplexPolySlab
 from tidy3d.web import run, run_async
@@ -1251,7 +1250,8 @@ def test_too_many_traced_structures(monkeypatch, use_emulated_run):
     def make_sim(*args):
         structure = make_structures(*args)[structure_key]
         return SIM_BASE.updated_copy(
-            structures=(MAX_NUM_TRACED_STRUCTURES + 1) * [structure], monitors=[monitor]
+            structures=(config.autograd.max_traced_structures + 1) * [structure],
+            monitors=[monitor],
         )
 
     def objective(*args):
@@ -1808,7 +1808,7 @@ def test_cylinder_discretization(eps_real):
     ):
         cylinder = td.Cylinder(axis=2, length=info.wavelength_min, radius=2 * info.wavelength_min)
 
-        expected_wvl_mat = info.wavelength_min * MIN_WVL_FRACTION_CYLINDER_DISCRETIZE
+        expected_wvl_mat = info.wavelength_min * config.autograd.min_wvl_fraction
         wvl_mat = cylinder._discretization_wavelength(derivative_info=info)
 
         assert np.isclose(expected_wvl_mat, wvl_mat), (

tidy3d/components/autograd/constants.py

@@ -13,12 +13,6 @@
 from tidy3d.constants import C_0, EPSILON_0, LARGE_NUMBER, MU_0
 from tidy3d.log import log
 
-from .constants import (
-    DEFAULT_WAVELENGTH_FRACTION,
-    GRADIENT_DTYPE_COMPLEX,
-    GRADIENT_DTYPE_FLOAT,
-    MINIMUM_SPACING_FRACTION,
-)
 from .types import PathType
 from .utils import get_static
 
@@ -27,6 +21,33 @@
 EpsType = Union[tidycomplex, FreqDataArray]
 
 
+GRADIENT_PRECISION = "single"  # Options: "single", "double"
+GRADIENT_DTYPE_FLOAT = np.float32
+GRADIENT_DTYPE_COMPLEX = np.complex64
+
+
+def set_gradient_precision(precision: str) -> None:
+    """Update global gradient precision and derived dtypes."""
+
+    if precision not in {"single", "double"}:
+        raise ValueError("gradient_precision must be 'single' or 'double'")
+
+    global GRADIENT_PRECISION, GRADIENT_DTYPE_FLOAT, GRADIENT_DTYPE_COMPLEX
+    GRADIENT_PRECISION = precision
+    if precision == "single":
+        GRADIENT_DTYPE_FLOAT = np.float32
+        GRADIENT_DTYPE_COMPLEX = np.complex64
+    else:
+        GRADIENT_DTYPE_FLOAT = np.float64
+        GRADIENT_DTYPE_COMPLEX = np.complex128
+
+
+MINIMUM_SPACING_FRACTION = 1e-2
+
+
+set_gradient_precision("single")
+
+
 class LazyInterpolator:
     """Lazy wrapper for interpolators that creates them on first access."""
 
@@ -711,7 +732,7 @@ def _project_in_basis(
 
     def adaptive_vjp_spacing(
         self,
-        wl_fraction: float = DEFAULT_WAVELENGTH_FRACTION,
+        wl_fraction: Optional[float] = None,
         min_allowed_spacing_fraction: float = MINIMUM_SPACING_FRACTION,
     ) -> float:
         """Compute adaptive spacing for finite-difference gradient evaluation.
@@ -721,8 +742,9 @@ def adaptive_vjp_spacing(
 
         Parameters
         ----------
-        wl_fraction : float = 0.1
-            Fraction of wavelength/skin depth to use as spacing.
+        wl_fraction : float, optional
+            Fraction of wavelength/skin depth to use as spacing. Defaults to the configured
+            ``autograd.default_wavelength_fraction`` when ``None``.
         min_allowed_spacing_fraction : float = 1e-2
             Minimum allowed spacing fraction of free space wavelength to
             prevent numerical issues.
@@ -732,6 +754,11 @@ def adaptive_vjp_spacing(
         float
             Adaptive spacing value for gradient evaluation.
         """
+        if wl_fraction is None:
+            from tidy3d.config import config
+
+            wl_fraction = config.autograd.default_wavelength_fraction
+
         # handle FreqDataArray or scalar eps_in
         if isinstance(self.eps_in, FreqDataArray):
             eps_real = np.asarray(self.eps_in.values, dtype=np.complex128).real

tidy3d/components/autograd/derivative_utils.py

@@ -17,6 +17,7 @@
 except ImportError:
     pass
 
+import tidy3d.components.autograd.derivative_utils as autograd_utils
 from tidy3d.compat import _shapely_is_older_than
 from tidy3d.components.autograd import (
     AutogradFieldMap,
@@ -25,7 +26,6 @@
     TracedSize,
     get_static,
 )
-from tidy3d.components.autograd.constants import GRADIENT_DTYPE_FLOAT
 from tidy3d.components.autograd.derivative_utils import (
     DerivativeInfo,
     FieldData,
@@ -2684,7 +2684,7 @@ def _derivative_face_dielectric(
 
         Parameters
         ----------
-        dtype : np.dtype = GRADIENT_DTYPE_FLOAT
+        dtype : np.dtype = autograd_utils.GRADIENT_DTYPE_FLOAT
             Data type for interpolation coordinates and values.
 
         dim_normal : str
@@ -2757,7 +2757,7 @@ def _derivative_face_pec(
 
         Parameters
         ----------
-        dtype : np.dtype = GRADIENT_DTYPE_FLOAT
+        dtype : np.dtype = autograd_utils.GRADIENT_DTYPE_FLOAT
             Data type for interpolation coordinates and values.
 
         dim_normal : str
@@ -3613,7 +3613,7 @@ 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(
-            dtype=GRADIENT_DTYPE_FLOAT
+            dtype=autograd_utils.GRADIENT_DTYPE_FLOAT
         )
 
         for field_path in derivative_info.paths: