diff --git a/changes/+78480439.improvement.rst b/changes/+78480439.improvement.rst new file mode 100644 index 00000000..15811ff0 --- /dev/null +++ b/changes/+78480439.improvement.rst @@ -0,0 +1 @@ +Improved spatial querying on lightcones, which should result in significant speedup for larger regions. diff --git a/python/opencosmo/dataset/dataset.py b/python/opencosmo/dataset/dataset.py index 52b0a6f9..ac3b62c6 100644 --- a/python/opencosmo/dataset/dataset.py +++ b/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) diff --git a/python/opencosmo/spatial/healpix.py b/python/opencosmo/spatial/healpix.py index 19a44adb..442b6db4 100644 --- a/python/opencosmo/spatial/healpix.py +++ b/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]) + } diff --git a/python/opencosmo/spatial/relations.py b/python/opencosmo/spatial/relations.py index 0aae3ff7..fd19481f 100644 --- a/python/opencosmo/spatial/relations.py +++ b/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))) # --------------------------------------------------------------------------- diff --git a/test/spatial/test_2d.py b/test/spatial/test_2d.py index 363017e7..5613e4b6 100644 --- a/test/spatial/test_2d.py +++ b/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)