diff --git a/environment.yml b/environment.yml index c34fbd4ff..40643c6e8 100644 --- a/environment.yml +++ b/environment.yml @@ -29,6 +29,7 @@ dependencies: - coloredlogs - cmocean - pip + - pygeos - pip: - motuclient - roaring-landmask diff --git a/opendrift/readers/basereader/unstructured.py b/opendrift/readers/basereader/unstructured.py index c6c7da3f9..8a5873991 100644 --- a/opendrift/readers/basereader/unstructured.py +++ b/opendrift/readers/basereader/unstructured.py @@ -159,6 +159,10 @@ 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 @@ -166,6 +170,13 @@ def __nearest_ckdtree__(self, idx, x, y): 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) diff --git a/opendrift/readers/reader_netCDF_CF_unstructured.py b/opendrift/readers/reader_netCDF_CF_unstructured.py index 617e6acd7..5fc5d892e 100644 --- a/opendrift/readers/reader_netCDF_CF_unstructured.py +++ b/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") diff --git a/opendrift/readers/unstructured/shyfem.py b/opendrift/readers/unstructured/shyfem.py index e28f98fda..72dcf0087 100644 --- a/opendrift/readers/unstructured/shyfem.py +++ b/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")