spd_learn.models.DeepGeoCCA#

class spd_learn.models.DeepGeoCCA(n_chans: int = 22, n_outputs: int = 2, bimap_dims: Sequence[int] = (16, 10), outdim_size: int = 10, eps: float = 0.0001, drop_prob: float = 0.0)[source]#

Bases: Module

Deep Geodesic Canonical Correlation Analysis (DeepGeoCCA).

This model implements the DeepGeoCCA architecture proposed in Ju et al. (2024) [Ju et al., 2024]. It performs multi-view learning on Symmetric Positive Definite (SPD) manifolds by maximizing the geodesic correlation between two multi-scale representations of multivariate time series (e.g., EEG signals).

Given a multi-channel signal, two view-specific temporal convolutions with complementary kernel sizes extract multi-scale temporal dynamics. Each view is then mapped to the SPD manifold via spatial covariance estimation, regularized with trace normalization and learnable shrinkage, and processed by stacked bilinear mapping (BiMap) and eigenvalue rectification (ReEig) layers, followed by Riemannian batch normalization (SPDBatchNormMean).

For classification, latent SPD representations from both views are mapped to Euclidean tangent vectors via LogEig, concatenated, and fed to a linear classifier. During self-supervised pre-training or multi-task optimization, GeodesicCCALoss maximizes the geodesic canonical correlation between the two views.

Parameters:

  • n_chans (int, default=22) – Number of input EEG channels.

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

  • bimap_dims (tuple of int, default=(16, 10)) – Sequence of output subspace dimensions for the BiMap layers in each encoder.

  • outdim_size (int, default=10) – Number of top canonical correlation dimensions to extract and maximize in GeodesicCCALoss.

  • eps (float, default=1e-4) – Regularization constant for numerical stability in eigenvalue clamping, covariance regularization, and CCA whitening.

  • drop_prob (float, default=0.0) – Dropout probability applied to the concatenated tangent features before the linear classifier.

compute_cca_loss(x1: Tensor | Sequence[Tensor], x2: Tensor | None = None) → Tensor[source]#

Compute the Geodesic CCA loss between the two views.

Parameters:

  • x1 (torch.Tensor or sequence of torch.Tensor) – Input EEG signals of shape (batch_size, n_chans, n_times) or paired views.

  • x2 (torch.Tensor, optional) – Second view if not provided in x1.

Returns:

Scalar Geodesic CCA loss (negative total canonical correlation).

Return type:

torch.Tensor

encode_views(x1: Tensor | Sequence[Tensor], x2: Tensor | None = None) → tuple[Tensor, Tensor][source]#

Encode raw signals into SPD representations for both views.

Parameters:

  • x1 (torch.Tensor or sequence of torch.Tensor) – If x2 is None, either a single 3D EEG tensor of shape (batch_size, n_chans, n_times) (which will be filtered into two multi-scale views via view1_conv and view2_conv), or a pair (v1, v2) of two view tensors. If x2 is provided, x1 is the first view tensor.

  • x2 (torch.Tensor, optional) – The second view tensor of shape (batch_size, n_chans, n_times).

Returns:

(Y1, Y2) where each is a batch of SPD matrices of shape (batch_size, d_out, d_out).

Return type:

tuple[torch.Tensor, torch.Tensor]

forward(x: Tensor | Sequence[Tensor], x2: Tensor | None = None) → Tensor[source]#

Forward pass of DeepGeoCCA.

Parameters:

  • x (torch.Tensor or sequence of torch.Tensor) – Input EEG signals of shape (batch_size, n_chans, n_times) or paired views (x1, x2).

  • x2 (torch.Tensor, optional) – Second view if not provided in x.

Returns:

Class logits of shape (batch_size, n_outputs).

Return type:

torch.Tensor