zipoop 0.1.0__py3-none-any.whl

zipoop/dataset.py

@@ -0,0 +1,119 @@
+"""The merged shard dataset: ``<root>/<chrom>.npz`` with members "<window>/<key>", plus ``MANIFEST.tsv``."""
+from __future__ import annotations
+
+import csv
+import hashlib
+import os
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from .shard import Shard, ShardView, read_shard
+from .teacher import TeacherWindow, load_window
+
+
+def chrom_of(window: str) -> str:
+    return window.rsplit("_", 2)[0]
+
+
+class Dataset:
+    """Opening a chromosome parses its whole zip directory (chr1: about 3.6 s and 0.9 GB), so handles are cached
+    per chromosome; forked DataLoader workers notice the new pid and reopen their own."""
+
+    def __init__(self, root):
+        self.root = Path(root)
+        if not self.root.is_dir():
+            raise FileNotFoundError(self.root)
+        self._lock = threading.Lock()
+        self._pid = os.getpid()
+        self._open = {}  # chrom -> (NpzFile, {window: [member keys]})
+
+    def __getstate__(self):
+        return {"root": self.root}
+
+    def __setstate__(self, state):
+        self.__init__(state["root"])
+
+    def chroms(self) -> list[str]:
+        return sorted(p.stem for p in self.root.glob("*.npz"))
+
+    def _chrom(self, chrom):
+        with self._lock:
+            if self._pid != os.getpid():
+                self._pid, self._open = os.getpid(), {}
+            if chrom not in self._open:
+                path = self.root / f"{chrom}.npz"
+                if not path.exists():
+                    raise FileNotFoundError(path)
+                npz = np.load(path, allow_pickle=False)
+                index = {}
+                for name in npz.files:
+                    window, key = name.split("/", 1)
+                    index.setdefault(window, []).append(key)
+                self._open[chrom] = (npz, index)
+            return self._open[chrom]
+
+    def windows(self, chrom: str) -> list[str]:
+        """Window names of one chromosome in genome order."""
+        return list(self._chrom(chrom)[1])
+
+    def read(self, window: str) -> Shard:
+        npz, index = self._chrom(chrom_of(window))
+        if window not in index:
+            raise KeyError(window)
+        return read_shard(ShardView(npz, window, index[window]))
+
+    def load_window(self, window: str, device="cpu", zero=None) -> TeacherWindow:
+        return load_window(self.read(window), device, zero)
+
+    def load_region(self, windows, device="cpu", zero=None) -> TeacherWindow:
+        """Adjacent windows, in the given order, concatenated along the position axis."""
+        ws = [self.load_window(w, device, zero) for w in windows]
+        return TeacherWindow("+".join(w.key for w in ws), *(torch.cat([getattr(w, f) for w in ws])
+                                                            for f in ("dna", "signal", "valid")))
+
+    def close(self):
+        with self._lock:
+            for npz, _ in self._open.values():
+                npz.close()
+            self._open.clear()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.close()
+
+
+def _check_file(root: Path, row: dict) -> str | None:
+    path = root / f"{row['chrom']}.npz"
+    if not path.exists():
+        return "missing"
+    if path.stat().st_size != int(row["bytes"]):
+        return f"size {path.stat().st_size} != {row['bytes']}"
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for block in iter(lambda: f.read(1 << 24), b""):
+            h.update(block)
+    return None if h.hexdigest() == row["sha256"] else "sha256 mismatch"
+
+
+def verify(root, chroms=None, progress=None, workers=4) -> list[tuple[str, str]]:
+    """Check every chromosome file against MANIFEST.tsv (size, then sha256). Returns [(chrom, problem)]."""
+    root = Path(root)
+    manifest = root / "MANIFEST.tsv"
+    if not manifest.exists():
+        raise FileNotFoundError(manifest)
+    rows = [r for r in csv.DictReader(open(manifest, encoding="utf-8"), delimiter="\t")
+            if not chroms or r["chrom"] in chroms]
+    problems = []
+    with ThreadPoolExecutor(max_workers=max(1, workers)) as pool:
+        for row, problem in zip(rows, pool.map(lambda r: _check_file(root, r), rows)):
+            if problem:
+                problems.append((row["chrom"], problem))
+            if progress:
+                progress(row["chrom"], problem)
+    return problems

