Skip to content

Commit

Permalink
Use healpix shim for healpy operations.
Browse files Browse the repository at this point in the history
  • Loading branch information
delucchi-cmu committed Sep 9, 2024
1 parent e3600fd commit 19f6579
Show file tree
Hide file tree
Showing 7 changed files with 17 additions and 16 deletions.
4 changes: 2 additions & 2 deletions src/lsdb/catalog/dataset/healpix_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion src/lsdb/core/plotting/skymap.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/lsdb/core/search/polygon_search.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/lsdb/dask/merge_catalog_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/lsdb/loaders/dataframe/margin_catalog_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
19 changes: 10 additions & 9 deletions tests/lsdb/catalog/test_catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion tests/lsdb/loaders/dataframe/test_from_dataframe.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down

0 comments on commit 19f6579

Please sign in to comment.