spectralbrain.backends.cpu#

CPU compute backend — NumPy/SciPy, Bayesian samplers, and parallelisation.

This module provides:

NumpyBackend — the default compute engine using SciPy sparse eigensolvers (ARPACK), sparse matrix operations, and NumPy array algebra. Every other backend mirrors this interface.
PyMCSampler / NutpieSampler — Bayesian MCMC backends for the statistics/bayesian.py module.
Joblib utilities — composable parallelisation helpers with Rich progress integration.
RAM management — memory monitoring, garbage collection, and estimation helpers for multi-subject pipelines.

All optional dependencies (PyMC, nutpie, joblib) are lazy-imported. Only NumPy and SciPy are hard requirements.

Functions

`batch_iterator`(data[, batch_size, axis])	Iterate over an array in memory-safe batches.
`estimate_array_memory`(shape[, dtype])	Estimate memory for an array in gigabytes.
`gc_collect`([generations])	Force Python garbage collection and return bytes freed.
`get_bayesian_sampler`([backend, config])	Factory for CPU Bayesian samplers.
`memory_guard`([min_available_gb, error_on_low])	Context manager that checks RAM before and after a block.
`parallel_batch`(func, data, *[, batch_size, ...])	Apply func to batches of an array in parallel.
`parallel_map`(func, items, *[, n_jobs, ...])	Apply func to each item in parallel with optional progress.
`ram_status`()	Return current system RAM usage.
`shrink_array`(arr[, target_dtype])	Downcast an array to save memory.

Classes

`MemoryInfo`(total_gb, available_gb, used_gb, ...)	Snapshot of system RAM usage.
`NumpyBackend`()	CPU compute backend using NumPy + SciPy.
`NutpieSampler`([config])	Bayesian sampler using nutpie (Rust-based NUTS).
`PyMCSampler`([config])	Bayesian sampler using PyMC's native NUTS implementation.
`SamplerConfig`([draws, tune, chains, cores, ...])	Configuration for Bayesian MCMC samplers.

class spectralbrain.backends.cpu.MemoryInfo(total_gb, available_gb, used_gb, percent_used)[source]#

Bases: object

Snapshot of system RAM usage.

Parameters:

total_gb (float)
available_gb (float)
used_gb (float)
percent_used (float)

total_gb#

Total physical RAM.

Type:: float

available_gb#

Available (free + cached) RAM.

Type:: float

used_gb#

Actively used RAM.

Type:: float

percent_used#

Usage percentage (0–100).

Type:: float

available_gb: float#

percent_used: float#

total_gb: float#

used_gb: float#

class spectralbrain.backends.cpu.NumpyBackend[source]#

Bases: object

CPU compute backend using NumPy + SciPy.

Provides the canonical interface that CupyBackend and JaxBackend mirror. All core/ and spectral/ modules call backend methods rather than importing NumPy or SciPy directly, enabling transparent GPU acceleration.

Examples

>>> from spectralbrain.backends.cpu import NumpyBackend
>>> be = NumpyBackend()
>>> evals, evecs = be.eigsh(L, M, k=100)
>>> hks = be.exp(-evals[None, :] * t[:, None])  # broadcasting

static argsort(x, axis=-1)[source]#

Indirect sort indices (mirrors numpy.argsort).

Parameters:

x (ndarray)
axis (int)

Return type:

ndarray

static array(data, dtype=<class 'numpy.float64'>)[source]#

Create a dense array.

Parameters:

data (Any)
dtype (dtype)

Return type:

ndarray

static clip(x, a_min, a_max)[source]#

Element-wise clip (mirrors numpy.clip).

Parameters:

x (ndarray)
a_min (float | None)
a_max (float | None)

Return type:

ndarray

static concatenate(arrays, axis=0)[source]#

Concatenate arrays along an axis.

Parameters:

arrays (Sequence[ndarray])
axis (int)

Return type:

ndarray

static eigsh(L, M=None, k=100, *, sigma=-0.01, which='LM', tol=0.0, maxiter=None)[source]#

Solve the generalised sparse eigenproblem L v = λ M v.

Uses SciPy’s ARPACK wrapper in shift-invert mode (default σ = −0.01) which is optimal for computing the smallest eigenvalues of the Laplacian.

Parameters:

L (sparse matrix, shape (N, N)) – Stiffness (Laplacian) matrix — symmetric positive semi-definite.
M (sparse matrix, shape (N, N), optional) – Mass matrix. If None, the standard eigenproblem L v = λ v is solved.
k (int) – Number of eigenpairs to compute.
sigma (float) – Shift for shift-invert mode. A small negative value avoids the singularity at λ = 0.
which (str) – Which eigenvalues to target ("LM" = largest magnitude of the shifted operator, yielding the smallest λ).
tol (float) – Convergence tolerance (0 = machine precision).
maxiter (int, optional) – Maximum ARPACK iterations.

Returns:

eigenvalues (ndarray, shape (k,)) – Sorted ascending, float64.
eigenvectors (ndarray, shape (N, k)) – Corresponding eigenvectors, M-orthonormal.

Raises:

