Merge branch 'ellipse' into 'main'

Add Ellipse comparison to QuadTree and OctTree classes

See merge request josidd/geospatialtools!2

GeoSpatialTools/octtree.py

 from datetime import datetime, timedelta
 from .distance_metrics import haversine
+from .distance_metrics import haversine, destination
+from math import degrees, sqrt
 
 
 class SpaceTimeRecord:
@@ -95,7 +97,7 @@ class SpaceTimeRectangle:
     h : float
         Height of the rectangle (latitude range).
     dt : timedelta
-        Time extent of the rectangle.
+        time extent of the rectangle.
     """
 
     def __init__(
@@ -182,6 +184,11 @@ class SpaceTimeRectangle:
         -------
         bool : True if the point is <= dist + max(dist(centre, corners))
         """
+        if (
+            point.datetime - t_dist > self.datetime + self.dt / 2
+            or point.datetime + t_dist < self.datetime - self.dt / 2
+        ):
+            return False
         # QUESTION: Is this sufficient? Possibly it is overkill
         corner_dist = max(
             haversine(
@@ -207,8 +214,115 @@ class SpaceTimeRectangle:
         return (
             haversine(self.lon, self.lat, point.lon, point.lat)
             <= dist + corner_dist
-            and point.datetime - t_dist <= self.datetime + self.dt / 2
-            and point.datetime + t_dist >= self.datetime - self.dt / 2
+        )
+
+
+class SpaceTimeEllipse:
+    """
+    A simple Ellipse Class for an ellipse on the surface of a sphere.
+
+    Parameters
+    ----------
+    lon : float
+        Horizontal centre of the ellipse
+    lat : float
+        Vertical centre of the ellipse
+    datetime : datetime
+        Datetime centre of the ellipse.
+    a : float
+        Length of the semi-major axis
+    b : float
+        Length of the semi-minor axis
+    theta : float
+        Angle of the semi-major axis from horizontal anti-clockwise in radians
+    dt : timedelta
+        (full) time extent of the ellipse.
+    """
+
+    def __init__(
+        self,
+        lon: float,
+        lat: float,
+        datetime: datetime,
+        a: float,
+        b: float,
+        theta: float,
+        dt: timedelta,
+    ) -> None:
+        self.a = a
+        self.b = b
+        self.lon = lon
+        self.lat = lat
+        self.datetime = datetime
+        self.dt = dt
+        # theta is anti-clockwise angle from horizontal in radians
+        self.theta = theta
+        # bearing is angle clockwise from north in degrees
+        self.bearing = (90 - degrees(self.theta)) % 360
+
+        a2 = self.a * self.a
+        b2 = self.b * self.b
+
+        self.c = sqrt(a2 - b2)
+        self.p1_lon, self.p1_lat = destination(
+            self.lon,
+            self.lat,
+            self.bearing,
+            self.c,
+        )
+        self.p2_lon, self.p2_lat = destination(
+            self.lon,
+            self.lat,
+            (self.bearing - 180) % 360,
+            self.c,
+        )
+
+    def contains(self, point: SpaceTimeRecord) -> bool:
+        """Test if a point is contained within the Ellipse"""
+        return (
+            (
+                haversine(self.p1_lon, self.p1_lat, point.lon, point.lat)
+                + haversine(self.p2_lon, self.p2_lat, point.lon, point.lat)
+            )
+            <= 2 * self.a
+            and point.datetime <= self.datetime + self.dt / 2
+            and point.datetime >= self.datetime - self.dt / 2
+        )
+
+    def nearby_rect(self, rect: SpaceTimeRectangle) -> bool:
+        """Test if a rectangle is near to the Ellipse"""
+        if (
+            rect.datetime - rect.dt / 2 > self.datetime + self.dt / 2
+            or rect.datetime + rect.dt / 2 < self.datetime - self.dt / 2
+        ):
+            return False
+        # TODO: Check corners, and 0 lat
+        corner_dist = max(
+            haversine(
+                rect.lon,
+                rect.lat,
+                rect.lon + rect.lon_range / 2,
+                rect.lat + rect.lat_range / 2,
+            ),
+            haversine(
+                rect.lon,
+                rect.lat,
+                rect.lon + rect.lon_range / 2,
+                rect.lat - rect.lat_range / 2,
+            ),
+        )
+        if (rect.lat + rect.lat_range / 2) * (
+            rect.lat - rect.lat_range / 2
+        ) < 0:
+            corner_dist = max(
+                corner_dist,
+                haversine(rect.lon, rect.lat, rect.lon + rect.lon_range / 2, 0),
+            )
+        return (
+            haversine(self.p1_lon, self.p1_lat, rect.lon, rect.lat)
+            <= corner_dist + self.a
+            and haversine(self.p2_lon, self.p2_lat, rect.lon, rect.lat)
+            <= corner_dist + self.a
         )
 
 
@@ -459,6 +573,33 @@ class OctTree:
 
         return points
 
+    def query_ellipse(
+        self,
+        ellipse: SpaceTimeEllipse,
+        points: SpaceTimeRecords | None = None,
+    ) -> SpaceTimeRecords:
+        """Get points that fall in an ellipse."""
+        if not points:
+            points = SpaceTimeRecords()
+        if not ellipse.nearby_rect(self.boundary):
+            return points
+
+        for point in self.points:
+            if ellipse.contains(point):
+                points.append(point)
+
+        if self.divided:
+            points = self.northwestfwd.query_ellipse(ellipse, points)
+            points = self.northeastfwd.query_ellipse(ellipse, points)
+            points = self.southwestfwd.query_ellipse(ellipse, points)
+            points = self.southeastfwd.query_ellipse(ellipse, points)
+            points = self.northwestback.query_ellipse(ellipse, points)
+            points = self.northeastback.query_ellipse(ellipse, points)
+            points = self.southwestback.query_ellipse(ellipse, points)
+            points = self.southeastback.query_ellipse(ellipse, points)
+
+        return points
+
     def nearby_points(
         self,
         point: SpaceTimeRecord,

GeoSpatialTools/quadtree.py

@@ -4,7 +4,8 @@ for detecting nearby records for example
 """
 
 from datetime import datetime
-from .distance_metrics import haversine
+from .distance_metrics import haversine, destination
+from math import degrees, sqrt
 
 
 class Record:
@@ -13,9 +14,9 @@ class Record:
 
     Parameters
     ----------
-    x : float
+    lon : float
         Horizontal coordinate
-    y : float
+    lat : float
         Vertical coordinate
     datetime : datetime | None
         Datetime of the record
@@ -58,13 +59,13 @@ class Rectangle:
 
     Parameters
     ----------
-    x : float
+    lon : float
         Horizontal centre of the rectangle
-    y : float
+    lat : float
         Vertical centre of the rectangle
-    w : float
+    lon_range : float
         Width of the rectangle
-    h : float
+    lat_range : float
         Height of the rectangle
     """
 
@@ -144,6 +145,97 @@ class Rectangle:
         )
 
 
+class Ellipse:
+    """
+    A simple Ellipse Class for an ellipse on the surface of a sphere.
+
+    Parameters
+    ----------
+    lon : float
+        Horizontal centre of the ellipse
+    lat : float
+        Vertical centre of the ellipse
+    a : float
+        Length of the semi-major axis
+    b : float
+        Length of the semi-minor axis
+    theta : float
+        Angle of the semi-major axis from horizontal anti-clockwise in radians
+    """
+
+    def __init__(
+        self,
+        lon: float,
+        lat: float,
+        a: float,
+        b: float,
+        theta: float,
+    ) -> None:
+        self.a = a
+        self.b = b
+        self.lon = lon
+        self.lat = lat
+        # theta is anti-clockwise angle from horizontal in radians
+        self.theta = theta
+        # bearing is angle clockwise from north in degrees
+        self.bearing = (90 - degrees(self.theta)) % 360
+
+        a2 = self.a * self.a
+        b2 = self.b * self.b
+
+        self.c = sqrt(a2 - b2)
+        self.p1_lon, self.p1_lat = destination(
+            self.lon,
+            self.lat,
+            self.bearing,
+            self.c,
+        )
+        self.p2_lon, self.p2_lat = destination(
+            self.lon,
+            self.lat,
+            (self.bearing - 180) % 360,
+            self.c,
+        )
+
+    def contains(self, point: Record) -> bool:
+        """Test if a point is contained within the Ellipse"""
+        return (
+            haversine(self.p1_lon, self.p1_lat, point.lon, point.lat)
+            + haversine(self.p2_lon, self.p2_lat, point.lon, point.lat)
+        ) <= 2 * self.a
+
+    def nearby_rect(self, rect: Rectangle) -> bool:
+        """Test if a rectangle is near to the Ellipse"""
+        # TODO: Check corners, and 0 lat
+        corner_dist = max(
+            haversine(
+                rect.lon,
+                rect.lat,
+                rect.lon + rect.lon_range / 2,
+                rect.lat + rect.lat_range / 2,
+            ),
+            haversine(
+                rect.lon,
+                rect.lat,
+                rect.lon + rect.lon_range / 2,
+                rect.lat - rect.lat_range / 2,
+            ),
+        )
+        if (rect.lat + rect.lat_range / 2) * (
+            rect.lat - rect.lat_range / 2
+        ) < 0:
+            corner_dist = max(
+                corner_dist,
+                haversine(rect.lon, rect.lat, rect.lon + rect.lon_range / 2, 0),
+            )
+        return (
+            haversine(self.p1_lon, self.p1_lat, rect.lon, rect.lat)
+            <= corner_dist + self.a
+            and haversine(self.p2_lon, self.p2_lat, rect.lon, rect.lat)
+            <= corner_dist + self.a
+        )
+
+
 class QuadTree:
     """
     A Simple QuadTree class for PyCOADS
@@ -288,6 +380,29 @@ class QuadTree:
 
         return points
 
+    def query_ellipse(
+        self,
+        ellipse: Ellipse,
+        points: list[Record] | None = None,
+    ) -> list[Record]:
+        """Get points that fall in an ellipse."""
+        if not points:
+            points = list()
+        if not ellipse.nearby_rect(self.boundary):
+            return points
+
+        for point in self.points:
+            if ellipse.contains(point):
+                points.append(point)
+
+        if self.divided:
+            points = self.northwest.query_ellipse(ellipse, points)
+            points = self.northeast.query_ellipse(ellipse, points)
+            points = self.southwest.query_ellipse(ellipse, points)
+            points = self.southeast.query_ellipse(ellipse, points)
+
+        return points
+
     def nearby_points(
         self,
         point: Record,

test/test_octtree.py

 import unittest
 from datetime import datetime, timedelta
 
+from GeoSpatialTools import haversine
 from GeoSpatialTools.octtree import (
     OctTree,
     SpaceTimeRecord as Record,
     SpaceTimeRectangle as Rectangle,
+    SpaceTimeEllipse as Ellipse,
 )
 
 
@@ -128,6 +130,65 @@ class TestOctTree(unittest.TestCase):
 
         assert res == expected
 
+    def test_ellipse_query(self):
+        d1 = haversine(0, 2.5, 1, 2.5)
+        d2 = haversine(0, 2.5, 0, 3.0)
+        theta = 0
+
+        d = datetime(2023, 3, 24, 12, 0)
+        dt = timedelta(days=10)
+
+        test_datetime = d + timedelta(hours=4)
+        test_timedelta = timedelta(hours=5)
+        ellipse = Ellipse(
+            12.5, 2.5, test_datetime, d1, d2, theta, test_timedelta
+        )
+        # TEST: distint locii
+        assert (ellipse.p1_lon, ellipse.p1_lat) != (
+            ellipse.p2_lon,
+            ellipse.p2_lat,
+        )
+
+        # TEST: Near Boundary Points
+        assert ellipse.contains(Record(13.49, 2.5, test_datetime))
+        assert ellipse.contains(Record(11.51, 2.5, test_datetime))
+        assert ellipse.contains(Record(12.5, 2.99, test_datetime))
+        assert ellipse.contains(Record(12.5, 2.01, test_datetime))
+        assert not ellipse.contains(Record(13.51, 2.5, test_datetime))
+        assert not ellipse.contains(Record(11.49, 2.5, test_datetime))
+        assert not ellipse.contains(Record(12.5, 3.01, test_datetime))
+        assert not ellipse.contains(Record(12.5, 1.99, test_datetime))
+
+        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(12.6, 2.1, d + timedelta(hours=2), "northeastfwd"),
+            Record(13.5, 2.6, test_datetime, "too north of eastern edge"),
+            Record(12.6, 3.0, test_datetime, "too east of northern edge"),
+            # Locii
+            Record(ellipse.p1_lon, ellipse.p1_lat, test_datetime, "locii1"),
+            Record(ellipse.p2_lon, ellipse.p2_lat, test_datetime, "locii2"),
+        ]
+        expected = [
+            Record(12.6, 2.1, d + timedelta(hours=2), "northeastfwd"),
+            Record(ellipse.p1_lon, ellipse.p1_lat, test_datetime, "locii1"),
+            Record(ellipse.p2_lon, ellipse.p2_lat, test_datetime, "locii2"),
+        ]
+
+        for point in points:
+            otree.insert(point)
+
+        res = otree.query_ellipse(ellipse)
+        assert res == expected
+
 
 if __name__ == "__main__":
     unittest.main()

test/test_quadtree.py

+from math import pi
 import unittest
-from GeoSpatialTools.quadtree import QuadTree, Record, Rectangle
+from GeoSpatialTools import haversine
+from GeoSpatialTools.quadtree import QuadTree, Record, Rectangle, Ellipse
 
 
 class TestRect(unittest.TestCase):
@@ -96,6 +98,54 @@ class TestQuadTree(unittest.TestCase):
 
         assert res == expected
 
+    def test_ellipse_query(self):
+        d1 = haversine(0, 2.5, 1, 2.5)
+        d2 = haversine(0, 2.5, 0, 3.0)
+        theta = 0
+
+        ellipse = Ellipse(12.5, 2.5, d1, d2, theta)
+        # TEST: distint locii
+        assert (ellipse.p1_lon, ellipse.p1_lat) != (
+            ellipse.p2_lon,
+            ellipse.p2_lat,
+        )
+
+        # TEST: Near Boundary Points
+        assert ellipse.contains(Record(13.49, 2.5))
+        assert ellipse.contains(Record(11.51, 2.5))
+        assert ellipse.contains(Record(12.5, 2.99))
+        assert ellipse.contains(Record(12.5, 2.01))
+        assert not ellipse.contains(Record(13.51, 2.5))
+        assert not ellipse.contains(Record(11.49, 2.5))
+        assert not ellipse.contains(Record(12.5, 3.01))
+        assert not ellipse.contains(Record(12.5, 1.99))
+
+        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(13.5, 2.6),  # Just North of Eastern edge
+            Record(12.6, 3.0),  # Just East of Northern edge
+            Record(12.8, 2.1),
+            # Locii
+            Record(ellipse.p1_lon, ellipse.p1_lat),
+            Record(ellipse.p2_lon, ellipse.p2_lat),
+        ]
+        expected = [
+            Record(12.8, 2.1),
+            Record(ellipse.p1_lon, ellipse.p1_lat),
+            Record(ellipse.p2_lon, ellipse.p2_lat),
+        ]
+
+        for point in points:
+            qtree.insert(point)
+
+        res = qtree.query_ellipse(ellipse)
+        assert res == expected
+
 
 if __name__ == "__main__":
     unittest.main()