healpix-geo 0.0.7__cp312-cp312-musllinux_1_2_armv7l.whl

healpix_geo/tests/test_index.py

@@ -0,0 +1,290 @@
+import pickle
+
+import numpy as np
+import pytest
+import shapely
+
+import healpix_geo
+
+
+class TestRangeMOCIndex:
+    @pytest.mark.parametrize("level", [0, 3, 6])
+    def test_full_domain(self, level):
+        index = healpix_geo.nested.RangeMOCIndex.full_domain(level)
+
+        expected = np.arange(12 * 4**level, dtype="uint64")
+
+        assert index.nbytes == 16
+        assert index.size == expected.size
+        assert index.depth == level
+
+    @pytest.mark.parametrize(
+        ["level", "cell_ids"],
+        (
+            (0, np.array([1, 2, 5], dtype="uint64")),
+            (3, np.array([12, 16, 17, 19, 22, 23, 71, 72, 73, 79], dtype="uint64")),
+            (6, np.arange(3 * 4**6, 5 * 4**6, dtype="uint64")),
+        ),
+    )
+    def test_from_cell_ids(self, level, cell_ids):
+        index = healpix_geo.nested.RangeMOCIndex.from_cell_ids(level, cell_ids)
+
+        assert index.size == cell_ids.size
+        assert index.depth == level
+
+    @pytest.mark.parametrize(
+        ["level", "cell_ids1", "cell_ids2", "expected"],
+        (
+            (
+                4,
+                np.arange(0, 6 * 4**4, dtype="uint64"),
+                np.arange(6 * 4**4, 12 * 4**4, dtype="uint64"),
+                np.arange(12 * 4**4, dtype="uint64"),
+            ),
+            (
+                1,
+                np.array([1, 2, 3, 4, 21, 22], dtype="uint64"),
+                np.array([23, 25, 26, 32, 33, 34, 35], dtype="uint64"),
+                np.array(
+                    [1, 2, 3, 4, 21, 22, 23, 25, 26, 32, 33, 34, 35], dtype="uint64"
+                ),
+            ),
+        ),
+    )
+    def test_union(self, level, cell_ids1, cell_ids2, expected):
+        index1 = healpix_geo.nested.RangeMOCIndex.from_cell_ids(level, cell_ids1)
+        index2 = healpix_geo.nested.RangeMOCIndex.from_cell_ids(level, cell_ids2)
+
+        actual = index1.union(index2)
+
+        assert isinstance(actual, healpix_geo.nested.RangeMOCIndex)
+        np.testing.assert_equal(actual.cell_ids(), expected)
+
+    @pytest.mark.parametrize(
+        ["level", "cell_ids1", "cell_ids2", "expected"],
+        (
+            (
+                4,
+                np.arange(2 * 4**4, 4 * 4**4, dtype="uint64"),
+                np.arange(3 * 4**4, 5 * 4**4, dtype="uint64"),
+                np.arange(3 * 4**4, 4 * 4**4, dtype="uint64"),
+            ),
+            (
+                1,
+                np.array([1, 2, 3, 4, 21, 22, 23, 24, 25], dtype="uint64"),
+                np.array([21, 22, 23, 25, 26, 32, 33, 34, 35], dtype="uint64"),
+                np.array([21, 22, 23, 25], dtype="uint64"),
+            ),
+        ),
+    )
+    def test_intersection(self, level, cell_ids1, cell_ids2, expected):
+        index1 = healpix_geo.nested.RangeMOCIndex.from_cell_ids(level, cell_ids1)
+        index2 = healpix_geo.nested.RangeMOCIndex.from_cell_ids(level, cell_ids2)
+
+        actual = index1.intersection(index2)
+
+        assert isinstance(actual, healpix_geo.nested.RangeMOCIndex)
+        np.testing.assert_equal(actual.cell_ids(), expected)
+
+    @pytest.mark.parametrize(
+        ["level", "cell_ids"],
+        (
+            pytest.param(0, np.arange(12, dtype="uint64"), id="base cells"),
+            pytest.param(
+                1,
+                np.array([0, 1, 2, 4, 5, 11, 12, 13, 25, 26, 27], dtype="uint64"),
+                id="list of level 1 cells",
+            ),
+            pytest.param(
+                4,
+                np.arange(1 * 4**4, 2 * 4**4, dtype="uint64"),
+                id="single level 4 base cell",
+            ),
+        ),
+    )
+    @pytest.mark.parametrize(
+        "indexer",
+        [
+            slice(None),
+            slice(None, 4),
+            slice(2, None),
+            slice(3, 7),
+            np.arange(5, dtype="uint64"),
+            np.array([1, 2, 4, 6, 8], dtype="uint64"),
+        ],
+    )
+    def test_isel(self, level, cell_ids, indexer):
+        expected = cell_ids[indexer]
+
+        index = healpix_geo.nested.RangeMOCIndex.from_cell_ids(level, cell_ids)
+
+        actual = index.isel(indexer)
+
+        np.testing.assert_equal(actual.cell_ids(), expected)
+
+    @pytest.mark.parametrize(
+        ["level", "cell_ids", "indexer"],
+        (
+            pytest.param(
+                0,
+                np.arange(12, dtype="uint64"),
+                slice(None),
+                id="base cells-slice-full",
+            ),
+            pytest.param(
+                0,
+                np.arange(12, dtype="uint64"),
+                slice(None, 6),
+                id="base cells-slice-left_open",
+            ),
+            pytest.param(
+                0,
+                np.arange(12, dtype="uint64"),
+                slice(2, None),
+                id="base cells-slice-right_open",
+            ),
+            pytest.param(
+                0,
+                np.arange(12, dtype="uint64"),
+                slice(2, 7),
+                id="base cells-slice-domain",
+            ),
+            pytest.param(
+                0,
+                np.arange(12, dtype="uint64"),
+                np.arange(12, dtype="uint64"),
+                id="base cells-array-full",
+            ),
+            pytest.param(
+                0,
+                np.arange(12, dtype="uint64"),
+                np.arange(2, 7, dtype="uint64"),
+                id="base cells-array-domain",
+            ),
+            pytest.param(
+                0,
+                np.arange(12, dtype="uint64"),
+                np.array([1, 2, 3, 7, 8, 9, 10], dtype="uint64"),
+                id="base cells-array-disconnected",
+            ),
+            pytest.param(
+                3,
+                np.arange(12 * 4**3, dtype="uint64"),
+                slice(None, 15),
+                id="level 3 cells-slice-left_open",
+            ),
+            pytest.param(
+                1,
+                np.array([0, 1, 2, 4, 5, 11, 12, 13, 25, 26, 27], dtype="uint64"),
+                np.array([2, 5, 11, 12, 25, 27], dtype="uint64"),
+                id="list of level 1 cells-array-disconnected",
+            ),
+            pytest.param(
+                4,
+                np.arange(1 * 4**4, 2 * 4**4, dtype="uint64"),
+                slice(260, 280),
+                id="single level 4 base cell-slice-domain",
+            ),
+        ),
+    )
+    def test_sel(self, level, cell_ids, indexer):
+        if isinstance(indexer, slice):
+            n = slice(
+                indexer.start if indexer.start is not None else 0,
+                (
+                    indexer.stop + 1
+                    if indexer.stop is not None
+                    else int(np.max(cell_ids)) + 1
+                ),
+                indexer.step if indexer.step is not None else 1,
+            )
+            range_ = np.arange(n.start, n.stop, n.step, dtype="uint64")
+            condition = np.isin(cell_ids, range_)
+        else:
+            condition = np.isin(cell_ids, indexer)
+        expected_cell_ids = cell_ids[condition]
+
+        index = healpix_geo.nested.RangeMOCIndex.from_cell_ids(level, cell_ids)
+
+        actual_indexer, actual_moc = index.sel(indexer)
+
+        np.testing.assert_equal(cell_ids[actual_indexer], expected_cell_ids)
+        np.testing.assert_equal(actual_moc.cell_ids(), expected_cell_ids)
+
+    @pytest.mark.parametrize(
+        ["depth", "cell_ids"],
+        (
+            (2, np.arange(1 * 4**2, 3 * 4**2, dtype="uint64")),
+            (5, np.arange(12 * 4**5, dtype="uint64")),
+        ),
+    )
+    def test_pickle_roundtrip(self, depth, cell_ids):
+        index = healpix_geo.nested.RangeMOCIndex.from_cell_ids(depth, cell_ids)
+
+        pickled = pickle.dumps(index)
+        assert isinstance(pickled, bytes)
+        unpickled = pickle.loads(pickled)
+
+        assert isinstance(unpickled, healpix_geo.nested.RangeMOCIndex)
+        assert index.depth == unpickled.depth
+        np.testing.assert_equal(unpickled.cell_ids(), index.cell_ids())
+
+    @pytest.mark.parametrize("depth", (0, 2, 10))
+    @pytest.mark.parametrize(
+        "geom",
+        (
+            pytest.param(shapely.Point(30, 30), id="point"),
+            pytest.param(shapely.box(-25, 15, 25, 35), id="polygon"),
+            pytest.param(
+                shapely.LineString([(30, 30), (31, 31), (32, 33)]), id="linestring"
+            ),
+            pytest.param(healpix_geo.geometry.Bbox(-25, 15, 25, 35), id="bbox"),
+        ),
+    )
+    @pytest.mark.parametrize("domain", ["full", "partial"])
+    def test_query(self, depth, domain, geom):
+        import cdshealpix.nested
+        from astropy.coordinates import Latitude, Longitude
+
+        if domain == "full":
+            index = healpix_geo.nested.RangeMOCIndex.full_domain(depth)
+            cell_ids = index.cell_ids()
+        else:
+            cell_ids = np.arange(4**depth, dtype="uint64")
+            index = healpix_geo.nested.RangeMOCIndex.from_cell_ids(depth, cell_ids)
+
+        if isinstance(geom, shapely.Point):
+            coords = geom.coords[0]
+            lon = Longitude([coords[0]], unit="deg")
+            lat = Latitude([coords[0]], unit="deg")
+            expected = cdshealpix.nested.lonlat_to_healpix(lon, lat, depth=depth)
+        elif isinstance(geom, shapely.LineString):
+            coords = np.asarray(geom.coords[:])
+            lon = Longitude(coords[:, 0], unit="deg")
+            lat = Latitude(coords[:, 1], unit="deg")
+
+            expected_ = np.unique(
+                cdshealpix.nested.lonlat_to_healpix(lon, lat, depth=depth)
+            )
+            expected = expected_[np.isin(expected_, cell_ids)]
+        elif isinstance(geom, shapely.Polygon):
+            coords = np.asarray(geom.exterior.coords[:])
+            lon = Longitude(coords[:, 0], unit="deg")
+            lat = Latitude(coords[:, 1], unit="deg")
+            expected_, _, _ = cdshealpix.nested.polygon_search(
+                lon, lat, depth=depth, flat=True
+            )
+            expected = expected_[np.isin(expected_, cell_ids)]
+        else:
+            expected = None
+
+        multi_slice, moc = index.query(geom)
+
+        reconstructed = np.concatenate(
+            [cell_ids[s.as_pyslice()] for s in multi_slice], axis=0
+        )
+        actual = moc.cell_ids()
+
+        if expected is not None:
+            np.testing.assert_equal(reconstructed, expected)
+        np.testing.assert_equal(actual, reconstructed)

