Repository: PMEAL/OpenPNM
Pull request #1184 (2018-12-31 23:05:29)

Speed up topological queries when flatten is True


Patch:
diff --git a/openpnm/network/GenericNetwork.py b/openpnm/network/GenericNetwork.py
index ed26bd707..a20fa7c1c 100644
--- a/openpnm/network/GenericNetwork.py
+++ b/openpnm/network/GenericNetwork.py
@@ -190,6 +190,10 @@ class GenericNetwork(Base, ModelsMixin):
 
         To obtain a matrix with weights other than ones at each non-zero
         location use ``create_adjacency_matrix``.
+
+        To obtain the non-directed graph, with only upper-triangular entries,
+        use ``sp.sparse.triu(am, k=1)``.
+
         """
         # Retrieve existing matrix if available
         if fmt in self._am.keys():
@@ -235,7 +239,7 @@ class GenericNetwork(Base, ModelsMixin):
         if fmt in self._im.keys():
             im = self._im[fmt]
         elif self._im.keys():
-            im = self._am[list(self._im.keys())[0]]
+            im = self._im[list(self._im.keys())[0]]
             tofmt = getattr(im, 'to'+fmt)
             im = tofmt()
             self._im[fmt] = im
@@ -687,11 +691,16 @@ class GenericNetwork(Base, ModelsMixin):
         pores = self._parse_indices(pores)
         if sp.size(pores) == 0:
             return sp.array([], ndmin=1, dtype=int)
-        if 'lil' not in self._im.keys():
-            self.get_incidence_matrix(fmt='lil')
-        neighbors = topotools.find_neighbor_bonds(sites=pores, logic=mode,
-                                                  im=self._im['lil'],
-                                                  flatten=flatten)
+        if flatten is False:
+            if 'lil' not in self._im.keys():
+                self.get_incidence_matrix(fmt='lil')
+            neighbors = topotools.find_neighbor_bonds(sites=pores, logic=mode,
+                                                      im=self._im['lil'],
+                                                      flatten=flatten)
+        else:
+            am = self.create_adjacency_matrix(fmt='coo', triu=True)
+            neighbors = topotools.find_neighbor_bonds(sites=pores, logic=mode,
+                                                      am=am, flatten=True)
         return neighbors
 
     def _find_neighbors(self, pores, element, **kwargs):
diff --git a/openpnm/topotools/topotools.py b/openpnm/topotools/topotools.py
index 5453b3040..d54b428e9 100644
--- a/openpnm/topotools/topotools.py
+++ b/openpnm/topotools/topotools.py
@@ -109,16 +109,25 @@ def find_neighbor_sites(sites, am, flatten=True, include_input=False,
     return neighbors
 
 
-def find_neighbor_bonds(sites, im, flatten=True, logic='or'):
+def find_neighbor_bonds(sites, im=None, am=None, flatten=True, logic='or'):
     r"""
     Given an incidence matrix, finds all sites that are connected to the
     input sites.
 
+    This function accepts either an incidence or adjacency matrix.
+
     Parameters
     ----------
     im : scipy.sparse matrix
         The incidence matrix of the network.  Must be shaped as (N-sites,
-        N-bonds), with non-zeros indicating which sites are connected.
+        N-bonds), with non-zeros indicating which sites are connected. Either
+        ``am`` or ``im`` must be given.  Passing in ``im`` is slower, but more
+        powerful as it allows for an unflattened list of neighbors.
+
+    am : scipy.sparse matrix (optional)
+        The adjacency matrix of the network. Either ``am`` or ``im`` must be
+        given.  Passing in ``am`` is faster, but does not allow for an
+        unflattened list.
 
     flatten : boolean (default is ``True``)
         Indicates whether the returned result is a compressed array of all
