From 6efc1b5cc3963f38f2478e64afe3460d7e78e95b Mon Sep 17 00:00:00 2001 From: Melissa DeLucchi <113376043+delucchi-cmu@users.noreply.github.com> Date: Thu, 12 Sep 2024 11:16:05 -0400 Subject: [PATCH] Use healpix shim for healpy operations. (#412) --- src/lsdb/catalog/dataset/healpix_dataset.py | 4 ++-- src/lsdb/core/plotting/skymap.py | 2 +- src/lsdb/core/search/polygon_search.py | 2 +- src/lsdb/dask/merge_catalog_functions.py | 2 +- .../dataframe/margin_catalog_generator.py | 2 +- tests/lsdb/catalog/test_catalog.py | 19 ++++++++++--------- .../loaders/dataframe/test_from_dataframe.py | 2 +- 7 files changed, 17 insertions(+), 16 deletions(-) diff --git a/src/lsdb/catalog/dataset/healpix_dataset.py b/src/lsdb/catalog/dataset/healpix_dataset.py index 56f7b603..bcf39bcf 100644 --- a/src/lsdb/catalog/dataset/healpix_dataset.py +++ b/src/lsdb/catalog/dataset/healpix_dataset.py @@ -5,8 +5,8 @@ import dask import dask.dataframe as dd -import healpy as hp import hipscat as hc +import hipscat.pixel_math.healpix_shim as hp import nested_dask as nd import nested_pandas as npd import numpy as np @@ -363,7 +363,7 @@ def skymap( self, func: Callable[[npd.NestedFrame, HealpixPixel], Any], order: int | None = None, - default_value: Any = hp.pixelfunc.UNSEEN, + default_value: Any = hp.unseen_pixel(), projection="moll", plotting_args: Dict | None = None, **kwargs, diff --git a/src/lsdb/core/plotting/skymap.py b/src/lsdb/core/plotting/skymap.py index 8bd4035c..4a464967 100644 --- a/src/lsdb/core/plotting/skymap.py +++ b/src/lsdb/core/plotting/skymap.py @@ -2,7 +2,7 @@ from typing import Any, Callable, Dict -import healpy as hp +import hipscat.pixel_math.healpix_shim as hp import nested_pandas as npd import numpy as np from dask import delayed diff --git a/src/lsdb/core/search/polygon_search.py b/src/lsdb/core/search/polygon_search.py index 52905723..ab5befcb 100644 --- a/src/lsdb/core/search/polygon_search.py +++ b/src/lsdb/core/search/polygon_search.py @@ -1,6 +1,6 @@ from typing import List, Tuple -import healpy as hp +import hipscat.pixel_math.healpix_shim as hp import nested_pandas as npd import numpy as np from hipscat.catalog.catalog_info import CatalogInfo diff --git a/src/lsdb/dask/merge_catalog_functions.py b/src/lsdb/dask/merge_catalog_functions.py index 632fa1ea..2e79eb79 100644 --- a/src/lsdb/dask/merge_catalog_functions.py +++ b/src/lsdb/dask/merge_catalog_functions.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Callable, List, Sequence, Tuple -import healpy as hp +import hipscat.pixel_math.healpix_shim as hp import nested_dask as nd import nested_pandas as npd import numpy as np diff --git a/src/lsdb/loaders/dataframe/margin_catalog_generator.py b/src/lsdb/loaders/dataframe/margin_catalog_generator.py index e22c4904..30a62097 100644 --- a/src/lsdb/loaders/dataframe/margin_catalog_generator.py +++ b/src/lsdb/loaders/dataframe/margin_catalog_generator.py @@ -2,8 +2,8 @@ from typing import Dict, List, Tuple -import healpy as hp import hipscat as hc +import hipscat.pixel_math.healpix_shim as hp import nested_dask as nd import nested_pandas as npd import numpy as np diff --git a/tests/lsdb/catalog/test_catalog.py b/tests/lsdb/catalog/test_catalog.py index 183011fd..fcd400a8 100644 --- a/tests/lsdb/catalog/test_catalog.py +++ b/tests/lsdb/catalog/test_catalog.py @@ -2,7 +2,8 @@ import dask.array as da import dask.dataframe as dd -import healpy as hp +import healpy +import hipscat.pixel_math.healpix_shim as hp import nested_dask as nd import nested_pandas as npd import numpy as np @@ -349,7 +350,7 @@ def func(df, healpix): pixel_map = small_sky_order1_catalog.skymap_data(func) pixel_map = {pixel: value.compute() for pixel, value in pixel_map.items()} max_order = max(pixel_map.keys(), key=lambda x: x.order).order - img = np.zeros(hp.nside2npix(hp.order2nside(max_order))) + img = np.zeros(hp.order2npix(max_order)) for pixel, value in pixel_map.items(): dorder = max_order - pixel.order start = pixel.pixel * (4**dorder) @@ -379,7 +380,7 @@ def func(df, _): order_3_pixels = hipscat_id_to_healpix(computed_catalog.index.to_numpy(), order) pixel_map = computed_catalog.groupby(order_3_pixels).apply(lambda x: func(x, None)) - img = np.full(hp.nside2npix(hp.order2nside(order)), default) + img = np.full(hp.order2npix(order), default) for pixel_num, row in pixel_map.items(): img[pixel_num] = row assert (small_sky_order1_catalog.skymap_histogram(func, order, default) == img).all() @@ -446,15 +447,15 @@ def func(df, healpix): pixel_map = small_sky_order1_catalog.skymap_data(func) pixel_map = {pixel: value.compute() for pixel, value in pixel_map.items()} max_order = max(pixel_map.keys(), key=lambda x: x.order).order - img = np.full(hp.nside2npix(hp.order2nside(max_order)), hp.pixelfunc.UNSEEN) + img = np.full(hp.order2npix(max_order), hp.unseen_pixel()) for pixel, value in pixel_map.items(): dorder = max_order - pixel.order start = pixel.pixel * (4**dorder) end = (pixel.pixel + 1) * (4**dorder) img_order_pixels = np.arange(start, end) img[img_order_pixels] = value - hp.mollview.assert_called_once() - assert (hp.mollview.call_args[0][0] == img).all() + healpy.mollview.assert_called_once() + assert (healpy.mollview.call_args[0][0] == img).all() # pylint: disable=no-member @@ -464,11 +465,11 @@ def test_plot_pixels(small_sky_order1_catalog, mocker): small_sky_order1_catalog.plot_pixels() # Everything will be empty, except the four pixels at order 1. - img = np.full(48, hp.pixelfunc.UNSEEN) + img = np.full(48, hp.unseen_pixel()) img[[44, 45, 46, 47]] = 1 - hp.mollview.assert_called_once() - assert (hp.mollview.call_args[0][0] == img).all() + healpy.mollview.assert_called_once() + assert (healpy.mollview.call_args[0][0] == img).all() def test_square_bracket_columns(small_sky_order1_catalog): diff --git a/tests/lsdb/loaders/dataframe/test_from_dataframe.py b/tests/lsdb/loaders/dataframe/test_from_dataframe.py index 048878d5..0a9da329 100644 --- a/tests/lsdb/loaders/dataframe/test_from_dataframe.py +++ b/tests/lsdb/loaders/dataframe/test_from_dataframe.py @@ -1,8 +1,8 @@ import math import astropy.units as u -import healpy as hp import hipscat as hc +import hipscat.pixel_math.healpix_shim as hp import nested_dask as nd import nested_pandas as npd import numpy as np