Code for align dynamics using CCA

aopy/analysis/base.py

@@ -235,6 +235,45 @@ def calc_task_rel_dims(neural_data, kin_data, conc_proj_data=False):
     else:    
         return task_subspace.T, projected_data
 
+def align_latent_dynamics(La, Lb, return_aligned_dynamics=False):
+    """
+    Aligns latent dynamics using Canonical Correlation Analysis (CCA) and computes pairwise Pearson correlation for both aligned and unaligned dynamics.
+    References: Gallego, J. A., Perich, M. G., Chowdhury, R. H., Solla, S. A. & Miller, L. E. Long-term stability of cortical population dynamics underlying consistent behavior. Nat Neurosci 23, 260–270 (2020).
+
+    Args:
+        La (ndarray): Latent dynamics of Dataset A with shape (m, n_timepoints). Usually first dimension is time, however Juancho's code has it as (m, n_t). Keeping it similar to his paper is easier for the computations below.
+        Lb (ndarray): Latent dynamics of Dataset B with shape (m, n_timepoints).
+
+    Returns:
+        CCs_unaligned (ndarray): Pairwise Pearson correlation between unaligned latent dynamics (La and Lb) with shape (m).
+        CCs_aligned (ndarray): Pairwise Pearson correlation between aligned latent dynamics (La_tilde and Lb_tilde) with shape (m).
+    """
+     # Step 1: QR decomposition
+    Qa, Ra = np.linalg.qr(La.T)  # QR decomposition of La transpose
+    Qb, Rb = np.linalg.qr(Lb.T)  # QR decomposition of Lb transpose
+
+    # Step 2: Construct the cross covariance matrix and perform SVD
+    QaT_Qb = Qa.T @ Qb  # Inner product matrix of Qa and Qb
+    U, S, Vt = np.linalg.svd(QaT_Qb)  # Singular value decomposition of QaT_Qb
+
+    # Step 3: Calculate projection matrices
+    Ma = np.linalg.pinv(Ra) @ U  # Projection matrix for La
+    Mb = np.linalg.pinv(Rb) @ Vt.T  # Projection matrix for Lb
+
+    # Step 4: Project latent dynamics onto new manifold axes
+    La_tilde = La.T @ Ma  # Latent dynamics projected onto new manifold axes for La
+    Lb_tilde = Lb.T @ Mb  # Latent dynamics projected onto new manifold axes for Lb
+
+    # Step 5: Calculate pairwise correlations between unaligned and aligned latent dynamics from S and pearson correlation
+    CCs_unaligned = np.abs(np.diag(np.corrcoef(La, Lb)[:La.shape[0], La.shape[0]:])) # Pairwise correlations between rows of La and Lb
+    CCs_aligned = S
+
+    if return_aligned_dynamics:
+        return CCs_unaligned, CCs_aligned, La_tilde.T, Lb_tilde.T
+    else:
+        return CCs_unaligned, CCs_aligned
+
+
 '''
 METRIC CALCULATIONS
 '''

tests/test_analysis.py

@@ -8,6 +8,7 @@
 
 import os
 import matplotlib.pyplot as plt
+from sklearn.decomposition import PCA
 from scipy import signal
 
 test_dir = os.path.dirname(__file__)
@@ -136,6 +137,49 @@ def test_get_unit_spiking_mean_variance(self):
         np.testing.assert_allclose(unit_mean, np.array([2, 0]))
         np.testing.assert_allclose(unit_var, np.array([0, 0]))
 
+class AlignDynamicsTests(unittest.TestCase):
+    def test_align_latent_dynamics(self):
+        # Generate dummy latent dynamics for testing
+        n_samples = 1000
+        n_features_x = 30
+        np.random.seed(0)
+
+        X1 = np.random.randn(n_samples,n_features_x).T  # data format used in Juan's 2020 paper is n_u x n_t # Assume nt is trial concatenated reach segments
+        X2 = np.random.randn(n_samples, n_features_x).T
+
+        X1 = (X1 - np.mean(X1, axis=1, keepdims=True)) / np.std(X1, axis=1, keepdims=True)  # mean across units
+        X2 = (X2 - np.mean(X2, axis=1, keepdims=True)) / np.std(X2, axis=1, keepdims=True)
+
+        # Perform PCA to extract latent dynamics
+        pca = PCA(n_components=10)  # Specify the number of components (1 in this example)
+        La = pca.fit_transform(X1.T).T  # Juan's paper dimensions of projected data is m x T (m = 10 for M1 assumed)
+        Lb = pca.fit_transform(X2.T).T
+
+        # Call the align_latent_dynamics function
+        CCs_unaligned, CCs_aligned = aopy.analysis.align_latent_dynamics(La, Lb, False)
+
+        # Assert the shapes of the computed correlations
+        assert CCs_unaligned.shape == (10,)
+        assert CCs_aligned.shape == (10,)
+
+    def test_align_latent_dynamics_samedata(self):
+
+        np.random.seed(42)
+        La = np.random.rand(1000, 10).T
+        Lb = La.copy()
+
+        # Call the align_latent_dynamics function
+        CCs_unaligned, CCs_aligned = aopy.analysis.align_latent_dynamics(La, Lb, False)
+
+        # Assert the shapes of the computed correlations
+        assert CCs_unaligned.shape == (10,)
+        assert CCs_aligned.shape == (10,)
+
+        # Assert that the pairwise correlation for aligned dynamics is approximately 0.99
+        assert np.allclose(CCs_aligned, 0.99, atol=0.01)
+
+
+
 class PCATests(unittest.TestCase):
     # test variance accounted for
     def test_get_pca_dimensions(self):