neural_mass

class JRTheta(A, B, a, b, v0, nu_max, r, J, a_1, a_2, a_3, a_4, mu, I)

Bases: tuple

A

Alias for field number 0

B

Alias for field number 1

I

Alias for field number 13

J

Alias for field number 7

a

Alias for field number 2

a_1

Alias for field number 8

a_2

Alias for field number 9

a_3

Alias for field number 10

a_4

Alias for field number 11

b

Alias for field number 3

mu

Alias for field number 12

nu_max

Alias for field number 5

r

Alias for field number 6

v0

Alias for field number 4

class JRState(y0, y1, y2, y3, y4, y5)

Bases: tuple

y0

Alias for field number 0

y1

Alias for field number 1

y2

Alias for field number 2

y3

Alias for field number 3

y4

Alias for field number 4

y5

Alias for field number 5

jr_dfun(ys, c, p)

class BVEPTheta(tau0, I1, x0)

Bases: tuple

I1

Alias for field number 1

tau0

Alias for field number 0

x0

Alias for field number 2

bvep_dfun(ys, c, p: BVEPTheta)

class MPRTheta(tau, I, Delta, J, eta, cr, cv)

Bases: tuple

Delta

Alias for field number 2

I

Alias for field number 1

J

Alias for field number 3

cr

Alias for field number 5

cv

Alias for field number 6

eta

Alias for field number 4

tau

Alias for field number 0

class MPRState(r, V)

Bases: tuple

V

Alias for field number 1

r

Alias for field number 0

mpr_dfun(ys, c, p)

mpr_r_positive(rv, _)

class BOLDTheta(tau_s, tau_f, tau_o, alpha, te, v0, e0, epsilon, nu_0, r_0, recip_tau_s, recip_tau_f, recip_tau_o, recip_alpha, recip_e0, k1, k2, k3)

Bases: tuple

alpha

Alias for field number 3

e0

Alias for field number 6

epsilon

Alias for field number 7

k1

Alias for field number 15

k2

Alias for field number 16

k3

Alias for field number 17

nu_0

Alias for field number 8

r_0

Alias for field number 9

recip_alpha

Alias for field number 13

recip_e0

Alias for field number 14

recip_tau_f

Alias for field number 11

recip_tau_o

Alias for field number 12

recip_tau_s

Alias for field number 10

tau_f

Alias for field number 1

tau_o

Alias for field number 2

tau_s

Alias for field number 0

te

Alias for field number 4

v0

Alias for field number 5

compute_bold_theta(tau_s=0.65, tau_f=0.41, tau_o=0.98, alpha=0.32, te=0.04, v0=4.0, e0=0.4, epsilon=0.5, nu_0=40.3, r_0=25.0)

bold_dfun(sfvq, x, p: BOLDTheta)

class DCMTheta(A, B, C)

Bases: tuple

A

Alias for field number 0

B

Alias for field number 1

C

Alias for field number 2

dcm_dfun(x, u, p: DCMTheta)

Implements the classical bilinear DCM dot{x} = (A + sum_j u_j B_j ) x + C u

DopaTheta

alias of dopaTheta

class DopaState(r, V, u, Sa, Sg, Dp)

Bases: tuple

Dp

Alias for field number 5

Sa

Alias for field number 3

Sg

Alias for field number 4

V

Alias for field number 1

r

Alias for field number 0

u

Alias for field number 2

dopa_dfun(y, cy, p: dopaTheta)

Adaptive QIF model with dopamine modulation.

dopa_net_dfun(y, p)

Canonical form for network of dopa nodes.

dopa_r_positive(y, _)

dopa_gfun_mulr(y, p)

Provide a multiplicative r, additive V gfun.

dopa_gfun_add(y, p)

Provides an additive noise gfun.

monitor

make_offline(step_fn, sample_fn, *args)

Compute monitor samples in an offline or batch fashion.

make_timeavg(shape)

Make a time average monitor.

compute_sarvas_gain(q, r, o, att, Ds=0, Dc=0) → Array

make_gain(gain, shape=None)

Make a gain-matrix monitor suitable for sEEG, EEG & MEG.

make_bold(shape, dt, p: BOLDTheta)

Make a BOLD fMRI monitor.

make_cov(shape)

make_fc(shape)

make_fft(shape, period)

loops

Functions for building time stepping loops.

euler_step(x, dfun, dt, *args, add=0, adhoc=None)

Use a Euler scheme to step state with a right hand sides dfun(.) and additional forcing term add.

heun_step(x, dfun, dt, *args, add=0, adhoc=None, return_euler=False)

Use a Heun scheme to step state with a right hand sides dfun(.) and additional forcing term add.

rk4_step(x, dfun, dt, *args, add=0, adhoc=None)

Use a Runge-Kutta 4 scheme to step state with a right hand sides dfun(.) and additional forcing term add.

make_sde(dt, dfun, gfun, adhoc=None, return_euler=False, unroll=10)

Use a stochastic Heun scheme to integrate autonomous stochastic differential equations (SDEs).

Parameters:

dtfloat

Time step

dfunfunction

Function of the form dfun(x, p) that computes drift coefficients of the stochastic differential equation.

gfunfunction or float

Function of the form gfun(x, p) that computes diffusion coefficients of the stochastic differential equation. If a numerical value is provided, this is used as a constant diffusion coefficient for additive linear SDE.

adhocfunction or None

Function of the form f(x, p) that allows making adhoc corrections to states after a step.

return_euler: bool, default False

Return solution with local Euler estimates.

unroll: int, default 10

Force unrolls the time stepping loop.

Returns:

stepfunction

Function of the form step(x, z_t, p) that takes one step in time according to the Heun scheme.

loopfunction

Function of the form loop(x0, zs, p) that iteratively calls step for all z.

Notes

In both cases, a Jax compatible parameter set p is provided, either an array or some pytree compatible structure.

Note that the integrator does not sample normally distributed noise, so this must be provided by the user.

>>> import vbjax as vb
>>> _, sde = vb.make_sde(1.0, lambda x, p: -x, 0.1)
>>> sde(1.0, vb.randn(4), None)
Array([ 0.5093468 ,  0.30794007,  0.07600437, -0.03876263], dtype=float32)

make_ode(dt, dfun, adhoc=None, method='heun')

Use a Heun scheme to integrate autonomous ordinary differential equations (ODEs).

Parameters:

dtfloat

Time step

dfunfunction

Function of the form dfun(x, p) that computes derivatives of the ordinary differential equations.

adhocfunction or None

Function of the form f(x, p) that allows making adhoc corrections to states after a step.

Returns:

stepfunction

Function of the form step(x, t, p) that takes one step in time according to the Heun scheme.

loopfunction

Function of the form loop(x0, ts, p) that iteratively calls step for all time steps ts.

Notes

In both cases, a Jax compatible parameter set p is provided, either an array or some pytree compatible structure.

>>> import vbjax as vb, jax.numpy as np
>>> _, ode = vb.make_ode(1.0, lambda x, p: -x)
>>> ode(1.0, np.r_[:4], None)
Array([0.5   , 0.25  , 0.125 , 0.0625], dtype=float32, weak_type=True)

make_dde(dt, nh, dfun, unroll=10, adhoc=None)

Invokes make_sdde w/ gfun 0.

make_sdde(dt, nh, dfun, gfun, unroll=1, zero_delays=False, adhoc=None)

Use a stochastic Heun scheme to integrate autonomous stochastic delay differential equations (SDEs).

Parameters:

dtfloat

Time step

nhint

Maximum delay in time steps.

dfunfunction

Function of the form dfun(xt, x, t, p) that computes drift coefficients of the stochastic differential equation.

gfunfunction or float

Function of the form gfun(x, p) that computes diffusion coefficients of the stochastic differential equation. If a numerical value is provided, this is used as a constant diffusion coefficient for additive linear SDE.

adhocfunction or None

Function of the form f(x,p) that allows making adhoc corrections after each step.

Returns:

stepfunction

Function of the form step((x_t,t), z_t, p) that takes one step in time according to the Heun scheme.

loopfunction