@@ -166,36 +175,62 @@ def find_neighbor_bonds(sites, im, flatten=True, logic='or'):
     if global sites are considered.
 
     """
-    if im.format != 'lil':
-        im = im.tolil(copy=False)
-    rows = [im.rows[i] for i in sp.array(sites, ndmin=1, dtype=sp.int64)]
-    if len(rows) == 0:
-        return []
-    neighbors = sp.hstack(rows).astype(sp.int64)
-    n_bonds = int(im.nnz/2)
-    if logic in ['or', 'union', 'any']:
-        neighbors = sp.unique(neighbors)
-    elif logic in ['xor', 'exclusive_or']:
-        neighbors = sp.unique(sp.where(sp.bincount(neighbors) == 1)[0])
-    elif logic in ['xnor', 'shared']:
-        neighbors = sp.unique(sp.where(sp.bincount(neighbors) > 1)[0])
-    elif logic in ['and', 'all', 'intersection']:
-        neighbors = set(neighbors)
-        [neighbors.intersection_update(i) for i in rows]
-        neighbors = sp.array(list(neighbors), dtype=int, ndmin=1)
-    else:
-        raise Exception('Specified logic is not implemented')
-    if (flatten is False):
-        if (neighbors.size > 0):
-            mask = sp.zeros(shape=n_bonds, dtype=bool)
-            mask[neighbors] = True
-            for i in range(len(rows)):
-                vals = sp.array(rows[i], dtype=sp.int64)
-                rows[i] = vals[mask[vals]]
-            neighbors = rows
+    if im is not None:
+        if im.format != 'lil':
+            im = im.tolil(copy=False)
+        rows = [im.rows[i] for i in sp.array(sites, ndmin=1, dtype=sp.int64)]
+        if len(rows) == 0:
+            return []
+        neighbors = sp.hstack(rows).astype(sp.int64)
+        n_bonds = int(im.nnz/2)
+        if logic in ['or', 'union', 'any']:
+            neighbors = sp.unique(neighbors)
+        elif logic in ['xor', 'exclusive_or']:
+            neighbors = sp.unique(sp.where(sp.bincount(neighbors) == 1)[0])
+        elif logic in ['xnor', 'shared']:
+            neighbors = sp.unique(sp.where(sp.bincount(neighbors) > 1)[0])
+        elif logic in ['and', 'all', 'intersection']:
+            neighbors = set(neighbors)
+            [neighbors.intersection_update(i) for i in rows]
+            neighbors = sp.array(list(neighbors), dtype=int, ndmin=1)
         else:
-            neighbors = [sp.array([], dtype=sp.int64) for i in range(len(sites))]
-    return neighbors
+            raise Exception('Specified logic is not implemented')
+        if (flatten is False):
+            if (neighbors.size > 0):
+                mask = sp.zeros(shape=n_bonds, dtype=bool)
+                mask[neighbors] = True
+                for i in range(len(rows)):
+                    vals = sp.array(rows[i], dtype=sp.int64)
+                    rows[i] = vals[mask[vals]]
+                neighbors = rows
+            else:
+                neighbors = [sp.array([], dtype=sp.int64) for i in range(len(sites))]
+        return neighbors
+    elif am is not None:
+        if am.format != 'coo':
+            am = am.tocoo(copy=False)
+        if not istriu(am):
+            am = sp.sparse.triu(am, k=1)
+        if flatten is False:
+            raise Exception('flatten cannot be used with an adjacency matrix')
+        Ps = sp.zeros(max(am.row.max(), am.col.max())+1, dtype=bool)
+        Ps[sites] = True
+        conns = sp.vstack((am.row, am.col)).T
+        if logic in ['or', 'union', 'any']:
+            neighbors = sp.any(Ps[conns], axis=1)
+        elif logic in ['xor', 'exclusive_or']:
+            neighbors = sp.sum(Ps[conns], axis=1) == 1
+        elif logic in ['xnor', 'shared']:
+            neighbors = sp.all(Ps[conns], axis=1)
+        elif logic in ['and', 'all', 'intersection']:
+            raise Exception('Specified logic is not implemented')
+        else:
+            raise Exception('Specified logic is not implemented')
+        neighbors = sp.where(neighbors)[0]
+        return neighbors
+    else:
+        raise Exception('Either the incidence or the adjacency matrix must ' +
+                        'must be specified')
 
 
 def find_connected_sites(bonds, am, flatten=True, logic='or'):


Test patch:
diff --git a/tests/unit/topotools/GraphToolsTest.py b/tests/unit/topotools/GraphToolsTest.py
index ee58a6f07..7e2bf670f 100644
--- a/tests/unit/topotools/GraphToolsTest.py
+++ b/tests/unit/topotools/GraphToolsTest.py
@@ -12,32 +12,34 @@ class GraphToolsTest:
 
         The Network:
 
-        3 ― 4 ― 5
-        | \ |
-        0 ― 1 ― 2
-
-        The Adjacency Matrix:    The Enumreated Adjacency Matrix:
-
-          | 0  1  2  3  4  5       | 0  1  2  3  4  5
-        ――|――――――――――――――――――    ――|――――――――――――――――――
-        0 | -  1     1           0 | -  0     1
-        1 | 1  -  1  1  1        1 | 0  -  3  2  5
-        2 |    1  -              2 |    3  -
-        3 | 1  1     -  1        3 | 1  2     -  4
-        4 |    1     1  -  1     4 |    5     4  -  6
-        5 |             1  -     5 |             6  -
-
-        The Incidence Matrix:    The Enumerated Incidence Matrix
-
-          | 0  1  2  3  4  5       | 0  1  2  3  4  5
-        ――|――――――――――――――――――    ――|――――――――――――――――――
-        0 | 1  1                 0 | 1  0
-        1 | 1        1           1 | 3        0
-        2 |    1     1           2 |    3     1
-        3 |    1  1              3 |    2  1
-        4 |          1  1        4 |          4  3
-        5 |    1        1        5 |    4        1
-        6 |             1  1     6 |             5  4
+        ::
+
+            3 ― 4 ― 5
+            | \ |
+            0 ― 1 ― 2
+
+            The Adjacency Matrix:    The Enumreated Adjacency Matrix:
+
+              | 0  1  2  3  4  5       | 0  1  2  3  4  5
+            ――|――――――――――――――――――    ――|――――――――――――――――――
+            0 | -  1     1           0 | -  0     1
+            1 | 1  -  1  1  1        1 | 0  -  3  2  5
+            2 |    1  -              2 |    3  -
+            3 | 1  1     -  1        3 | 1  2     -  4
+            4 |    1     1  -  1     4 |    5     4  -  6
+            5 |             1  -     5 |             6  -
+
+            The Incidence Matrix:    The Enumerated Incidence Matrix
+
+              | 0  1  2  3  4  5       | 0  1  2  3  4  5
+            ――|――――――――――――――――――    ――|――――――――――――――――――
+            0 | 1  1                 0 | 1  0
+            1 | 1        1           1 | 3        0
+            2 |    1     1           2 |    3     1
+            3 |    1  1              3 |    2  1
+            4 |          1  1        4 |          4  3
+            5 |    1        1        5 |    4        1
+            6 |             1  1     6 |             5  4
 
         """
         self.ws = op.Workspace()
