add numerical_structures hook to custom run functions

[groberts-flex]

The numerical_structures function is another hook to go alongside custom_vjp. In custom_vjp we offer the ability to define a vjp function to compute gradients for a structure or set of structures that already exist in the simulation. With numerical_structures we allow the creation of a structure from a set of traced parameters based on a provided function. This structure is inserted into the simulation and on the backwards pass we require a vjp function to compute the gradients based on the derivative_info.

An example use case for this hook is if we have a generation function for a tidy3d structure based on a set of parameters where the generation function can't be traced through directly or is difficult to trace through. We can use finite differences in this case to compute gradients for our original parameters based on the adjoint fields and created simulation geometry/permittivity.

---

Note

Medium Risk
Touches core autograd execution and adjoint postprocessing paths, adding new parameter/structure routing and validation logic; mistakes could surface as incorrect gradients or run-time failures across simulation and batch workflows.

Overview
Adds a new numerical_structures extension point alongside custom_vjp, allowing callers of run_custom/run_async_custom (and ComponentModeler _run_local) to create and insert extra structures from traced parameters and provide a compute_derivatives(parameters, derivative_info) VJP for those parameters.

Updates the autograd pipeline to (1) validate config/signatures and parameter shapes, (2) insert numerical structures (even when not traced) for non-autograd runs, (3) route traced numerical parameters through a new ("numerical", structure_index, param_index) namespace, and (4) compute their gradients during adjoint postprocessing via the supplied numerical VJP while ensuring custom_vjp only targets the original ("structures", ...) namespace.

Adds changelog entry plus comprehensive numerical tests comparing adjoint vs finite-difference gradients for both Simulation and ComponentModeler workflows, and new negative tests for invalid configs and unsupported workflow types.

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

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

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

[cursor]

Wrong custom_vjp structure index targets numerical structure

High Severity

CustomVJPConfig(structure=3, ...) targets the wrong structure. The base simulation has structures [input_waveguide, output_waveguide, sphere_structure] at indices 0, 1, 2. After the numerical ring structure is appended by _insert_numerical_structures, the ring lands at index 3. Since vjp_sphere is designed to handle the sphere's geometry paths (radius, center), the config needs structure=2. With structure=3, verify_custom_vjp will raise an AdjointError because index 3 only exists in the "numerical" path namespace, not the "structures" namespace that custom VJP validation checks against.

 

[cursor]

DerivativeInfo paths incorrect when both VJP types active

Medium Severity

When both structure_paths and use_numerical_vjp are true for the same structure index, the DerivativeInfo is created with paths=structure_paths (since a non-empty tuple is truthy), but the numerical VJP function at line 514 then receives a derivative_info containing structure-namespace paths instead of numerical-namespace paths. The current setup_run design prevents overlapping indices, but the logic is fragile — a single DerivativeInfo can't correctly serve both consumers with different path formats.

 

[github-actions]

Diff Coverage

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

• tidy3d/plugins/smatrix/run.py (57.1%): Missing lines 165,196-197,214-215,227
• tidy3d/web/api/autograd/autograd.py (87.0%): Missing lines 78,82-83,88,129,131-133,135,155,158-159,206,589,729,758,766,783,871,1576
• tidy3d/web/api/autograd/backward.py (82.6%): Missing lines 184,519,525,534,536-537,540,544
• tidy3d/web/api/autograd/types.py (100%)

Summary

• Total: 224 lines
• Missing: 34 lines
• Coverage: 84%

tidy3d/plugins/smatrix/run.py

Lines 161-169

  161 
  162     sims = modeler.sim_dict
  163 
  164     if isinstance(numerical_structures, NumericalStructureConfig):
! 165         numerical_structures = (numerical_structures,)
  166 
  167     traced_numerical_structures = numerical_structures and web_ag.has_traced_numerical_structures(
  168         numerical_structures
  169     )

---

Lines 192-201

  192         if not local_gradient:
  193             if custom_vjp is not None:
  194                 raise AdjointError("custom_vjp specified for a remote gradient not supported.")
  195 
