refine and test corr mat function

.github/workflows/python-package-conda.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       max-parallel: 5
       matrix:
-        python-version: [3.5, 3.6, 3.7, 3.8, 3.9]
+        python-version: [3.7, 3.8, 3.9,3.10]
 
     steps:
     - uses: actions/checkout@v2

README.md

@@ -1,4 +1,4 @@
-[![Build with Conda and python 3.5-3.9](https://github.com/bouromain/hippocampy/actions/workflows/python-package-conda.yml/badge.svg)](https://github.com/bouromain/hippocampy/actions/workflows/python-package-conda.yml)
+[![Build with Conda and python 3.5-3.10](https://github.com/bouromain/hippocampy/actions/workflows/python-package-conda.yml/badge.svg)](https://github.com/bouromain/hippocampy/actions/workflows/python-package-conda.yml)
 
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 

environment.yml

@@ -1,7 +1,7 @@
 name: hippocampy
 
 dependencies:
-  - python>=3.8,<3.10
+  - python>=3.8,<3.11
   - pip
   - mkl
   - numpy

hippocampy/binning.py

@@ -371,7 +371,7 @@ def psth_2d(
         raise NotImplementedError(f"Method {method} not implemented")
 
     # to be sure we have floats here, this can create problems with nan later
-    mat = mat.astype(np.float, casting="safe")
+    mat = mat.astype(float, casting="safe")
 
     # initialise values
     sz = mat.shape

hippocampy/matrix_utils.py

@@ -1,13 +1,16 @@
+from typing import Union
+
 import bottleneck as bn
 import numpy as np
-from astropy.convolution import convolve
-from skimage import measure
 import pandas as pd
 import tqdm as tqdm
-from hippocampy.utils.nan import remove_nan
+from astropy.convolution import convolve
+from skimage import measure
 
+from hippocampy.utils.nan import remove_nan
 from hippocampy.utils.type_utils import float_to_int
 
+
 #%% SMOOTH
 def smooth_1d(
     data: np.ndarray,
@@ -187,14 +190,18 @@ def smooth_2d(
 
 
 #%% STATISTICS
-def corr_mat(a: np.ndarray, axis=-1) -> np.ndarray:
+def corr_mat(a: np.ndarray, b: Union[None, np.ndarray] = None, axis=-1) -> np.ndarray:
     """
-    Compute correlation between all the rows (or column) of a given matrix
+    Compute correlation between a two matrices in a particular dimension. If only 
+    one matrix is provided an autocorrelation will be returned. 
 
     Parameters
     ----------
     a : np.ndarray
         input matrix for example [unit, samples]
+
+    b : np.ndarray
+        input matrix for example [other_unit, samples]
     axis : int, optional
         axis along which the function is performed, by default -1
 
@@ -203,21 +210,32 @@ def corr_mat(a: np.ndarray, axis=-1) -> np.ndarray:
     np.ndarray
         correlation matrix
 
-    TODO modify this function in order to take two array as an input
-
     Reference
     ---------
     https://en.wikipedia.org/wiki/Pearson_correlation_coefficient
     """
-    a = np.asarray(a)
 
-    a_z = zscore(a, axis=axis)
+    a = np.array(a, ndmin=2)
     n = a.shape[axis]
 
+    if a.ndim > 2:
+        raise ValueError("Input should have a max 2 dimensions")
+
+    a_z = zscore(a, axis=axis)
+
+    if b is None:
+        b_z = a_z
+    else:
+        if a.ndim > 2:
+            raise ValueError("Input should have a max 2 dimensions")
+
+        b = np.array(b, ndmin=2)
+        b_z = zscore(b, axis=axis)
+
     if axis == 0:
-        return (1 / (n - 1)) * (a_z.T @ a_z)
+        return np.squeeze((1 / (n - 1)) * (a_z.T @ b_z))
     else:
-        return (1 / (n - 1)) * (a_z @ a_z.T)
+        return np.squeeze((1 / (n - 1)) * (a_z @ b_z.T))
 
 
 def zscore(matrix, axis=-1):
@@ -939,7 +957,7 @@ def find_peaks(M, min_amplitude=None):
     return peaks, peaks_idx
 
 
-def fill_diag_slice(mat: np.array, val: np.float = np.nan):
+def fill_diag_slice(mat: np.array, val: float = np.nan):
     """
     Fill the diagonal 
 

setup.py

@@ -7,7 +7,7 @@
     author="Bourboulou Romain",
     author_email="[email protected]",
     packages=["hippocampy"],
-    python_requires=">=3.7, <3.10",
+    python_requires=">=3.7, <3.11",
     install_requires=[
         "numpy",
         "scipy",

tests/test_matrix_utils.py

@@ -107,3 +107,56 @@ def test_smooth(self):
     # fmt: on
         assert (np.testing.assert_array_almost_equal(v_s,exp_v) == None)
 
+class TestCorrMat(unittest.TestCase):
+    def test_vect_corr(self):
+        # this example is taken from numpy corrcoef function docstring
+        rng = np.random.default_rng(seed=42)
+        xarr = rng.random((3, 3))
+
+        a = matrix_utils.corr_mat(xarr[0,:],xarr[1,:],axis = 1)
+        b = np.corrcoef(xarr[0,:],xarr[1,:])[0,1]
+
+        self.assertAlmostEqual(a,b,places=9)
+
+    def test_vect_rows(self):
+        # this example is taken from numpy corrcoef function docstring
+        rng = np.random.default_rng(seed=42)
+        xarr = rng.random((3, 3))
+
+        a = matrix_utils.corr_mat(xarr, axis = 1)
+        b = np.corrcoef(xarr)
+
+        np.testing.assert_array_almost_equal(a,b)
+
+    def test_vect_col(self):
+        # this example is taken from numpy corrcoef function docstring
+        rng = np.random.default_rng(seed=42)
+        xarr = rng.random((3, 3))
+
+        a = matrix_utils.corr_mat(xarr, axis = 0)
+        b = np.corrcoef(xarr.T)
+
+        np.testing.assert_array_almost_equal(a,b.T)
+
+
+
+
+
+
+
+# %%
+
+# import numpy as np
+# from hippocampy.matrix_utils import corr_mat
+# rng = np.random.default_rng(seed=42)
+# xarr = rng.random((3, 3))
+
+# a = corr_mat(xarr, axis = 1)
+# b = np.corrcoef(xarr)
+
+# sel
+
+
+
+
+# %%