Function of the form loop((xs, t), p) that iteratively calls step for each xs[nh:] and starting time index t.

Notes

• A Jax compatible parameter set p is provided, either an array or some pytree compatible structure.

• The integrator does not sample normally distributed noise, so this must be provided by the user.

• The history buffer passed to the user functions, on the corrector stage of the Heun method, does not contain the predictor stage, for performance reasons, unless zero_delays is set to True. A good compromise can be to set all zero delays to dt.

>>> import vbjax as vb, jax.numpy as np
>>> _, sdde = vb.make_sdde(1.0, 2, lambda xt, x, t, p: -xt[t-2], 0.0)
>>> x,t = sdde(np.ones(6)+10, None)
>>> x
Array([ 11.,  11.,  11.,   0., -11., -22.], dtype=float32)

make_continuation(run_chunk, chunk_len, max_lag, n_from, n_svar, stochastic=True)

Helper function to lower memory usage for longer simulations with time delays. WIP

Takes a function

run_chunk(buf, params) -> (buf, chunk_states)

and returns another

continue_chunk(buf, params, rng_key) -> (buf, chunk_states)

The continue_chunk function wraps run_chunk and manages moving the latest states to the first part of buf and filling the rest with samples from N(0,1) if required.

connectome

Utilities for connectomes.

make_conn_latent_mvnorm(SCs, nc=10, return_full=False)

Make a latent multivariate normal distribution over connectomes.

Parameters:

SCs(nconn, n, n) array

Array of connectomes in a given parcellation.

ncint, optional

Number of components to use for the latent space.

return_full: bool, optional

Whether or not to return extra information on the SVD.

Returns:

u_mean(nc, ), array_like

Mean of distribution in latent space.

u_cov(nc, nc), array_like

Covariance of distribution in latent space.

xfmfunction

Maps latent vector to full connectome.

u_cov(nc, nc) array

Covariance of the multivariate normal. Returned if return_full=True.

u(nconn, nconn) array

Left singular vectors corresponding to connectomes embedded. Returned if return_full=True.

s(nconn) array

Singular values. Returned if return_full=True.

vt(nconn, n*n) array

Right singular vectors. Returned if return_full=True.

nconfint

Number of confusions induced by dimensionality reduction. Returned if return_full=True.

crosscoder

CrossCoder for amortized inference of whole-brain connectomes.

A single CrossCoder maps multiple views (e.g. different parcellations or imaging modalities) of a connectome cohort into a shared low-dimensional latent space via strictly linear encoders and decoders. The variational mode places a Gaussian over the latent and regularises towards a standard normal, which helps when training data pools heterogeneous cohorts.

class TrainedArch(nlat: int, wbs: Any, history: dict = <factory>, variational: bool = False)

Bases: object

One trained cross-coder architecture.

nlat: int

wbs: Any

history: dict

variational: bool = False

triu_to_mat(triu: Any) → Any

Fold flat upper-triangular vectors into symmetric square matrices.

triu_to_mat_np(triu: Any) → Any

NumPy equivalent of triu_to_mat() for post-training paths.

class MvNorm(us: Any, mean: Any, cov: Any, key: Any | None = None)

Bases: object

Multivariate normal with persistent PRNG key for sampling.

sample(n: int) → Any

class CrossCoder(variational: bool = True, chunked_training: bool = True)

Bases: object

Multi-view linear auto-encoder over flat upper-triangular connectomes.

Parameters:

variationalbool

If True, the encoder emits a Gaussian over the latent rather than a point estimate, and training minimises MSE + β·KL.

chunked_trainingbool

If True, each train call compiles inner steps into a lax.scan and returns to Python only for logging. Disable for step-wise debug.

archs: List[TrainedArch]

to_pkl(fname: str) → None

classmethod from_pkl(fname: str) → CrossCoder

classmethod from_numpy_array(weights: ndarray, tts: int | None = None, parc: str = 'Schaefer-17Networks', variational: bool = False, chunked_training: bool = True, normalize: str = 'center') → CrossCoder

Build a single-view CrossCoder from an (ns, nn, nn) connectome stack.