! 196             if traced_numerical_structures:
! 197                 raise AdjointError(
  198                     "ComponentModeler autograd with traced numerical structures requires local_gradient=True."
  199                 )
  200 
  201         expanded_custom_vjp_dict = None

---

Lines 210-219

  210                     custom_vjp_entry, sims[sim_key]
  211                 )
  212 
  213         if numerical_structures:
! 214             web_ag.validate_numerical_structure_parameters(numerical_structures)
! 215             numerical_structures = dict.fromkeys(sims, numerical_structures)
  216 
  217         sim_data_map = _run_async(
  218             simulations=sims,
  219             numerical_structures=numerical_structures,

---

Lines 223-231

  223 
  224         return compose_modeler_data_from_batch_data(modeler=modeler, batch_data=sim_data_map)
  225 
  226     if numerical_structures is not None:
! 227         modeler = modeler.updated_copy(
  228             simulation=web_ag.insert_numerical_structures_static(
  229                 simulation=modeler.simulation, numerical_structures=numerical_structures
  230             )
  231         )

---

tidy3d/web/api/autograd/autograd.py

Lines 74-92

  74     """Convert numerical structure parameters to a static 1D NumPy array."""
  75     from tidy3d.components.autograd import get_static
  76 
  77     if isinstance(parameters, dict):
! 78         raise AdjointError("NumericalStructureConfig.parameters must not be a dict.")
  79 
  80     try:
  81         parameters_array = np.asarray(parameters)
! 82     except Exception as exc:
! 83         raise AdjointError(
  84             "NumericalStructureConfig.parameters must be coercible to a 1D numpy array."
  85         ) from exc
  86 
  87     if parameters_array.ndim != 1:
! 88         raise AdjointError("Parameters for each numerical structure must be 1D array-like.")
  89 
  90     return np.asarray([get_static(param) for param in parameters_array])
  91 

---

Lines 125-139

  125 
  126     if isinstance(simulations, dict):
  127         return simulations
  128 
! 129     normalized: dict[str, td.Simulation] = {}
  130 
! 131     for idx, sim in enumerate(simulations):
! 132         task_name = Tidy3dStub(simulation=sim).get_default_task_name() + f"_{idx + 1}"
! 133         normalized[task_name] = sim
  134 
! 135     return normalized
  136 
  137 
  138 def has_traced_numerical_structures(
  139     numerical_structures: Union[

---

Lines 151-163

  151     )
  152     for cfg in iterable_structures:
  153         parameters = cfg.parameters
  154         if hasbox(parameters):
! 155             return True
  156         try:
  157             parameters_array = np.asarray(parameters, dtype=object)
! 158         except Exception:
! 159             continue
  160         if hasbox(tuple(parameters_array.flat)):
  161             return True
  162 
  163     return False

---

Lines 202-210

  202             raise AdjointError(
  203                 "Entries in 'numerical_structures' must be NumericalStructureConfig instances."
  204             )
  205         if isinstance(numerical_config.parameters, dict):
! 206             raise AdjointError("NumericalStructureConfig.parameters must not be a dict.")
  207 
  208         _validate_numerical_structure_create_signature(numerical_config.create)
  209         _validate_numerical_structure_vjp_signature(numerical_config.compute_derivatives)

---

Lines 585-593

  585     simulation_static = simulation
  586     if isinstance(simulation, td.Simulation) and (numerical_structures is not None):
  587         # if there are numerical_structures without traced parameters, we still want
  588         # to insert them into the simulation
! 589         simulation_static = insert_numerical_structures_static(
  590             simulation=simulation,
  591             numerical_structures=numerical_structures,
  592         )

---

Lines 725-733

  725         values: Sequence[Any], expected_type: type, arg_name: str, key_name: str
  726     ) -> None:
  727         for idx, value in enumerate(values):
  728             if not isinstance(value, expected_type):
! 729                 raise AdjointError(
  730                     f"{arg_name}[{key_name}][{idx}] must be {expected_type.__name__}, got {type(value)}."
  731                 )
  732 
  733     def _expand_spec(

---

Lines 754-762

  754                 )
  755 
  756         if isinstance(orig_sim_arg, dict):
  757             if fn_arg.keys() != sim_dict.keys():
! 758                 raise AdjointError(f"{arg_name} keys do not match simulations keys")
  759             for key, val in fn_arg.items():
  760                 if isinstance(val, item_type):
  761                     expanded[key] = (val,)
  762                 elif isinstance(val, (list, tuple)):

---

Lines 762-770

  762                 elif isinstance(val, (list, tuple)):
  763                     _validate_sequence_elements(val, item_type, arg_name, key)
  764                     expanded[key] = tuple(val)
  765                 else:
! 766                     raise AdjointError(
  767                         f"{arg_name}[{key}] must be {item_type.__name__} or a sequence of them, got {type(val)}."
  768                     )
  769         else:
  770             if len(fn_arg) != len(orig_sim_arg):

---

Lines 779-787

  779                 elif isinstance(val, (list, tuple)):
  780                     _validate_sequence_elements(val, item_type, arg_name, key)
  781                     expanded[key] = tuple(val)
  782                 else:
! 783                     raise AdjointError(
  784                         f"{arg_name}[{idx}] must be {item_type.__name__} or a sequence of them, got {type(val)}."
  785                     )
  786 
  787         return expanded

---

Lines 867-875

  867         )
  868 
  869     # insert numerical_structures even if not traced
  870     if numerical_structures is not None:
