Merge pull request #164 from nvictus/feat-mark-compress-runs

Add mark_runs and merge_runs

.github/workflows/ci.yml

@@ -27,6 +27,4 @@ jobs:
  - run: |
  python -m pip install --upgrade pip hatch
  pip install -e .[dev]
- # Stop the build if there are Python syntax errors or undefined names
- ruff . --select=E9,F63,F7,F82 --show-source
  hatch run test

bioframe/__init__.py

@@ -65,6 +65,8 @@
  "expand",
  "merge",
  "overlap",
+ "mark_runs",
+ "merge_runs",
  "select",
  "select_indices",
  "select_labels",
@@ -107,6 +109,8 @@
  frac_gene_coverage,
  frac_mapped,
  make_chromarms,
+ mark_runs,
+ merge_runs,
  pair_by_distance,
  seq_gc,
 )

bioframe/extras.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy as np
 import pandas as pd
 
@@ -13,6 +15,8 @@
  "seq_gc",
  "frac_gene_coverage",
  "pair_by_distance",
+ "mark_runs",
+ "merge_runs"
 ]
 
 
@@ -538,3 +542,199 @@ def pair_by_distance(
  out_df.reset_index(drop=True, inplace=True)
 
  return out_df
+
+
+def mark_runs(
+ df: pd.DataFrame,
+ col: str,
+ *,
+ allow_overlaps: bool = False,
+ reset_counter: bool = True,
+ run_col: str = 'run',
+ cols: tuple[str, str, str] | None = None,
+) -> pd.DataFrame:
+ """
+ Mark runs of spatially consecutive intervals sharing the same value of
+ ``col``.
+
+ Parameters
+ ----------
+ df : DataFrame
+ A bioframe dataframe.
+ col : str
+ The column to mark runs of values for.
+ allow_overlaps : bool, optional [default: False]
+ If True, allow intervals in ``df`` to overlap. This may cause
+ unexpected results.
+ reset_counter : bool, optional [default: True]
+ If True, reset the run counter for each chromosome.
+ run_col : str, optional [default: 'run']
+ The name of the column to store the run numbers in.
+
+ Returns
+ -------
+ pandas.DataFrame
+ A reordered copy the input dataframe with an additional column 'run'
+ marking runs of values in the input column.
+
+ Notes
+ -----
+ This is similar to :func:`cluster`, but only clusters intervals sharing
+ the same value of ``col``.
+
+ Examples
+ --------
+
+ >>> df = pd.DataFrame({
+ ... 'chrom': ['chr1', 'chr1', 'chr1', 'chr1', 'chr1', 'chr1'],
+ ... 'start': [0, 100, 200, 300, 400, 500],
+ ... 'end': [100, 200, 300, 400, 500, 600],
+ ... 'value': [1, 1, 1, 2, 2, 2],
+ ... })
+
+ >>> mark_runs(df, 'value')
+ chrom start end value run
+ 0 chr1 0 100 1 0
+ 1 chr1 100 200 1 0
+ 2 chr1 200 300 1 0
+ 3 chr1 300 400 2 1
+ 4 chr1 400 500 2 1
+ 5 chr1 500 600 2 1
+
+ See Also
+ --------
+ merge_runs
+ cluster
+ merge
+ """
+ ck, sk, ek = _get_default_colnames() if cols is None else cols
+
+ if not allow_overlaps and len(ops.overlap(df, df)) > len(df):
+ raise ValueError("Not a proper bedGraph: found overlapping intervals.")
+
+ result = []
+ n_runs = 0
+
+ for _, group in df.groupby(ck, sort=False):
+ group = group.sort_values([sk, ek])
+ starts = group[sk].to_numpy()
+ ends = group[ek].to_numpy()
+
+ # Extend ends by running max
+ ends = np.maximum.accumulate(ends)
+
+ # Find borders of interval clusters and assign cluster ids
+ is_cluster_border = np.r_[True, starts[1:] > ends[:-1], False]
+
+ # Find borders of consecutive equal values
+ values = group[col].to_numpy()
+ if values.dtype.kind == 'f':
+ is_value_border = np.r_[
+ True,
+ ~np.isclose(values[1:], values[:-1], equal_nan=True),
+ False
+ ]
+ else:
+ is_value_border = np.r_[True, values[1:] != values[:-1], False]
+
+ # Find index extents of runs
+ is_border = is_cluster_border | is_value_border
+ sum_borders = np.cumsum(is_border)
+ run_ids = sum_borders[:-1] - 1
+
+ # Assign run numbers to intervals
+ if reset_counter:
+ n_runs = 0
+ group[run_col] = n_runs + run_ids
+ n_runs += sum_borders[-1]
+
+ result.append(group)
+
+ return pd.concat(result)
+
+
+def merge_runs(
+ df: pd.DataFrame,
+ col: str,
+ *,
+ allow_overlaps: bool = False,
+ agg: dict | None = None,
+ cols: tuple[str, str, str] | None = None,
+) -> pd.DataFrame:
+ """
+ Merge runs of spatially consecutive intervals sharing the same value of
+ ``col``.
+
+ Parameters
+ ----------
+ df : DataFrame
+ A bioframe dataframe.
+ col : str
+ The column to compress runs of values for.
+ allow_overlaps : bool, optional [default: False]
+ If True, allow intervals in ``df`` to overlap. This may cause
+ unexpected results.
+ agg : dict, optional [default: None]
+ A dictionary of additional column names and aggregation functions to
+ apply to each run. Takes the format:
+ {'agg_name': ('column_name', 'agg_func')}
+
+ Returns
+ -------
+ pandas.DataFrame
+ Dataframe with consecutive intervals in the same run merged.
+
+ Notes
+ -----
+ This is similar to :func:`merge`, but only merges intervals sharing
+ the same value of ``col``.
+
+ Examples
+ --------
+ >>> df = pd.DataFrame({
+ ... 'chrom': ['chr1', 'chr1', 'chr1', 'chr1', 'chr1', 'chr1'],
+ ... 'start': [0, 100, 200, 300, 400, 500],
+ ... 'end': [100, 200, 300, 400, 500, 600],
+ ... 'value': [1, 1, 1, 2, 2, 2],
+ ... })
+
+ >>> merge_runs(df, 'value')
+ chrom start end value
+ 0 chr1 0 300 1
+ 1 chr1 300 600 2
+
+ >>> merge_runs(df, 'value', agg={'sum': ('value', 'sum')})
+ chrom start end value sum
+ 0 chr1 0 300 1 3
+ 1 chr1 300 600 2 6
+
+ See Also
+ --------
+ mark_runs
+ cluster
+ merge
+ """
+ ck, sk, ek = _get_default_colnames() if cols is None else cols
+
+ if agg is None:
+ agg = {}
+
+ df_runs = mark_runs(
+ df,
+ col,
+ allow_overlaps=allow_overlaps,
+ reset_counter=False,
+ run_col='_run',
+ )
+ df_merged = (
+ df_runs
+ .groupby('_run')
+ .agg(**{
+ ck: (ck, 'first'),
+ sk: (sk, 'min'),
+ ek: (ek, 'max'),
+ col: (col, 'first'),
+ **agg
+ })
+ )
+ return df_merged.reset_index(drop=True)

pyproject.toml

@@ -35,6 +35,7 @@ classifiers = [
  "Programming Language :: Python :: 3.11",
 ]
 readme = "README.md"
+requires-python = ">=3.8"
 dependencies = [
  "matplotlib",
  "numpy>=1.10, <2",
@@ -49,6 +50,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
  "biopython",
+ "pre-commit",
  "pysam",
  "pybbi",
  "pytest",

tests/test_extras.py

@@ -304,3 +304,102 @@ def test_pair_by_distance():
  bioframe.pair_by_distance(
  df, min_sep=0, max_sep=9, min_intervening=10, max_intervening=9
  )
+
+
+def test_mark_merge_runs():
+ df1 = pd.DataFrame([
+ # chr1
+ # consecutive run of "c"
+ ["chr1", 85563, 129897, "c", 0.2],
+ ["chr1", 129897, 508340, "c", 0.8],
+ ["chr1", 508340, 620903, "c", 0.5],
+
+ # singleton run of "c" separated by 1bp from previous run
+ ["chr1", 620904, 688020, "c", 0.7],
+
+ # consecutive with previous interval but different value of "name"
+ ["chr1", 688020, 858415, "b", 0.8],
+
+ # chr2
+ ["chr2", 548402, 639680, "a", 0.6],
+ ["chr2", 639680, 1026586, "b", 0.8],
+
+ # chr3
+ ["chr3", 260538, 272930, "c", 0.5],
+ ["chr3", 272930, 470969, "c", 0.5],
+ ["chr3", 470969, 502336, "c", 0.5],
+ ], columns=["chrom", "start", "end", "name", "score"])
+
+ runs = bioframe.mark_runs(df1, "name")
+ assert (
+ runs["name"].to_numpy()
+ == np.array(["c", "c", "c", "c", "b", "a", "b", "c", "c", "c"])
+ ).all()
+ assert (
+ runs["run"].to_numpy()
+ == np.array([0, 0, 0, 1, 2, 0, 1, 0, 0, 0])
+ ).all()
+
+ runs = bioframe.mark_runs(df1, "name", reset_counter=False)
+ assert (
+ runs["run"].to_numpy()
+ == np.array([0, 0, 0, 1, 2, 3, 4, 5, 5, 5])
+ ).all()
+
+ runs = bioframe.mark_runs(df1, "name", run_col="foo", reset_counter=False)
+ assert (
+ runs["foo"].to_numpy()
+ == np.array([0, 0, 0, 1, 2, 3, 4, 5, 5, 5])
+ ).all()
+
+ merged = bioframe.merge_runs(
+ df1, "name", agg={"score_mean": ("score", "mean")}
+ )
+ assert (
+ merged["name"].to_numpy()
+ == np.array(["c", "c", "b", "a", "b", "c"])
+ ).all()
+ assert np.allclose(
+ merged["score_mean"].to_numpy(),
+ np.array([0.5, 0.7, 0.8, 0.6, 0.8, 0.5]),
+ )
+
+
+def test_mark_merge_runs__with_overlaps():
+ df1 = pd.DataFrame([
+ # chr1
+ # consecutive run of "c"
+ ["chr1", 85563, 129897, "c", 0.2],
+ ["chr1", 129897, 508340, "c", 0.8],
+ ["chr1", 508340, 620903, "c", 0.5],
+
+ # singleton run of "c" separated by 1bp from previous run
+ ["chr1", 620904, 688020, "c", 0.7],
+
+ # consecutive with previous interval but different value of "name"
+ ["chr1", 688020, 858415, "b", 0.8],
+ # overlapping with previous interval
+ ["chr1", 700000, 900000, "b", 0.8],
+
+ # chr2
+ ["chr2", 548402, 639680, "a", 0.6],
+ ["chr2", 639680, 1026586, "b", 0.8],
+
+ # chr3
+ ["chr3", 260538, 272930, "c", 0.5],
+ ["chr3", 272930, 470969, "c", 0.5],
+ ["chr3", 470969, 502336, "c", 0.5],
+ ], columns=["chrom", "start", "end", "name", "score"])
+
+ with pytest.raises(ValueError):
+ bioframe.mark_runs(df1, "name")
+
+ runs = bioframe.mark_runs(df1, "name", allow_overlaps=True)
+ assert (
+ runs["name"].to_numpy()
+ == np.array(["c", "c", "c", "c", "b", "b", "a", "b", "c", "c", "c"])
+ ).all()
+ assert (
+ runs["run"].to_numpy()
+ == np.array([0, 0, 0, 1, 2, 2, 0, 1, 0, 0, 0])
+ ).all()