feat: add enclose_poles to antimeridian module

Author: gadomski

src/stactools/core/utils/antimeridian.py

@@ -235,3 +235,62 @@ def normalize_multipolygon(multi_polygon: MultiPolygon) -> Optional[MultiPolygon
         return MultiPolygon(polygons)
     else:
         return None
+
+
+def enclose_poles(polygon: Polygon) -> Polygon:
+    """Updates an anti-meridian-crossing polygon to enclose the poles.
+
+    This works by detecting anti-meridian crossings and adding points to extend
+    the geometry up to the north (or down to the south) pole. This is useful for
+    (e.g.) polar-orbiting satellites who have swaths that go over the poles.
+
+    This function will raise a value error if the polygon has any interior rings.
+
+    Args:
+        polygon (:py:class:`shapely.geometry.Polygon`): An input polygon.
+
+    Returns:
+        :py:class:`shapely.geometry.Polygon`: The same polygon, modified to
+        enclose the poles.
+    """
+    if bool(polygon.interiors):
+        raise ValueError("cannot enclose poles if there is an interior ring")
+    coords = list(polygon.exterior.coords)
+    new_coords = list()
+    for start, end in zip(coords, coords[1:]):
+        new_coords.append(start)
+        longitude_delta = end[0] - start[0]
+        if abs(longitude_delta) > 180:
+            latitude_delta = end[1] - start[1]
+            if start[1] > 0:
+                latitude = 90
+            else:
+                latitude = -90
+            if start[0] > 0:
+                split_latitude = round(
+                    start[1]
+                    + (180.0 - start[0]) * latitude_delta / (end[0] + 360 - start[0]),
+                    7,
+                )
+                box = (
+                    (180.0, split_latitude),
+                    (180.0, latitude),
+                    (-180.0, latitude),
+                    (-180.0, split_latitude),
+                )
+                new_coords.extend(box)
+            else:
+                split_latitude = round(
+                    start[1]
+                    + (start[0] + 180.0) * latitude_delta / (start[0] + 360 - end[0]),
+                    7,
+                )
+                box = (
+                    (-180.0, split_latitude),
+                    (-180.0, latitude),
+                    (180.0, latitude),
+                    (180.0, split_latitude),
+                )
+                new_coords.extend(box)
+    new_coords.append(coords[-1])
+    return Polygon(new_coords)

tests/core/utils/test_antimeridian.py

@@ -3,7 +3,6 @@
 import shapely.geometry
 from pystac import Item
 from shapely.geometry import MultiPolygon, Polygon
-
 from stactools.core.utils import antimeridian
 
 
@@ -16,7 +15,7 @@ def test_antimeridian_split() -> None:
         (
             shapely.geometry.box(170, 40, 180, 50),
             shapely.geometry.box(-180, 40, -170, 50),
-        )
+        ),
     )
     for actual, expected in zip(split.geoms, expected.geoms):
         assert actual.exterior.is_ccw
@@ -27,15 +26,15 @@ def test_antimeridian_split() -> None:
     assert split is None
 
     canonical_other_way = Polygon(
-        ((-170, 40), (170, 40), (170, 50), (-170, 50), (-170, 40))
+        ((-170, 40), (170, 40), (170, 50), (-170, 50), (-170, 40)),
     )
     split = antimeridian.split(canonical_other_way)
     assert split
     expected = MultiPolygon(
         (
             shapely.geometry.box(-180, 40, -170, 50),
             shapely.geometry.box(170, 40, 180, 50),
-        )
+        ),
     )
     for actual, expected in zip(split.geoms, expected.geoms):
         assert actual.exterior.is_ccw
@@ -44,7 +43,7 @@ def test_antimeridian_split() -> None:
 
 def test_antimeridian_split_complicated() -> None:
     complicated = Polygon(
-        ((170, 40), (170, 50), (-170, 50), (170, 45), (-170, 40), (170, 40))
+        ((170, 40), (170, 50), (-170, 50), (170, 45), (-170, 40), (170, 40)),
     )
     split = antimeridian.split(complicated)
     assert split
