Repository: PMEAL/OpenPNM
Pull request #2137 (2021-11-03 20:37:32)

Add asmask argument to neighbor lookup functions


Patch:
diff --git a/example.py b/example.py
index 5fe4baa53..c781347fe 100644
--- a/example.py
+++ b/example.py
@@ -25,6 +25,7 @@ mip.run()
 hg.update(mip.results(Pc=70000))
 # mip.plot_intrusion_curve()
 
+
 # %% Perform Stokes flow simulation
 sf = op.algorithms.StokesFlow(network=pn, phase=water)
 sf.set_value_BC(pores=pn.pores('right'), values=0)
diff --git a/openpnm/network/_generic.py b/openpnm/network/_generic.py
index 685565a4f..eaa421244 100644
--- a/openpnm/network/_generic.py
+++ b/openpnm/network/_generic.py
@@ -502,7 +502,7 @@ class GenericNetwork(ParamMixin, Base, ModelsMixin, LabelMixin):
         return Ts
 
     def find_neighbor_pores(self, pores, mode='union', flatten=True,
-                            include_input=False):
+                            include_input=False, asmask=False):
         r"""
         Returns a list of pores that are direct neighbors to the given pore(s)
 
@@ -540,6 +540,11 @@ class GenericNetwork(ParamMixin, Base, ModelsMixin, LabelMixin):
                 (somtimes) as 'all' in boolean logic. Both keywords are
                 accepted and treated as 'and'.
 
+        asmask : boolean
+            If ``False`` (default), the returned result is a list of the
+            neighboring pores as indices. If ``True``, the returned result is a
+            boolean mask. (Useful for labelling)
+
         Returns
         -------
         If ``flatten`` is ``True``, returns a 1D array of pore indices
@@ -591,9 +596,16 @@ class GenericNetwork(ParamMixin, Base, ModelsMixin, LabelMixin):
                                                   am=self._am['lil'],
                                                   flatten=flatten,
                                                   include_input=include_input)
-        return neighbors
+        if asmask is False:
+            return neighbors
+        elif flatten is True:
+            neighbors = self._tomask(element='pore', indices=neighbors)
+            return neighbors
+        else:
+            raise Exception('Cannot create mask on an unflattened output')
 
-    def find_neighbor_throats(self, pores, mode='union', flatten=True):
+    def find_neighbor_throats(self, pores, mode='union', flatten=True,
+                              asmask=False):
         r"""
         Returns a list of throats neighboring the given pore(s)
 
@@ -625,6 +637,11 @@ class GenericNetwork(ParamMixin, Base, ModelsMixin, LabelMixin):
                 (somtimes) as 'all' in boolean logic. Both keywords are
                 accepted and treated as 'and'.
 
+        asmask : boolean
+            If ``False`` (default), the returned result is a list of the
+            neighboring throats as indices. If ``True``, the returned result is a
+            boolean mask. (Useful for labelling)
+
         Returns
         -------
         If ``flatten`` is ``True``, returns a 1D array of throat indices
@@ -667,7 +684,13 @@ class GenericNetwork(ParamMixin, Base, ModelsMixin, LabelMixin):
             am = self.create_adjacency_matrix(fmt='coo', triu=True)
             neighbors = topotools.find_neighbor_bonds(sites=pores, logic=mode,
                                                       am=am, flatten=True)
-        return neighbors
+        if asmask is False:
+            return neighbors
+        elif flatten is True:
+            neighbors = self._tomask(element='throat', indices=neighbors)
+            return neighbors
+        else:
+            raise Exception('Cannot create mask on an unflattened output')
 
     def _find_neighbors(self, pores, element, **kwargs):
         element = self._parse_element(element=element, single=True)


Test patch:
diff --git a/tests/unit/network/GenericNetworkTest.py b/tests/unit/network/GenericNetworkTest.py
index 1c3d859ae..3bbd9f88f 100644
--- a/tests/unit/network/GenericNetworkTest.py
+++ b/tests/unit/network/GenericNetworkTest.py
@@ -1,5 +1,6 @@
 import numpy as np
 import openpnm as op
+import pytest
 
 
 class GenericNetworkTest:
@@ -58,6 +59,15 @@ class GenericNetworkTest:
                                          include_input=True)
         assert np.all(a == [0, 1, 2, 10, 11, 100, 101])
 
+    def test_find_neighbor_pores_asmask(self):
+        a = self.net.find_neighbor_pores(pores=[0, 2], flatten=True, asmask=True)
+        assert sum(a) == 6
+        assert a.dtype == 'bool'
+
+    def tes_find_neighbor_pores_asmask_unflattened(self):
+        with pytest.raises():
+            a = self.net.find_neighbor_pores(pores=[0, 2], flatten=False, asmask=True)
+
     def test_find_neighbor_pores_numeric_intersection_include_input(self):
         a = self.net.find_neighbor_pores(pores=[0, 2], mode='and',
                                          include_input=True)