Merge pull request #249 from astronomy-commons/sean/interval-pixel-tree

Interval Pixel Tree

src/hipscat/catalog/catalog.py

@@ -25,6 +25,7 @@
     validate_polygon,
     validate_radius,
 )
+from hipscat.pixel_tree.negative_tree import compute_negative_tree_pixels
 
 
 class Catalog(HealpixDataset):
@@ -130,4 +131,4 @@ def generate_negative_tree_pixels(self) -> List[HealpixPixel]:
         Returns:
             List of HealpixPixels representing the 'negative tree' for the catalog.
         """
-        return self.pixel_tree.get_negative_pixels()
+        return compute_negative_tree_pixels(self.pixel_tree)

src/hipscat/catalog/healpix_dataset/healpix_dataset.py

@@ -60,12 +60,7 @@ def _get_partition_info_from_pixels(pixels: PixelInputTypes) -> PartitionInfo:
         if isinstance(pixels, PartitionInfo):
             return pixels
         if isinstance(pixels, PixelTree):
-            return PartitionInfo.from_healpix(
-                [
-                    HealpixPixel(node.hp_order, node.hp_pixel)
-                    for node in pixels.root_pixel.get_all_leaf_descendants()
-                ]
-            )
+            return PartitionInfo.from_healpix(pixels.get_healpix_pixels())
         if pd.api.types.is_list_like(pixels):
             return PartitionInfo.from_healpix(pixels)
         raise TypeError("Pixels must be of type PartitionInfo, PixelTree, or List[HealpixPixel]")

src/hipscat/pixel_math/healpix_pixel.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-import math
 from dataclasses import dataclass
 from typing import List
 
+import numpy as np
+
 from hipscat.pixel_math.hipscat_id import HIPSCAT_ID_HEALPIX_ORDER
 
 
@@ -32,6 +33,13 @@ def __str__(self) -> str:
     def __repr__(self):
         return self.__str__()
 
+    def __getitem__(self, key: int) -> int:
+        if key == 0:
+            return self.order
+        if key == 1:
+            return self.pixel
+        raise IndexError("Invalid healpix index")
+
     def convert_to_lower_order(self, delta_order: int) -> HealpixPixel:
         """Returns the HEALPix pixel that contains the pixel at a lower order
 
@@ -48,12 +56,8 @@ def convert_to_lower_order(self, delta_order: int) -> HealpixPixel:
                 generated at a negative order. Or if delta_order is negative
 
         """
-        if self.order - delta_order < 0:
-            raise ValueError("Pixel Order cannot be below zero")
-        if delta_order < 0:
-            raise ValueError("delta order cannot be below zero")
+        new_pixel = get_lower_order_pixel(self.order, self.pixel, delta_order)
         new_order = self.order - delta_order
-        new_pixel = math.floor(self.pixel / 4**delta_order)
         return HealpixPixel(new_order, new_pixel)
 
     def convert_to_higher_order(self, delta_order: int) -> List[HealpixPixel]:
@@ -71,14 +75,9 @@ def convert_to_higher_order(self, delta_order: int) -> List[HealpixPixel]:
             ValueError: If delta_order + current order is greater than the maximum HEALPix order,
                 pixels cannot be  generated. Or if delta_order is negative
         """
-        if self.order + delta_order > HIPSCAT_ID_HEALPIX_ORDER:
-            raise ValueError(f"Pixel Order cannot be above maximum order {HIPSCAT_ID_HEALPIX_ORDER}")
-        if delta_order < 0:
-            raise ValueError("delta order cannot be below zero")
-        pixels = []
+        new_pixels = get_higher_order_pixels(self.order, self.pixel, delta_order)
         new_order = self.order + delta_order
-        for new_pixel in range(self.pixel * 4**delta_order, (self.pixel + 1) * 4**delta_order):
-            pixels.append(HealpixPixel(new_order, new_pixel))
+        pixels = [HealpixPixel(new_order, new_pixel) for new_pixel in new_pixels]
         return pixels
 
     @property
@@ -98,3 +97,43 @@ def dir(self) -> int:
 
 
 INVALID_PIXEL = HealpixPixel(-1, -1)
