geokdtree.geokdtree

  1from .core import kdtree
  2from math import cos, sin, pi
  3
  4
  5# Speed Utils for the 3D GeoKDTree
  6def __squared_distance_3d__(p1, p2):
  7    """
  8    Function:
  9
 10    - Calculate the squared distance between two 3D points.
 11
 12    Required Arguments:
 13
 14    - `p1`
 15        - Type: tuple
 16        - What: The first point in 3D space
 17    - `p2`
 18        - Type: tuple
 19        - What: The second point in 3D space
 20
 21    Returns:
 22
 23    - The squared distance between the two 3D points.
 24    """
 25    return sum(
 26        [(p1[0] - p2[0]) ** 2, (p1[1] - p2[1]) ** 2, (p1[2] - p2[2]) ** 2]
 27    )
 28
 29
 30def __closest_point_3d__(node, point, best=None, best_dist=float("inf")):
 31    """
 32    Function:
 33
 34    - Find the closest point in a 3d cartesian system using a KDTree.
 35
 36    Required Arguments:
 37
 38    - `node`
 39        - Type: tuple
 40        - What: The node of the KDTree
 41    - `point`
 42        - Type: tuple
 43        - What: The point to find the closest point to
 44
 45    Optional Arguments:
 46
 47    - `best`
 48        - Type: tuple or None
 49        - What: The best point found so far (default is None)
 50    - `best_dist`
 51        - Type: float
 52        - What: The best distance found so far (default is infinity)
 53
 54    Returns:
 55
 56    - The closest point found in the KDTree to the given point.
 57    """
 58    if node == 0:
 59        return best, best_dist
 60    # Get the median node and its distance
 61    median_node = node[0]
 62    median_node_dist = __squared_distance_3d__(point, median_node)
 63    # Update the best point and distance if necessary
 64    if best is None or median_node_dist < best_dist:
 65        best = median_node
 66        best_dist = median_node_dist
 67    # Calculate the difference for node selection given the current axis
 68    axis = node[1]
 69    diff = point[axis] - median_node[axis]
 70    # Choose side to search
 71    close, away = (node[2], node[3]) if diff < 0 else (node[3], node[2])
 72    # Search the close side first
 73    best, best_dist = __closest_point_3d__(close, point, best, best_dist)
 74    # Check the other side if needed
 75    if diff**2 < best_dist:
 76        best, best_dist = __closest_point_3d__(away, point, best, best_dist)
 77    return best, best_dist
 78
 79
 80# Special Serializer to convert lat,lon,index to x,y,z,index
 81def __lat_lon_idx_to_xyz_idx__(
 82    lat: int | float, lon: int | float, idx: int = 0
 83):
 84    """
 85    Function:
 86
 87    - Convert latitude and longitude to Cartesian coordinates (x, y, z) and include an index.
 88
 89    Required Arguments:
 90
 91    - `lat`
 92        - Type: int or float
 93        - What: The latitude in degrees
 94    - `lon`
 95        - Type: int or float
 96        - What: The longitude in degrees
 97
 98    Optional Arguments:
 99
100    - `idx`
101        - Type: int
102        - What: An index to include with the coordinates (default is 0)
103    """
104    lat_rad = lat * pi / 180
105    lon_rad = lon * pi / 180
106    cos_lat = cos(lat_rad)
107    x = cos_lat * cos(lon_rad)
108    y = cos_lat * sin(lon_rad)
109    z = sin(lat_rad)
110    return (x, y, z, idx)
111
112
113class GeoKDTree:
114    def __init__(self, points: list[tuple]):
115        """
116        Function:
117
118        - Build a GeoKDTree from a list of latitude and longitude points or an existing KDTree.
119
120        Required Arguments:
121
122        - `points`
123            - Type: list of tuples
124            - What: A list of latitude and longitude points to build the GeoKDTree from
125            - The points should be in the format [(lat1, lon1), (lat2, lon2), ...].
126
127        """
128        # Store the original points
129        self.points = points
130        # Store the KDTree built from the converted points
131        self.tree = kdtree(
132            [
133                __lat_lon_idx_to_xyz_idx__(point[0], point[1], idx)
134                for idx, point in enumerate(points)
135            ],
136            depth=0,
137            axis_count=3,
138        )
139
140    def closest_idx(self, point: tuple):
141        """
142        Function:
143
144        - Find the index of the closest point in the GeoKDTree to a given latitude and longitude point.
145
146        Required Arguments:
147
148        - `point`
149            - Type: tuple
150            - What: The latitude and longitude point to find the closest point to
151
152        Returns:
153
154        - The index of the closest point found in the GeoKDTree to the given latitude and longitude point.
155        """
156        best, best_dist = __closest_point_3d__(
157            self.tree,
158            __lat_lon_idx_to_xyz_idx__(point[0], point[1]),
159        )
160        return best[3]  # Return the index of the closest point
161
162    def closest_point(self, point: tuple):
163        """
164        Function:
165
166        - Find the closest latitude and longitude point in the GeoKDTree to a given latitude and longitude point.
167
168        Required Arguments:
169
170        - `point`
171            - Type: tuple
172            - What: The latitude and longitude point to find the closest point to
173
174        Returns:
175
176        - The closest latitude and longitude point found in the GeoKDTree to the given latitude and longitude point.
177        """
178        return self.points[self.closest_idx(point)]