solve

experimental.network_dynamics.solve

Solving system for network architecture.

This module provides the prepare-solve pattern for Network with multi-coupling support. The prepare() function sets up the integration with all coupling state management, and returns a pure function for execution.

Functions

Name	Description
prepare	Compile a model into a pure JAX solve function and a config PyTree.
run_scan	Run the integration scan, dispatching on the solver’s gradient/memory knobs.
solve	Main entry point for simulation.

prepare

experimental.network_dynamics.solve.prepare(
    dynamics,
    solver,
    t0=0.0,
    t1=1.0,
    dt=0.1,
    n_nodes=1,
    noise=None,
    externals=None,
    reduce=None,
)

Compile a model into a pure JAX solve function and a config PyTree.

Builds per-dispatch data (coupling buffers, noise samples, external inputs) and returns (solve_fn, config) where solve_fn(config) runs the integration. Dispatches on the first two arguments via plum: Network/AbstractDynamics paired with NativeSolver/DiffraxSolver.

Parameters

Name	Type	Description	Default
t0	float	Integration interval and step size. `dt` is the fixed step for native solvers and the initial step for Diffrax.	`0.0`
t1	float	Integration interval and step size. `dt` is the fixed step for native solvers and the initial step for Diffrax.	`0.0`
dt	float	Integration interval and step size. `dt` is the fixed step for native solvers and the initial step for Diffrax.	`0.0`

Returns

Name	Type	Description
	(Callable, Bunch)	Pure solve function and its runtime configuration PyTree.
	See `help(prepare)` or `prepare.__doc__` for the full reference,
	including per-dispatch parameters (`n_nodes`, `noise`, `externals`
	for bare dynamics) and Diffrax limitations (no delays, no auxiliaries,
	no VOI filtering).

run_scan

experimental.network_dynamics.solve.run_scan(
    op,
    state0,
    scan_inputs,
    n_steps,
    solver,
    fold=None,
    noise_gen=None,
)

Run the integration scan, dispatching on the solver’s gradient/memory knobs.

The single seam where scan-level features live. Independent, nullable knobs select the path:

noise_gen (streaming noise): noise_gen(block_idx, block_len) or None. Set only when block_size is set and the network has noise with no injected tensor; the per-block noise is generated in-scan from the absolute block ordinal, so scan_inputs carries only the time signal (no noise leaf) and the block step combines them.
fold (the reduce output handler): (acc0, update) or None. When set with block_size, the trajectory is folded block-wise into acc and the final carry exposes acc at index 1; the caller reads it and applies finalize. Requires block_size; with block_size=None the caller folds the stacked trajectory once instead (the degenerate single-block / post-hoc case).
grad_horizon (gradient horizon): if set, run a windowed scan that severs the carry gradient every W steps. Snapped to a multiple of block_size when both are set so window and block boundaries align.
block_size (block granularity): with no truncation, None is the plain single jax.lax.scan (the monolithic default, no-regression path) and an int is the outer block scan that trades recompute for backward memory and (with noise_gen) streams the per-block noise.

op consumes its per-step driving signals (time, and for SDEs the noise slice) from its block inputs and is agnostic to how they were produced.

solve

experimental.network_dynamics.solve.solve(
    model,
    solver,
    t0=0.0,
    t1=100.0,
    dt=0.1,
    **kwargs,
)

Main entry point for simulation.

Accepts either a Network or a bare AbstractDynamics instance. Dispatches to the appropriate prepare() overload via plum.

Args: model: Network or AbstractDynamics instance solver: NativeSolver or DiffraxSolver instance t0: Start time t1: End time (inclusive for native solvers — see note on time grid) dt: Time step **kwargs: Additional arguments forwarded to prepare() (e.g. n_nodes for bare dynamics)

Returns: Simulation results wrapped in result object

Notes: Native solvers use the half-open scan grid arange(t0, t1, dt) and emit the post-step state on each iteration, so the returned save grid is (t0, t1]: result.ts = [t0 + dt, t0 + 2*dt, ..., t1], with the initial state at t0 excluded and the endpoint t1 included. The number of saved samples is (t1 - t0) / dt. t1 - t0 must be an integer multiple of dt for the grid to land exactly on t1.

Examples: >>> # With Network >>> result = solve(network, Euler(), t0=0, t1=10, dt=0.01)

>>> # With bare dynamics (single node)
>>> result = solve(JansenRit(), Heun(), t0=0, t1=1.0, dt=0.001)

>>> # With bare dynamics (multi-node uncoupled)
>>> result = solve(JansenRit(), Heun(), t0=0, t1=1.0, dt=0.001, n_nodes=3)

