add numerical_structures hook to custom run functions#3308
add numerical_structures hook to custom run functions#3308groberts-flex wants to merge 3 commits intodevelopfrom
Conversation
3b5e20a to
501636a
Compare
501636a to
5185a41
Compare
5185a41 to
0aadf33
Compare
tests/test_components/autograd/numerical/test_autograd_cm_custom_vjp_numerical_structures.py
Show resolved
Hide resolved
Diff CoverageDiff: origin/develop...HEAD, staged and unstaged changes
Summary
tidy3d/plugins/smatrix/run.pyLines 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 = NoneLines 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.pyLines 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 FalseLines 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 expandedLines 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.pyLines 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_structureLines 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 |
0aadf33 to
0758cca
Compare
733c31c to
461c0da
Compare
marcorudolphflex
left a comment
There was a problem hiding this comment.
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)) |
There was a problem hiding this comment.
this check is not done for run_async_custom?
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
I would argue for some gating here as the API origin comes from web.run which has WorkflowType as input type.
There was a problem hiding this comment.
sounds good, adding that in! better to catch something like that then have some silent failures!
tests/test_components/autograd/numerical/test_autograd_cm_custom_vjp_numerical_structures.py
Show resolved
Hide resolved
There was a problem hiding this comment.
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.
…for user-defined structure gradients
3fda6a1 to
2c8988d
Compare
The
numerical_structuresfunction is another hook to go alongsidecustom_vjp. Incustom_vjpwe offer the ability to define a vjp function to compute gradients for a structure or set of structures that already exist in the simulation. Withnumerical_structureswe 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 thederivative_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_structureshook torun_custom/run_async_customand ComponentModeler local runs, allowing users to create and append structures from traced parameters and provide acompute_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_vjpinvocation to passderivative_infoby 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/tidy3dlayout: path-scoped.gitattributes(LFS + linguist + changelog merge), CODEOWNERS updates, pre-commit scoping, and major GitHub Actions changes (new/renamedpublic_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.