Merge branch 'master' into output_chunk_explanation_tutorial

acoat · Jan 12, 2024 · a5c44bb · a5c44bb
2 parents b1d7c20 + 2d17966
commit a5c44bb
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/parcels/particleset.py b/parcels/particleset.py
@@ -450,11 +450,14 @@ def populate_indices(self):
                     continue
 
                 tree_data = np.stack((grid.lon.flat, grid.lat.flat), axis=-1)
-                tree = KDTree(tree_data)
+                IN = np.all(~np.isnan(tree_data), axis=1)
+                tree = KDTree(tree_data[IN, :])
                 # stack all the particle positions for a single query
                 pts = np.stack((self.particledata.data['lon'], self.particledata.data['lat']), axis=-1)
                 # query datatype needs to match tree datatype
-                _, idx = tree.query(pts.astype(tree_data.dtype))
+                _, idx_nan = tree.query(pts.astype(tree_data.dtype))
+
+                idx = np.where(IN)[0][idx_nan]
                 yi, xi = np.unravel_index(idx, grid.lon.shape)
 
                 self.particledata.data['xi'][:, i] = xi