add_view(data: Any, parc_name: str, normalize: str = 'zscore', nonneg: bool = False) → None

Register a view, normalizing its flat upper-tri connectomes.

shuffle(seed: int | None = None) → ndarray

Shuffle all views with a common permutation and return it.

make_wbs(nlat: int, key: Any | None = None) → list

Initialize weights/biases for all views at given latent size.

make_loss() → tuple

Build the cross-prediction loss and its gradient.

train(nlat: int, lr: float = 0.0003, niter: int = 2000, tts: int | None = None, mb: int = 64, beta_start: float = 0.0, beta_end: float = 0.001, anneal_steps: int = 1500, key: Any | None = None) → tuple

Fit a single architecture. Appends the learned weights and trace into self.wbs / self.history and returns (trace, wbs, confusion_rate).

property arch

Latent sizes for each trained architecture.

calc_confusion_rate(arch: int, tts: int | None = None, self_recon_only: bool = True, n_samples: int = 0) → float

Fraction of test subjects not fingerprinted correctly.

confusion_matrix(arch: int, tts: int | None = None, n_samples: int = 0) → ndarray

Return (n_views, n_views) confusion rate matrix.

Entry (i, j) is the fraction of test subjects from view i whose reconstruction in view j is closest to the wrong subject.

Parameters:

archint

Latent dimension (nlat value)

ttsint, optional

Train/test split index

n_samplesint

Number of posterior samples for variational mode. 0 = point estimate.

Returns:

numpy.ndarray of shape (n_views, n_views)

encode(arch: int, parc: str, tts: int | None = None, sample: bool = False, key: Any | None = None) → Any

Encode the normalized connectomes of a view into latent space.

encode_all(arch: int, tts: int | None = None, sample: bool = False, key: Any | None = None) → dict[str, Any]

Encode all views into latent space.

Returns:

dict[str, Array]

Mapping of parcellation name to latent vectors of shape (n_subjects, nlat).

decode(arch: int, parc: str, z: Any, raw: bool = False) → Any

Decode latent vectors into flat upper-tri connectomes.

decode_conn(arch: int, parc: str, z: Any, clip_positive: bool | None = None) → Any

Decode latents into full symmetric connectomes (ns, nn, nn).

get_triu(parc: str, tts: int | None = None) → Any

Return normalized flat upper-tri connectomes for a view.

get_conn(parc: str, tts: int | None = None) → Any

Return empirical connectomes (ns, nn, nn) for a view.

calc_mvn(arch: int, tts: int | None = None) → MvNorm

Total-variance multivariate normal over the cohort latents.

decompose_latent(arch: int, tts: int | None = None) → dict

SVD of the centered cohort latents for a given architecture.

classmethod combine(cc1: CrossCoder, cc2: CrossCoder, shuffle: bool = True) → CrossCoder

Concatenate two CrossCoders with identical views and normalizations.

sweep_crosscoder(model: CrossCoder, dims: list[int], n_trials: int = 20, seed: int = 42, keep_best: bool = False, lr_range: tuple[float, float] = (1e-05, 0.01), niter_range: tuple[int, int] = (500, 5000), mb_choices: tuple[int, ...] = (32, 64, 128), beta_end_range: tuple[float, float] = (1e-06, 0.001), anneal_range: tuple[int, int] = (500, 3000), score_fn: Any | None = None) → tuple[list[dict], dict | None]

Random hyperparameter sweep over CrossCoder.train.

Trained weights are discarded after each trial; only summary statistics are kept. When keep_best is True the best-scoring trial’s weights are re-attached to the model after the sweep. Returns (results_sorted_by_score, best).

Visualisation helpers for CrossCoder models.

Matplotlib is imported lazily so the core package has no hard dependency. All plotting routines expect a trained CrossCoder with at least one architecture in model.wbs.

plot_training(results, ncols=3, figsize=None)

Plot train/test traces for a list of training results.

plot_identifiability(model, arch, parc, n_subs=50, cmap='inferno', ax=None)

Pairwise empirical/reconstructed distances on the test set.

