add invert idx

hippocampy/matrix_utils.py

@@ -1112,11 +1112,63 @@ def average_diag(mat: np.ndarray):
     if mat.shape[0] != mat.shape[1]:
         raise ValueError("This function only work for square matrices")
 
-    l = mat.shape[0]
+    L = mat.shape[0]
     # np trace calculate the sum over a diagonal with an offset k
     # we do it over all the possible diagonals
-    b = np.array([np.trace(mat, k) for k in np.arange(-l + 1, l)])
+    b = np.array([np.trace(mat, k) for k in np.arange(-L + 1, L)])
     # calculate the number of element in each diag
-    n = -np.abs(np.arange(-l + 1, l)) + l
+    n = -np.abs(np.arange(-L + 1, L)) + L
     # calculate the average
     return b / n
+
+
+def invert_indices(idx: int, size: int):
+    """
+    Inverts a list of indices within a given range.
+
+    This function takes a list or array of indices (`idx`) and the total number of elements (`size`),
+    and returns an array of indices that are not included in the input list within the specified range.
+    Essentially, it calculates the complement set of indices in the range [0, size).
+
+    Parameters:
+    idx (array-like): An array, list, or sequence of integer indices to be inverted.
+    size (int): The size of the array in which the indices are to be considered.
+                This is the upper bound of the range of indices to invert.
+
+    Returns:
+    numpy.ndarray: An array of indices that represent the complement of the input `idx` within the range [0, size).
+
+    Raises:
+    ValueError: If any index in `idx` is negative or greater than or equal to `size`.
+
+    Examples:
+    >>> invert_indices([1, 3, 5], 10)
+    array([0, 2, 4, 6, 7, 8, 9])
+
+    >>> invert_indices([0, 2, 4, 6, 8], 10)
+    array([1, 3, 5, 7, 9])
+
+    Note:
+    This function does not handle duplicate values in `idx`. If `idx` contains duplicates,
+    the returned array will simply ignore these without raising an error.
+    """
+    idx = np.asarray(idx)
+
+    # if idx is empty return a range
+    if idx.size == 0:
+        return np.arange(size)
+
+    # Check if any index is outside the allowed range [0, size)
+    if any(idx < 0) or any(idx >= size):
+        # If an index is out of bounds, raise a ValueError with an explanatory message
+        raise ValueError(f"Indices should be >= 0 and smaller than {size}")
+
+    # Create a dense boolean array ('d_mat') with 'size' elements initialized to False
+    d_mat = np.zeros(size, dtype=bool)
+
+    # Set the elements at the indices given by 'idx' to True
+    d_mat[idx] = True
+
+    # Invert the boolean array ('d_mat') and get the indices of the elements
+    # that are now True
+    return np.nonzero(~d_mat)[0]

tests/test_matrix_utils.py

@@ -1,10 +1,11 @@
-#%% Import
+# %% Import
 import warnings
 from hippocampy import matrix_utils
 import unittest
 import numpy as np
 
-#%%
+
+# %%
 class TestLabel(unittest.TestCase):
     def test_label_bool(self):
         # test boolean vector
@@ -122,67 +123,124 @@ def test_smooth_float(self):
             v_s = matrix_utils.smooth_1d(v,2.1)
             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]
+        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)
+        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)
+        a = matrix_utils.corr_mat(xarr, axis=1)
         b = np.corrcoef(xarr)
 
-        np.testing.assert_array_almost_equal(a,b)
+        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)
+        a = matrix_utils.corr_mat(xarr, axis=0)
         b = np.corrcoef(xarr.T)
 
-        np.testing.assert_array_almost_equal(a,b.T)
+        np.testing.assert_array_almost_equal(a, b.T)
 
-
 
 class Testdiagonality(unittest.TestCase):
     def test_mat(self):
-        a = np.array([[6,1,0],[1,5,2],[1,3,6]])
+        a = np.array([[6, 1, 0], [1, 5, 2], [1, 3, 6]])
         d = matrix_utils.diagonality(a)
