geokdtree.geokdtree API documentation

# geokdtree.geokdtree

View Source

      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)]

class GeoKDTree: View Source

    114class GeoKDTree:
    115    def __init__(self, points: list[tuple]):
    116        """
    117        Function:
    118
    119        - Build a GeoKDTree from a list of latitude and longitude points or an existing KDTree.
    120
    121        Required Arguments:
    122
    123        - `points`
    124            - Type: list of tuples
    125            - What: A list of latitude and longitude points to build the GeoKDTree from
    126            - The points should be in the format [(lat1, lon1), (lat2, lon2), ...].
    127
    128        """
    129        # Store the original points
    130        self.points = points
    131        # Store the KDTree built from the converted points
    132        self.tree = kdtree(
    133            [
    134                __lat_lon_idx_to_xyz_idx__(point[0], point[1], idx)
    135                for idx, point in enumerate(points)
    136            ],
    137            depth=0,
    138            axis_count=3,
    139        )
    140
    141    def closest_idx(self, point: tuple):
    142        """
    143        Function:
    144
    145        - Find the index of the closest point in the GeoKDTree to a given latitude and longitude point.
    146
    147        Required Arguments:
    148
    149        - `point`
    150            - Type: tuple
    151            - What: The latitude and longitude point to find the closest point to
    152
    153        Returns:
    154
    155        - The index of the closest point found in the GeoKDTree to the given latitude and longitude point.
    156        """
    157        best, best_dist = __closest_point_3d__(
    158            self.tree,
    159            __lat_lon_idx_to_xyz_idx__(point[0], point[1]),
    160        )
    161        return best[3]  # Return the index of the closest point
    162
    163    def closest_point(self, point: tuple):
    164        """
    165        Function:
    166
    167        - Find the closest latitude and longitude point in the GeoKDTree to a given latitude and longitude point.
    168
    169        Required Arguments:
    170
    171        - `point`
    172            - Type: tuple
    173            - What: The latitude and longitude point to find the closest point to
    174
    175        Returns:
    176
    177        - The closest latitude and longitude point found in the GeoKDTree to the given latitude and longitude point.
    178        """
    179        return self.points[self.closest_idx(point)]

GeoKDTree(points: list[tuple]) View Source

    115    def __init__(self, points: list[tuple]):
    116        """
    117        Function:
    118
    119        - Build a GeoKDTree from a list of latitude and longitude points or an existing KDTree.
    120
    121        Required Arguments:
    122
    123        - `points`
    124            - Type: list of tuples
    125            - What: A list of latitude and longitude points to build the GeoKDTree from
    126            - The points should be in the format [(lat1, lon1), (lat2, lon2), ...].
    127
    128        """
    129        # Store the original points
    130        self.points = points
    131        # Store the KDTree built from the converted points
    132        self.tree = kdtree(
    133            [
    134                __lat_lon_idx_to_xyz_idx__(point[0], point[1], idx)
    135                for idx, point in enumerate(points)
    136            ],
    137            depth=0,
    138            axis_count=3,
    139        )

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), ...].

points
tree
def closest_idx(self, point: tuple): View Source

    141    def closest_idx(self, point: tuple):
    142        """
    143        Function:
    144
    145        - Find the index of the closest point in the GeoKDTree to a given latitude and longitude point.
    146
    147        Required Arguments:
    148
    149        - `point`
    150            - Type: tuple
    151            - What: The latitude and longitude point to find the closest point to
    152
    153        Returns:
    154
    155        - The index of the closest point found in the GeoKDTree to the given latitude and longitude point.
    156        """
    157        best, best_dist = __closest_point_3d__(
    158            self.tree,
    159            __lat_lon_idx_to_xyz_idx__(point[0], point[1]),
    160        )
    161        return best[3]  # Return the index of the closest point

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.

def closest_point(self, point: tuple): View Source

    163    def closest_point(self, point: tuple):
    164        """
    165        Function:
    166
    167        - Find the closest latitude and longitude point in the GeoKDTree to a given latitude and longitude point.
    168
    169        Required Arguments:
    170
    171        - `point`
    172            - Type: tuple
    173            - What: The latitude and longitude point to find the closest point to
    174
    175        Returns:
    176
    177        - The closest latitude and longitude point found in the GeoKDTree to the given latitude and longitude point.
    178        """
    179        return self.points[self.closest_idx(point)]

Function:

- Find the closest latitude and longitude 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 closest latitude and longitude point found in the GeoKDTree to the given latitude and longitude point.