geokdtree.core
1from math import pi, cos, sin 2 3 4def kdtree(points, depth=0, axis_count=2): 5 """ 6 Function: 7 8 - Build a KDTree from a list of points. 9 10 Required Arguments: 11 12 - `points` 13 - Type: list of tuples 14 - What: A list of points to build the KDTree from 15 16 Optional Arguments: 17 18 - `depth` 19 - Type: int 20 - What: The current depth in the tree (used for axis selection) 21 - `axis_count` 22 - Type: int 23 - What: The number of dimensions in the points (default is 2 for 2D points) 24 25 Returns: 26 27 - The constructed KDTree as a tuple in the format (point, axis, left, right). 28 - Where left and right are subtrees. 29 """ 30 if not points: 31 return 0 32 axis = depth % axis_count 33 points.sort(key=lambda p: p[axis]) 34 median = len(points) // 2 35 return ( 36 points[median], 37 axis, 38 kdtree(points[:median], depth + 1), 39 kdtree(points[median + 1 :], depth + 1), 40 ) 41 42 43def squared_distance(p1, p2, axis_count=2): 44 """ 45 Function: 46 47 - Calculate the squared distance between two points. 48 49 Required Arguments: 50 51 - `p1` 52 - Type: tuple 53 - What: The first point 54 - `p2` 55 - Type: tuple 56 - What: The second point 57 58 Optional Arguments: 59 60 - `axis_count` 61 - Type: int 62 - What: The number of dimensions in the points 63 - Default: 2 (for 2D points) 64 65 Returns: 66 67 - The squared distance between the two points. 68 """ 69 return sum([(p1[i] - p2[i]) ** 2 for i in range(axis_count)]) 70 71 72def closest_point( 73 node, point, depth=0, best=None, axis_count=2, best_dist=float("inf") 74): 75 """ 76 Function: 77 78 - Find the closest point in the KDTree to a given point. 79 80 Required Arguments: 81 82 - `node` 83 - Type: tuple 84 - What: The node of the KDTree 85 - `point` 86 - Type: tuple 87 - What: The point to find the closest point to 88 89 Optional Arguments: 90 91 - `depth` 92 - Type: int 93 - What: The current depth in the tree (used for axis selection) 94 - `best` 95 - Type: tuple or None 96 - What: The best point found so far (default is None) 97 - `axis_count` 98 - Type: int 99 - What: The number of dimensions in the points (default is 2 for 2D points) 100 - `best_dist` 101 - Type: float 102 - What: The best distance found so far (default is infinity) 103 104 Returns: 105 106 - The closest point found in the KDTree to the given point. 107 """ 108 if node == 0: 109 return best, best_dist 110 # Get the median node and its distance 111 median_node = node[0] 112 median_node_dist = squared_distance( 113 point, median_node, axis_count=axis_count 114 ) 115 # Update the best point and distance if necessary 116 if best is None or median_node_dist < best_dist: 117 best = median_node 118 best_dist = median_node_dist 119 # Calculate the difference for node selection given the current axis 120 axis = node[1] 121 diff = point[axis] - median_node[axis] 122 # Choose side to search 123 close, away = (node[2], node[3]) if diff < 0 else (node[3], node[2]) 124 # Search the close side first 125 best, best_dist = closest_point( 126 close, 127 point, 128 depth + 1, 129 best, 130 axis_count=axis_count, 131 best_dist=best_dist, 132 ) 133 # Check the other side if needed 134 if diff**2 < best_dist: 135 best, best_dist = closest_point( 136 away, 137 point, 138 depth + 1, 139 best, 140 axis_count=axis_count, 141 best_dist=best_dist, 142 ) 143 return best, best_dist 144 145 146def squared_distance_3d(p1, p2): 147 """ 148 Function: 149 150 - Calculate the squared distance between two 3D points. 151 152 Required Arguments: 153 154 - `p1` 155 - Type: tuple 156 - What: The first point in 3D space 157 - `p2` 158 - Type: tuple 159 - What: The second point in 3D space 160 161 Returns: 162 163 - The squared distance between the two 3D points. 164 """ 165 return sum( 166 [(p1[0] - p2[0]) ** 2, (p1[1] - p2[1]) ** 2, (p1[2] - p2[2]) ** 2] 167 ) 168 169 170def closest_point_3d(node, point, depth=0, best=None, best_dist=float("inf")): 171 """ 172 Function: 173 174 - Find the closest point in the KDTree to a given point. 175 176 Required Arguments: 177 178 - `node` 179 - Type: tuple 180 - What: The node of the KDTree 181 - `point` 182 - Type: tuple 183 - What: The point to find the closest point to 184 185 Optional Arguments: 186 187 - `depth` 188 - Type: int 189 - What: The current depth in the tree (used for axis selection) 190 - `best` 191 - Type: tuple or None 192 - What: The best point found so far (default is None) 193 - `best_dist` 194 - Type: float 195 - What: The best distance found so far (default is infinity) 196 197 Returns: 198 199 - The closest point found in the KDTree to the given point. 200 """ 201 if node == 0: 202 return best, best_dist 203 # Get the median node and its distance 204 median_node = node[0] 205 median_node_dist = squared_distance_3d(point, median_node) 206 # Update the best point and distance if necessary 207 if best is None or median_node_dist < best_dist: 208 best = median_node 209 best_dist = median_node_dist 210 # Calculate the difference for node selection given the current axis 211 axis = node[1] 212 diff = point[axis] - median_node[axis] 213 # Choose side to search 214 close, away = (node[2], node[3]) if diff < 0 else (node[3], node[2]) 215 # Search the close side first 216 best, best_dist = closest_point_3d(close, point, depth + 1, best, best_dist) 217 # Check the other side if needed 218 if diff**2 < best_dist: 219 best, best_dist = closest_point_3d( 220 away, point, depth + 1, best, best_dist 221 ) 222 return best, best_dist 223 224 225class KDTree: 226 def __init__(self, points): 227 """ 228 Function: 229 230 - Build a KDTree from a list of points. 231 232 Required Arguments: 233 234 - `points` 235 - Type: list of tuples 236 - What: A list of points to build the KDTree from 237 238 Returns: 239 240 - A KDTree object that can be used to find the closest point to a given point. 241 """ 242 self.tree = kdtree(points, axis_count=len(points[0])) 243 244 def closest_point(self, point): 245 """ 246 Function: 247 248 - Find the closest point in the KDTree to a given point. 249 250 Required Arguments: 251 252 - `point` 253 - Type: tuple 254 - What: The point to find the closest point to 255 256 Returns: 257 258 - The closest point found in the KDTree to the given point. 259 """ 260 return closest_point(self.tree, point)[ 261 0 262 ] # Return only the point, not the distance 263 264 265class GeoKDTree: 266 def __init__(self, points: list[tuple]): 267 """ 268 Function: 269 270 - Build a GeoKDTree from a list of latitude and longitude points or an existing KDTree. 271 272 Required Arguments: 273 274 - `points` 275 - Type: list of tuples 276 - What: A list of latitude and longitude points to build the GeoKDTree from 277 - The points should be in the format [(lat1, lon1), (lat2, lon2), ...]. 278 279 Returns: 280 281 - A GeoKDTree object that can be used to find the closest point to a given latitude and longitude. 282 - The points should be in the format [(lat1, lon1), (lat2, lon2), ...]. 283 """ 284 self.tree = kdtree( 285 [ 286 GeoKDTree.lat_lon_idx_to_xyz_idx(point[0], point[1], idx) 287 for idx, point in enumerate(points) 288 ], 289 axis_count=3, 290 ) 291 292 def closest_idx(self, point: tuple): 293 """ 294 Function: 295 296 - Find the index of the closest point in the GeoKDTree to a given latitude and longitude point. 297 298 Required Arguments: 299 300 - `point` 301 - Type: tuple 302 - What: The latitude and longitude point to find the closest point to 303 304 Returns: 305 306 - The index of the closest point found in the GeoKDTree to the given latitude and longitude point. 307 """ 308 return closest_point_3d( 309 self.tree, 310 GeoKDTree.lat_lon_idx_to_xyz_idx(point[0], point[1]), 311 )[0][ 312 3 313 ] # Return the point [0] and the index of the point [3] 314 315 @staticmethod 316 def lat_lon_idx_to_xyz_idx( 317 lat: int | float, lon: int | float, idx: int = 0 318 ): 319 """ 320 Function: 321 322 - Convert latitude and longitude to Cartesian coordinates (x, y, z) and include an index. 323 324 Required Arguments: 325 326 - `lat` 327 - Type: int or float 328 - What: The latitude in degrees 329 - `lon` 330 - Type: int or float 331 - What: The longitude in degrees 332 333 Optional Arguments: 334 335 - `idx` 336 - Type: int 337 - What: An index to include with the coordinates (default is 0) 338 """ 339 lat_rad = lat * pi / 180 340 lon_rad = lon * pi / 180 341 cos_lat = cos(lat_rad) 342 x = cos_lat * cos(lon_rad) 343 y = cos_lat * sin(lon_rad) 344 z = sin(lat_rad) 345 return (x, y, z, idx)
def
kdtree(points, depth=0, axis_count=2):
5def kdtree(points, depth=0, axis_count=2): 6 """ 7 Function: 8 9 - Build a KDTree from a list of points. 10 11 Required Arguments: 12 13 - `points` 14 - Type: list of tuples 15 - What: A list of points to build the KDTree from 16 17 Optional Arguments: 18 19 - `depth` 20 - Type: int 21 - What: The current depth in the tree (used for axis selection) 22 - `axis_count` 23 - Type: int 24 - What: The number of dimensions in the points (default is 2 for 2D points) 25 26 Returns: 27 28 - The constructed KDTree as a tuple in the format (point, axis, left, right). 29 - Where left and right are subtrees. 30 """ 31 if not points: 32 return 0 33 axis = depth % axis_count 34 points.sort(key=lambda p: p[axis]) 35 median = len(points) // 2 36 return ( 37 points[median], 38 axis, 39 kdtree(points[:median], depth + 1), 40 kdtree(points[median + 1 :], depth + 1), 41 )
Function:
- Build a KDTree from a list of points.
Required Arguments:
points- Type: list of tuples
- What: A list of points to build the KDTree from
Optional Arguments:
depth- Type: int
- What: The current depth in the tree (used for axis selection)
axis_count- Type: int
- What: The number of dimensions in the points (default is 2 for 2D points)
Returns:
- The constructed KDTree as a tuple in the format (point, axis, left, right).
- Where left and right are subtrees.
def
squared_distance(p1, p2, axis_count=2):
44def squared_distance(p1, p2, axis_count=2): 45 """ 46 Function: 47 48 - Calculate the squared distance between two points. 49 50 Required Arguments: 51 52 - `p1` 53 - Type: tuple 54 - What: The first point 55 - `p2` 56 - Type: tuple 57 - What: The second point 58 59 Optional Arguments: 60 61 - `axis_count` 62 - Type: int 63 - What: The number of dimensions in the points 64 - Default: 2 (for 2D points) 65 66 Returns: 67 68 - The squared distance between the two points. 69 """ 70 return sum([(p1[i] - p2[i]) ** 2 for i in range(axis_count)])
Function:
- Calculate the squared distance between two points.
Required Arguments:
p1- Type: tuple
- What: The first point
p2- Type: tuple
- What: The second point
Optional Arguments:
axis_count- Type: int
- What: The number of dimensions in the points
- Default: 2 (for 2D points)
Returns:
- The squared distance between the two points.
def
closest_point(node, point, depth=0, best=None, axis_count=2, best_dist=inf):
73def closest_point( 74 node, point, depth=0, best=None, axis_count=2, best_dist=float("inf") 75): 76 """ 77 Function: 78 79 - Find the closest point in the KDTree to a given point. 80 81 Required Arguments: 82 83 - `node` 84 - Type: tuple 85 - What: The node of the KDTree 86 - `point` 87 - Type: tuple 88 - What: The point to find the closest point to 89 90 Optional Arguments: 91 92 - `depth` 93 - Type: int 94 - What: The current depth in the tree (used for axis selection) 95 - `best` 96 - Type: tuple or None 97 - What: The best point found so far (default is None) 98 - `axis_count` 99 - Type: int 100 - What: The number of dimensions in the points (default is 2 for 2D points) 101 - `best_dist` 102 - Type: float 103 - What: The best distance found so far (default is infinity) 104 105 Returns: 106 107 - The closest point found in the KDTree to the given point. 108 """ 109 if node == 0: 110 return best, best_dist 111 # Get the median node and its distance 112 median_node = node[0] 113 median_node_dist = squared_distance( 114 point, median_node, axis_count=axis_count 115 ) 116 # Update the best point and distance if necessary 117 if best is None or median_node_dist < best_dist: 118 best = median_node 119 best_dist = median_node_dist 120 # Calculate the difference for node selection given the current axis 121 axis = node[1] 122 diff = point[axis] - median_node[axis] 123 # Choose side to search 124 close, away = (node[2], node[3]) if diff < 0 else (node[3], node[2]) 125 # Search the close side first 126 best, best_dist = closest_point( 127 close, 128 point, 129 depth + 1, 130 best, 131 axis_count=axis_count, 132 best_dist=best_dist, 133 ) 134 # Check the other side if needed 135 if diff**2 < best_dist: 136 best, best_dist = closest_point( 137 away, 138 point, 139 depth + 1, 140 best, 141 axis_count=axis_count, 142 best_dist=best_dist, 143 ) 144 return best, best_dist
Function:
- Find the closest point in the KDTree to a given point.
Required Arguments:
node- Type: tuple
- What: The node of the KDTree
point- Type: tuple
- What: The point to find the closest point to
Optional Arguments:
depth- Type: int
- What: The current depth in the tree (used for axis selection)
best- Type: tuple or None
- What: The best point found so far (default is None)
axis_count- Type: int
- What: The number of dimensions in the points (default is 2 for 2D points)
best_dist- Type: float
- What: The best distance found so far (default is infinity)
Returns:
- The closest point found in the KDTree to the given point.
def
squared_distance_3d(p1, p2):
147def squared_distance_3d(p1, p2): 148 """ 149 Function: 150 151 - Calculate the squared distance between two 3D points. 152 153 Required Arguments: 154 155 - `p1` 156 - Type: tuple 157 - What: The first point in 3D space 158 - `p2` 159 - Type: tuple 160 - What: The second point in 3D space 161 162 Returns: 163 164 - The squared distance between the two 3D points. 165 """ 166 return sum( 167 [(p1[0] - p2[0]) ** 2, (p1[1] - p2[1]) ** 2, (p1[2] - p2[2]) ** 2] 168 )
Function:
- Calculate the squared distance between two 3D points.
Required Arguments:
p1- Type: tuple
- What: The first point in 3D space
p2- Type: tuple
- What: The second point in 3D space
Returns:
- The squared distance between the two 3D points.
def
closest_point_3d(node, point, depth=0, best=None, best_dist=inf):
171def closest_point_3d(node, point, depth=0, best=None, best_dist=float("inf")): 172 """ 173 Function: 174 175 - Find the closest point in the KDTree to a given point. 176 177 Required Arguments: 178 179 - `node` 180 - Type: tuple 181 - What: The node of the KDTree 182 - `point` 183 - Type: tuple 184 - What: The point to find the closest point to 185 186 Optional Arguments: 187 188 - `depth` 189 - Type: int 190 - What: The current depth in the tree (used for axis selection) 191 - `best` 192 - Type: tuple or None 193 - What: The best point found so far (default is None) 194 - `best_dist` 195 - Type: float 196 - What: The best distance found so far (default is infinity) 197 198 Returns: 199 200 - The closest point found in the KDTree to the given point. 201 """ 202 if node == 0: 203 return best, best_dist 204 # Get the median node and its distance 205 median_node = node[0] 206 median_node_dist = squared_distance_3d(point, median_node) 207 # Update the best point and distance if necessary 208 if best is None or median_node_dist < best_dist: 209 best = median_node 210 best_dist = median_node_dist 211 # Calculate the difference for node selection given the current axis 212 axis = node[1] 213 diff = point[axis] - median_node[axis] 214 # Choose side to search 215 close, away = (node[2], node[3]) if diff < 0 else (node[3], node[2]) 216 # Search the close side first 217 best, best_dist = closest_point_3d(close, point, depth + 1, best, best_dist) 218 # Check the other side if needed 219 if diff**2 < best_dist: 220 best, best_dist = closest_point_3d( 221 away, point, depth + 1, best, best_dist 222 ) 223 return best, best_dist
Function:
- Find the closest point in the KDTree to a given point.
Required Arguments:
node- Type: tuple
- What: The node of the KDTree
point- Type: tuple
- What: The point to find the closest point to
Optional Arguments:
depth- Type: int
- What: The current depth in the tree (used for axis selection)
best- Type: tuple or None
- What: The best point found so far (default is None)
best_dist- Type: float
- What: The best distance found so far (default is infinity)
Returns:
- The closest point found in the KDTree to the given point.
class
KDTree:
226class KDTree: 227 def __init__(self, points): 228 """ 229 Function: 230 231 - Build a KDTree from a list of points. 232 233 Required Arguments: 234 235 - `points` 236 - Type: list of tuples 237 - What: A list of points to build the KDTree from 238 239 Returns: 240 241 - A KDTree object that can be used to find the closest point to a given point. 242 """ 243 self.tree = kdtree(points, axis_count=len(points[0])) 244 245 def closest_point(self, point): 246 """ 247 Function: 248 249 - Find the closest point in the KDTree to a given point. 250 251 Required Arguments: 252 253 - `point` 254 - Type: tuple 255 - What: The point to find the closest point to 256 257 Returns: 258 259 - The closest point found in the KDTree to the given point. 260 """ 261 return closest_point(self.tree, point)[ 262 0 263 ] # Return only the point, not the distance
KDTree(points)
227 def __init__(self, points): 228 """ 229 Function: 230 231 - Build a KDTree from a list of points. 232 233 Required Arguments: 234 235 - `points` 236 - Type: list of tuples 237 - What: A list of points to build the KDTree from 238 239 Returns: 240 241 - A KDTree object that can be used to find the closest point to a given point. 242 """ 243 self.tree = kdtree(points, axis_count=len(points[0]))
Function:
- Build a KDTree from a list of points.
Required Arguments:
points- Type: list of tuples
- What: A list of points to build the KDTree from
Returns:
- A KDTree object that can be used to find the closest point to a given point.
def
closest_point(self, point):
245 def closest_point(self, point): 246 """ 247 Function: 248 249 - Find the closest point in the KDTree to a given point. 250 251 Required Arguments: 252 253 - `point` 254 - Type: tuple 255 - What: The point to find the closest point to 256 257 Returns: 258 259 - The closest point found in the KDTree to the given point. 260 """ 261 return closest_point(self.tree, point)[ 262 0 263 ] # Return only the point, not the distance
Function:
- Find the closest point in the KDTree to a given point.
Required Arguments:
point- Type: tuple
- What: The point to find the closest point to
Returns:
- The closest point found in the KDTree to the given point.
class
GeoKDTree:
266class GeoKDTree: 267 def __init__(self, points: list[tuple]): 268 """ 269 Function: 270 271 - Build a GeoKDTree from a list of latitude and longitude points or an existing KDTree. 272 273 Required Arguments: 274 275 - `points` 276 - Type: list of tuples 277 - What: A list of latitude and longitude points to build the GeoKDTree from 278 - The points should be in the format [(lat1, lon1), (lat2, lon2), ...]. 279 280 Returns: 281 282 - A GeoKDTree object that can be used to find the closest point to a given latitude and longitude. 283 - The points should be in the format [(lat1, lon1), (lat2, lon2), ...]. 284 """ 285 self.tree = kdtree( 286 [ 287 GeoKDTree.lat_lon_idx_to_xyz_idx(point[0], point[1], idx) 288 for idx, point in enumerate(points) 289 ], 290 axis_count=3, 291 ) 292 293 def closest_idx(self, point: tuple): 294 """ 295 Function: 296 297 - Find the index of the closest point in the GeoKDTree to a given latitude and longitude point. 298 299 Required Arguments: 300 301 - `point` 302 - Type: tuple 303 - What: The latitude and longitude point to find the closest point to 304 305 Returns: 306 307 - The index of the closest point found in the GeoKDTree to the given latitude and longitude point. 308 """ 309 return closest_point_3d( 310 self.tree, 311 GeoKDTree.lat_lon_idx_to_xyz_idx(point[0], point[1]), 312 )[0][ 313 3 314 ] # Return the point [0] and the index of the point [3] 315 316 @staticmethod 317 def lat_lon_idx_to_xyz_idx( 318 lat: int | float, lon: int | float, idx: int = 0 319 ): 320 """ 321 Function: 322 323 - Convert latitude and longitude to Cartesian coordinates (x, y, z) and include an index. 324 325 Required Arguments: 326 327 - `lat` 328 - Type: int or float 329 - What: The latitude in degrees 330 - `lon` 331 - Type: int or float 332 - What: The longitude in degrees 333 334 Optional Arguments: 335 336 - `idx` 337 - Type: int 338 - What: An index to include with the coordinates (default is 0) 339 """ 340 lat_rad = lat * pi / 180 341 lon_rad = lon * pi / 180 342 cos_lat = cos(lat_rad) 343 x = cos_lat * cos(lon_rad) 344 y = cos_lat * sin(lon_rad) 345 z = sin(lat_rad) 346 return (x, y, z, idx)
GeoKDTree(points: list[tuple])
267 def __init__(self, points: list[tuple]): 268 """ 269 Function: 270 271 - Build a GeoKDTree from a list of latitude and longitude points or an existing KDTree. 272 273 Required Arguments: 274 275 - `points` 276 - Type: list of tuples 277 - What: A list of latitude and longitude points to build the GeoKDTree from 278 - The points should be in the format [(lat1, lon1), (lat2, lon2), ...]. 279 280 Returns: 281 282 - A GeoKDTree object that can be used to find the closest point to a given latitude and longitude. 283 - The points should be in the format [(lat1, lon1), (lat2, lon2), ...]. 284 """ 285 self.tree = kdtree( 286 [ 287 GeoKDTree.lat_lon_idx_to_xyz_idx(point[0], point[1], idx) 288 for idx, point in enumerate(points) 289 ], 290 axis_count=3, 291 )
Function:
- Build a GeoKDTree from a list of latitude and longitude points or an existing KDTree.
Required Arguments:
points- Type: list of tuples
- What: A list of latitude and longitude points to build the GeoKDTree from
- The points should be in the format [(lat1, lon1), (lat2, lon2), ...].
Returns:
- A GeoKDTree object that can be used to find the closest point to a given latitude and longitude.
- The points should be in the format [(lat1, lon1), (lat2, lon2), ...].
def
closest_idx(self, point: tuple):
293 def closest_idx(self, point: tuple): 294 """ 295 Function: 296 297 - Find the index of the closest point in the GeoKDTree to a given latitude and longitude point. 298 299 Required Arguments: 300 301 - `point` 302 - Type: tuple 303 - What: The latitude and longitude point to find the closest point to 304 305 Returns: 306 307 - The index of the closest point found in the GeoKDTree to the given latitude and longitude point. 308 """ 309 return closest_point_3d( 310 self.tree, 311 GeoKDTree.lat_lon_idx_to_xyz_idx(point[0], point[1]), 312 )[0][ 313 3 314 ] # Return the point [0] and the index of the point [3]
Function:
- Find the index of the closest point in the GeoKDTree to a given latitude and longitude point.
Required Arguments:
point- Type: tuple
- What: The latitude and longitude point to find the closest point to
Returns:
- The index of the closest point found in the GeoKDTree to the given latitude and longitude point.
@staticmethod
def
lat_lon_idx_to_xyz_idx(lat: int | float, lon: int | float, idx: int = 0):
316 @staticmethod 317 def lat_lon_idx_to_xyz_idx( 318 lat: int | float, lon: int | float, idx: int = 0 319 ): 320 """ 321 Function: 322 323 - Convert latitude and longitude to Cartesian coordinates (x, y, z) and include an index. 324 325 Required Arguments: 326 327 - `lat` 328 - Type: int or float 329 - What: The latitude in degrees 330 - `lon` 331 - Type: int or float 332 - What: The longitude in degrees 333 334 Optional Arguments: 335 336 - `idx` 337 - Type: int 338 - What: An index to include with the coordinates (default is 0) 339 """ 340 lat_rad = lat * pi / 180 341 lon_rad = lon * pi / 180 342 cos_lat = cos(lat_rad) 343 x = cos_lat * cos(lon_rad) 344 y = cos_lat * sin(lon_rad) 345 z = sin(lat_rad) 346 return (x, y, z, idx)
Function:
- Convert latitude and longitude to Cartesian coordinates (x, y, z) and include an index.
Required Arguments:
lat- Type: int or float
- What: The latitude in degrees
lon- Type: int or float
- What: The longitude in degrees
Optional Arguments:
idx- Type: int
- What: An index to include with the coordinates (default is 0)