Code Generation & Export

Generate standalone simulation code in multiple languages and export experiments to standard formats.

Code Generation

View the generated source code for any backend:

from tvbo import Dynamics, SimulationExperiment

model = Dynamics.from_db("Generic2dOscillator")
exp = SimulationExperiment(dynamics=model)

JAX

print(exp.render_code('jax'))


import logging

import jax
from tvbo.data.types import TimeSeries
from tvbo.utils import Bunch
import jax.numpy as jnp

logger = logging.getLogger("tvbo.run")



def cfun(weights, history, current_state, p, delay_indices, t):
    n_node = weights.shape[0]
    b, a = p.b, p.a


    x_j = jnp.array([

    current_state[0, delay_indices[1]],
    
    ])

    V = x_j[0]

    pre = x_j
    pre = pre.reshape(-1, n_node ,n_node)

    def op(x): return jnp.sum(weights * x, axis=-1)
    gx = jax.vmap(op, in_axes=0)(pre)
    return b + a*gx





import jax.numpy as jnp
import jax.scipy as jsp



def dfun(current_state, t, cX, _p):
    # Parameters
    I = _p.I
    a = _p.a
    alpha = _p.alpha
    b = _p.b
    beta = _p.beta
    c = _p.c
    d = _p.d
    e = _p.e
    f = _p.f
    g = _p.g
    gamma = _p.gamma
    tau = _p.tau

    # Coupling
    c_glob = cX[0]
    local_coupling = 0


    # State variables
    V = current_state[0]
    W = current_state[1]


    # Derivatives
    dV_dt = d*tau*(I*gamma + V*g + V*local_coupling + W*alpha + c_glob*gamma + e*V**2 - f*V**3)
    dW_dt = d*(a + V*b + c*V**2 - W*beta)/tau

    derivatives = jnp.array(jnp.broadcast_arrays(dV_dt, dW_dt))

    return derivatives







def integrate(state, weights, dt, params_integrate, delay_indices, external_input):
    """
    Heun Integration
    ================
    """
    t, _ = external_input
    noise = 0


    params_dfun, params_cfun, params_stimulus = params_integrate

    history, current_state = state
    stimulus = 0


    cX = jax.vmap(cfun, in_axes=(None, -1, -1, None, None, None), out_axes=-1)(weights, history, current_state, params_cfun, delay_indices, t)

    dX0 = dfun(current_state, t, cX, params_dfun)

    X = current_state



    # Calculate intermediate step X1
    X1 = X + dX0 * dt + noise + stimulus * dt

    # Calculate derivative X1
    dX1 = dfun(X1, t, cX, params_dfun)
    # Calculate the state change dX
    dX = (dX0 + dX1) * (dt / 2)
    next_state = current_state + (dX)


    return (history, next_state), next_state



import jax
import jax.numpy as jnp

def g(dt, nt, n_svar, n_nodes, n_modes, seed=0, sigma_vec=None, sigma=0.0, state=None):
    """Standard Gaussian white noise using xi ~ N(0,1).

    Returns (nt, n_svar, n_nodes, n_modes): sqrt(dt) * sigma * xi.

    - sigma_vec: optional per-state sigma (length n_svar).
    - sigma: scalar fallback when sigma_vec is None.
    - state: optional current state placeholder for future correlative noise.
    """
    key = jax.random.PRNGKey(int(seed))
    xi = jax.random.normal(key, (nt, n_svar, n_nodes, n_modes))

    if sigma_vec is not None:
        sigma_b = jnp.asarray(sigma_vec)[None, ..., None, None]
    else:
        sigma_b = jnp.asarray(sigma)

    noise = jnp.sqrt(dt) * sigma_b * xi
    return noise






def monitor_raw(time_steps, trace, params, t_offset = 0):
    dt = 0.01220703125
    return TimeSeries(time=(time_steps + t_offset) * dt, data=trace, title = "Raw")


def transform_parameters(_p):
    I, a, alpha, b, beta, c, d, e, f, g, gamma, tau = _p.I, _p.a, _p.alpha, _p.b, _p.beta, _p.c, _p.d, _p.e, _p.f, _p.g, _p.gamma, _p.tau
    
    return _p

