Add routine for mapping unique entries in array potentially containin…

…g duplicates/nonconsecutive entries to consecutive integers.

neurokernel/tools/misc.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 from functools import wraps
+import itertools
 import numbers
 import re
 import subprocess
@@ -219,3 +220,37 @@ def openmpi_cuda_support(path='ompi_info'):
                     return False
         return False
 
+def renumber_in_order(arr):
+    """
+    Map unique array elements to in-order integers.
+
+    Maps an array of elements that may contain duplicates and might not be
+    monotonically ordered to an array of equivalent length where the unique elements
+    are respectively replaced by 0, 1, etc.
+
+    Parameters
+    ----------
+    arr : array_like
+        1D array of elements.
+
+    Returns
+    -------
+    result : list
+        Array of mapped integers.
+
+    Examples
+    --------
+    >>> arr = np.array([0, 2, 2, 3, 5, 5])
+    >>> result = renumber_in_order(arr)
+    >>> np.allclose(arr, [0, 1, 1, 2, 3, 3])
+    True
+    """
+
+    c = itertools.count()
+    result = []
+    already_seen = {}
+    for e in arr:
+        if e not in already_seen:
+            already_seen[e] = c.next()
+        result.append(already_seen[e])
+    return result

tests/test_misc.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+
+from unittest import main, TestCase
+
+import nk.tools.misc as misc
+
+class test_misc(TestCase):
+    def test_no_duplicates_consecutive_in_order(self):
+        result = misc.renumber_in_order([0, 1, 2])
+        self.assertSequenceEqual(result,
+                                 [0, 1, 2])
+
+    def test_no_duplicates_nonconsecutive_in_order(self):
+        result = misc.renumber_in_order([0, 2, 4])
+        self.assertSequenceEqual(result,
+                                 [0, 1, 2])
+
+    def test_duplicates_consecutive_in_order(self):
+        result = misc.renumber_in_order([0, 0, 1, 2, 2, 3])
+        self.assertSequenceEqual(result,
+                                 [0, 0, 1, 2, 2, 3])
+
+    def test_duplicates_nonconsecutive_in_order(self):
+        result = misc.renumber_in_order([0, 0, 2, 3, 3, 5])
+        self.assertSequenceEqual(result,
+                                 [0, 0, 1, 2, 2, 3])
+
+    def test_duplicates_nonconsecutive_out_of_order(self):
+        result = misc.renumber_in_order([0, 3, 1, 1, 2, 3])
+        self.assertSequenceEqual(result,
+                                 [0, 1, 2, 2, 3, 1])
+
+if __name__ == '__main__':
+    main()
+