spd_learn.models.CorNet#

class spd_learn.models.CorNet(n_chans: int = 22, n_outputs: int = 2, bimap_dims: tuple[int, ...]=(16, 12), metric: str = 'OLM', threshold: float = 0.0001, drop_prob: float = 0.0, input_type: str = 'raw', cov_method=<function sample_covariance>, device: device | None = None, dtype: dtype | None = None)[source]#

Bases: Module

Riemannian Neural Network on the Correlation Manifold (CorNet).

This class implements the CorNet model [Chen et al., 2026], an architecture operating directly on the manifold of full-rank correlation matrices \(\mathcal{C}_{++}^n\).

The model consists of:

  1. Sample covariance estimation (CovLayer) and shrinkage regularization (Shrinkage).

  2. Conversion to correlation matrices via cov_to_corr().

  3. Stacked correlation blocks, each comprising a bilinear mapping (BiMap), correlation normalization (cov_to_corr()), rectified eigenvalue activation (ReEig), and correlation normalization.

  4. Riemannian tangent projection to hollow symmetric space (log_map_olm() or log_map_lsm()).

  5. Isometric lower-triangular embedding (tril_embed()).

  6. Dropout and linear classification head.

Parameters:

  • n_chans (int, default=22) – Number of input channels (matrix dimension before first BiMap).

  • n_outputs (int, default=2) – Number of output classes.

  • bimap_dims (tuple of int, default=(16, 12)) – Subspace dimensions for successive correlation BiMap layers.

  • metric ({"OLM", "LSM"}, default="OLM") – Riemannian metric on the correlation manifold used for tangent projection.

  • threshold (float, default=1e-4) – Eigenvalue rectification threshold for ReEig.

  • drop_prob (float, default=0.0) – Dropout probability before the final classification head.

  • input_type ({"raw", "cov"}, default="raw") – Format of input data (‘raw’ time series or precomputed ‘cov’ matrices).

  • cov_method (callable, default=sample_covariance) – Method to compute covariance when input_type is ‘raw’.

  • device (torch.device, optional) – Device on which to allocate parameters.

  • dtype (torch.dtype, optional) – Data type for model parameters.

forward(x: Tensor) → Tensor[source]#

Forward pass of CorNet.

Parameters:

x (torch.Tensor) – Input tensor. If input_type is ‘raw’, shape is (batch_size, n_chans, n_times). If input_type is ‘cov’, shape is (batch_size, n_chans, n_chans).

Returns:

Output class logits of shape (batch_size, n_outputs).

Return type:

torch.Tensor