c_vars = jnp.array([0]).astype(jnp.int32)

def kernel(state):
    # problem dimensions
    n_nodes = 1
    n_svar = 2
    n_cvar = 1
    n_modes = 1
    nh = 1

    current_state, history = (state.initial_conditions.data[-1], None) ## history = current_state

    ics = (history, current_state)
    weights = state.network.weights_matrix

    dn = jnp.arange(int(n_nodes)) * jnp.ones((int(n_nodes), int(n_nodes))).astype(jnp.int32)
    idelays = jnp.round(state.network.lengths_matrix / state.network.conduction_speed.value / state.dt).astype(jnp.int32) if state.network.conduction_speed.value > 0 else jnp.zeros((int(n_nodes), int(n_nodes)), dtype=jnp.int32)
    di = -1 * idelays - 1
    delay_indices = (di, dn)

    dt = state.dt
    nt = state.nt
    time_steps = jnp.arange(0, nt)

    # Generate batch noise using xi with per-state sigma_vec.
    # Prefer state-provided sigma_vec (supports vmapped sweeps); fallback to experiment-level constants.
    seed = getattr(state.noise, 'seed', 0) if hasattr(state.noise, 'seed') else 0
    try:
        sigma_vec_runtime = getattr(state.noise, 'sigma_vec', None)
    except Exception:
        sigma_vec_runtime = None
    sigma_vec = sigma_vec_runtime if sigma_vec_runtime is not None else jnp.array([0.,0.])
    noise = g(dt, nt, n_svar, n_nodes, n_modes, seed=seed, sigma_vec=sigma_vec)


    p = transform_parameters(state.parameters.dynamics)
    params_integrate = (p, state.parameters.coupling, state.stimulus)

    op = lambda ics, external_input: integrate(ics, weights, dt, params_integrate, delay_indices, external_input)
    latest_carry, res = jax.lax.scan(op, ics, (time_steps, noise))

    trace = res


    



    trace = jnp.hstack((
        trace[:, [0], :],
        trace[:, [1], :],
        ))

    t_offset = 0
    time_steps = time_steps + 1

    
    labels_dimensions = {
        "Time": None,
        "State Variable": ['V', 'W'],
        "Space": ['0'],
        "Mode": ['m0'],
    }
    return TimeSeries(time=(time_steps + t_offset) * dt, data=trace, title = "Raw", sample_period=dt, labels_dimensions=labels_dimensions)


def run_experiment(state):
    """Run simulation, apply observations, and return a fully typed ExperimentResult.

    Mirrors the tvboptim backend: the generated script constructs
    ``SimulationResult`` / ``ExperimentResult`` objects directly, without any
    post-hoc ``from_timeseries`` transformation in the Python caller.
    """
    from tvbo.data.types import SimulationResult, ExperimentResult, _to_dataarray

    ts = kernel(state)

    # ── Integration result ─────────────────────────────────────────────────
    ld = ts.labels_dimensions if isinstance(ts.labels_dimensions, dict) else {}
    integration_da = _to_dataarray(
        ts.data,
        raw_time=ts.time,
        state_names=ld.get("State Variable"),
        nodes=ld.get("Region"),
    )

    # ── Observations ───────────────────────────────────────────────────────
    observations = {}

    integration = SimulationResult(data=integration_da, observations=observations)
    return ExperimentResult(integration=integration, name='')


