multigeom

multigeom: the multi-scale Fisher information geometry of neural population codes, estimated with conditional diffusion models (Azeglio, Laquitaine, Ferrari, Chalk; NeurIPS 2026).

Each group of functions is a subpackage (and the main functions are also available
directly as `multigeom.<name>`).

`multigeom.encoders`: use one that is already fitted, build your own, or bring any
PyTorch model that follows the contract of `multigeom.encoders.contract`:

    load_encoder(name, area)
        an encoder of the paper, fitted to macaque V1 or V4 (downloaded)
    NeuralEncoder(backbone, n_neurons)
        build your own encoder: fit_readout, predict, save, load
    validate_encoder(encoder)
        check that an encoder follows the contract

`multigeom.geometry`: the Fisher metric and the multi-scale metric (the latter needs the
diffusion model of `multigeom.diffusion`):

    compute_fisher(image, encoder, sigma)
        the Fisher information J(x): its top eigenvalues and eigenvectors
    multigeom.diffusion
        the conditional diffusion model of an encoder: train it on its responses
    compute_g(image, encoder, unet, ...)
        the multi-scale Fisher metric G(x) (needs the trained diffusion model)

`multigeom.analysis`: information-geometric analyses of the metrics:

    subspace_alignment(eigvecs_a, eigvecs_b)
        the subspace alignment of two sets of eigenvectors
    bures_wasserstein(...)
        the Bures-Wasserstein distance between two reconstructed metrics

Install the core with `pip install multigeom`; the extra `[paper]` adds what reproducing
the paper needs (the ON-OFF encoder and the TVSD loaders), and `[all]` also Jupyter and
the Modal launcher. Subpackages: `encoders`, `geometry`, `analysis`, `diffusion`,
`pipeline`, `datasets`, `fit`.

Modules

analysis	Analysis (core: torch): measures that compare the geometries of two encoders, from the eigenvectors and eigenvalues that multigeom.geometry returns: the subspace alignment of two sets of eigenvectors, the Bures-Wasserstein distance between two reconstructed metrics, and the matching of two sets of eigenvectors for display.
cli	Command line: multigeom <command> [options] (or python -m multigeom.cli).
config	Configuration: loading YAML files, the default configurations shipped with the package, and the paths of a project.
datasets	Data: TinyImageNet and the canonical image transform of the pipeline (images 64 x 64, one channel, in [-1, 1]), the functions that download the datasets (TinyImageNet, TVSD, THINGS images), and synthetic neural data with a known ground truth.
diffusion	Diffusion (core, needs diffusers): the conditional diffusion model of an encoder, everything G(x) needs besides the encoder.
encoders	Encoders, the models that predict neural responses from images: a backbone followed by a readout fitted to neural data.
fit	Fitting to the paper's recordings (extra [paper]: h5py, plenoptic).
geometry	Geometry (core: torch, numpy, scipy): the Fisher information J(x) (compute_fisher), the multi-scale Fisher metric G(x) of the conditional-pair estimator (compute_g), and the local metric field of Fig.
pipeline	Pipeline (core): the orchestration used by the notebooks and the command line: G(x) and J(x) for a set of images (per_image), the comparison between configurations (alignment), the local metric field (metric_field) and running them on several workers (distribute).
utils	Shared utilities: RNG seeding, provenance records, device selection, a progress line and showing a figure in a notebook.