geokdtree.core

  1from operator import itemgetter
  2
  3_GETTERS = [itemgetter(i) for i in range(16)]
  4
  5
  6def kdtree(points, depth, axis_count):
  7    """
  8    Function:
  9
 10    - Build a KDTree from a list of points.
 11
 12    Required Arguments:
 13
 14    - `points`
 15        - Type: list of tuples
 16        - What: A list of points to build the KDTree from
 17
 18    Optional Arguments:
 19
 20    - `depth`
 21        - Type: int
 22        - What: The current depth in the tree (used for axis selection)
 23    - `axis_count`
 24        - Type: int
 25        - What: The number of dimensions in the points (default is 2 for 2D points)
 26
 27    Returns:
 28
 29    - The constructed KDTree as a tuple in the format (point, axis, left, right).
 30    - Where left and right are subtrees.
 31    """
 32    if not points:
 33        return 0
 34    axis = depth % axis_count
 35    getter = _GETTERS[axis] if axis < 16 else itemgetter(axis)
 36    points.sort(key=getter)
 37    median = len(points) // 2
 38    return (
 39        points[median],
 40        axis,
 41        kdtree(points=points[:median], depth=depth + 1, axis_count=axis_count),
 42        kdtree(
 43            points=points[median + 1 :], depth=depth + 1, axis_count=axis_count
 44        ),
 45    )
 46
 47
 48def squared_distance(p1, p2, axis_count=2):
 49    """
 50    Function:
 51
 52    - Calculate the squared distance between two points.
 53
 54    Required Arguments:
 55
 56    - `p1`
 57        - Type: tuple
 58        - What: The first point
 59    - `p2`
 60        - Type: tuple
 61        - What: The second point
 62
 63    Optional Arguments:
 64
 65    - `axis_count`
 66        - Type: int
 67        - What: The number of dimensions in the points
 68        - Default: 2 (for 2D points)
 69
 70    Returns:
 71
 72    - The squared distance between the two points.
 73    """
 74    if axis_count == 2:
 75        dx = p1[0] - p2[0]
 76        dy = p1[1] - p2[1]
 77        return dx * dx + dy * dy
 78    elif axis_count == 3:
 79        dx = p1[0] - p2[0]
 80        dy = p1[1] - p2[1]
 81        dz = p1[2] - p2[2]
 82        return dx * dx + dy * dy + dz * dz
 83    return sum((p1[i] - p2[i]) ** 2 for i in range(axis_count))
 84
 85
 86def closest_point(node, point, best=None, axis_count=2, best_dist=float("inf")):
 87    """
 88    Function:
 89
 90    - Find the closest point in the KDTree to a given point.
 91
 92    Required Arguments:
 93
 94    - `node`
 95        - Type: tuple
 96        - What: The node of the KDTree
 97    - `point`
 98        - Type: tuple
 99        - What: The point to find the closest point to
100
101    Optional Arguments:
102
103    - `best`
104        - Type: tuple or None
105        - What: The best point found so far (default is None)
106    - `axis_count`
107        - Type: int
108        - What: The number of dimensions in the points (default is 2 for 2D points)
109    - `best_dist`
110        - Type: float
111        - What: The best distance found so far (default is infinity)
112
113    Returns:
114
115    - The closest point found in the KDTree to the given point.
116    """
117    if node == 0:
118        return best, best_dist
119    # Get the median node and its distance
120    median_node = node[0]
121    median_node_dist = squared_distance(
122        point, median_node, axis_count=axis_count
123    )
124    # Update the best point and distance if necessary
125    if best is None or median_node_dist < best_dist:
126        best = median_node
127        best_dist = median_node_dist
128    # Calculate the difference for node selection given the current axis
129    axis = node[1]
130    diff = point[axis] - median_node[axis]
131    # Choose side to search
132    close, away = (node[2], node[3]) if diff < 0 else (node[3], node[2])
133    # Search the close side first
134    best, best_dist = closest_point(
135        close,
136        point,
137        best,
138        axis_count=axis_count,
139        best_dist=best_dist,
140    )
141    # Check the other side if needed
142    if diff**2 < best_dist:
143        best, best_dist = closest_point(
144            away,
145            point,
146            best,
147            axis_count=axis_count,
148            best_dist=best_dist,
149        )
150    return best, best_dist
151
152
153def closest_point_per_quadrant_2d(node, point):
154    """
155    Function:
156
157    - Find the closest point in each of the 4 quadrants (ne, nw, se, sw)
158      for 2D Cartesian points.
159
160    Required Arguments:
161
162    - `node`
163        - Type: tuple
164        - What: The root node of the KDTree
165    - `point`
166        - Type: tuple
167        - What: The 2D query point (x, y)
168
169    Returns:
170
171    - A dictionary mapping 'ne', 'nw', 'se', 'sw' to their closest point tuple or None.
172    """
173    bests = [None, None, None, None]
174    best_dists = [float("inf"), float("inf"), float("inf"), float("inf")]
175
176    px, py = point[0], point[1]
177
178    def _search(curr, active_mask):
179        if curr == 0 or curr is None or active_mask == 0:
180            return
181
182        median_node = curr[0]
183        mx, my = median_node[0], median_node[1]
184        dist = (px - mx) * (px - mx) + (py - my) * (py - my)
185
186        is_east = mx >= px
187        is_west = mx <= px
188        is_north = my >= py
189        is_south = my <= py
190
191        if is_north and is_east and (active_mask & 1):
192            if dist < best_dists[0]:
193                bests[0] = median_node
194                best_dists[0] = dist
195
196        if is_north and is_west and (active_mask & 2):
197            if dist < best_dists[1]:
198                bests[1] = median_node
199                best_dists[1] = dist
200
201        if is_south and is_east and (active_mask & 4):
202            if dist < best_dists[2]:
203                bests[2] = median_node
204                best_dists[2] = dist
205
206        if is_south and is_west and (active_mask & 8):
207            if dist < best_dists[3]:
208                bests[3] = median_node
209                best_dists[3] = dist
210
211        axis = curr[1]
212        diff = point[axis] - median_node[axis]
213
214        if diff < 0:
215            close, away = curr[2], curr[3]
216            close_mask = active_mask
217            away_mask = (
218                active_mask & (1 | 4) if axis == 0 else active_mask & (1 | 2)
219            )
220        else:
221            close, away = curr[3], curr[2]
222            close_mask = active_mask
223            away_mask = (
224                active_mask & (2 | 8) if axis == 0 else active_mask & (4 | 8)
225            )
226
227        _search(close, close_mask)
228
229        if away_mask:
230            diff_sq = diff * diff
231            max_away = 0.0
232            if (away_mask & 1) and best_dists[0] > max_away:
233                max_away = best_dists[0]
234            if (away_mask & 2) and best_dists[1] > max_away:
235                max_away = best_dists[1]
236            if (away_mask & 4) and best_dists[2] > max_away:
237                max_away = best_dists[2]
238            if (away_mask & 8) and best_dists[3] > max_away:
239                max_away = best_dists[3]
240
241            if diff_sq < max_away:
242                _search(away, away_mask)
243
244    _search(node, 15)
245    return {"ne": bests[0], "nw": bests[1], "se": bests[2], "sw": bests[3]}