! 871         simulations_static = {
  872             name: (
  873                 insert_numerical_structures_static(
  874                     simulation=simulations_norm[name],
  875                     numerical_structures=numerical_structures[name],

---

Lines 1572-1580

  1572                 task_custom_vjp = custom_vjp.get(task_name)
  1573                 task_numerical_structure_map = numerical_structures.get(task_name, {})
  1574 
  1575                 if isinstance(task_custom_vjp, CustomVJPConfig):
! 1576                     task_custom_vjp = (task_custom_vjp,)
  1577 
  1578                 vjp_results[adj_task_name] = postprocess_adj(
  1579                     sim_data_adj=sim_data_adj,
  1580                     sim_data_orig=sim_data_orig,

---

tidy3d/web/api/autograd/backward.py

Lines 180-188

  180 
  181     def lookup_numerical_structure(structure_index: int) -> NumericalStructureConfig:
  182         numerical_structure = numerical_structure_map.get(structure_index)
  183         if numerical_structure is None:
! 184             raise AdjointError(
  185                 "No NumericalStructureConfig found for numerical structure index "
  186                 f"{structure_index}. Available indices: {sorted(numerical_structure_map.keys())}."
  187             )
  188         return numerical_structure

---

Lines 515-523

  515                     numerical_params_static, derivative_info=derivative_info
  516                 )
  517 
  518                 if not isinstance(gradients, dict):
! 519                     raise AdjointError(
  520                         "Numerical structure VJP function must return a dict mapping paths to gradients."
  521                     )
  522 
  523                 missing_paths = set(numerical_paths_ordered) - set(gradients.keys())

---

Lines 521-529

  521                     )
  522 
  523                 missing_paths = set(numerical_paths_ordered) - set(gradients.keys())
  524                 if missing_paths:
! 525                     raise AdjointError(
  526                         "Numerical structure VJP function did not return gradients for paths: "
  527                         f"{sorted(missing_paths)}."
  528                     )

---

Lines 530-548

  530                 gradient_items = ((path, gradients.get(path)) for path in numerical_paths_ordered)
  531 
  532                 for path, grad_value in gradient_items:
  533                     if grad_value is None:
! 534                         continue
  535                     if path in numerical_value_map:
! 536                         existing = numerical_value_map[path]
! 537                         if isinstance(existing, (list, tuple)) and isinstance(
  538                             grad_value, (list, tuple)
  539                         ):
! 540                             numerical_value_map[path] = type(existing)(
  541                                 x + y for x, y in zip(existing, grad_value)
  542                             )
  543                         else:
! 544                             numerical_value_map[path] = existing + grad_value
  545                     else:
  546                         numerical_value_map[path] = grad_value
  547 
  548         # store vjps in output map

---