plot-io 0.12.0__py3-none-win_amd64.whl

plot_io/genome/series.py

@@ -0,0 +1,344 @@
+"""The base of the genome series: the region read and the axes width the data are binned
+and packed for, and the ``bins`` setting.
+
+Each series reads its data when the figure is drawn, for the region drawn (the series'
+``region``, else its x axis'), from a reader (see :mod:`plot_io.genome.sources`) or from
+plain data.  x values are positions in bp; they ask their x axis for the genome scale
+(:mod:`plot_io.genome.scale`), which draws them less its ``origin``.
+
+On an x axis over several regions (:class:`~plot_io.genome.scale.Segments`), a series
+reads and lays out each region as if it were drawn alone, on axes as wide as the region's
+segment: ``layers`` is made once per segment (the series' ``region()`` and ``axes_px()``
+those of the segment) and the calls are merged into one, each layer drawn through its
+segment's x map and cut at its edges (:class:`SegmentScope`).  ``track_height`` and
+``lane_rows`` are the largest of the segments'.  A series that places its data across
+the regions itself returns ``False`` from ``segmented()``.
+"""
+
+import contextlib
+import functools
+import math
+from array import array
+from collections.abc import Generator, Mapping
+from typing import Any, NamedTuple
+
+from vec_io import Path
+
+from ..settings import Setting, number
+from ..layers.base import Series
+from ..layers.core import Call
+from .regions import Region, parse_bp, parse_region, pending, whole
+
+__all__ = ["GenomeSeries", "Bins", "check_bins", "axes_width", "region_setting",
+           "SegmentScope", "in_segment", "axis_region", "known_sizes"]
+
+
+class Bins(NamedTuple):
+    """A ``bins`` value: ``kind`` ``"count"`` (``value`` bins over the region), ``"bp"``
+    (bins of ``value`` bp) or ``"px"`` (bins ``value`` px wide on the axes)."""
+
+    kind: str
+    value: float
+
+    def __repr__(self) -> str:
+        return f"Bins({self.kind}={self.value:g})"
+
+
+def check_bins(v: Any) -> Bins:
+    """A bin count (``200``), a bin size in bp (``"10kb"``, ``"5000bp"``, ``"5000d"``) or a
+    bin width on the axes (``"0.5px"``)."""
+    if isinstance(v, Bins):
+        return v
+    if isinstance(v, str):
+        t = v.strip()
+        if t.lower().endswith("px"):
+            try:
+                w = float(t[:-2])
+            except ValueError:
+                w = math.nan
+            if not (math.isfinite(w) and w > 0):
+                raise ValueError(f"expected a positive bin width such as \"0.5px\", got {v!r}")
+            return Bins("px", w)
+        size = t[:-1] if t.endswith("d") and not t.lower().endswith("bp") else t
+        return Bins("bp", float(parse_bp(size, "bin size", positive=True)))
+    return Bins("count", float(number(positive=True, integer=True)(v)))
+
+
+def _region(v: Any) -> Any:
+    return parse_region(v, whole=True)
+
+
+def region_setting() -> Setting:
+    return Setting(_region,
+                   doc="The region read: `\"chr2:143,120,001-143,780,000\"` (1-based, "
+                       "inclusive), `(chr, start, end)` (0-based, half-open), a mapping, "
+                       "or a chromosome name (the whole chromosome, as the x axis' "
+                       "`chr_sizes` gives it); unset: the x axis' `region`.")
+
+
+def known_sizes(owner: Any, axis: Any = None) -> dict[str, int]:
+    """The chromosome sizes known to the figure of ``owner`` (a plot or a figure): the
+    largest each reader of its genome series gives (``chr_sizes``), then ``axis``'
+    ``chr_sizes``, which win."""
+    fig = owner if hasattr(owner, "plots") else getattr(owner, "figure", None)
+    plots = fig.plots if fig is not None else [owner]
+    sizes: dict[str, int] = {}
+    for p in plots:
+        for s in getattr(p, "series", ()):
+            got = getattr(getattr(s, "reader", None), "chr_sizes", None)
+            if isinstance(got, Mapping):
+                for k, n in got.items():
+                    sizes[str(k)] = max(sizes.get(str(k), 0), int(n))
+    given = axis.resolve("chr_sizes") if axis is not None else None
+    if given:
+        sizes.update(given)
+    return sizes
+
+
+def axis_region(axis: Any) -> Any:
+    """The ``region`` of a genome axis, a whole chromosome (``"chr1"``) made the region
+    of its size (:func:`known_sizes`)."""
+    r = axis.resolve("region")
+    if not pending(r):
+        return r
+    return whole(r, lambda: known_sizes(axis._owner, axis))
+
+
+def axes_width(plot: Any) -> float:
+    """The axes width (px) genome data are binned and packed for: in a render, the width
+    of the plot's axes as laid out when its calls were made (the layout makes them again
+    when it places the axes at another width); outside a render, the figure's ``width``."""
+    from ..spec import NeedsBox, current
+    main = getattr(plot, "twin_of", None)
+    plot = main[0] if main else plot
+    fig = getattr(plot, "figure", None)
+    r = current()
+    if r is not None and fig is not None and r.root is fig:
+        w = r.widths.get(id(plot))
+        if w is None:
+            boxes = r.boxes.get(id(plot))
+            if not boxes:
+                raise NeedsBox("the axes width genome data are binned for is known once the "
+                               "layout sizes the axes")
+            w = r.widths[id(plot)] = boxes["axes_width"]
+        return w
+    if fig is None or "width" not in fig.SETTINGS:
+        return 480.0
+    return float(fig.resolve("width"))
+
+
+@contextlib.contextmanager
+def in_segment(series: Any, seg: Any) -> Generator[None]:
+    """``series``, and the other genome series of its plot (a series may read another's
+    data), reading and drawing the region of segment ``seg``."""
+    plot = getattr(series, "plot", None)
+    every = [s for s in (plot.series if plot is not None else ()) if
+             isinstance(s, GenomeSeries)] or [series]
+    if series not in every:
+        every.append(series)
+    before = [s.__dict__.get("_segment") for s in every]
+    for s in every:
+        s.__dict__["_segment"] = seg
+    try:
+        yield
+    finally:
+        for s, b in zip(every, before):
+            s.__dict__["_segment"] = b
+
+
+class SegmentScope:
+    """A layer of a series drawn in a segment of an axis over several regions:
+    :meth:`scope` gives the drawing context whose x map draws the segment's positions, the
+    series reading the segment's region, and the canvas cut at the segment's edges."""
+
+    __slots__ = ("series", "segment")
+
+    def __init__(self, series: Any, segment: Any) -> None:
+        self.series, self.segment = series, segment
+
+    def context(self, ctx: Any) -> Any:
+        """``ctx`` (a call's context) with the segment's x map; one per plot context and
+        segment (its own ``root``: values a layer keeps per context are per segment)."""
+        seg = self.segment
+        made = ctx.root.__dict__.setdefault("_segments", {})
+        key = (id(ctx), seg.number)
+        got = made.get(key)
+        if got is None:
+            import copy
+            c = copy.copy(ctx)
+            for k in ("_segments", "_lazy", "_seams"):
+                c.__dict__.pop(k, None)
+            c.xmap = ctx.xmap.affine(seg.region.start - seg.u0 / seg.k, seg.k, "x")
+            c.root = c
+            got = made[key] = (ctx, c)          # ctx kept alive while its id is a key
+        return got[1]
+
+    @contextlib.contextmanager
+    def scope(self, ctx: Any, clip: bool = False) -> Generator[Any]:
+        seg = self.segment
+        c = self.context(ctx)
+        with in_segment(self.series, seg), contextlib.ExitStack() as stack:
+            if clip:
+                canvas = ctx.canvas
+                stack.enter_context(canvas.state())
+                a, b = sorted((ctx.root.xmap(seg.u0), ctx.root.xmap(seg.u1)))
+                y0, y1 = ctx.rect[1], ctx.rect[3]
+                canvas.clip(Path().rect(a, y0 - 1e5, b - a, y1 - y0 + 2e5))
+            yield c
+
+
+def _merged(calls: list[Any], series: Any, segs: Any) -> Any:
+    """One call of the calls a series made for each segment: their layers, tables and
+    color and size sources, the extents of all (legend entries, colors taken and the rest of
+    the first)."""
+    out = calls[0]
+    for c, seg in zip(calls, segs):
+        scope = SegmentScope(series, seg)
+        for la in c.layers:
+            la.options["segment"] = scope
+        for t in c.tables:
+            t.segment = seg
+    for c in calls[1:]:
+        out.layers += c.layers
+        out.tables += c.tables
+        out.color_sources += c.color_sources
+        out.size_sources += c.size_sources
+        out.after += c.after
+        for role, v in c.extra.items():
+            have = out.extra.get(role)
+            out.extra[role] = v if have is None else array("d", list(have) + list(v))
+        for role, v in c.sticky_values.items():
+            out.sticky_values[role] = out.sticky_values.get(role, []) + list(v)
+        for role, v in c.ticks.items():
+            out.ticks[role] = out.ticks.get(role, []) + list(v)
+        for role, f in c.extent_of.items():
+            g = out.extent_of.get(role)
+            out.extent_of[role] = f if g is None else _union(g, f)
+        for name in ("zero", "exact", "invert", "whole", "hide"):
+            getattr(out, name).update(getattr(c, name))
+    return out
+
+
+def _union(f: Any, g: Any) -> Any:
+    def extent(scale: Any) -> tuple[int, float, float]:
+        n1, a1, b1 = f(scale)
+        n2, a2, b2 = g(scale)
+        if not n1:
+            return n2, a2, b2
+        if not n2:
+            return n1, a1, b1
+        return n1 + n2, min(a1, a2), max(b1, b2)
+    return extent
+
+
+def _per_segment(name: str, method: Any) -> Any:
+    """``method`` of a genome series run for each segment of its x axis when it has
+    several: ``layers`` merged (:func:`_merged`), other methods the largest result."""
+    @functools.wraps(method)
+    def run(self: Any, *args: Any) -> Any:
+        if self.__dict__.get("_segment") is not None or self.plot is None \
+                or not self.segmented():
+            return method(self, *args)
+        from .scale import segments_of
+        width = args[0] if name == "track_height" else None
+        segs = segments_of(self.plot, width)
+        if segs is None:
+            return method(self, *args)
+        got = []
+        for seg in segs:
+            with in_segment(self, seg):
+                got.append(method(self, *((seg.px,) if name == "track_height" else args)))
+        if name == "layers":
+            return _merged(got, self, segs)
+        known = [g for g in got if g is not None]
+        return max(known) if known else None
+    return run
+
+
+class GenomeSeries(Series):
+    """Base of the genome series: the region and the width the data are binned for."""
+
+    #: The data argument (a reader or plain data).
+    reader: Any = None
+    #: The methods made per segment (see :func:`_per_segment`).
+    _PER_SEGMENT = ("layers", "track_height", "lane_rows")
+
+    def __init_subclass__(cls, **kw: Any) -> None:
+        super().__init_subclass__(**kw)
+        for name in GenomeSeries._PER_SEGMENT:
+            m = cls.__dict__.get(name)
+            if m is not None and not getattr(m, "_per_segment", False):
+                w = _per_segment(name, m)
+                w._per_segment = True
+                setattr(cls, name, w)
+
+    def segmented(self) -> bool:
+        """Whether the series is made once per region of an axis over several (see the
+        module); ``False``: it places its data across the regions itself."""
+        return True
+
+    def segment(self) -> Any:
+        """The segment being read or drawn (:class:`~plot_io.genome.scale.Segment`), or
+        ``None``."""
+        return self.__dict__.get("_segment")
+
+    def cached(self, key: Any, compute: Any, reads: tuple[str, ...] = (),
+               kind: Any = None) -> Any:
+        seg = self.__dict__.get("_segment")
+        if seg is not None and kind is not None:      # one result per kind and segment
+            kind = (kind, seg.number)
+        return super().cached(key, compute, reads, kind)
+
+    def _given_region(self) -> Region | None:
+        seg = self.__dict__.get("_segment")
+        if seg is not None:
+            return seg.region
+        r = self.resolve("region")
+        if r is None and self.plot is not None:
+            return axis_region(self.plot.x_axis)
+        if pending(r):
+            plot = self.plot
+            return whole(r, lambda: known_sizes(plot if plot is not None else self,
+                                                None if plot is None else plot.x_axis))
+        return r
+
+    def region(self) -> Region:
+        r = self._given_region()
+        if r is None:
+            raise ValueError(f"{self!r} has no region: give region= to the x axis "
+                             f"(plt.Tracks(region), figure.x_axis) or the series")
+        return r
+
+    def axes_px(self) -> float:
+        """The width (px) the data are binned for (see :func:`axes_width`; in a segment, the
+        segment's)."""
+        seg = self.__dict__.get("_segment")
+        return axes_width(self.plot) if seg is None else seg.px
+
+    def genome_call(self) -> Call:
+        """A call asking its x axis for the genome scale, over the series' region (``None``:
+        neither the series nor its x axis has one)."""
+        call = Call(self)
+        call.scales["x"] = "genome"
+        call.region = self._given_region()
+        return call
+
+    def _what(self) -> str:
+        label = self.plot.y_axis.resolve("label") if self.plot is not None else None
+        name = self.resolve("name") if "name" in self.SETTINGS else None
+        tag = label or name
+        return f"{self.kind}" + (f" {str(tag)!r}" if tag else "")
+
+    def _guard(self, read: Any) -> Any:
+        """``read()``, an error naming the track."""
+        try:
+            return read()
+        except Exception as e:
+            if self._what() in str(e):
+                raise
+            raise type(e)(f"{self._what()}: {e}") from e
+
+    def track_height(self, width: float) -> float | None:
+        """The height (px) the series asks of its track's value lane when the track's
+        ``height`` is ``"auto"``; ``None``: the track's ``auto_height``."""
+        return None