# ---------------------------------------------------------------------------
# Standalone entry point (executed when this script is run directly).
# Reconstructs the experiment from the frozen YAML spec and invokes kernel().
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import argparse
    from pathlib import Path as _Path
    from tvbo.log import configure_logging

    # Standalone run: progress on stderr, controlled by TVBO_LOG_LEVEL (default INFO).
    configure_logging()

    _parser = argparse.ArgumentParser(description="Run JAX-generated TVBO simulation")
    _parser.add_argument("--spec", type=_Path, default=None,
                         help="YAML experiment spec (default: ../spec/*.yaml next to this script)")
    _parser.add_argument("-o", "--output", type=_Path, default=None,
                         help="Output directory for the result")
    _args = _parser.parse_args()

    _spec = _args.spec
    if _spec is None:
        _candidates = sorted((_Path(__file__).resolve().parent.parent / "spec").glob("*.yaml"))
        if not _candidates:
            raise SystemExit("No spec found; pass --spec PATH")
        _spec = _candidates[0]

    from tvbo import SimulationExperiment
    _experiment = SimulationExperiment.from_yaml(str(_spec))
    _state = _experiment.collect_state()
    _result = kernel(_state)
    logger.info("Done: %s, shape=%s", type(_result).__name__, getattr(_result, "shape", None))

    if _args.output is not None:
        _args.output.mkdir(parents=True, exist_ok=True)
        if hasattr(_result, "save"):
            _result.save(str(_args.output))
        else:
            import numpy as _np
            _np.savez(_args.output / "result.npz", data=getattr(_result, "data", _result))
        logger.info("Wrote results to %s", _args.output)

Python (SciPy)

print(model.render_code('python'))
import numpy as np
import scipy.special


def Generic2dOscillator(
    current_state,
    t,
    I=0.0,
    a=-2.0,
    alpha=1.0,
    b=-10.0,
    beta=1.0,
    c=0.0,
    d=0.02,
    e=3.0,
    f=1.0,
    g=0.0,
    gamma=1.0,
    tau=1.0,
    c_glob=0.0,
    stimulus=False,
):
    stim_t = stimulus(t) if stimulus else 0.0
    local_coupling = 0.0

    # State variables
    V = current_state[0]
    W = current_state[1]

    # Derivatives
    dV_dt = (
        d
        * tau
        * (
            I * gamma
            + V * g
            + V * local_coupling
            + W * alpha
            + c_glob * gamma
            + e * V**2
            - f * V**3
        )
    )
    dW_dt = d * (a + V * b + c * V**2 - W * beta) / tau

    derivatives = np.array(np.broadcast_arrays(dV_dt, dW_dt))

    return derivatives

Julia

print(model.render_code('julia'))





function Generic2dOscillator!(dx, x, p, t = 0)

    (;I, a, alpha, b, beta, c, d, e, f, g, gamma, tau, c_glob, local_coupling) = p

    V, W = x




    dx[1] = d .* tau .* (I .* gamma .+ V .* g .+ V .* local_coupling .+ W .* alpha .+ c_glob .* gamma .+ e .* V .^ 2 .- f .* V .^ 3)
    dx[2] = d .* (a .+ V .* b .+ c .* V .^ 2 .- W .* beta) ./ tau
    dx
end

# Parameter values
p = (I = 0.0, a = -2.0, alpha = 1.0, b = -10.0, beta = 1.0, c = 0.0, d = 0.02, e = 3.0, f = 1.0, g = 0.0, gamma = 1.0, tau = 1.0, c_glob = 0.0, local_coupling = 0.0)


using OrdinaryDiffEqTsit5

# Initial conditions (flat state vector; multi-mode SVs are length-n_modes blocks)
u0 = [
        0.1,
        0.1,
    ]

# Define time span
tspan = (0.0, 1000) # Adjust time span as needed

prob = ODEProblem(Generic2dOscillator!, u0, tspan, p)


# Solve
sol = solve(prob, Tsit5(); saveat=0.01)

All Code Generation Targets

From Dynamics.render_code(format)

Format Output
'tvb' TVB model class
'python' / 'scipy' SciPy ODE function
'autodiff' / 'jax' JAX dfun kernel
'julia' DifferentialEquations.jl
'bifurcation-julia' BifurcationKit.jl
'bifurcation-numcont' MatCont-style
'bifurcation-auto7p' AUTO-07p
'pde-fem' scikit-fem PDE solver

From SimulationExperiment.render_code(format)

Format Output
'jax' / 'autodiff' Full JAX simulation
'tvboptim' tvboptim experiment
'tvb' TVB model class
'networkdynamics' NetworkDynamics.jl
'mtk' ModelingToolkit.jl
'rateml' RateML Python + Numba
'cuda' CUDA kernel
'pde' scikit-fem PDE

