feat: quadtree and octtree

GeoSpatialTools/distance_metrics.py

+"""
+Functions for computing navigational information. Can be used to add
+navigational information to DataFrames.
+"""
+
+from math import acos, asin, atan2, cos, sin, degrees, radians, sqrt
+
+
+def gcd_slc(
+    lon0: float,
+    lat0: float,
+    lon1: float,
+    lat1: float,
+) -> float:
+    """
+    Compute great circle distance on earth surface between two locations.
+
+    Parameters
+    ----------
+    lon0 : float
+        Longitude of position 0
+    lat0 : float
+        Latitude of position 0
+    lon1 : float
+        Longitude of position 1
+    lat1 : float
+        Latitude of position 1
+
+    Returns
+    -------
+    dist : float
+        Great circle distance between position 0 and position 1.
+
+    """
+    if abs(lat0 - lat1) <= 1e-6 and abs(lon0 - lon1) <= 1e-6:
+        return 0
+
+    r_earth = 6371
+
+    # Convert to radians
+    lat0, lat1, lon0, lon1 = map(radians, [lat0, lat1, lon0, lon1])
+
+    return r_earth * acos(
+        sin(lat0) * sin(lat1) + cos(lat0) * cos(lat1) * cos(lon1 - lon0)
+    )
+
+
+def haversine(
+    lon0: float,
+    lat0: float,
+    lon1: float,
+    lat1: float,
+) -> float:
+    """
+    Compute Haversine distance between two points.
+
+    Parameters
+    ----------
+    lon0 : float
+        Longitude of position 0
+    lat0 : float
+        Latitude of position 0
+    lon1 : float
+        Longitude of position 1
+    lat1 : float
+        Latitude of position 1
+
+    Returns
+    -------
+    dist : float
+        Haversine distance between position 0 and position 1.
+
+    """
+    lat0, lat1, dlon, dlat = map(
+        radians, [lat0, lat1, lon1 - lon0, lat1 - lat0]
+    )
+    if abs(dlon) < 1e-6 and abs(dlat) < 1e-6:
+        return 0
+
+    r_earth = 6371
+
+    a = sin(dlat / 2) ** 2 + cos(lat0) * cos(lat1) * sin(dlon / 2) ** 2
+    c = 2 * asin(sqrt(a))
+    return c * r_earth
+
+
+def bearing(
+    lon0: float,
+    lat0: float,
+    lon1: float,
+    lat1: float,
+) -> float:
+    """
+    Compute the bearing of a track from (lon0, lat0) to (lon1, lat1).
+
+    Duplicated from geo-py
+
+    Parameters
+    ----------
+    lon0 : float,
+        Longitude of start point
+    lat0 : float,
+        Latitude of start point
+    lon1 : float,
+        Longitude of target point
+    lat1 : float,
+        Latitude of target point
+
+    Returns
+    -------
+    bearing : float
+        The bearing from point (lon0, lat0) to point (lon1, lat1) in degrees.
+    """
+    lon0, lat0, lon1, lat1 = map(radians, [lon0, lat0, lon1, lat1])
+
+    dlon = lon1 - lon0
+    numerator = sin(dlon) * cos(lat1)
+    denominator = cos(lat0) * sin(lat1) - (sin(lat0) * cos(lat1) * cos(dlon))
+
+    theta = atan2(numerator, denominator)
+    theta_deg = (degrees(theta) + 360) % 360
+    return theta_deg
+
+
+def destination(
+    lon: float, lat: float, bearing: float, distance: float
+) -> tuple[float, float]:
+    """
+    Compute destination of a great circle path.
+
+    Compute the destination of a track started from 'lon', 'lat', with
+    'bearing'. Distance is in units of km.
+
+    Duplicated from geo-py
+
+    Parameters
+    ----------
+    lon : float
+        Longitude of initial position
+    lat : float
+        Latitude of initial position
+    bearing : float
+        Direction of track
+    distance : float
+        Distance to travel
+
+    Returns
+    -------
+    destination : tuple[float, float]
+        Longitude and Latitude of final position
+    """
+    lon, lat = radians(lon), radians(lat)
+    radians_bearing = radians(bearing)
+    r_earth = 6371
+    delta = distance / r_earth
+
+    lat2 = asin(
+        sin(lat) * cos(delta) + cos(lat) * sin(delta) * cos(radians_bearing)
+    )
+    numerator = sin(radians_bearing) * sin(delta) * cos(lat)
+    denominator = cos(delta) - sin(lat) * sin(lat2)
+
+    lon2 = lon + atan2(numerator, denominator)
+
+    lon2_deg = (degrees(lon2) + 540) % 360 - 180
+    lat2_deg = degrees(lat2)
+
+    return lon2_deg, lat2_deg
+
+
+def midpoint(
+    lon0: float,
+    lat0: float,
+    lon1: float,
+    lat1: float,
+) -> tuple[float, float]:
+    """
+    Compute the midpoint of a great circle track
+
+    Parameters
+    ----------
+    lon0 : float
+        Longitude of position 0
+    lat0 : float
+        Latitude of position 0
+    lon1 : float
+        Longitude of position 1
+    lat1 : float
+        Latitude of position 1
+
+    Returns
+    -------
+    lon, lat
+        Positions of midpoint between position 0 and position 1
+
+    """
+    bear = bearing(lon0, lat0, lon1, lat1)
+    dist = haversine(lon0, lat0, lon1, lat1)
+
+    return destination(lon0, lat0, bear, dist / 2)

