Merge pull request #268 from astronomy-commons/sean/moc-alignment

Add MOC filter and alignment methods

src/hipscat/catalog/healpix_dataset/healpix_dataset.py

@@ -13,6 +13,8 @@
 from hipscat.catalog.partition_info import PartitionInfo
 from hipscat.io import FilePointer, file_io, paths
 from hipscat.pixel_math import HealpixPixel
+from hipscat.pixel_tree import PixelAlignment, PixelAlignmentType
+from hipscat.pixel_tree.pixel_alignment import align_with_mocs
 from hipscat.pixel_tree.pixel_tree import PixelTree
 
 PixelInputTypes = Union[PartitionInfo, PixelTree, List[HealpixPixel]]
@@ -141,3 +143,35 @@ def filter_from_pixel_list(self, pixels: List[HealpixPixel]) -> Self:
         """
         filtered_catalog_info = dataclasses.replace(self.catalog_info, total_rows=None)
         return self.__class__(filtered_catalog_info, pixels)
+
+    def align(
+        self, other_cat: Self, alignment_type: PixelAlignmentType = PixelAlignmentType.INNER
+    ) -> PixelAlignment:
+        """Performs an alignment to another catalog, using the pixel tree and mocs if available
+
+        An alignment compares the pixel structures of the two catalogs, checking which pixels overlap.
+        The alignment includes the mapping of all pairs of pixels in each tree that overlap with each other,
+        and the aligned tree which consists of the overlapping pixels in the two input catalogs, using the
+        higher order pixels where there is overlap with differing orders.
+
+        For more information, see this document:
+        https://docs.google.com/document/d/1gqb8qb3HiEhLGNav55LKKFlNjuusBIsDW7FdTkc5mJU/edit?usp=sharing
+
+        Args:
+            other_cat (Catalog): The catalog to align to
+            alignment_type (PixelAlignmentType): The type of alignment describing how to handle nodes which
+            exist in one tree but not the other. Mirrors the 'how' argument of a pandas/sql join. 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
+                - "outer" - use all pixels from both catalogs
+
+        Returns (PixelAlignment):
+            A `PixelAlignment` object with the alignment from the two catalogs
+        """
+        left_moc = self.moc if self.moc is not None else self.pixel_tree.to_moc()
+        right_moc = other_cat.moc if other_cat.moc is not None else other_cat.pixel_tree.to_moc()
+        return align_with_mocs(
+            self.pixel_tree, other_cat.pixel_tree, left_moc, right_moc, alignment_type=alignment_type
+        )

src/hipscat/pixel_tree/moc_filter.py

@@ -0,0 +1,68 @@
+import numba
+import numpy as np
+from mocpy import MOC
+from numba import njit
+
+from hipscat.pixel_tree.pixel_tree import PixelTree
+
+
+def filter_by_moc(
+    tree: PixelTree,
+    moc: MOC,
+) -> PixelTree:
+    """Filters a pixel tree to only include the pixels that overlap with the pixels in the moc
+
+    Args:
+        tree (PixelTree): The tree to perform the filtering on
+        moc (mocpy.MOC): The moc to use to filter
+
+    Returns:
+        A new PixelTree object with only the pixels from the input tree that overlap with the moc.
+    """
+    moc_ranges = moc.to_depth29_ranges
+    # Convert tree intervals to order 29 to match moc intervals
+    tree_29_ranges = tree.tree << (2 * (29 - tree.tree_order))
+    tree_mask = perform_filter_by_moc(tree_29_ranges, moc_ranges)
+    return PixelTree(tree.tree[tree_mask], tree.tree_order)
+
+
+@njit(
+    numba.bool_[::1](
+        numba.int64[:, :],
+        numba.uint64[:, :],
+    )
+)
+def perform_filter_by_moc(
+    tree: np.ndarray,
+    moc: np.ndarray,
+) -> np.ndarray:  # pragma: no cover
+    """Performs filtering with lists of pixel intervals
+
+    Input interval lists must be at the same order.
+
+    Args:
+        tree (np.ndarray): Array of pixel intervals to be filtered
+        moc (np.ndarray): Array of pixel intervals to be used to filter
+
+    Returns:
+        A boolean array of dimension tree.shape[0] which masks which pixels in tree overlap with the pixels in
+        moc
+    """
+    output = np.full(tree.shape[0], fill_value=False, dtype=np.bool_)
+    tree_index = 0
+    moc_index = 0
+    while tree_index < len(tree) and moc_index < len(moc):
+        tree_pix = tree[tree_index]
+        moc_pix = moc[moc_index]
+        if tree_pix[0] >= moc_pix[1]:
+            # Don't overlap, tree pixel ahead so move onto next MOC pixel
+            moc_index += 1
+            continue
+        if moc_pix[0] >= tree_pix[1]:
+            # Don't overlap, MOC pixel ahead so move onto next tree pixel
+            tree_index += 1
+            continue
+        # Pixels overlap, so include current tree pixel and check next tree pixel
+        output[tree_index] = True
+        tree_index += 1
+    return output

src/hipscat/pixel_tree/pixel_alignment.py

@@ -1,11 +1,13 @@
-from typing import List
+from typing import Callable, Dict, List
 
 import numba
 import numpy as np
 import pandas as pd
+from mocpy import MOC
 from numba import njit
 
 from hipscat.pixel_math.healpix_pixel_function import get_pixels_from_intervals
+from hipscat.pixel_tree.moc_filter import perform_filter_by_moc
 from hipscat.pixel_tree.pixel_alignment_types import PixelAlignmentType
 from hipscat.pixel_tree.pixel_tree import PixelTree
 
@@ -386,3 +388,58 @@ def perform_align_trees(
     if include_all_left and left_index < len(left):
         _add_remaining_pixels(added_until, left, left_index, LEFT_SIDE, mapping)
     return mapping
+
+
+def filter_alignment_by_moc(alignment: PixelAlignment, moc: MOC) -> PixelAlignment:
+    """Filters an alignment by a moc to only include pixels in the aligned_tree that overlap with the moc,
+    and the corresponding rows in the mapping.
+
+    Args:
+        alignment (PixelAlignment): The pixel alignment to filter
+        moc (mocpy.MOC): The moc to filter by
+
+    Returns:
+        PixelAlignment object with the filtered mapping and tree
+    """
+    moc_ranges = moc.to_depth29_ranges
+    tree_29_ranges = alignment.pixel_tree.tree << (2 * (29 - alignment.pixel_tree.tree_order))
+    tree_mask = perform_filter_by_moc(tree_29_ranges, moc_ranges)
+    new_tree = PixelTree(alignment.pixel_tree.tree[tree_mask], alignment.pixel_tree.tree_order)
+    return PixelAlignment(new_tree, alignment.pixel_mapping.iloc[tree_mask], alignment.alignment_type)
+
+
+def align_with_mocs(
+    left_tree: PixelTree,
+    right_tree: PixelTree,
+    left_moc: MOC,
+    right_moc: MOC,
+    alignment_type: PixelAlignmentType = PixelAlignmentType.INNER,
+) -> PixelAlignment:
+    """Aligns two pixel trees and mocs together, resulting in a pixel alignment with only aligned pixels that
+    have coverage in the mocs.
+
+    Args:
+        left_tree (PixelTree): The left tree to align
+        right_tree (PixelTree): The right tree to align
+        left_moc (mocpy.MOC): the moc with the coverage of the left catalog
+        right_moc (mocpy.MOC): the moc with the coverage of the right catalog
+        alignment_type (PixelAlignmentType): 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
+                - outer - use all pixels from both catalogs
+
+    Returns:
+        The PixelAlignment object with the aligned trees filtered by the coverage in the catalogs.
+    """
+    moc_intersection_methods: Dict[PixelAlignmentType, Callable[[MOC, MOC], MOC]] = {
+        PixelAlignmentType.INNER: lambda l, r: l.intersection(r),
+        PixelAlignmentType.LEFT: lambda l, r: l,
+        PixelAlignmentType.RIGHT: lambda l, r: r,
+        PixelAlignmentType.OUTER: lambda l, r: l.union(r),
+    }
+    filter_moc = moc_intersection_methods[alignment_type](left_moc, right_moc)
+    alignment = align_trees(left_tree, right_tree, alignment_type=alignment_type)
+    return filter_alignment_by_moc(alignment, filter_moc)

src/hipscat/pixel_tree/pixel_tree.py

@@ -3,6 +3,7 @@
 from typing import List, Sequence
 
 import numpy as np
+from mocpy import MOC
 
 from hipscat.pixel_math import HealpixInputTypes, HealpixPixel
 from hipscat.pixel_math.healpix_pixel_convertor import get_healpix_tuple
@@ -84,7 +85,11 @@ def get_healpix_pixels(self) -> List[HealpixPixel]:
         Returns (List[HealpixPixel]):
             A list of the HEALPix pixels in the tree
         """
-        return np.vectorize(HealpixPixel)(self.pixels.T[0], self.pixels.T[1])
+        return [HealpixPixel(p[0], p[1]) for p in self.pixels]
+
+    def to_moc(self) -> MOC:
+        """Returns the MOC object that covers the same pixels as the tree"""
+        return MOC.from_healpix_cells(self.pixels.T[1], self.pixels.T[0], self.tree_order)
 
     @classmethod
     def from_healpix(cls, healpix_pixels: Sequence[HealpixInputTypes], tree_order=None) -> PixelTree:

tests/hipscat/pixel_tree/conftest.py

@@ -17,20 +17,23 @@ def pixel_tree_1():
 
 
 @pytest.fixture
-def pixel_tree_2():
-    return PixelTree.from_healpix(
-        [
-            HealpixPixel(0, 10),
-            HealpixPixel(1, 33),
-            HealpixPixel(1, 35),
-            HealpixPixel(1, 44),
-            HealpixPixel(1, 45),
-            HealpixPixel(1, 46),
-            HealpixPixel(2, 128),
-            HealpixPixel(2, 130),
-            HealpixPixel(2, 131),
-        ]
-    )
+def pixel_tree_2_pixels():
+    return [
+        HealpixPixel(2, 128),
+        HealpixPixel(2, 130),
+        HealpixPixel(2, 131),
+        HealpixPixel(1, 33),
+        HealpixPixel(1, 35),
+        HealpixPixel(0, 10),
+        HealpixPixel(1, 44),
+        HealpixPixel(1, 45),
+        HealpixPixel(1, 46),
+    ]
+
+
+@pytest.fixture
+def pixel_tree_2(pixel_tree_2_pixels):
+    return PixelTree.from_healpix(pixel_tree_2_pixels)
 
 
 @pytest.fixture

