Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use Alignment moc in crossmatch/join operations #438

Merged
merged 3 commits into from
Oct 15, 2024
Merged

Use Alignment moc in crossmatch/join operations #438

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
34974b4
use alignment moc in crossmatched/joined catalogs
smcguire-cmu Oct 14, 2024
fef8f8b
remove unused import
smcguire-cmu Oct 14, 2024
c7dd949
update docs hipscat branch
smcguire-cmu Oct 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,4 +10,4 @@ sphinx-autoapi
sphinx-copybutton
sphinx-book-theme
sphinx-design
git+https://github.com/astronomy-commons/hipscat.git@main
git+https://github.com/astronomy-commons/hipscat.git@development
18 changes: 13 additions & 5 deletions src/lsdb/catalog/catalog.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,9 @@ def crossmatch(
dec_column=self.hc_structure.catalog_info.dec_column + suffixes[0],
total_rows=None,
)
hc_catalog = hc.catalog.Catalog(new_catalog_info, alignment.pixel_tree, schema=get_arrow_schema(ddf))
hc_catalog = hc.catalog.Catalog(
new_catalog_info, alignment.pixel_tree, schema=get_arrow_schema(ddf), moc=alignment.moc
)
return Catalog(ddf, ddf_map, hc_catalog)

def cone_search(self, ra: float, dec: float, radius_arcsec: float, fine: bool = True) -> Catalog:
Expand Down Expand Up @@ -441,7 +443,9 @@ def merge_asof(
dec_column=self.hc_structure.catalog_info.dec_column + suffixes[0],
total_rows=None,
)
hc_catalog = hc.catalog.Catalog(new_catalog_info, alignment.pixel_tree, schema=get_arrow_schema(ddf))
hc_catalog = hc.catalog.Catalog(
new_catalog_info, alignment.pixel_tree, schema=get_arrow_schema(ddf), moc=alignment.moc
)
return Catalog(ddf, ddf_map, hc_catalog)

