code review and bugbot

Author: Gregory Roberts

tests/test_components/autograd/test_autograd.py

@@ -1386,7 +1386,10 @@ def polyslab_custom_vjp_bad_arg_name(polyslab, d_info):
 
     with pytest.raises(
         td.exceptions.AdjointError,
-        match="CustomVJPConfig compute_derivatives function should accept two arguments and it currently accepts 3 arguments.",
+        match=(
+            "CustomVJPConfig compute_derivatives function should accept two arguments "
+            r"\(target, derivative_info\), and it currently accepts 3 arguments\."
+        ),
     ):
         CustomVJPConfig(
             structure=td.PolySlab,

tidy3d/web/api/autograd/autograd.py

@@ -105,23 +105,6 @@ def insert_numerical_structures_static(
     return updated_simulation
 
 
-def _normalize_simulations_input(
-    simulations: Union[dict[str, td.Simulation], tuple[td.Simulation], list[td.Simulation]],
-) -> dict[str, td.Simulation]:
-    """Normalize simulations to a dict keyed by task name."""
-
-    if isinstance(simulations, dict):
-        return simulations
-
-    normalized: dict[str, td.Simulation] = {}
-
-    for idx, sim in enumerate(simulations):
-        task_name = Tidy3dStub(simulation=sim).get_default_task_name() + f"_{idx + 1}"
-        normalized[task_name] = sim
-
-    return normalized
-
-
 def has_traced_numerical_structures(
     numerical_structures: Union[
         tuple[NumericalStructureConfig, ...],
@@ -137,6 +120,10 @@ def has_traced_numerical_structures(
         else numerical_structures
     )
     for cfg in iterable_structures:
+        if not isinstance(cfg, NumericalStructureConfig):
+            raise AdjointError(
+                "Entries in 'numerical_structures' must be NumericalStructureConfig instances."
+            )
         parameters = cfg.parameters
         if hasbox(parameters):
             return True
@@ -746,6 +733,13 @@ def _expand_spec(
             validate_numerical_structure_parameters(
                 numerical_structures=numerical_structures_configs
             )
+        if any(numerical_structures.values()) and not all(
+            isinstance(sim, td.Simulation) for sim in sim_dict.values()
+        ):
+            raise AdjointError(
+                "numerical_structures is only supported for 'Simulation' workflows in "
+                "run_async_custom."
+            )
 
     custom_vjp = _expand_spec(
         fn_arg=custom_vjp,
@@ -755,18 +749,14 @@ def _expand_spec(
         arg_name="custom_vjp",
     )
 
-    simulations = sim_dict
-
     path_dir = Path(path_dir)
 
-    simulations_norm = _normalize_simulations_input(simulations)
-
     traced_numerical_structures = bool(numerical_structures) and any(
         has_traced_numerical_structures(numerical_structure)
         for _, numerical_structure in numerical_structures.items()
     )
     should_use_autograd_async = (
-        is_valid_for_autograd_async(simulations_norm) or traced_numerical_structures
+        is_valid_for_autograd_async(sim_dict) or traced_numerical_structures
     )
 
     if should_use_autograd_async:
@@ -782,13 +772,13 @@ def _expand_spec(
             expanded_custom_vjp_dict = {}
             for sim_key, custom_vjp_entry in custom_vjp.items():
                 expanded_custom_vjp_dict[sim_key] = expand_custom_vjp(
-                    custom_vjp_entry, simulations_norm[sim_key]
+                    custom_vjp_entry, sim_dict[sim_key]
                 )
         else:
             expanded_custom_vjp_dict = None
 
         return _run_async(
-            simulations=simulations_norm,
+            simulations=sim_dict,
             folder_name=folder_name,
             path_dir=path_dir,
             callback_url=callback_url,
@@ -811,16 +801,16 @@ def _expand_spec(
         simulations_static = {
             name: (
                 insert_numerical_structures_static(
-                    simulation=simulations_norm[name],
+                    simulation=sim_dict[name],
                     numerical_structures=numerical_structures[name],
                 )
                 if numerical_structures[name]
-                else simulations_norm[name]
+                else sim_dict[name]
             )
-            for name in simulations_norm
+            for name in sim_dict
         }
     else:
-        simulations_static = simulations_norm
+        simulations_static = sim_dict
 
     return asynchronous_webapi.run_async(
         simulations=simulations_static,

tidy3d/web/api/autograd/types.py

@@ -3,17 +3,20 @@
 import inspect
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, NamedTuple, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, NamedTuple, Optional, Union
 
 import numpy as np
 
 from tidy3d.exceptions import AdjointError
 
 if TYPE_CHECKING:
+    from tidy3d.components.autograd.derivative_utils import DerivativeInfo
+    from tidy3d.components.autograd.types import PathType
     from tidy3d.components.autograd import AutogradFieldMap
     from tidy3d.components.geometry.utils import GeometryType
     from tidy3d.components.medium import MediumType
     from tidy3d.components.simulation import Simulation
+    from tidy3d.components.structure import Structure
     from tidy3d.components.types import ArrayLike
 
 
@@ -45,10 +48,10 @@ def cylinder_vjp(parameters, derivative_info):
         )
     """
 
-    create: Callable
+    create: Callable[[ArrayLike], Structure]
     """Function that creates the structure from static ``parameters``."""
 
-    compute_derivatives: Callable
+    compute_derivatives: Callable[[ArrayLike, DerivativeInfo], dict[PathType, Any]]
     """Function that computes numerical gradients for ``("numerical", index, param_i)`` paths.
     Signature: ``compute_derivatives(parameters, derivative_info) -> dict[path, gradient]``.
     """
@@ -57,11 +60,18 @@ def cylinder_vjp(parameters, derivative_info):
     """1D parameter vector consumed by ``create`` and ``compute_derivatives``."""
 
     def __post_init__(self) -> None:
+        self._validate_callables()
+        self._validate_create_signature()
+        self._validate_compute_derivatives_signature()
+        self._validate_parameters()
+
+    def _validate_callables(self) -> None:
         if not callable(self.create):
             raise AdjointError("NumericalStructureConfig.create must be callable.")
         if not callable(self.compute_derivatives):
             raise AdjointError("NumericalStructureConfig.compute_derivatives must be callable.")
 
+    def _validate_create_signature(self) -> None:
         create_sig = inspect.signature(self.create)
         create_arg_names = list(create_sig.parameters.keys())
         if len(create_arg_names) != 1:
@@ -71,6 +81,7 @@ def __post_init__(self) -> None:
                 f"accepts {len(create_arg_names)} arguments."
             )
 
+    def _validate_compute_derivatives_signature(self) -> None:
         vjp_sig = inspect.signature(self.compute_derivatives)
         vjp_arg_names = list(vjp_sig.parameters.keys())
         if len(vjp_arg_names) != 2:
@@ -87,6 +98,7 @@ def __post_init__(self) -> None:
                 f"{vjp_arg_names[1]} but it should be derivative_info."
             )
 
+    def _validate_parameters(self) -> None:
         try:
             array_params = np.asarray(self.parameters)
         except Exception as exc:
@@ -126,7 +138,7 @@ def polyslab_vjp(polyslab, derivative_info):
     Can be an index or a geometry/medium type (expanded to matching indices).
     """
 
-    compute_derivatives: Callable
+    compute_derivatives: Callable[[GeometryType | MediumType, DerivativeInfo], dict[PathType, Any]]
     """Function for computing the targeted vjp value. The function should accept the geometry or medium in the
     structure depending on if this is a geometry or medium path (see path_key) as the first argument. The second
     argument should accept a DerivativeInfo object that contains important for computing the gradient. The function
@@ -141,15 +153,23 @@ def polyslab_vjp(polyslab, derivative_info):
     """
 
     def __post_init__(self) -> None:
+        self._validate_callable()
+        self._validate_compute_derivatives_signature()
+
+    def _validate_callable(self) -> None:
         if not callable(self.compute_derivatives):
             raise AdjointError("CustomVJPConfig.compute_derivatives must be callable.")
 
+    def _validate_compute_derivatives_signature(self) -> None:
         vjp_sig = inspect.signature(self.compute_derivatives)
         vjp_arg_names = list(vjp_sig.parameters.keys())
         if len(vjp_arg_names) != 2:
             raise AdjointError(
                 "CustomVJPConfig compute_derivatives function should accept two arguments "
-                f"and it currently accepts {len(vjp_arg_names)} arguments."
+                "(target, derivative_info), and it currently accepts "
+                f"{len(vjp_arg_names)} arguments. The target is the geometry or medium "
+                "instance selected by path_key, and derivative_info contains the field "
+                "data and path metadata needed to compute the VJP."
             )
         if vjp_arg_names[1] != "derivative_info":
             raise AdjointError(