@@ -222,6 +224,34 @@ class GraphToolsTest:
         with pytest.raises(Exception):
             topotools.find_neighbor_bonds(sites=[0], im=im, logic='nand')
 
+    def test_find_neighbor_bonds_with_am_and_logic(self):
+        am = self.net.get_adjacency_matrix(fmt='coo')
+        im = self.net.get_incidence_matrix(fmt='coo')
+        Ts1 = topotools.find_neighbor_bonds(sites=[1], am=am,
+                                            flatten=True, logic='or')
+        Ts2 = topotools.find_neighbor_bonds(sites=[1], im=im,
+                                            flatten=True, logic='or')
+        assert np.all(Ts1 == Ts2)
+        Ts1 = topotools.find_neighbor_bonds(sites=[1], am=am,
+                                            flatten=True, logic='xor')
+        Ts2 = topotools.find_neighbor_bonds(sites=[1], im=im,
+                                            flatten=True, logic='xor')
+        assert np.all(Ts1 == Ts2)
+        Ts1 = topotools.find_neighbor_bonds(sites=[1], am=am,
+                                            flatten=True, logic='xnor')
+        Ts2 = topotools.find_neighbor_bonds(sites=[1], im=im,
+                                            flatten=True, logic='xnor')
+        assert np.all(Ts1 == Ts2)
+
+    def test_find_neighbor_bonds_with_am_exceptions(self):
+        am = self.net.get_adjacency_matrix(fmt='coo')
+        with pytest.raises(Exception):
+            topotools.find_neighbor_bonds(sites=[1], am=am, flatten=True,
+                                          logic='intersection')
+        with pytest.raises(Exception):
+            topotools.find_neighbor_bonds(sites=[1], am=am, flatten=False,
+                                          logic='or')
+
     def test_find_neighbor_sites(self):
         am = self.net.create_adjacency_matrix(fmt='lil')
         Ps = topotools.find_neighbor_sites(sites=[0], am=am)