from __future__ import annotations import os import re from collections import OrderedDict, defaultdict from collections.abc import Generator, Iterable, Iterator from contextlib import contextmanager from typing import IO, Any import h5py import numpy as np import pandas as pd from pandas.api.types import is_integer, is_scalar from ._typing import GenomicRangeSpecifier, GenomicRangeTuple def partition(start: int, stop: int, step: int) -> Iterator[tuple[int, int]]: """Partition an integer interval into equally-sized subintervals. Like builtin :py:func:`range`, but yields pairs of end points. Examples -------- >>> for lo, hi in partition(0, 9, 2): print(lo, hi) 0 2 2 4 4 6 6 8 8 9 """ return ((i, min(i + step, stop)) for i in range(start, stop, step)) def parse_cooler_uri(s: str) -> tuple[str, str]: """ Parse a Cooler URI string e.g. /path/to/mycoolers.cool::/path/to/cooler """ parts = s.split("::") if len(parts) == 1: file_path, group_path = parts[0], "/" elif len(parts) == 2: file_path, group_path = parts if not group_path.startswith("/"): group_path = "/" + group_path else: raise ValueError("Invalid Cooler URI string") return file_path, group_path def atoi(s: str) -> int: return int(s.replace(",", "")) def parse_humanized(s: str) -> int: _NUMERIC_RE = re.compile("([0-9,.]+)") _, value, unit = _NUMERIC_RE.split(s.replace(",", "")) if not len(unit): return int(value) value = float(value) unit = unit.upper().strip() if unit in ("K", "KB"): value *= 1000 elif unit in ("M", "MB"): value *= 1000000 elif unit in ("G", "GB"): value *= 1000000000 else: raise ValueError(f"Unknown unit '{unit}'") return int(value) def parse_region_string(s: str) -> tuple[str, int | None, int | None]: """ Parse a UCSC-style genomic region string into a triple. Parameters ---------- s : str UCSC-style string, e.g. "chr5:10,100,000-30,000,000". Ensembl and FASTA style sequence names are allowed. End coordinate must be greater than or equal to start. Returns ------- (str, int or None, int or None) """ def _tokenize(s): token_spec = [ ("HYPHEN", r"-"), ("COORD", r"[0-9,]+(\.[0-9]*)?(?:[a-z]+)?"), ("OTHER", r".+"), ] pattern = r"|\s*".join([rf"(?P<{pair[0]}>{pair[1]})" for pair in token_spec]) tok_regex = re.compile(rf"\s*{pattern}", re.IGNORECASE) for match in tok_regex.finditer(s): typ = match.lastgroup yield typ, match.group(typ) def _check_token(typ, token, expected): if typ is None: raise ValueError("Expected {} token missing".format(" or ".join(expected))) else: if typ not in expected: raise ValueError(f'Unexpected token "{token}"') def _expect(tokens): typ, token = next(tokens, (None, None)) _check_token(typ, token, ["COORD"]) start = parse_humanized(token) typ, token = next(tokens, (None, None)) _check_token(typ, token, ["HYPHEN"]) typ, token = next(tokens, (None, None)) if typ is None: return start, None _check_token(typ, token, ["COORD"]) end = parse_humanized(token) if end < start: raise ValueError("End coordinate less than start") return start, end parts = s.split(":") chrom = parts[0].strip() if not len(chrom): raise ValueError("Chromosome name cannot be empty") if len(parts) < 2: return (chrom, None, None) start, end = _expect(_tokenize(parts[1])) return (chrom, start, end) def parse_region( reg: GenomicRangeSpecifier, chromsizes: dict | pd.Series | None = None ) -> GenomicRangeTuple: """ Genomic regions are represented as half-open intervals (0-based starts, 1-based ends) along the length coordinate of a contig/scaffold/chromosome. Parameters ---------- reg : str or tuple UCSC-style genomic region string, or Triple (chrom, start, end), where ``start`` or ``end`` may be ``None``. chromsizes : mapping, optional Lookup table of scaffold lengths to check against ``chrom`` and the ``end`` coordinate. Required if ``end`` is not supplied. Returns ------- A well-formed genomic region triple (str, int, int) """ if isinstance(reg, str): chrom, start, end = parse_region_string(reg) else: chrom, start, end = reg start = int(start) if start is not None else start end = int(end) if end is not None else end try: clen = chromsizes[chrom] if chromsizes is not None else None except KeyError as e: raise ValueError(f"Unknown sequence label: {chrom}") from e start = 0 if start is None else start if end is None: if clen is None: # TODO --- remove? raise ValueError("Cannot determine end coordinate.") end = clen if end < start: raise ValueError("End cannot be less than start") if start < 0 or (clen is not None and end > clen): raise ValueError(f"Genomic region out of bounds: [{start}, {end})") return chrom, start, end def natsort_key(s: str, _NS_REGEX=re.compile(r"(\d+)", re.U)) -> tuple: return tuple([int(x) if x.isdigit() else x for x in _NS_REGEX.split(s) if x]) def natsorted(iterable: Iterable[str]) -> list[str]: return sorted(iterable, key=natsort_key) def argnatsort(array: Iterable[str]) -> np.ndarray: array = np.asarray(array) if not len(array): return np.array([], dtype=int) cols = tuple(zip(*(natsort_key(x) for x in array))) return np.lexsort(cols[::-1]) def read_chromsizes( filepath_or: str | IO[str], name_patterns: tuple[str, ...] = (r"^chr[0-9]+$", r"^chr[XY]$", r"^chrM$"), all_names: bool = False, **kwargs, ) -> pd.Series: """ Parse a ``.chrom.sizes`` or ``.chromInfo.txt`` file from the UCSC database, where ``db`` is a genome assembly name. Parameters ---------- filepath_or : str or file-like Path or url to text file, or buffer. name_patterns : sequence, optional Sequence of regular expressions to capture desired sequence names. Each corresponding set of records will be sorted in natural order. all_names : bool, optional Whether to return all contigs listed in the file. Default is ``False``. Returns ------- :py:class:`pandas.Series` Series of integer bp lengths indexed by sequence name. References ---------- * `UCSC assembly terminology `_ * `GRC assembly terminology `_ """ if isinstance(filepath_or, str) and filepath_or.endswith(".gz"): kwargs.setdefault("compression", "gzip") chromtable = pd.read_csv( filepath_or, sep="\t", usecols=[0, 1], names=["name", "length"], dtype={"name": str}, **kwargs, ) if not all_names: parts = [] for pattern in name_patterns: part = chromtable[chromtable["name"].str.contains(pattern)] part = part.iloc[argnatsort(part["name"])] parts.append(part) chromtable = pd.concat(parts, axis=0) chromtable.index = chromtable["name"].values return chromtable["length"] def fetch_chromsizes(db: str, **kwargs) -> pd.Series: """ Download chromosome sizes from UCSC as a :py:class:`pandas.Series`, indexed by chromosome label. """ return read_chromsizes( f"http://hgdownload.soe.ucsc.edu/goldenPath/{db}/database/chromInfo.txt.gz", **kwargs, ) def load_fasta(names: list[str], *filepaths: str) -> OrderedDict[str, Any]: """ Load lazy FASTA records from one or multiple files without reading them into memory. Parameters ---------- names : sequence of str Names of sequence records in FASTA file or files. filepaths : str Paths to one or more FASTA files to gather records from. Returns ------- OrderedDict of sequence name -> sequence record """ import pyfaidx if len(filepaths) == 0: raise ValueError("Need at least one file") if len(filepaths) == 1: fa = pyfaidx.Fasta(filepaths[0], as_raw=True) else: fa = {} for filepath in filepaths: fa.update(pyfaidx.Fasta(filepath, as_raw=True).records) records = OrderedDict((chrom, fa[chrom]) for chrom in names) return records def binnify(chromsizes: pd.Series, binsize: int) -> pd.DataFrame: """ Divide a genome into evenly sized bins. Parameters ---------- chromsizes : Series pandas Series indexed by chromosome name with chromosome lengths in bp. binsize : int size of bins in bp Returns ------- bins : :py:class:`pandas.DataFrame` Dataframe with columns: ``chrom``, ``start``, ``end``. """ def _each(chrom): clen = chromsizes[chrom] n_bins = int(np.ceil(clen / binsize)) binedges = np.arange(0, (n_bins + 1)) * binsize binedges[-1] = clen return pd.DataFrame( {"chrom": [chrom] * n_bins, "start": binedges[:-1], "end": binedges[1:]}, columns=["chrom", "start", "end"], ) bintable = pd.concat(map(_each, chromsizes.keys()), axis=0, ignore_index=True) bintable["chrom"] = pd.Categorical( bintable["chrom"], categories=list(chromsizes.index), ordered=True ) return bintable make_bintable = binnify def digest(fasta_records: OrderedDict[str, Any], enzyme: str) -> pd.DataFrame: """ Divide a genome into restriction fragments. Parameters ---------- fasta_records : OrderedDict Dictionary of chromosome names to sequence records. enzyme: str Name of restriction enzyme (e.g., 'DpnII'). Returns ------- frags : :py:class:`pandas.DataFrame` Dataframe with columns: ``chrom``, ``start``, ``end``. """ try: import Bio.Restriction as biorst import Bio.Seq as bioseq except ImportError: raise ImportError( "Biopython is required to find restriction fragments." ) from None # http://biopython.org/DIST/docs/cookbook/Restriction.html#mozTocId447698 chroms = fasta_records.keys() try: cut_finder = getattr(biorst, enzyme).search except AttributeError as e: raise ValueError(f"Unknown enzyme name: {enzyme}") from e def _each(chrom): seq = bioseq.Seq(str(fasta_records[chrom][:])) cuts = np.r_[0, np.array(cut_finder(seq)) + 1, len(seq)].astype(np.int64) n_frags = len(cuts) - 1 frags = pd.DataFrame( {"chrom": [chrom] * n_frags, "start": cuts[:-1], "end": cuts[1:]}, columns=["chrom", "start", "end"], ) return frags return pd.concat(map(_each, chroms), axis=0, ignore_index=True) def get_binsize(bins: pd.DataFrame) -> int | None: """ Infer bin size from a bin DataFrame. Assumes that the last bin of each contig is allowed to differ in size from the rest. Returns ------- int or None if bins are non-uniform """ sizes = set() for _chrom, group in bins.groupby("chrom", observed=True): sizes.update((group["end"] - group["start"]).iloc[:-1].unique()) if len(sizes) > 1: return None if len(sizes) == 1: return next(iter(sizes)) else: return None def get_chromsizes(bins: pd.DataFrame) -> pd.Series: """ Infer chromsizes Series from a bin DataFrame. Assumes that the last bin of each contig is allowed to differ in size from the rest. Returns ------- int or None if bins are non-uniform """ chromtable = ( bins.drop_duplicates(["chrom"], keep="last")[["chrom", "end"]] .reset_index(drop=True) .rename(columns={"chrom": "name", "end": "length"}) ) chroms, lengths = list(chromtable["name"]), list(chromtable["length"]) return pd.Series(index=chroms, data=lengths) def bedslice( grouped, chromsizes: pd.Series | dict, region: GenomicRangeSpecifier, ) -> pd.DataFrame: """ Range query on a BED-like dataframe with non-overlapping intervals. """ chrom, start, end = parse_region(region, chromsizes) result = grouped.get_group(chrom) if start > 0 or end < chromsizes[chrom]: lo = result["end"].values.searchsorted(start, side="right") hi = lo + result["start"].values[lo:].searchsorted(end, side="left") result = result.iloc[lo:hi] return result def asarray_or_dataset(x: Any) -> np.ndarray | h5py.Dataset: return x if isinstance(x, h5py.Dataset) else np.asarray(x) def rlencode( array: np.ndarray, chunksize: int | None = None ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ Run length encoding. Based on http://stackoverflow.com/a/32681075, which is based on the rle function from R. Parameters ---------- x : 1D array_like Input array to encode dropna: bool, optional Drop all runs of NaNs. Returns ------- start positions, run lengths, run values """ where = np.flatnonzero array = asarray_or_dataset(array) n = len(array) if n == 0: return ( np.array([], dtype=int), np.array([], dtype=int), np.array([], dtype=array.dtype), ) if chunksize is None: chunksize = n starts, values = [], [] last_val = np.nan for i in range(0, n, chunksize): x = array[i : i + chunksize] locs = where(x[1:] != x[:-1]) + 1 if x[0] != last_val: locs = np.r_[0, locs] starts.append(i + locs) values.append(x[locs]) last_val = x[-1] starts = np.concatenate(starts) lengths = np.diff(np.r_[starts, n]) values = np.concatenate(values) return starts, lengths, values def cmd_exists(cmd: str) -> bool: return any( os.access(os.path.join(path, cmd), os.X_OK) for path in os.environ["PATH"].split(os.pathsep) ) def mad(data: np.ndarray, axis: int | None = None) -> np.ndarray: return np.median(np.abs(data - np.median(data, axis)), axis) @contextmanager def open_hdf5( fp: str | h5py.Group, mode: str = "r", *args, **kwargs ) -> Generator[h5py.Group, None, None]: """ Context manager like ``h5py.File`` but accepts already open HDF5 file handles which do not get closed on teardown. Parameters ---------- fp : str or ``h5py.File`` object If an open file object is provided, it passes through unchanged, provided that the requested mode is compatible. If a filepath is passed, the context manager will close the file on tear down. mode : str * r Readonly, file must exist * r+ Read/write, file must exist * a Read/write if exists, create otherwise * w Truncate if exists, create otherwise * w- or x Fail if exists, create otherwise """ if isinstance(fp, str): own_fh = True fh = h5py.File(fp, mode, *args, **kwargs) else: own_fh = False if mode == "r" and fp.file.mode == "r+": # warnings.warn("File object provided is writeable but intent is read-only") pass elif mode in ("r+", "a") and fp.file.mode == "r": raise ValueError("File object provided is not writeable") elif mode == "w": raise ValueError("Cannot truncate open file") elif mode in ("w-", "x"): raise ValueError("File exists") fh = fp try: yield fh finally: if own_fh: fh.close() class closing_hdf5(h5py.Group): def __init__(self, grp: h5py.Group): super().__init__(grp.id) def __enter__(self) -> h5py.Group: return self def __exit__(self, *exc_info) -> None: return self.file.close() def close(self) -> None: self.file.close() def attrs_to_jsonable(attrs: h5py.AttributeManager) -> dict: out = dict(attrs) for k, v in attrs.items(): try: out[k] = v.item() except ValueError: out[k] = v.tolist() except AttributeError: out[k] = v return out def infer_meta(x, index=None): # pragma: no cover """ Extracted and modified from dask/dataframe/utils.py : make_meta (BSD licensed) Create an empty pandas object containing the desired metadata. Parameters ---------- x : dict, tuple, list, pd.Series, pd.DataFrame, pd.Index, dtype, scalar To create a DataFrame, provide a `dict` mapping of `{name: dtype}`, or an iterable of `(name, dtype)` tuples. To create a `Series`, provide a tuple of `(name, dtype)`. If a pandas object, names, dtypes, and index should match the desired output. If a dtype or scalar, a scalar of the same dtype is returned. index : pd.Index, optional Any pandas index to use in the metadata. If none provided, a `RangeIndex` will be used. Examples -------- >>> make_meta([('a', 'i8'), ('b', 'O')]) Empty DataFrame Columns: [a, b] Index: [] >>> make_meta(('a', 'f8')) Series([], Name: a, dtype: float64) >>> make_meta('i8') 1 """ _simple_fake_mapping = { "b": np.bool_(True), "V": np.void(b" "), "M": np.datetime64("1970-01-01"), "m": np.timedelta64(1), "S": np.str_("foo"), "a": np.str_("foo"), "U": np.str_("foo"), "O": "foo", } UNKNOWN_CATEGORIES = "__UNKNOWN_CATEGORIES__" def _scalar_from_dtype(dtype): if dtype.kind in ("i", "f", "u"): return dtype.type(1) elif dtype.kind == "c": return dtype.type(complex(1, 0)) elif dtype.kind in _simple_fake_mapping: o = _simple_fake_mapping[dtype.kind] return o.astype(dtype) if dtype.kind in ("m", "M") else o else: raise TypeError(f"Can't handle dtype: {dtype}") def _nonempty_scalar(x): if isinstance(x, (pd.Timestamp, pd.Timedelta, pd.Period)): return x elif np.isscalar(x): dtype = x.dtype if hasattr(x, "dtype") else np.dtype(type(x)) return _scalar_from_dtype(dtype) else: raise TypeError("Can't handle meta of type " f"'{type(x).__name__}'") def _empty_series(name, dtype, index=None): if isinstance(dtype, str) and dtype == "category": return pd.Series( pd.Categorical([UNKNOWN_CATEGORIES]), name=name, index=index ).iloc[:0] return pd.Series([], dtype=dtype, name=name, index=index) if hasattr(x, "_meta"): return x._meta if isinstance(x, (pd.Series, pd.DataFrame)): return x.iloc[0:0] elif isinstance(x, pd.Index): return x[0:0] index = index if index is None else index[0:0] if isinstance(x, dict): return pd.DataFrame( {c: _empty_series(c, d, index=index) for (c, d) in x.items()}, index=index ) if isinstance(x, tuple) and len(x) == 2: return _empty_series(x[0], x[1], index=index) elif isinstance(x, (list, tuple)): if not all(isinstance(i, tuple) and len(i) == 2 for i in x): raise ValueError( "Expected iterable of tuples of (name, dtype), " f"got {x}" ) return pd.DataFrame( {c: _empty_series(c, d, index=index) for (c, d) in x}, columns=[c for c, d in x], index=index, ) elif not hasattr(x, "dtype") and x is not None: # could be a string, a dtype object, or a python type. Skip `None`, # because it is implictly converted to `dtype('f8')`, which we don't # want here. try: dtype = np.dtype(x) return _scalar_from_dtype(dtype) except: # noqa # Continue on to next check pass if is_scalar(x): return _nonempty_scalar(x) raise TypeError(f"Don't know how to create metadata from {x}") def get_meta( columns, dtype=None, index_columns=None, index_names=None, default_dtype=np.object_ ): # pragma: no cover """ Extracted and modified from pandas/io/parsers.py : _get_empty_meta (BSD licensed). """ columns = list(columns) # Convert `dtype` to a defaultdict of some kind. # This will enable us to write `dtype[col_name]` # without worrying about KeyError issues later on. if not isinstance(dtype, dict): # if dtype == None, default will be default_dtype. dtype = defaultdict(lambda: dtype or default_dtype) else: # Save a copy of the dictionary. _dtype = dtype.copy() dtype = defaultdict(lambda: default_dtype) # Convert column indexes to column names. for k, v in _dtype.items(): col = columns[k] if is_integer(k) else k dtype[col] = v if index_columns is None or index_columns is False: index = pd.Index([]) else: data = [pd.Series([], dtype=dtype[name]) for name in index_names] if len(data) == 1: index = pd.Index(data[0], name=index_names[0]) else: index = pd.MultiIndex.from_arrays(data, names=index_names) index_columns.sort() for i, n in enumerate(index_columns): columns.pop(n - i) col_dict = {col_name: pd.Series([], dtype=dtype[col_name]) for col_name in columns} return pd.DataFrame(col_dict, columns=columns, index=index) def check_bins(bins: pd.DataFrame, chromsizes: pd.Series) -> pd.DataFrame: is_cat = isinstance(bins["chrom"].dtype, pd.CategoricalDtype) bins = bins.copy() if not is_cat: bins["chrom"] = pd.Categorical( bins.chrom, categories=list(chromsizes.index), ordered=True ) else: assert (bins["chrom"].cat.categories == chromsizes.index).all() return bins def balanced_partition( gs: GenomeSegmentation, n_chunk_max: int, file_contigs: list[str], loadings: list[int | float] | None = None ) -> list[GenomicRangeTuple]: # n_bins = len(gs.bins) grouped = gs._bins_grouped chrom_nbins = grouped.size() if loadings is None: loadings = chrom_nbins chrmax = loadings.idxmax() loadings = loadings / loadings.loc[chrmax] const = chrom_nbins.loc[chrmax] / n_chunk_max granges = [] for chrom, group in grouped: if chrom not in file_contigs: continue clen = gs.chromsizes[chrom] step = int(np.ceil(const / loadings.loc[chrom])) anchors = group.start.values[::step] if anchors[-1] != clen: anchors = np.r_[anchors, clen] granges.extend( (chrom, start, end) for start, end in zip(anchors[:-1], anchors[1:]) ) return granges class GenomeSegmentation: def __init__(self, chromsizes: pd.Series, bins: pd.DataFrame): bins = check_bins(bins, chromsizes) self._bins_grouped = bins.groupby("chrom", observed=True, sort=False) nbins_per_chrom = self._bins_grouped.size().values self.chromsizes = chromsizes self.binsize = get_binsize(bins) self.contigs = list(chromsizes.keys()) self.bins = bins self.idmap = pd.Series(index=chromsizes.keys(), data=range(len(chromsizes))) self.chrom_binoffset = np.r_[0, np.cumsum(nbins_per_chrom)] self.chrom_abspos = np.r_[0, np.cumsum(chromsizes.values)] self.start_abspos = ( self.chrom_abspos[bins["chrom"].cat.codes] + bins["start"].values ) def fetch(self, region: GenomicRangeSpecifier) -> pd.DataFrame: chrom, start, end = parse_region(region, self.chromsizes) result = self._bins_grouped.get_group(chrom) if start > 0 or end < self.chromsizes[chrom]: lo = result["end"].values.searchsorted(start, side="right") hi = lo + result["start"].values[lo:].searchsorted(end, side="left") result = result.iloc[lo:hi] return result def buffered( chunks: Iterable[pd.DataFrame], size: int = 10000000 ) -> Iterator[pd.DataFrame]: """ Take an incoming iterator of small data frame chunks and buffer them into an outgoing iterator of larger chunks. Parameters ---------- chunks : iterator of :py:class:`pandas.DataFrame` Each chunk should have the same column names. size : int Minimum length of output chunks. Yields ------ Larger outgoing :py:class:`pandas.DataFrame` chunks made from concatenating the incoming ones. """ buf = [] n = 0 for chunk in chunks: n += len(chunk) buf.append(chunk) if n > size: yield pd.concat(buf, axis=0) buf = [] n = 0 if len(buf): yield pd.concat(buf, axis=0)