@@ -57,7 +56,7 @@ def test_antimeridian_split_complicated() -> None:
                     (170.0, 45.0),
                     (170.0, 40.0),
                     (180.0, 40.0),
-                ]
+                ],
             ),
             Polygon([(-180.0, 42.5), (-180.0, 40.0), (-170.0, 40.0), (-180.0, 42.5)]),
             Polygon(
@@ -67,10 +66,10 @@ def test_antimeridian_split_complicated() -> None:
                     (170.0, 50.0),
                     (170.0, 45.0),
                     (180.0, 47.5),
-                ]
+                ],
             ),
             Polygon([(-180.0, 50.0), (-180.0, 47.5), (-170.0, 50.0), (-180.0, 50.0)]),
-        )
+        ),
     )
     for actual, expected in zip(split.geoms, expected.geoms):
         assert actual.exterior.is_ccw
@@ -86,7 +85,7 @@ def test_antimeridian_normalize() -> None:
     assert normalized.equals(expected), f"actual={normalized}, expected={expected}"
 
     canonical_other_way = Polygon(
-        ((-170, 40), (170, 40), (170, 50), (-170, 50), (-170, 40))
+        ((-170, 40), (170, 40), (170, 50), (-170, 50), (-170, 40)),
     )
     normalized = antimeridian.normalize(canonical_other_way)
     assert normalized
@@ -129,10 +128,11 @@ def test_item_fix_antimeridian_split() -> None:
         (
             shapely.geometry.box(170, 40, 180, 50),
             shapely.geometry.box(-180, 40, -170, 50),
-        )
+        ),
     )
     for actual, expected in zip(
-        shapely.geometry.shape(fix.geometry).geoms, expected.geoms
+        shapely.geometry.shape(fix.geometry).geoms,
+        expected.geoms,
     ):
         assert actual.equals(expected)
     assert fix.bbox == [170, 40, -170, 50]
@@ -159,7 +159,7 @@ def test_item_fix_antimeridian_multipolygon_ok() -> None:
         (
             shapely.geometry.box(170, 40, 180, 50),
             shapely.geometry.box(-180, 40, -170, 50),
-        )
+        ),
     )
     item = Item(
         "an-id",
@@ -177,7 +177,7 @@ def test_antimeridian_multipolygon() -> None:
         [
             Polygon(((170, 40), (170, 42), (-170, 42), (-170, 40), (170, 40))),
             Polygon(((170, 48), (170, 50), (-170, 50), (-170, 48), (170, 48))),
-        ]
+        ],
     )
     split = antimeridian.split_multipolygon(multi_polygon)
     assert split
@@ -187,7 +187,7 @@ def test_antimeridian_multipolygon() -> None:
             shapely.geometry.box(-180, 40, -170, 42),
             shapely.geometry.box(170, 48, 180, 50),
             shapely.geometry.box(-180, 48, -170, 50),
-        )
+        ),
     )
     for actual, expected in zip(split.geoms, expected.geoms):
         assert actual.exterior.is_ccw
@@ -199,8 +199,30 @@ def test_antimeridian_multipolygon() -> None:
         (
             Polygon(((170, 40), (170, 42), (190, 42), (190, 40), (170, 40))),
             Polygon(((170, 48), (170, 50), (190, 50), (190, 48), (170, 48))),
-        )
+        ),
     )
     for actual, expected in zip(normalized.geoms, expected.geoms):
         assert actual.exterior.is_ccw
         assert actual.equals(expected), f"actual={actual}, expected={expected}"
+
+
+def test_antimeridian_enclose_poles() -> None:
+    before = Polygon(((170, 40), (-170, 50), (-170, -50), (170, -40), (170, 40)))
+    after = antimeridian.enclose_poles(before)
+    assert after == Polygon(
+        (
+            (170, 40),
+            (180, 45),
+            (180, 90),
+            (-180, 90),
+            (-180, 45),
+            (-170, 50),
+            (-170, -50),
+            (-180, -45),
+            (-180, -90),
+            (180, -90),
+            (180, -45),
+            (170, -40),
+            (170, 40),
+        )
+    )