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:
else:
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,
point: Record,
current_best: Record | None = None,
best_distance: float = inf,
) -> tuple[Record | None, float]:
"""Find the nearest Record within the KDTree to a query Record"""
if not self.split:
for p in self.points:
dist = point.distance(p)
......@@ -98,25 +112,25 @@ class KDTree:
return current_best, best_distance
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
)
if (
point.distance(self._get_partition_record(point))
< best_distance
):
current_best, best_distance = self.child_right.query(
current_best, best_distance = self.child_right._query(
point, current_best, best_distance
)
else:
current_best, best_distance = self.child_right.query(
current_best, best_distance = self.child_right._query(
point, current_best, best_distance
)
if (
point.distance(self._get_partition_record(point))
< best_distance
):
current_best, best_distance = self.child_left.query(
current_best, best_distance = self.child_left._query(
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