Merge pull request #7 from lincc-frameworks-mask-incubator/issue/5/al…

…ignment

Methods for alignment generation

pyproject.toml

@@ -18,7 +18,9 @@ dynamic = ["version"]
 requires-python = ">=3.9, <3.12"
 dependencies = [
     "lsdb",
+    "hipscat",
     "gundam",
+    "numpy<2",
 ]
 
 [project.urls]

src/corrgi/__init__.py

@@ -1,3 +0,0 @@
-from .example_module import greetings, meaning
-
-__all__ = ["greetings", "meaning"]

src/corrgi/alignment.py

@@ -0,0 +1,63 @@
+import itertools
+
+import pandas as pd
+from hipscat.catalog import Catalog
+from hipscat.pixel_tree.pixel_alignment import PixelAlignment
+from hipscat.pixel_tree.pixel_alignment_types import PixelAlignmentType
+
+column_names = [
+    PixelAlignment.PRIMARY_ORDER_COLUMN_NAME,
+    PixelAlignment.PRIMARY_PIXEL_COLUMN_NAME,
+    PixelAlignment.JOIN_ORDER_COLUMN_NAME,
+    PixelAlignment.JOIN_PIXEL_COLUMN_NAME,
+    PixelAlignment.ALIGNED_ORDER_COLUMN_NAME,
+    PixelAlignment.ALIGNED_PIXEL_COLUMN_NAME,
+]
+
+
+def autocorrelation_alignment(catalog: Catalog) -> PixelAlignment:
+    """Determine all pairs of partitions that should be correlated within the same catalog.
+
+    This considers all combinations, without duplicates between the "primary" and "join"
+    pixels in the alignment.
+
+    Args:
+        catalog (Catalog): catalog for autocorrelation
+    Returns:
+        alignment object where the `aligned` columns simply match the left pixel.
+    """
+    upper_triangle = [
+        [left.order, left.pixel, right.order, right.pixel, left.order, left.pixel]
+        for (left, right) in itertools.combinations(catalog.get_healpix_pixels(), 2)
+    ]
+    upper_triangle = pd.DataFrame(upper_triangle, columns=column_names)
+    diagonal = pd.DataFrame(
+        [
+            [pix.order, pix.pixel, pix.order, pix.pixel, pix.order, pix.pixel]
+            for pix in catalog.get_healpix_pixels()
+        ],
+        columns=column_names,
+    )
+    result_mapping = pd.concat([upper_triangle, diagonal])
+    return PixelAlignment(catalog.pixel_tree, result_mapping, PixelAlignmentType.OUTER)
+
+
+def crosscorrelation_alignment(catalog_left: Catalog, catalog_right: Catalog) -> PixelAlignment:
+    """Determine all pairs of partitions that should be correlated between two catalogs.
+
+    This considers the full cross-product of pixels.
+
+    Args:
+        catalog_left (Catalog): left side of the cross-correlation
+        catalog_right (Catalog): right side of the cross-correlation
+    Returns:
+        alignment object where the `aligned` columns simply match the left pixel.
+    """
+    full_product = [
+        [left.order, left.pixel, right.order, right.pixel, left.order, left.pixel]
+        for (left, right) in itertools.product(
+            catalog_left.get_healpix_pixels(), catalog_right.get_healpix_pixels()
+        )
+    ]
+    result_mapping = pd.DataFrame(full_product, columns=column_names)
+    return PixelAlignment(catalog_left.pixel_tree, result_mapping, PixelAlignmentType.OUTER)

tests/corrgi/conftest.py

@@ -0,0 +1,28 @@
+from pathlib import Path
+
+import pytest
+
+
+@pytest.fixture
+def test_data_dir():
+    return Path(__file__).parent.parent / "data"
+
+
+@pytest.fixture
+def data_catalog_dir(test_data_dir):
+    return test_data_dir / "DATA"
+
+
+@pytest.fixture
+def raw_catalog_dir(test_data_dir):
+    return test_data_dir / "RAW"
+
+
+@pytest.fixture
+def dr7_lrg_catalog_dir(test_data_dir):
+    return test_data_dir / "DR7-lrg"
+
+
+@pytest.fixture
+def dr7_lrg_rand_catalog_dir(test_data_dir):
+    return test_data_dir / "DR7-lrg-rand"

tests/corrgi/test_alignment.py

@@ -0,0 +1,20 @@
+import hipscat
+from corrgi.alignment import autocorrelation_alignment, crosscorrelation_alignment
+
+
+def test_autocorrelation_alignment(data_catalog_dir):
+    data_catalog = hipscat.read_from_hipscat(data_catalog_dir)
+    alignment = autocorrelation_alignment(data_catalog)
+    assert len(alignment.pixel_mapping) == 28
+    assert len(alignment.pixel_mapping.columns) == 6
+
+
+def test_crosscorrelation_alignment(dr7_lrg_catalog_dir, dr7_lrg_rand_catalog_dir):
+    dr7_catalog = hipscat.read_from_hipscat(dr7_lrg_catalog_dir)
+    dr7_rand_catalog = hipscat.read_from_hipscat(dr7_lrg_rand_catalog_dir)
+    alignment = crosscorrelation_alignment(dr7_catalog, dr7_rand_catalog)
+    ## dr7_catalog has 12 partitions
+    ## dr7_rand_catalog has 22 partitions
+    ## 12*21 = 252
+    assert len(alignment.pixel_mapping) == 252
+    assert len(alignment.pixel_mapping.columns) == 6

tests/ruff.toml

@@ -0,0 +1,3 @@
+lint.ignore = [
+    "D103", # Allow Missing docstring in public function
+]