plot_fidelity(model, arch, parc, cmap='inferno', gridsize=50, mask_zeros=True, hide_masked=True, ax=None)

Hexbin of empirical vs reconstructed edge weights with regression line.

plot_obs_vs_pred(model, arch, parc=None, n_subs=3, cmap='inferno', cmap_res='PuOr', subject_ids=None, start_idx=None)

Observed/predicted/residual panels for a few random test subjects.

plot_latent(model, arch, parc, color_vals=None, cohorts=None, method='pca', n_components=2, tts=None, dims=(0, 1), per_cohort_color=True, pca_params=None, umap_params=None, cohort_markers=None, fallback_markers=None, ax=None, **scatter_kw)

2-D latent scatter via PCA or UMAP, optionally colored by a covariate.

plot_generative(model, arch, parc, n_samples=500, alpha=0.85, c_emp='slategray', c_gen='tan')

Compare empirical and model-sampled edge-weight summaries.

plot_traversal(model, arch, parc, dim_idx=0, range_sd=3.0, n_steps=7)

Sweep a single latent dimension and visualise the decoded connectome.

layers

make_dense_layers(in_dim, latent_dims=[10], out_dim=None, init_scl=0.1, extra_in=0, act_fn=<PjitFunction of <function leaky_relu>>, key=Array([ 0, 42], dtype=uint32))

Make a dense neural network with the given latent layer sizes.

sparse

make_spmv(A, is_symmetric=False, use_scipy=False)

Make a closure for a general sparse matrix-vector multiplication.

Parameters:

Ascipy.sparse.csr_matrix

Constant sparse matrix.

is_symmetricbool, optional, default False

Whether matrix is symmetric.

use_scipy: bool, optional, default False

Use scipy.

Returns:

spmvfunction

Function implementing spase matrix vector multiply with support for gradients in Jax.

csr_to_jax_bcoo(A: csr_matrix)

Convert CSR format to batched COO format.

make_sg_spmv(A: csr_matrix, use_pmap=False, sharding: Sharding = None)

Make a SpMV kernel w/ generic scatter-gather operations.

util

to_np(x: Array) → ndarray

to_jax(x: ndarray)

Move NumPy array to JAX via DLPack.

tuple_meshgrid(tup)

Applies meshgrid to arrays in a named tuple.

tuple_ravel(tup)

Flatten arrays in fields of tuple.

tuple_shard(tup, n)

Shard arrays in fields of tuple.

shtlc

SHT based local coupling functions.

sph_harm(m, n, theta, phi)

make_grid_shtns(lmax, nlat, nlon, D)

Create shtns object and grid as in make_grid.

make_lm(lmax: int)

make_grid(nlat, nlon)

Create grid for SHT, phi latitude, theta longitude.

grid_pairwise_distance(theta, phi)

Compute pairwise distances on grid. Memory intensive for large grids.

randn_relaxed(sht)

For shtns sht return random spatial array captured by sht.

kernel_diff(D, l)

Compute diffusion kernel, l in shtns order.

kernel_dist_origin(theta, phi)

Compute distance to origin on grid.

kernel_sh_normalized(sht, k)

kernel_laplace(sht, theta, phi, size)

Compute spatial & spectral coefficients for Laplacian kernel.

kernel_gaussian(sht, theta, phi, size)

Compute spatial & spectral coefficients for Gaussian kernel.

kernel_mexican_hat(sht, theta, phi, size)

Compute spatial & spectral coefficients for Mexican hat kernel.

kernel_conv_prep(sht, k)

Prepares evaluated kernel for SHT convolution.

kernel_estimate_shtns(sht, k, x0)

Estimate effective kernel for kernel k with state x0 for shtns object sht.

make_shtdiff_np(lmax, nlat, nlon, D, return_L=False, np=<module 'numpy' from '/opt/hostedtoolcache/Python/3.14.4/x64/lib/python3.14/site-packages/numpy/__init__.py'>)

Construct SHT diff implementation in plain NumPy.

make_shtdiff(nlat, lmax=None, nlon=None, D=0.0004, return_L=False)

Construct SHT diff implementation with Jax.