def join(
Expand Down Expand Up @@ -492,7 +496,7 @@ def join(
total_rows=None,
)
hc_catalog = hc.catalog.Catalog(
new_catalog_info, alignment.pixel_tree, schema=get_arrow_schema(ddf)
new_catalog_info, alignment.pixel_tree, schema=get_arrow_schema(ddf), moc=alignment.moc
)
return Catalog(ddf, ddf_map, hc_catalog)
if left_on is None or right_on is None:
Expand All @@ -515,7 +519,9 @@ def join(
dec_column=self.hc_structure.catalog_info.dec_column + suffixes[0],
total_rows=None,
)
hc_catalog = hc.catalog.Catalog(new_catalog_info, alignment.pixel_tree, schema=get_arrow_schema(ddf))
hc_catalog = hc.catalog.Catalog(
new_catalog_info, alignment.pixel_tree, schema=get_arrow_schema(ddf), moc=alignment.moc
)
return Catalog(ddf, ddf_map, hc_catalog)

def join_nested(
Expand Down Expand Up @@ -573,7 +579,9 @@ def join_nested(
catalog_name=output_catalog_name,
total_rows=None,
)
hc_catalog = hc.catalog.Catalog(new_catalog_info, alignment.pixel_tree)
hc_catalog = hc.catalog.Catalog(
new_catalog_info, alignment.pixel_tree, schema=get_arrow_schema(ddf), moc=alignment.moc
)
return Catalog(ddf, ddf_map, hc_catalog)

def nest_lists(
Expand Down
5 changes: 5 additions & 0 deletions tests/lsdb/catalog/test_crossmatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from lsdb.core.crossmatch.abstract_crossmatch_algorithm import AbstractCrossmatchAlgorithm
from lsdb.core.crossmatch.bounded_kdtree_match import BoundedKdTreeCrossmatch
from lsdb.core.crossmatch.kdtree_match import KdTreeCrossmatch
from lsdb.dask.merge_catalog_functions import align_catalogs


@pytest.mark.parametrize("algo", [KdTreeCrossmatch])
Expand All @@ -24,6 +25,10 @@ def test_kdtree_crossmatch(algo, small_sky_catalog, small_sky_xmatch_catalog, xm
)
assert isinstance(xmatched_cat._ddf, nd.NestedFrame)
xmatched = xmatched_cat.compute()
alignment = align_catalogs(small_sky_catalog, small_sky_xmatch_catalog)
assert xmatched_cat.hc_structure.moc == alignment.moc
assert xmatched_cat.get_healpix_pixels() == alignment.pixel_tree.get_healpix_pixels()

assert isinstance(xmatched, npd.NestedFrame)
assert len(xmatched) == len(xmatch_correct)
for _, correct_row in xmatch_correct.iterrows():
Expand Down
32 changes: 32 additions & 0 deletions tests/lsdb/catalog/test_join.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
import pytest
from hipscat.pixel_math.hipscat_id import HIPSCAT_ID_COLUMN, hipscat_id_to_healpix

from lsdb.dask.merge_catalog_functions import align_catalogs


def test_small_sky_join_small_sky_order1(
small_sky_catalog, small_sky_order1_catalog, assert_divisions_are_correct
Expand All @@ -21,6 +23,9 @@ def test_small_sky_join_small_sky_order1(
assert (col_name + suffixes[1], dtype) in joined.dtypes.items()
assert joined._ddf.index.name == HIPSCAT_ID_COLUMN
assert joined._ddf.index.dtype == np.uint64
alignment = align_catalogs(small_sky_catalog, small_sky_order1_catalog)
assert joined.hc_structure.moc == alignment.moc
assert joined.get_healpix_pixels() == alignment.pixel_tree.get_healpix_pixels()

joined_compute = joined.compute()
assert isinstance(joined_compute, npd.NestedFrame)
Expand All @@ -46,6 +51,11 @@ def test_small_sky_join_small_sky_order1_source(
assert (col_name + suffixes[0], dtype) in joined.dtypes.items()
for col_name, dtype in small_sky_order1_source_with_margin.dtypes.items():
assert (col_name + suffixes[1], dtype) in joined.dtypes.items()

alignment = align_catalogs(small_sky_catalog, small_sky_order1_source_with_margin)
assert joined.hc_structure.moc == alignment.moc
assert joined.get_healpix_pixels() == alignment.pixel_tree.get_healpix_pixels()

joined_compute = joined.compute()
small_sky_order1_compute = small_sky_order1_source_with_margin.compute()
assert len(joined_compute) == len(small_sky_order1_compute)
Expand Down Expand Up @@ -74,6 +84,10 @@ def test_join_association(small_sky_catalog, small_sky_xmatch_catalog, small_sky
)
assert isinstance(joined._ddf, nd.NestedFrame)
assert joined._ddf.npartitions == len(small_sky_to_xmatch_catalog.hc_structure.join_info.data_frame)
alignment = align_catalogs(small_sky_catalog, small_sky_xmatch_catalog)
assert joined.hc_structure.moc == alignment.moc
assert joined.get_healpix_pixels() == alignment.pixel_tree.get_healpix_pixels()

joined_data = joined.compute()
assert isinstance(joined_data, npd.NestedFrame)
association_data = small_sky_to_xmatch_catalog.compute()
Expand Down Expand Up @@ -120,6 +134,10 @@ def test_join_association_source_margin(
small_sky_order1_source_with_margin, through=small_sky_to_o1source_catalog, suffixes=suffixes
)
assert joined._ddf.npartitions == len(small_sky_to_o1source_catalog.hc_structure.join_info.data_frame)
alignment = align_catalogs(small_sky_catalog, small_sky_order1_source_with_margin)
assert joined.hc_structure.moc == alignment.moc
assert joined.get_healpix_pixels() == alignment.pixel_tree.get_healpix_pixels()

joined_data = joined.compute()
association_data = small_sky_to_o1source_catalog.compute()
assert len(joined_data) == 17161
Expand Down Expand Up @@ -156,6 +174,9 @@ def test_join_association_soft(small_sky_catalog, small_sky_xmatch_catalog, smal
small_sky_xmatch_catalog, through=small_sky_to_xmatch_soft_catalog, suffixes=suffixes
)
assert joined._ddf.npartitions == len(small_sky_to_xmatch_soft_catalog.hc_structure.join_info.data_frame)
alignment = align_catalogs(small_sky_catalog, small_sky_xmatch_catalog)
assert joined.hc_structure.moc == alignment.moc
assert joined.get_healpix_pixels() == alignment.pixel_tree.get_healpix_pixels()

with pytest.warns(match="margin"):
joined_on = small_sky_catalog.join(
Expand All @@ -177,6 +198,9 @@ def test_join_source_margin_soft(
assert joined._ddf.npartitions == len(
small_sky_to_o1source_soft_catalog.hc_structure.join_info.data_frame
)
alignment = align_catalogs(small_sky_catalog, small_sky_order1_source_with_margin)
assert joined.hc_structure.moc == alignment.moc
assert joined.get_healpix_pixels() == alignment.pixel_tree.get_healpix_pixels()

joined_on = small_sky_catalog.join(
small_sky_order1_source_with_margin,
Expand All @@ -200,6 +224,10 @@ def test_join_nested(small_sky_catalog, small_sky_order1_source_with_margin, ass
if col_name != "object_id":
assert (col_name, dtype.pyarrow_dtype) in joined["sources"].dtypes.fields.items()
assert_divisions_are_correct(joined)
alignment = align_catalogs(small_sky_catalog, small_sky_order1_source_with_margin)
assert joined.hc_structure.moc == alignment.moc
assert joined.get_healpix_pixels() == alignment.pixel_tree.get_healpix_pixels()

joined_compute = joined.compute()
source_compute = small_sky_order1_source_with_margin.compute()
assert isinstance(joined_compute, npd.NestedFrame)
Expand All @@ -224,6 +252,10 @@ def test_merge_asof(small_sky_catalog, small_sky_xmatch_catalog, assert_division
)
assert isinstance(joined._ddf, nd.NestedFrame)
assert_divisions_are_correct(joined)
alignment = align_catalogs(small_sky_catalog, small_sky_xmatch_catalog)
assert joined.hc_structure.moc == alignment.moc
assert joined.get_healpix_pixels() == alignment.pixel_tree.get_healpix_pixels()

joined_compute = joined.compute()
assert isinstance(joined_compute, npd.NestedFrame)
small_sky_compute = small_sky_catalog.compute().rename(
Expand Down