Skip to content

add numerical_structures hook to custom run functions#3308

Closed
groberts-flex wants to merge 3 commits intodevelopfrom
groberts-flex/custom_numerical_structures_hook_FXC-3730
Closed

add numerical_structures hook to custom run functions#3308
groberts-flex wants to merge 3 commits intodevelopfrom
groberts-flex/custom_numerical_structures_hook_FXC-3730

Conversation

@groberts-flex
Copy link
Contributor

@groberts-flex groberts-flex commented Feb 18, 2026

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
Adds a new gradient hook that injects user-defined structures into autograd runs and changes how local ComponentModeler runs route gradients; mistakes could yield incorrect gradients or new runtime errors. CI/workflow refactors (LFS/sparse checkout, disabled checks, removed release workflows) also carry operational risk if misconfigured.

Overview
Adds a new numerical_structures hook to run_custom/run_async_custom and ComponentModeler local runs, allowing users to create and append structures from traced parameters and provide a compute_derivatives(parameters, derivative_info) VJP for synthetic ("numerical", ...) paths; this is enforced as local-gradient only and includes signature/parameter validation.

Updates autograd plumbing to insert numerical structures (static insertion when untraced; autograd path when traced), tightens custom_vjp invocation to pass derivative_info by keyword, and adds extensive numerical/finite-difference tests plus a Towncrier fragment documenting the new capability.

Refactors repo/CI metadata for the public flex/public/tidy3d layout: path-scoped .gitattributes (LFS + linguist + changelog merge), CODEOWNERS updates, pre-commit scoping, and major GitHub Actions changes (new/renamed public_tidy3d-* workflows, sparse checkouts + LFS, CI user token usage, new testmon cache promotion job, and removal of older release/docs-sync workflows).

Written by Cursor Bugbot for commit 2c8988d. This will update automatically on new commits. Configure here.

@groberts-flex groberts-flex force-pushed the groberts-flex/custom_numerical_structures_hook_FXC-3730 branch 2 times, most recently from 3b5e20a to 501636a Compare February 19, 2026 00:04
@groberts-flex groberts-flex marked this pull request as ready for review February 19, 2026 00:11
@groberts-flex groberts-flex force-pushed the groberts-flex/custom_numerical_structures_hook_FXC-3730 branch from 501636a to 5185a41 Compare February 19, 2026 00:18
@groberts-flex groberts-flex force-pushed the groberts-flex/custom_numerical_structures_hook_FXC-3730 branch from 5185a41 to 0aadf33 Compare February 19, 2026 17:05
@github-actions
Copy link
Contributor

github-actions bot commented Feb 19, 2026

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 (81.2%): Missing lines 176,191,517,523,532,534-535,538,542
  • tidy3d/web/api/autograd/types.py (100%)

Summary

  • Total: 226 lines
  • Missing: 35 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 172-180

  172 
  173     def lookup_numerical_structure(structure_index: int) -> NumericalStructureConfig:
  174         numerical_structure = numerical_structure_map.get(structure_index)
  175         if numerical_structure is None:
! 176             raise AdjointError(
  177                 "No NumericalStructureConfig found for numerical structure index "
  178                 f"{structure_index}. Available indices: {sorted(numerical_structure_map.keys())}."
  179             )
  180         return numerical_structure

Lines 187-195

  187         structure_paths = tuple(sim_vjp_map.get(structure_index, ()))
  188 
  189         use_numerical_vjp = structure_index in numerical_vjp_map
  190         if structure_paths and use_numerical_vjp:
! 191             raise AdjointError(
  192                 "Invalid autograd field mapping: structure index "
  193                 f"{structure_index} has both 'structures' and 'numerical' traced paths. "
  194                 "A structure index must be handled by exactly one VJP path."
  195             )

Lines 513-521

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

Lines 519-527

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

Lines 528-546

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

@groberts-flex groberts-flex force-pushed the groberts-flex/custom_numerical_structures_hook_FXC-3730 branch from 0aadf33 to 0758cca Compare February 23, 2026 23:53
@groberts-flex groberts-flex force-pushed the groberts-flex/custom_numerical_structures_hook_FXC-3730 branch 2 times, most recently from 733c31c to 461c0da Compare February 24, 2026 00:18
Copy link
Contributor

@marcorudolphflex marcorudolphflex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great overall, thanks!
Got some questions, minor comments and nits.


if (numerical_structures is not None) and not (
isinstance(simulation, td.Simulation)
or isinstance(simulation, get_args(ComponentModelerType))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this check is not done for run_async_custom?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's not currently done in run_async_custom but this is because the currently the typing on run_async_custom says it accepts Simulation objects versus run_custom accepts the broader WorkflowType. Do you think it would be safer to add it in anyways to run_async_custom just in case?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would argue for some gating here as the API origin comes from web.run which has WorkflowType as input type.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sounds good, adding that in! better to catch something like that then have some silent failures!

Copy link

@cursor cursor bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@groberts-flex groberts-flex force-pushed the groberts-flex/custom_numerical_structures_hook_FXC-3730 branch from 3fda6a1 to 2c8988d Compare February 27, 2026 14:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants