utils
templates.tvboptim.utils
TVB-Optim Template Utilities
Reusable Python functions for tvboptim Mako templates. Import these in template blocks to avoid code duplication.
Usage in templates
<% from tvbo.templates.tvboptim.utils import ( safe_name, as_list, is_network_observation, parse_loss_function, get_observation_refs ) %>
Functions
| Name | Description |
|---|---|
| adapt_class_reference_for_tvboptim | Translate schema class references to native tvboptim monitor classes. |
| as_list | Convert dict or list to list of values. |
| dist_expr | Render a Distribution as a numpyro dist.* constructor string. |
| format_bounds_array | Render a list of SymPy bound values as a code-level list literal. |
| get_all_hyperparams | Get all hyperparameters including from COMBINED included algorithms. |
| get_all_observations_from_algo | Get all observation names including from COMBINED included algorithms. |
| get_attr | Safe attribute access. |
| get_domain_bounds | Lookup domain bounds from model.parameters or coupling.parameters. |
| get_include_info | Extract algorithm name and argument overrides from AlgorithmInclude. |
| get_mode_layout | Compute the folded scalar-state layout for a (possibly multi-mode) model. |
| get_node_param_overrides | Scan network.nodes for per-node parameter overrides. |
| get_node_state_overrides | Scan network.nodes for per-node initial state overrides. |
| get_observation_dependencies | Observations that obs_name derives from — its source entries that are |
| get_observation_refs | Categorize observations into network vs simulation-derived. |
| get_output_channels | Resolve the sv.record-honoring output channels for the presented result. |
| get_param_info | Extract parameter names, defaults, and shapes from a parameters collection. |
| get_recorded_variable_names | Compute the variable layout recorded by tvboptim’s solver. |
| get_state_bounds | Extract state variable bounds as SymPy expressions. |
| is_external_observation | Check if observation is external (has data_source or network.observations source). |
| is_network_observation | Check if observation is a network observation (static data from network). |
| iter_parameter_values | Yield (name, value) pairs from schema Parameter collections. |
| normalize_coupling_aliases | Collapse duplicate keys that point to the same coupling object. |
| obs_has_all_args | Check if observation has all required arguments satisfied. |
| parameter_value | Return a named Parameter value from a schema collection. |
| parse_exploration | Parse exploration specification from YAML. |
| parse_free_param | Parse a free_parameter entry. |
| parse_list_elements | Split a [a, b, c] list-literal string into top-level element strings, |
| parse_loss_arguments | Parse loss function call arguments. |
| parse_loss_function | Parse optimization loss function specification. |
| pipeline_argument | Return the named pipeline argument object (arguments are keyed by name). |
| pipeline_equation_parameters | Collect equation parameters from all observation pipeline steps. |
| render_analysis_observations | Render the body of the generated compute_analysis_observations() function. |
| render_inference | Render the body of one Bayesian inference (numpyro NUTS/MCMC), 8-space indented. |
| render_jax_default | Render a parameter default as a JAX-ready source literal. |
| render_recorded_observable | Render the body of an exploration observable_fn that records a record: list. |
| resolve_config_access | Dotted state-config path for a <scope>.<param> parameter reference. |
| resolve_coupling_input_map | Map coupling-input names to coupling functions for the tvboptim network dict. |
| resolve_coupling_spec | Resolve every derived field a tvboptim coupling class needs from a Coupling. |
| resolve_optimizer_mode | Map the backend-neutral integration.differentiation.mode onto the native |
| resolve_solver_kwargs | Map the backend-neutral integration.differentiation strategy onto |
| safe_name | Convert name to valid Python identifier (preserves case). |
| time_argument_ms | Resolve a time-valued schema argument to milliseconds. |
| to_numeric | Convert string to numeric if possible. |
| toposort_observations | Dependency-order observations so any that lists another as a source is |
adapt_class_reference_for_tvboptim
templates.tvboptim.utils.adapt_class_reference_for_tvboptim(class_info, obs, dt)Translate schema class references to native tvboptim monitor classes.
Database observation metadata may point at TVB monitor classes because the same schema object is used by the TVB backend. The tvboptim backend should consume the equivalent tvboptim monitor API when one exists.
as_list
templates.tvboptim.utils.as_list(obj)Convert dict or list to list of values.
dist_expr
templates.tvboptim.utils.dist_expr(dist_obj)Render a Distribution as a numpyro dist.* constructor string.
Normal -> dist.Normal(mean, std); Uniform -> dist.Uniform(lo, hi) (from parameters or domain). Reuses the standard Distribution vocabulary (name + parameters/domain) as both prior and likelihood-noise family.
format_bounds_array
templates.tvboptim.utils.format_bounds_array(bounds, format='jax')Render a list of SymPy bound values as a code-level list literal.
Uses the appropriate TVBO code printer so infinity is rendered correctly for any backend (jnp.inf, np.inf, Inf, …).
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| bounds | List | list of sympy expressions (Float / oo / -oo) | required |
| format | str | target backend ('jax', 'numpy', 'julia', …) |
'jax' |
Returns
| Name | Type | Description |
|---|---|---|
| str | String like [-10.0, -jnp.inf] ready for code generation. |
get_all_hyperparams
templates.tvboptim.utils.get_all_hyperparams(algo, algorithms_dict)Get all hyperparameters including from COMBINED included algorithms.
Nested includes are skipped — their hyperparameters are passed directly to the inner algorithm’s run_
get_all_observations_from_algo
templates.tvboptim.utils.get_all_observations_from_algo(algo, algorithms_dict)Get all observation names including from COMBINED included algorithms.
Nested includes are skipped — their observations are computed inside the inner algorithm’s own loop, not the outer one.
get_attr
templates.tvboptim.utils.get_attr(obj, name, default=None)Safe attribute access.
get_domain_bounds
templates.tvboptim.utils.get_domain_bounds(param_name, model, all_couplings)Lookup domain bounds from model.parameters or coupling.parameters.
Returns (lo, hi) tuple, where None means unbounded.
get_include_info
templates.tvboptim.utils.get_include_info(inc)Extract algorithm name and argument overrides from AlgorithmInclude.
Returns (algo_name, {param_name: value}) tuple.
get_mode_layout
templates.tvboptim.utils.get_mode_layout(model)Compute the folded scalar-state layout for a (possibly multi-mode) model.
tvboptim’s solver carries a 2-D state (n_states, n_nodes) and its coupling contracts the node axis with a plain matmul, so it has no place for a third per-node mode axis. A model with number_of_modes > 1 (the Stefanescu-Jirsa ReducedSet models) folds that mode axis into the state axis: each state variable v occupies n_modes contiguous scalar slots v__mode0 .. v__mode{M-1}. The dfun reconstructs the (n_nodes, n_modes) mode-vector for each variable from its slots, evaluates the mode-aware equations (mode_dot/mode_sum), and scatters the per-mode derivatives back into those slots; per-mode coupling falls out of the existing 2-D matmul because each (var, mode) slot couples to the same slot across nodes.
For a single-mode model this is the identity (one slot per variable), so the generated code is byte-for-byte unchanged.
Returns (n_modes, slot_names, var_slots) where slot_names is the flat solver state ordering (grouped by variable, then mode) and var_slots maps each variable name to the slot indices of its modes.
get_node_param_overrides
templates.tvboptim.utils.get_node_param_overrides(
network,
n_nodes,
dyn_param_defaults,
)Scan network.nodes for per-node parameter overrides.
When nodes define parameters that differ from the dynamics defaults, build per-node arrays. Only parameters that differ on at least one node are returned.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| network | Any | Network object with .nodes list | required |
| n_nodes | int | number of nodes | required |
| dyn_param_defaults | Dict[str, float] | dict of param_name -> scalar default from dynamics | required |
Returns
| Name | Type | Description |
|---|---|---|
| Dict[str, List[float]] | dict of param_name -> list of per-node values (length n_nodes) |
get_node_state_overrides
templates.tvboptim.utils.get_node_state_overrides(
network,
n_nodes,
state_names,
default_initial_state,
)Scan network.nodes for per-node initial state overrides.
When nodes define state: {theta: {value: 0.8}} in the YAML, build per-node arrays for state variables that differ across nodes.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| network | Any | Network object with .nodes list | required |
| n_nodes | int | number of nodes | required |
| state_names | List[str] | ordered list of state variable names | required |
| default_initial_state | List[float] | default initial value per state variable | required |
Returns
| Name | Type | Description |
|---|---|---|
| Dict[str, List[float]] | dict of sv_name -> list of per-node values (length n_nodes) |
get_observation_dependencies
templates.tvboptim.utils.get_observation_dependencies(
obs_name,
derived_obs_dict,
all_observations,
)Observations that obs_name derives from — its source entries that are themselves observations (edges in the observation dependency graph).
all_observations is the full observation collection (its membership test filters sources down to observation references, ignoring result/state sources).
get_observation_refs
templates.tvboptim.utils.get_observation_refs(observations_dict)Categorize observations into network vs simulation-derived.
Returns
| Name | Type | Description |
|---|---|---|
| Tuple[Set[str], List[str]] | (network_observation_names, observation_names_with_all_args) |
get_output_channels
templates.tvboptim.utils.get_output_channels(model, experiment=None)Resolve the sv.record-honoring output channels for the presented result.
tvboptim’s solver records ALL states (VARIABLES_OF_INTEREST = states + recorded aux) because the full trajectory is needed for observations and the algorithm warmup. The user-facing SimulationResult should instead present only record=True state channels (+ recorded auxiliaries), matching the tvb backend’s variables_of_interest.
Returns (output_indices, output_names, is_subset) — the indices/names of the kept channels within the full recorded ordering (:func:get_recorded_variable_names), and whether that is a strict subset. For the common all-record model this is the identity (is_subset False), so the template emits the result unsliced. Modes are honored: each v__mode{m} slot inherits v’s record flag.
get_param_info
templates.tvboptim.utils.get_param_info(parameters)Extract parameter names, defaults, and shapes from a parameters collection.
Works for both model.parameters and coupling.parameters.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| parameters | dict | dict-like of Parameter objects | required |
Returns
| Name | Type | Description |
|---|---|---|
| tuple | Tuple[List[str], Dict[str, float], Dict[str, str]] | (param_names, param_defaults, param_shapes) - param_names: list of parameter names - param_defaults: dict of name -> scalar value (for DEFAULT_PARAMS) - param_shapes: dict of name -> shape string (only for params with shape attribute) |
get_recorded_variable_names
templates.tvboptim.utils.get_recorded_variable_names(model, experiment=None)Compute the variable layout recorded by tvboptim’s solver.
The generated dynamics class declares VARIABLES_OF_INTEREST = state_names + recorded_aux where recorded_aux is the union of: * derived variables with record: true, * model.output entries that are derived (auxiliary) variables, and * derived variables referenced as the source of any experiment observation (so observations of auxiliaries work without requiring users to also list them in model.output).
Returns (state_names, recorded_aux, all_var_names) where all_var_names is the runtime ordering on axis 1 of solution.ys / result.data and matches solution.variable_names produced by tvboptim >= 0.2.7.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| model | Any | Dynamics object (with state_variables and derived_variables). | required |
| experiment | Any | Optional SimulationExperiment; observations are scanned when present. | None |
get_state_bounds
templates.tvboptim.utils.get_state_bounds(model)Extract state variable bounds as SymPy expressions.
Uses sympy.oo for unbounded dimensions so that code printers automatically render the correct backend literal (jnp.inf, np.inf, Inf, etc.).
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| model | Any | Dynamics instance with .state_variables |
required |
Returns
| Name | Type | Description |
|---|---|---|
| tuple | Tuple[List, List, bool] | (bounds_lo, bounds_hi, has_finite_bounds) - bounds_lo: list of sympy expressions (Float or -oo) - bounds_hi: list of sympy expressions (Float or oo) - has_finite_bounds: True if any bound is finite |
is_external_observation
templates.tvboptim.utils.is_external_observation(obs)Check if observation is external (has data_source or network.observations source).
is_network_observation
templates.tvboptim.utils.is_network_observation(obs)Check if observation is a network observation (static data from network).
Network observations have source starting with ‘network.observations’ or ‘network.edges’. The slot is multivalued; for raw network observations there is exactly one entry. Accept both scalar and list forms.
iter_parameter_values
templates.tvboptim.utils.iter_parameter_values(parameters)Yield (name, value) pairs from schema Parameter collections.
normalize_coupling_aliases
templates.tvboptim.utils.normalize_coupling_aliases(all_couplings, model=None)Collapse duplicate keys that point to the same coupling object.
tvbo can expose one coupling under multiple names such as the coupling function name, a coupling-input key, or an explicit CouplingInput.source. tvboptim only needs one key per distinct object, so prefer stable, user-meaningful aliases before attempting input-to-coupling mapping.
obs_has_all_args
templates.tvboptim.utils.obs_has_all_args(obs)Check if observation has all required arguments satisfied.
Returns True if all pipeline step arguments either have values or are implicitly satisfied by source.
parameter_value
templates.tvboptim.utils.parameter_value(parameters, name, default=None)Return a named Parameter value from a schema collection.
parse_exploration
templates.tvboptim.utils.parse_exploration(
expl,
all_couplings,
get_pipeline_output_key_fn=None,
)Parse exploration specification from YAML.
Returns dict with: name, label, mode, n_parallel, axes, observable_*
parse_free_param
templates.tvboptim.utils.parse_free_param(
fp,
coupling_keys,
model=None,
all_couplings=None,
)Parse a free_parameter entry.
Handles: str, dotted notation, stringified dict, dict, and Parameter objects.
name, heterogeneous, shape, coupling_key, dynamics_key,
lower_bound, upper_bound
parse_list_elements
templates.tvboptim.utils.parse_list_elements(rhs_str)Split a [a, b, c] list-literal string into top-level element strings, respecting nested brackets/parens (so [f(x, y), g(z)] yields two elements).
parse_loss_arguments
templates.tvboptim.utils.parse_loss_arguments(loss_call)Parse loss function call arguments.
Returns
| Name | Type | Description |
|---|---|---|
| List[Dict] | (parsed_args, obs_refs) where: | |
| Set[str] | - parsed_args: list of dicts with ‘name’, ‘type’, and type-specific keys | |
| Tuple[List[Dict], Set[str]] | - obs_refs: set of observation names referenced |
parse_loss_function
templates.tvboptim.utils.parse_loss_function(opt)Parse optimization loss function specification.
Returns dict with: opt_name, func_name, args, obs_refs, agg_over, agg_type or None if no loss defined.
pipeline_argument
templates.tvboptim.utils.pipeline_argument(pipeline, name)Return the named pipeline argument object (arguments are keyed by name).
pipeline_equation_parameters
templates.tvboptim.utils.pipeline_equation_parameters(pipeline)Collect equation parameters from all observation pipeline steps.
render_analysis_observations
templates.tvboptim.utils.render_analysis_observations(
analysis_obs,
coupling_keys,
solver_class,
transient_time,
t1_default,
dt,
)Render the body of the generated compute_analysis_observations() function.
Analysis observations ANALYZE the solve/loss (Lyapunov spectrum, autodiff and finite-difference gradients) rather than transforming result.data. Each is emitted from its declarative analysis metadata (type + target + wrt + parameters). This lives in the adapter/Python layer — NOT the mako template — so the per-type branching can be deduped/harmonized and reused across backends; the template only interpolates the returned block. Analysis solves use a plain solver_class() (the truncation window is an optimization knob, not part of these diagnostics). Returns a string whose lines are indented for a function body (4 spaces), empty string if there are no analysis observations.
render_inference
templates.tvboptim.utils.render_inference(
inf,
coupling_keys,
external_keys,
derived_names,
network_obs_names,
)Render the body of one Bayesian inference (numpyro NUTS/MCMC), 8-space indented.
Mirrors the tvboptim workflow’s make_model + MCMC(NUTS(...)).run: sample each prior, inject it into the forward config at its resolved path, run the SAME differentiable model_fn, score the observed observable under the likelihood. Config injection uses eqx.tree_at (functionally identical to the reference’s in-place mutation). The observed data comes from the likelihood.source observation — a runtime binding or a loaded network measure — so synthetic ground-truth generation stays out of the schema.
render_jax_default
templates.tvboptim.utils.render_jax_default(value)Render a parameter default as a JAX-ready source literal.
Array-valued constants (mode-coupling matrices, Gaussian-quadrature vectors) must be wrapped in jnp.array(...) so the generated dfun’s arithmetic broadcasts; emitting the bare Python list would make scalar * list raise TypeError at runtime. Scalars render as their full-precision repr literal (str/repr of a float are equivalent in Python 3, so no precision is lost).
render_recorded_observable
templates.tvboptim.utils.render_recorded_observable(
record_names,
derived_names,
network_obs_names,
analysis_names,
)Render the body of an exploration observable_fn that records a record: list.
Each recorded name resolves to compute_all_observations (derived / network / simulated observations) or compute_analysis_observations (the analysis diagnostics — Lyapunov, gradients). The observable returns a Bunch of the named values, which the exploration stacks over the grid into one array per name. Kept in the adapter (not the template) so the same routing serves any backend. Returns the function-body string (8-space indented for def observable_fn(s): inside the exploration function).
resolve_config_access
templates.tvboptim.utils.resolve_config_access(
dotted,
coupling_keys,
external_keys=frozenset(),
)Dotted state-config path for a <scope>.<param> parameter reference.
One addressing grammar, shared by optimization free_parameters, analysis wrt, and inference priors — so “which knob” reads the same everywhere:
<coupling_key>.<param>->coupling.<key>.<param>(prefix ∈ coupling_keys)<event_name>.<param>->external.<name>.<param>(prefix ∈ external_keys)<Dynamics>.<param>or bare<param>->dynamics.<param>(default scope)
external_keys are stimulus/external-input event names (the keys of the network’s external_input dict, e.g. stimulus).
resolve_coupling_input_map
templates.tvboptim.utils.resolve_coupling_input_map(
model,
all_couplings,
coupling_inputs_dict,
)Map coupling-input names to coupling functions for the tvboptim network dict.
tvboptim keys coupling by coupling-input name; the schema keys by function name. Resolution order: (1) explicit CouplingInput.source, (2) same name, (3) a single unmapped function broadcasts to all remaining inputs, (4) equal counts zip positionally. LOCAL inputs (CouplingInput.local=True, e.g. local_coupling) are then dropped from the network mapping — a local term is TVB’s surface/local coupling, zero for the region-based simulations tvboptim supports, so it must not be wired to the long-range connectome (the dfun binds it to 0 via its fallback).
Returns (ci_coupling_map, func_to_first_ci) where ci_coupling_map maps ci_name -> (func_name, coupling_obj) and func_to_first_ci maps func_name to the first ci_name using it (for state-access translation).
resolve_coupling_spec
templates.tvboptim.utils.resolve_coupling_spec(
coupling,
coupling_key,
model,
coupling_inputs_info,
func_to_ci,
n_modes=1,
)Resolve every derived field a tvboptim coupling class needs from a Coupling.
Keeps the cfun mako template emission-only (resolution lives here, per the resolve-in-Python-not-mako convention). Covers: output dimension, incoming/local states (explicit, inferred from the pre-expression, or the model’s coupling_variable states), the mode fold (a multi-mode cvar → its per-node mode slots, one output per mode), pre-expression term parsing (list decomposition + n_pre), vectorized-vs-per-edge selection, class/base names, the differentiable- delay kwarg, state-subscript aliases ({state}_j/_i) and post-recombination symbols, plus the symbol list for the JAX expression printer. Expression rendering (jaxcode) stays in the template.
resolve_optimizer_mode
templates.tvboptim.utils.resolve_optimizer_mode(integration)Map the backend-neutral integration.differentiation.mode onto the native optimizer differentiation mode.
reverse -> "rev" (reverse-mode BPTT; pairs with a grad_horizon window for truncated BPTT); forward -> "fwd" (forward-mode AD, the exact untruncated gradient for a scalar parameter). Defaults to "rev" when no differentiation strategy is declared.
resolve_solver_kwargs
templates.tvboptim.utils.resolve_solver_kwargs(
integration,
dt,
is_diffrax=False,
)Map the backend-neutral integration.differentiation strategy onto native-solver kwargs, returned as a ready-to-emit string (e.g. "grad_horizon=100, block_size=50").
truncation_window / checkpoint_interval are in ms of simulated time; the native JAX solver counts integration steps, so they are converted with dt. Diffrax has no such knobs, so is_diffrax=True yields "". Shared by the experiment and solver templates so the mapping lives in one place rather than being duplicated in both mako blocks.
safe_name
templates.tvboptim.utils.safe_name(name)Convert name to valid Python identifier (preserves case).
Python identifiers are case-sensitive, and result keys must match the user’s YAML keys verbatim so that res.explorations.C_sweep_fig3 works for a YAML entry named C_sweep_fig3. Only characters that are invalid in identifiers (spaces, hyphens) are replaced.
time_argument_ms
templates.tvboptim.utils.time_argument_ms(argument, default)Resolve a time-valued schema argument to milliseconds.
to_numeric
templates.tvboptim.utils.to_numeric(val)Convert string to numeric if possible.
toposort_observations
templates.tvboptim.utils.toposort_observations(
obs_names,
derived_obs_dict,
all_observations,
)Dependency-order observations so any that lists another as a source is emitted AFTER that source — the same dependency-graph principle used for derived variables/parameters (see tvbo.classes.equation). Independent observations keep their input order (stable / deterministic). Lives in the tvboptim adapter so the mako templates only call it rather than redefining the sort inline.