perf(reader): replace the scipy.spatial.cKDTree with pygeos.STRtree to find the nearest node and element

environment.yml

@@ -29,6 +29,7 @@ dependencies:
   - coloredlogs
   - cmocean
   - pip
+  - pygeos
   - pip:
     - motuclient
     - roaring-landmask

opendrift/readers/basereader/unstructured.py

@@ -159,13 +159,24 @@ def _build_ckdtree_(self, x, y):
         P = np.vstack((x, y)).T
         return cKDTree(P)
 
+    def _build_strtree_(self, x, y):
+        from pygeos import STRtree, points
+        return STRtree(points(x, y))
+
     def __nearest_ckdtree__(self, idx, x, y):
         """
         Return index of nearest point in cKDTree
         """
         q = np.vstack((x, y)).T
         return idx.query(q, k = 1, workers = self.PARALLEL_WORKERS)[1]
 
+    def __nearest_strtree__(self, idx, x, y):
+        """
+        Return index of nearest point in STRtree
+        """
+        from pygeos import points
+        return idx.nearest(points(x, y))[1]
+
     @staticmethod
     def __nearest_rtree__(idx, x, y):
         """
@@ -177,10 +188,10 @@ def _nearest_node_(self, x, y):
         """
         Return nearest node (id) for x and y
         """
-        return self.__nearest_ckdtree__(self.nodes_idx, x, y)
+        return self.__nearest_strtree__(self.nodes_idx, x, y)
 
     def _nearest_face_(self, xc, yc):
         """
         Return nearest element or face (id) for xc and yc
         """
-        return self.__nearest_ckdtree__(self.faces_idx, xc, yc)
+        return self.__nearest_strtree__(self.faces_idx, xc, yc)

opendrift/readers/reader_netCDF_CF_unstructured.py

@@ -171,10 +171,10 @@ def __init__(self, filename=None, name=None, proj4=None):
 
         self.timer_start("build index")
         logger.debug("building index of nodes..")
-        self.nodes_idx = self._build_ckdtree_(self.x, self.y)
+        self.nodes_idx = self._build_strtree_(self.x, self.y)
 
         logger.debug("building index of faces..")
-        self.faces_idx = self._build_ckdtree_(self.xc, self.yc)
+        self.faces_idx = self._build_strtree_(self.xc, self.yc)
 
         self.timer_end("build index")
 

opendrift/readers/unstructured/shyfem.py

@@ -125,7 +125,7 @@ def __init__(self, filename=None, name=None):
 
         self.timer_start("build index")
         logger.debug("building index of nodes..")
-        self.nodes_idx = self._build_ckdtree_(self.x, self.y)
+        self.nodes_idx = self._build_strtree_(self.x, self.y)
         self.timer_end("build index")
 
         self.timer_end("open dataset")