Export Formats

YAML

yaml_str = exp.to_yaml()
print(yaml_str[:300])
id: 1
model: Generic2dOscillator
dynamics:
  name: Generic2dOscillator
  iri: tvbo:Generic2dOscillator
  parameters:
    I:
      name: I
      definition: Baseline shift of the cubic nullcline
      value: 0.0
      domain:
        enforce: none
        lo: -5.0
        hi: 5.0
        step: 0.01

BIDS

A run writes its own BIDS derivative: tvbo run puts each result and its YAML sidecar in the study’s derivatives/tvbo/, named by the layout record. There is no separate export call.

OpenMINDS

exp.to_openminds("output/experiment.jsonld")

Model Reports

Generate formatted documentation of a model:

from IPython.display import Markdown
report = model.generate_report(format="markdown")
Markdown(report)

Generic2dOscillator

The Generic 2-Dimensional Oscillator (G2D) is a phenomenological, coupled, nonlinear two-dimensional (i.e., two state-variables (‘V’, ‘W’)) oscillatory, neural mass model. The G2D is a generalization of the well-known FitzHugh-Nagumo model (FitzHugh, 1961; Nagumo et. al, 1962), adapted here for reproducing a wilder class of dynamical configurations of physiological phenomena as observed in neuronal population using phase-portrait method.

type: generic; autonomous: True; modes: 1; state variables: 2; parameters: 12.

State Equations

\[ \dot{V} = d*\tau*\left(I*\gamma + V*g + V*c_{local} + W*\alpha + c_{glob}*\gamma + e*V^{2} - f*V^{3}\right) \] \[ \dot{W} = \frac{d*\left(a + V*b + c*V^{2} - W*\beta\right)}{\tau} \]

State Variables

Variable Initial Value Equation Domain / Sampling Flags Description
\(V\) 0.1 differential (order 1) [-2, 4] coupling, recorded V_G2D is the first state-variable of the Generic 2-Dimensional Oscillator model, that can be associated with membrane potentials and sodium activation (FitzHugh, 1961)
\(W\) 0.1 differential (order 1) [-6, 6] recorded W_G2D is the second state-variable of the Generic 2-Dimensional Oscillator model, that can be associated with accommodation and refractoriness and considered to represent potassium activation, sodium inactivation, or both (FitzHugh, 1961)

Parameters

Parameter Value Domain / Sampling Description
\(I\) 0 [-5, 5], step=0.01 Baseline shift of the cubic nullcline
\(a\) -2 [-5, 5], step=0.01 Vertical shift of the configurable nullcline
\(\alpha\) 1 [-5, 5], step=0.0001 Constant parameter to scale the rate of feedback from the slow variable to the fast variable.
\(b\) -10 [-20, 15], step=0.01 Linear slope of the configurable nullcline
\(\beta\) 1 [-5, 5], step=0.0001 Constant parameter to scale the rate of feedback from the slow variable to itself
\(c\) 0 [-10, 10], step=0.01 Parabolic term of the configurable nullcline
\(d\) 0.02 [0.0001, 1], step=0.0001 Temporal scale factor
\(e\) 3 [-5, 5], step=0.0001 Coefficient of the quadratic term of the cubic nullcline
\(f\) 1 [-5, 5], step=0.0001 Coefficient of the cubic term of the cubic nullcline
\(g\) 0 [-5, 5], step=0.5 Coefficient of the linear term of the cubic nullcline
\(\gamma\) 1 [-1, 1], step=0.1 Constant parameter to reproduce FHN dynamics where excitatory input currents are negative
\(\tau\) 1 [1, 5], step=0.01 A time-scale hierarchy can be introduced for the state variables :math:V and :math:W

Coupling Inputs

c_glob, local_coupling

References

Nagumo, J., Arimoto, S., & Yoshizawa, S. (1962). An active pulse transmission line simulating nerve axon. Proceedings of the IRE, 50(10), 2061-2070.

FitzHugh, R. (1961). Impulses and physiological states in theoretical models of nerve membrane. Biophysical Journal, 1(6), 445-466.

See Also