-        self.assertAlmostEqual(d,0.674149,places=6)
-    
+        self.assertAlmostEqual(d, 0.674149, places=6)
+
     def test_nan(self):
-        a_n = np.array([[6,1,np.nan],[1,5,2],[1,3,6]])
+        a_n = np.array([[6, 1, np.nan], [1, 5, 2], [1, 3, 6]])
         with self.assertRaises(ValueError):
             matrix_utils.diagonality(a_n)
 
     def test_tridiag(self):
-        a = np.array([[2,1,0,0,0],[1,3,2,0,0],[0,2,3,4,0],[0,0,1,2,3],[0,0,0,1,1]])
+        a = np.array(
+            [
+                [2, 1, 0, 0, 0],
+                [1, 3, 2, 0, 0],
+                [0, 2, 3, 4, 0],
+                [0, 0, 1, 2, 3],
+                [0, 0, 0, 1, 1],
+            ]
+        )
         d = matrix_utils.diagonality(a)
-        self.assertAlmostEqual(d,0.812383,places=6)
-    
+        self.assertAlmostEqual(d, 0.812383, places=6)
+
     def test_transpose(self):
-        a = np.array([[6,1,0],[1,5,2],[1,3,6]])
+        a = np.array([[6, 1, 0], [1, 5, 2], [1, 3, 6]])
         d = matrix_utils.diagonality(a)
         d_t = matrix_utils.diagonality(a.T)
-        self. assertEqual(d,d_t)
+        self.assertEqual(d, d_t)
 
     def test_opositdiag(self):
-        a = np.array([[6,1,0],[1,5,2],[1,3,6]])
+        a = np.array([[6, 1, 0], [1, 5, 2], [1, 3, 6]])
         d = matrix_utils.diagonality(a)
         d_op = matrix_utils.diagonality(np.rot90(a))
-        self. assertEqual(d,-d_op)
+        self.assertEqual(d, -d_op)
+
+
+# Define a test case for invert_indices function
+class TestInvertIndices(unittest.TestCase):
+    def test_basic_functionality(self):
+        np.testing.assert_array_almost_equal(
+            matrix_utils.invert_indices([1, 3, 5], 10), [0, 2, 4, 6, 7, 8, 9]
+        )
+        np.testing.assert_array_almost_equal(
+            matrix_utils.invert_indices([0, 2, 4, 6, 8], 10), [1, 3, 5, 7, 9]
+        )
+
+    def test_empty_indices(self):
+        self.assertTrue(
+            np.array_equal(matrix_utils.invert_indices([], 10), np.arange(10))
+        )
+
+    def test_full_range(self):
+        self.assertEqual(len(matrix_utils.invert_indices(np.arange(10), 10)), 0)
+
+    def test_single_index(self):
+        self.assertTrue(
+            np.array_equal(
+                matrix_utils.invert_indices([5], 10), [0, 1, 2, 3, 4, 6, 7, 8, 9]
+            )
+        )
+
+    def test_index_out_of_bounds(self):
+        with self.assertRaises(ValueError):
+            matrix_utils.invert_indices([10], 10)
+        with self.assertRaises(ValueError):
+            matrix_utils.invert_indices([-1], 10)
+
+    def test_duplicate_indices(self):
+        # Assuming function should handle duplicates by ignoring them
+        self.assertTrue(
+            np.array_equal(
+                matrix_utils.invert_indices([1, 1, 3, 5], 10), [0, 2, 4, 6, 7, 8, 9]
+            )
+        )
 
+    def test_large_range(self):
+        # This will also test the function's performance on a larger scale
+        size = 10000
+        idx = np.random.choice(size, size // 2, replace=False)
+        result = matrix_utils.invert_indices(idx, size)
+        # Check that there are no indices in result that are in idx
+        self.assertTrue(np.intersect1d(result, idx).size == 0)
+        # Check that the result and idx together make up the full range
+        self.assertEqual(len(np.union1d(result, idx)), size)
 
 
 # %%