Make array correspondence simpler, slightly safer, and make API more numpy-like

.gitignore

@@ -9,4 +9,5 @@ build-linux/
 .vscode/
 site/
 .eggs/
-ignore/
+ignore/
+*.egg-info

README.md

@@ -153,10 +153,8 @@ library. We are thankful to them:
 
 ### Contributors
 
-We would like to thank [Michael Jäger](https://github.com/EmJay276) for being
-our Gpytoolbox's first external contributor (see [PR
-#45](https://github.com/sgsellan/gpytoolbox/pull/45)).
-
+- We would like to thank [Michael Jäger](https://github.com/EmJay276) for being Gpytoolbox's first external contributor (see [PR #45](https://github.com/sgsellan/gpytoolbox/pull/45)).
+- [Towaki Takikawa](https://github.com/tovacinni) ([PR #49](https://github.com/sgsellan/gpytoolbox/pull/49))
 
 
 <!-- Most of the functionality in this library is python-only, and it requires no

docs/index.md

@@ -152,4 +152,5 @@ library. We are thankful to them:
 
 ### Contributors
 
-We would like to thank [Michael Jäger](https://github.com/EmJay276) for being our Gpytoolbox's first external contributor (see [PR #45](https://github.com/sgsellan/gpytoolbox/pull/45)).
+- We would like to thank [Michael Jäger](https://github.com/EmJay276) for being our Gpytoolbox's first external contributor (see [PR #45](https://github.com/sgsellan/gpytoolbox/pull/45)).
+- [Towaki Takikawa](https://github.com/tovacinni) ([PR #49](https://github.com/sgsellan/gpytoolbox/pull/49))

src/gpytoolbox/array_correspondence.py

@@ -22,62 +22,52 @@ def array_correspondence(A,B,axis=None):
 
     Examples
     --------
-    TODO
+    Example with simple array correspondence:
+    ```python
+    >>> A = np.array([1,7,2,7,9,3,7,0])
+    >>> B = np.array([7,7,3,2,7,8])
+    >>> f = gpy.array_correspondence(A, B)
+    >>> f
+    array([-1,  0,  3,  0, -1,  2,  0, -1])
+    ```
+    Example with row correspondence:
+    ```python
+    >>> A = np.array([[1,3],[5,2],[1,2],[5,2]])
+    >>> B = np.array([[1,2],[6,9],[5,2]])
+    >>> f = gpy.array_correspondence(A, B, axis=0)
+    >>> f
+    array([-1,  2,  0,  2])
+    ```
 
     """
+    if axis not in (None, 0, 1, -1, -2):
+        raise Exception("Axis can only be None, 0, 1, -1, -2")
+    if len(A.shape) > 2 or len(B.shape) > 2:
+        raise Exception("Inputs A, B can only be up to 2 dimensional")
 
-    if axis is None:
-        A = A.ravel()
-        B = B.ravel()
+    if axis==1 or axis==-1:
+        # np.unique behaves weird with axis==1... but we only work with
+        # up to dim 2, so always simply use the axis==0 case.
+        A = A.T
+        B = B.T
+        axis = 0
 
-        # Slow for loop
-        # f = np.full(A.size, -1, dtype=np.int64)
-        # for a in range(A.size):
-        #     for b in range(B.size):
-        #         if A[a]==B[b]:
-        #             f[a] = b
+    # While we have to keep track of duplicates in A (to map them to the
+    # correct place), we do not care about duplicates in B
 
-        # While we have to keep track of duplicates in A (to map them to the
-        # correct place), we do not care about duplicates in B
-        uB,mapB = np.unique(B, return_index=True)
-        _,idx,inv = np.unique(np.concatenate((uB,A)),
-            return_index=True, return_inverse=True)
-        imap = idx[inv[uB.size:]]
-        imap[imap>=uB.size] = -1
-        f = np.where(imap<0, -1, mapB[imap])
+    # uB is deduplicated B. mapB allows us to map to the first occurence of each vector.
+    # (contribution by Towaki Takikawa)
+    uB, mapB = np.unique(B, return_index=True, axis=axis)
 
-    else:
-        assert len(A.shape) == 2
-        assert len(B.shape) == 2
-        assert axis==-2 or axis==-1 or axis==0 or axis==1
-        assert A.shape[axis] == B.shape[axis]
+    # We concatenate uB with A. Any entries from A that get de-duped is a 'hit'.
+    _, idx, inv = np.unique(np.concatenate((uB,A), axis=axis),
+        return_index=True, return_inverse=True, axis=axis)
 
-        # This function compares rows, so we reduce the problem to rows here.
-        if axis==-2 or axis==0:
-                A = A.transpose()
-                B = B.transpose()
-
-        # # https://stackoverflow.com/a/64930992 is too memory-intensive
-        # inds = (A[None,:,:] == B[:,None,:])
-        # i,j = np.nonzero(inds.sum(axis=2) == A.shape[1])
-        # f = np.full(A.shape[0], -1, dtype=np.int64)
-        # f[j] = i
-
-        # # Slower (but less memory-intensive) with for loops
-        # f = np.full(A.shape[0], -1, dtype=np.int64)
-        # for a in range(A.shape[0]):
-        #     for b in range(B.shape[0]):
-        #         if (A[a,:]==B[b,:]).all():
-        #             f[a] = b
-
-        # Convert each row to bytes, intersect byte arrays in 1d
-        def to_byte_array(x):
-            # Adapted from https://stackoverflow.com/a/54683422
-            dt = np.dtype('S{:d}'.format(x.shape[1] * x.dtype.itemsize))
-            return np.frombuffer(x.tobytes(), dtype=dt)
-        bytesA = to_byte_array(A)
-        bytesB = to_byte_array(B.astype(A.dtype))
-        f = array_correspondence(bytesA,bytesB,axis=None)
+    # We don't care about the range of the output that is about uB- that was a decoy to
+    # grill out the hits from A.
+    imap = idx[inv[uB.shape[0]:]]
+    imap[imap>=uB.shape[0]] = -1
+    f = np.where(imap<0, -1, mapB[imap])
 
     return f
 

src/gpytoolbox/halfedge_edge_map.py

@@ -75,9 +75,9 @@ def unflat_he(flat_he):
             E_to_he[0:n_b,0,:] = unflat_he(bdry_E_to_flat_he)
         if n_i>0:
             E_to_he[n_b:(n_b+n_i),0,:] = \
-            unflat_he(array_correspondence(interior_E,flat_he,axis=1))
+            unflat_he(array_correspondence(interior_E,flat_he,axis=0))
             E_to_he[n_b:(n_b+n_i),1,:] = \
-            unflat_he(array_correspondence(interior_E,np.flip(flat_he, axis=-1),axis=1))
+            unflat_he(array_correspondence(interior_E,np.flip(flat_he, axis=-1),axis=0))
     else:
         E_to_he = []
         bdry_E_to_he = unflat_he(bdry_E_to_flat_he)

src/gpytoolbox/triangle_triangle_adjacency.py

@@ -40,7 +40,7 @@ def triangle_triangle_adjacency(F):
     he_flat = np.concatenate((he[:,0,:], he[:,1,:], he[:,2,:]), axis=0)
     he_flip_flat = np.flip(he_flat, axis=-1)
 
-    map_to_flip = array_correspondence(he_flat,he_flip_flat,axis=1)
+    map_to_flip = array_correspondence(he_flat,he_flip_flat,axis=0)
     TT = np.reshape(np.where(map_to_flip<0, -1, map_to_flip % m), F.shape, order='F')
     TTi = np.reshape(np.where(map_to_flip<0, -1, map_to_flip // m), F.shape, order='F')
 

test/test_array_correspondence.py

@@ -27,30 +27,48 @@ def test_random_arrays(self):
     def test_matrix1(self):
         A = np.array([[1,2],[1,3],[1,4],[2,0],[3,2],[4,3],[0,1],[2,1],[3,1]])
         B = np.array([[2,1],[3,1],[1,4],[1,4],[3,4],[1,0],[1,2],[1,3]])
-        f = gpy.array_correspondence(A,B,axis=1)
+        f = gpy.array_correspondence(A,B,axis=0)
         self.mapping_condition(A, B, f)
-        fi = gpy.array_correspondence(B,A,axis=1)
+        fi = gpy.array_correspondence(B,A,axis=0)
         self.mapping_condition(B, A, fi)
-        ft = gpy.array_correspondence(A.transpose(),B.transpose(),axis=0)
+        ft = gpy.array_correspondence(A.transpose(),B.transpose(),axis=1)
         self.mapping_condition(A, B, ft)
-        fti = gpy.array_correspondence(B.transpose(),A.transpose(),axis=0)
+        fti = gpy.array_correspondence(B.transpose(),A.transpose(),axis=1)
         self.mapping_condition(B, A, fti)
 
     def test_random_matrices(self):
         rng = default_rng()
         dims = np.concatenate((rng.integers(1, 10000, size=(20,2)),
-            rng.integers(1, 6, size=(20,1))), axis=1)
+                               rng.integers(1, 6, size=(20,1))), axis=1)
         for i in range(dims.shape[0]):
             A = rng.integers(0, 100, size=(dims[i,0],dims[i,2]))
             B = rng.integers(0, 100, size=(dims[i,1],dims[i,2]))
-            f = gpy.array_correspondence(A,B,axis=1)
+            f = gpy.array_correspondence(A,B,axis=0)
             self.mapping_condition(A, B, f)
-            fi = gpy.array_correspondence(B,A,axis=1)
+            fi = gpy.array_correspondence(B,A,axis=0)
             self.mapping_condition(B, A, fi)
-            ft = gpy.array_correspondence(A.transpose(),B.transpose(),axis=0)
+            ft = gpy.array_correspondence(A.transpose(),B.transpose(),axis=1)
             self.mapping_condition(A, B, ft)
-            fti = gpy.array_correspondence(B.transpose(),A.transpose(),axis=0)
+            fti = gpy.array_correspondence(B.transpose(),A.transpose(),axis=1)
             self.mapping_condition(B, A, fti)
+
+    def test_objects(self):
+        class Aobj():
+            def __init__(self, val=5):
+                self.val = val
+            def __eq__(self, other):
+                return self.val == other.val
+            def __lt__(self, other):
+                return self.val < other.val
+        a_ptr_0 = Aobj(9)
+        a_ptr_1 = a_ptr_0
+        a_ptr_2 = a_ptr_1
+        a_ptr_4 = Aobj(5)
+        A = np.array([[a_ptr_0], [a_ptr_1], [a_ptr_2], [a_ptr_0], 
+                      [Aobj(2)], [Aobj(2)], [Aobj(8)], [a_ptr_4]])
+        B = np.array([[Aobj(5)], [Aobj(2)], [Aobj(4)], [a_ptr_0]])
+        f0 = gpy.array_correspondence(A[..., 0], B[..., 0])
+        self.mapping_condition(A, B, f0)
 
     def mapping_condition(self, A, B, f):
         self.assertTrue(f.shape[0] == A.shape[0])