Prepared Observations for Long Simulations

Returning BOLD and sampled signals without retaining the neural trajectory

Try this notebook interactively:

Download .ipynb Download .qmd Open in Colab

Introduction

Gradient Checkpointing controls activation storage for the backward pass. A prepared observation addresses a separate forward cost: the full neural trajectory that a solver would otherwise return.

If the raw trajectory fits in memory, prefer post-hoc observation: calling a monitor on a stored solution is the most flexible and usually the fastest route, and it allows arbitrary analysis afterwards. Use observe= when retaining the trajectory is the bottleneck, keep the streaming monitor to the reduction that removes it, and do further analysis post-hoc on the reduced output.

Pass a supported monitor through observe= and the native solver computes that signal inside each integration block. A low-rate BOLD result can then be inspected, reused for several losses, or transformed with ordinary functions such as compute_fc.

The public controls have distinct roles:

control	meaning
observe=monitor	choose the scientific time series returned by the solver
block_size=K	choose integration/checkpoint blocks and enable block-wise noise generation
grad_horizon=W	truncate credit assignment every W raw integration steps

Blocking alone still returns and retains the full trajectory. Observation removes that output when the requested signal is smaller. Exact reverse mode still saves checkpoint carries at block boundaries, so total gradient memory is not constant in simulation duration.

NoteScope and grid requirements

• Prepared temporal observations currently require a native solver.
• SubSampling, TemporalAverage, BalloonWindkesselBold, and HRFBold are built in. Experimental custom monitors implement the public preparation contract described below.
• Sampling periods must be integer multiples of the input interval. Every regular block and truncation window must contain complete observation periods.
• Only an incomplete final tail may be dropped. A run shorter than one period returns an empty NativeSolution with the correct shape and metadata.
• HRFBold supports its default temporal averaging, explicit point sampling, and Identity() to use the incoming grid directly; online convolution currently requires convolution_mode="valid".

Custom temporal monitors (experimental)

A custom regular causal monitor prepares one recipe and can then run through both observe= and its post-hoc __call__. A per-sample recurrence subclasses SampledMonitor; a block algorithm implements prepare, usually by subclassing StreamingMonitor; a type you do not own is wrapped in a small StreamingMonitor subclass. The supported authoring imports are:

from tvboptim.observations import (
    SampledMonitor,
    StreamingMonitor,
    SimulationGrid,
    ObservationOutput,
    PreparedObservation,
    prepare_observation,
    apply_observation,
    sampling_stride,
)

Preparation receives the simulation grid, an abstract [channels, nodes] input sample, and optional channel names. It returns live parameters, an init(params) callback, a block update, and output metadata. The update has this signature:

next_state, samples = update(state, block, params)

