Add healpix_to_hipscat_id conversion

src/hipscat/pixel_math/hipscat_id.py

@@ -53,13 +53,18 @@ def compute_hipscat_id(ra_values, dec_values):
     offset_counter = boring_number_index - unique_indices[unique_inverse]
 
     ## Add counter to shifted pixel, and map back to the original, unsorted, values
-    shifted_pixels = mapped_pixels.astype(np.uint64) << (64 - (4 + 2 * HIPSCAT_ID_HEALPIX_ORDER))
-    shifted_pixels = shifted_pixels + offset_counter
+    shifted_pixels = _compute_hipscat_id_from_mapped_pixels(mapped_pixels, offset_counter)
 
     unsort_index = np.argsort(sort_index, kind="stable")
     return shifted_pixels[unsort_index]
 
 
+def _compute_hipscat_id_from_mapped_pixels(mapped_pixels, offset_counter):
+    shifted_pixels = mapped_pixels.astype(np.uint64) << np.uint64(64 - (4 + 2 * HIPSCAT_ID_HEALPIX_ORDER))
+    shifted_pixels = shifted_pixels + offset_counter
+    return shifted_pixels
+
+
 def hipscat_id_to_healpix(ids):
     """Convert some hipscat ids to the healpix pixel at order 19.
     This is just bit-shifting the counter away.
@@ -70,3 +75,26 @@ def hipscat_id_to_healpix(ids):
         list of order 19 pixels from the hipscat id
     """
     return np.asarray(ids, dtype=np.uint64) >> (64 - (4 + 2 * HIPSCAT_ID_HEALPIX_ORDER))
+
+
+def healpix_to_hipscat_id(order: int, pixel: int, counter: int = 0) -> int:
+    """Convert a healpix pixel to a hipscat_id
+
+    This maps the healpix pixel to the lowest pixel number within that pixel at order 19,
+    then shifts and adds the given counter to get a hipscat_id.
+
+    Useful for operations such as filtering by hipscat_id.
+
+    Args:
+        order (int64 | list[int64]): order of pixel to convert
+        pixel (int64 | list[int64]): pixel number in nested ordering of pixel to convert
+        counter (int64 | list[int64]) (Default: 0): counter value in converted hipscat id
+    Returns:
+        list of hipscat ids
+    """
+    order = np.uint64(order)
+    pixel = np.uint64(pixel)
+    counter = np.uint64(counter)
+    pixel_higher_order = pixel * (4 ** (HIPSCAT_ID_HEALPIX_ORDER - order))
+    hipscat_id = _compute_hipscat_id_from_mapped_pixels(pixel_higher_order, counter)
+    return hipscat_id

src/hipscat/pixel_tree/pixel_alignment.py

@@ -73,7 +73,7 @@ def align_trees(
         right: The right tree to align
         alignment_type: The type of alignment describing how to handle nodes which exist in one tree
             but not the other. Options are:
-            
+
                 - inner - only use pixels that appear in both catalogs
                 - left - use all pixels that appear in the left catalog and any overlapping from the right
                 - right - use all pixels that appear in the right catalog and any overlapping from the left

tests/hipscat/pixel_math/test_hipscat_id.py

@@ -1,10 +1,16 @@
 """Test construction (and de-construction) of the hipscat ID"""
 
 import numpy as np
+import healpy as hp
 import numpy.testing as npt
 import pytest
 
-from hipscat.pixel_math.hipscat_id import compute_hipscat_id, hipscat_id_to_healpix
+from hipscat.pixel_math.hipscat_id import (
+    compute_hipscat_id,
+    hipscat_id_to_healpix,
+    HIPSCAT_ID_HEALPIX_ORDER,
+    healpix_to_hipscat_id,
+)
 
 
 def test_single():
@@ -109,3 +115,46 @@ def test_hipscat_id_to_healpix():
     ]
 
     npt.assert_array_equal(result, expected)
+
+
+def test_healpix_to_hipscat_id_single():
+    orders = [3, 3, 4, 1]
+    pixels = [0, 12, 1231, 11]
+    pixels_at_high_order = [
+        p * (4 ** (HIPSCAT_ID_HEALPIX_ORDER - o)) for o, p in zip(orders, pixels)
+    ]
+    lon, lat = hp.pix2ang(
+        [2**HIPSCAT_ID_HEALPIX_ORDER] * len(orders), pixels_at_high_order, nest=True, lonlat=True
+    )
+    actual_hipscat_ids = compute_hipscat_id(lon, lat)
+    test_hipscat_ids = [healpix_to_hipscat_id(o, p) for o, p in zip(orders, pixels)]
+    assert np.all(test_hipscat_ids == actual_hipscat_ids)
+
+
+def test_healpix_to_hipscat_id_array():
+    orders = [3, 3, 4, 1]
+    pixels = [0, 12, 1231, 11]
+    pixels_at_high_order = [
+        p * (4 ** (HIPSCAT_ID_HEALPIX_ORDER - o)) for o, p in zip(orders, pixels)
+    ]
+    lon, lat = hp.pix2ang(
+        [2**HIPSCAT_ID_HEALPIX_ORDER] * len(orders), pixels_at_high_order, nest=True, lonlat=True
+    )
+    actual_hipscat_ids = compute_hipscat_id(lon, lat)
+    test_hipscat_ids = healpix_to_hipscat_id(orders, pixels)
+    assert np.all(test_hipscat_ids == actual_hipscat_ids)
+
+
+def test_healpix_to_hipscat_id_offset():
+    orders = [3, 3, 4, 1]
+    pixels = [0, 0, 1231, 11]
+    offsets = [0, 1, 0, 0]
+    pixels_at_high_order = [
+        p * (4 ** (HIPSCAT_ID_HEALPIX_ORDER - o)) for o, p in zip(orders, pixels)
+    ]
+    lon, lat = hp.pix2ang(
+        [2**HIPSCAT_ID_HEALPIX_ORDER] * len(orders), pixels_at_high_order, nest=True, lonlat=True
+    )
+    actual_hipscat_ids = compute_hipscat_id(lon, lat)
+    test_hipscat_ids = healpix_to_hipscat_id(orders, pixels, offsets)
+    assert np.all(test_hipscat_ids == actual_hipscat_ids)