Add polygon-to-morton coverage function

.github/workflows/codspeed.yml

@@ -40,7 +40,7 @@ jobs:
         uses: CodSpeedHQ/action@v4
         with:
           mode: simulation
-          run: pytest benchmarks/test_bench_cpu.py --codspeed
+          run: pytest benchmarks/ --codspeed -v
 
   rust-benchmarks:
     name: Rust benchmarks

Cargo.toml

@@ -8,7 +8,7 @@ license = "BSD-3-Clause"
 
 [lib]
 name = "mortie_rustie"
-crate-type = ["cdylib"]
+crate-type = ["cdylib", "rlib"]
 path = "src_rust/src/lib.rs"
 
 [dependencies]
@@ -30,3 +30,8 @@ strip = true
 name = "morton_bench"
 harness = false
 path = "src_rust/benches/morton_bench.rs"
+
+[[bench]]
+name = "coverage_bench"
+harness = false
+path = "src_rust/benches/coverage_bench.rs"

README.md

@@ -56,6 +56,23 @@ expanded = np.union1d(cells, border)
 
 All input indices must be at the same order. The function returns only the new border cells, not the input cells themselves.
 
+## Polygon Coverage
+
+`morton_coverage` computes the set of morton indices that completely cover a polygon defined by lat/lon vertices. The algorithm builds a contiguous boundary via great-circle interpolation, classifies inner/outer buffer cells using connected components + point-in-polygon testing, then recursively fills the interior.
+
+```python
+import mortie
+
+# Define polygon vertices (lat, lon in degrees)
+lats = [40.0, 40.0, 50.0, 50.0]
+lons = [-125.0, -115.0, -115.0, -125.0]
+
+# Get all morton cells covering the polygon at order 6
+cells = mortie.morton_coverage(lats, lons, order=6)
+```
+
+The function handles concave polygons, antimeridian-crossing polygons, and polar regions. Polygon holes are not yet supported.
+
 ## Dependencies
 
 **numpy**. All HEALPix operations use the Rust-native `healpix` crate bundled in the compiled extension — no external HEALPix library is needed.

benchmarks/test_bench_coverage.py

@@ -0,0 +1,92 @@
+"""
+Performance benchmarks for morton_coverage (CodSpeed-compatible).
+
+Run locally:
+    pytest benchmarks/test_bench_coverage.py -v
+
+Run with CodSpeed:
+    pytest benchmarks/test_bench_coverage.py --codspeed
+"""
+
+import numpy as np
+import pytest
+
+from mortie import morton_coverage
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+@pytest.fixture
+def triangle():
+    """Simple ~10° × 10° triangle."""
+    return np.array([40.0, 50.0, 45.0]), np.array([-120.0, -120.0, -110.0])
+
+
+@pytest.fixture
+def square():
+    """~10° × 10° square."""
+    return np.array([40.0, 40.0, 50.0, 50.0]), np.array([-125.0, -115.0, -115.0, -125.0])
+
+
+@pytest.fixture
+def circle_100():
+    """Circular polygon with 100 vertices (~5° radius, southern hemisphere)."""
+    n = 100
+    center_lat, center_lon, radius = -75.0, 0.0, 5.0
+    angles = np.linspace(0, 2 * np.pi, n, endpoint=False)
+    lats = center_lat + radius * np.cos(angles)
+    lons = center_lon + radius * np.sin(angles)
+    return lats, lons
+
+
+@pytest.fixture
+def circle_500():
+    """Circular polygon with 500 vertices (~5° radius, southern hemisphere)."""
+    n = 500
+    center_lat, center_lon, radius = -75.0, 0.0, 5.0
+    angles = np.linspace(0, 2 * np.pi, n, endpoint=False)
+    lats = center_lat + radius * np.cos(angles)
+    lons = center_lon + radius * np.sin(angles)
+    return lats, lons
+
+
+# ---------------------------------------------------------------------------
+# Benchmarks
+# ---------------------------------------------------------------------------
+
+def test_coverage_triangle_order4(benchmark, triangle):
+    """morton_coverage: triangle at order 4."""
+    lats, lons = triangle
+    benchmark(morton_coverage, lats, lons, order=4)
+
+
+def test_coverage_triangle_order6(benchmark, triangle):
+    """morton_coverage: triangle at order 6."""
+    lats, lons = triangle
+    benchmark(morton_coverage, lats, lons, order=6)
+
+
+def test_coverage_triangle_order8(benchmark, triangle):
+    """morton_coverage: triangle at order 8."""
+    lats, lons = triangle
+    benchmark(morton_coverage, lats, lons, order=8)
+
+
+def test_coverage_square_order6(benchmark, square):
+    """morton_coverage: square at order 6."""
+    lats, lons = square
+    benchmark(morton_coverage, lats, lons, order=6)
+
+
+def test_coverage_circle100_order6(benchmark, circle_100):
+    """morton_coverage: 100-vertex circle at order 6."""
+    lats, lons = circle_100
+    benchmark(morton_coverage, lats, lons, order=6)
+
+
+def test_coverage_circle500_order6(benchmark, circle_500):
+    """morton_coverage: 500-vertex circle at order 6."""
+    lats, lons = circle_500
+    benchmark(morton_coverage, lats, lons, order=6)