+
+
+def get_lower_order_pixel(order: int, pixel: int, delta_order: int) -> int:
+    """Returns the pixel number at a lower order
+
+    Args:
+        order (int): the order of the pixel
+        pixel (int): the pixel number of the pixel in NESTED ordering
+        delta_order (int): the change in order to the new lower order
+
+    Returns:
+        The pixel number at order (order - delta_order) for the pixel that contains the given pixel
+    """
+    if order - delta_order < 0:
+        raise ValueError("Pixel Order cannot be below zero")
+    if delta_order < 0:
+        raise ValueError("delta order cannot be below zero")
+    new_pixel = pixel >> (2 * delta_order)
+    return new_pixel
+
+
+def get_higher_order_pixels(order: int, pixel: int, delta_order: int) -> np.ndarray:
+    """Returns the pixel numbers at a higher order
+
+    Args:
+        order (int): the order of the pixel
+        pixel (int): the pixel number of the pixel in NESTED ordering
+        delta_order (int): the change in order to the new higher order
+
+    Returns:
+        The list of pixel numbers at order (order + delta_order) for the pixels contained by the given pixel
+    """
+    if order == -1:
+        return np.arange(0, 12)
+    if order + delta_order > HIPSCAT_ID_HEALPIX_ORDER:
+        raise ValueError(f"Pixel Order cannot be above maximum order {HIPSCAT_ID_HEALPIX_ORDER}")
+    if delta_order < 0:
+        raise ValueError("delta order cannot be below zero")
+    new_pixels = np.arange(pixel << (2 * delta_order), (pixel + 1) << (2 * delta_order))
+    return new_pixels

src/hipscat/pixel_math/healpix_pixel_convertor.py

@@ -22,3 +22,20 @@ def get_healpix_pixel(pixel: HealpixInputTypes) -> HealpixPixel:
     if isinstance(pixel, HealpixPixel):
         return pixel
     raise TypeError("pixel must either be of type `HealpixPixel` or tuple (order, pixel)")
+
+
+def get_healpix_tuple(pixel: HealpixInputTypes) -> Tuple[int, int]:
+    """Function to convert argument of either HealpixPixel or a tuple of (order, pixel) to a
+    tuple of (order, pixel)
+
+    Args:
+        pixel: an object to be converted to a HealpixPixel object
+    """
+
+    if isinstance(pixel, tuple):
+        if len(pixel) != 2:
+            raise ValueError("Tuple must contain two values: HEALPix order and HEALPix pixel number")
+        return pixel
+    if isinstance(pixel, HealpixPixel):
+        return (pixel.order, pixel.pixel)
+    raise TypeError("pixel must either be of type `HealpixPixel` or tuple (order, pixel)")

src/hipscat/pixel_math/healpix_pixel_function.py

@@ -27,3 +27,27 @@ def get_pixel_argsort(pixels: List[HealpixPixel]):
 
     # Get the argsort of the higher order array.
     return np.argsort(constant_breadth_pixel, kind="stable")
+
+
+def get_pixels_from_intervals(intervals: np.ndarray, tree_order: int) -> np.ndarray:
+    """Computes an array of HEALPix [order, pixel] for an array of intervals
+
+    Args:
+        intervals (np.ndarray): Array of intervals of the start and end pixel numbers of HEALPix pixels.
+            Must be NESTED numbering scheme.
+        tree_order (int): The order of the pixel numbers in the interval array
+
+    Returns (np.ndarray):
+        An array of [order, pixel] in NESTED numbering scheme for each interval in the array.
+    """
+    if intervals.shape[0] == 0:
+        return np.empty((0, 2), dtype=np.int64)
+    orders = np.full(intervals.shape[0], fill_value=-1)
+    pixels = np.full(intervals.shape[0], fill_value=-1)
+    # alignment uses (-1, -1) as a missing pixel, so we can't use the HEALPix math on these elements
+    non_negative_mask = intervals.T[0] >= 0
+    start_intervals = intervals.T[0][non_negative_mask]
+    end_intervals = intervals.T[1][non_negative_mask]
+    orders[non_negative_mask] = tree_order - (np.int64(np.log2(end_intervals - start_intervals)) >> 1)
+    pixels[non_negative_mask] = start_intervals >> 2 * (tree_order - orders[non_negative_mask])
+    return np.array([orders, pixels]).T

src/hipscat/pixel_tree/negative_tree.py

@@ -0,0 +1,22 @@
+from typing import List
+
+from hipscat.pixel_math import HealpixPixel
+from hipscat.pixel_tree import PixelAlignmentType, align_trees
+from hipscat.pixel_tree.pixel_tree import PixelTree
+
+
+def compute_negative_tree_pixels(tree: PixelTree) -> List[HealpixPixel]:
+    """Computes a 'negative pixel tree' consisting of the pixels needed to cover the full sky not in the tree
+
+    Args:
+        tree (PixelTree): the input tree to compute the negative of
+
+    Returns (List[HealpixPixel]):
+        A list of HEALPix pixels needed to cover the part of the sky not covered by the tree, using the least
+        number of pixels possible.
+    """
+    tree_pixels = set(tree.get_healpix_pixels())
+    full_tree = PixelTree.from_healpix([(0, i) for i in range(12)])
+    alignment = align_trees(full_tree, tree, alignment_type=PixelAlignmentType.OUTER)
+    aligned_tree = alignment.pixel_tree
+    return [p for p in aligned_tree.get_healpix_pixels() if p not in tree_pixels]