bugfix(snipping): features with empty regions return snips filled with NaNs (#216)

* bugfix(snipping): features with empty regions return snips filled with NaNs.

Related to 1a63791
which turned out to be not complete fix.

* bugfix(snipping): Docstrings added, code simplified.

* bugfix+tests(snipping): Snipping adapted for pandas failing sometimes while grouping np.nans.

Pandas fails sometimes on groupby with np.nan in the dataframe:
pandas-dev/pandas#35202

With this fix, we replace NaN regions (unannotated) with empty
string, and then do the grouping.

Tests of snipping added, covering NaN in "region" column of the
windows dataframe for snipping.

* style(snipping test): improved

* formatting imporoved

* Update cooltools/snipping.py

Co-authored-by: Nezar Abdennur <[email protected]>

Author: agalitsyna

cooltools/snipping.py

@@ -134,6 +134,16 @@ def assign_regions(features, supports):
 
 def _pileup(data_select, data_snip, arg):
     support, feature_group = arg
+    
+    # check if region is unannotated
+    if len(support) == 0:
+        lo = feature_group["lo"].values[0]
+        hi = feature_group["hi"].values[0]
+        s = hi-lo # Shape of individual snip
+        stack = np.full((s, s, len(feature_group)), np.nan)
+        # return empty snippets if region is not annotated
+        return stack, feature_group["_rank"].values
+
     # check if support region is on- or off-diagonal
     if len(support) == 2:
         region1, region2 = map(bioframe.region.parse_region_string, support)
@@ -183,13 +193,15 @@ def pileup(features, data_select, data_snip, map=map):
         warnings.warn("Some features do not have regions assigned! Some snips will be empty.")
 
     features = features.copy()
+    features['region'] = features.region.fillna('') # fill in unanotated regions with empty string
     features["_rank"] = range(len(features))
 
     # cumul_stack = []
     # orig_rank = []
     cumul_stack, orig_rank = zip(
         *map(
             partial(_pileup, data_select, data_snip),
+            # Note that unannotated regions will form a separate group
             features.groupby("region", sort=False),
         )
     )

tests/test_snipping.py

@@ -0,0 +1,42 @@
+import cooler
+import os.path as op
+
+import cooltools.snipping
+import numpy as np
+
+
+def test_snipper(request):
+    """
+    Test the snipping on matrix:
+    """
+    # Read cool file and create regions out of it:
+    clr = cooler.Cooler(op.join(request.fspath.dirname, "data/CN.mm9.1000kb.cool"))
+
+    regions = [(chrom, 0, clr.chromsizes[chrom]) for chrom in clr.chromnames]
+
+    # Example region with windows, two regions from annotated genomic regions:
+    windows = cooltools.snipping.make_bin_aligned_windows(
+        1_000_000, ["chr10", "chr10"], [10_000_000, 20_000_000], flank_bp=2_000_000
+    )
+
+    windows = cooltools.snipping.assign_regions(windows, regions).reset_index(drop=True)
+
+    snipper = cooltools.snipping.CoolerSnipper(clr)
+    stack = cooltools.snipping.pileup(windows, snipper.select, snipper.snip, map=map)
+
+    # Check that the size of snips is OK and there are two of them:
+    assert stack.shape == (5, 5, 2)
+
+    # Example region with windows, second window comes from unannotated genomic region:
+    windows = cooltools.snipping.make_bin_aligned_windows(
+        1_000_000, ["chr10", "chr10"], [10_000_000, 150_000_000], flank_bp=2_000_000
+    )
+
+    windows = cooltools.snipping.assign_regions(windows, regions).reset_index(drop=True)
+
+    snipper = cooltools.snipping.CoolerSnipper(clr)
+    stack = cooltools.snipping.pileup(windows, snipper.select, snipper.snip, map=map)
+
+    assert stack.shape == (5, 5, 2)
+    assert np.all(np.isfinite(stack[:, :, 0]))
+    assert np.all(np.isnan(stack[:, :, 1]))