[pyfr] 20/32: Improve the point sampler.

[Ghislain Vaillant]

This is an automated email from the git hooks/post-receive script.

ghisvail-guest pushed a commit to branch master
in repository pyfr.

commit 5eb1e515b1f6a9feb1c5ccca536956143c67d3a3
Author: Freddie Witherden <freddie at witherden.org>
Date:   Sun Mar 27 11:18:57 2016 -0700

    Improve the point sampler.
    
    The sampler plugin will now use scipi.spatial.cKDTree if
    available in order to accelerate the point location phase.  This
    should allow the sampler to be used with a larger number of points.
---
 doc/src/user_guide.rst  |  8 ++++++-
 pyfr/plugins/sampler.py | 63 ++++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/doc/src/user_guide.rst b/doc/src/user_guide.rst
index 92b3aa2..afe24cf 100644
--- a/doc/src/user_guide.rst
+++ b/doc/src/user_guide.rst
@@ -855,7 +855,13 @@ Example::
 ^^^^^^^^^^^^^^^^^^^^^
 
 Periodically samples specific points in the volume and writes them out
-to a CSV file. Parameterised with
+to a CSV file.  The plugin actually samples the solution point
+closest to each sample point a slight discrepancy in the output
+sampling locations is to be expected.  A nearest-neighbour search is
+used to locate the closest solution point to the sample point.  The
+location process automatically takes advantage of
+`scipy.spatial.cKDTree <http://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.cKDTree.html>`_
+where available.  Parameterised with
 
 1. ``nsteps`` --- sample every ``nsteps``:
 
diff --git a/pyfr/plugins/sampler.py b/pyfr/plugins/sampler.py
index 9327641..1df2436 100644
--- a/pyfr/plugins/sampler.py
+++ b/pyfr/plugins/sampler.py
@@ -8,19 +8,53 @@ from pyfr.mpiutil import get_comm_rank_root, get_mpi
 from pyfr.plugins.base import BasePlugin, init_csv
 
 
-def _closest_upt(etypes, eupts, p):
-    # Compute the distances between each point and p
-    dists = [np.linalg.norm(e - p, axis=2) for e in eupts]
+def _closest_upts_bf(etypes, eupts, pts):
+    for p in pts:
+        # Compute the distances between each point and p
+        dists = [np.linalg.norm(e - p, axis=2) for e in eupts]
 
-    # Get the index of the closest point to p for each element type
-    amins = [np.unravel_index(np.argmin(d), d.shape) for d in dists]
+        # Get the index of the closest point to p for each element type
+        amins = [np.unravel_index(np.argmin(d), d.shape) for d in dists]
 
-    # Dereference to get the actual distances and locations
-    dmins = [d[a] for d, a in zip(dists, amins)]
-    plocs = [e[a] for e, a in zip(eupts, amins)]
+        # Dereference to get the actual distances and locations
+        dmins = [d[a] for d, a in zip(dists, amins)]
+        plocs = [e[a] for e, a in zip(eupts, amins)]
 
-    # Find the minimum across all element types
-    return min(zip(dmins, plocs, etypes, amins))
+        # Find the minimum across all element types
+        yield min(zip(dmins, plocs, etypes, amins))
+
+
+def _closest_upts_kd(etypes, eupts, pts):
+    from scipy.spatial import cKDTree
+
+    # Flatten the physical location arrays
+    feupts = [e.reshape(-1, e.shape[-1]) for e in eupts]
+
+    # For each element type construct a KD-tree of the upt locations
+    trees = [cKDTree(f) for f in feupts]
+
+    for p in pts:
+        # Query the distance/index of the closest upt to p
+        dmins, amins = zip(*[t.query(p) for t in trees])
+
+        # Unravel the indices
+        amins = [np.unravel_index(i, e.shape[:2])
+                 for i, e in zip(amins, eupts)]
+
+        # Dereference to obtain the precise locations
+        plocs = [e[a] for e, a in zip(eupts, amins)]
+
+        # Reduce across element types
+        yield min(zip(dmins, plocs, etypes, amins))
+
+
+def _closest_upts(etypes, eupts, pts):
+    try:
+        # Attempt to use a KD-tree based approach
+        yield from _closest_upts_kd(etypes, eupts, pts)
+    except ImportError:
+        # Otherwise fall back to brute force
+        yield from _closest_upts_bf(etypes, eupts, pts)
 
 
 class SamplerPlugin(BasePlugin):
@@ -50,11 +84,12 @@ class SamplerPlugin(BasePlugin):
         # Physical location of the solution points
         plocs = [p.swapaxes(1, 2) for p in intg.system.ele_ploc_upts]
 
-        for p in self.pts:
-            # Find the nearest point in our partition
-            cp = _closest_upt(intg.system.ele_types, plocs, p)
+        # Locate the closest solution points in our partition
+        closest = _closest_upts(intg.system.ele_types, plocs, self.pts)
 
-            # Reduce over all partitions
+        # Process these points
+        for cp in closest:
+            # Reduce cp over all partitions
             mcp, mrank = comm.allreduce(cp, op=get_mpi('minloc'))
 
             # Store the rank responsible along with the info

-- 
Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/debian-science/packages/pyfr.git