GeoSpatialTools/quadtree.py

+"""
+Constuctors for QuadTree classes that can decrease the number of comparisons
+for detecting nearby records for example
+"""
+
+from datetime import datetime
+from .distance_metrics import haversine
+
+
+class Record:
+    """
+    ICOADS Record class
+
+    Parameters
+    ----------
+    x : float
+        Horizontal coordinate
+    y : float
+        Vertical coordinate
+    datetime : datetime | None
+        Datetime of the record
+    uid : str | None
+        Unique Identifier
+    """
+
+    def __init__(
+        self,
+        lon: float,
+        lat: float,
+        datetime: datetime | None = None,
+        uid: str | None = None,
+        **data,
+    ) -> None:
+        self.lon = lon
+        self.lat = lat
+        self.datetime = datetime
+        self.uid = uid
+        for var, val in data.items():
+            setattr(self, var, val)
+        return None
+
+    def __str__(self) -> str:
+        return f"Record(x = {self.lon}, y = {self.lat}, datetime = {self.datetime}, uid = {self.uid})"
+
+    def __eq__(self, other: object) -> bool:
+        return (
+            isinstance(other, Record)
+            and self.lon == other.lon
+            and self.lat == other.lat
+            and self.datetime == other.datetime
+            and (not (self.uid or other.uid) or self.uid == other.uid)
+        )
+
+
+class Rectangle:
+    """
+    A simple Rectangle class
+
+    Parameters
+    ----------
+    x : float
+        Horizontal centre of the rectangle
+    y : float
+        Vertical centre of the rectangle
+    w : float
+        Width of the rectangle
+    h : float
+        Height of the rectangle
+    """
+
+    def __init__(
+        self,
+        lon: float,
+        lat: float,
+        lon_range: float,
+        lat_range: float,
+    ) -> None:
+        self.lon = lon
+        self.lat = lat
+        self.lon_range = lon_range
+        self.lat_range = lat_range
+
+    def __str__(self) -> str:
+        return f"Rectangle(x = {self.lon}, y = {self.lat}, w = {self.lon_range}, h = {self.lat_range})"
+
+    def __eq__(self, other: object) -> bool:
+        return (
+            isinstance(other, Rectangle)
+            and self.lon == other.lon
+            and self.lat == other.lat
+            and self.lon_range == other.lon_range
+            and self.lat_range == other.lat_range
+        )
+
+    def contains(self, point: Record) -> bool:
+        """Test if a point is contained within the Rectangle"""
+        return (
+            point.lon <= self.lon + self.lon_range / 2
+            and point.lon >= self.lon - self.lon_range / 2
+            and point.lat <= self.lat + self.lat_range / 2
+            and point.lat >= self.lat - self.lat_range / 2
+        )
+
+    def intersects(self, other: object) -> bool:
+        """Test if another Rectangle object intersects this Rectangle"""
+        return isinstance(other, Rectangle) and not (
+            self.lon - self.lon_range / 2 > other.lon + other.lon_range / 2
+            or self.lon + self.lon_range / 2 < other.lon - other.lon_range / 2
+            or self.lat - self.lat_range / 2 > other.lat + other.lat_range / 2
+            or self.lat + self.lat_range / 2 < other.lat - other.lat_range / 2
+        )
+
+    def nearby(
+        self,
+        point: Record,
+        dist: float,
+    ) -> bool:
+        """Check if point is nearby the Rectangle"""
+        # QUESTION: Is this sufficient? Possibly it is overkill
+        corner_dist = max(
+            haversine(
+                self.lon,
+                self.lat,
+                self.lon + self.lon_range / 2,
+                self.lat + self.lat_range / 2,
+            ),
+            haversine(
+                self.lon,
+                self.lat,
+                self.lon + self.lon_range / 2,
+                self.lat - self.lat_range / 2,
+            ),
+        )
+        if (self.lat + self.lat_range / 2) * (
+            self.lat - self.lat_range / 2
+        ) < 0:
+            corner_dist = max(
+                corner_dist,
+                haversine(self.lon, self.lat, self.lon + self.lon_range / 2, 0),
+            )
+        return (
+            haversine(self.lon, self.lat, point.lon, point.lat)
+            <= dist + corner_dist
+        )
+
+
+class QuadTree:
+    """
+    A Simple QuadTree class for PyCOADS
+
+    Parameters
+    ----------
+    boundary : Rectangle
+        The bounding Rectangle of the QuadTree
+    capacity : int
+        The capacity of each cell, if max_depth is set then a cell at the
+        maximum depth may contain more points than the capacity.
+    depth : int
+        The current depth of the cell. Initialises to zero if unset.
+    max_depth : int | None
+        The maximum depth of the QuadTree. If set, this can override the
+        capacity for cells at the maximum depth.
+    """
+
+    def __init__(
+        self,
+        boundary: Rectangle,
+        capacity: int = 5,
+        depth: int = 0,
+        max_depth: int | None = None,
+    ) -> None:
+        self.boundary = boundary
+        self.capacity = capacity
+        self.depth = depth
+        self.max_depth = max_depth
+        self.points: list[Record] = list()
+        self.divided: bool = False
+        return None
+
+    def __str__(self) -> str:
+        indent = "    " * self.depth
+        out = f"{indent}QuadTree:\n"
+        out += f"{indent}- boundary: {self.boundary}\n"
+        out += f"{indent}- capacity: {self.capacity}\n"
+        out += f"{indent}- depth: {self.depth}\n"
+        if self.max_depth:
+            out += f"{indent}- max_depth: {self.max_depth}\n"
+        out += f"{indent}- contents: {self.points}\n"
+        if self.divided:
+            out += f"{indent}- with children:\n"
+            out += f"{self.northwest}"
+            out += f"{self.northeast}"
+            out += f"{self.southwest}"
+            out += f"{self.southeast}"
+        return out
+
+    def divide(self):
+        """Divide the QuadTree"""
+        self.northwest = QuadTree(
+            Rectangle(
+                self.boundary.lon - self.boundary.lon_range / 4,
+                self.boundary.lat + self.boundary.lat_range / 4,
+                self.boundary.lon_range / 2,
+                self.boundary.lat_range / 2,
+            ),
+            capacity=self.capacity,
+            depth=self.depth + 1,
+            max_depth=self.max_depth,
+        )
+        self.northeast = QuadTree(
+            Rectangle(
+                self.boundary.lon + self.boundary.lon_range / 4,
+                self.boundary.lat + self.boundary.lat_range / 4,
+                self.boundary.lon_range / 2,
+                self.boundary.lat_range / 2,
+            ),
+            capacity=self.capacity,
+            depth=self.depth + 1,
+            max_depth=self.max_depth,
+        )
+        self.southwest = QuadTree(
+            Rectangle(
+                self.boundary.lon - self.boundary.lon_range / 4,
+                self.boundary.lat - self.boundary.lat_range / 4,
+                self.boundary.lon_range / 2,
+                self.boundary.lat_range / 2,
+            ),
+            capacity=self.capacity,
+            depth=self.depth + 1,
+            max_depth=self.max_depth,
+        )
+        self.southeast = QuadTree(
+            Rectangle(
+                self.boundary.lon + self.boundary.lon_range / 4,
+                self.boundary.lat - self.boundary.lat_range / 4,
+                self.boundary.lon_range / 2,
+                self.boundary.lat_range / 2,
+            ),
+            capacity=self.capacity,
+            depth=self.depth + 1,
+            max_depth=self.max_depth,
+        )
+        self.divided = True
+
+    def insert(self, point: Record) -> bool:
+        """Insert a point into the QuadTree"""
+        if not self.boundary.contains(point):
+            return False
+        elif self.max_depth and self.depth == self.max_depth:
+            self.points.append(point)
+            return True
+        elif len(self.points) < self.capacity:
+            self.points.append(point)
+            return True
+        else:
+            if not self.divided:
+                self.divide()
+            if self.northwest.insert(point):
+                return True
+            elif self.northeast.insert(point):
+                return True
+            elif self.southwest.insert(point):
+                return True
+            elif self.southeast.insert(point):
+                return True
+            return False
+
+    def query(
+        self,
+        rect: Rectangle,
+        points: list[Record] | None = None,
+    ) -> list[Record]:
+        """Get points that fall in a rectangle"""
+        if not points:
+            points = list()
+        if not self.boundary.intersects(rect):
+            return points
+
+        for point in self.points:
+            if rect.contains(point):
+                points.append(point)
+
+        if self.divided:
+            points = self.northwest.query(rect, points)
+            points = self.northeast.query(rect, points)
+            points = self.southwest.query(rect, points)
+            points = self.southeast.query(rect, points)
+
+        return points
+
+    def nearby_points(
+        self,
+        point: Record,
+        dist: float,
+        points: list[Record] | None = None,
+    ) -> list[Record]:
+        """Get all points that are nearby another point"""
+        if not points:
+            points = list()
+        if not self.boundary.nearby(point, dist):
+            return points
+
+        for test_point in self.points:
+            if (
+                haversine(point.lon, point.lat, test_point.lon, test_point.lat)
+                <= dist
+            ):
+                points.append(test_point)
+
+        if self.divided:
+            points = self.northwest.nearby_points(point, dist, points)
+            points = self.northeast.nearby_points(point, dist, points)
+            points = self.southwest.nearby_points(point, dist, points)
+            points = self.southeast.nearby_points(point, dist, points)
+
+        return points

