Initial index search improvement

[angus-g]

From what I can tell, the index search routines (both rectilinear and curvilinear, but in particular the latter) use an iterative algorithm to find the cell indices containing a given point. In the rectilinear case, this is just a step search in each direction orthogonally, which is probably bounded by max(xdim, ydim) iterations. For a curvilinear grid, however, the algorithm equates to finding the target position in the local coordinates of each cell, stepping to an adjacent cell if it ends up outside [0,1]x[0,1]. I assume this is then bounded by xdim * ydim iterations, however each iteration involves a pretty expensive calculation for the local coordinate transformation.

I've been playing around with data on a curvilinear grid of about 6500x3000 points, and the time to calculate the initial indices for all particles is a fairly significant cost. To get around this, I used a k-d tree, e.g. Scipy's cKDTree. I actually used a C implementation, and called it as part of search_indices_curvilinear, but upon further reflection, I don't think this is necessary. Given the previous cell indices, the search algorithm should only take a couple of steps to find the new cell, assuming the particle didn't move too far.

To help out the index search, doing a 4-neighbour query for each point (hopefully) gives us the indices of the cell bounds. I found that this worked for most of the domain, but some areas required an extra iteration or two of the regular index search to converge on the right place. I dropped the maximum iteration limit maxIterSearch to 3 and never exceeded it. Anyway, perhaps this could be a useful addition to ParticleSet initialisation: the k-d tree is fairly cheap to construct and query, and would probably give a sufficiently "close enough" answer even for nested domains to drastically reduce the number of search iterations.

[erikvansebille]

Thanks for these thoughts, @angus-g! Do you also have some code that you want to share?

Just to confirm: the index-searching is only slow for the initialisation of the particles, right? Because, as we keep track of particle.xi and particle.yi (their grid index locations), updating locations should be fairly fast (and as you say should only take a few iterations), perhaps the maxIterSearch = 1e6 in https://github.com/OceanParcels/parcels/blob/master/parcels/field.py#L794 and https://github.com/OceanParcels/parcels/blob/master/parcels/include/index_search.h#L306 are complete overkill and should be reduced to 10?

Note that we also recently merged https://github.com/OceanParcels/parcels/pull/940, which implements a local search too in SciPy mode. Did you use that already?

Finally, note that @delandmeterp already raised this point about tree-based initialisation in #599; also it has't been further implemented yet. SO very keen to see your implementation!

[CKehl]

Erik, see comment about trees and neighbours in Slack. Sounds certainly practical, but you just can't guarantee that this reliably provides you with the N nearest neighbours. It's the basic problem that all ANN (-> approximate nearest-neighbour) algorithms such as libANN, libFLANN and SciPy have (-> scipy using flann in the background): it just works on fairly coherently distributed samples. So, if approximate is enough for what we need: fine. If you really need reliably THE N NEAREST neighbours, we need to look at the natural-neighbour algorithm, computing the Voronoi diagram for all particles - and that's quite a bit slower.

just food for thought.

[angus-g]

I haven't seen the local search from #940, though I don't use scipy mode anyway.

As far as I know (at least with the pykdtree library), this isn't an approximate lookup (but you can perform an approximate lookup if desired). Since we're using the iterative algorithm anyway, it doesn't seem to be a big problem if we don't land in the exact cell: the aim is just to speed up the initial index search.

My tree search just looks like this, I call populate_indices() after making a ParticleSet. It's sufficient to let me drop the max curvilinear index search indices to 3. Obviously this doesn't handle a lot of functionality that it'd need to, but it works for my use case at the moment.

diff --git a/parcels/particleset.py b/parcels/particleset.py
index ef214659..51729bca 100644
--- a/parcels/particleset.py
+++ b/parcels/particleset.py
@@ -9,11 +9,13 @@ import xarray as xr
 from operator import attrgetter
 import progressbar

+from pykdtree.kdtree import KDTree
+
 from parcels.compiler import GNUCompiler
 from parcels.field import Field
 from parcels.field import NestedField
 from parcels.field import SummedField
-from parcels.grid import GridCode
+from parcels.grid import GridCode, CurvilinearGrid
 from parcels.kernel import Kernel
 from parcels.kernels.advection import AdvectionRK4
 from parcels.particle import JITParticle
@@ -300,6 +302,26 @@ class ParticleSet(object):
         else:
             raise ValueError("Latitude and longitude required for generating ParticleSet")

+    def populate_indices(self):
+        """Pre-populate guesses of particle xi/yi indices using a kdtree"""
+
+        if self.fieldset is None:
+            # we need a fieldset for its gridset
+            return
+
+        for i, grid in enumerate(self.fieldset.gridset.grids):
+            # don't bother on non-curvilinear grids
+            if not isinstance(grid, CurvilinearGrid):
+                continue
+
+            tree = KDTree(np.stack((grid.lon.flat, grid.lat.flat), axis=-1))
+            # datatype needs to match between tree and query
+            _, idx = tree.query(np.stack((self.particle_data['lon'], self.particle_data['lat']), axis=-1).astype(np.float32))
+            yi, xi = np.unravel_index(idx, grid.lon.shape)
+
+            self.particle_data['xi'][:, i] = xi
+            self.particle_data['yi'][:, i] = yi
+
     def data_accessor(self):
         return ParticleAccessor(self)

[CKehl]

K, nice - yeah, now I see that for this purpose, the tree-based search works well as intended.

[erikvansebille]

Thank you @angus-g, for providing this code. Would you mind proposing this as a PR? In that way, you'll also get the appropriate attribution for this idea ;-)