Commit 9496c161 authored by Joseph Siddons's avatar Joseph Siddons
Browse files

fix(kdtree): account for lon wrapping

parent 56498568
...@@ -82,13 +82,27 @@ class KDTree: ...@@ -82,13 +82,27 @@ class KDTree:
else: else:
return self.child_right.delete(point) return self.child_right.delete(point)
def query( def query(self, point) -> tuple[Record | None, float]:
"""Find the nearest Record within the KDTree to a _query Record"""
if point.lon < 0:
point2 = Record(point.lon + 360, point.lat)
else:
point2 = Record(point.lon - 360, point.lat)
r1, d1 = self._query(point)
r2, d2 = self._query(point2)
if d1 <= d2:
r = r1
else:
r = r2
return r, point.distance(r)
def _query(
self, self,
point: Record, point: Record,
current_best: Record | None = None, current_best: Record | None = None,
best_distance: float = inf, best_distance: float = inf,
) -> tuple[Record | None, float]: ) -> tuple[Record | None, float]:
"""Find the nearest Record within the KDTree to a query Record"""
if not self.split: if not self.split:
for p in self.points: for p in self.points:
dist = point.distance(p) dist = point.distance(p)
...@@ -98,25 +112,25 @@ class KDTree: ...@@ -98,25 +112,25 @@ class KDTree:
return current_best, best_distance return current_best, best_distance
if getattr(point, self.variable) < self.partition_value: if getattr(point, self.variable) < self.partition_value:
current_best, best_distance = self.child_left.query( current_best, best_distance = self.child_left._query(
point, current_best, best_distance point, current_best, best_distance
) )
if ( if (
point.distance(self._get_partition_record(point)) point.distance(self._get_partition_record(point))
< best_distance < best_distance
): ):
current_best, best_distance = self.child_right.query( current_best, best_distance = self.child_right._query(
point, current_best, best_distance point, current_best, best_distance
) )
else: else:
current_best, best_distance = self.child_right.query( current_best, best_distance = self.child_right._query(
point, current_best, best_distance point, current_best, best_distance
) )
if ( if (
point.distance(self._get_partition_record(point)) point.distance(self._get_partition_record(point))
< best_distance < best_distance
): ):
current_best, best_distance = self.child_left.query( current_best, best_distance = self.child_left._query(
point, current_best, best_distance point, current_best, best_distance
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment