diff --git a/python/genvarloader/_dataset/_concat_io.py b/python/genvarloader/_dataset/_concat_io.py index 44340757..e6b92d63 100644 --- a/python/genvarloader/_dataset/_concat_io.py +++ b/python/genvarloader/_dataset/_concat_io.py @@ -80,10 +80,17 @@ def copy_runs( merged = np.empty(n_slots + 1, np.int64) merged[0] = 0 lengths = merged[1:] - for dst_start, ds_vec, slots in plan.slot_batches(): + for dst_start, ds_vec, slots, ds_uniq in plan.slot_batches(): n = len(slots) out = lengths[dst_start : dst_start + n] - for d in np.unique(ds_vec): + if len(ds_uniq) == 1: + # Single-sourced batch -- every region-axis batch and every + # ExplicitRunPlan batch. Skip the mask entirely: it would be all-True + # over as many as _SLOT_BATCH_SLOTS elements. + off = src_offsets[int(ds_uniq[0])] + out[:] = off[slots + 1] - off[slots] + continue + for d in ds_uniq: m = ds_vec == d off = src_offsets[int(d)] sl = slots[m] diff --git a/python/genvarloader/_dataset/_concat_plan.py b/python/genvarloader/_dataset/_concat_plan.py index 4cdc8139..93844dba 100644 --- a/python/genvarloader/_dataset/_concat_plan.py +++ b/python/genvarloader/_dataset/_concat_plan.py @@ -282,20 +282,31 @@ def n_slots(self) -> int: def slot_batches( self, - ) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]": - """Yield ``(dst_start, src_ds, src_slots)`` batches in destination order. + ) -> ( + "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64], NDArray[np.int64]]]" + ): + """Yield ``(dst_start, src_ds, src_slots, src_ds_uniq)`` in destination order. Each batch describes a destination-contiguous span: ``src_ds[i]`` and ``src_slots[i]`` are the origin of merged slot ``dst_start + i``. - Concatenating every batch in order rebuilds :func:`provenance`'s output - exactly, which is what pins this method. + Concatenating every batch's first three fields in order rebuilds + :func:`provenance`'s output exactly, which is what pins this method. + + ``src_ds_uniq`` is the sorted distinct values of ``src_ds``. It is + yielded rather than left to the consumer because this method already + knows it for free: on the regions axis a batch is single-sourced, and on + the samples axis ``src_ds`` is loop-invariant, so a consumer calling + ``np.unique`` per batch would re-sort the identical array once per + region -- 3,734 sorts of a 1.07e6-element array at the chr22 grid. Yields: - ``(dst_start, src_ds, src_slots)``, where the two arrays are int64 - and equal in length. + ``(dst_start, src_ds, src_slots, src_ds_uniq)``. The first three are + int64, with ``src_ds`` and ``src_slots`` equal in length. """ if self.axis == "regions": for run in self: + # Every batch of a region run is single-sourced. + uniq = np.array([run.src], np.int64) pos, dst = run.src_start, run.dst_start while pos < run.src_stop: n = min(_SLOT_BATCH_SLOTS, run.src_stop - pos) @@ -303,6 +314,7 @@ def slot_batches( dst, np.full(n, run.src, np.int64), np.arange(pos, pos + n, dtype=np.int64), + uniq, ) pos += n dst += n @@ -318,10 +330,13 @@ def slot_batches( p = np.arange(self.ploidy, dtype=np.int64) # `order` is per merged SAMPLE; each contributes `ploidy` adjacent slots. ds_vec = np.repeat(ds, self.ploidy) + # Loop-invariant: the merged sample axis, hence its source-dataset + # grouping, is identical for every region. + ds_uniq = np.unique(ds_vec) for r in range(n_regions): base = (r * s_d + w) * self.ploidy slots = (base[:, None] + p[None, :]).reshape(-1) - yield (r * n_merged * self.ploidy, ds_vec, slots) + yield (r * n_merged * self.ploidy, ds_vec, slots, ds_uniq) def _segments(self) -> tuple[NDArray[np.int64], NDArray[np.int64]]: """Half-open ``[start, stop)`` spans of ``order`` with no break inside.""" @@ -413,14 +428,21 @@ def n_slots(self) -> int: def slot_batches( self, - ) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]": - """Yield one ``(dst_start, src_ds, src_slots)`` batch per run.""" + ) -> ( + "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64], NDArray[np.int64]]]" + ): + """Yield one ``(dst_start, src_ds, src_slots, src_ds_uniq)`` batch per run. + + A run is drawn from exactly one source dataset, so ``src_ds_uniq`` is + always a single element here. + """ for r in self._runs: n = r.src_stop - r.src_start yield ( r.dst_start, np.full(n, r.src, np.int64), np.arange(r.src_start, r.src_stop, dtype=np.int64), + np.array([r.src], np.int64), ) diff --git a/tests/conftest.py b/tests/conftest.py index 2435b1f5..102b83b8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -176,6 +176,33 @@ def _svar2_slot_src(tmp_path_factory) -> tuple[Path, Path]: return bcf, ref +@pytest.fixture(scope="session") +def svar2_slot_store(_svar2_slot_src, tmp_path_factory) -> Path: + """The ``.svar2`` store built from `_svar2_slot_src`. + + Shared and session-scoped, so consumers must treat it as READ-ONLY: copy it + into ``tmp_path`` before mutating anything inside it, or the corruption + leaks into every later consumer of this fixture. + """ + from genoray import _core + + bcf, ref = _svar2_slot_src + store = tmp_path_factory.mktemp("svar2_slot_store") / "store.svar2" + _core.run_conversion_pipeline( + str(bcf), + str(ref), + ["chr1"], + str(store), + ["S0", "S1"], + 25_000, + 2, + 1, + 8 * 1024 * 1024, + ) + assert (store / "meta.json").exists(), "svar2 conversion did not finish" + return store + + @pytest.fixture(scope="session") def svar2_slot_reference(_svar2_slot_src): """Opened ``gvl.Reference`` matching `phased_svar2_gvl`'s tiny chr1. @@ -192,7 +219,7 @@ def svar2_slot_reference(_svar2_slot_src): @pytest.fixture(scope="session") -def phased_svar2_gvl(_svar2_slot_src, tmp_path_factory) -> Path: +def phased_svar2_gvl(svar2_slot_store, tmp_path_factory) -> Path: """A gvl dataset written from a ``.svar2`` source -> opens as ``Svar2Haps``. Reproduces the Phase-0-pinned record class for #315 (see @@ -215,24 +242,11 @@ def phased_svar2_gvl(_svar2_slot_src, tmp_path_factory) -> Path: empirically before picking 80 (green up to ~40 regions, red by 60+). """ import polars as pl - from genoray import SparseVar2, _core + from genoray import SparseVar2 import genvarloader as gvl - bcf, ref = _svar2_slot_src - store = tmp_path_factory.mktemp("svar2_slot_store") / "store.svar2" - _core.run_conversion_pipeline( - str(bcf), - str(ref), - ["chr1"], - str(store), - ["S0", "S1"], - 25_000, - 2, - 1, - 8 * 1024 * 1024, - ) - assert (store / "meta.json").exists(), "svar2 conversion did not finish" + store = svar2_slot_store n_regions = 80 bed = pl.DataFrame( @@ -325,7 +339,7 @@ def _build_svar2(vcf_text: str, samples: list[str], d: Path, name: str) -> Path: return out -@pytest.fixture(scope="module") +@pytest.fixture(scope="session") def svar2_store_2s(tmp_path_factory) -> Path: """A two-sample (S0, S1) .svar2 store. @@ -333,6 +347,14 @@ def svar2_store_2s(tmp_path_factory) -> Path: six of this fixture's former copies were private duplicates, three of which lived in ``tests/dataset/`` and silently shadowed ``tests/dataset/conftest.py``'s same-named three-sample fixture. + + Session-scoped, and therefore SHARED: treat it as READ-ONLY. A test that + needs to mutate the store must ``shutil.copytree`` it into ``tmp_path`` and + mutate the copy. Mutating it in place corrupts every later consumer, and the + resulting failures surface in other modules -- when + ``test_fingerprint_detects_mutated_store`` appended one byte to a ``.bin`` + here, it produced 17 failures plus a Rust ``cast_slice`` panic. That is why + this fixture was module-scoped until the mutating test was fixed to copy. """ d = tmp_path_factory.mktemp("svar2_2s") return _build_svar2(_SVAR2_VCF_2S, ["S0", "S1"], d, "store") diff --git a/tests/dataset/test_svar2_readbound_tracks.py b/tests/dataset/test_svar2_readbound_tracks.py index a8b4ddf1..30cb703c 100644 --- a/tests/dataset/test_svar2_readbound_tracks.py +++ b/tests/dataset/test_svar2_readbound_tracks.py @@ -12,58 +12,21 @@ from __future__ import annotations -import subprocess -from pathlib import Path import numpy as np import pytest -# 40 bp reference (chr1). VCF POS (1-based) -> 0-based: SNP@2 (A>G, low-carrier, -# routes to var_key), INS@6 (C>CAT), SNP@9 (G>C, carried by 3 haps -> dense/snp -# per the cost model used in test_svar2_readbound_haps.py), DEL@11 (GTA>G, -# ilen -2). Exercises both var_key and dense/snp + dense/indel channels. -_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" -_VCF = """\ -##fileformat=VCFv4.2 -##contig= -##FORMAT= -#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1 -chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0 -chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1 -chr1\t10\t.\tG\tC\t.\t.\t.\tGT\t1|1\t1|0 -chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1 -""" - - -@pytest.fixture(scope="module") -def svar2_store(tmp_path_factory) -> Path: - from genoray import _core - - d = tmp_path_factory.mktemp("svar2_readbound_tracks") - ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") - subprocess.run(["samtools", "faidx", str(ref)], check=True) - - vcf = d / "in.vcf" - vcf.write_text(_VCF) - bcf = d / "in.bcf" - subprocess.run(["bcftools", "view", "-Ob", "-o", str(bcf), str(vcf)], check=True) - subprocess.run(["bcftools", "index", str(bcf)], check=True) - - out = d / "store" - _core.run_conversion_pipeline( - str(bcf), - str(ref), - ["chr1"], - str(out), - ["S0", "S1"], - 25_000, - 2, - 1, - 8 * 1024 * 1024, - ) - assert (out / "meta.json").exists(), "conversion did not finish" - return out +# The store comes from the shared session-scoped `svar2_slot_store` fixture in +# tests/conftest.py. Its reference and VCF used to be duplicated verbatim here: +# SNP@2 (A>G, low-carrier, routes to var_key), INS@6 (C>CAT), SNP@9 (G>C, +# carried by 3 haps -> dense/snp per the cost model used in +# test_svar2_readbound_haps.py), DEL@11 (GTA>G, ilen -2) -- so both var_key and +# the dense/snp + dense/indel channels are exercised. The fixture is shared and +# must be treated as read-only. +# +# Used under its own name rather than aliased: this module lives under +# tests/dataset/, where tests/dataset/conftest.py defines a DIFFERENT 3-sample +# fixture named `svar2_store`, and an alias would silently bind to that instead. def _synthetic_track_inputs(regions, seed=0): @@ -87,14 +50,14 @@ def _synthetic_track_inputs(regions, seed=0): [(0, 40), (2, 2), (20, 25)], # empty region + a variant-free window ], ) -def test_readbound_tracks_match_union_oracle(svar2_store, regions): +def test_readbound_tracks_match_union_oracle(svar2_slot_store, regions): import genoray from tests._oracles.svar2_source import SparseVar2Source from tests._oracles.svar2_readbound_inputs import build_readbound_tracks contig = "chr1" - sv = genoray.SparseVar2(str(svar2_store)) + sv = genoray.SparseVar2(str(svar2_slot_store)) S, P = sv.n_samples, sv.ploidy assert (S, P) == (2, 2) @@ -155,7 +118,7 @@ def test_readbound_tracks_match_union_oracle(svar2_store, regions): pytest.fail("data mismatch but no single hap slice differed (offset bug?)") -def test_readbound_tracks_match_union_oracle_with_shifts(svar2_store): +def test_readbound_tracks_match_union_oracle_with_shifts(svar2_slot_store): """Non-trivial per-hap jitter shifts must also match byte-for-byte.""" import genoray @@ -164,7 +127,7 @@ def test_readbound_tracks_match_union_oracle_with_shifts(svar2_store): contig = "chr1" regions = [(0, 40), (5, 20)] - sv = genoray.SparseVar2(str(svar2_store)) + sv = genoray.SparseVar2(str(svar2_slot_store)) S, P = sv.n_samples, sv.ploidy n_q = len(regions) * S diff --git a/tests/unit/dataset/test_concat_plan.py b/tests/unit/dataset/test_concat_plan.py index 870b5140..5cb99140 100644 --- a/tests/unit/dataset/test_concat_plan.py +++ b/tests/unit/dataset/test_concat_plan.py @@ -257,9 +257,12 @@ def test_slot_batches_reproduce_the_provenance_map(): want = provenance(axis, shapes, ploidy, order=order) got = np.zeros_like(want) seen = np.zeros(len(want), bool) - for dst_start, ds_vec, slots in plan.slot_batches(): + for dst_start, ds_vec, slots, ds_uniq in plan.slot_batches(): n = len(slots) assert len(ds_vec) == n + # The yielded uniques must be exactly what a consumer would + # have computed itself; copy_runs trusts them to pick files. + np.testing.assert_array_equal(ds_uniq, np.unique(ds_vec)) got[dst_start : dst_start + n, 0] = ds_vec got[dst_start : dst_start + n, 1] = slots seen[dst_start : dst_start + n] = True @@ -295,12 +298,12 @@ def test_slot_batches_chunk_large_region_runs(monkeypatch): assert len(list(plan)) == 1 batches = list(plan.slot_batches()) - assert [len(slots) for _, _, slots in batches] == [5, 5, 5, 5, 4] + assert [len(slots) for _, _, slots, _ in batches] == [5, 5, 5, 5, 4] assert plan.n_slots == 24 # Destination stays contiguous across every chunk boundary. - assert [dst for dst, _, _ in batches] == [0, 5, 10, 15, 20] + assert [dst for dst, _, _, _ in batches] == [0, 5, 10, 15, 20] np.testing.assert_array_equal( - np.concatenate([slots for _, _, slots in batches]), np.arange(24) + np.concatenate([slots for _, _, slots, _ in batches]), np.arange(24) ) @@ -314,6 +317,9 @@ def test_explicit_run_plan_round_trips_a_hand_built_list(): assert [b[0] for b in batches] == [0, 2] np.testing.assert_array_equal(batches[0][2], [0, 1]) np.testing.assert_array_equal(batches[1][1], [1, 1]) + # A run is single-sourced, so each batch reports exactly one dataset. + np.testing.assert_array_equal(batches[0][3], [0]) + np.testing.assert_array_equal(batches[1][3], [1]) def test_as_plan_rejects_a_generator(): diff --git a/tests/unit/dataset/test_svar2_link.py b/tests/unit/dataset/test_svar2_link.py index 5649f923..4462e3fa 100644 --- a/tests/unit/dataset/test_svar2_link.py +++ b/tests/unit/dataset/test_svar2_link.py @@ -44,8 +44,9 @@ def test_verify_none_link_is_noop(tmp_path: Path): def test_fingerprint_detects_mutated_store(svar2_store_2s: Path, tmp_path: Path): - # `svar2_store_2s` is a module-scoped fixture shared with five other test - # modules -- mutate a private copy, never the shared store itself. + # `svar2_store_2s` is a session-scoped fixture shared with five other test + # modules -- mutate a private copy, never the shared store itself. Mutating + # it in place corrupts every later consumer across the whole session. store = shutil.copytree(svar2_store_2s, tmp_path / "store.svar2") gvl_path = tmp_path / "ds.gvl"