healpix_geo/tests/test_inference.py

@@ -0,0 +1,185 @@
+import cdshealpix
+import numpy as np
+import pytest
+from astropy.coordinates import Latitude, Longitude
+
+import healpix_geo
+
+
+class TestHealpixToGeographic:
+    @pytest.mark.parametrize(
+        ["cell_ids", "depth", "indexing_scheme"],
+        (
+            pytest.param(np.array([0, 4, 5, 7, 9]), 0, "ring", id="level0-ring"),
+            pytest.param(np.array([1, 2, 3, 8]), 0, "nested", id="level0-nested"),
+            pytest.param(
+                np.array([3, 19, 54, 63, 104, 127]), 4, "ring", id="level4-ring"
+            ),
+            pytest.param(
+                np.array([22, 89, 134, 154, 190]), 4, "nested", id="level4-nested"
+            ),
+        ),
+    )
+    def test_spherical(self, cell_ids, depth, indexing_scheme):
+        if indexing_scheme == "ring":
+            param_cds = 2**depth
+            hg_healpix_to_lonlat = healpix_geo.ring.healpix_to_lonlat
+            cds_healpix_to_lonlat = cdshealpix.ring.healpix_to_lonlat
+        else:
+            param_cds = depth
+            hg_healpix_to_lonlat = healpix_geo.nested.healpix_to_lonlat
+            cds_healpix_to_lonlat = cdshealpix.nested.healpix_to_lonlat
+
+        actual_lon, actual_lat = hg_healpix_to_lonlat(
+            cell_ids, depth, ellipsoid="sphere"
+        )
+        expected_lon_, expected_lat_ = cds_healpix_to_lonlat(cell_ids, param_cds)
+        expected_lon = np.asarray(expected_lon_.to("degree"))
+        expected_lat = np.asarray(expected_lat_.to("degree"))
+
+        np.testing.assert_allclose(actual_lon, expected_lon)
+        np.testing.assert_allclose(actual_lat, expected_lat)
+
+    @pytest.mark.parametrize("ellipsoid", ["unitsphere", "sphere", "WGS84", "bessel"])
+    @pytest.mark.parametrize("depth", [0, 1, 9])
+    @pytest.mark.parametrize("indexing_scheme", ["ring", "nested"])
+    def test_ellipsoidal(self, depth, indexing_scheme, ellipsoid):
+        cell_ids = np.arange(12)
+        if indexing_scheme == "ring":
+            param_cds = 2**depth
+            hg_healpix_to_lonlat = healpix_geo.ring.healpix_to_lonlat
+            cds_healpix_to_lonlat = cdshealpix.ring.healpix_to_lonlat
+        else:
+            param_cds = depth
+            hg_healpix_to_lonlat = healpix_geo.nested.healpix_to_lonlat
+            cds_healpix_to_lonlat = cdshealpix.nested.healpix_to_lonlat
+
+        actual_lon, actual_lat = hg_healpix_to_lonlat(
+            cell_ids, depth, ellipsoid=ellipsoid
+        )
+        expected_lon_, expected_lat_ = cds_healpix_to_lonlat(cell_ids, param_cds)
+        expected_lon = np.asarray(expected_lon_.to("degree"))
+        expected_lat = np.asarray(expected_lat_.to("degree"))
+
+        np.testing.assert_allclose(actual_lon, expected_lon)
+
+        diff_lat = actual_lat - expected_lat
+        assert np.all(abs(diff_lat) < 0.3)
+
+        signs = np.array([-1, 1])
+        actual = signs[(actual_lat >= 0).astype(int)]
+        expected_ = np.sign(diff_lat)
+        expected = np.where(expected_ == 0, 1, expected_)
+        assert np.all(diff_lat == 0) or np.all(actual == expected)
+
+
+class TestGeographicToHealpix:
+    @pytest.mark.parametrize(
+        ["lon", "lat", "depth", "indexing_scheme"],
+        (
+            pytest.param(
+                np.array([-170.0, 10.0, 30.0, 124.0, 174.0]),
+                np.array([-48.0, -30.0, -5.0, 15.0, 30.0]),
+                0,
+                "ring",
+                id="level0-ring",
+            ),
+            pytest.param(
+                np.array([-170.0, 10.0, 30.0, 124.0, 174.0]),
+                np.array([-48.0, -30.0, -5.0, 15.0, 30.0]),
+                0,
+                "nested",
+                id="level0-nested",
+            ),
+            pytest.param(
+                np.array([-70.0, 135.0, 150.0]),
+                np.array([-65.0, 0.0, 65.0]),
+                4,
+                "ring",
+                id="level4-ring",
+            ),
+            pytest.param(
+                np.array([-70.0, 135.0, 150.0]),
+                np.array([-65.0, 0.0, 65.0]),
+                4,
+                "nested",
+                id="level4-nested",
+            ),
+        ),
+    )
+    def test_spherical(self, lon, lat, depth, indexing_scheme):
+        if indexing_scheme == "ring":
+            param_cds = 2**depth
+            hg_lonlat_to_healpix = healpix_geo.ring.lonlat_to_healpix
+            cds_lonlat_to_healpix = cdshealpix.ring.lonlat_to_healpix
+        else:
+            param_cds = depth
+            hg_lonlat_to_healpix = healpix_geo.nested.lonlat_to_healpix
+            cds_lonlat_to_healpix = cdshealpix.nested.lonlat_to_healpix
+
+        actual = hg_lonlat_to_healpix(lon, lat, depth, ellipsoid="sphere")
+        lon_ = Longitude(lon, unit="degree")
+        lat_ = Latitude(lat, unit="degree")
+        expected = cds_lonlat_to_healpix(lon_, lat_, param_cds)
+
+        np.testing.assert_equal(actual, expected)
+
+    @pytest.mark.parametrize("ellipsoid", ["unitsphere", "sphere", "WGS84", "bessel"])
+    @pytest.mark.parametrize("depth", [0, 1, 9])
+    @pytest.mark.parametrize("indexing_scheme", ["ring", "nested"])
+    def test_ellipsoidal(self, ellipsoid, depth, indexing_scheme):
+        lat = np.linspace(-90, 90, 50)
+        lon = np.full_like(lat, fill_value=45.0)
+
+        if indexing_scheme == "ring":
+            param_cds = 2**depth
+            hg = healpix_geo.ring.lonlat_to_healpix
+            cds = cdshealpix.ring.lonlat_to_healpix
+        else:
+            param_cds = depth
+            hg = healpix_geo.nested.lonlat_to_healpix
+            cds = cdshealpix.nested.lonlat_to_healpix
+
+        actual = hg(lon, lat, depth, ellipsoid=ellipsoid)
+
+        lon_ = Longitude(lon, unit="degree")
+        lat_ = Latitude(lat, unit="degree")
+        expected = cds(lon_, lat_, param_cds)
+
+        assert actual.dtype == "uint64"
+        assert expected.dtype == "uint64"
+
+        # TODO: this is currently a smoke check, try more thorough checks
+
+
+class TestVertices:
+    @pytest.mark.parametrize(
+        ["cell_ids", "depth", "indexing_scheme"],
+        (
+            pytest.param(np.array([0, 4, 5, 7, 9]), 0, "ring", id="level0-ring"),
+            pytest.param(np.array([1, 2, 3, 8]), 0, "nested", id="level0-nested"),
+            pytest.param(
+                np.array([3, 19, 54, 63, 104, 127]), 4, "ring", id="level4-ring"
+            ),
+            pytest.param(
+                np.array([22, 89, 134, 154, 190]), 4, "nested", id="level4-nested"
+            ),
+        ),
+    )
+    def test_spherical(self, cell_ids, depth, indexing_scheme):
+        if indexing_scheme == "ring":
+            param_cds = 2**depth
+            hg_vertices = healpix_geo.ring.vertices
+            cds_vertices = cdshealpix.ring.vertices
+        else:
+            param_cds = depth
+            hg_vertices = healpix_geo.nested.vertices
+            cds_vertices = cdshealpix.nested.vertices
+
+        actual_lon, actual_lat = hg_vertices(cell_ids, depth, ellipsoid="sphere")
+        expected_lon_, expected_lat_ = cds_vertices(cell_ids, param_cds)
+        expected_lon = np.asarray(expected_lon_.to("degree"))
+        expected_lat = np.asarray(expected_lat_.to("degree"))
+
+        np.testing.assert_allclose(actual_lon, expected_lon)
+        np.testing.assert_allclose(actual_lat, expected_lat)