Here block is [input_samples, channels, nodes], while samples must be [input_samples // stride, output_channels, nodes]. Regular block boundaries must align with the declared cadence; an incomplete final tail is dropped. The external LowPass implementation in docs/advanced/low_pass_monitor.py demonstrates a SampledMonitor with state, a live parameter, and dtype promotion without importing solver internals.

The complete Authoring Temporal Observations tutorial builds that monitor, differentiates its live parameter, composes named outputs, and documents the callback contract and current publication limits.

Named observations with shared preprocessing

JointObservation applies one optional shared monitor once, then feeds its single output to a flat mapping of named monitors. Each result is its own NativeSolution, so branches can use different periods and retain different numbers of samples:

from tvboptim.observations import Identity, JointObservation
from tvboptim.observations.tvb_monitors import HRFBold, TemporalAverage

observations = JointObservation(
    preprocess=TemporalAverage(period=4.0, voi=0),
    outputs={
        "neural": Identity(),
        "bold": HRFBold(period=1000.0, downsample=Identity()),
    },
)

result = solve(
    network,
    Heun(block_size=2000),
    t0=0.0,
    t1=60_000.0,
    dt=1.0,
    observe=observations,
)
neural = result["neural"]  # dt == 4 ms
bold = result["bold"]       # dt == 1000 ms

Omitting preprocess passes the raw monitor input to every branch. Identity() preserves the incoming values, interval, timestamp phase, names, and dtype. It also explicitly bypasses the default internal averaging in HRFBold; the HRF kernel and warm history are then resolved on the actual incoming interval.

The composite block alignment is the least common multiple of the shared and branch requirements in raw integration steps. For independent 6-step and 10-step branches, for example, block_size must be a multiple of 30. An incomplete final tail is handled independently for every branch, without padding coarse outputs.

After prepare, live parameters use explicit namespaces:

config.monitor.preprocess
config.monitor.outputs.neural
config.monitor.outputs.bold

The shared and branch states initialize once per model call. Gradients from several branch losses add through the shared operation. Nested joints and mapping-valued preprocessors or branches are outside this bounded composition.

import gc
import time
import threading

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np

try:
    import psutil
    _HAS_PSUTIL = True
except ImportError:
    _HAS_PSUTIL = False

jax.config.update("jax_enable_x64", True)

from tvboptim.data import load_structural_connectivity
from tvboptim.experimental.network_dynamics import Network, solve
from tvboptim.experimental.network_dynamics.coupling import DelayedLinearCoupling
from tvboptim.experimental.network_dynamics.dynamics.tvb import ReducedWongWang
from tvboptim.experimental.network_dynamics.graph import DenseDelayGraph
from tvboptim.experimental.network_dynamics.noise import AdditiveNoise
from tvboptim.experimental.network_dynamics.solvers import Heun
from tvboptim.observations import compute_fc
from tvboptim.observations.tvb_monitors import HRFBold
from tvboptim.utils import cache, set_cache_path

set_cache_path("./prepared_observations")

Workload: delayed RWW observed as BOLD

The workload uses 68 delayed Reduced Wong-Wang nodes and a two-minute rollout. Both compared functions use the same blocked solver, so stochastic noise is generated from the same block keys. One function returns the neural trajectory and applies HRF BOLD afterward. The other returns BOLD directly.

DT = 1.0
T1 = 120_000.0
N_STEPS = int(T1 / DT)
BLOCK_SIZE = 2000

weights, lengths, region_labels = load_structural_connectivity(name="dk_average")
weights = weights / np.max(weights)
graph = DenseDelayGraph(
    weights=jnp.asarray(weights),
    delays=jnp.asarray(lengths / 4.0),
    region_labels=region_labels,
)
network = Network(
    dynamics=ReducedWongWang(w=0.5, I_o=0.32, INITIAL_STATE=(0.3,)),
    coupling={
        "delayed": DelayedLinearCoupling(
            source="S", G=0.5, buffer_strategy="roll"
        )
    },
    graph=graph,
    noise=AdditiveNoise(sigma=0.00283, apply_to="S", key=jax.random.key(0)),
)

# The default 4 ms TemporalAverage is preserved as part of HRFBold's
# scientific definition. The 1000 ms TR divides BLOCK_SIZE * DT exactly.
monitor = HRFBold(period=1000.0, voi=0)
solver = Heun(block_size=BLOCK_SIZE)
n_nodes = weights.shape[0]

raw_output_mb = N_STEPS * n_nodes * 8 / 1e6
bold_output_mb = int(T1 / monitor.period) * n_nodes * 8 / 1e6
print(f"raw output ~{raw_output_mb:.1f} MB")
print(f"BOLD output ~{bold_output_mb:.3f} MB")

Forward-memory comparison

RSSPeakMonitor records the peak process RSS above its entry baseline. It is a practical CPU measure; device runtimes should use their allocator memory tools.

class RSSPeakMonitor:
    def __init__(self, sample_interval=0.02):
        self.sample_interval = sample_interval
        self.peak_delta_bytes = None

    def __enter__(self):
        if not _HAS_PSUTIL:
            return self
        self._process = psutil.Process()
        self._baseline = self._process.memory_info().rss
        self._peak = self._baseline
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._sample, daemon=True)
        self._thread.start()
        return self

    def __exit__(self, *exc):
        if not _HAS_PSUTIL:
            return False
        self._stop.set()
        self._thread.join()
        self.peak_delta_bytes = max(0, self._peak - self._baseline)
        return False

    def _sample(self):
        while not self._stop.is_set():
            try:
                self._peak = max(self._peak, self._process.memory_info().rss)
            except Exception:
                break
            self._stop.wait(self.sample_interval)


