plot-io 0.12.0__py3-none-win_amd64.whl

plot_io/genome/contacts.py

@@ -0,0 +1,508 @@
+"""Hi-C contact matrices (:class:`HiC`): the upper triangle of a contact matrix, the
+diagonal along the x axis, or (``shape="square"``) the whole matrix of ``region`` (x)
+against ``region2`` (y), both axes on the genome scale.  Series called with
+``lane="over"`` after one draw in its coordinates (:func:`under_hic`): intervals as
+outlined triangles or squares on the diagonal (TADs, :func:`tads`), links as loop marks
+(:mod:`plot_io.genome.links`)."""
+
+import math
+from array import array
+from typing import TYPE_CHECKING, Any, Self, Unpack
+
+from .. import kernels
+from ..color.values import color, with_alpha
+from ..settings import DECLINE, Setting, number, one_of, table
+from ..transform import settings_both as coord_settings
+from ..layers import add_layer
+from ..layers.base import common
+from ..layers.core import Call, Layer, Table
+from ..layers.geoms import GEOMS
+from ..layers.paint import Paint, role
+from ..mapped import color_axis_of, trained
+from ..presets._util import gradient_entry
+from ..validators import length
+from . import readers, sources
+from .lanes import is_lane, value_height
+from .adapters import Matrix
+from .regions import Region, Regions, parse_region, pending, whole
+from .series import GenomeSeries, check_bins, known_sizes, region_setting
+
+if TYPE_CHECKING:
+    from ..types import HiCSettings
+
+__all__ = ["HiC", "under_hic", "tads"]
+
+
+def _max_height(v: Any) -> Any:
+    """``None``, a fraction of the axes width (``"33%"``, kept as ``("%", 0.33)``) or a
+    length (px)."""
+    if v is None:
+        return None
+    if isinstance(v, str) and v.strip().endswith("%"):
+        f = float(v.strip()[:-1]) / 100
+        if not f > 0:
+            raise ValueError(f"expected a positive percentage, got {v!r}")
+        return ("%", f)
+    if isinstance(v, tuple) and len(v) == 2 and v[0] == "%":
+        return v
+    return length(positive=True)(v)
+
+
+def _region2(v: Any) -> Any:
+    return parse_region(v, whole=True)
+
+
+def _oe(rule: Any) -> Any:
+    """A class rule of ``matrix_type="oe"`` (declines for other matrices)."""
+    def r(s: "HiC") -> Any:
+        return rule(s) if s.resolve("matrix_type") == "oe" else DECLINE
+    return r
+
+
+def _oe_limits(s: "HiC") -> Any:
+    """Limits symmetric about 1 on a log scale: ``1 / m`` to ``m``, ``log m`` the 0.99
+    quantile of the values' ``|log v|``."""
+    def compute() -> Any:
+        vals = s.cell_values()
+        logs = array("d", [abs(math.log(v)) for v in vals if v > 0 and math.isfinite(v)])
+        q = kernels.quantiles(logs, [0.99])[0] if logs else math.nan
+        m = math.exp(q) if math.isfinite(q) and q > 0 else 2.0
+        return (1 / m, m)
+    return s.cached(s.data_key(), compute, kind="oe_limits")
+
+
+class HiC(GenomeSeries):
+    """The upper triangle of a contact matrix, the diagonal along the x axis (each cell a
+    diamond), or with ``shape="square"`` the whole matrix of ``region`` (x) against
+    ``region2`` (y, its start at the top).  ``reader`` is a Hi-C reader, a path (opened
+    with gwseq_io), or a square matrix over the region (bins of ``span / n``).
+
+    On the triangle, the matrix is drawn as an image under a 45° transform and clipped to
+    the track.  Readers and paths: see `readers <docs/api.md#readers>`_.  ``bins="0.5px"`` takes the finest of the
+    reader's ``bin_sizes`` whose cells are at least 0.5 px wide; the cells of a size
+    (``bins="5kb"``) sit on the file's grid, those past the region cut by the track's edges,
+    and ``bins=n`` reads n cells over the region exactly.  The color axis' limits default to
+    quantiles of the cells drawn (``color_limits="q:0.01/q:0.99"``), on a log scale (cells of
+    0 take the map's low end; cells without a value take ``nan_color``, transparent),
+    through the ``sequential_map``.  A region without a finite positive value draws no
+    cells, its color axis at its default limits.  With ``matrix_type="oe"`` the defaults are
+    a diverging map about 1 (``center=1``: the ``diverging_map``) on a log scale, the limits
+    ``1/m`` to ``m``, ``log m`` the 0.99 quantile of the cells' ``|log v|``; values given
+    (``center``, ``color_map``, ``color_limits``, ``color_scale``, or a style's) win.
+
+    Series called on the track after the triangle with ``lane="over"`` draw in its
+    coordinates: a ``bed`` draws each interval as an outlined triangle (a TAD: its sides
+    from the interval's ends up to the apex at their distance, cut at the track's depth), in
+    the features' colors and ``line_width``; ``links`` draw a mark at each pair's contact
+    cell, an outlined square one bin wide (at least half an em) at the cell of the anchors'
+    centres::
+
+        h = fig.hic(reader, depth=2_000_000)
+        h.plot.bed("tads.bed", lane="over")
+        h.plot.links(loops, lane="over")
+
+    ``shape="square"`` draws the whole matrix of ``region`` (x) against ``region2`` (y;
+    default ``region``, another chromosome draws a rectangle) as cells, row 0 at the top,
+    the diagonal from the top left.  Both axes take the genome scale.  In a tracks figure
+    the track is ``width × span2 / span`` high; in an ordinary plot the aspect is equal
+    (``data_ratio``).  ``depth``, ``direction``, ``max_height`` and ``fit`` do not apply;
+    ``bins``, ``normalization``, ``matrix_type`` and the color settings do.  Over a square, a
+    ``bed`` draws each interval as an outlined square on the diagonal (none when
+    ``region2`` is on another chromosome) and ``links`` a mark at each pair's cell (first
+    anchor on x, second on y) and at its mirror.  A square cannot span several regions side
+    by side::
+
+        fig = plt.Figure()
+        h = fig.plot().hic("cells.hic", region="chr2:30M-70M", shape="square", colorbar=True)
+        h.plot.bed("tads.bed", region="chr2:30M-70M", lane="over")
+
+    Example::
+
+        rng = random.Random(1)
+        n = 48
+        fig = plt.Tracks(("chr1", 0, 2_400_000), width=300)
+        fig.hic([[100 / (1 + abs(i - j)) * (3 if i // 16 == j // 16 else 1) * rng.uniform(0.7, 1.3)
+                  for j in range(n)] for i in range(n)])
+    """
+
+    SETTINGS = table(
+        region=region_setting(),
+        shape=Setting(one_of("triangle", "square"),
+                      doc="`triangle`: the upper triangle, the diagonal along x; `square`: "
+                          "the matrix of `region` (x) against `region2` (y, from the top), "
+                          "the diagonal from the top left, in a track as high as "
+                          "`width × span2 / span`, elsewhere at equal bp per px on both "
+                          "axes (`depth`, `direction`, `max_height` and `fit` do not "
+                          "apply)."),
+        region2=Setting(_region2,
+                        doc="The region of the y axis of `shape=\"square\"`, in the forms "
+                            "of `region`; unset: `region`."),
+        depth=Setting(number(positive=True, integer=True), doc="Largest distance shown (bp); default the region's span "
+                                  "(the whole triangle)."),
+        bins=Setting(check_bins,
+                     doc="A bin count over the region (exactly, from the coarsest "
+                         "resolution giving that many); one of the reader's `bin_sizes` "
+                         "(`\"10kb\"`), read on its grid (the region grown to it, then cut "
+                         "by the track); or a cell width on the axes (`\"0.5px\"`): the "
+                         "finest resolution whose cells are at least that wide."),
+        normalization=Setting(str, doc="`\"none\"` or one of the reader's "
+                                       "`normalizations`."),
+        matrix_type=Setting(one_of("observed", "oe", "expected"),
+                            doc="`observed`, `oe` (observed / expected) or `expected`."),
+        **common("color_map"),
+        nan_color=Setting(color, doc="Color of cells without a value; default: the map's "
+                                     "`bad` (transparent)."),
+        direction=Setting(one_of("up", "down"),
+                          doc="`up`: the apex up; `down`: the triangle hangs from the top "
+                              "(a triangle)."),
+        max_height=Setting(_max_height,
+                           doc="Largest height of the track: a percentage of its width "
+                               "(`\"33%\"`) or a length; `None`: the triangle's (square "
+                               "cells)."),
+        fit=Setting(one_of("clip", "squash"),
+                    doc="How the triangle fits a height below its own (`max_height`, or the "
+                        "track's `height`): `clip` keeps square cells and cuts the "
+                        "distances above; `squash` flattens the whole triangle."),
+        **common("name", "opacity", "z", "rasterize", "simplify", "font_size", "font_family",
+                 "text_color"),
+        **coord_settings(),
+    )
+    #: ``matrix_type="oe"``: a diverging map about 1 on a log scale
+    RULES = {"center": _oe(lambda s: 1), "color_scale": _oe(lambda s: "log"),
+             "color_limits": _oe(_oe_limits)}
+
+    def __init__(self, reader: Any, plot: Any = None, **settings: Any) -> None:
+        self.reader = readers.adapt(reader, "contacts", "HiC", owner=self)
+        super().__init__(plot=plot, **settings)
+
+    def set(self, **settings: "Unpack[HiCSettings]") -> Self:  # type: ignore[override]
+        super().set(**settings)
+        if getattr(self, "plot", None) is not None:
+            self._attach_bar()
+        return self
+
+    def _bar_source(self) -> Any:
+        return self
+
+    # -- data ------------------------------------------------------------------------
+
+    def square(self) -> bool:
+        """Whether the whole matrix is drawn (``shape="square"``)."""
+        return self.resolve("shape") == "square"
+
+    def segmented(self) -> bool:
+        return not self.square()
+
+    def region(self) -> Region:
+        r = super().region()
+        if isinstance(r, Regions) and self.square():
+            raise ValueError(f"{self._what()}: shape=\"square\" needs one region on the x "
+                             f"axis, got {r}")
+        return r
+
+    def region2(self) -> Region:
+        """The region of the y axis of a square (``region2``, else ``region``)."""
+        r = self.resolve("region2")
+        if r is None:
+            return self.region()
+        if pending(r):
+            plot = self.plot
+            r = whole(r, lambda: known_sizes(plot if plot is not None else self,
+                                             None if plot is None else plot.x_axis))
+        if isinstance(r, Regions):
+            raise ValueError(f"{self._what()}: region2 is one region, got {r}")
+        return r
+
+    def equal_aspect(self) -> bool:
+        """Whether the plot's ``data_ratio`` is 1 (a square outside a tracks figure)."""
+        return self.square() and self.plot is not None and "height" not in self.plot.SETTINGS
+
+    def full_depth(self) -> float:
+        """The largest distance asked (``depth``, at most the span)."""
+        span = self.region().span
+        d = self.resolve("depth")
+        return float(span if d is None else min(d, span))
+
+    def _square_height(self, width: float) -> float:
+        return width * self.full_depth() / (2 * self.region().span)
+
+    def _height_cap(self, width: float) -> float | None:
+        """The height (px) the triangle must fit: ``max_height`` or its lane in the track's
+        ``height``, the lower."""
+        caps = []
+        m = self.resolve("max_height")
+        if m is not None:
+            caps.append(m[1] * width if isinstance(m, tuple) else float(m))
+        h = self.plot.get("height") if "height" in self.plot.SETTINGS else "auto"
+        if h != "auto" and not isinstance(h, tuple):          # not a share ("fr")
+            caps.append(value_height(self.plot, float(self.plot.resolve("height"))))
+        return min(caps) if caps else None
+
+    def depth(self) -> float:
+        """The largest distance drawn: ``full_depth``, less what ``fit="clip"`` cuts."""
+        d = self.full_depth()
+        if self.resolve("fit") != "clip":
+            return d
+        width = self.axes_px()
+        cap = self._height_cap(width)
+        natural = self._square_height(width)
+        if cap is None or cap >= natural:
+            return d
+        return d * cap / natural
+
+    def binning(self) -> tuple[int | None, int | None]:
+        """``(bin count, bin size)``: one of them."""
+        b = self.resolve("bins")
+        if b.kind == "count":
+            return int(b.value), None
+        if b.kind == "bp":
+            return None, int(b.value)
+        return None, sources.bin_size_for(self.reader, self.region().span,
+                                          self.axes_px() / b.value)
+
+    def data(self) -> tuple[int, array, float, float]:
+        """``(n, values, start, bin)``: an n × n matrix (rows first) whose bins start at
+        ``start`` (bp) and are ``bin`` bp wide."""
+        region = self.region()
+        return self.cached(self.data_key(), lambda: self._guard(lambda: self._read(region)),
+                           kind="data")
+
+    def data_key(self) -> Any:
+        """What the matrix read depends on."""
+        region = self.region()
+        if self.square():
+            region2 = self.region2()
+            if isinstance(self.reader, Matrix):
+                return (region, region2, "square")
+            return (region, region2, self.binning(), self.resolve("normalization"),
+                    self.resolve("matrix_type"), "square")
+        if isinstance(self.reader, Matrix):
+            return (region,)
+        return (region, self.binning(), self.resolve("normalization"),
+                self.resolve("matrix_type"), self.depth())
+
+    def square_data(self) -> tuple[int, int, array, float, float, float, float]:
+        """``(h, w, values, start, bin, start2, bin2)``: the h × w matrix (rows first) of
+        ``region`` (rows: x) against ``region2`` (columns: y), rows from ``start`` in bins
+        of ``bin`` bp, columns from ``start2`` in bins of ``bin2``."""
+        return self.cached(self.data_key(), lambda: self._guard(self._read_square),
+                           kind="data")
+
+    def _read_square(self) -> tuple[int, int, array, float, float, float, float]:
+        region, region2 = self.region(), self.region2()
+        if isinstance(self.reader, Matrix):     # over the region, whatever it is
+            if region2 != region:
+                raise ValueError("a plain matrix has no coordinates: region2 needs a reader")
+            n, b = self.reader.n, region.span / self.reader.n
+            return n, n, self.reader.values, float(region.start), b, float(region.start), b
+        count, size = self.binning()
+        counts = None if count is None else \
+            (count, max(1, round(count * region2.span / region.span)))
+        m, start, start2 = sources.read_square(self.reader, region, region2, size,
+                                               self.resolve("normalization"),
+                                               self.resolve("matrix_type"), counts)
+        h, w, _, flat = sources.matrix(m)
+        if counts is not None:
+            return h, w, flat, start, region.span / counts[0], start2, \
+                region2.span / counts[1]
+        return h, w, flat, start, float(size or 0), start2, float(size or 0)
+
+    def _read(self, region: Region) -> tuple[int, array, float, float]:
+        if isinstance(self.reader, Matrix):         # over the region, whatever it is
+            n, flat = self.reader.n, self.reader.values
+            return n, flat, float(region.start), region.span / n
+        count, size = self.binning()
+        b = float(size) if size is not None else region.span / (count or 1)
+        d = self.depth()
+        m, start = sources.read_contacts(self.reader, region, size,
+                                            self.resolve("normalization"),
+                                            self.resolve("matrix_type"),
+                                            None if d >= region.span else int(d + b),
+                                            bin_count=count)
+        h, w, _, flat = sources.matrix(m)
+        if h != w:
+            raise ValueError(f"the reader gave a {h} × {w} matrix for one region")
+        return h, flat, float(start), b
+
+    def cell_values(self) -> array:
+        """The values of the cells drawn: the upper triangle within ``depth``, or the whole
+        matrix of a square."""
+        if self.square():
+            return self.square_data()[2]
+        n, flat, _, b = self.data()
+        reach = int(self.depth() / b) + 1
+        out = array("d")
+        for i in range(n):
+            out.extend(flat[i * n + i:i * n + min(n, i + reach + 1)])
+        return out
+
+    def band(self) -> array:
+        """The matrix with NaN outside the cells that can show: the lower triangle and the
+        cells past ``depth`` (they fall outside the axes)."""
+        n, flat, _, b = self.data()
+        reach = int(self.depth() / b) + 2
+        out = array("d")
+        nan = array("d", [math.nan])
+        for i in range(n):
+            end = min(n, i + reach + 1)
+            out.extend(nan * i)
+            out.extend(flat[i * n + i:i * n + end])
+            out.extend(nan * (n - end))
+        return out
+
+    def track_height(self, width: float) -> float | None:
+        if self.square():
+            return width * self.region2().span / self.region().span
+        natural = self._square_height(width)
+        cap = self._height_cap(width)
+        return natural if cap is None else min(natural, cap)
+
+    # -- the preset (stat contact_matrix) ---------------------------------------------
+
+    def _cells(self, call: Call, values: array, grid: Any) -> None:
+        """The cells layer of ``call``: ``grid(ctx)`` gives the values drawn (NaN: none),
+        the transform placing the grid and its columns and rows."""
+        axis = color_axis_of(self)
+        cmap = axis.color_map()
+        if kernels.extent(values, "log")[0]:      # no finite positive value: no cells
+            call.layers.append(Layer(GEOMS["cells"], Table(),
+                                     paint=Paint(self, "contacts.cells", grid=grid, cmap=cmap),
+                                     stat="contact_matrix", occupy="all", key=False))
+        name = self.resolve("name")
+        if name is not None:
+            call.entries.append(gradient_entry(self, name, axis.color_map))
+        call.color_sources = [(axis, values)]
+
+    def layers(self) -> Call:
+        call = self.genome_call()
+        if self.square():
+            return self._square_layers(call)
+        n, flat, start, b = self.data()
+
+        def grid(ctx: Any) -> tuple[array, Any, int, int]:
+            x0 = ctx.xmap(start)
+            xb = ctx.xmap(start + b) - x0              # px per bin along x
+            y0 = ctx.ymap(0.0)
+            yb = ctx.ymap(b) - y0                      # px per bin of distance
+            # cell (row v, col u) → x = start + b (u + v) / 2 + …, distance b (u − v)
+            return self.band(), (xb / 2, yb, xb / 2, -yb, x0, y0), n, n
+        self._cells(call, self.cell_values(), grid)
+        r = self.region()
+        call.extra = {"x": array("d", [r.start, r.end]),
+                      "y": array("d", [0.0, self.depth()])}
+        call.hide = {"y"}                         # distances: no y axis
+        if self.resolve("direction") == "down":
+            call.invert = {"y"}
+            call.hide.add("x")                    # hanging from the top
+        call.exact = {"x", "y"}
+        return call
+
+    def _square_layers(self, call: Call) -> Call:
+        h, w, flat, start, b, start2, b2 = self.square_data()
+
+        def grid(ctx: Any) -> tuple[array, Any, int, int]:
+            x0, y0 = ctx.xmap(start), ctx.ymap(start2)
+            # cell (row v, column u): x from start + b v, y from start2 + b2 u
+            return flat, (0.0, ctx.ymap(start2 + b2) - y0, ctx.xmap(start + b) - x0, 0.0,
+                          x0, y0), w, h
+        self._cells(call, flat, grid)
+        r, r2 = self.region(), self.region2()
+        call.scales["y"] = "genome"
+        call.regions["y"] = r2
+        call.extra = {"x": array("d", [r.start, r.end]), "y": array("d", [r2.start, r2.end])}
+        call.invert = {"y"}                       # row 0 at the top
+        call.exact = {"x", "y"}
+        return call
+
+
+def under_hic(series: Any) -> HiC | None:
+    """The Hi-C triangle a series with ``lane="over"`` is drawn over: the last series
+    with a y scale called before it, when no lane series of its own came between them;
+    ``None`` otherwise."""
+    if series.plot is None or series.resolve("lane") != "over":
+        return None
+    under = None
+    for s in series.plot.series:
+        if s is series:
+            break
+        if not is_lane(s):
+            under = s
+        elif s.resolve("lane") != "over":
+            under = None
+    return under if isinstance(under, HiC) else None
+
+
+def tads(bed: Any, hic: HiC) -> Call:
+    """The call of a ``Bed`` over ``hic``: each interval an outlined triangle (a TAD) in
+    the triangle's coordinates, its apex at the interval's span, cut at ``hic``'s depth;
+    over a square, an outlined square on the diagonal."""
+    call = bed.genome_call()
+    cols = bed.entries()
+
+    def squares(ctx: Any) -> list[Any]:
+        if hic.region2().chr != bed.region().chr:
+            return []                                   # no diagonal in view
+        cs = bed._colors(ctx, cols)
+        opacity = bed.resolve("opacity")
+        opacity = 1.0 if opacity is None else opacity
+        lw = bed.resolve("line_width")
+        return [(array("d", [ctx.xmap(x) for x in (s, e, e, s, s)]),
+                 array("d", [ctx.ymap(y) for y in (s, s, e, e, s)]),
+                 with_alpha(cs[i], opacity), lw)
+                for i, (s, e) in enumerate(zip(cols["start"], cols["end"]))]
+
+    def lines(ctx: Any) -> list[Any]:
+        if hic.square():
+            return squares(ctx)
+        depth = hic.depth()
+        cs = bed._colors(ctx, cols)
+        opacity = bed.resolve("opacity")
+        opacity = 1.0 if opacity is None else opacity
+        lw = bed.resolve("line_width")
+        out = []
+        for i, (s, e) in enumerate(zip(cols["start"], cols["end"])):
+            h = min(e - s, depth)                       # the sides cut at the depth
+            xs = [s, s + h / 2, math.nan, e - h / 2, e] if h < e - s else [s, (s + e) / 2, e]
+            ys = [0.0, h, math.nan, h, 0.0] if h < e - s else [0.0, h, 0.0]
+            out.append((array("d", [ctx.xmap(x) for x in xs]),     # NaN: a break
+                        array("d", [ctx.ymap(y) for y in ys]),
+                        with_alpha(cs[i], opacity), lw))
+        return out
+
+    t = Table()
+    t.lazy["page_lines"] = lines
+    call.layers.append(Layer(GEOMS["path"], t, paint={}, stat="intervals", occupy=False,
+                             key=False))
+    if bed.resolve("color") == "score":
+        scores = [v for v in cols.get("score", []) if v is not None]
+        call.color_sources = [(color_axis_of(bed), array("d", scores))]
+    r = bed.region()
+    call.extra = {"x": array("d", [r.start, r.end])}
+    call.hide = {"y"}
+    call.exact = {"x"}
+    return call
+
+
+add_layer("hic", HiC)
+
+
+@role("contacts.cells")
+def _cells_paint(s: Any, ctx: Any, glyph: bool, grid: Any = None,
+                 cmap: Any = None) -> dict[str, Any]:
+    """The cells of ``grid(ctx)`` (values, the transform placing them, columns, rows)."""
+    st = trained(color_axis_of(s))
+    assert st.scale is not None
+    flo, fhi = st.flo, st.fhi
+    fn = st.scale.function
+    f, transform, cols, rows = grid(ctx)
+    if st.scale.base is not None:
+        # no contacts is the low end of the map, not a missing value
+        f = kernels.scaled(f, low=5e-324)
+    vals = kernels.scaled(f, 1.0, 0.0, fn, low=flo) if fn else f
+    fill, table, _ = cmap.cells(vals, flo, fhi, "image")
+    nan = ctx.color(s.resolve("nan_color"))
+    opacity = s.resolve("opacity")
+    return {"transform": transform, "x_edges": range(cols + 1),
+            "y_edges": range(rows + 1), "fill": fill, "palette": table,
+            "bad": nan if nan is not None else cmap.bad, "mode": "image",
+            "opacity": 1.0 if opacity is None else opacity}