A geometric and topological deep learning framework for uncovering spatiotemporal signatures in neural activity.
Implementation of Neurospectrum: A Geometric and Topological Deep Learning Framework for Uncovering Spatiotemporal Signatures in Neural Activity (Bhaskar, Zhang, Moore et al.).
Neural activity is modelled as a time-varying signal on a graph. At each timepoint the spatial pattern over the graph is encoded with multiscale diffusion wavelets modulated by learnable attention. The sequence of those embeddings is compressed by a T-PHATE-regularized autoencoder into a low-dimensional latent trajectory that preserves the similarity and ordering of brain states. That trajectory is then summarized with geometric and topological descriptors — curvature, path signatures, Betti curves, recurrent embeddings — which feed a downstream classifier or regressor.
graph signal scattering + attention T-PHATE-regularized AE descriptors task
X(t) ∈ R^N → S̃(X(t)) ∈ R^{N×C} → z(t) ∈ R^d → κ, PS, BC, RNN → ŷ
Two packages, deliberately separate:
neurospectrum/ representation learning (paper Sec. 3)
├── config.py dataclass configuration for every stage
├── data/
│ ├── graphs.py hexagonal lattice, kNN / radius / Gaussian graphs, lazy random walk
│ ├── kuramoto.py Kuramoto oscillator simulator + benchmark dataset
│ └── datasets.py scattering + T-PHATE precomputation, group-aware folds
├── spatial/ Sec. 3.1 — spatial module
│ ├── scattering.py diffusion wavelet scattering (orders 0–2), differentiable
│ └── attention.py learnable attention α_{ℓ,j} over nodes and scales
└── temporal/ Sec. 3.2 — temporal module
├── tphate.py dual-view diffusion operator and potential distances
├── autoencoder.py encoder / decoder over per-timepoint coefficients
├── losses.py reconstruction + T-PHATE distance matching
└── trainer.py stage-1 training loop, checkpointing, latent extraction
trajectory_descriptors/ geometry and topology of latent trajectories (Appendix A)
├── curvature.py circle-fit curvature (paper) + differentiable Frenet estimator
├── path_signature.py iterated integrals via Chen's identity, differentiable
├── persistence.py Vietoris-Rips diagrams, Betti curves, diagram statistics
├── rnn.py recurrent trajectory summary
├── pipeline.py descriptors → fixed-length feature matrix
├── models.py post-hoc predictor and end-to-end differentiable predictor
└── evaluate.py group-aware 5-fold cross-validation and per-block ablation
scripts/ four-step command-line pipeline
tests/ 98 tests, including closed-form checks of each descriptor
legacy/ the pre-refactor fNIRS/fMRI scripts (see legacy/README.md)
data/ DiFuMo atlas and the OCD / control fMRI datasets
trajectory_descriptors does not import neurospectrum. It takes latent
trajectories as plain (S, T, d) arrays and can be used on trajectories from any
source.
pip install -e ".[topology,viz,dev]"Or, without installing, run the scripts straight from a checkout — they add the
repository root to sys.path themselves.
ripser is optional. Without a homology backend, H_0 diagrams are still
computed exactly (from the Euclidean minimum spanning tree, which is equivalent
for Vietoris-Rips H_0); loops and higher dimensions need ripser or gudhi
and a warning is raised.
100 oscillators on a hexagonal grid; the task is the inverse problem of recovering the coupling strength η from the observed node time series (paper Sec. 4.2).
# 1. simulate — 10 coupling values x 10 replicates, 2000 steps at dt = 0.1
python scripts/simulate_kuramoto.py --output data/kuramoto/coupling.npz --transient
# 2. stage 1 — scattering + attention + T-PHATE-regularized autoencoder
python scripts/train_autoencoder.py \
--dataset data/kuramoto/coupling.npz \
--output results/kuramoto \
--config configs/kuramoto.json
# 3. stage 2 — geometric and topological descriptors of the latent trajectories
python scripts/extract_descriptors.py \
--trajectories results/kuramoto/latent_trajectories.npz \
--output results/kuramoto/descriptors.npz --rnn
# 4. stage 3 — downstream regression of η, with the per-descriptor ablation
python scripts/run_downstream.py \
--descriptors results/kuramoto/descriptors.npz --ablateStep 2 writes autoencoder.pt, latent_trajectories.npz, attention.npz (the
learned importance over nodes and scales) and the resolved config.json.
import numpy as np
from neurospectrum.data import ScatteringTrajectoryDataset, generate_coupling_dataset
from neurospectrum.temporal import (
build_autoencoder, extract_latent_trajectories, train_autoencoder,
)
from trajectory_descriptors import (
build_feature_matrix, cross_validate, extract_descriptors,
)
raw = generate_coupling_dataset() # signals (S, T, N), targets (S,)
dataset = ScatteringTrajectoryDataset( # scattering + T-PHATE distances
signals=raw["signals"],
adjacency=raw["adjacency"],
targets=raw["targets"],
groups=raw["group"],
)
model = build_autoencoder(dataset)
model, history = train_autoencoder(dataset, model=model)
trajectories = extract_latent_trajectories(model, dataset) # (S, T, d)
descriptors = extract_descriptors(trajectories)
features, _ = build_feature_matrix(descriptors)
print(cross_validate(features, raw["targets"], groups=raw["group"]))The pipeline needs two things: a graph and a signal on it.
from neurospectrum.data import ScatteringTrajectoryDataset, graph_from_coordinates
# calcium imaging: connect cells within 2 um of each other
adjacency = graph_from_coordinates(cell_xyz, method="radius", radius=2.0)
# fMRI: a graph over atlas parcels (e.g. the DiFuMo coordinates under data/)
adjacency = graph_from_coordinates(parcel_xyz, method="knn", k=8)
dataset = ScatteringTrajectoryDataset(
signals=bold, # (n_subjects, n_timepoints, n_parcels)
adjacency=adjacency,
targets=diagnosis, # (n_subjects,)
groups=subject_id, # keeps all runs of a subject inside one fold
)Everything downstream is unchanged. Node counts, timepoint counts and the number of wavelet scales are all read from the data.
Scattering. The filter bank is Ψ_0 = I − R and Ψ_j = R^{2^{j-1}} − R^{2^j}
for 1 ≤ j ≤ J, built from the lazy random walk R = ½(I + AD⁻¹). Channels per
node are 1 + (J+1) + (J+1)J/2; the default J = 4 gives 16, matching the
768 = 48 × 16 feature dimension of the original fNIRS code. The transform is a
torch.nn.Module, so gradients flow back to the attention weights, and it is
verified to be permutation equivariant.
Attention. ScatteringAttention learns α over (node, channel) by default —
the paper's α_{ℓ,j} — with node- and channel-only variants for ablations.
Weights are non-negative (softplus) and initialized near 1, so training starts
from the plain scattering transform. importance() returns the map broken down
by node and by scale.
T-PHATE regularization. P_D is the density-normalized anisotropic kernel
over timepoint feature vectors; P_T comes from the autocorrelation function,
truncated at its first zero crossing. The dual operator P = P_D P_T is powered
to t_D (chosen from the von Neumann entropy knee) and the log-potential
distances become the regularization target. Distances are precomputed per
recording — the paper notes this term "can be precomputed and stored" — and
normalized to unit mean so lambda_tphate means the same thing across datasets.
Training samples a contiguous block of timepoints per recording, bounding the
quadratic term at O(timepoints_per_batch²).
Descriptors. Every one is checked against a case with a known answer:
curvature of a radius-r circle is 1/r (and exactly 0 on a straight line);
the level-1 signature is net displacement and the antisymmetric part of level 2
is the signed area (π for the unit circle); a circular point cloud has exactly
one persistent 1-cycle; the MST fallback for H_0 matches ripser bit for bit.
Two ways to consume descriptors. DescriptorPredictor takes precomputed
features from frozen trajectories. EndToEndPredictor computes the
differentiable ones (curvature, signatures, RNN) inside the network, so the
encoder, latent representation and predictor are optimized jointly — the paper's
end-to-end mode. Betti curves are not differentiable and enter there as a
precomputed side input.
Cross-validation. grouped_kfold keeps every recording sharing a group id
inside one fold. For the Kuramoto dataset the group is the coupling value,
which is what the paper specifies ("no replicate of a given η appeared in both
training and validation") — note that this makes the task extrapolative for the
two folds holding out the range endpoints. Pass groups=np.arange(n) for
per-simulation splits instead.
The modular design is the point; each piece can be switched off.
| Ablation | How |
|---|---|
| No T-PHATE regularization | train_autoencoder.py --no-tphate (or --lambda-tphate 0) |
| No spatial encoding | build the dataset with ScatteringConfig(max_order=0) |
| Node-only / scale-only attention | --attention node / --attention channel |
| Single descriptor | run_downstream.py --blocks betti |
| Per-descriptor table (Table D.4) | run_downstream.py --ablate |
pytest tests — 98 tests covering the graph construction, the Kuramoto
integrator, the scattering transform (shapes, zeroth order, permutation
equivariance), the T-PHATE operators (row-stochasticity, metric properties of
the potential distances), the autoencoder and its checkpointing, every descriptor
against its closed form, and the cross-validation machinery.
An end-to-end run of the four scripts on 50 simulations (10 coupling values × 5 replicates, 400 timepoints, latent dim 8, 200 epochs) recovers η well above chance under the paper's η-holdout protocol:
| Descriptor | MSE |
|---|---|
| Betti curves | 0.028 ± 0.018 |
| RNN | 0.032 ± 0.019 |
| Curvature | 0.047 ± 0.018 |
| All | 0.040 ± 0.030 |
| Path signatures | 0.094 ± 0.051 |
| (predicting the mean) | 0.102 |
The area under the H_0 Betti curve of the latent trajectory also falls
monotonically with η (193 → 82 from η = 0 to η = 1), reproducing the
qualitative result of Fig. 3C: as oscillators entrain, the trajectory collapses
and its components merge earlier in the filtration. These are numbers from one
reduced run, not a reproduction of the paper's Table 1 — that needs the full 100
simulations and tuning.
Eq. (1) normalizes the coupling sum by N, the mean-field convention, which
assumes all-to-all coupling. The hexagonal grid of Fig. 3A couples each
oscillator to at most six neighbours, and dividing a six-term sum by N = 100
makes the effective coupling ~20× too weak: the population then never
synchronizes anywhere in η ∈ [0, 1] (measured order parameter stays at ~0.1
across the whole range), so there is nothing for the model to recover.
simulate_kuramoto therefore normalizes by node degree when coupling through
the graph (normalization="auto", the default), which reproduces the regimes
the paper describes — desynchronized at η = 0.1, partial clusters at η = 0.5,
global entrainment at η = 0.9:
| η | 0.0 | 0.2 | 0.4 | 0.6 | 0.8 | 1.0 |
|---|---|---|---|---|---|---|
order parameter r |
0.10 | 0.37 | 0.54 | 0.80 | 0.88 | 0.93 |
Pass normalization="nodes" for the literal Eq. (1), and coupling_mode="global"
for genuine all-to-all coupling (where the two normalizations agree).
The RK4 integrator at dt = 0.1 resolves η up to roughly 2; much larger coupling
needs a smaller timestep.
- The calcium-imaging and fMRI experiments of Sec. 4.3–4.4. The data loaders
are generic and
data/holds the fMRI arrays and the DiFuMo atlas, but the preprocessing and task definitions for those datasets are not wired up. - The baseline methods of Appendix B (PCA, ICA, HMM, Granger causality, dFC, GCNs, zig-zag persistence) used for Table 1.
- CROCKER matrices (
GraphPH-CROCKER, Appendix D).