# solve { #tvboptim.experimental.network_dynamics.solve } `experimental.network_dynamics.solve` Solving system for network architecture. This module provides the prepare-solve pattern for Network with multi-coupling support. The prepare() function sets up the integration with all coupling state management, and returns a pure function for execution. ## Functions | Name | Description | | --- | --- | | [prepare](#tvboptim.experimental.network_dynamics.solve.prepare) | Compile a model into a pure JAX solve function and a config PyTree. | | [run_scan](#tvboptim.experimental.network_dynamics.solve.run_scan) | Run the integration scan, dispatching on the solver's gradient/memory knobs. | | [solve](#tvboptim.experimental.network_dynamics.solve.solve) | Main entry point for simulation. | ### prepare { #tvboptim.experimental.network_dynamics.solve.prepare } ```python experimental.network_dynamics.solve.prepare( dynamics, solver, t0=0.0, t1=1.0, dt=0.1, n_nodes=1, noise=None, externals=None, reduce=None, ) ``` Compile a model into a pure JAX solve function and a config PyTree. Builds per-dispatch data (coupling buffers, noise samples, external inputs) and returns ``(solve_fn, config)`` where ``solve_fn(config)`` runs the integration. Dispatches on the first two arguments via ``plum``: ``Network``/``AbstractDynamics`` paired with ``NativeSolver``/``DiffraxSolver``. #### Parameters {.doc-section .doc-section-parameters} | Name | Type | Description | Default | |--------|------------------|-------------------------------------------------------------------------------------------------------------------|-----------| | t0 | [float](`float`) | Integration interval and step size. ``dt`` is the fixed step for native solvers and the initial step for Diffrax. | `0.0` | | t1 | [float](`float`) | Integration interval and step size. ``dt`` is the fixed step for native solvers and the initial step for Diffrax. | `0.0` | | dt | [float](`float`) | Integration interval and step size. ``dt`` is the fixed step for native solvers and the initial step for Diffrax. | `0.0` | #### Returns {.doc-section .doc-section-returns} | Name | Type | Description | |--------|-----------------------------------------------------------------------------------------------------|-----------------------------------------------------------| | | ([Callable](`typing.Callable`), [Bunch](`tvboptim.experimental.network_dynamics.core.bunch.Bunch`)) | Pure solve function and its runtime configuration PyTree. | | | See ``help(prepare)`` or ``prepare.__doc__`` for the full reference, | | | | including per-dispatch parameters (``n_nodes``, ``noise``, ``externals`` | | | | for bare dynamics) and Diffrax limitations (no delays, no auxiliaries, | | | | no VOI filtering). | | ### run_scan { #tvboptim.experimental.network_dynamics.solve.run_scan } ```python experimental.network_dynamics.solve.run_scan( op, state0, scan_inputs, n_steps, solver, fold=None, noise_gen=None, ) ``` Run the integration scan, dispatching on the solver's gradient/memory knobs. The single seam where scan-level features live. Independent, nullable knobs select the path: - ``noise_gen`` (streaming noise): ``noise_gen(block_idx, block_len)`` or None. Set only when ``block_size`` is set and the network has noise with no injected tensor; the per-block noise is generated in-scan from the absolute block ordinal, so ``scan_inputs`` carries only the time signal (no noise leaf) and the block step combines them. - ``fold`` (the reduce output handler): ``(acc0, update)`` or None. When set with ``block_size``, the trajectory is folded block-wise into ``acc`` and the final carry exposes ``acc`` at index 1; the caller reads it and applies ``finalize``. Requires ``block_size``; with ``block_size=None`` the caller folds the stacked trajectory once instead (the degenerate single-block / post-hoc case). - ``grad_horizon`` (gradient horizon): if set, run a windowed scan that severs the carry gradient every ``W`` steps. Snapped to a multiple of ``block_size`` when both are set so window and block boundaries align. - ``block_size`` (block granularity): with no truncation, None is the plain single ``jax.lax.scan`` (the monolithic default, no-regression path) and an int is the outer block scan that trades recompute for backward memory and (with ``noise_gen``) streams the per-block noise. ``op`` consumes its per-step driving signals (time, and for SDEs the noise slice) from its block inputs and is agnostic to how they were produced. ### solve { #tvboptim.experimental.network_dynamics.solve.solve } ```python experimental.network_dynamics.solve.solve( model, solver, t0=0.0, t1=100.0, dt=0.1, **kwargs, ) ``` Main entry point for simulation. Accepts either a Network or a bare AbstractDynamics instance. Dispatches to the appropriate prepare() overload via plum. Args: model: Network or AbstractDynamics instance solver: NativeSolver or DiffraxSolver instance t0: Start time t1: End time (inclusive for native solvers — see note on time grid) dt: Time step **kwargs: Additional arguments forwarded to prepare() (e.g. n_nodes for bare dynamics) Returns: Simulation results wrapped in result object Notes: Native solvers use the half-open scan grid ``arange(t0, t1, dt)`` and emit the post-step state on each iteration, so the returned save grid is ``(t0, t1]``: ``result.ts = [t0 + dt, t0 + 2*dt, ..., t1]``, with the initial state at ``t0`` excluded and the endpoint ``t1`` included. The number of saved samples is ``(t1 - t0) / dt``. ``t1 - t0`` must be an integer multiple of ``dt`` for the grid to land exactly on ``t1``. Examples: >>> # With Network >>> result = solve(network, Euler(), t0=0, t1=10, dt=0.01) >>> # With bare dynamics (single node) >>> result = solve(JansenRit(), Heun(), t0=0, t1=1.0, dt=0.001) >>> # With bare dynamics (multi-node uncoupled) >>> result = solve(JansenRit(), Heun(), t0=0, t1=1.0, dt=0.001, n_nodes=3)