zipoop/decoder.py

@@ -0,0 +1,62 @@
+"""Teacher decoder: sparse-FFT base signal + DNA-conditioned network residual."""
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+from .network import TeacherNet
+from .shard import Shard
+
+
+def _unscale(x: torch.Tensor, mean: torch.Tensor, power: torch.Tensor) -> torch.Tensor:
+    magnitude = x.abs()
+    magnitude = torch.where(magnitude > 10, (magnitude + 10).square() / 40, magnitude)
+    return x.sign() * magnitude.pow(1 / power) * mean
+
+
+class Decoder:
+    def __init__(self, shard: Shard, device="cpu"):
+        self.shard = shard
+        self.device = torch.device(device)
+        self.shape = shard.shape
+        self.valid = shard.valid
+        self.net = TeacherNet(shard.config, shard.shape[1]).to(self.device)
+        if set(shard.parameters) != set(self.net.state_dict()):
+            raise ValueError("Decoder parameters are incomplete or incompatible")
+        self.net.load_state_dict({k: torch.from_numpy(v).to(self.device) for k, v in shard.parameters.items()},
+                                 strict=True)
+        self.net.eval()
+        t = lambda a: torch.from_numpy(a).float().to(self.device)  # noqa: E731
+        self._scales = tuple(t(getattr(shard, n)) for n in ("mean", "power", "gain", "offset"))
+        self._final = (t(shard.final_gain), t(shard.final_offset))
+
+    @torch.no_grad()
+    def decode(self, start=0, end=None, tracks=None, missing_nan=True) -> torch.Tensor:
+        """(end - start, len(tracks)) float32. Positions the mask marks missing are NaN unless ``missing_nan`` is
+        False (then the raw network output, which training never constrained there, is returned)."""
+        length, count = self.shape
+        end = length if end is None else end
+        if not 0 <= start < end <= length:
+            raise ValueError("Invalid position slice")
+        ids = np.arange(count, dtype=np.int64) if tracks is None else np.asarray(tracks)
+        if ids.ndim != 1 or not np.issubdtype(ids.dtype, np.integer) or np.any((ids < 0) | (ids >= count)):
+            raise ValueError("Invalid track selection")
+        out = self._fft(ids, start, end) + self._residual(ids, start, end)
+        out = out * self._final[0][ids] + self._final[1][ids]
+        if missing_nan:
+            out = out.masked_fill(~torch.from_numpy(self.valid[start:end][:, ids]).to(self.device), float("nan"))
+        return out
+
+    def _fft(self, ids, start, end):
+        s = self.shard
+        spectrum = torch.zeros((len(ids), self.shape[0] // 2 + 1), dtype=torch.complex64, device=self.device)
+        spectrum.scatter_(1, torch.from_numpy(s.fft_indices[ids].astype(np.int64)).to(self.device),
+                          torch.from_numpy(s.fft_coefficients[ids]).to(self.device))
+        return torch.fft.irfft(spectrum, n=self.shape[0], dim=1)[:, start:end].T
+
+    def _residual(self, ids, start, end):
+        mean, power, gain, offset = self._scales
+        idx = torch.as_tensor(ids, device=self.device)
+        dna = torch.from_numpy(self.shard.dna.astype(np.int64)).unsqueeze(0).to(self.device)
+        pred = self.net.head(self.net.features(dna)[0, start:end], idx)
+        return _unscale(pred, mean[idx], power[idx]) * gain[idx] + offset[idx]

zipoop/evaluation.py

@@ -0,0 +1,77 @@
+"""Evaluation convention: per 32 kb window, per track, 1 bp Pearson, then the mean over evaluable
+(window, track) pairs. Never concatenate windows into one Pearson: between-window level differences
+would leak into the correlation, inflating it with region length and making lengths incomparable."""
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+from .panel import panel_tracks, zero_tracks
+
+WINDOW = 32768
+
+
+def per_track_pearson(pred, truth, mask):
+    """Pearson over observed positions per column; second value marks evaluable columns (constant targets are not)."""
+    counts = mask.sum(0).clamp_min(1)
+    a = (pred - (pred * mask).sum(0) / counts) * mask
+    b = (truth - (truth * mask).sum(0) / counts) * mask
+    var_a, var_b = a.square().sum(0), b.square().sum(0)
+    return (a * b).sum(0) / (var_a * var_b).clamp_min(1e-24).sqrt(), (counts >= 2) & (var_b > 1e-12)
+
+
+def calibrate(pred, truth, valid):
+    """Per-track affine map of a Pearson-trained output into truth's units (Pearson unchanged)."""
+    counts = valid.sum(0).clamp_min(1)
+    pred_mean = (pred * valid).sum(0) / counts
+    truth_mean = (truth * valid).sum(0) / counts
+    a, b = (pred - pred_mean) * valid, (truth - truth_mean) * valid
+    variance = a.square().sum(0)
+    gain = ((a * b).sum(0) / variance.clamp_min(1e-12)).clamp_min(0)
+    gain = torch.where(variance > 1e-12, gain, torch.zeros_like(gain))
+    return pred * gain + (truth_mean - gain * pred_mean)
+
+
+def restore_real(signal, valid, zero=None):
+    """Apply the teacher's missing-value rule to real data; tracks on the last axis, numpy or torch."""
+    zero = zero_tracks() if zero is None else np.asarray(zero, dtype=bool)
+    if torch.is_tensor(signal):
+        zero = torch.as_tensor(zero, device=signal.device)
+        return torch.where(valid, signal, torch.zeros_like(signal)), valid | zero
+    return np.where(valid, signal, 0).astype(signal.dtype, copy=False), valid | zero
+
+
+def windowed(x, window=WINDOW):
+    """(L, C) -> (window, (L // window) * C): one column per (window, track), so column Pearson is per-window."""
+    L, C = x.shape
+    if L % window:
+        raise ValueError(f"length {L} is not a multiple of window {window}")
+    return x.reshape(L // window, window, C).permute(1, 0, 2).reshape(window, (L // window) * C)
+
+
+def window_pearson(pred, target, valid, window=WINDOW, chunk=256):
+    """(L // window, T) Pearson, NaN where not evaluable. Inputs are (L, T) with L a multiple of ``window``."""
+    L, T = target.shape
+    cols = []
+    for c in range(0, T, chunk):
+        d = slice(c, min(T, c + chunk))
+        p, ok = per_track_pearson(windowed(pred[:, d].float(), window), windowed(target[:, d].float(), window),
+                                  windowed(valid[:, d], window))
+        cols.append(torch.where(ok, p, torch.full_like(p, float("nan"))).reshape(L // window, -1))
+    return torch.cat(cols, 1)
+
+
+def summarize(p, tracks=None):
+    """Mean/median over evaluable pairs, plus per-modality and per-window means."""
+    tracks = panel_tracks() if tracks is None else tracks
+    p = p.detach().float().cpu().numpy()
+    modality = np.array([t.get("modality", "?") for t in tracks])
+    ok = np.isfinite(p)
+    out = {"mean": float(p[ok].mean()) if ok.any() else None, "median": float(np.median(p[ok])) if ok.any() else None,
+           "pairs": int(ok.sum()), "by_modality": {},
+           "per_window": [float(np.nanmean(r)) if np.isfinite(r).any() else None for r in p]}
+    for m in sorted(set(modality)):
+        sel = ok & (modality == m)[None, :]
+        if sel.any():
+            out["by_modality"][m] = {"mean": float(p[sel].mean()), "pairs": int(sel.sum())}
+    return out

zipoop/masks.py

@@ -0,0 +1,73 @@
+"""Decoder for the validity-mask wire format ``track-rle-lzma-v2``.
+
+A mask is a (length, tracks) boolean field. Tracks are classified as all-absent, all-present or mixed;
+only mixed tracks carry data, as concatenated run lengths (uvarint, first byte = value of the first run)
+inside one LZMA stream.
+"""
+from __future__ import annotations
+
+import lzma
+
+import numpy as np
+
+MASK_FORMAT = "track-rle-lzma-v2"
+_MAX_CELLS = 1_000_000_000
+_LZMA_MEMLIMIT = 128 * 1024 * 1024
+
+
+def _uvarint(data: memoryview, pos: int) -> tuple[int, int]:
+    value = shift = 0
+    while True:
+        if pos >= len(data) or shift >= 64:
+            raise ValueError("Invalid or truncated varint")
+        byte = data[pos]
+        pos += 1
+        value |= (byte & 0x7F) << shift
+        if not byte & 0x80:
+            return value, pos
+        shift += 7
+
+
+def _bit_plane(data: memoryview, pos: int, count: int) -> tuple[np.ndarray, int]:
+    nbytes = (count + 7) // 8
+    bits = np.unpackbits(np.frombuffer(data[pos:pos + nbytes], dtype=np.uint8), bitorder="little", count=count)
+    return bits.astype(bool), pos + nbytes
+
+
+def _run_column(data: memoryview, pos: int, length: int) -> tuple[np.ndarray, int]:
+    value, pos = _uvarint(data, pos)
+    if value not in (0, 1):
+        raise ValueError("Invalid run value")
+    runs, total = [], 0
+    while total < length:
+        run, pos = _uvarint(data, pos)
+        if run <= 0 or run > length - total:
+            raise ValueError("Invalid run length")
+        runs.append(run)
+        total += run
+    values = (np.arange(len(runs)) % 2 == 0) == bool(value)
+    return np.repeat(values, runs), pos
+
+
+def decode_mask(data: bytes, shape: tuple[int, int]) -> np.ndarray:
+    length, tracks = (int(n) for n in shape)
+    if min(length, tracks) <= 0 or length * tracks > _MAX_CELLS:
+        raise ValueError("Invalid mask shape")
+    view = memoryview(data)
+    count, pos = _uvarint(view, 0)
+    if count != tracks:
+        raise ValueError("class vector length mismatch")
+    mixed, pos = _bit_plane(view, pos, count)
+    present, pos = _bit_plane(view, pos, count)
+    mask = np.zeros((length, tracks), dtype=bool)
+    mask[:, ~mixed & present] = True
+    lz = lzma.LZMADecompressor(memlimit=_LZMA_MEMLIMIT)
+    payload = memoryview(lz.decompress(data[pos:], max_length=length * tracks * 2 + 1024))
+    if not lz.eof or lz.unused_data:
+        raise ValueError("Oversized or trailing mask payload")
+    cursor = 0
+    for track in np.flatnonzero(mixed):
+        mask[:, track], cursor = _run_column(payload, cursor, length)
+    if cursor != len(payload):
+        raise ValueError("Trailing mask runs")
+    return mask

zipoop/network.py

@@ -0,0 +1,102 @@
+"""Inference-only teacher network: DNA -> per-base rank features -> per-track head."""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    width: int
+    rank: int
+    transformer_layers: int
+    attention_heads: int
+    downsample_levels: int
+
+    # The only variant present in Zipoop shards; anything else is rejected rather than silently mis-decoded.
+    # Older archives that omit "pooling" used average pooling, so a missing key must not count as "max".
+    _SUPPORTED = {"pooling": "max", "position_encoding": "sinusoidal", "pair_channels": 0, "head_nonlinear_rank": 0}
+    _LEGACY_DEFAULT = {"pooling": "avg", "position_encoding": "sinusoidal", "pair_channels": 0, "head_nonlinear_rank": 0}
+
+    @classmethod
+    def from_json(cls, raw: dict) -> "ModelConfig":
+        for key, expected in cls._SUPPORTED.items():
+            got = raw.get(key, cls._LEGACY_DEFAULT[key])
+            if got != expected:
+                raise ValueError(f"Unsupported model option {key}={got!r} (zipoop decodes {expected!r} only)")
+        cfg = cls(*(int(raw[name]) for name in ("width", "rank", "transformer_layers", "attention_heads",
+                                                "downsample_levels")))
+        if min(cfg.width, cfg.rank, cfg.transformer_layers, cfg.attention_heads) < 1:
+            raise ValueError("Model dimensions must be positive")
+        if cfg.width % cfg.attention_heads or not 1 <= cfg.downsample_levels <= 7:
+            raise ValueError("Width must divide attention heads; levels must be 1..7")
+        return cfg
+
+
+class _Residual(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.net = nn.Sequential(nn.GroupNorm(1, channels), nn.GELU(),
+                                 nn.Conv1d(channels, channels, 5, padding=2),
+                                 nn.GroupNorm(1, channels), nn.GELU(),
+                                 nn.Conv1d(channels, channels, 1))
+
+    def forward(self, x):
+        return x + self.net(x)
+
+
+class _TrackHead(nn.Module):
+    def __init__(self, rank: int, tracks: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(tracks, rank))
+        self.bias = nn.Parameter(torch.empty(tracks))
+        self.scale_logits = nn.Parameter(torch.empty(tracks))
+        self.register_buffer("signed", torch.zeros(tracks, dtype=torch.bool))
+
+    def forward(self, z, tracks):
+        x = F.linear(z, self.weight[tracks], self.bias[tracks])
+        return torch.where(self.signed[tracks], x, F.softplus(x)) * F.softplus(self.scale_logits[tracks])
+
+
+def _sinusoidal(length: int, width: int, like: torch.Tensor) -> torch.Tensor:
+    p = torch.arange(length, device=like.device, dtype=like.dtype)[:, None]
+    freq = torch.exp(torch.arange(0, width, 2, device=like.device, dtype=like.dtype) * (-math.log(10000) / width))
+    table = torch.zeros(length, width, device=like.device, dtype=like.dtype)
+    table[:, 0::2] = torch.sin(p * freq)
+    table[:, 1::2] = torch.cos(p * freq[:table[:, 1::2].shape[1]])
+    return table
+
+
+class TeacherNet(nn.Module):
+    """Parameter names match the stored shards (stem, down, attention, up, trunk_projection, head)."""
+
+    def __init__(self, config: ModelConfig, tracks: int):
+        super().__init__()
+        c = config.width
+        self.config = config
+        self.stem = nn.Conv1d(4, c, 15, padding=7)
+        self.down = nn.ModuleList(_Residual(c) for _ in range(config.downsample_levels))
+        layer = nn.TransformerEncoderLayer(c, config.attention_heads, c * 4, dropout=0, activation="gelu",
+                                           batch_first=True, norm_first=True)
+        self.attention = nn.TransformerEncoder(layer, config.transformer_layers, enable_nested_tensor=False)
+        self.up = nn.ModuleList(_Residual(c) for _ in range(config.downsample_levels))
+        self.trunk_projection = nn.Conv1d(c, config.rank, 1)
+        self.head = _TrackHead(config.rank, tracks)
+
+    def features(self, dna: torch.Tensor) -> torch.Tensor:
+        """(B, L) int bases (0-3 = ACGT, 4 = N -> all-zero input) -> (B, L, rank)."""
+        x = self.stem(F.one_hot(dna.long(), 5)[..., :4].float().transpose(1, 2))
+        skips = []
+        for block in self.down:
+            x = block(x)
+            skips.append(x)
+            x = F.max_pool1d(x, 2)
+        x = x.transpose(1, 2)
+        x = self.attention(x + _sinusoidal(x.shape[1], x.shape[2], x)).transpose(1, 2)
+        for block, skip in zip(self.up, reversed(skips)):
+            x = block(F.interpolate(x, size=skip.shape[-1], mode="nearest") + skip)
+        return self.trunk_projection(x).transpose(1, 2)

zipoop/panel.py

@@ -0,0 +1,26 @@
+"""The 1,885-track ENCODE panel the shards were fitted to (track order = last axis of every signal array)."""
+from __future__ import annotations
+
+import json
+from functools import lru_cache
+from pathlib import Path
+
+import numpy as np
+
+PANEL = Path(__file__).resolve().parent / "data" / "encode-panel_1885.json"
+ZERO_MODALITIES = frozenset({"rna"})
+
+
+@lru_cache(maxsize=None)
+def _default_tracks():
+    return json.load(open(PANEL))["tracks"]
+
+
+def panel_tracks(panel=None):
+    return json.load(open(panel))["tracks"] if panel else _default_tracks()
+
+
+def zero_tracks(tracks=None) -> np.ndarray:
+    """Bool per track: unrecorded means 0 (RNA-seq bigWigs never store zeros)."""
+    tracks = _default_tracks() if tracks is None else tracks
+    return np.array([t.get("modality") in ZERO_MODALITIES for t in tracks], dtype=bool)

zipoop/shard.py

@@ -0,0 +1,105 @@
+"""Parse and validate one teacher shard (a ``pooper-dna-v3`` hybrid archive) into a typed record."""
+from __future__ import annotations
+
+import json
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+from .masks import MASK_FORMAT, decode_mask
+from .network import ModelConfig
+
+SHARD_FORMAT = "pooper-dna-v3"
+HYBRID_FORMAT = "fft-plus-dna-residual-v1"
+
+
+class ShardView:
+    """Read-only mapping over one window inside a merged per-chromosome npz (members "<window>/<key>")."""
+
+    def __init__(self, npz, key: str, members):
+        self._npz, self.key, self._members = npz, key, frozenset(members)
+
+    def __contains__(self, name):
+        return name in self._members
+
+    def __getitem__(self, name):
+        if name not in self._members:
+            raise KeyError(name)
+        return self._npz[f"{self.key}/{name}"]
+
+
+@dataclass
+class Shard:
+    key: str
+    shape: tuple[int, int]          # (length, tracks)
+    metadata: dict
+    config: ModelConfig
+    dna: np.ndarray                 # (L,) uint8
+    valid: np.ndarray               # (L, T) bool
+    parameters: dict[str, np.ndarray]
+    mean: np.ndarray                # per-track target scaling and calibration, (T,)
+    power: np.ndarray
+    gain: np.ndarray
+    offset: np.ndarray
+    fft_indices: np.ndarray         # (T, K) retained rfft bins
+    fft_coefficients: np.ndarray    # (T, K) complex64
+    final_gain: np.ndarray
+    final_offset: np.ndarray
+
+
+@contextmanager
+def _opened(source):
+    if isinstance(source, (str, os.PathLike)):
+        with np.load(source, allow_pickle=False) as z:
+            yield z, Path(source).stem
+    else:
+        yield source, source.key
+
+
+def read_shard(source) -> Shard:
+    """``source``: a shard path, or a ShardView from :class:`zipoop.Dataset`."""
+    with _opened(source) as (z, key):
+        if "format" not in z or str(z["format"]) != SHARD_FORMAT:
+            raise ValueError(f"Unsupported shard format: {z['format'] if 'format' in z else None}")
+        if "hybrid_format" not in z or str(z["hybrid_format"]) != HYBRID_FORMAT:
+            raise ValueError("Unsupported hybrid archive")
+        if "valid_encoding" not in z or str(z["valid_encoding"]) != MASK_FORMAT:
+            raise ValueError("Unsupported mask encoding")
+        shape = tuple(int(n) for n in z["shape"])
+        if len(shape) != 2:
+            raise ValueError("Invalid shape")
+        length, tracks = shape
+        dna = z["dna"].copy()
+        if dna.shape != (length,) or np.any(dna > 4):
+            raise ValueError("Invalid DNA data in archive")
+        names = json.loads(z["parameter_names"].tobytes())
+        if len(names) != len(set(names)):
+            raise ValueError("Duplicate parameter names")
+        shard = Shard(
+            key=key, shape=shape, metadata=json.loads(z["metadata"].tobytes()),
+            config=ModelConfig.from_json(json.loads(z["model_config"].tobytes())),
+            dna=dna, valid=decode_mask(z["valid_v2"].tobytes(), shape),
+            parameters={name: z[f"parameter_{i}"].copy() for i, name in enumerate(names)},
+            **{name: z[name].copy() for name in ("mean", "power", "gain", "offset")},
+            fft_indices=z["fft_indices"].copy(), fft_coefficients=z["fft_coefficients"].copy(),
+            final_gain=z["final_gain"].copy(), final_offset=z["final_offset"].copy())
+    _validate(shard)
+    return shard
+
+
+def _validate(s: Shard) -> None:
+    length, tracks = s.shape
+    if any(a.shape != (tracks,) for a in (s.mean, s.power, s.gain, s.offset)):
+        raise ValueError("Invalid track scales")
+    i, c = s.fft_indices, s.fft_coefficients
+    if (i.ndim != 2 or i.shape[0] != tracks or c.shape != i.shape or not np.issubdtype(i.dtype, np.integer)
+            or (i < 0).any() or (i > length // 2).any() or not np.isfinite(c).all()):
+        raise ValueError("Invalid FFT payload")
+    for a in (s.final_gain, s.final_offset):
+        if a.shape != (tracks,) or not np.isfinite(a).all():
+            raise ValueError("Invalid final-signal calibration")
+    if (s.final_gain < 0).any():
+        raise ValueError("Invalid final-signal calibration")

zipoop/teacher.py

@@ -0,0 +1,43 @@
+"""Teacher signal restoration. Rules (shared by teacher and real data, see evaluation.restore_real):
+
+1. Valid positions use the decoded value. Missing positions: RNA-seq tracks are 0 and count as valid
+   (unrecorded = zero expression, and the archive mask agrees with "unrecorded" position by position);
+   all other tracks stay missing (ChIP/DNase/ATAC bigWigs record zeros explicitly, so missing is real).
+2. The raw network output at missing positions is never used as a target: training never constrained it.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from .decoder import Decoder
+from .panel import zero_tracks
+from .shard import Shard, read_shard
+
+
+@dataclass
+class TeacherWindow:
+    key: str
+    dna: torch.Tensor      # (L,) int64, 0-3 = ACGT, 4 = N
+    signal: torch.Tensor   # (L, T) float32, 0 where missing
+    valid: torch.Tensor    # (L, T) bool, all True for RNA tracks
+
+
+def restore(decoder: Decoder, zero) -> tuple[torch.Tensor, torch.Tensor]:
+    y = decoder.decode(missing_nan=False)
+    valid = torch.from_numpy(np.ascontiguousarray(decoder.valid)).to(y.device)
+    zero = torch.as_tensor(np.asarray(zero, dtype=bool), device=y.device)
+    return torch.where(valid, y, torch.zeros_like(y)), valid | zero[None, :]
+
+
+def load_window(source, device="cpu", zero=None) -> TeacherWindow:
+    """``source``: shard path, ShardView, or an already parsed Shard."""
+    shard = source if isinstance(source, Shard) else read_shard(source)
+    zero = zero_tracks() if zero is None else np.asarray(zero, dtype=bool)
+    if zero.shape != (shard.shape[1],):
+        raise ValueError(f"zero_tracks length {zero.shape} does not match {shard.shape[1]} tracks in {shard.key}")
+    signal, valid = restore(Decoder(shard, device), zero)
+    return TeacherWindow(shard.key, torch.from_numpy(shard.dna.astype(np.int64)).to(signal.device),
+                         signal.float(), valid)

zipoop-0.1.0.dist-info/METADATA

@@ -0,0 +1,51 @@
+Metadata-Version: 2.4
+Name: zipoop
+Version: 0.1.0
+Summary: Decoder and reader for Zipoop teacher shards (per-window neural compression of ENCODE signal tracks)
+Author: zatumout
+License-Expression: MIT
+Project-URL: Dataset, https://huggingface.co/datasets/zatumout/zipoop-teacher-shards
+Keywords: genomics,ENCODE,neural compression,distillation
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.26
+Requires-Dist: torch>=2.2
+Provides-Extra: test
+Requires-Dist: pytest>=8; extra == "test"
+Dynamic: license-file
+
+# zipoop
+
+Decoder and reader for the Zipoop teacher shards (dataset: https://huggingface.co/datasets/zatumout/zipoop-teacher-shards): 89,438 windows of 32 kb, 1,885 ENCODE tracks, stored as one `<chrom>.npz` per chromosome (`weights_npz`, with `MANIFEST.tsv`).
+
+## Data layout
+
+`<root>/<chrom>.npz` (25 files) plus `MANIFEST.tsv` (chrom, windows, members, bytes, sha256). Each file is a deflate zip with members `<window>/<key>.npy` (161 arrays per window), windows in genome order. Window names are `<chrom>_<start>_<end>`.
+
+## Usage
+
+```python
+from zipoop import Dataset
+with Dataset("/path/to/weights_npz") as ds:
+    w = ds.load_window(ds.windows("chr21")[0])   # TeacherWindow: key, dna (L,), signal (L, T), valid (L, T)
+```
+
+```
+zipoop verify /path/to/weights_npz      # size + sha256 against MANIFEST.tsv
+zipoop show   /path/to/weights_npz chr21_10027008_10059776
+```
+
+Lower level: `Dataset.read(window)` gives the parsed `Shard`, `Decoder(shard).decode(start, end, tracks)` decodes any slice or track subset. Evaluation helpers (`window_pearson`, `summarize`, `calibrate`, `restore_real`) are in `zipoop.evaluation`.
+
+Rules baked into `load_window` (see `zipoop/teacher.py`): RNA-seq tracks are 0 and valid where unrecorded; other tracks stay missing there. Opening a chromosome file parses its whole zip directory (chr1 about 3.6 s and 0.9 GB), so keep the `Dataset` open; forked DataLoader workers reopen their own handles.
+
+Only the shard variant in this dataset is supported (hybrid `pooper-dna-v3`, max-pool, sinusoidal positions, mask `track-rle-lzma-v2`); other variants raise a clear error.
+
+Tests: `ZIPOOP_DATA=<weights_npz> pytest tests` (the pooper-equivalence test also needs the original shards and the pooper package).
+
+## License and data terms
+
+The code is MIT-licensed (see `LICENSE`; keep the copyright notice "zatumout"). The shard dataset is released separately under CC BY 4.0 (credit: zatumout), commercial use and model distillation allowed. The shards and `zipoop/data/encode-panel_1885.json` are derived from public ENCODE tracks: when you use them, also cite and acknowledge the ENCODE Consortium as described at https://www.encodeproject.org/help/citing-encode.