fix: Pixel annotation with partial bin table

CHANGES.md

@@ -1,5 +1,10 @@
 # Release notes #
 
+## [v0.10.1](https://github.com/open2c/cooler/compare/v0.10.0...v0.10.1)
+
+## Bug fixes
+* fix: Pixel annotation with partial bin table in https://github.com/open2c/cooler/pull/422
+
 ## [v0.10.0](https://github.com/open2c/cooler/compare/v0.9.3...v0.10.0)
 
 ### New features

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "cooler"
-version = "0.10.0"
+version = "0.10.1"
 description = "Sparse binary format for genomic interaction matrices."
 requires-python = ">=3.8"
 license = {text = "BSD-3-Clause"}

src/cooler/api.py

@@ -439,7 +439,7 @@ def chroms(
     lo: int = 0,
     hi: int | None = None,
     fields: list[str] | None = None,
-    **kwargs
+    **kwargs,
 ) -> pd.DataFrame:
     """
     Table describing the chromosomes/scaffolds/contigs used.
@@ -473,7 +473,7 @@ def bins(
     lo: int = 0,
     hi: int | None = None,
     fields: list[str] | None = None,
-    **kwargs
+    **kwargs,
 ) -> pd.DataFrame:
     """
     Table describing the genomic bins that make up the axes of the heatmap.
@@ -530,7 +530,7 @@ def pixels(
     hi: int | None = None,
     fields: list[str] | None = None,
     join: bool = True,
-    **kwargs
+    **kwargs,
 ) -> pd.DataFrame:
     """
     Table describing the nonzero upper triangular pixels of the Hi-C contact
@@ -570,9 +570,7 @@ def pixels(
 
 
 def annotate(
-    pixels: pd.DataFrame,
-    bins: pd.DataFrame | RangeSelector1D,
-    replace: bool = False
+    pixels: pd.DataFrame, bins: pd.DataFrame | RangeSelector1D, replace: bool = False
 ) -> pd.DataFrame:
     """
     Add bin annotations to a data frame of pixels.
@@ -609,6 +607,7 @@ def annotate(
         def _loc_slice(sel, beg, end):
             # slicing a range selector is end-exclusive like iloc
             return sel[beg : end + 1 if end is not None else None]
+
     else:
 
         def _loc_slice(df, beg, end):
@@ -631,7 +630,7 @@ def _loc_slice(df, beg, end):
             bmin, bmax = 0, None
         ann1 = _loc_slice(bins, bmin, bmax)
         anns.append(
-            ann1.iloc[bin1 - bmin]
+            ann1.iloc[bin1 - ann1.index[0]]
             .rename(columns=lambda x: x + "1")
             .reset_index(drop=True)
         )
@@ -647,7 +646,7 @@ def _loc_slice(df, beg, end):
             bmin, bmax = 0, None
         ann2 = _loc_slice(bins, bmin, bmax)
         anns.append(
-            ann2.iloc[bin2 - bmin]
+            ann2.iloc[bin2 - ann2.index[0]]
             .rename(columns=lambda x: x + "2")
             .reset_index(drop=True)
         )

src/cooler/core/__init__.py

@@ -5,5 +5,18 @@
     region_to_extent,
     region_to_offset,
 )
-from ._selectors import RangeSelector1D, RangeSelector2D, _IndexingMixin
+from ._selectors import RangeSelector1D, RangeSelector2D, _IndexingMixin  # noqa
 from ._tableops import delete, get, put
+
+__all__ = [
+    "CSRReader",
+    "DirectRangeQuery2D",
+    "FillLowerRangeQuery2D",
+    "RangeSelector1D",
+    "RangeSelector2D",
+    "region_to_extent",
+    "region_to_offset",
+    "delete",
+    "get",
+    "put",
+]

tests/test_api.py

@@ -99,6 +99,28 @@ def test_annotate_with_partial_bins():
         assert out[col].notnull().all()
 
 
+def test_annotate_with_partial_bins_and_partial_pixels():
+    # Addresses a bug where annotating only certain pixels with a partial
+    # bin-table dataframe would lead to IndexError
+
+    clr = api.Cooler(op.join(datadir, "hg19.GM12878-MboI.matrix.2000kb.cool"))
+    region = "chr2:0-1000000"
+
+    pix = clr.matrix(as_pixels=True, balance=False).fetch(region)
+
+    # Use partial bin table
+    bins_region = clr.bins().fetch(region)
+    out = api.annotate(pix, bins_region)
+    for col in ["chrom1", "start1", "end1", "chrom2", "start2", "end2"]:
+        assert out[col].notnull().all()
+
+    # Compare with full bin table as selector
+    pd.testing.assert_frame_equal(out, api.annotate(pix, clr.bins()))
+
+    # Compare with full bin table materialized
+    pd.testing.assert_frame_equal(out, api.annotate(pix, clr.bins()[:]))
+
+
 def test_matrix():
     clr = api.Cooler(op.join(datadir, "yeast.10kb.cool"))
     region = ("chrI:100345-220254", "chrII:200789-813183")