test/test_octtree.py

+import unittest
+from datetime import datetime, timedelta
+
+from GeoSpatialTools.octtree import (
+    OctTree,
+    SpaceTimeRecord as Record,
+    SpaceTimeRectangle as Rectangle,
+)
+
+
+class TestRect(unittest.TestCase):
+    def test_contains(self):
+        d = datetime(2009, 1, 1, 0, 0)
+        dt = timedelta(days=14)
+        rect = Rectangle(10, 5, d, 20, 10, dt)
+        points: list[Record] = [
+            Record(10, 5, d),
+            Record(20, 10, d + timedelta(hours=4)),
+            Record(20, 10, datetime(2010, 4, 12, 13, 15)),
+            Record(0, 0, d - timedelta(days=6)),
+            Record(12.8, 2.1, d + timedelta(days=-1)),
+            Record(-2, -9.2, d),
+        ]
+        expected = [True, True, False, True, True, False]
+        res = list(map(rect.contains, points))
+        assert res == expected
+
+
+class TestOctTree(unittest.TestCase):
+    def test_divides(self):
+        d = datetime(2023, 3, 24, 12, 0)
+        dt = timedelta(days=1)
+
+        d1 = datetime(2023, 3, 24, 6, 0)
+        d2 = datetime(2023, 3, 24, 18, 0)
+        dt2 = timedelta(hours=12)
+
+        boundary = Rectangle(10, 4, d, 20, 8, dt)
+        otree = OctTree(boundary)
+        expected: list[Rectangle] = [
+            Rectangle(5, 6, d1, 10, 4, dt2),
+            Rectangle(15, 6, d1, 10, 4, dt2),
+            Rectangle(5, 2, d1, 10, 4, dt2),
+            Rectangle(15, 2, d1, 10, 4, dt2),
+            Rectangle(5, 6, d2, 10, 4, dt2),
+            Rectangle(15, 6, d2, 10, 4, dt2),
+            Rectangle(5, 2, d2, 10, 4, dt2),
+            Rectangle(15, 2, d2, 10, 4, dt2),
+        ]
+        otree.divide()
+        res = [
+            otree.northwestback.boundary,
+            otree.northeastback.boundary,
+            otree.southwestback.boundary,
+            otree.southeastback.boundary,
+            otree.northwestfwd.boundary,
+            otree.northeastfwd.boundary,
+            otree.southwestfwd.boundary,
+            otree.southeastfwd.boundary,
+        ]
+        assert res == expected
+
+    def test_insert(self):
+        d = datetime(2023, 3, 24, 12, 0)
+        dt = timedelta(days=10)
+        boundary = Rectangle(10, 4, d, 20, 8, dt)
+        otree = OctTree(boundary, capacity=3)
+        points: list[Record] = [
+            Record(10, 4, d, "main"),
+            Record(12, 1, d + timedelta(hours=3), "main2"),
+            Record(3, 7, d - timedelta(days=3), "main3"),
+            Record(13, 2, d + timedelta(hours=17), "southeastfwd"),
+            Record(3, 6, d - timedelta(days=1), "northwestback"),
+            Record(10, 4, d, "northwestback"),
+            Record(18, 2, d + timedelta(days=23), "not added"),
+            Record(11, 7, d + timedelta(hours=2), "northeastfwd"),
+        ]
+        for point in points:
+            otree.insert(point)
+        assert otree.divided
+        expected = [
+            points[:3],
+            points[4:6],
+            [],
+            [],
+            [],
+            [],
+            [points[-1]],
+            [],
+            [points[3]],
+        ]
+        res = [
+            otree.points,
+            otree.northwestback.points,
+            otree.northeastback.points,
+            otree.southwestback.points,
+            otree.southeastback.points,
+            otree.northwestfwd.points,
+            otree.northeastfwd.points,
+            otree.southwestfwd.points,
+            otree.southeastfwd.points,
+        ]
+        assert res == expected
+
+    def test_query(self):
+        d = datetime(2023, 3, 24, 12, 0)
+        dt = timedelta(days=10)
+        boundary = Rectangle(10, 4, d, 20, 8, dt)
+        otree = OctTree(boundary, capacity=3)
+        points: list[Record] = [
+            Record(10, 4, d, "main"),
+            Record(12, 1, d + timedelta(hours=3), "main2"),
+            Record(3, 7, d - timedelta(days=3), "main3"),
+            Record(13, 2, d + timedelta(hours=17), "southeastfwd"),
+            Record(3, 6, d - timedelta(days=1), "northwestback"),
+            Record(10, 4, d, "northwestback"),
+            Record(18, 2, d + timedelta(days=23), "not added"),
+            Record(11, 7, d + timedelta(hours=2), "northeastfwd"),
+        ]
+        for point in points:
+            otree.insert(point)
+        test_point = Record(11, 6, d + timedelta(hours=4))
+        expected = [Record(11, 7, d + timedelta(hours=2), "northeastfwd")]
+
+        res = otree.nearby_points(
+            test_point, dist=200, t_dist=timedelta(hours=5)
+        )
+
+        assert res == expected
+
+
+if __name__ == "__main__":
+    unittest.main()

