Updated healpix containment algorithm to improve spatial query speed

changes/+78480439.improvement.rst

@@ -0,0 +1 @@
+Improved spatial querying on lightcones, which should result in significant speedup for larger regions.

python/opencosmo/dataset/dataset.py

@@ -406,8 +406,10 @@ def bound(self, region: Region, select_by: Optional[str] = None):
         else:
             new_intersects_index = np.array([], dtype=np.int64)
 
-        new_index = np.concatenate(
-            [into_array(contained_index), into_array(new_intersects_index)]
+        new_index = np.sort(
+            np.concatenate(
+                [into_array(contained_index), into_array(new_intersects_index)]
+            )
         )
 
         new_state = st.with_region(st.take_rows(self.__state, new_index), check_region)

python/opencosmo/spatial/healpix.py

@@ -2,7 +2,9 @@
 
 from typing import TYPE_CHECKING
 
+import healpy as hp
 import numpy as np
+from astropy.coordinates import SkyCoord
 
 from opencosmo.index import into_array
 from opencosmo.spatial.region import HealpixRegion
@@ -47,4 +49,18 @@ def query(
             raise ValueError("Didn't recieve a 2D region!")
         nside = 2**level
         intersects = region.get_healpix_intersections(nside)
-        return {level: (np.array([]), intersects)}
+        boundaries = (
+            hp.boundaries(nside, intersects, nest=True)
+            .transpose(
+                0,
+                2,
+                1,
+            )
+            .reshape(-1, 3)
+        )
+        coords = SkyCoord(*hp.vec2ang(boundaries, lonlat=True), unit="deg")
+        coord_is_contained = region.contains(coords)
+        pixel_is_contained = np.all(coord_is_contained.reshape(-1, 4), axis=1)
+        return {
+            level: (intersects[pixel_is_contained], intersects[~pixel_is_contained])
+        }

python/opencosmo/spatial/relations.py

@@ -160,7 +160,7 @@ def __healpix_contains_other(region: HealpixRegion, other) -> bool:
         intersections = other.get_healpix_intersections(region.nside)
     except AttributeError:
         raise ValueError(f"Expected a 2D Sky Region but received {type(other)}")
-    return len(np.intersect1d(intersections, region.pixels)) == len(intersections)
+    return bool(np.all(np.isin(intersections, region.pixels)))
 
 
 # ---------------------------------------------------------------------------

test/spatial/test_2d.py

@@ -1,9 +1,9 @@
 import astropy.units as u
 import numpy as np
 from astropy.coordinates import SkyCoord
+from opencosmo.spatial.relations import contains_2d, intersects_2d
 
 import opencosmo as oc
-from opencosmo.spatial.relations import contains_2d, intersects_2d
 
 # ---------------------------------------------------------------------------
 # Cone search (bound with ConeRegion)