scipy.sparse.linalg.ArpackNoConvergence – If ARPACK fails to converge within maxiter iterations.

Return type:

tuple[ndarray[tuple[Any, …], dtype[floating]], ndarray[tuple[Any, …], dtype[floating]]]

static exp(x)[source]#

Element-wise exponential (mirrors numpy.exp).

Parameters:: x (ndarray)
Return type:: ndarray

static eye(n, dtype=<class 'numpy.float64'>)[source]#

Create an identity matrix (mirrors numpy.eye).

Parameters:

n (int)
dtype (dtype)

Return type:

ndarray

static linspace(start, stop, num)[source]#

Linearly spaced values (mirrors numpy.linspace).

Parameters:

start (float)
stop (float)
num (int)

Return type:

ndarray

static log(x)[source]#

Element-wise safe log with clamp at 1e-300.

Parameters:: x (ndarray)
Return type:: ndarray

static logspace(start, stop, num)[source]#

Log-spaced values (mirrors numpy.logspace).

Parameters:

start (float)
stop (float)
num (int)

Return type:

ndarray

static matmul(a, b)[source]#

Matrix multiply (sparse- and dense-aware).

Parameters:

a (ndarray)
b (ndarray)

Return type:

ndarray

static mean(x, axis=None)[source]#

Mean reduction (mirrors numpy.mean).

Parameters:

x (ndarray)
axis (int | None)

Return type:

ndarray

static norm(x, axis=None, ord=None)[source]#

Vector/matrix norm (mirrors numpy.linalg.norm).

Parameters:

x (ndarray)
axis (int | None)
ord (int | None)

Return type:

ndarray

static ones(shape, dtype=<class 'numpy.float64'>)[source]#

Create a ones-filled array (mirrors numpy.ones).

Parameters:

shape (tuple[int, ...])
dtype (dtype)

Return type:

ndarray

static sparse_matrix(data, row, col, shape, *, format='csc')[source]#

Build a sparse matrix from COO triplets.

Parameters:

data (ndarray) – Non-zero values.
row (ndarray) – Row and column indices.
col (ndarray) – Row and column indices.
shape ((int, int)) – Matrix dimensions.
format (str) – Output format ("csc", "csr", "coo").

Returns:

SparseMatrix

Return type:

spmatrix

static sqrt(x)[source]#

Element-wise safe sqrt with clamp at 0.

Parameters:: x (ndarray)
Return type:: ndarray

static stack(arrays, axis=0)[source]#

Stack arrays along a new axis.

Parameters:

arrays (Sequence[ndarray])
axis (int)

Return type:

ndarray

static sum(x, axis=None)[source]#

Sum reduction (mirrors numpy.sum).

Parameters:

x (ndarray)
axis (int | None)

Return type:

ndarray

static to_numpy(x)[source]#

Convert any array-like to a NumPy ndarray.

Parameters:: x (Any)
Return type:: ndarray

static zeros(shape, dtype=<class 'numpy.float64'>)[source]#

Create a zero-filled array (mirrors numpy.zeros).

Parameters:

shape (tuple[int, ...])
dtype (dtype)

Return type:

ndarray

name: str = 'numpy'#

class spectralbrain.backends.cpu.NutpieSampler(config=None)[source]#

Bases: object

Bayesian sampler using nutpie (Rust-based NUTS).

nutpie is a high-performance drop-in replacement for PyMC’s default sampler. It compiles the PyMC model to Rust and runs NUTS 2–10× faster on CPU.

Parameters:: config (SamplerConfig, optional) – Sampling configuration.

Examples

>>> sampler = NutpieSampler(SamplerConfig(draws=2000))
>>> trace = sampler.sample(model)

sample(model, **kwargs)[source]#

Run nutpie NUTS on a PyMC model.

Parameters:

model (pymc.Model) – A fully specified PyMC model.
**kwargs – Overrides passed to nutpie.sample().

Returns:

arviz.InferenceData

Return type:

Any

name: str = 'nutpie'#

class spectralbrain.backends.cpu.PyMCSampler(config=None)[source]#

Bases: object

Bayesian sampler using PyMC’s native NUTS implementation.

This is the default CPU sampler. It wraps pymc.sample() with SpectralBrain-compatible configuration and logging.

Parameters:: config (SamplerConfig, optional) – Sampling configuration.

Examples

>>> sampler = PyMCSampler(SamplerConfig(draws=1000, chains=2))
>>> with pm.Model() as model:
...     mu = pm.Normal("mu", 0, 1)
...     obs = pm.Normal("obs", mu, 1, observed=data)
>>> trace = sampler.sample(model)

sample(model, **kwargs)[source]#

Run NUTS sampling on a PyMC model.

Parameters:

model (pymc.Model) – A fully specified PyMC model.
**kwargs – Overrides passed to pymc.sample().

Returns:

arviz.InferenceData – Posterior samples with diagnostics.

Return type:

Any

name: str = 'nuts'#

class spectralbrain.backends.cpu.SamplerConfig(draws=2000, tune=1000, chains=4, cores=4, target_accept=0.95, random_seed=42)[source]#

Bases: object

Configuration for Bayesian MCMC samplers.

Parameters:

