Merge pull request #4 from CINPLA/add_pv_correlation

population vector correlation

spatial_maps/stats.py

@@ -179,3 +179,65 @@ def prob_dist_1d(x, bins):
 
     H, _ = np.histogram(x, bins=bins, normed=False)
     return (H / len(x)).T
+
+
+def population_vector_correlation(rmaps1, rmaps2,
+                                  mask_nans=False,
+                                  return_corr_coeff_map=False):
+    """
+    Calcualte population vector correlation between two
+    stacks of rate maps.
+    
+    Parameters
+    ----------
+    rmaps1 : ndarray
+    Array of the shape [n_units, n_bins_dim1, n_bins_dim2]
+    rmaps2 : ndarray
+    Array of the shape [n_units, n_bins_dim1, n_bins_dim2]
+    mask_nans : bool
+    If mask_nans, nan-values will be excluded for x-, y-bin
+    from the correlation calculation
+    return_corrcoeffs : bool
+    If return_corr_coeff_map, return the correlation coefficient
+    map instead for the mean.
+
+    Returns
+    -------
+    pop_vec_corr : float
+    Population vector correlation
+    """
+    bins_x = rmaps1.shape[1]
+    bins_y = rmaps1.shape[2]
+
+    assert rmaps1.shape == rmaps2.shape
+    # correlation coefficient requires at least 2x2 values
+    assert rmaps1.shape[0] > 1
+
+    corr_coeff_map = np.zeros((bins_x, bins_y))
+    corr_coeff_map[:] = np.nan
+
+    for i in range(bins_x):
+        for j in range(bins_y):
+            xy1 = rmaps1[:, i, j]
+            xy2 = rmaps2[:, i, j]
+            if mask_nans:
+                bool_nan_xy1 = np.isnan(xy1)
+                bool_nan_xy2 = np.isnan(xy2)
+
+                mask_invalid = np.logical_or(
+                    bool_nan_xy1,
+                    bool_nan_xy2)
+                mask_valid = ~mask_invalid
+                xy1 = xy1[mask_valid]
+                xy2 = xy2[mask_valid]
+
+            corr_coeff_map[i, j] = np.corrcoef(
+                xy1,
+                xy2)[0, 1]
+
+    pop_vec_corr = np.nanmean(corr_coeff_map)
+
+    if return_corr_coeff_map:
+        return pop_vec_corr, corr_coeff_map
+    else:
+        return pop_vec_corr

spatial_maps/tests/test_stats.py

@@ -0,0 +1,39 @@
+import numpy as np
+import pytest
+
+
+def test_calc_population_vector_correlation():
+    from spatial_maps.stats import population_vector_correlation as pvcorr
+    rmaps1 = np.array([
+        [
+            [1, 0.1],
+            [0.1, 4]
+        ],
+        [
+            [6, 0.1],
+            [0.1, 2]
+        ],
+        [
+            [2, 0.1],
+            [0.1, 3]
+        ]])
+    rmaps2 = np.array([
+        [
+            [2, 0.2],
+            [0.2, 8]
+        ],
+        [
+            [12, 0.2],
+            [0.2, 4]
+        ],
+        [
+            [4, 0.2],
+            [0.2, 6]
+        ]])
+    rmaps2 += 10e-5
+    pv = pvcorr(rmaps1, rmaps2)
+    err = pv-1
+    assert err < 10e-5
+
+
+