Package architecture#

This page describes the internal structure of the Latents package for contributors and developers looking to extend it.

Package structure#

src/latents/
├── __init__.py
├── _version.py           # Auto-generated by hatch-vcs
├── base.py               # ArrayContainer base class
├── data.py               # Observation containers: ObsStatic, ObsTimeSeries
├── tracking.py           # Base FitTracker, FitFlags classes
├── callbacks.py          # ProgressCallback, LoggingCallback, CheckpointCallback
│
├── observation/          # Observation model components
│   ├── priors.py         # ObsParamsHyperPrior, ObsParamsPrior
│   ├── posteriors.py     # ObsParamsPosterior, LoadingPosterior, etc.
│   └── realizations.py   # ObsParamsRealization, ObsParamsPoint
│
├── state/                # State model components
│   ├── priors.py         # LatentsPriorStatic, LatentsPriorGP
│   ├── posteriors.py     # LatentsPosteriorStatic, LatentsPosteriorDelayed
│   └── realizations.py   # LatentsRealization
│
├── plotting/             # Visualization utilities
│   ├── hinton.py         # hinton_diagram()
│   ├── latents.py        # Latent variable visualizations
│   └── observation.py    # plot_dimensionalities(), plot_var_exp(), etc.
│
├── gfa/                  # Group Factor Analysis
│   ├── model.py          # GFAModel class
│   ├── config.py         # GFAFitConfig, GFASimConfig
│   ├── inference.py      # fit(), infer_latents()
│   ├── tracking.py       # GFAFitTracker, GFAFitFlags, GFAFitContext
│   ├── simulation.py     # simulate(), sample_observations()
│   └── analysis.py       # predictive_performance()
│
├── mdlag/                # mDLAG (parallel structure)
│
└── _internal/            # Internal utilities
    ├── logging.py        # Structured logging
    └── numerics.py       # stability_floor(), validate_tolerance()

Method subpackages#

Each method subpackage (gfa/, mdlag/) follows a consistent structure:

Module

Purpose

model.py

Model class with user-facing API (fit(), infer_latents())

config.py

Frozen dataclass for fitting configuration

inference.py

Fitting algorithm implementation

tracking.py

FitTracker and FitFlags subclasses

simulation.py

Forward model for data generation

analysis.py

Method-specific post-hoc analysis

The model class stores:

  • Priors — set at construction

  • Posteriors — populated after fitting

  • Config — fitting hyperparameters

  • Tracker — ELBO and runtime history

  • Flags — convergence status

Observation model#

The observation/ subpackage contains components organized by probabilistic level:

priors.py

  • ObsParamsHyperPrior — homogeneous scalar hyperpriors

  • ObsParamsHyperPriorStructured — per-group, per-latent hyperpriors

  • ObsParamsPrior — prior with sample() method

posteriors.py

  • ObsParamsPosterior — bundle of all parameter posteriors

  • LoadingPosterior, ARDPosterior, ObsMeanPosterior, ObsPrecPosterior

realizations.py

  • ObsParamsRealization — sampled parameter values (Bayesian)

  • ObsParamsPoint — point estimates (non-Bayesian)

State model#

The state/ subpackage follows the same organization:

priors.py

  • LatentsPriorStatic — i.i.d. Gaussian prior with sample()

  • LatentsPriorGP — GP prior with sample() (for time series methods)

posteriors.py

  • LatentsPosteriorStatic — posterior for static latents

  • LatentsPosteriorDelayed — posterior for delayed latents (mDLAG)

realizations.py

  • LatentsRealization — sampled latent values

Data containers#

data.py provides observation containers with group structure:

  • ObsStatic — static observations, shape (y_dim, n_samples)

  • ObsTimeSeries — variable-length sequences per trial

Both support get_groups() to split by group boundaries.

Internal utilities#

_internal/ contains utilities not exposed in the public API:

  • numerics.py — stability_floor(), validate_tolerance() for numerical stability

  • logging.py — structured logging via FitEvent and log_event()

Extending Latents#

The architecture enables adding new methods by composing existing components.

Adding a new method#

  1. Identify components — which building blocks from observation/ and state/ apply? What’s new?

  2. Create method subpackage — follow the established structure

  3. Compose in the model class — wire together appropriate priors

  4. Implement inference — the method-specific fitting algorithm

Example: Factor Analysis#

FA is non-Bayesian single-group factor analysis:

  • State prior: LatentsPriorStatic (reuse)

  • Observation model: point estimates, no ARD → ObsParamsPoint

  • Inference: EM without ARD updates

class FAModel:
    latents_prior: LatentsPriorStatic
    obs_params: ObsParamsPoint      # Point estimates, not posterior
    latents_posterior: LatentsPosteriorStatic

Example: GPFA#

GPFA adds temporal dynamics via GP prior:

  • State prior: LatentsPriorGP (new)

  • Observation model: point estimates → ObsParamsPoint (reuse)

  • Inference: EM with gradient updates for GP hyperparameters

class GPFAModel:
    latents_prior: LatentsPriorGP
    obs_params: ObsParamsPoint
    latents_posterior: LatentsPosteriorTimeSeries

Shared components eliminate redundant implementation while the method subpackage structure keeps each method’s interface clean.