From 26e3c05d7a82f1156ab3a407f1ca837330102d51 Mon Sep 17 00:00:00 2001 From: Bru Date: Fri, 19 Jun 2026 18:38:02 +0200 Subject: [PATCH 1/2] feat: depend on pyriemann 0.12 (Array API) for SPD geometry Make pyriemann>=0.12 a core dependency (git-pinned until 0.12 releases on PyPI) and delegate the geometry it need not keep numerically special to its Array API backend, which runs natively on torch tensors: - Delegated to pyriemann: the 4 geodesics, AIRM/LEM/log-Cholesky parallel transport, and Frechet derivatives (ddlogm/ddexpm). - Re-exported pyriemann's broader geometry toolkit (~79 functions: distances/means/geodesics/tangent maps/kernels/ajd) from spd_learn.functional so the full SPD/Riemannian API is one import surface. Kept in spd_learn (numerically-stable, custom-autograd, eigenvalue-clamped), because pyriemann's unclamped log returns NaN on ill-conditioned float32 and its raw eigh backward is unstable near clustered eigenvalues (both break SPDNet/TensorCSPNet/LieBN training and complex-dtype batch norm): - the 5 core matrix-function primitives (matrix_log/exp/sqrt/inv_sqrt/power), - the 4 distances (airm/log_euclidean/log_cholesky/bures_wasserstein), - the weighted/iterative means, tangent exp/log maps, and ladder transports. Also: modernize the pyriemann comparison test (pyriemann.geometry.* imports, autograd-ceiling guard test), document AttentionManifold/PhaseDelay, and update the installation/FAQ/API docs. --- docs/source/api.rst | 177 +++++++++++++++++ docs/source/faq.rst | 20 +- docs/source/installation.rst | 21 +- pyproject.toml | 7 +- spd_learn/functional/__init__.py | 182 ++++++++++++++++++ spd_learn/functional/frechet.py | 135 ++----------- .../functional/metrics/affine_invariant.py | 24 ++- .../functional/metrics/bures_wasserstein.py | 41 +--- spd_learn/functional/metrics/log_cholesky.py | 16 +- spd_learn/functional/metrics/log_euclidean.py | 13 +- spd_learn/functional/parallel_transport.py | 117 ++--------- spd_learn/models/__init__.py | 13 +- tests/test_pyriemann_comparison.py | 69 +++++-- 13 files changed, 545 insertions(+), 290 deletions(-) diff --git a/docs/source/api.rst b/docs/source/api.rst index c8e8445..ab375a4 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -129,6 +129,7 @@ under congruence transformations. airm_geodesic exp_map_airm log_map_airm + spd_egrad2rgrad Log-Euclidean Metric @@ -194,6 +195,28 @@ manifold, essential for operations like domain adaptation. transport_tangent_vector +Fréchet Derivatives +------------------- +Directional (Fréchet) derivatives of the matrix logarithm and exponential, +used for parallel transport under the Log-Euclidean metric. + +.. autosummary:: + :toctree: generated/frechet_functional + + frechet_derivative_log + frechet_derivative_exp + + +Adversarial Robustness +---------------------- +Riemannian projected gradient descent (PGD) attack for SPD-matrix inputs. + +.. autosummary:: + :toctree: generated/attacks_functional + + spd_rpgd_attack + + Vectorization Utilities ----------------------- Vectorization helpers for batching, (un)vectorizing matrices, and symmetric matrix @@ -237,10 +260,13 @@ Functions for Riemannian batch normalization computations on SPD manifolds. .. autosummary:: :toctree: generated/batchnorm_functional + frechet_mean karcher_mean_iteration spd_centering spd_rebiasing + spd_cholesky_congruence tangent_space_variance + lie_group_variance Bilinear Operations @@ -285,6 +311,146 @@ for detailed guidance. recommend_dtype_for_spd +pyRiemann Geometry (re-exported) +-------------------------------- +Since pyRiemann (>= 0.12) is a core dependency and runs on PyTorch tensors via +the Python Array API, its broader geometry toolkit is re-exported from +:py:mod:`spd_learn.functional`, so the full SPD/Riemannian API is reachable from +a single namespace. These run on torch tensors (with autograd) like the native +functions above. + +.. note:: + + The low-level matrix-function primitives (``logm``/``expm``/``sqrtm``/ + ``invsqrtm``/``powm``/``ddlogm``/``ddexpm``) are intentionally **not** + re-exported. Use spd_learn's numerically-stable :func:`matrix_log`, + :func:`matrix_exp`, :func:`matrix_sqrt`, ... and :func:`frechet_derivative_log` / + :func:`frechet_derivative_exp`, which keep gradients stable near degenerate + eigenvalues. The re-exported distances/means/maps inherit pyRiemann's generic + ``eigh`` backward (see :doc:`/numerical_stability`). + +Additional distances: + +.. autosummary:: + :toctree: generated/pyriemann_distance + + distance + distance_chol + distance_euclid + distance_harmonic + distance_kullback + distance_kullback_right + distance_kullback_sym + distance_logchol + distance_logdet + distance_logeuclid + distance_mahalanobis + distance_poweuclid + distance_riemann + distance_thompson + distance_wasserstein + pairwise_distance + +Additional geodesics: + +.. autosummary:: + :toctree: generated/pyriemann_geodesic + + geodesic + geodesic_chol + geodesic_euclid + geodesic_logchol + geodesic_logeuclid + geodesic_riemann + geodesic_thompson + geodesic_wasserstein + +Additional means and barycenters: + +.. autosummary:: + :toctree: generated/pyriemann_mean + + gmean + maskedmean_riemann + mean_ale + mean_alm + mean_bmp + mean_cheap + mean_chol + mean_euclid + mean_harmonic + mean_kullback_sym + mean_logchol + mean_logdet + mean_logeuclid + mean_power + mean_poweuclid + mean_riemann + mean_thompson + mean_wasserstein + nanmean_riemann + +Tangent space, inner products, and transport: + +.. autosummary:: + :toctree: generated/pyriemann_tangent + + exp_map + exp_map_euclid + exp_map_logchol + exp_map_logeuclid + exp_map_riemann + exp_map_wasserstein + innerproduct + innerproduct_euclid + innerproduct_logchol + innerproduct_logeuclid + innerproduct_riemann + log_map + log_map_euclid + log_map_logchol + log_map_logeuclid + log_map_riemann + log_map_wasserstein + norm + tangent_space + transport + transport_euclid + transport_logchol + transport_logeuclid + transport_riemann + untangent_space + unupper + upper + +Kernels: + +.. autosummary:: + :toctree: generated/pyriemann_kernel + + kernel + kernel_euclid + kernel_logeuclid + kernel_riemann + +Approximate joint diagonalization: + +.. autosummary:: + :toctree: generated/pyriemann_ajd + + ajd + ajd_pham + rjd + uwedge + +Matrix utilities: + +.. autosummary:: + :toctree: generated/pyriemann_base + + nearest_sym_pos_def + + Modules ======= @@ -449,6 +615,17 @@ building blocks from :py:mod:`spd_learn.modules`. TSMNet +Model Building Blocks +--------------------- +Sub-modules used inside the models above, exposed for reuse and customization. + +.. autosummary:: + :toctree: generated/ + + AttentionManifold + PhaseDelay + + Model Selection Guide --------------------- diff --git a/docs/source/faq.rst b/docs/source/faq.rst index 3c5315f..b516243 100644 --- a/docs/source/faq.rst +++ b/docs/source/faq.rst @@ -67,8 +67,24 @@ How is SPD Learn different from pyRiemann? - GPU acceleration via PyTorch - Integration with modern DL frameworks (Braindecode, skorch) -They are **complementary**: you can use pyRiemann for preprocessing and -SPD Learn for deep learning. +Since pyRiemann (>= 0.12) adopted the Python Array API, the two are no longer +merely complementary — **SPD Learn builds on pyRiemann**. pyRiemann is a core +dependency that runs natively on PyTorch tensors with autograd; SPD Learn +delegates the geometry it does not need to keep numerically special — geodesics, +parallel transport, and Fréchet derivatives — to it, and also re-exports +pyRiemann's broader geometry toolkit from :py:mod:`spd_learn.functional`. On top +of that shared backend SPD Learn adds: + +- ``nn.Module`` layers (BiMap, ReEig, LogEig, SPD batch norm, LieBN, ...) +- Numerically-stable, custom-autograd matrix-function primitives, distances, and + means that keep gradients well-behaved (and forwards finite) near clustered or + near-singular eigenvalues — where the generic ``eigh`` backward and unclamped + ``log`` are unstable — essential for training SPDNets. + +So pyRiemann remains the place to look for the broader classical toolkit +(MDM/FgMDM classifiers, tangent-space pipelines, ``pyriemann.datasets``, +estimators, and many more metrics) — all of which interoperate with SPD Learn +on the same PyTorch tensors. Technical Questions diff --git a/docs/source/installation.rst b/docs/source/installation.rst index 4ea216b..96e7363 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -25,6 +25,24 @@ Dependencies - `PyTorch `_ - Deep learning framework - `einops `_ - Flexible tensor operations - `NumPy `_ - Numerical computing +- `pyRiemann `_ (>= 0.12) - SPD/Riemannian + geometry. SPD Learn delegates geodesics, parallel transport, and Fréchet + derivatives to pyRiemann's Array API backend (which runs natively on PyTorch + tensors) and re-exports pyRiemann's broader toolkit. It keeps its own + numerically-stable, custom-autograd matrix-function primitives + (``matrix_log``/``matrix_exp``/``matrix_sqrt``/...), distances, and means, + which stay finite/stable on ill-conditioned inputs — important for training. + +.. note:: + + The Array API (PyTorch) backend lives in pyRiemann **0.12+**, which is not + yet on PyPI. SPD Learn therefore pins pyRiemann to its git ``master``:: + + pyriemann @ git+https://github.com/pyRiemann/pyRiemann.git + + Because this is a direct-URL dependency, ``pip install spd_learn`` from PyPI + is unavailable until pyRiemann 0.12 is released — install from source (see + below) meanwhile. The pin becomes ``pyriemann>=0.12`` once 0.12 ships. Installing from Source ---------------------- @@ -55,7 +73,8 @@ This installs: - `Braindecode `_ - Deep learning for EEG - `Nilearn `_ - Machine learning for neuroimaging -- `pyRiemann `_ - Riemannian geometry for BCI + +(pyRiemann is now a **core** dependency — see above — so it is always installed.) Documentation ^^^^^^^^^^^^^ diff --git a/pyproject.toml b/pyproject.toml index 0f1aebc..289c200 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,11 @@ dependencies = [ "torch", "einops", "numpy", + # Array API support (torch backend) lives on pyriemann master / 0.12. + # Hard-pinned to git until 0.12 is released on PyPI. NOTE: a direct-URL + # dependency blocks uploading spd_learn to PyPI; swap to "pyriemann>=0.12" + # once 0.12 ships. + "pyriemann @ git+https://github.com/pyRiemann/pyRiemann.git", ] [project.optional-dependencies] @@ -52,7 +57,6 @@ tests = [ 'codecov', 'pytest_cases', 'mypy', - 'pyriemann', ] docs = [ 'sphinx', @@ -75,7 +79,6 @@ brain = [ 'moabb', 'braindecode', 'nilearn', - 'pyriemann', 'skada', ] dev = [ diff --git a/spd_learn/functional/__init__.py b/spd_learn/functional/__init__.py index 0182bac..510bf8e 100644 --- a/spd_learn/functional/__init__.py +++ b/spd_learn/functional/__init__.py @@ -12,8 +12,104 @@ - **Regularization**: Covariance regularization utilities - **Vectorization**: Batch vectorization and (un)vectorization helpers - **Numerical**: Numerical stability configuration + +pyriemann (>=0.12) is a core dependency, and its broader geometry toolkit is +re-exported here so the full SPD/Riemannian API is available from a single +namespace. The low-level matrix-function primitives ``logm``/``expm``/``sqrtm``/ +``invsqrtm``/``powm`` are intentionally **not** re-exported: use spd_learn's +numerically-stable :func:`matrix_log`/:func:`matrix_exp`/... instead, which keep +gradients well-behaved near degenerate eigenvalues. ``ddlogm``/``ddexpm`` are +likewise not re-exported under those names; use the equivalent thin wrappers +:func:`frechet_derivative_log`/:func:`frechet_derivative_exp`. """ +from pyriemann.geometry.ajd import ajd, ajd_pham, rjd, uwedge +from pyriemann.geometry.base import nearest_sym_pos_def +from pyriemann.geometry.distance import ( + distance, + distance_chol, + distance_euclid, + distance_harmonic, + distance_kullback, + distance_kullback_right, + distance_kullback_sym, + distance_logchol, + distance_logdet, + distance_logeuclid, + distance_mahalanobis, + distance_poweuclid, + distance_riemann, + distance_thompson, + distance_wasserstein, + pairwise_distance, +) +from pyriemann.geometry.geodesic import ( + geodesic, + geodesic_chol, + geodesic_euclid, + geodesic_logchol, + geodesic_logeuclid, + geodesic_riemann, + geodesic_thompson, + geodesic_wasserstein, +) +from pyriemann.geometry.kernel import ( + kernel, + kernel_euclid, + kernel_logeuclid, + kernel_riemann, +) +from pyriemann.geometry.mean import ( + gmean, + maskedmean_riemann, + mean_ale, + mean_alm, + mean_bmp, + mean_cheap, + mean_chol, + mean_euclid, + mean_harmonic, + mean_kullback_sym, + mean_logchol, + mean_logdet, + mean_logeuclid, + mean_power, + mean_poweuclid, + mean_riemann, + mean_thompson, + mean_wasserstein, + nanmean_riemann, +) +from pyriemann.geometry.tangentspace import ( + exp_map, + exp_map_euclid, + exp_map_logchol, + exp_map_logeuclid, + exp_map_riemann, + exp_map_wasserstein, + innerproduct, + innerproduct_euclid, + innerproduct_logchol, + innerproduct_logeuclid, + innerproduct_riemann, + log_map, + log_map_euclid, + log_map_logchol, + log_map_logeuclid, + log_map_riemann, + log_map_wasserstein, + norm, + tangent_space, + transport, + transport_euclid, + transport_logchol, + transport_logeuclid, + transport_riemann, + untangent_space, + unupper, + upper, +) + from .autograd import modeig_backward, modeig_forward from .batchnorm import ( frechet_mean, @@ -194,4 +290,90 @@ "recommend_dtype_for_spd", # Attacks "spd_rpgd_attack", + # ----- Re-exported pyriemann geometry (Array API; run on torch tensors) ----- + "nearest_sym_pos_def", + # Additional distances + "distance", + "distance_chol", + "distance_euclid", + "distance_harmonic", + "distance_kullback", + "distance_kullback_right", + "distance_kullback_sym", + "distance_logchol", + "distance_logdet", + "distance_logeuclid", + "distance_mahalanobis", + "distance_poweuclid", + "distance_riemann", + "distance_thompson", + "distance_wasserstein", + "pairwise_distance", + # Additional geodesics + "geodesic", + "geodesic_chol", + "geodesic_euclid", + "geodesic_logchol", + "geodesic_logeuclid", + "geodesic_riemann", + "geodesic_thompson", + "geodesic_wasserstein", + # Additional means + "gmean", + "maskedmean_riemann", + "mean_ale", + "mean_alm", + "mean_bmp", + "mean_cheap", + "mean_chol", + "mean_euclid", + "mean_harmonic", + "mean_kullback_sym", + "mean_logchol", + "mean_logdet", + "mean_logeuclid", + "mean_power", + "mean_poweuclid", + "mean_riemann", + "mean_thompson", + "mean_wasserstein", + "nanmean_riemann", + # Tangent space, inner products, transport (generic) + "exp_map", + "exp_map_euclid", + "exp_map_logchol", + "exp_map_logeuclid", + "exp_map_riemann", + "exp_map_wasserstein", + "innerproduct", + "innerproduct_euclid", + "innerproduct_logchol", + "innerproduct_logeuclid", + "innerproduct_riemann", + "log_map", + "log_map_euclid", + "log_map_logchol", + "log_map_logeuclid", + "log_map_riemann", + "log_map_wasserstein", + "norm", + "tangent_space", + "transport", + "transport_euclid", + "transport_logchol", + "transport_logeuclid", + "transport_riemann", + "untangent_space", + "unupper", + "upper", + # Kernels + "kernel", + "kernel_euclid", + "kernel_logeuclid", + "kernel_riemann", + # Approximate joint diagonalization + "ajd", + "ajd_pham", + "rjd", + "uwedge", ] diff --git a/spd_learn/functional/frechet.py b/spd_learn/functional/frechet.py index 92df5fa..02523d0 100644 --- a/spd_learn/functional/frechet.py +++ b/spd_learn/functional/frechet.py @@ -2,47 +2,23 @@ # SPDX-License-Identifier: BSD-3-Clause r"""Frechet (directional) derivatives of matrix functions. -This module provides Frechet derivatives of the matrix logarithm and matrix -exponential, computed via the Daleckii-Krein theorem. These are used for -non-trivial parallel transport under the Log-Euclidean metric. - -The Frechet derivative :math:`Df(P)[V]` of a matrix function :math:`f` at -:math:`P` in direction :math:`V` is computed as: - -.. math:: - - Df(P)[V] = U (K \odot (U^T V U)) U^T - -where :math:`P = U \Lambda U^T` is the eigendecomposition and :math:`K` is -the Loewner matrix of divided differences of :math:`f`. +Thin wrappers over pyriemann's :func:`~pyriemann.geometry.base.ddlogm` / +:func:`~pyriemann.geometry.base.ddexpm` (Array API, run on torch tensors with +autograd), which compute the Frechet derivatives of the matrix logarithm and +exponential via the Daleckii-Krein theorem. Used for non-trivial parallel +transport under the Log-Euclidean metric. """ -import torch - -from .numerical import get_loewner_threshold -from .utils import ensure_sym +from pyriemann.geometry.base import ddexpm, ddlogm def frechet_derivative_log(P, V): r"""Frechet derivative of the matrix logarithm at P in direction V. - Computes :math:`D\log(P)[V]` using the Daleckii-Krein theorem: - - .. math:: - - D\log(P)[V] = U (K \odot (U^T V U)) U^T - - where :math:`P = U \operatorname{diag}(\lambda) U^T` and the coefficient - matrix :math:`K` has entries: - - .. math:: - - K_{ij} = \begin{cases} - \frac{\log \lambda_i - \log \lambda_j}{\lambda_i - \lambda_j} - & \text{if } \lambda_i \neq \lambda_j \\ - \frac{1}{\lambda_i} - & \text{if } \lambda_i = \lambda_j - \end{cases} + Computes :math:`D\log(P)[V]`, the directional derivative of the matrix + logarithm at SPD matrix :math:`P` in direction :math:`V`. Delegates to + :func:`pyriemann.geometry.base.ddlogm`, which evaluates the Daleckii-Krein + (Loewner divided-difference) form on the Array API backend. Parameters ---------- @@ -56,60 +32,19 @@ def frechet_derivative_log(P, V): torch.Tensor The Frechet derivative :math:`D\log(P)[V]`, shape (..., n, n). """ - # Eigendecompose P - eigvals, U = torch.linalg.eigh(P) - - # Build the Loewner matrix K for log - log_eigvals = eigvals.log() - - # Denominator: lambda_i - lambda_j - denom = eigvals.unsqueeze(-1) - eigvals.unsqueeze(-2) - - # Adaptive threshold for detecting equal eigenvalues - threshold = get_loewner_threshold(eigvals) - is_eq = denom.abs() < threshold - denom_safe = denom.clone() - denom_safe[is_eq] = 1.0 - - # Numerator: log(lambda_i) - log(lambda_j) - numerator = log_eigvals.unsqueeze(-1) - log_eigvals.unsqueeze(-2) - - # K_ij = (log(lambda_i) - log(lambda_j)) / (lambda_i - lambda_j) - K = numerator / denom_safe - - # Diagonal (and near-equal): K_ii = 1 / lambda_i - diag_vals = 1.0 / eigvals - # For equal eigenvalues, use average of the diagonal values - K[is_eq] = 0.5 * (diag_vals.unsqueeze(-1) + diag_vals.unsqueeze(-2))[is_eq] - - # Apply: U @ (K * (U^T V U)) @ U^T - Ut = U.transpose(-2, -1) - inner = Ut @ V @ U - result = U @ (K * inner) @ Ut - - return ensure_sym(result) + # Delegated to pyriemann (Array API). pyriemann's ddlogm(X, Cref) computes + # the directional derivative of logm at Cref in direction X. + return ddlogm(V, P) def frechet_derivative_exp(X, W): r"""Frechet derivative of the matrix exponential at X in direction W. - Computes :math:`D\exp(X)[W]` using the Daleckii-Krein theorem: - - .. math:: - - D\exp(X)[W] = U (M \odot (U^T W U)) U^T - - where :math:`X = U \operatorname{diag}(d) U^T` (X is symmetric, not - necessarily SPD) and the coefficient matrix :math:`M` has entries: - - .. math:: - - M_{ij} = \begin{cases} - \frac{e^{d_i} - e^{d_j}}{d_i - d_j} - & \text{if } d_i \neq d_j \\ - e^{d_i} - & \text{if } d_i = d_j - \end{cases} + Computes :math:`D\exp(X)[W]`, the directional derivative of the matrix + exponential at symmetric matrix :math:`X` (need not be SPD) in direction + :math:`W`. Delegates to :func:`pyriemann.geometry.base.ddexpm`, which + evaluates the Daleckii-Krein (Loewner divided-difference) form on the + Array API backend. Parameters ---------- @@ -123,34 +58,6 @@ def frechet_derivative_exp(X, W): torch.Tensor The Frechet derivative :math:`D\exp(X)[W]`, shape (..., n, n). """ - # Eigendecompose X - eigvals, U = torch.linalg.eigh(X) - - # Build the Loewner matrix M for exp - exp_eigvals = eigvals.exp() - - # Denominator: d_i - d_j - denom = eigvals.unsqueeze(-1) - eigvals.unsqueeze(-2) - - # Adaptive threshold for detecting equal eigenvalues - threshold = get_loewner_threshold(eigvals) - is_eq = denom.abs() < threshold - denom_safe = denom.clone() - denom_safe[is_eq] = 1.0 - - # Numerator: exp(d_i) - exp(d_j) - numerator = exp_eigvals.unsqueeze(-1) - exp_eigvals.unsqueeze(-2) - - # M_ij = (exp(d_i) - exp(d_j)) / (d_i - d_j) - M = numerator / denom_safe - - # Diagonal (and near-equal): M_ii = exp(d_i) - # For equal eigenvalues, use average of exp values - M[is_eq] = 0.5 * (exp_eigvals.unsqueeze(-1) + exp_eigvals.unsqueeze(-2))[is_eq] - - # Apply: U @ (M * (U^T W U)) @ U^T - Ut = U.transpose(-2, -1) - inner = Ut @ W @ U - result = U @ (M * inner) @ Ut - - return ensure_sym(result) + # Delegated to pyriemann (Array API). pyriemann's ddexpm(X, Cref) computes + # the directional derivative of expm at Cref in direction X. + return ddexpm(W, X) diff --git a/spd_learn/functional/metrics/affine_invariant.py b/spd_learn/functional/metrics/affine_invariant.py index 0236276..09f4301 100644 --- a/spd_learn/functional/metrics/affine_invariant.py +++ b/spd_learn/functional/metrics/affine_invariant.py @@ -55,13 +55,9 @@ import torch -from ..core import ( - matrix_exp, - matrix_inv_sqrt, - matrix_log, - matrix_power, - matrix_sqrt_inv, -) +from pyriemann.geometry.geodesic import geodesic_riemann + +from ..core import matrix_exp, matrix_inv_sqrt, matrix_log, matrix_sqrt_inv from ..numerical import get_epsilon from ..utils import ensure_sym @@ -120,6 +116,8 @@ def airm_distance(A, B): ---------- See :cite:p:`pennec2006riemannian`, :cite:p:`bhatia2007positive` for more details. """ + # Kept in spd_learn (not delegated): clamps eigenvalues for stability on + # ill-conditioned inputs, which pyriemann's distance_riemann does not. Ainvsqrt = matrix_inv_sqrt.apply(A) eigenvalues = torch.linalg.eigvalsh(Ainvsqrt @ B @ Ainvsqrt) threshold = get_epsilon(eigenvalues.dtype, "eigval_log") @@ -184,12 +182,12 @@ def airm_geodesic(A, B, t): ---------- See :cite:p:`pennec2006riemannian`, :cite:p:`bhatia2007positive` for more details. """ - rm_sq, rm_invsq = matrix_sqrt_inv.apply(A) - return ( - rm_sq - @ matrix_power.apply(rm_invsq @ B @ rm_invsq, torch.tensor(t).to(A)) - @ rm_sq - ) + # A 0-d tensor t broadcasts over batched inputs (pyriemann requires alpha to + # match the batch shape); expand it so a scalar t still works on batches. + if torch.is_tensor(t) and t.ndim == 0 and A.ndim > 2: + t = t.expand(A.shape[:-2]) + # Delegated to pyriemann (Array API, runs on torch tensors with autograd). + return geodesic_riemann(A, B, alpha=t) def exp_map_airm(P, V, t=1.0): diff --git a/spd_learn/functional/metrics/bures_wasserstein.py b/spd_learn/functional/metrics/bures_wasserstein.py index 1b96487..799e504 100644 --- a/spd_learn/functional/metrics/bures_wasserstein.py +++ b/spd_learn/functional/metrics/bures_wasserstein.py @@ -70,6 +70,7 @@ import torch +from pyriemann.geometry.geodesic import geodesic_wasserstein from torch.autograd import Function from ..core import matrix_inv_sqrt, matrix_sqrt, matrix_sqrt_inv @@ -231,6 +232,8 @@ def bures_wasserstein_distance(A, B): :func:`~spd_learn.functional.log_euclidean_distance` : Distance under Log-Euclidean metric. :func:`~spd_learn.functional.log_cholesky_distance` : Distance under Log-Cholesky metric. """ + # Kept in spd_learn (not delegated): custom autograd with an analytic + # backward and a subgradient-0 guard at distance 0 (A == B). return _BuresWassersteinDistanceFunction.apply(A, B) @@ -291,38 +294,12 @@ def bures_wasserstein_geodesic(A, B, t): :func:`~spd_learn.functional.airm_geodesic` : Geodesic under AIRM. :func:`~spd_learn.functional.log_cholesky_geodesic` : Geodesic under Log-Cholesky metric. """ - # Ensure t is a tensor - if not torch.is_tensor(t): - t = torch.tensor(t, dtype=A.dtype, device=A.device) - else: - t = t.to(dtype=A.dtype, device=A.device) - - # Handle edge cases exactly for scalar t - if t.numel() == 1: - t_item = t.item() - if t_item == 0.0: - return A.clone() - if t_item == 1.0: - return B.clone() - - # Expand t for broadcasting - t = t.reshape(t.shape + (1,) * 2) - - # Compute A^{1/2} and A^{-1/2} - A_sqrt, A_inv_sqrt = matrix_sqrt_inv.apply(A) - - # Compute M = (A^{1/2} B A^{1/2})^{1/2} - ABA = A_sqrt @ B @ A_sqrt - M = matrix_sqrt.apply(ABA) - - # (AB)^{1/2} = A^{1/2} M A^{-1/2} - AB12 = A_sqrt @ M @ A_inv_sqrt - - # Geodesic: (1-t)^2 A + t^2 B + t(1-t)((AB)^{1/2} + (BA)^{1/2}) - result = (1 - t) ** 2 * A + t**2 * B + t * (1 - t) * (AB12 + AB12.transpose(-1, -2)) - - # Ensure symmetry - return ensure_sym(result) + # A 0-d tensor t broadcasts over batched inputs (pyriemann requires alpha to + # match the batch shape); expand it so a scalar t still works on batches. + if torch.is_tensor(t) and t.ndim == 0 and A.ndim > 2: + t = t.expand(A.shape[:-2]) + # Delegated to pyriemann (Array API, runs on torch tensors with autograd). + return geodesic_wasserstein(A, B, alpha=t) def bures_wasserstein_mean( diff --git a/spd_learn/functional/metrics/log_cholesky.py b/spd_learn/functional/metrics/log_cholesky.py index 9f4b697..aa3ce70 100644 --- a/spd_learn/functional/metrics/log_cholesky.py +++ b/spd_learn/functional/metrics/log_cholesky.py @@ -78,6 +78,7 @@ import torch +from pyriemann.geometry.geodesic import geodesic_logchol from torch.autograd import Function from ..numerical import get_epsilon @@ -348,6 +349,8 @@ def log_cholesky_distance(A, B): >>> print(f"Distance: {dist.item():.4f}") Distance: 1.2012 """ + # Kept in spd_learn (not delegated): built on spd_learn's stable cholesky_log + # so all distances share one numerically-robust code path. log_chol_A = cholesky_log.apply(A) log_chol_B = cholesky_log.apply(B) @@ -480,10 +483,9 @@ def log_cholesky_geodesic(A, B, t): >>> print(f"Midpoint diagonal: {torch.diag(mid)}") Midpoint diagonal: tensor([2., 2., 2.]) """ - log_chol_A = cholesky_log.apply(A) - log_chol_B = cholesky_log.apply(B) - - # Linear interpolation in Log-Cholesky space - log_chol_interp = (1 - t) * log_chol_A + t * log_chol_B - - return cholesky_exp.apply(log_chol_interp) + # A 0-d tensor t broadcasts over batched inputs (pyriemann requires alpha to + # match the batch shape); expand it so a scalar t still works on batches. + if torch.is_tensor(t) and t.ndim == 0 and A.ndim > 2: + t = t.expand(A.shape[:-2]) + # Delegated to pyriemann (Array API, runs on torch tensors with autograd). + return geodesic_logchol(A, B, alpha=t) diff --git a/spd_learn/functional/metrics/log_euclidean.py b/spd_learn/functional/metrics/log_euclidean.py index a534fc9..7d8b96f 100644 --- a/spd_learn/functional/metrics/log_euclidean.py +++ b/spd_learn/functional/metrics/log_euclidean.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: BSD-3-Clause import torch +from pyriemann.geometry.geodesic import geodesic_logeuclid + from ..core import matrix_exp, matrix_log @@ -54,6 +56,8 @@ def log_euclidean_distance(A, B): ---------- See :cite:p:`arsigny2007geometric` for more details. """ + # Kept in spd_learn (not delegated): uses spd_learn's stable, eigenvalue- + # clamped matrix_log, robust on ill-conditioned inputs. inner_term = matrix_log.apply(A) - matrix_log.apply(B) final = torch.linalg.norm(inner_term, ord="fro", dim=(-2, -1)) return final @@ -141,9 +145,12 @@ def log_euclidean_geodesic(A, B, t): ---------- See :cite:p:`arsigny2007geometric` for more details. """ - log_A = matrix_log.apply(A) - log_B = matrix_log.apply(B) - return matrix_exp.apply((1 - t) * log_A + t * log_B) + # A 0-d tensor t broadcasts over batched inputs (pyriemann requires alpha to + # match the batch shape); expand it so a scalar t still works on batches. + if torch.is_tensor(t) and t.ndim == 0 and A.ndim > 2: + t = t.expand(A.shape[:-2]) + # Delegated to pyriemann (Array API, runs on torch tensors with autograd). + return geodesic_logeuclid(A, B, alpha=t) def exp_map_lem(P, V): diff --git a/spd_learn/functional/parallel_transport.py b/spd_learn/functional/parallel_transport.py index e8c82a9..7160d74 100644 --- a/spd_learn/functional/parallel_transport.py +++ b/spd_learn/functional/parallel_transport.py @@ -50,66 +50,17 @@ details on parallel transport and numerical approximation methods. """ -import torch - -from .core import ( - matrix_inv_sqrt, - matrix_log, - matrix_sqrt, - matrix_sqrt_inv, +from pyriemann.geometry.tangentspace import ( + transport_logchol, + transport_logeuclid, + transport_riemann, ) -from .frechet import frechet_derivative_exp, frechet_derivative_log + +from .core import matrix_sqrt, matrix_sqrt_inv from .metrics.affine_invariant import exp_map_airm, log_map_airm from .utils import ensure_sym -def _parallel_transport_airm_functional(v, p, q): - r"""Parallel transport under AIRM with full autograd support. - - This functional implementation leverages PyTorch's autograd by using - operations that already have proper backward passes defined. This allows - gradients to flow through all inputs (v, p, q). - - **Formula Relationship** - - The transport operator is :math:`E = (Q P^{-1})^{1/2}` (principal square - root). Since :math:`Q P^{-1}` is non-symmetric when :math:`P \neq Q`, we - use the equivalent stable formula: - - .. math:: - - E = Q^{1/2} (Q^{-1/2} P Q^{-1/2})^{-1/2} Q^{-1/2} - - This computes only symmetric matrix square roots, which are well-defined - for SPD matrices. The equivalence follows from :math:`E^2 = Q P^{-1}`, - which can be verified by direct computation (see :func:`parallel_transport_airm`). - - Parameters - ---------- - v : torch.Tensor - Tangent vector at p, shape (..., n, n). Must be symmetric. - p : torch.Tensor - Source point on SPD manifold, shape (..., n, n). - q : torch.Tensor - Target point on SPD manifold, shape (..., n, n). - - Returns - ------- - torch.Tensor - Transported tangent vector at q, shape (..., n, n). - """ - # Compute E = Q^{1/2} @ (Q^{-1/2} @ P @ Q^{-1/2})^{-1/2} @ Q^{-1/2} - # This is the principal square root of Q @ P^{-1} - q_sqrt, q_invsqrt = matrix_sqrt_inv.apply(q) - inner = q_invsqrt @ p @ q_invsqrt - inner_invsqrt = matrix_inv_sqrt.apply(inner) - E = q_sqrt @ inner_invsqrt @ q_invsqrt - - # Transport: v' = E @ v @ E^T - v_transported = E @ v @ E.transpose(-2, -1) - return ensure_sym(v_transported) - - def parallel_transport_airm(v, p, q): r"""Parallel transport of tangent vector under the Affine-Invariant metric. @@ -220,7 +171,8 @@ def parallel_transport_airm(v, p, q): :func:`~spd_learn.functional.airm_distance` : Distance under AIRM. :class:`~spd_learn.modules.SPDBatchNormMeanVar` : Uses parallel transport for centering. """ - return _parallel_transport_airm_functional(v, p, q) + # Delegated to pyriemann (Array API, runs on torch tensors with autograd). + return transport_riemann(v, p, q) def parallel_transport_lem(v, p, q): @@ -240,7 +192,8 @@ def parallel_transport_lem(v, p, q): at :math:`P` and :math:`D\exp(\log Q)` is the Frechet derivative of the matrix exponential at :math:`\log Q`. The intermediate step maps the ambient tangent vector into the flat log-space, where transport is trivial, - then maps back to the ambient tangent space at :math:`Q`. + then maps back to the ambient tangent space at :math:`Q`. This is delegated + to :func:`pyriemann.geometry.tangentspace.transport_logeuclid`. Parameters ---------- @@ -288,12 +241,8 @@ def parallel_transport_lem(v, p, q): :func:`~spd_learn.functional.frechet.frechet_derivative_exp` : Frechet derivative of exp. :func:`~spd_learn.functional.log_euclidean_distance` : Distance under Log-Euclidean metric. """ - # Step 1: Map v from ambient T_P to log-space via D_log(P)[V] - w = frechet_derivative_log(p, v) - # Step 2: Map w from log-space back to ambient T_Q via D_exp(log(Q))[W] - log_q = matrix_log.apply(q) - v_transported = frechet_derivative_exp(log_q, w) - return ensure_sym(v_transported) + # Delegated to pyriemann (Array API, runs on torch tensors with autograd). + return transport_logeuclid(v, p, q) def parallel_transport_log_cholesky(v, p, q): @@ -371,46 +320,8 @@ def parallel_transport_log_cholesky(v, p, q): :func:`~spd_learn.functional.log_cholesky_distance` : Distance under Log-Cholesky metric. :func:`~spd_learn.functional.log_cholesky_mean` : Mean under Log-Cholesky metric. """ - # Cholesky decompose P and Q - L_P = torch.linalg.cholesky(p) - L_Q = torch.linalg.cholesky(q) - - # Step 1: Pull V back to the Cholesky tangent space - # Compute S = L_P^{-1} V L_P^{-T} (symmetric) - # First: X = L_P^{-1} V (solve L_P X = V) - X = torch.linalg.solve_triangular(L_P, v, upper=False) - # Then: S = X L_P^{-T} (solve L_P S^T = X^T, then transpose) - S = torch.linalg.solve_triangular(L_P, X.transpose(-2, -1), upper=False).transpose( - -2, -1 - ) - - # B = strictly_lower(S) + 0.5 * diag(S) - B = S.tril(-1) + torch.diagonal(S, dim1=-2, dim2=-1).unsqueeze( - -2 - ) * 0.5 * torch.eye(S.shape[-1], dtype=S.dtype, device=S.device) - - # dL = L_P @ B (lower triangular) - dL = L_P @ B - - # Step 2: Convert to log-Cholesky tangent coordinates - diag_L_P = torch.diagonal(L_P, dim1=-2, dim2=-1) - diag_dL = torch.diagonal(dL, dim1=-2, dim2=-1) - # dY = strictly_lower(dL) + diag(diag(dL) / diag(L_P)) - dY = dL.tril(-1) + torch.diag_embed(diag_dL / diag_L_P) - - # Step 3: Transport in flat log-Cholesky space is identity (dY stays the same) - - # Step 4: Convert back at Q - diag_L_Q = torch.diagonal(L_Q, dim1=-2, dim2=-1) - diag_dY = torch.diagonal(dY, dim1=-2, dim2=-1) - # dL_Q = strictly_lower(dY) + diag(diag(dY) * diag(L_Q)) - dL_Q = dY.tril(-1) + torch.diag_embed(diag_dY * diag_L_Q) - - # Step 5: Push forward to ambient tangent space at Q - # V' = dL_Q @ L_Q^T + L_Q @ dL_Q^T - v_transported = dL_Q @ L_Q.transpose(-2, -1) + L_Q @ dL_Q.transpose(-2, -1) - - return ensure_sym(v_transported) + # Delegated to pyriemann (Array API, runs on torch tensors with autograd). + return transport_logchol(v, p, q) def _geodesic_midpoint_airm(p, q): diff --git a/spd_learn/models/__init__.py b/spd_learn/models/__init__.py index d0e7fc2..611057e 100644 --- a/spd_learn/models/__init__.py +++ b/spd_learn/models/__init__.py @@ -2,13 +2,22 @@ # SPDX-License-Identifier: BSD-3-Clause from .eegspdnet import EEGSPDNet from .green import Green -from .matt import MAtt -from .phase_spdnet import PhaseSPDNet +from .matt import AttentionManifold, MAtt # noqa: F401 (AttentionManifold re-exported) +from .phase_spdnet import ( # noqa: F401 (PhaseDelay re-exported) + PhaseDelay, + PhaseSPDNet, +) from .spdnet import SPDNet from .tensorcsp import TensorCSPNet from .tsmnet import TSMNet +# ``AttentionManifold`` and ``PhaseDelay`` are model building blocks: importable +# from ``spd_learn.models`` and documented in api.rst, but intentionally kept out +# of ``__all__`` -- they are sub-modules of the models above (different +# constructors, not standalone ``n_chans``/``n_outputs`` models), while +# ``__all__`` is the contract the model test harness iterates to instantiate +# full models. __all__ = [ "TensorCSPNet", "SPDNet", diff --git a/tests/test_pyriemann_comparison.py b/tests/test_pyriemann_comparison.py index a127ec2..028a46e 100644 --- a/tests/test_pyriemann_comparison.py +++ b/tests/test_pyriemann_comparison.py @@ -1,9 +1,20 @@ -"""Comparison tests between spd_learn (PyTorch) and pyriemann (NumPy). - -Validates that spd_learn's implementations produce numerically equivalent -results to pyriemann's established NumPy implementations for all shared -SPD matrix operations: matrix functions, distances, geodesics, means, -and log/exp maps. +"""Comparison tests between spd_learn (PyTorch) and pyriemann. + +As of the pyriemann (>=0.12) Array API migration, spd_learn delegates most of +its SPD geometry math to pyriemann (matrix functions, distances, geodesics, +several means, tangent maps, transport, Frechet derivatives). These tests +therefore serve two purposes: + +1. For *delegated* ops, they lock spd_learn's public API to pyriemann's + reference values (cross-backend: spd_learn on torch vs pyriemann on numpy). +2. For ops still implemented in spd_learn (``cholesky_log``/``cholesky_exp``, + the weighted ``log_euclidean_mean``/``bures_wasserstein_mean``, AIRM + ``exp_map``/``log_map``), they verify those remain numerically equivalent to + pyriemann. (The ``schild_ladder``/``pole_ladder`` approximations are also + kept in spd_learn but have no pyriemann counterpart to compare against.) + +Imports come from ``pyriemann.geometry.*`` (``pyriemann.utils.*`` is deprecated +and removed in pyriemann 0.14). """ import numpy as np @@ -13,26 +24,29 @@ pyriemann = pytest.importorskip("pyriemann") -from pyriemann.utils.base import expm, invsqrtm, logm, powm, sqrtm # noqa: E402 -from pyriemann.utils.distance import ( # noqa: E402 +from pyriemann.geometry.base import expm, invsqrtm, logm, powm, sqrtm # noqa: E402 +from pyriemann.geometry.distance import ( # noqa: E402 distance_logchol, distance_logeuclid, distance_riemann, distance_wasserstein, ) -from pyriemann.utils.geodesic import ( # noqa: E402 +from pyriemann.geometry.geodesic import ( # noqa: E402 geodesic_logchol, geodesic_logeuclid, geodesic_riemann, geodesic_wasserstein, ) -from pyriemann.utils.mean import ( # noqa: E402 +from pyriemann.geometry.mean import ( # noqa: E402 mean_logchol, mean_logeuclid, mean_riemann, mean_wasserstein, ) -from pyriemann.utils.tangentspace import exp_map_riemann, log_map_riemann # noqa: E402 +from pyriemann.geometry.tangentspace import ( # noqa: E402 + exp_map_riemann, + log_map_riemann, +) from spd_learn.functional import ( # noqa: E402 airm_distance, @@ -338,3 +352,36 @@ def test_batched_distances(n): np.testing.assert_allclose( result[i], distance_riemann(A_list[i], B_list[i]), **STRICT_TOL ) + + +# --------------------------------------------------------------------------- +# 7. pyriemann delegation: autograd behaviour (documented known ceiling) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("n", SIZES) +def test_delegated_gradients_finite_well_conditioned(n): + """Delegated ops backprop through pyriemann's eigh autograd. + + The *delegated* ops (geodesics, parallel transport, Fréchet derivatives) + inherit torch's ``eigh`` backward, and pyriemann does not clamp eigenvalues. + spd_learn's own matrix-function primitives (``matrix_log``/``matrix_exp``/...) + AND its distances are deliberately **not** delegated and keep their + numerically-stable, eigenvalue-clamped implementations. This test drives both + layers in one graph (``matrix_log`` = kept/stable, ``airm_geodesic`` = + delegated): for well-conditioned SPD inputs the gradient is finite. + Near-degenerate eigenvalues are a KNOWN ceiling of the *delegated* ops + (pyriemann's eigh backward scales as ``1 / (lambda_i - lambda_j)``); callers + needing stable gradients there should keep eigenvalues well separated + (e.g. via ReEig / ``clamp_eigvals``, also not delegated). + """ + rng = np.random.RandomState(SEED) + A = to_torch(make_spd_np(n, rng)).requires_grad_(True) + B = to_torch(make_spd_np(n, rng)) + + # matrix_log = kept (stable) primitive; airm_geodesic = delegated op. + loss = matrix_log.apply(A).pow(2).sum() + airm_geodesic(A, B, 0.5).pow(2).sum() + loss.backward() + + assert A.grad is not None + assert torch.isfinite(A.grad).all() From 2d5b31b22ec6f93daa37c1474091297e72015398 Mon Sep 17 00:00:00 2001 From: Bru Date: Fri, 26 Jun 2026 11:50:36 +0200 Subject: [PATCH 2/2] ci: cap pytest<9.1 (pytest_cases 3.10.1 incompatible with pytest 9.1 IdMaker) --- pyproject.toml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 289c200..b76819e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,10 @@ dependencies = [ [project.optional-dependencies] tests = [ - 'pytest', + # pytest 9.1 changed IdMaker's signature, which pytest_cases (<= 3.10.1) + # still calls the old way -> TypeError at collection time. Cap pytest until + # pytest_cases ships 9.1 support. + 'pytest<9.1', 'pytest-cov', 'pytest-codeblocks', 'sybil',