healpix_geo/tests/test_neighbours.py

@@ -0,0 +1,40 @@
+import cdshealpix.nested
+import cdshealpix.ring
+import numpy as np
+import pytest
+
+import healpix_geo
+
+
+@pytest.mark.parametrize("depth", [2, 8])
+@pytest.mark.parametrize("ring", [0, 1])
+@pytest.mark.parametrize(
+    "indexing_scheme",
+    [
+        "nested",
+        "ring",
+    ],
+)
+def test_kth_neighbourhood(depth, ring, indexing_scheme):
+    if indexing_scheme == "nested":
+        kth_neighbourhood = healpix_geo.nested.kth_neighbourhood
+        neighbours = cdshealpix.nested.neighbours
+    elif indexing_scheme == "ring":
+        kth_neighbourhood = healpix_geo.ring.kth_neighbourhood
+
+        def neighbours(ipix, depth):
+            return cdshealpix.to_ring(
+                cdshealpix.nested.neighbours(cdshealpix.from_ring(ipix, depth), depth),
+                depth=depth,
+            )
+
+    ipixels = np.array([50, 100], dtype="int64")
+
+    actual = kth_neighbourhood(depth=depth, ipix=ipixels, ring=ring)
+    if ring == 0:
+        expected = np.reshape(ipixels, (-1, 1))
+    else:
+        expected = neighbours(ipix=ipixels, depth=depth)
+
+    np.testing.assert_equal(np.sort(actual, axis=-1), np.sort(expected, axis=-1))
+    np.testing.assert_equal(actual[:, 0], ipixels)

