Contributing#
Thank you for your interest in contributing to SPD Learn! This guide will help you get started with contributing to the project.
Getting Started#
Setting Up Development Environment#
Fork and clone the repository:
git clone https://github.com/spdlearn/spd_learn.git cd spd_learn
Create a virtual environment:
python -m venv venv source venv/bin/activate # On Windows: venv\Scripts\activate
Install in development mode with all dependencies:
pip install -e ".[all]"
Install pre-commit hooks (required for consistent formatting):
pip install pre-commit pre-commit install pre-commit run --all-files
Development Workflow#
Creating a Branch#
Create a new branch for your feature or bug fix:
git checkout -b feature/your-feature-name
# or
git checkout -b fix/your-bug-fix
Running Tests#
Run the test suite to ensure your changes don’t break existing functionality:
pytest
Run tests with coverage:
pytest --cov=spd_learn --cov-report=html
Code Style#
We use ruff for linting and black for formatting. Run before committing:
ruff check spd_learn/
black spd_learn/
Types of Contributions#
Bug Reports#
If you find a bug, please open an issue on GitHub with:
A clear, descriptive title
Steps to reproduce the bug
Expected behavior vs. actual behavior
Your environment (Python version, PyTorch version, OS)
Minimal code example that reproduces the issue
Feature Requests#
We welcome feature requests! Please open an issue describing:
The problem you’re trying to solve
Your proposed solution
Any alternatives you’ve considered
Code Contributions#
Pull Requests#
Create an issue first for significant changes
Write tests for new functionality
Update documentation if needed
Follow the code style guidelines
Keep PRs focused - one feature/fix per PR
PR Checklist#
Before submitting a PR, ensure:
[ ] All tests pass (
pytest)[ ] Code is formatted (
black,ruff)[ ] Documentation is updated
[ ] Commit messages are clear and descriptive
[ ]
pre-commit run --all-filespasses
Adding New Features#
Adding a New Layer#
To add a new neural network layer:
Create the layer in the appropriate file under
spd_learn/modules/Add comprehensive docstrings following NumPy format
Export in
spd_learn/modules/__init__.pyAdd to
docs/source/api.rstWrite unit tests in
tests/
Example layer structure:
import torch
import torch.nn as nn
class MyNewLayer(nn.Module):
"""Short description of the layer.
Longer description explaining what the layer does,
its mathematical formulation, and when to use it.
Parameters
----------
in_features : int
Input dimension.
out_features : int
Output dimension.
References
----------
.. [1] Author, A. (Year). Paper Title. Journal.
Examples
--------
>>> layer = MyNewLayer(64, 32)
>>> x = torch.randn(16, 64, 64)
>>> output = layer(x)
>>> output.shape
torch.Size([16, 32, 32])
"""
def __init__(self, in_features: int, out_features: int):
super().__init__()
self.in_features = in_features
self.out_features = out_features
# Initialize parameters...
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward pass.
Parameters
----------
x : torch.Tensor
Input tensor of shape ``(batch, in_features, in_features)``.
Returns
-------
torch.Tensor
Output tensor of shape ``(batch, out_features, out_features)``.
"""
# Implementation...
return x
Adding a New Model#
To add a new model architecture:
Create the model in
spd_learn/models/Include a docstring with:
Architecture description
Figure reference (if available)
All parameters documented
Original paper reference
Export in
spd_learn/models/__init__.pyAdd to
docs/source/api.rstConsider adding an example script
Adding a Functional Operation#
For low-level operations:
Add to
spd_learn/functional/Implement both forward and backward passes if using custom autograd
Export in
spd_learn/functional/__init__.pyDocument in
docs/source/api.rst
Documentation#
Building Documentation#
Build the documentation locally:
cd docs
make html
View the built documentation:
open build/html/index.html # macOS
xdg-open build/html/index.html # Linux
Writing Documentation#
Use NumPy-style docstrings for all public functions and classes
Include examples in docstrings where helpful
Add references to papers when implementing published methods
Update the User Guide for conceptual changes
Code Organization#
Architecture Overview#
SPD Learn follows PyTorch’s design philosophy with a clear separation between functional operations and stateful modules. This architecture enables flexibility, composability, and ease of extension.
┌─────────────────────────────────────────┐
│ models/ │
│ Pre-built architectures (SPDNet, etc.) │
│ Compose modules into complete models │
└──────────────────┬──────────────────────┘
│ uses
┌──────────────────▼──────────────────────┐
│ modules/ │
│ Stateful nn.Module layers (BiMap, etc.)│
│ Learnable parameters + forward() │
└──────────────────┬──────────────────────┘
│ calls
┌──────────────────▼──────────────────────┐
│ functional/ │
│ Pure functions (matrix_log, etc.) │
│ Custom autograd for SPD operations │
└─────────────────────────────────────────┘
Design Principles#
SPD Learn follows PyTorch’s design philosophy [Paszke et al., 2019] with a clear separation between stateless operations and stateful modules. This pattern, also adopted by JAX [Frostig et al., 2018], reflects classical systems design principles [Saltzer et al., 1984], [Waldo et al., 1994].
Functional First: Mathematical operations are implemented as pure functions in
spd_learn.functional. No internal state, usable outside neural networks.Modules Wrap Functional:
nn.Moduleclasses inspd_learn.moduleswrap functional operations, adding learnable parameters and state management.Models Compose Modules: Complete architectures in
spd_learn.modelscompose modules into end-to-end trainable networks.Trivialization for Constraints: Manifold-valued parameters (Stiefel, SPD) use trivialization [Lezcano-Casado, 2019]—mappings from unconstrained space to the manifold—enabling standard optimizers.
Package Structure#
spd_learn/
├── __init__.py # Package initialization and public API
├── version.py # Version string
├── logging.py # Logging utilities
│
├── functional/ # Pure functions (no learnable parameters)
│ ├── __init__.py # Public functional API
│ ├── autograd.py # Custom autograd: modeig_forward/backward
│ ├── core.py # Core spectral operations (matrix_log, matrix_exp, etc.)
│ ├── covariance.py # Covariance estimators
│ ├── metrics/ # Riemannian metrics subpackage
│ │ ├── __init__.py # Metrics API aggregation
│ │ ├── airm.py # Affine-Invariant Riemannian Metric
│ │ ├── log_euclidean.py # Log-Euclidean Metric
│ │ ├── bures_wasserstein.py # Bures-Wasserstein Metric
│ │ └── log_cholesky.py # Log-Cholesky Metric
│ ├── transport.py # Parallel transport operations
│ ├── numerical.py # Numerical stability configuration
│ ├── regularize.py # Shrinkage estimators
│ ├── dropout.py # SPD-aware dropout
│ └── utils.py # Helper functions (ensure_sym, etc.)
│
├── modules/ # nn.Module layers (stateful, learnable)
│ ├── __init__.py # Public modules API
│ ├── bilinear.py # BiMap, BiMapIncreaseDim (Stiefel-constrained)
│ ├── modeig.py # LogEig, ReEig, ExpEig (spectral layers)
│ ├── covariance.py # CovLayer (covariance pooling)
│ ├── batchnorm.py # SPDBatchNormMeanVar, SPDBatchNormMean
│ ├── dropout.py # SPDDropout
│ ├── regularize.py # Shrinkage, TraceNorm
│ ├── wavelet.py # WaveletConv (Gabor wavelets)
│ ├── manifold.py # SPD parametrization helpers
│ ├── residual.py # Residual connections on SPD
│ └── utils.py # PatchEmbeddingLayer, Vec, Vech
│
└── models/ # Pre-built architectures
├── __init__.py # Public models API
├── spdnet.py # SPDNet (Huang et al., 2017)
├── tsmnet.py # TSMNet (Kobler et al., 2022)
├── tensorcsp.py # TensorCSPNet (Ju et al., 2022)
├── eegspdnet.py # EEGSPDNet (Wilson et al., 2025)
├── green.py # GREEN (Paillard et al., 2025)
├── matt.py # MAtt (Pan et al., 2022)
└── phase_spdnet.py # PhaseSPDNet (Carrara et al., 2025)
Functional Layer (spd_learn.functional)#
The functional layer provides pure functions for SPD matrix operations. These functions:
Have no internal state or learnable parameters
Are differentiable via custom autograd (when needed)
Can be used standalone outside neural networks
Follow the pattern:
output = function(input, *args)
Key components:
Spectral operations (
functional.py):matrix_log,matrix_exp,matrix_sqrt,matrix_power— implemented via eigendecomposition with custom backward passes using the Loewner matrix formulation.Autograd (
autograd.py):modeig_forwardandmodeig_backwardimplement the forward and backward passes for spectral matrix functions, caching eigendecompositions for efficient gradient computation.Riemannian operations (
metrics/):airm_distance,airm_geodesic,log_euclidean_distance,log_euclidean_meanfor operations under various Riemannian metrics (AIRM, Log-Euclidean, etc.).
Example usage:
import torch
from spd_learn.functional import matrix_log, matrix_exp, log_euclidean_mean
# Create SPD matrices
X = torch.randn(32, 16, 16)
X = X @ X.mT + 0.1 * torch.eye(16)
# Apply matrix logarithm (differentiable)
log_X = matrix_log.apply(X)
# Compute Log-Euclidean mean (unweighted)
mean = matrix_exp.apply(log_X.mean(dim=0))
Modules Layer (spd_learn.modules)#
The modules layer provides stateful ``nn.Module`` classes that wrap functional operations. These modules:
Inherit from
torch.nn.ModuleManage learnable parameters (weights, biases)
Handle manifold constraints via trivialization
Follow the pattern:
output = module(input)
Key components:
BiMap (
bilinear.py): Bilinear mappingY = W^T X WwhereWis constrained to the Stiefel manifold viatorch.nn.utils.parametrizations.orthogonal.LogEig/ReEig/ExpEig (
modeig.py): Spectral layers that apply functions to eigenvalues.ReEigrectifies eigenvalues (non-linearity),LogEigmaps to tangent space,ExpEigmaps back to manifold.SPDBatchNormMeanVar (
batchnorm.py): Batch normalization on the SPD manifold using the Fréchet mean and variance, with learnable SPD bias via trivialization.CovLayer (
covariance.py): Computes covariance matrices from time series, supporting various estimators (sample, shrinkage, real-valued).
Example of trivialization in BiMap:
from torch.nn.utils import parametrizations
class BiMap(nn.Module):
def __init__(self, in_features, out_features):
super().__init__()
# Raw unconstrained parameter
self.weight = nn.Parameter(torch.empty(in_features, out_features))
# Apply orthogonal parametrization (Stiefel constraint)
parametrizations.orthogonal(self, "weight")
def forward(self, X):
# W is automatically orthogonalized
return self.weight.T @ X @ self.weight
Models Layer (spd_learn.models)#
The models layer provides complete architectures that compose modules into end-to-end trainable networks. These models:
Accept raw data (EEG signals, fMRI time series, or precomputed SPD matrices)
Output predictions (class logits, features)
Are ready for training with standard PyTorch optimizers
Architecture pattern (most models follow this):
Input → [CovLayer] → BiMap → ReEig → BiMap → ReEig → LogEig → Vec → Linear → Output
(optional) └────────── SPD Processing ──────────┘ └─ Classifier ─┘
Example model usage:
from spd_learn.models import SPDNet
# Create model for 22-channel EEG, 4 classes
model = SPDNet(n_chans=22, n_outputs=4, subspacedim=16)
# Input: raw EEG (batch, channels, time)
X = torch.randn(32, 22, 500)
logits = model(X) # (32, 4)
Trivialization Approach#
SPD Learn uses trivialization to handle manifold-valued parameters, enabling unconstrained optimization with standard gradient descent.
What is trivialization?
A trivialization is a smooth mapping Φ: ℝ^d → M from an unconstrained
Euclidean space to the target manifold M. Instead of optimizing directly on
the manifold (which requires Riemannian optimization), we optimize the
unconstrained parameters and map them to the manifold.
Stiefel manifold (orthogonal matrices):
For BiMap weights W ∈ St(n, m) (matrices with orthonormal columns):
# PyTorch's built-in orthogonal parametrization
parametrizations.orthogonal(module, "weight")
# Internally uses Cayley map or matrix exponential:
# W = cayley(A) where A is skew-symmetric
SPD manifold (positive definite matrices):
For SPD parameters (e.g., bias in SPDBatchNormMeanVar):
# Map symmetric matrix S to SPD via matrix exponential
# X = exp(S), where S is unconstrained symmetric
X = matrix_exp(S) # Always SPD for any symmetric S
This approach:
Enables standard optimizers (Adam, SGD) without modification
Provides automatic constraint satisfaction by construction
Supports stable gradients via the chain rule through the mapping
Naming Conventions#
Modules:
CamelCase(e.g.,BiMap,SPDBatchNormMeanVar)Functions:
snake_case(e.g.,matrix_log(),log_euclidean_mean())Private methods: prefix with
_(e.g.,_compute_mean)Type hints: Use throughout for clarity and IDE support
Testing Guidelines#
Test Structure#
Tests are organized in tests/ mirroring the package structure:
tests/
├── test_functional.py
├── test_modules.py
└── test_models.py
Writing Tests#
Test both forward and backward passes
Test edge cases (empty batches, single samples)
Test numerical stability with extreme values
Use parametrized tests for multiple configurations
Example test:
import pytest
import torch
from spd_learn.modules import BiMap
@pytest.mark.parametrize(
"in_features,out_features",
[
(64, 32),
(32, 16),
(16, 8),
],
)
def test_bimap_output_shape(in_features, out_features):
layer = BiMap(in_features, out_features)
x = torch.randn(8, in_features, in_features)
x = x @ x.transpose(-1, -2) # Make SPD
y = layer(x)
assert y.shape == (8, out_features, out_features)
def test_bimap_preserves_spd():
layer = BiMap(32, 16)
x = torch.randn(8, 32, 32)
x = x @ x.transpose(-1, -2) + torch.eye(32) # SPD
y = layer(x)
# Check positive definiteness via eigenvalues
eigvals = torch.linalg.eigvalsh(y)
assert (eigvals > 0).all()
Community#
GitHub Issues: Bug reports and feature requests
Pull Requests: Code contributions
Discussions: Questions and ideas
Thank you for contributing to SPD Learn!