Repository: PMEAL/OpenPNM
Pull request #382 (2015-06-10 01:45:36)

Feature/find nearby pores
This method replaces the find_nearest_pores (nearby is more apt) and is now complete, include the option to flatten output, exclude self and so on.  There are also numerous unit tests.  


Patch:
diff --git a/OpenPNM/Network/__GenericNetwork__.py b/OpenPNM/Network/__GenericNetwork__.py
index ba2d1b9c2..d73730dd9 100644
--- a/OpenPNM/Network/__GenericNetwork__.py
+++ b/OpenPNM/Network/__GenericNetwork__.py
@@ -722,17 +722,21 @@ class GenericNetwork(Core):
 
         return (p_clusters, t_clusters)
 
-    def find_nearest_pores(self, pores, distance=0):
+    def find_nearby_pores(self, pores, distance, flatten=False, excl_self=True):
         r"""
-        Find all pores within a given distance of the input pore(s) regardless
-        of whether or not they are toplogically connected.
+        Find all pores within a given radial distance of the input pore(s)
+        regardless of whether or not they are toplogically connected.
 
         Parameters
         ----------
         pores : array_like
             The list of pores for whom nearby neighbors are to be found
         distance : scalar
-            The within which the nearby pores should be found
+            The maximum distance within which the nearby should be found
+        excl_self : bool
+            Controls whether the input pores should be included in the returned
+            list.  The default is True which means they are not included.
+        flatten :
 
         Returns
         -------
@@ -743,18 +747,49 @@ class GenericNetwork(Core):
 
         Examples
         --------
-        >>> import OpenPNM as op
-        >>> pn = op.Network.TestNet()
-        >>> pn.find_nearest_pores(pores=[0, 1],distance=1)
-        array([[0, 1, 25, 5], [0, 1, 26, 6, 2]], dtype=object)
-        >>> pn.find_nearest_pores(pores=[0, 1],distance=.5)
-        array([[0], [1]], dtype=object)
+        >>> import OpenPNM
+        >>> pn = OpenPNM.Network.TestNet()
+        >>> pn.find_nearby_pores(pores=[0, 1], distance=1)
+        array([array([ 1,  5, 25]), array([ 0,  2,  6, 26])], dtype=object)
+        >>> pn.find_nearby_pores(pores=[0, 1], distance=0.5)
+        array([], shape=(2, 0), dtype=int64)
         """
+        # Convert to ND-array
+        pores = sp.array(pores, ndmin=1)
+        # Convert boolean mask to indices if necessary
+        if pores.dtype == bool:
+            pores = self.Ps[pores]
+        # Handle an empty array if given
+        if sp.size(pores) == 0:
+            return sp.array([], dtype=sp.int64)
+        if distance <= 0:
+            logger.error('Provided distances should be greater than 0')
+            if flatten:
+                Pn = sp.array([])
+            else:
+                Pn = sp.array([sp.array([]) for i in range(0, len(pores))])
+            return Pn.astype(sp.int64)
+        # Create kdTree objects
         kd = sptl.cKDTree(self['pore.coords'])
-        if distance == 0:
-            pass
-        elif distance > 0:
-            Pn = kd.query_ball_point(self['pore.coords'][pores], r=distance)
+        kd_pores = sptl.cKDTree(self['pore.coords'][pores])
+        # Perform search
+        Pn = kd_pores.query_ball_tree(kd, r=distance)
+        # Sort the indices in each list
+        [Pn[i].sort() for i in range(0, sp.size(pores))]
+        if flatten:  # Convert list of lists to a flat nd-array
+            temp = []
+            [temp.extend(Ps) for Ps in Pn]
+            Pn = sp.unique(temp)
+            if excl_self:  # Remove inputs if necessary
+                Pn = Pn[~sp.in1d(Pn, pores)]
+        else:  # Convert list of lists to an nd-array of nd-arrays
+            if excl_self:  # Remove inputs if necessary
+                [Pn[i].remove(pores[i]) for i in range(0, sp.size(pores))]
+            temp = []
+            [temp.append(sp.array(Pn[i])) for i in range(0, sp.size(pores))]
+            Pn = sp.array(temp)
+        if Pn.dtype == float:
+            Pn = Pn.astype(sp.int64)
         return Pn
 
     def extend(self, pore_coords=[], throat_conns=[], labels=[]):


Test patch:
diff --git a/test/unit/Network/GenericNetworkTest.py b/test/unit/Network/GenericNetworkTest.py
index f35c8f163..9e2d5db53 100644
--- a/test/unit/Network/GenericNetworkTest.py
+++ b/test/unit/Network/GenericNetworkTest.py
@@ -176,3 +176,33 @@ class GenericNetworkTest:
         a = net.check_network_health()
         assert sp.size(a['bidirectional_throats']) == 1
         assert sp.size(a['duplicate_throats']) == 0
+
+    def test_find_nearby_pores_distance_1(self):
+        a = self.net.find_nearby_pores(pores=[0, 1], distance=1)
+        b = self.net.find_neighbor_pores(pores=[0, 1], flatten=False)
+        assert sp.all([sp.all(a[i] == b[i]) for i in range(0, len(a))])
+
+    def test_find_nearby_pores_distance_2(self):
+        a = self.net.find_nearby_pores(pores=[0, 1], distance=2)
+        assert sp.all([sp.size(a[i]) for i in [0, 1]] == [10, 14])
+
+    def test_find_nearby_pores_distance_0(self):
+        a = self.net.find_nearby_pores(pores=[0, 1], distance=0)
+        assert sp.shape(a) == (2, 0)
+
+    def test_find_nearby_pores_distance_1_flattened(self):
+        a = self.net.find_nearby_pores(pores=[0, 1], distance=1, flatten=True)
+        b = self.net.find_neighbor_pores(pores=[0, 1])
+        assert sp.all(a == b)
+
+    def test_find_nearby_pores_distance_2_flattened(self):
+        a = self.net.find_nearby_pores(pores=[0, 1], distance=2, flatten=True)
+        assert sp.size(a) == 15
+
+    def test_find_nearby_pores_distance_2_flattened_inclself(self):
+        a = self.net.find_nearby_pores(pores=[0, 1],
+                                       distance=2,
+                                       flatten=True,
+                                       excl_self=False)
+        assert sp.size(a) == 17
+        assert sp.all(sp.in1d([0, 1], a))