def posthoc():
    raw = solve(network, solver, t0=0.0, t1=T1, dt=DT)
    return monitor(raw)


def prepared():
    return solve(
        network,
        solver,
        t0=0.0,
        t1=T1,
        dt=DT,
        observe=monitor,
    )


def measure(fn, n_rep=3):
    first = fn()
    jax.block_until_ready(first.ys)
    gc.collect()
    rss = RSSPeakMonitor()
    with rss:
        out = fn()
        jax.block_until_ready(out.ys)
    best = float("inf")
    for _ in range(n_rep):
        start = time.perf_counter()
        jax.block_until_ready(fn().ys)
        best = min(best, time.perf_counter() - start)
    return {
        "peak_mb": rss.peak_delta_bytes / 1e6 if rss.peak_delta_bytes else None,
        "best_s": best,
        "result": out,
    }


@cache("prepared_observation_comparison", redo=True)
def run_comparison():
    return {"post-hoc": measure(posthoc), "prepared": measure(prepared)}


results = run_comparison()

names = ["post-hoc", "prepared"]
mems = [results[name]["peak_mb"] for name in names]
if all(value is not None for value in mems):
    plt.bar(names, mems, color=["darkorange", "seagreen"])
    plt.ylabel("peak RSS delta (MB)")
    plt.title("Forward memory")
    plt.grid(alpha=0.3, axis="y")
    plt.show()
else:
    print("Peak memory unavailable because psutil is not installed")

posthoc_bold = results["post-hoc"]["result"]
prepared_bold = results["prepared"]["result"]
np.testing.assert_allclose(posthoc_bold.ys, prepared_bold.ys, rtol=1e-5, atol=1e-8)
np.testing.assert_allclose(posthoc_bold.ts, prepared_bold.ts)
print(f"BOLD shape: {prepared_bold.ys.shape}")
print({name: results[name]["best_s"] for name in names})

Figure 1: Peak process RSS for blocked raw output followed by BOLD and prepared BOLD computed inside the same integration blocks.

BOLD shape: (120, 1, 84)
{'post-hoc': 1.504632516996935, 'prepared': 1.452867159008747}

Reuse the returned signal

The observation returns a regular NativeSolution. Compute statistics after the solve so the same scientific signal can support several analyses:

bold = prepared_bold
fc = compute_fc(bold, skip_t=20)
mean_amplitude = bold.ys.mean(axis=0)
print(bold.dt, bold.variable_names, fc.shape, mean_amplitude.shape)

1000.0 ('BOLD(S)',) (84, 84) (1, 84)

For optimization, prepare(..., observe=monitor) places live monitor parameters under config.monitor. Network, readout, and monitor leaves remain ordinary JAX inputs. Partition trainable Parameter leaves before calling jax.grad, as in the optimization workflows.

Forward memory consists roughly of resident inputs and prepared data, the simulation and monitor carry, one integration block’s workspace, and the sampled result. Exact reverse mode additionally retains a checkpoint carry per block and one block’s backward workspace. Delayed-coupling histories and HRF convolution history are part of those carries. Larger blocks therefore trade workspace for fewer saved boundaries; measure the actual workload rather than assuming one universal optimum.

Heterogeneous networks

A heterogeneous simulation needs an explicit common readout before its temporal monitor:

from tvboptim.experimental.network_dynamics import GroupObservation

common = GroupObservation(
    {"cortex": "S", "relay": relay_readout},
    channels=("activity",),
    monitor=HRFBold(period=1000.0, voi=0),
)
bold = solve(
    heterogeneous_network,
    Heun(block_size=2000),
    t0=0.0,
    t1=120_000.0,
    dt=1.0,
    observe=common,
)

The group readout runs first, followed by its temporal monitor. The older observe=(group_observation, monitor) tuple remains supported. Full node coverage is required by default. If allow_partial_coverage=True, the configured fill value becomes real monitor input and later statistics are not automatically fill-aware.

Reducer migration

The temporary reduce= keyword, welford_cov, and streaming_hrf_bold are deprecated in 0.5.0 and scheduled for removal in 0.6.0. Pass BOLD monitors through observe= and compute FC or covariance from the returned signal. A custom terminal online statistic should remain bespoke; observe= always means that the solver returns a scientific time series.