test/test_quadtree.py

+import unittest
+from GeoSpatialTools.quadtree import QuadTree, Record, Rectangle
+
+
+class TestRect(unittest.TestCase):
+    def test_contains(self):
+        rect = Rectangle(10, 5, 20, 10)
+        points: list[Record] = [
+            Record(10, 5),
+            Record(20, 10),
+            Record(0, 0),
+            Record(12.8, 2.1),
+            Record(-2, -9.2),
+        ]
+        expected = [True, True, True, True, False]
+        res = list(map(rect.contains, points))
+        assert res == expected
+
+    def test_intersection(self):
+        rect = Rectangle(10, 5, 20, 10)
+        test_rects: list[Rectangle] = [
+            Rectangle(10, 5, 18, 8),
+            Rectangle(25, 5, 9, 12),
+            Rectangle(15, 8, 12, 7),
+        ]
+        expected = [True, False, True]
+        res = list(map(rect.intersects, test_rects))
+        assert res == expected
+
+
+class TestQuadTree(unittest.TestCase):
+    def test_divides(self):
+        boundary = Rectangle(10, 4, 20, 8)
+        qtree = QuadTree(boundary)
+        expected: list[Rectangle] = [
+            Rectangle(5, 6, 10, 4),
+            Rectangle(15, 6, 10, 4),
+            Rectangle(5, 2, 10, 4),
+            Rectangle(15, 2, 10, 4),
+        ]
+        qtree.divide()
+        res = [
+            qtree.northwest.boundary,
+            qtree.northeast.boundary,
+            qtree.southwest.boundary,
+            qtree.southeast.boundary,
+        ]
+        assert res == expected
+
+    def test_insert(self):
+        boundary = Rectangle(10, 4, 20, 8)
+        qtree = QuadTree(boundary, capacity=3)
+        points: list[Record] = [
+            Record(10, 5),
+            Record(19, 1),
+            Record(0, 0),
+            Record(-2, -9.2),
+            Record(12.8, 2.1),
+        ]
+        expected = [
+            points[:3],
+            [],
+            [],
+            [],
+            [points[-1]],
+        ]
+        for point in points:
+            qtree.insert(point)
+        assert qtree.divided
+        res = [
+            qtree.points,
+            qtree.northwest.points,
+            qtree.northeast.points,
+            qtree.southwest.points,
+            qtree.southeast.points,
+        ]
+        assert res == expected
+
+    def test_query(self):
+        boundary = Rectangle(10, 4, 20, 8)
+        qtree = QuadTree(boundary, capacity=3)
+        points: list[Record] = [
+            Record(10, 5),
+            Record(19, 1),
+            Record(0, 0),
+            Record(-2, -9.2),
+            Record(12.8, 2.1),
+        ]
+        test_rect = Rectangle(12.5, 2.5, 1, 1)
+        expected = [Record(12.8, 2.1)]
+
+        for point in points:
+            qtree.insert(point)
+
+        res = qtree.query(test_rect)
+
+        assert res == expected
+
+
+if __name__ == "__main__":
+    unittest.main()