mortie/__init__.py

@@ -35,6 +35,9 @@
     morton_buffer,
 )
 
+# Import coverage functions
+from .coverage import morton_coverage
+
 # Import prefix trie functions
 from .prefix_trie import (
     MortonChild,
@@ -66,6 +69,7 @@
     'generate_morton_children',
     'mort2healpix',
     'morton_buffer',
+    'morton_coverage',
     'MortonChild',
     'split_children',
     'split_children_geo',

mortie/coverage.py

@@ -0,0 +1,107 @@
+"""Polygon-to-morton coverage.
+
+Compute the set of morton indices at a given order that completely cover
+a polygon defined by lat/lon vertices.  Supports single and multipart
+polygons.
+"""
+
+import numpy as np
+
+
+def _is_multipart(lats):
+    """Check if *lats* is a sequence of sequences (multipart polygon)."""
+    if isinstance(lats, np.ndarray):
+        return lats.ndim == 2
+    if isinstance(lats, (list, tuple)) and len(lats) > 0:
+        return isinstance(lats[0], (list, tuple, np.ndarray))
+    return False
+
+
+def _single_coverage(lats, lons, order):
+    """Coverage for one polygon ring."""
+    lats = np.asarray(lats, dtype=np.float64).ravel()
+    lons = np.asarray(lons, dtype=np.float64).ravel()
+
+    if lats.shape != lons.shape:
+        raise ValueError("lats and lons must have the same length")
+    if lats.size < 3:
+        raise ValueError("Need at least 3 vertices for a polygon")
+    if not np.all(np.isfinite(lats)) or not np.all(np.isfinite(lons)):
+        raise ValueError("lats and lons must not contain NaN or infinity")
+
+    # Strip duplicate closing vertex (first == last) if present
+    if lats[0] == lats[-1] and lons[0] == lons[-1] and lats.size > 3:
+        lats = lats[:-1].copy()
+        lons = lons[:-1].copy()
+
+    from . import _rustie
+
+    return np.asarray(_rustie.rust_polygon_coverage(lats, lons, order))
+
+
+def morton_coverage(lats, lons, order=18):
+    """Compute morton indices covering a polygon defined by lat/lon vertices.
+
+    Given a polygon (as arrays of vertex latitudes and longitudes), returns the
+    set of morton indices at the requested HEALPix order that completely cover
+    the polygon interior.  The coverage is optimally compact — it includes all
+    boundary cells plus every cell whose centre lies inside the polygon.
+
+    For **multipart polygons**, pass *lats* and *lons* as lists of arrays
+    (one per part).  The coverage of all parts is unioned.
+
+    Parameters
+    ----------
+    lats : array_like or list of array_like
+        Vertex latitudes in degrees.  For a single polygon, a 1-D array
+        with at least 3 vertices.  For multipart polygons, a list of
+        such arrays.
+    lons : array_like or list of array_like
+        Vertex longitudes in degrees.  Must match the structure of *lats*.
+    order : int, optional
+        HEALPix depth / tessellation order (1–18).  Default 18.
+
+    Returns
+    -------
+    numpy.ndarray
+        Sorted 1-D array of unique morton indices (dtype ``int64``).
+
+    Raises
+    ------
+    ValueError
+        If fewer than 3 vertices, mismatched lengths, invalid order,
+        or coordinates containing NaN/infinity.
+
+    Notes
+    -----
+    - Self-intersecting polygons produce undefined results.
+    - Polygons with holes are not supported; pass the outer ring only.
+    - The algorithm uses gnomonic projection centered on each test point
+      with a winding-number PIP test, which works correctly for polygons
+      in any hemisphere.
+
+    Examples
+    --------
+    Single polygon:
+
+    >>> import mortie
+    >>> lats = [40.0, 50.0, 45.0]
+    >>> lons = [-120.0, -120.0, -110.0]
+    >>> cells = mortie.morton_coverage(lats, lons, order=6)
+
+    Multipart polygon:
+
+    >>> lats_parts = [[40.0, 50.0, 45.0], [10.0, 20.0, 15.0]]
+    >>> lons_parts = [[-120.0, -120.0, -110.0], [-80.0, -80.0, -70.0]]
+    >>> cells = mortie.morton_coverage(lats_parts, lons_parts, order=6)
+    """
+    if not 1 <= order <= 18:
+        raise ValueError("Order must be between 1 and 18")
+
+    if _is_multipart(lats):
+        if len(lats) != len(lons):
+            raise ValueError("lats and lons must have the same number of parts")
+        parts = [_single_coverage(la, lo, order) for la, lo in zip(lats, lons)]
+        return np.unique(np.concatenate(parts))
+
+    return _single_coverage(lats, lons, order)