Skip to content

Commit

Permalink
Patch sky map when catalog has empty partitions (#474)
Browse files Browse the repository at this point in the history
* Handle empty partitions in inner skymap routine

* Improve test readability
  • Loading branch information
camposandro authored Oct 29, 2024
1 parent c9779aa commit 6888b52
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 3 deletions.
3 changes: 2 additions & 1 deletion src/lsdb/catalog/dataset/healpix_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -336,13 +336,14 @@ def skymap_data(
have the default_value as its result, as well as any pixels for which the aggregate
function returns None.
"""
results = {}
partitions = self.to_delayed()
if order is None:
results = {
pixel: delayed(func)(partitions[index], pixel, **kwargs)
for pixel, index in self._ddf_pixel_map.items()
}
else:
elif len(self.hc_structure.pixel_tree) > 0:
if order < self.hc_structure.pixel_tree.get_max_depth():
raise ValueError(
f"order must be greater than or equal to max order in catalog "
Expand Down
8 changes: 6 additions & 2 deletions src/lsdb/core/plotting/skymap.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,12 @@ def perform_inner_skymap(
**kwargs,
) -> np.ndarray:
"""Splits a partition into pixels at a target order and performs a given function on the new pixels"""
delta_order = target_order - pixel.order
img = np.full(1 << 2 * delta_order, fill_value=default_value)

if len(partition) == 0:
return img

spatial_index = partition.index.to_numpy()
order_pixels = spatial_index_to_healpix(spatial_index, target_order=target_order)

Expand All @@ -28,8 +34,6 @@ def apply_func(df):
return func(df, HealpixPixel(target_order, p), **kwargs)

gb = partition.groupby(order_pixels, sort=False).apply(apply_func)
delta_order = target_order - pixel.order
img = np.full(1 << 2 * delta_order, fill_value=default_value)
min_pixel_value = pixel.pixel << 2 * delta_order
img[gb.index.to_numpy() - min_pixel_value] = gb.to_numpy(na_value=default_value)
return img
Expand Down
39 changes: 39 additions & 0 deletions tests/lsdb/catalog/test_catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -458,6 +458,45 @@ def func(df, healpix):
assert np.array_equal(expected_arr, arr)


def test_skymap_histogram_order_empty(small_sky_order1_catalog):
order = 3

def func(df, healpix):
return len(df) / hp.nside2pixarea(hp.order2nside(healpix.order), degrees=True)

catalog = small_sky_order1_catalog.cone_search(0, 0, 1)
_, non_empty_partitions = catalog._get_non_empty_partitions()
assert len(non_empty_partitions) == 0

img = catalog.skymap_histogram(func, order)
expected_img = np.zeros(hp.order2npix(order))
assert (img == expected_img).all()


def test_skymap_histogram_order_some_partitions_empty(small_sky_order1_catalog):
order = 3

def func(df, healpix):
return len(df) / hp.nside2pixarea(hp.order2nside(healpix.order), degrees=True)

catalog = small_sky_order1_catalog.query("ra > 350 and dec < -50")
_, non_empty_partitions = catalog._get_non_empty_partitions()
assert 0 < len(non_empty_partitions) < catalog._ddf.npartitions

img = catalog.skymap_histogram(func, order)

pixel_map = catalog.skymap_data(func, order)
pixel_map = {pixel: value.compute() for pixel, value in pixel_map.items()}
expected_img = np.zeros(hp.order2npix(order))
for pixel, value in pixel_map.items():
dorder = order - pixel.order
start = pixel.pixel * (4**dorder)
end = (pixel.pixel + 1) * (4**dorder)
img_order_pixels = np.arange(start, end)
expected_img[img_order_pixels] = value
assert (img == expected_img).all()


# pylint: disable=no-member
def test_skymap_plot(small_sky_order1_catalog, mocker):
mocker.patch("lsdb.catalog.dataset.healpix_dataset.plot_healpix_map")
Expand Down

0 comments on commit 6888b52

Please sign in to comment.