draws (int) – Number of posterior draws per chain.
tune (int) – Number of tuning (burn-in) samples.
chains (int) – Number of independent chains.
cores (int) – CPU cores for parallel chains.
target_accept (float) – Target acceptance probability for NUTS.
random_seed (int or None) – RNG seed for reproducibility.

chains: int = 4#

cores: int = 4#

draws: int = 2000#

random_seed: int | None = 42#

target_accept: float = 0.95#

tune: int = 1000#

spectralbrain.backends.cpu.batch_iterator(data, batch_size=1000, *, axis=0)[source]#

Iterate over an array in memory-safe batches.

Unlike parallel_batch(), this is a sequential generator suitable for GPU-offloading loops where only one batch should be in memory at a time.

Parameters:

data (ndarray) – Array to iterate.
batch_size (int) – Rows per batch.
axis (int) – Axis to split along.

Yields:

ndarray – A view (not copy) of the batch.

Return type:

Iterator[ndarray]

Examples

>>> for batch in batch_iterator(big_array, batch_size=500):
...     result = expensive_compute(batch)
...     accumulate(result)

spectralbrain.backends.cpu.estimate_array_memory(shape, dtype=<class 'numpy.float64'>)[source]#

Estimate memory for an array in gigabytes.

Parameters:

shape (tuple of int)
dtype (numpy dtype)

Returns:

float – Estimated size in GB.

Return type:

float

Examples

>>> estimate_array_memory((160_000, 300), np.float64)
0.358  # ~358 MB for cortical eigenvectors

spectralbrain.backends.cpu.gc_collect(generations=2)[source]#

Force Python garbage collection and return bytes freed.

Parameters:: generations (int) – GC generations to collect (0, 1, or 2).
Returns:: int – Number of unreachable objects collected.
Return type:: int

Examples

>>> del large_array
>>> freed = gc_collect()
>>> logger.info("GC freed %d objects", freed)

spectralbrain.backends.cpu.get_bayesian_sampler(backend='nuts', config=None)[source]#

Factory for CPU Bayesian samplers.

Parameters:

backend ("nuts" or "nutpie") – Which sampler to use.
config (SamplerConfig, optional) – Sampling parameters.

Returns:

PyMCSampler or NutpieSampler

Return type:

PyMCSampler | NutpieSampler

spectralbrain.backends.cpu.memory_guard(min_available_gb=2.0, error_on_low=False)[source]#

Context manager that checks RAM before and after a block.

Parameters:

min_available_gb (float) – Minimum free RAM required to proceed.
error_on_low (bool) – Raise MemoryError if RAM is below threshold. If False (default), logs a warning instead.

Return type:

Generator[None, None, None]

Examples

>>> with memory_guard(min_available_gb=4.0):
...     big_result = compute_all_subjects()

spectralbrain.backends.cpu.parallel_batch(func, data, *, batch_size=1000, n_jobs=-1, axis=0, progress=True, description='Batch processing')[source]#

Apply func to batches of an array in parallel.

Splits data along axis into chunks of batch_size, applies func to each chunk in parallel, and concatenates the results. Useful for operations that are O(N²) per vertex but can be batched (e.g. geodesic distance computation).

Parameters:

func (callable) – Function that accepts an ndarray batch and returns an ndarray.
data (ndarray) – Full array to process.
batch_size (int) – Number of rows per batch.
n_jobs (int) – Parallel workers.
axis (int) – Axis along which to split.
progress (bool) – Show progress bar.
description (str) – Progress label.

Returns:

ndarray – Concatenated results.

Return type:

ndarray

spectralbrain.backends.cpu.parallel_map(func, items, *, n_jobs=-1, backend='loky', progress=True, description='Processing', **func_kwargs)[source]#

Apply func to each item in parallel with optional progress.

A thin wrapper around joblib.Parallel that integrates with SpectralBrain’s Rich progress bars.

Parameters:

func (callable) – Function to apply. Must accept each item as its first positional argument.
items (sequence) – Items to process.
n_jobs (int) – Number of parallel workers (-1 = all cores).
backend (str) – Joblib backend ("loky", "threading", "multiprocessing").
progress (bool) – Show a Rich progress bar.
description (str) – Progress bar label.
**func_kwargs – Extra keyword arguments passed to func.

Returns:

list – Results in the same order as items.

Return type:

list[R]

Examples

>>> def compute(subj_id, k=100):
...     mesh = load(subj_id)
...     return decompose(mesh, k=k)
>>> results = parallel_map(compute, subject_ids, n_jobs=8, k=50)

spectralbrain.backends.cpu.ram_status()[source]#

Return current system RAM usage.

Uses /proc/meminfo on Linux and psutil as fallback.

Returns:: MemoryInfo
Return type:: MemoryInfo

spectralbrain.backends.cpu.shrink_array(arr, target_dtype=None)[source]#

Downcast an array to save memory.

If target_dtype is None, applies safe downcasting rules: float64 → float32, int64 → int32 (if values fit).

Parameters:

arr (ndarray)
target_dtype (dtype, optional)

Returns:

ndarray – A (possibly) smaller copy.

Return type:

ndarray