geokdtree.geokdtree

  1from .core import kdtree
  2from math import cos, sin, pi
  3
  4# Speed Utils for the 3D GeoKDTree
  5def __squared_distance_3d__(p1, p2):
  6    """
  7    Function:
  8
  9    - Calculate the squared distance between two 3D points.
 10
 11    Required Arguments:
 12
 13    - `p1`
 14        - Type: tuple
 15        - What: The first point in 3D space
 16    - `p2`
 17        - Type: tuple
 18        - What: The second point in 3D space
 19
 20    Returns:
 21
 22    - The squared distance between the two 3D points.
 23    """
 24    return sum(
 25        [(p1[0] - p2[0]) ** 2, (p1[1] - p2[1]) ** 2, (p1[2] - p2[2]) ** 2]
 26    )
 27
 28
 29def __closest_point_3d__(node, point, best=None, best_dist=float("inf")):
 30    """
 31    Function:
 32
 33    - Find the closest point in a 3d cartesian system using a KDTree.
 34
 35    Required Arguments:
 36
 37    - `node`
 38        - Type: tuple
 39        - What: The node of the KDTree
 40    - `point`
 41        - Type: tuple
 42        - What: The point to find the closest point to
 43
 44    Optional Arguments:
 45
 46    - `best`
 47        - Type: tuple or None
 48        - What: The best point found so far (default is None)
 49    - `best_dist`
 50        - Type: float
 51        - What: The best distance found so far (default is infinity)
 52
 53    Returns:
 54
 55    - The closest point found in the KDTree to the given point.
 56    """
 57    if node == 0:
 58        return best, best_dist
 59    # Get the median node and its distance
 60    median_node = node[0]
 61    median_node_dist = __squared_distance_3d__(point, median_node)
 62    # Update the best point and distance if necessary
 63    if best is None or median_node_dist < best_dist:
 64        best = median_node
 65        best_dist = median_node_dist
 66    # Calculate the difference for node selection given the current axis
 67    axis = node[1]
 68    diff = point[axis] - median_node[axis]
 69    # Choose side to search
 70    close, away = (node[2], node[3]) if diff < 0 else (node[3], node[2])
 71    # Search the close side first
 72    best, best_dist = __closest_point_3d__(close, point, best, best_dist)
 73    # Check the other side if needed
 74    if diff**2 < best_dist:
 75        best, best_dist = __closest_point_3d__(
 76            away, point, best, best_dist
 77        )
 78    return best, best_dist
 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)]