healpix_geo/tests/test_slices.py

@@ -0,0 +1,116 @@
+import pytest
+
+from healpix_geo import slices
+
+missing = object()
+
+
+class TestSlice:
+    @pytest.mark.parametrize("start", [0, None, 10])
+    @pytest.mark.parametrize("stop", [3, None, 10])
+    @pytest.mark.parametrize("step", [1, None, 2, missing])
+    def test_init(self, start, stop, step):
+        if step is missing:
+            vals = (start, stop)
+            step = None
+        else:
+            vals = (start, stop, step)
+
+        actual = slices.Slice(*vals)
+
+        assert actual.start == start
+        assert actual.stop == stop
+        assert actual.step == step
+
+    @pytest.mark.parametrize("start", [0, None, 10])
+    @pytest.mark.parametrize("stop", [3, None, 10])
+    @pytest.mark.parametrize("step", [1, None, 2, -1])
+    def test_repr(self, start, stop, step):
+        slice_ = slices.Slice(start, stop, step)
+        actual = repr(slice_)
+
+        expected = repr(slice(start, stop, step)).title()
+
+        assert actual == expected
+
+    @pytest.mark.parametrize(
+        "pyslice",
+        (
+            slice(None),
+            slice(None, 2),
+            slice(3, None),
+            slice(None, None, -1),
+            slice(10, None, -1),
+            slice(3, -2),
+        ),
+    )
+    def test_from_pyslice(self, pyslice):
+        actual = slices.Slice.from_pyslice(pyslice)
+
+        assert actual.start == pyslice.start
+        assert actual.start == pyslice.start
+        assert actual.start == pyslice.start
+
+    @pytest.mark.parametrize(
+        "pyslice",
+        (
+            slice(None),
+            slice(None, 2),
+            slice(3, None),
+            slice(None, None, -1),
+            slice(10, None, -1),
+            slice(3, -2),
+        ),
+    )
+    @pytest.mark.parametrize("size", (3, 8, 13))
+    def test_roundtrip_pyslice(self, pyslice, size):
+        slice_ = slices.Slice.from_pyslice(pyslice)
+
+        actual = slice_.as_pyslice()
+        assert actual == pyslice
+
+    @pytest.mark.parametrize(
+        "pyslice",
+        (
+            slice(None),
+            slice(None, 2),
+            slice(3, None),
+            slice(None, None, -1),
+            slice(10, None, -1),
+            slice(3, -2),
+        ),
+    )
+    @pytest.mark.parametrize("size", (3, 8, 13))
+    def test_as_concrete(self, pyslice, size):
+        slice_ = slices.Slice(pyslice.start, pyslice.stop, pyslice.step)
+
+        actual = slice_.as_concrete(size)
+        expected = slice(*pyslice.indices(size))
+
+        assert actual.start == expected.start
+        assert actual.stop == expected.stop
+        assert actual.step == expected.step
+
+    def test_in_dict(self):
+        slice1 = slices.Slice(0, 4)
+        slice2 = slices.Slice(4, 6)
+
+        d = {slice1: 1, slice2: 2}
+
+        assert d[slice1] == 1
+        assert d[slice2] == 2
+
+    @pytest.mark.parametrize(
+        ["vals", "expected"],
+        (
+            ((0, 4), True),
+            ((1, 4), False),
+            ((0, 4, 1), False),
+        ),
+    )
+    def test_compare(self, vals, expected):
+        slice_ = slices.Slice(0, 4)
+        other = slices.Slice(*vals)
+
+        actual = slice_ == other
+        assert actual == expected