tests/hipscat/pixel_tree/test_moc_filter.py

@@ -0,0 +1,59 @@
+import numpy as np
+from mocpy import MOC
+
+from hipscat.pixel_math import HealpixPixel
+from hipscat.pixel_tree.moc_filter import filter_by_moc
+
+
+def test_moc_filter(pixel_tree_2):
+    orders = np.array([1, 1, 2])
+    pixels = np.array([45, 46, 128])
+    moc = MOC.from_healpix_cells(pixels, orders, 2)
+    filtered_tree = filter_by_moc(pixel_tree_2, moc)
+    assert filtered_tree.get_healpix_pixels() == [
+        HealpixPixel(2, 128),
+        HealpixPixel(1, 45),
+        HealpixPixel(1, 46),
+    ]
+
+
+def test_moc_filter_lower_order(pixel_tree_2):
+    orders = np.array([0, 1])
+    pixels = np.array([11, 32])
+    moc = MOC.from_healpix_cells(pixels, orders, 1)
+    filtered_tree = filter_by_moc(pixel_tree_2, moc)
+    assert filtered_tree.get_healpix_pixels() == [
+        HealpixPixel(2, 128),
+        HealpixPixel(2, 130),
+        HealpixPixel(2, 131),
+        HealpixPixel(1, 44),
+        HealpixPixel(1, 45),
+        HealpixPixel(1, 46),
+    ]
+
+
+def test_moc_filter_higher_order(pixel_tree_2):
+    orders = np.array([1, 3])
+    pixels = np.array([40, 520])
+    moc = MOC.from_healpix_cells(pixels, orders, 3)
+    filtered_tree = filter_by_moc(pixel_tree_2, moc)
+    assert filtered_tree.get_healpix_pixels() == [
+        HealpixPixel(2, 130),
+        HealpixPixel(0, 10),
+    ]
+
+
+def test_moc_filter_empty_moc(pixel_tree_2):
+    orders = np.array([])
+    pixels = np.array([])
+    moc = MOC.from_healpix_cells(pixels, orders, 0)
+    filtered_tree = filter_by_moc(pixel_tree_2, moc)
+    assert filtered_tree.get_healpix_pixels() == []
+
+
+def test_moc_filter_empty_result(pixel_tree_2):
+    orders = np.array([0])
+    pixels = np.array([1])
+    moc = MOC.from_healpix_cells(pixels, orders, 0)
+    filtered_tree = filter_by_moc(pixel_tree_2, moc)
+    assert filtered_tree.get_healpix_pixels() == []