diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 42319117..97a8c1b6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,7 +10,7 @@ repos: - id: check-case-conflict - id: check-yaml - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.12 + rev: v0.16.8 hooks: - id: ruff-check - id: ruff-format @@ -34,7 +34,7 @@ repos: always_run: true stages: [pre-push] - repo: https://github.com/commitizen-tools/commitizen - rev: v4.15.1 + rev: v4.18.1 hooks: - id: commitizen stages: [commit-msg] diff --git a/docs/roadmaps/phase-3-getitem-glue-audit.md b/docs/roadmaps/phase-3-getitem-glue-audit.md index c16e573b..2937ca99 100644 --- a/docs/roadmaps/phase-3-getitem-glue-audit.md +++ b/docs/roadmaps/phase-3-getitem-glue-audit.md @@ -244,7 +244,7 @@ pub fn reconstruct_haplotypes_fused<'py>( out_data, out_offsets = gvl_rust.reconstruct_haplotypes_fused( regions=req.regions, geno_offset_idx=req.geno_offset_idx, - geno_offsets=self.genotypes.offsets, # already (2,n) or 1-D; Rust normalizes + geno_offsets=self.genotypes.offsets, # already (2,n) or 1-D; Rust normalizes geno_v_idxs=self.genotypes.data, v_starts=self.variants.start, ilens=self.variants.ilen, diff --git a/docs/source/dataset.md b/docs/source/dataset.md index 47503881..b115a9bb 100644 --- a/docs/source/dataset.md +++ b/docs/source/dataset.md @@ -130,7 +130,11 @@ Set `realign_tracks=False` in two cases: ds = gvl.get_dummy_dataset() # Reference-coordinate float tracks alongside haplotypes -ds_ref_tracks = ds.with_seqs("haplotypes").with_tracks(["read-depth"]).with_settings(realign_tracks=False) +ds_ref_tracks = ( + ds.with_seqs("haplotypes") + .with_tracks(["read-depth"]) + .with_settings(realign_tracks=False) +) # Interval tracks alongside haplotypes (realign_tracks=False is required) ds_itvs = ( @@ -164,7 +168,9 @@ ds = gvl.Dataset.open("ds.gvl", reference="ref.fa", var_fields=["AF"]) rv = ds.with_seqs("variants")[0, 0] rv["AF"] # per-variant AF values, aligned with rv.alt/.start/.ilen -win = ds.with_seqs("variant-windows", gvl.VarWindowOpt(...)).with_output_format("flat")[0, 0] +win = ds.with_seqs("variant-windows", gvl.VarWindowOpt(...)).with_output_format("flat")[ + 0, 0 +] win.fields["AF"] # same field, alongside win.fields["start"]/["ilen"] ``` diff --git a/docs/source/faq.md b/docs/source/faq.md index 576c4dd4..41418f5a 100644 --- a/docs/source/faq.md +++ b/docs/source/faq.md @@ -76,7 +76,7 @@ GVL's read path (haplotype reconstruction and track re-alignment) is parallelize Environment variables configure a whole process, which means a script's parallelism can't be determined by reading the script — a value in a shell profile, a Dockerfile, or a SLURM template changes how it runs. To state the policy where a reader can see it, set it on the dataset: ```python -ds = ds.with_settings(parallel=False) # True | False | "auto" +ds = ds.with_settings(parallel=False) # True | False | "auto" ``` - `True` — always hand batches to rayon, whatever their size. diff --git a/docs/superpowers/REGRESSIONS.md b/docs/superpowers/REGRESSIONS.md index d077d7e5..bbbfef74 100644 --- a/docs/superpowers/REGRESSIONS.md +++ b/docs/superpowers/REGRESSIONS.md @@ -148,19 +148,35 @@ cells fit within a ≤256 GB allocation. The OOM-kills observed in the new pipel ```python import genvarloader as gvl from time import perf_counter + # tracks-only: with_seqs(None) -> batch is a single track tensor (no RaggedVariants) -ds = gvl.Dataset.open(DS_PATH, FASTA).with_seqs(None).with_tracks("read-depth","tracks").with_len(16384) +ds = ( + gvl.Dataset.open(DS_PATH, FASTA) + .with_seqs(None) + .with_tracks("read-depth", "tracks") + .with_len(16384) +) for bs in (8, 32): dl = ds.to_dataloader(batch_size=bs, shuffle=False) - ny=nnuc=0; burn=5; nb=150; t0=perf_counter(); esz=4; done=False + ny = nnuc = 0 + burn = 5 + nb = 150 + t0 = perf_counter() + esz = 4 + done = False while not done: for b in dl: - trk = b[1] if isinstance(b,(list,tuple)) else b # track tensor - if ny==burn: t0=perf_counter() - if ny>=burn: nnuc+=trk.numel(); esz=trk.element_size() - ny+=1 - if ny>=nb: done=True; break - print(bs, nnuc/(perf_counter()-t0)/2**20*esz, "MiB/s") + trk = b[1] if isinstance(b, (list, tuple)) else b # track tensor + if ny == burn: + t0 = perf_counter() + if ny >= burn: + nnuc += trk.numel() + esz = trk.element_size() + ny += 1 + if ny >= nb: + done = True + break + print(bs, nnuc / (perf_counter() - t0) / 2**20 * esz, "MiB/s") ``` Run with `NUMBA_NUM_THREADS=1` for the single-thread numbers. The 0.6.1 side uses the equivalent old API (`gvl.Dataset.open(ds, fasta, return_sequences=False)`; restored in @@ -171,12 +187,18 @@ the *same* BED + variants + BigWig table to keep regions/samples identical. ```python import numpy as np, numba as nb + + @nb.njit(parallel=True) def f(x): - s=0.0 - for i in nb.prange(x.size): s+=x[i] + s = 0.0 + for i in nb.prange(x.size): + s += x[i] return s -f(np.ones(1000)); print(nb.threading_layer()) # 'tbb' if installed, else 'omp' + + +f(np.ones(1000)) +print(nb.threading_layer()) # 'tbb' if installed, else 'omp' ``` 0.6.1 pulled `tbb` transitively; 0.24.1 made it optional, so fresh installs report `omp`. diff --git a/docs/superpowers/plans/2026-05-08-get-splice-bed.md b/docs/superpowers/plans/2026-05-08-get-splice-bed.md index 75b824d7..e8944275 100644 --- a/docs/superpowers/plans/2026-05-08-get-splice-bed.md +++ b/docs/superpowers/plans/2026-05-08-get-splice-bed.md @@ -259,11 +259,13 @@ def get_splice_bed( if contigs is not None: lf = lf.filter(pl.col("seqname").is_in(contigs)) - lf = lf.filter(pl.col("feature") == "CDS").rename({ - "seqname": "chrom", - "start": "chromStart", - "end": "chromEnd", - }) + lf = lf.filter(pl.col("feature") == "CDS").rename( + { + "seqname": "chrom", + "start": "chromStart", + "end": "chromEnd", + } + ) lf = lf.with_columns( pl.col("chromStart") - 1, diff --git a/docs/superpowers/plans/2026-05-08-tracks-generalize.md b/docs/superpowers/plans/2026-05-08-tracks-generalize.md index ed578306..52f26c66 100644 --- a/docs/superpowers/plans/2026-05-08-tracks-generalize.md +++ b/docs/superpowers/plans/2026-05-08-tracks-generalize.md @@ -214,20 +214,24 @@ the implementation; delete afterward. import polars as pl import polars_bio as pb -queries = pl.DataFrame({ - "chrom": ["chr1", "chr1", "chr2"], - "start": [0, 50, 0], - "end": [10, 60, 10], - "_q": [0, 1, 2], - "sample_id": ["s0", "s0", "s0"], -}) -table = pl.DataFrame({ - "chrom": ["chr1", "chr1"], - "start": [2, 100], - "end": [5, 105], - "value": [1.0, 2.0], - "sample_id": ["s0", "s0"], -}) +queries = pl.DataFrame( + { + "chrom": ["chr1", "chr1", "chr2"], + "start": [0, 50, 0], + "end": [10, 60, 10], + "_q": [0, 1, 2], + "sample_id": ["s0", "s0", "s0"], + } +) +table = pl.DataFrame( + { + "chrom": ["chr1", "chr1"], + "start": [2, 100], + "end": [5, 105], + "value": [1.0, 2.0], + "sample_id": ["s0", "s0"], + } +) print("=== overlap ===") ov = pb.overlap( @@ -300,13 +304,15 @@ from genvarloader._table import Table def make_long_df(): - return pl.DataFrame({ - "sample_id": ["s0", "s0", "s1", "s1"], - "chrom": ["chr1", "chr1", "chr1", "chr2"], - "start": [10, 100, 20, 0], - "end": [20, 110, 30, 5], - "value": [1.0, 2.0, 3.0, 4.0], - }) + return pl.DataFrame( + { + "sample_id": ["s0", "s0", "s1", "s1"], + "chrom": ["chr1", "chr1", "chr1", "chr2"], + "start": [10, 100, 20, 0], + "end": [20, 110, 30, 5], + "value": [1.0, 2.0, 3.0, 4.0], + } + ) def test_table_init_from_long_df(): @@ -378,19 +384,20 @@ class Table: ) -> None: self.name = name df = self._normalize_input(data, column_map) - df = df.cast({ - "sample_id": pl.Utf8, - "chrom": pl.Utf8, - "start": pl.Int64, - "end": pl.Int64, - "value": pl.Float32, - }).sort("chrom", "sample_id", "start") + df = df.cast( + { + "sample_id": pl.Utf8, + "chrom": pl.Utf8, + "start": pl.Int64, + "end": pl.Int64, + "value": pl.Float32, + } + ).sort("chrom", "sample_id", "start") self._df = df self.samples = sorted(df["sample_id"].unique().to_list()) self.contigs = { row["chrom"]: int(row["max_end"]) - for row in df - .group_by("chrom") + for row in df.group_by("chrom") .agg(pl.col("end").max().alias("max_end")) .iter_rows(named=True) } @@ -477,18 +484,22 @@ git commit -m "feat: add Table skeleton with long-form DataFrame init" ```python def test_table_init_from_dict_of_dfs(): per_sample = { - "s0": pl.DataFrame({ - "chrom": ["chr1"], - "start": [10], - "end": [20], - "value": [1.0], - }), - "s1": pl.DataFrame({ - "chrom": ["chr2"], - "start": [0], - "end": [5], - "value": [2.0], - }), + "s0": pl.DataFrame( + { + "chrom": ["chr1"], + "start": [10], + "end": [20], + "value": [1.0], + } + ), + "s1": pl.DataFrame( + { + "chrom": ["chr2"], + "start": [0], + "end": [5], + "value": [2.0], + } + ), } t = Table("signal", per_sample) assert t.samples == ["s0", "s1"] @@ -496,13 +507,15 @@ def test_table_init_from_dict_of_dfs(): def test_table_column_map_renames_long_form(): - df = pl.DataFrame({ - "donor": ["s0"], - "chrom": ["chr1"], - "chromStart": [10], - "chromEnd": [20], - "signal": [1.5], - }) + df = pl.DataFrame( + { + "donor": ["s0"], + "chrom": ["chr1"], + "chromStart": [10], + "chromEnd": [20], + "signal": [1.5], + } + ) t = Table( "signal", df, @@ -519,12 +532,14 @@ def test_table_column_map_renames_long_form(): def test_table_column_map_per_sample_dict(): per_sample = { - "s0": pl.DataFrame({ - "chrom": ["chr1"], - "chromStart": [10], - "chromEnd": [20], - "signal": [1.5], - }), + "s0": pl.DataFrame( + { + "chrom": ["chr1"], + "chromStart": [10], + "chromEnd": [20], + "signal": [1.5], + } + ), } t = Table( "signal", @@ -694,13 +709,15 @@ def _brute_count(df: pl.DataFrame, contig: str, starts, ends, samples): def test_table_count_intervals_matches_brute_force(): - df = pl.DataFrame({ - "sample_id": ["s0", "s0", "s0", "s1", "s1"], - "chrom": ["chr1", "chr1", "chr1", "chr1", "chr1"], - "start": [0, 50, 200, 10, 60], - "end": [10, 60, 210, 20, 70], - "value": [1.0, 2.0, 3.0, 4.0, 5.0], - }) + df = pl.DataFrame( + { + "sample_id": ["s0", "s0", "s0", "s1", "s1"], + "chrom": ["chr1", "chr1", "chr1", "chr1", "chr1"], + "start": [0, 50, 200, 10, 60], + "end": [10, 60, 210, 20, 70], + "value": [1.0, 2.0, 3.0, 4.0, 5.0], + } + ) t = Table("signal", df) starts = np.array([0, 55, 100, 200], dtype=np.int32) ends = np.array([15, 65, 150, 205], dtype=np.int32) @@ -801,12 +818,14 @@ def _build_queries( samples: list[str], ) -> pl.DataFrame: n = len(starts) - return pl.DataFrame({ - "_q": np.arange(n, dtype=np.int64), - "chrom": np.repeat(np.array([contig], dtype=object), n), - "start": starts, - "end": ends, - }).join(pl.DataFrame({"sample_id": samples}), how="cross") + return pl.DataFrame( + { + "_q": np.arange(n, dtype=np.int64), + "chrom": np.repeat(np.array([contig], dtype=object), n), + "start": starts, + "end": ends, + } + ).join(pl.DataFrame({"sample_id": samples}), how="cross") ``` - [ ] **Step 4: Run tests** @@ -841,13 +860,15 @@ from genvarloader._utils import lengths_to_offsets def test_table_intervals_from_offsets_roundtrip(): - df = pl.DataFrame({ - "sample_id": ["s0", "s0", "s1"], - "chrom": ["chr1", "chr1", "chr1"], - "start": [0, 50, 10], - "end": [10, 60, 20], - "value": [1.5, 2.5, 3.5], - }) + df = pl.DataFrame( + { + "sample_id": ["s0", "s0", "s1"], + "chrom": ["chr1", "chr1", "chr1"], + "start": [0, 50, 10], + "end": [10, 60, 20], + "value": [1.5, 2.5, 3.5], + } + ) t = Table("signal", df) starts = np.array([0, 40], dtype=np.int32) ends = np.array([15, 70], dtype=np.int32) @@ -1196,22 +1217,26 @@ ddir = Path(__file__).parents[1] / "data" def _make_bed(tmp_path: Path) -> pl.DataFrame: - bed = pl.DataFrame({ - "chrom": ["chr1", "chr1"], - "chromStart": [0, 100], - "chromEnd": [50, 200], - }) + bed = pl.DataFrame( + { + "chrom": ["chr1", "chr1"], + "chromStart": [0, 100], + "chromEnd": [50, 200], + } + ) return bed def _make_table_df() -> pl.DataFrame: - return pl.DataFrame({ - "sample_id": ["s0", "s0", "s1", "s1"], - "chrom": ["chr1", "chr1", "chr1", "chr1"], - "start": [10, 110, 5, 150], - "end": [20, 130, 15, 160], - "value": [1.0, 2.0, 3.0, 4.0], - }) + return pl.DataFrame( + { + "sample_id": ["s0", "s0", "s1", "s1"], + "chrom": ["chr1", "chr1", "chr1", "chr1"], + "start": [10, 110, 5, 150], + "end": [20, 130, 15, 160], + "value": [1.0, 2.0, 3.0, 4.0], + } + ) def test_write_with_table_only_roundtrip(tmp_path): @@ -1267,11 +1292,13 @@ Append to `tests/dataset/test_write_tracks.py`: ```python def test_write_with_mixed_bigwigs_and_table(tmp_path): - bed = pl.DataFrame({ - "chrom": ["chr1"], - "chromStart": [0], - "chromEnd": [200], - }) + bed = pl.DataFrame( + { + "chrom": ["chr1"], + "chromStart": [0], + "chromEnd": [200], + } + ) bw_dir = ddir / "bigwig" bw = gvl.BigWigs( "bw_signal", @@ -1283,13 +1310,15 @@ def test_write_with_mixed_bigwigs_and_table(tmp_path): # Table sample IDs match the BigWigs sample IDs so the intersection is non-empty. table = Table( "tab_signal", - pl.DataFrame({ - "sample_id": ["sample_0", "sample_1"], - "chrom": ["chr1", "chr1"], - "start": [0, 50], - "end": [10, 60], - "value": [9.0, 8.0], - }), + pl.DataFrame( + { + "sample_id": ["sample_0", "sample_1"], + "chrom": ["chr1", "chr1"], + "start": [0, 50], + "end": [10, 60], + "value": [9.0, 8.0], + } + ), ) out = tmp_path / "mixed.gvl" @@ -1305,23 +1334,27 @@ def test_write_duplicate_track_names_rejected(tmp_path): bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [0], "chromEnd": [100]}) t1 = Table( "dup", - pl.DataFrame({ - "sample_id": ["s0"], - "chrom": ["chr1"], - "start": [0], - "end": [10], - "value": [1.0], - }), + pl.DataFrame( + { + "sample_id": ["s0"], + "chrom": ["chr1"], + "start": [0], + "end": [10], + "value": [1.0], + } + ), ) t2 = Table( "dup", - pl.DataFrame({ - "sample_id": ["s0"], - "chrom": ["chr1"], - "start": [50], - "end": [60], - "value": [2.0], - }), + pl.DataFrame( + { + "sample_id": ["s0"], + "chrom": ["chr1"], + "start": [50], + "end": [60], + "value": [2.0], + } + ), ) with pytest.raises(ValueError, match="[Dd]uplicate"): gvl.write(path=tmp_path / "x.gvl", bed=bed, tracks=[t1, t2]) diff --git a/docs/superpowers/plans/2026-05-11-track-insertion-options.md b/docs/superpowers/plans/2026-05-11-track-insertion-options.md index cae288bb..5348b37e 100644 --- a/docs/superpowers/plans/2026-05-11-track-insertion-options.md +++ b/docs/superpowers/plans/2026-05-11-track-insertion-options.md @@ -870,10 +870,12 @@ def test_with_insertion_fill_dict_partial_falls_back(): def test_with_tracks_prunes_insertion_fill(): - tracks = _make_tracks(["a", "b"]).with_insertion_fill({ - "a": Constant(0.0), - "b": FlankSample(), - }) + tracks = _make_tracks(["a", "b"]).with_insertion_fill( + { + "a": Constant(0.0), + "b": FlankSample(), + } + ) new = tracks.with_tracks("a") assert set(new.insertion_fill) == {"a"} assert isinstance(new.insertion_fill["a"], Constant) diff --git a/docs/superpowers/plans/2026-05-12-1kg-bcftools-parity.md b/docs/superpowers/plans/2026-05-12-1kg-bcftools-parity.md index 5578904b..3027b2cb 100644 --- a/docs/superpowers/plans/2026-05-12-1kg-bcftools-parity.md +++ b/docs/superpowers/plans/2026-05-12-1kg-bcftools-parity.md @@ -224,15 +224,17 @@ def normalize_bcf(source_bcf: Path) -> Path: filtered = ONE_KG_DIR / "filtered.bcf" # Step A: left-align - result = run_shell([ - "bcftools", - "norm", - "-f", - str(REF), - "-O", - "u", - str(source_bcf), - ]) + result = run_shell( + [ + "bcftools", + "norm", + "-f", + str(REF), + "-O", + "u", + str(source_bcf), + ] + ) logger.info("bcftools norm (left-align) done") # Step B: atomize + split multiallelics; emit as bgzipped BCF @@ -303,16 +305,18 @@ Insert below `normalize_bcf`: ```python def make_pgen(filtered_bcf: Path) -> Path: out_prefix = ONE_KG_DIR / "filtered" - _ = run_shell([ - "plink2", - "--bcf", - str(filtered_bcf), - "--make-pgen", - "--vcf-half-call", - "r", - "--out", - str(out_prefix), - ]) + _ = run_shell( + [ + "plink2", + "--bcf", + str(filtered_bcf), + "--make-pgen", + "--vcf-half-call", + "r", + "--out", + str(out_prefix), + ] + ) return out_prefix.with_suffix(".pgen") ``` @@ -402,28 +406,32 @@ def pick_regions(filtered_bcf: Path) -> Path: bed_path = ONE_KG_DIR / "regions.bed" # Pull (chrom, pos) for chr21/chr22 only via bcftools query. - proc = run_shell([ - "bcftools", - "query", - "-f", - "%CHROM\t%POS\n", - "-r", - "chr21,chr22", - str(filtered_bcf), - ]) - raw = proc.stdout.decode().strip() - if not raw: - # The Zenodo dataset uses GRCh38 contig names without the "chr" prefix - # on some 1KG releases. Fall back to bare contig names. - proc = run_shell([ + proc = run_shell( + [ "bcftools", "query", "-f", "%CHROM\t%POS\n", "-r", - "21,22", + "chr21,chr22", str(filtered_bcf), - ]) + ] + ) + raw = proc.stdout.decode().strip() + if not raw: + # The Zenodo dataset uses GRCh38 contig names without the "chr" prefix + # on some 1KG releases. Fall back to bare contig names. + proc = run_shell( + [ + "bcftools", + "query", + "-f", + "%CHROM\t%POS\n", + "-r", + "21,22", + str(filtered_bcf), + ] + ) raw = proc.stdout.decode().strip() if not raw: @@ -453,14 +461,16 @@ def pick_regions(filtered_bcf: Path) -> Path: ends = starts + REGION_LEN strand = rng.choice(["+", "-"], size=N_REGIONS, replace=True) - out = pl.DataFrame({ - "chrom": chosen["chrom"].to_numpy(), - "start": starts, - "end": ends, - "name": ["."] * N_REGIONS, - "score": ["."] * N_REGIONS, - "strand": strand, - }) + out = pl.DataFrame( + { + "chrom": chosen["chrom"].to_numpy(), + "start": starts, + "end": ends, + "name": ["."] * N_REGIONS, + "score": ["."] * N_REGIONS, + "strand": strand, + } + ) out.write_csv(bed_path, include_header=False, separator="\t") logger.info(f"Wrote {N_REGIONS} regions to {bed_path}") return bed_path @@ -592,12 +602,14 @@ def generate_consensus_fastas(filtered_bcf: Path, bed_path: Path) -> None: for row_nr, chrom, start, end in bed.select( "index", "chrom", "start", "end" ).iter_rows(): - subseq = run_shell([ - "samtools", - "faidx", - str(REF), - f"{chrom}:{start + 1}-{end}", - ]) + subseq = run_shell( + [ + "samtools", + "faidx", + str(REF), + f"{chrom}:{start + 1}-{end}", + ] + ) for sample in samples: for hap in (0, 1): out_fa = CONS_DIR / f"1kg_{sample}_nr{row_nr}_h{hap}.fa" @@ -679,8 +691,7 @@ pytestmark = pytest.mark.slow def dataset_bcf(): return ( - gvl.Dataset - .open(data_dir / "1kg" / "phased_1kg.bcf.gvl", ref, rc_neg=False) + gvl.Dataset.open(data_dir / "1kg" / "phased_1kg.bcf.gvl", ref, rc_neg=False) .with_len("ragged") .with_seqs("haplotypes") .with_tracks(False) @@ -689,8 +700,7 @@ def dataset_bcf(): def dataset_pgen(): return ( - gvl.Dataset - .open(data_dir / "1kg" / "phased_1kg.pgen.gvl", ref, rc_neg=False) + gvl.Dataset.open(data_dir / "1kg" / "phased_1kg.pgen.gvl", ref, rc_neg=False) .with_len("ragged") .with_seqs("haplotypes") .with_tracks(False) @@ -699,8 +709,7 @@ def dataset_pgen(): def dataset_svar(): return ( - gvl.Dataset - .open(data_dir / "1kg" / "phased_1kg.svar.gvl", ref, rc_neg=False) + gvl.Dataset.open(data_dir / "1kg" / "phased_1kg.svar.gvl", ref, rc_neg=False) .with_len("ragged") .with_seqs("haplotypes") .with_tracks(False) diff --git a/docs/superpowers/plans/2026-05-21-refdataset-splicing.md b/docs/superpowers/plans/2026-05-21-refdataset-splicing.md index 96129fe1..6ad235e4 100644 --- a/docs/superpowers/plans/2026-05-21-refdataset-splicing.md +++ b/docs/superpowers/plans/2026-05-21-refdataset-splicing.md @@ -228,8 +228,7 @@ class SpliceMap: """ if isinstance(splice_info, str): sp_bed = ( - full_bed - .rename({splice_info: "splice_id"}) + full_bed.rename({splice_info: "splice_id"}) .with_row_index() .group_by("splice_id", maintain_order=True) .agg(pl.all()) @@ -241,8 +240,7 @@ class SpliceMap: "names for splice IDs and element ordering." ) sp_bed = ( - full_bed - .rename({splice_info[0]: "splice_id"}) + full_bed.rename({splice_info[0]: "splice_id"}) .with_row_index() .group_by("splice_id", maintain_order=True) .agg(pl.all().sort_by(splice_info[1])) @@ -390,13 +388,15 @@ Add a temporary script `/tmp/smoke_splicemap.py`: import polars as pl from genvarloader._dataset._splice import SpliceMap -bed = pl.DataFrame({ - "chrom": ["chr1"] * 4, - "chromStart": [0, 100, 200, 300], - "chromEnd": [10, 110, 210, 310], - "transcript_id": ["T1", "T1", "T2", "T2"], - "exon_number": [1, 2, 1, 2], -}) +bed = pl.DataFrame( + { + "chrom": ["chr1"] * 4, + "chromStart": [0, 100, 200, 300], + "chromEnd": [10, 110, 210, 310], + "transcript_id": ["T1", "T1", "T2", "T2"], + "exon_number": [1, 2, 1, 2], + } +) sm, sp_bed = SpliceMap.from_bed("transcript_id", bed) assert sm.n_rows == 2 print("ok", sm.n_rows, sp_bed.height) @@ -700,11 +700,13 @@ Append to `tests/test_ref_ds.py`: ```python def test_refdataset_unspliced_defaults(reference: gvl.Reference): - bed = pl.DataFrame({ - "chrom": ["chr1", "chr1"], - "chromStart": [0, 100], - "chromEnd": [100, 150], - }) + bed = pl.DataFrame( + { + "chrom": ["chr1", "chr1"], + "chromStart": [0, 100], + "chromEnd": [100, 150], + } + ) ds = gvl.RefDataset(reference, bed) assert ds.is_spliced is False assert ds.splice_info is None @@ -852,14 +854,16 @@ def reference() -> gvl.Reference: @pytest.fixture def two_transcript_bed() -> pl.DataFrame: # Two transcripts, both on '+' strand. T1 has 2 exons; T2 has 1 exon. - return pl.DataFrame({ - "chrom": ["chr1", "chr1", "chr1"], - "chromStart": [1000, 2000, 5000], - "chromEnd": [1010, 2010, 5010], - "strand": [1, 1, 1], - "transcript_id": ["T1", "T1", "T2"], - "exon_number": [1, 2, 1], - }) + return pl.DataFrame( + { + "chrom": ["chr1", "chr1", "chr1"], + "chromStart": [1000, 2000, 5000], + "chromEnd": [1010, 2010, 5010], + "strand": [1, 1, 1], + "transcript_id": ["T1", "T1", "T2"], + "exon_number": [1, 2, 1], + } + ) def test_spliced_single_col(reference: gvl.Reference, two_transcript_bed: pl.DataFrame): @@ -940,14 +944,16 @@ Append to `tests/test_ref_ds_splicing.py`: ```python def test_spliced_two_col_reorders_exons(reference: gvl.Reference): # Exons stored out-of-order; exon_number column dictates splice order. - bed = pl.DataFrame({ - "chrom": ["chr1", "chr1"], - "chromStart": [2000, 1000], - "chromEnd": [2010, 1010], - "strand": [1, 1], - "transcript_id": ["T1", "T1"], - "exon_number": [2, 1], - }) + bed = pl.DataFrame( + { + "chrom": ["chr1", "chr1"], + "chromStart": [2000, 1000], + "chromEnd": [2010, 1010], + "strand": [1, 1], + "transcript_id": ["T1", "T1"], + "exon_number": [2, 1], + } + ) ds = gvl.RefDataset(reference, bed, splice_info=("transcript_id", "exon_number")) spliced = ds[0] @@ -960,14 +966,16 @@ def test_spliced_two_col_reorders_exons(reference: gvl.Reference): def test_spliced_mixed_strand(reference: gvl.Reference): # T1 has both exons on '-' strand; rc_neg=True means per-exon RC, then concat. - bed = pl.DataFrame({ - "chrom": ["chr1", "chr1"], - "chromStart": [1000, 2000], - "chromEnd": [1010, 2010], - "strand": [-1, -1], - "transcript_id": ["T1", "T1"], - "exon_number": [1, 2], - }) + bed = pl.DataFrame( + { + "chrom": ["chr1", "chr1"], + "chromStart": [1000, 2000], + "chromEnd": [1010, 2010], + "strand": [-1, -1], + "transcript_id": ["T1", "T1"], + "exon_number": [1, 2], + } + ) ds = gvl.RefDataset(reference, bed, splice_info="transcript_id") spliced = ds[0] diff --git a/docs/superpowers/plans/2026-05-21-splice-exonic-var-filter.md b/docs/superpowers/plans/2026-05-21-splice-exonic-var-filter.md index 1ed63272..bee37c44 100644 --- a/docs/superpowers/plans/2026-05-21-splice-exonic-var-filter.md +++ b/docs/superpowers/plans/2026-05-21-splice-exonic-var-filter.md @@ -224,8 +224,7 @@ def svar_gvl_path(tmp_path): def test_with_settings_var_filter_propagates_to_recon(svar_gvl_path): ref_path = _DATA_DIR / "fasta" / "hg38.fa.bgz" ds = ( - gvl.Dataset - .open(svar_gvl_path, reference=ref_path) + gvl.Dataset.open(svar_gvl_path, reference=ref_path) .with_seqs("haplotypes") .with_settings(var_filter="exonic") ) @@ -240,8 +239,7 @@ def test_with_settings_var_filter_propagates_to_recon(svar_gvl_path): def test_with_settings_var_filter_false_clears_recon(svar_gvl_path): ref_path = _DATA_DIR / "fasta" / "hg38.fa.bgz" ds = ( - gvl.Dataset - .open(svar_gvl_path, reference=ref_path) + gvl.Dataset.open(svar_gvl_path, reference=ref_path) .with_seqs("haplotypes") .with_settings(var_filter="exonic") .with_settings(var_filter=False) @@ -551,8 +549,7 @@ def test_open_vs_with_settings_parity_state(spliced_svar_ds_path): ).with_seqs("haplotypes") ds_b = ( - gvl.Dataset - .open(spliced_svar_ds_path, reference=ref_path) + gvl.Dataset.open(spliced_svar_ds_path, reference=ref_path) .with_seqs("haplotypes") .with_settings( splice_info=("transcript_id", "exon_number"), @@ -578,8 +575,7 @@ def test_open_vs_with_settings_parity_output(spliced_svar_ds_path): ).with_seqs("haplotypes") ds_b = ( - gvl.Dataset - .open(spliced_svar_ds_path, reference=ref_path) + gvl.Dataset.open(spliced_svar_ds_path, reference=ref_path) .with_seqs("haplotypes") .with_settings( splice_info=("transcript_id", "exon_number"), diff --git a/docs/superpowers/plans/2026-05-21-svar-link-replacement.md b/docs/superpowers/plans/2026-05-21-svar-link-replacement.md index cf0a357d..ac0606f4 100644 --- a/docs/superpowers/plans/2026-05-21-svar-link-replacement.md +++ b/docs/superpowers/plans/2026-05-21-svar-link-replacement.md @@ -165,12 +165,14 @@ from pydantic_extra_types.semantic_version import SemanticVersion def test_metadata_version_parses_existing_strings(): - payload = json.dumps({ - "samples": ["s1"], - "contigs": ["1"], - "n_regions": 1, - "version": "0.18.0", - }) + payload = json.dumps( + { + "samples": ["s1"], + "contigs": ["1"], + "n_regions": 1, + "version": "0.18.0", + } + ) m = Metadata.model_validate_json(payload) assert isinstance(m.version, SemanticVersion) assert m.version == SemanticVersion.parse("0.18.0") diff --git a/docs/superpowers/plans/2026-05-22-splice-zero-copy.md b/docs/superpowers/plans/2026-05-22-splice-zero-copy.md index ac7e411f..2739c2c6 100644 --- a/docs/superpowers/plans/2026-05-22-splice-zero-copy.md +++ b/docs/superpowers/plans/2026-05-22-splice-zero-copy.md @@ -345,10 +345,12 @@ def build_splice_plan( cell_lengths = np.repeat(pair_lengths.astype(np.int64), E) # cell_lengths length = n_pairs * E. group_offsets indexes the # *permuted_lengths* array at cell boundaries. - cell_starts = np.concatenate(( - [0], - np.cumsum(cell_lengths, dtype=np.int64), - )) # length n_pairs*E + 1 + cell_starts = np.concatenate( + ( + [0], + np.cumsum(cell_lengths, dtype=np.int64), + ) + ) # length n_pairs*E + 1 # group_offsets[i] = permuted_out_offsets[cell_starts[i]] group_offsets = permuted_out_offsets[cell_starts] @@ -418,14 +420,16 @@ def test_ref_call_with_plan_matches_current_behavior(tmp_path, request): DDIR = Path(request.config.rootpath) / "tests" / "data" ref = gvl.Reference.from_path(DDIR / "fasta" / "hg38.fa.bgz", in_memory=False) - bed = pl.DataFrame({ - "chrom": ["chr1", "chr1", "chr1"], - "chromStart": [1000, 2000, 5000], - "chromEnd": [1010, 2010, 5010], - "strand": [1, 1, 1], - "transcript_id": ["T1", "T1", "T2"], - "exon_number": [1, 2, 1], - }) + bed = pl.DataFrame( + { + "chrom": ["chr1", "chr1", "chr1"], + "chromStart": [1000, 2000, 5000], + "chromEnd": [1010, 2010, 5010], + "strand": [1, 1, 1], + "transcript_id": ["T1", "T1", "T2"], + "exon_number": [1, 2, 1], + } + ) sp_ds = gvl.RefDataset(ref, bed, splice_info="transcript_id") new_path = sp_ds[:] # exercises the production path once Task 3 lands. @@ -433,10 +437,12 @@ def test_ref_call_with_plan_matches_current_behavior(tmp_path, request): # Legacy path replicated inline. plain = gvl.RefDataset(ref, bed) unsp = plain[:] - t1 = np.concatenate([ - np.asarray(unsp[0], dtype="S1"), - np.asarray(unsp[1], dtype="S1"), - ]) + t1 = np.concatenate( + [ + np.asarray(unsp[0], dtype="S1"), + np.asarray(unsp[1], dtype="S1"), + ] + ) t2 = np.asarray(unsp[2], dtype="S1") np.testing.assert_equal(np.asarray(new_path[0], dtype="S1").ravel(), t1) @@ -458,12 +464,14 @@ def test_ref_call_with_plan_writes_grouped_layout(tmp_path, request): DDIR = Path(request.config.rootpath) / "tests" / "data" ref = gvl.Reference.from_path(DDIR / "fasta" / "hg38.fa.bgz", in_memory=False) - bed = pl.DataFrame({ - "chrom": ["chr1", "chr1", "chr1"], - "chromStart": [1000, 2000, 5000], - "chromEnd": [1010, 2010, 5010], - "strand": [1, 1, 1], - }) + bed = pl.DataFrame( + { + "chrom": ["chr1", "chr1", "chr1"], + "chromStart": [1000, 2000, 5000], + "chromEnd": [1010, 2010, 5010], + "strand": [1, 1, 1], + } + ) plain = gvl.RefDataset(ref, bed).with_len("ragged") # Manually drive Ref.__call__ as RefDataset._getitem_spliced will: from genvarloader._dataset._utils import bed_to_regions @@ -1513,8 +1521,7 @@ def test_spliced_tracks_round_trip(multi_exon_ds_path: Path): # invariant that the splice plan path produces sane buffers. try: ds = ( - gvl.Dataset - .open(multi_exon_ds_path, ref_path) + gvl.Dataset.open(multi_exon_ds_path, ref_path) .with_tracks("dummy") .with_settings(splice_info=("transcript_id", "exon_number")) ) @@ -1606,8 +1613,7 @@ def test_haptracks_splicing_raises(multi_exon_ds_path: Path): # If the fixture lacks tracks, skip. try: ds = ( - gvl.Dataset - .open(multi_exon_ds_path, ref_path) + gvl.Dataset.open(multi_exon_ds_path, ref_path) .with_seqs("haplotypes") .with_tracks("dummy") .with_settings(splice_info=("transcript_id", "exon_number")) diff --git a/docs/superpowers/plans/2026-05-24-issue-191-var-fields-loading.md b/docs/superpowers/plans/2026-05-24-issue-191-var-fields-loading.md index 070be4ec..c8c6627e 100644 --- a/docs/superpowers/plans/2026-05-24-issue-191-var-fields-loading.md +++ b/docs/superpowers/plans/2026-05-24-issue-191-var-fields-loading.md @@ -114,8 +114,7 @@ def test_dosage_absent_when_not_requested(svar_with_dosages_ds): The output RaggedVariants must not contain a `dosage` field. """ ds = ( - gvl.Dataset - .open(svar_with_dosages_ds, _REF, rc_neg=False) + gvl.Dataset.open(svar_with_dosages_ds, _REF, rc_neg=False) .with_len("ragged") .with_seqs("variants") .with_settings(var_fields=["alt", "ref", "start"]) @@ -129,8 +128,7 @@ def test_dosage_absent_when_not_requested(svar_with_dosages_ds): def test_dosage_present_when_requested(svar_with_dosages_ds): """Sanity: opting in adds the field.""" ds = ( - gvl.Dataset - .open(svar_with_dosages_ds, _REF, rc_neg=False) + gvl.Dataset.open(svar_with_dosages_ds, _REF, rc_neg=False) .with_len("ragged") .with_seqs("variants") .with_settings(var_fields=["alt", "ref", "start", "dosage"]) diff --git a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase4.md b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase4.md index 6e213299..75ea66e1 100644 --- a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase4.md +++ b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase4.md @@ -189,11 +189,13 @@ Expected (last block of the file): ```python def test_refdataset_unspliced_defaults(reference: gvl.Reference): - bed = pl.DataFrame({ - "chrom": ["chr1", "chr1"], - "chromStart": [0, 100], - "chromEnd": [100, 150], - }) + bed = pl.DataFrame( + { + "chrom": ["chr1", "chr1"], + "chromStart": [0, 100], + "chromEnd": [100, 150], + } + ) ds = gvl.RefDataset(reference, bed) assert ds.is_spliced is False assert ds.splice_info is None diff --git a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-splice.md b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-splice.md index f08d6c6f..13cc808c 100644 --- a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-splice.md +++ b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-splice.md @@ -133,14 +133,16 @@ def reference(ref_fasta) -> gvl.Reference: @pytest.fixture def two_transcript_bed() -> pl.DataFrame: # Two transcripts, both '+' strand. T1 has 2 exons; T2 has 1 exon. - return pl.DataFrame({ - "chrom": ["chr1", "chr1", "chr1"], - "chromStart": [1000, 2000, 5000], - "chromEnd": [1010, 2010, 5010], - "strand": [1, 1, 1], - "transcript_id": ["T1", "T1", "T2"], - "exon_number": [1, 2, 1], - }) + return pl.DataFrame( + { + "chrom": ["chr1", "chr1", "chr1"], + "chromStart": [1000, 2000, 5000], + "chromEnd": [1010, 2010, 5010], + "strand": [1, 1, 1], + "transcript_id": ["T1", "T1", "T2"], + "exon_number": [1, 2, 1], + } + ) def test_with_settings_disable_splice(reference, two_transcript_bed): @@ -225,14 +227,16 @@ def reference(ref_fasta) -> gvl.Reference: @pytest.fixture def two_transcript_bed() -> pl.DataFrame: # Two transcripts, both '+' strand. T1 has 2 exons; T2 has 1 exon. - return pl.DataFrame({ - "chrom": ["chr1", "chr1", "chr1"], - "chromStart": [1000, 2000, 5000], - "chromEnd": [1010, 2010, 5010], - "strand": [1, 1, 1], - "transcript_id": ["T1", "T1", "T2"], - "exon_number": [1, 2, 1], - }) + return pl.DataFrame( + { + "chrom": ["chr1", "chr1", "chr1"], + "chromStart": [1000, 2000, 5000], + "chromEnd": [1010, 2010, 5010], + "strand": [1, 1, 1], + "transcript_id": ["T1", "T1", "T2"], + "exon_number": [1, 2, 1], + } + ) def _as_s1(x) -> np.ndarray: diff --git a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-svar-link.md b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-svar-link.md index a2d023f8..becf7a9c 100644 --- a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-svar-link.md +++ b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-svar-link.md @@ -110,12 +110,14 @@ def test_svar_link_rejects_malformed_fingerprint(): def test_metadata_version_parses_existing_strings(): - payload = json.dumps({ - "samples": ["s1"], - "contigs": ["1"], - "n_regions": 1, - "version": "0.18.0", - }) + payload = json.dumps( + { + "samples": ["s1"], + "contigs": ["1"], + "n_regions": 1, + "version": "0.18.0", + } + ) m = Metadata.model_validate_json(payload) assert isinstance(m.version, SemanticVersion) assert m.version == SemanticVersion.parse("0.18.0") diff --git a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-tracks-broader.md b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-tracks-broader.md index ee97798e..0f29906f 100644 --- a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-tracks-broader.md +++ b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-tracks-broader.md @@ -194,23 +194,27 @@ def test_write_duplicate_track_names_rejected(tmp_path): bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [0], "chromEnd": [100]}) t1 = gvl.Table( "dup", - pl.DataFrame({ - "sample_id": ["s0"], - "chrom": ["chr1"], - "start": [0], - "end": [10], - "value": [1.0], - }), + pl.DataFrame( + { + "sample_id": ["s0"], + "chrom": ["chr1"], + "start": [0], + "end": [10], + "value": [1.0], + } + ), ) t2 = gvl.Table( "dup", - pl.DataFrame({ - "sample_id": ["s0"], - "chrom": ["chr1"], - "start": [50], - "end": [60], - "value": [2.0], - }), + pl.DataFrame( + { + "sample_id": ["s0"], + "chrom": ["chr1"], + "start": [50], + "end": [60], + "value": [2.0], + } + ), ) with pytest.raises(ValueError, match="[Dd]uplicate"): gvl.write(path=tmp_path / "x.gvl", bed=bed, tracks=[t1, t2]) diff --git a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-variants.md b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-variants.md index c8763497..efb60af7 100644 --- a/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-variants.md +++ b/docs/superpowers/plans/2026-05-24-test-suite-overhaul-phase5-variants.md @@ -227,8 +227,7 @@ def test_dosage_absent_when_not_requested(svar_with_dosages_ds, ref_fasta): The output RaggedVariants must not contain a `dosage` field. """ ds = ( - gvl.Dataset - .open(svar_with_dosages_ds, ref_fasta, rc_neg=False) + gvl.Dataset.open(svar_with_dosages_ds, ref_fasta, rc_neg=False) .with_len("ragged") .with_seqs("variants") .with_settings(var_fields=["alt", "ref", "start"]) @@ -242,8 +241,7 @@ def test_dosage_absent_when_not_requested(svar_with_dosages_ds, ref_fasta): def test_dosage_present_when_requested(svar_with_dosages_ds, ref_fasta): """Sanity: opting in adds the field.""" ds = ( - gvl.Dataset - .open(svar_with_dosages_ds, ref_fasta, rc_neg=False) + gvl.Dataset.open(svar_with_dosages_ds, ref_fasta, rc_neg=False) .with_len("ragged") .with_seqs("variants") .with_settings(var_fields=["alt", "ref", "start", "dosage"]) diff --git a/docs/superpowers/plans/2026-05-25-test-coverage-deeper-implementation.md b/docs/superpowers/plans/2026-05-25-test-coverage-deeper-implementation.md index cde29ef7..2d954ec2 100644 --- a/docs/superpowers/plans/2026-05-25-test-coverage-deeper-implementation.md +++ b/docs/superpowers/plans/2026-05-25-test-coverage-deeper-implementation.md @@ -79,8 +79,7 @@ def base_dataset(request, phased_vcf_gvl, phased_pgen_gvl, phased_svar_gvl, ref_ "svar": phased_svar_gvl, }[request.param] return ( - gvl.Dataset - .open(gvl_path, ref_fasta, rc_neg=False) + gvl.Dataset.open(gvl_path, ref_fasta, rc_neg=False) .with_len("ragged") .with_tracks(False) ) @@ -186,8 +185,7 @@ import seqpro as sp def _open_haps(path, ref): return ( - gvl.Dataset - .open(path, ref, rc_neg=False) + gvl.Dataset.open(path, ref, rc_neg=False) .with_tracks(False) .with_seqs("haplotypes") ) @@ -890,11 +888,13 @@ import genvarloader as gvl @pytest.fixture def tiny_bed(): - return pl.DataFrame({ - "chrom": ["chr1"], - "chromStart": [1000], - "chromEnd": [1100], - }) + return pl.DataFrame( + { + "chrom": ["chr1"], + "chromStart": [1000], + "chromEnd": [1100], + } + ) def test_empty_bed_either_succeeds_or_raises_clearly(tmp_path, ref_fasta, source_vcf): @@ -920,11 +920,13 @@ def test_empty_bed_either_succeeds_or_raises_clearly(tmp_path, ref_fasta, source def test_overlapping_bed_regions_succeed_or_raise(tmp_path, ref_fasta, source_vcf): """Overlapping BED regions: should succeed (regions are independent) OR raise a clear error documenting the constraint.""" - overlapping = pl.DataFrame({ - "chrom": ["chr1", "chr1"], - "chromStart": [1000, 1050], - "chromEnd": [1100, 1150], - }) + overlapping = pl.DataFrame( + { + "chrom": ["chr1", "chr1"], + "chromStart": [1000, 1050], + "chromEnd": [1100, 1150], + } + ) out = tmp_path / "overlap.gvl" try: gvl.write( @@ -938,11 +940,13 @@ def test_overlapping_bed_regions_succeed_or_raise(tmp_path, ref_fasta, source_vc def test_bed_with_missing_contig_raises(tmp_path, ref_fasta, source_vcf): """A BED entry on a contig not in the reference must raise.""" - bad_bed = pl.DataFrame({ - "chrom": ["chrZZZ_not_real"], - "chromStart": [0], - "chromEnd": [100], - }) + bad_bed = pl.DataFrame( + { + "chrom": ["chrZZZ_not_real"], + "chromStart": [0], + "chromEnd": [100], + } + ) out = tmp_path / "bad_contig.gvl" with pytest.raises((ValueError, KeyError, RuntimeError)): gvl.write(out_dir=out, bed=bad_bed, variants=source_vcf, reference=ref_fasta) diff --git a/docs/superpowers/plans/2026-05-28-prefetching-dataloader-implementation.md b/docs/superpowers/plans/2026-05-28-prefetching-dataloader-implementation.md index 918e2a57..6010730a 100644 --- a/docs/superpowers/plans/2026-05-28-prefetching-dataloader-implementation.md +++ b/docs/superpowers/plans/2026-05-28-prefetching-dataloader-implementation.md @@ -63,6 +63,7 @@ The helper returns the exact total bytes of REF or ALT allele payloads for the variants selected by each (region, sample, ploid) entry, computed in O(|V|) by differencing the RaggedAlleles offsets array (no payload read). """ + import numpy as np import pytest import genvarloader as gvl @@ -85,6 +86,7 @@ def test_allele_bytes_sum_matches_materialized_alt(ds): ragv = haps._get_variants(idx) # ragv.alt has shape (b, p, ~v, ~length). Sum length across innermost two ragged dims. import awkward as ak + expected = ak.sum(ak.num(ragv.alt, axis=-1), axis=-1).to_numpy().ravel() np.testing.assert_array_equal(got, expected) @@ -98,6 +100,7 @@ def test_allele_bytes_sum_ref(ds): got = haps._allele_bytes_sum(idx, "ref") ragv = haps._get_variants(idx) import awkward as ak + expected = ak.sum(ak.num(ragv.ref, axis=-1), axis=-1).to_numpy().ravel() np.testing.assert_array_equal(got, expected) ``` @@ -128,6 +131,7 @@ def _allele_bytes_sum( r, s = np.unravel_index(idx, self.genotypes.shape[:2]) # type: ignore[no-matching-overload] genos = self.genotypes[r, s] import awkward as ak + genos = ak.to_packed(genos) v_idxs = genos.data @@ -186,6 +190,7 @@ Create `tests/unit/dataset/test_output_bytes_per_instance.py`: Invariant: Dataset._output_bytes_per_instance(r, s) == nbytes of the actual dataset[r, s] output, summed over arrays returned for that instance. """ + import numpy as np import pytest import genvarloader as gvl @@ -201,6 +206,7 @@ def _materialized_nbytes_per_instance(ds, r_arr, s_arr): # Each ndarray/Ragged contributes its data nbytes per instance. For Ragged, # we sum the per-instance data nbytes via the offsets. from seqpro.rag import Ragged + n_inst = len(r_arr) totals = np.zeros(n_inst, dtype=np.int64) for arr in out: @@ -279,7 +285,9 @@ def _output_bytes_per_instance( ds_idx, squeeze, out_reshape = self._idxer.parse_idx(idx) r_idx, s_idx = np.unravel_index(ds_idx, self.full_shape) - seq_kind = self.sequence_type # "reference" | "haplotypes" | "annotated" | "variants" | None + seq_kind = ( + self.sequence_type + ) # "reference" | "haplotypes" | "annotated" | "variants" | None total = np.zeros(len(r_idx), dtype=np.int64) # --- seqs payload --- @@ -392,12 +400,12 @@ Expected: FAIL with `NotImplementedError: annotated branch added in Task 3`. In `_output_bytes_per_instance`, replace `raise NotImplementedError("annotated branch added in Task 3")` with: ```python - # annotated: add ref_coords (int32, length=hap_len_sum) and var_idxs (int32, length=n_variants_sum) - n_vars = self.n_variants(regions, samples) - n_vars_flat = n_vars.reshape(-1, n_vars.shape[-1]).astype(np.int64) - n_vars_sum = n_vars_flat.sum(-1) - total += hap_len_sum * 4 # ref_coords int32 - total += n_vars_sum * 4 # var_idxs int32 +# annotated: add ref_coords (int32, length=hap_len_sum) and var_idxs (int32, length=n_variants_sum) +n_vars = self.n_variants(regions, samples) +n_vars_flat = n_vars.reshape(-1, n_vars.shape[-1]).astype(np.int64) +n_vars_sum = n_vars_flat.sum(-1) +total += hap_len_sum * 4 # ref_coords int32 +total += n_vars_sum * 4 # var_idxs int32 ``` - [ ] **Step 4: Verify pass** @@ -454,8 +462,11 @@ def test_variants_with_ref_exact(): def test_variants_with_info_column_exact(): ds = gvl.get_dummy_dataset().with_seqs("variants").with_tracks(False) - info_cols = [c for c in ds._seqs.available_var_fields - if c not in {"alt", "ref", "ilen", "start", "dosage"}] + info_cols = [ + c + for c in ds._seqs.available_var_fields + if c not in {"alt", "ref", "ilen", "start", "dosage"} + ] if not info_cols: pytest.skip("dummy dataset has no INFO columns") ds = ds.with_settings(var_fields=["alt", "start", "ilen", info_cols[0]]) @@ -476,6 +487,7 @@ def _materialized_nbytes_per_instance(ds, r_arr, s_arr): from seqpro.rag import Ragged import awkward as ak from genvarloader._dataset._rag_variants import RaggedVariants + n_inst = len(r_arr) totals = np.zeros(n_inst, dtype=np.int64) for arr in out: @@ -582,7 +594,11 @@ Append: ```python def test_haplotypes_plus_tracks_exact(): - ds = gvl.get_dummy_dataset().with_seqs("haplotypes").with_settings(deterministic=True) + ds = ( + gvl.get_dummy_dataset() + .with_seqs("haplotypes") + .with_settings(deterministic=True) + ) # Default ds has tracks active; if not, with_tracks(True). if not ds.active_tracks: pytest.skip("dummy dataset has no tracks") @@ -617,21 +633,23 @@ Expected: track tests FAIL with `NotImplementedError: tracks branch added in Tas Replace `if self.active_tracks: raise NotImplementedError(...)` with: ```python - if self.active_tracks: - # Track length per instance equals haplotype length (haplotypes/annotated) - # or region length (reference). Compute that base length. - if seq_kind in ("haplotypes", "annotated"): - base_len = hap_len_sum # already includes ploidy sum - else: - # reference or no-seq: tracks span region length × ploidy if haplotypes - # not active. With reference, tracks have no ploidy axis → just region length. - regions_arr = self._full_regions[r_idx].copy() - regions_arr[:, 1] -= self.jitter - regions_arr[:, 2] += self.jitter - base_len = (regions_arr[:, 2] - regions_arr[:, 1]).astype(np.int64) - for track_name in self.active_tracks: - track_dtype = self._tracks.intervals[track_name].dtype # adjust attr to actual storage - total += base_len * track_dtype.itemsize +if self.active_tracks: + # Track length per instance equals haplotype length (haplotypes/annotated) + # or region length (reference). Compute that base length. + if seq_kind in ("haplotypes", "annotated"): + base_len = hap_len_sum # already includes ploidy sum + else: + # reference or no-seq: tracks span region length × ploidy if haplotypes + # not active. With reference, tracks have no ploidy axis → just region length. + regions_arr = self._full_regions[r_idx].copy() + regions_arr[:, 1] -= self.jitter + regions_arr[:, 2] += self.jitter + base_len = (regions_arr[:, 2] - regions_arr[:, 1]).astype(np.int64) + for track_name in self.active_tracks: + track_dtype = self._tracks.intervals[ + track_name + ].dtype # adjust attr to actual storage + total += base_len * track_dtype.itemsize ``` ⚠ Verify the attribute path to track dtypes (`self._tracks.intervals[track_name].dtype` is a placeholder). Open `_dataset/_tracks.py` and confirm how track dtypes are exposed; adjust this line accordingly. If the dtype is uniform across tracks (typically `float32`), substitute that constant rather than per-track lookup. @@ -667,6 +685,7 @@ Create `tests/unit/test_chunk_planner.py`: ```python """ChunkPlanner unit tests. Pure logic, no Dataset dependency.""" + import numpy as np import pytest from genvarloader._chunked import ChunkPlanner @@ -681,8 +700,11 @@ def test_plan_respects_slot_bytes(): r = flat_idx // 10 s = flat_idx % 10 planner = ChunkPlanner( - r_idx=r, s_idx=s, batch_size=5, - bytes_per_instance=bytes_per_instance, slot_bytes=200, + r_idx=r, + s_idx=s, + batch_size=5, + bytes_per_instance=bytes_per_instance, + slot_bytes=200, ) chunks = list(planner) # Each chunk's total bytes ≤ 200; each chunk is a multiple of batch_size. @@ -699,8 +721,11 @@ def test_plan_single_batch_chunks_when_tight(): bytes_per_instance = np.full((4, 1), 100, dtype=np.int64) flat = np.arange(4) planner = ChunkPlanner( - r_idx=flat, s_idx=np.zeros_like(flat), batch_size=2, - bytes_per_instance=bytes_per_instance, slot_bytes=200, + r_idx=flat, + s_idx=np.zeros_like(flat), + batch_size=2, + bytes_per_instance=bytes_per_instance, + slot_bytes=200, ) chunks = list(planner) assert len(chunks) == 2 # 200 bytes per batch fits exactly one chunk @@ -712,10 +737,15 @@ def test_plan_raises_when_batch_exceeds_slot(): bytes_per_instance = np.full((2, 1), 300, dtype=np.int64) flat = np.arange(2) with pytest.raises(ValueError, match="exceeds slot"): - list(ChunkPlanner( - r_idx=flat, s_idx=np.zeros_like(flat), batch_size=2, - bytes_per_instance=bytes_per_instance, slot_bytes=200, - )) + list( + ChunkPlanner( + r_idx=flat, + s_idx=np.zeros_like(flat), + batch_size=2, + bytes_per_instance=bytes_per_instance, + slot_bytes=200, + ) + ) def test_peak_chunk_bytes_reported(): @@ -724,8 +754,11 @@ def test_peak_chunk_bytes_reported(): r = flat // 2 s = flat % 2 planner = ChunkPlanner( - r_idx=r, s_idx=s, batch_size=2, - bytes_per_instance=bytes_per_instance, slot_bytes=1000, + r_idx=r, + s_idx=s, + batch_size=2, + bytes_per_instance=bytes_per_instance, + slot_bytes=1000, ) chunks = list(planner) # Single chunk of 4 instances, total bytes = 10+20+30+40 = 100. @@ -748,6 +781,7 @@ Create `python/genvarloader/_chunked.py`: ```python """Chunk planner: groups (r, s) pairs into per-slot chunks aligned to mini-batch boundaries.""" + from __future__ import annotations from typing import Iterator @@ -802,13 +836,18 @@ class ChunkPlanner: self._batch_totals = batch_totals self.peak_chunk_bytes: int = 0 - def __iter__(self) -> Iterator[tuple[NDArray[np.integer], NDArray[np.integer], int]]: + def __iter__( + self, + ) -> Iterator[tuple[NDArray[np.integer], NDArray[np.integer], int]]: n_batches = len(self._batch_totals) i = 0 while i < n_batches: running = 0 j = i - while j < n_batches and running + int(self._batch_totals[j]) <= self.slot_bytes: + while ( + j < n_batches + and running + int(self._batch_totals[j]) <= self.slot_bytes + ): running += int(self._batch_totals[j]) j += 1 # j-i batches go into this chunk; at least one (guaranteed by the per-batch check). @@ -849,8 +888,10 @@ def slice_chunk(chunk_output, batch_size: int): is_tuple = isinstance(chunk_output, tuple) arrs = chunk_output if is_tuple else (chunk_output,) - n = len(arrs[0]) if not isinstance(arrs[0], (ak.Array, Ragged)) else ( - arrs[0].shape[0] if isinstance(arrs[0], ak.Array) else arrs[0].shape[0] + n = ( + len(arrs[0]) + if not isinstance(arrs[0], (ak.Array, Ragged)) + else (arrs[0].shape[0] if isinstance(arrs[0], ak.Array) else arrs[0].shape[0]) ) if n is None: raise ValueError("slice_chunk: cannot determine chunk length") @@ -898,6 +939,7 @@ def _compare(a, b): from seqpro.rag import Ragged from genvarloader._types import AnnotatedHaps import awkward as ak + if isinstance(a, tuple): assert isinstance(b, tuple) and len(a) == len(b) for x, y in zip(a, b): @@ -917,7 +959,9 @@ def _compare(a, b): raise AssertionError(f"unsupported {type(a)}") -@pytest.mark.parametrize("seq_kind", ["reference", "haplotypes", "annotated", "variants"]) +@pytest.mark.parametrize( + "seq_kind", ["reference", "haplotypes", "annotated", "variants"] +) def test_slice_chunk_matches_direct(seq_kind): ds = gvl.get_dummy_dataset().with_seqs(seq_kind) if seq_kind in ("haplotypes", "annotated"): @@ -931,7 +975,7 @@ def test_slice_chunk_matches_direct(seq_kind): sliced = list(slice_chunk(chunk, batch_size=n_s)) assert len(sliced) == n_r for i, mini in enumerate(sliced): - direct = ds[r[i * n_s:(i + 1) * n_s], s[i * n_s:(i + 1) * n_s]] + direct = ds[r[i * n_s : (i + 1) * n_s], s[i * n_s : (i + 1) * n_s]] _compare(mini, direct) ``` @@ -965,12 +1009,15 @@ Create `tests/unit/test_buffered_loader.py`: ```python """End-to-end tests for mode='buffered'.""" + import numpy as np import pytest import genvarloader as gvl -@pytest.mark.parametrize("seq_kind", ["reference", "haplotypes", "annotated", "variants"]) +@pytest.mark.parametrize( + "seq_kind", ["reference", "haplotypes", "annotated", "variants"] +) def test_buffered_iter_matches_direct(seq_kind): ds = gvl.get_dummy_dataset().with_seqs(seq_kind).with_tracks(False) if seq_kind in ("haplotypes", "annotated"): @@ -987,8 +1034,10 @@ def test_buffered_iter_matches_direct(seq_kind): ) seen = 0 for batch in loader: - seen += batch_size if not isinstance(batch, tuple) else ( - batch[0].shape[0] if hasattr(batch[0], "shape") else len(batch[0]) + seen += ( + batch_size + if not isinstance(batch, tuple) + else (batch[0].shape[0] if hasattr(batch[0], "shape") else len(batch[0])) ) assert seen == (n_total // batch_size) * batch_size @@ -1006,7 +1055,11 @@ def test_buffered_rejects_oversized_batch(): def test_buffered_rejects_nondeterministic_for_haplotypes(): - ds = gvl.get_dummy_dataset().with_seqs("haplotypes").with_settings(deterministic=False) + ds = ( + gvl.get_dummy_dataset() + .with_seqs("haplotypes") + .with_settings(deterministic=False) + ) with pytest.raises(ValueError, match="deterministic"): ds.to_dataloader(mode="buffered", batch_size=2) ``` @@ -1025,6 +1078,7 @@ Create `python/genvarloader/_buffered_loader.py`: ```python """mode='buffered' dataloader path: synchronous chunked fetch in main process.""" + from __future__ import annotations from typing import TYPE_CHECKING, Any @@ -1052,8 +1106,11 @@ def make_buffered_dataset( self._dataset = dataset self._batch_size = batch_size self._planner = ChunkPlanner( - r_idx=flat_r, s_idx=flat_s, batch_size=batch_size, - bytes_per_instance=bytes_per_instance, slot_bytes=slot_bytes, + r_idx=flat_r, + s_idx=flat_s, + batch_size=batch_size, + bytes_per_instance=bytes_per_instance, + slot_bytes=slot_bytes, ) def __iter__(self): @@ -1086,9 +1143,12 @@ def _resolve_buffered_inputs( ): """Compute flat (r_idx, s_idx) order, bytes_per_instance, and slot_bytes.""" import numpy as np + # 1) Resolve full epoch order from the sampler. if sampler is None: - sampler = get_sampler(len(dataset), batch_size, shuffle, drop_last, generator=generator) + sampler = get_sampler( + len(dataset), batch_size, shuffle, drop_last, generator=generator + ) flat = [] for batch in sampler: flat.extend(batch) @@ -1131,7 +1191,7 @@ def get_dataloader( persistent_workers: bool = False, pin_memory_device: str = "", mode: str | None = None, - buffer_bytes: int = 2 * 1024 ** 3, + buffer_bytes: int = 2 * 1024**3, copy: bool = True, heartbeat_seconds: float = 60.0, ): @@ -1143,7 +1203,9 @@ def get_dataloader( " multithreading which has lower overhead than multiprocessing." ) if sampler is None: - sampler = get_sampler(len(dataset), batch_size, shuffle, drop_last, generator=generator) + sampler = get_sampler( + len(dataset), batch_size, shuffle, drop_last, generator=generator + ) return td.DataLoader( dataset, batch_size=None, @@ -1162,25 +1224,49 @@ def get_dataloader( ) if mode not in {"buffered", "double_buffered"}: - raise ValueError(f"unknown mode={mode!r}; expected None, 'buffered', or 'double_buffered'") + raise ValueError( + f"unknown mode={mode!r}; expected None, 'buffered', or 'double_buffered'" + ) if num_workers > 0: - raise ValueError(f"mode={mode!r} is incompatible with num_workers>0; the loader IS the concurrency strategy") + raise ValueError( + f"mode={mode!r} is incompatible with num_workers>0; the loader IS the concurrency strategy" + ) n_slots = 1 if mode == "buffered" else 2 r_idx, s_idx, bpi, slot_bytes, _sampler = _resolve_buffered_inputs( - dataset, batch_size, shuffle, drop_last, sampler, generator, buffer_bytes, n_slots, + dataset, + batch_size, + shuffle, + drop_last, + sampler, + generator, + buffer_bytes, + n_slots, ) if mode == "buffered": from ._buffered_loader import make_buffered_dataset + inner_ds = make_buffered_dataset( - dataset, batch_size, slot_bytes, bpi, r_idx, s_idx, + dataset, + batch_size, + slot_bytes, + bpi, + r_idx, + s_idx, ) else: from ._double_buffered_loader import make_double_buffered_dataset + inner_ds = make_double_buffered_dataset( - dataset, batch_size, slot_bytes, bpi, r_idx, s_idx, - copy=copy, heartbeat_seconds=heartbeat_seconds, + dataset, + batch_size, + slot_bytes, + bpi, + r_idx, + s_idx, + copy=copy, + heartbeat_seconds=heartbeat_seconds, ) return td.DataLoader( @@ -1226,6 +1312,7 @@ Create `tests/unit/test_shm_layout.py`: ```python """Round-trip tests for the shm slot layout.""" + import multiprocessing as mp import numpy as np import pytest @@ -1322,6 +1409,7 @@ Header layout (little-endian throughout): u64 inner_offsets_offset (used by ragged_alleles; 0 otherwise) u64 inner_offsets_nbytes """ + from __future__ import annotations import struct @@ -1362,8 +1450,9 @@ def write_chunk( for a in arrays: cursor = _align(cursor) payload_offsets.append(cursor) - np.frombuffer(buf, dtype=a.dtype, count=a.size, - offset=cursor).reshape(a.shape)[...] = a + np.frombuffer(buf, dtype=a.dtype, count=a.size, offset=cursor).reshape(a.shape)[ + ... + ] = a cursor += a.nbytes payload_bytes = cursor - HEADER_RESERVED # 2) Write header. @@ -1373,15 +1462,15 @@ def write_chunk( hdr += _DESCRIPTOR_FIXED.pack(0, a.dtype.num, a.ndim) for d in a.shape: hdr += struct.pack(" HEADER_RESERVED: raise ValueError(f"header too large: {len(hdr)} > {HEADER_RESERVED}") - buf[:len(hdr)] = bytes(hdr) + buf[: len(hdr)] = bytes(hdr) return cursor @@ -1398,13 +1487,22 @@ def read_chunk(buf: memoryview) -> tuple[int, list[np.ndarray]]: (dim,) = struct.unpack_from(" 255: raise ValueError("at most 255 arrays per chunk") descriptors = [] @@ -1509,34 +1614,50 @@ def write_chunk(buf, arrays, n_instances): cursor = _align(cursor) if isinstance(a, np.ndarray): data_off = cursor - np.frombuffer(buf, dtype=a.dtype, count=a.size, - offset=data_off).reshape(a.shape)[...] = a + np.frombuffer(buf, dtype=a.dtype, count=a.size, offset=data_off).reshape( + a.shape + )[...] = a cursor += a.nbytes - descriptors.append({ - "kind": 0, "dtype_num": a.dtype.num, "shape": a.shape, - "data_offset": data_off, "data_nbytes": a.nbytes, - "lengths_offset": 0, "lengths_nbytes": 0, - "inner_offsets_offset": 0, "inner_offsets_nbytes": 0, - }) + descriptors.append( + { + "kind": 0, + "dtype_num": a.dtype.num, + "shape": a.shape, + "data_offset": data_off, + "data_nbytes": a.nbytes, + "lengths_offset": 0, + "lengths_nbytes": 0, + "inner_offsets_offset": 0, + "inner_offsets_nbytes": 0, + } + ) elif isinstance(a, Ragged): data_off = cursor data_arr = np.ascontiguousarray(a.data) - np.frombuffer(buf, dtype=data_arr.dtype, count=data_arr.size, - offset=data_off)[...] = data_arr.ravel() + np.frombuffer( + buf, dtype=data_arr.dtype, count=data_arr.size, offset=data_off + )[...] = data_arr.ravel() cursor += data_arr.nbytes cursor = _align(cursor) off_off = cursor off_arr = np.ascontiguousarray(a.offsets) - np.frombuffer(buf, dtype=off_arr.dtype, count=off_arr.size, - offset=off_off)[...] = off_arr + np.frombuffer(buf, dtype=off_arr.dtype, count=off_arr.size, offset=off_off)[ + ... + ] = off_arr cursor += off_arr.nbytes - descriptors.append({ - "kind": 1, "dtype_num": data_arr.dtype.num, - "shape": [data_arr.size], # flat - "data_offset": data_off, "data_nbytes": data_arr.nbytes, - "lengths_offset": off_off, "lengths_nbytes": off_arr.nbytes, - "inner_offsets_offset": 0, "inner_offsets_nbytes": 0, - }) + descriptors.append( + { + "kind": 1, + "dtype_num": data_arr.dtype.num, + "shape": [data_arr.size], # flat + "data_offset": data_off, + "data_nbytes": data_arr.nbytes, + "lengths_offset": off_off, + "lengths_nbytes": off_arr.nbytes, + "inner_offsets_offset": 0, + "inner_offsets_nbytes": 0, + } + ) else: raise TypeError(f"write_chunk: unsupported array type {type(a)}") payload_bytes = cursor - HEADER_RESERVED @@ -1546,13 +1667,18 @@ def write_chunk(buf, arrays, n_instances): hdr += _DESCRIPTOR_FIXED.pack(d["kind"], d["dtype_num"], len(d["shape"])) for dim in d["shape"]: hdr += struct.pack(" HEADER_RESERVED: raise ValueError(f"header too large: {len(hdr)} > {HEADER_RESERVED}") - buf[:len(hdr)] = bytes(hdr) + buf[: len(hdr)] = bytes(hdr) return cursor ``` @@ -1561,6 +1687,7 @@ Extend `read_chunk`: ```python def read_chunk(buf): from seqpro.rag import Ragged + n_inst, payload_bytes, n_arrays = _HEADER_PREAMBLE.unpack_from(buf, 0) cursor = _HEADER_PREAMBLE.size views = [] @@ -1569,19 +1696,21 @@ def read_chunk(buf): cursor += _DESCRIPTOR_FIXED.size shape = [] for _ in range(ndim): - (dim,) = struct.unpack_from(" subprocess.CompletedProcess: def choose_samples() -> list[str]: """Deterministically pick N samples present in both genotypes and bigwigs.""" - psam = pl.read_csv( - PLINK_PREFIX.with_suffix(".psam"), separator="\t" - ) + psam = pl.read_csv(PLINK_PREFIX.with_suffix(".psam"), separator="\t") # plink psam first column is "#IID" or "IID". iid_col = "#IID" if "#IID" in psam.columns else "IID" geno_samples = set(psam[iid_col].to_list()) @@ -358,11 +356,15 @@ def slice_pgen(samples: list[str]) -> Path: run( [ "plink2", - "--pfile", str(PLINK_PREFIX), - "--chr", "chr22", - "--keep", str(keep), + "--pfile", + str(PLINK_PREFIX), + "--chr", + "chr22", + "--keep", + str(keep), "--make-pgen", - "--out", str(out_prefix), + "--out", + str(out_prefix), ] ) keep.unlink() @@ -441,7 +443,7 @@ def build_dataset(samples: list[str], pgen: Path, bed_path: Path) -> Path: if not bw.exists(): raise SystemExit(f"Missing chr22 bigwig for {sample}: {bw}") paths[sample] = str(bw) - assert set(paths) == set(samples), (set(samples) - set(paths)) + assert set(paths) == set(samples), set(samples) - set(paths) tracks = gvl.BigWigs("read-depth", paths) @@ -653,10 +655,20 @@ Expected: no import errors (conftest is imported during collection). No tests co Create a temporary check (do NOT commit) `tests/benchmarks/_smoke.py`: ```python -def test_smoke(captured_haplotypes, captured_diffs, captured_intervals_to_tracks, - captured_realign_tracks, captured_germline_ccfs): - for c in (captured_haplotypes, captured_diffs, captured_intervals_to_tracks, - captured_realign_tracks, captured_germline_ccfs): +def test_smoke( + captured_haplotypes, + captured_diffs, + captured_intervals_to_tracks, + captured_realign_tracks, + captured_germline_ccfs, +): + for c in ( + captured_haplotypes, + captured_diffs, + captured_intervals_to_tracks, + captured_realign_tracks, + captured_germline_ccfs, + ): assert c.args or c.kwargs ``` @@ -731,9 +743,7 @@ def test_reconstruct_haplotypes_from_sparse(benchmark, captured_haplotypes): def test_intervals_to_tracks(benchmark, captured_intervals_to_tracks): - result = _warm_and_run( - benchmark, intervals_to_tracks, captured_intervals_to_tracks - ) + result = _warm_and_run(benchmark, intervals_to_tracks, captured_intervals_to_tracks) # intervals_to_tracks returns a dense track array. assert result is not None @@ -934,7 +944,9 @@ def build(ds, mode: str): def main() -> None: p = argparse.ArgumentParser() - p.add_argument("--mode", choices=["haplotypes", "tracks", "variants"], required=True) + p.add_argument( + "--mode", choices=["haplotypes", "tracks", "variants"], required=True + ) p.add_argument("--n-batches", type=int, default=N_BATCHES) args = p.parse_args() @@ -952,8 +964,10 @@ def main() -> None: regions = [i % n_regions for i in range(n)] samples = [(i // n_regions) % n_samples for i in range(n)] - print(f"mode={args.mode} threads={os.environ['NUMBA_NUM_THREADS']} " - f"batches={args.n_batches} batch={n}") + print( + f"mode={args.mode} threads={os.environ['NUMBA_NUM_THREADS']} " + f"batches={args.n_batches} batch={n}" + ) for i in range(args.n_batches + BURN_IN): _ = ds[regions, samples] print("done") diff --git a/docs/superpowers/plans/2026-05-29-dataloader-bench.md b/docs/superpowers/plans/2026-05-29-dataloader-bench.md index 7f098ef5..c3b15b20 100644 --- a/docs/superpowers/plans/2026-05-29-dataloader-bench.md +++ b/docs/superpowers/plans/2026-05-29-dataloader-bench.md @@ -230,7 +230,8 @@ def test_new_mode_cell_count_is_25_per_mode_output(): for mode in C.MODES_NEW: for output in C.OUTPUTS: cells = [ - c for c in C.enumerate_cells() + c + for c in C.enumerate_cells() if c.mode == mode and c.with_seqs == output ] # 16 factorial + 9 fan (12 raw − 3 shared midpoints) = 25 @@ -240,8 +241,7 @@ def test_new_mode_cell_count_is_25_per_mode_output(): def test_baseline_cell_count_is_15_per_output_and_has_no_buffer(): for output in C.OUTPUTS: cells = [ - c for c in C.enumerate_cells() - if c.mode is None and c.with_seqs == output + c for c in C.enumerate_cells() if c.mode is None and c.with_seqs == output ] # 8 factorial corners + 7 fan (9 raw − 2 shared midpoints) = 15 assert len(cells) == 15, (output, len(cells)) @@ -260,9 +260,12 @@ def test_total_cell_count_is_195_and_all_unique(): def test_baseline_fan_cells_sit_at_midpoints(): # the threads fan for baseline pins region=MID, batch=MID - base = [c for c in C.enumerate_cells() if c.mode is None and c.with_seqs == "variants"] + base = [ + c for c in C.enumerate_cells() if c.mode is None and c.with_seqs == "variants" + ] threads_fan = [ - c for c in base + c + for c in base if c.region_length == C.REGION_MID and c.batch_size == C.BATCH_MID ] assert sorted(c.threads for c in threads_fan) == [2, 4, 16] @@ -270,8 +273,10 @@ def test_baseline_fan_cells_sit_at_midpoints(): def test_new_mode_buffer_fan_pins_other_axes_at_midpoint(): buf_fan = [ - c for c in C.enumerate_cells() - if c.mode == "buffered" and c.with_seqs == "haplotypes" + c + for c in C.enumerate_cells() + if c.mode == "buffered" + and c.with_seqs == "haplotypes" and c.threads == C.THREADS_MID and c.region_length == C.REGION_MID and c.batch_size == C.BATCH_MID @@ -617,17 +622,26 @@ def test_measure_cell_returns_a_complete_row(tmp_path): paths = C.prepare_datasets([1_000], svar, regions, tmp_path) cell = C.Cell( - mode=None, with_seqs="variants", - threads=1, region_length=1_000, batch_size=16, buffer_bytes=None, + mode=None, + with_seqs="variants", + threads=1, + region_length=1_000, + batch_size=16, + buffer_bytes=None, ) # tiny stop conditions so the test is fast row = C.measure_cell( - cell, paths[1_000], ref, min_epochs=1, min_seconds=0.0, hard_cap_s=10.0, + cell, + paths[1_000], + ref, + min_epochs=1, + min_seconds=0.0, + hard_cap_s=10.0, ) for col in C.CSV_COLUMNS: assert col in row, col - assert row["mode"] == "" # None serialized as empty + assert row["mode"] == "" # None serialized as empty assert row["with_seqs"] == "variants" assert row["n_epochs"] >= 1 assert row["instances"] == 100 * 5 * row["n_epochs"] @@ -650,9 +664,23 @@ import resource import time CSV_COLUMNS = [ - "mode", "with_seqs", "threads", "region_length", "batch_size", "buffer_bytes", - "n_epochs", "instances", "bytes", "wall_s", "instances_per_s", "MiB_per_s", - "peak_rss_MiB", "timed_out", "git_sha", "host", "started_at", + "mode", + "with_seqs", + "threads", + "region_length", + "batch_size", + "buffer_bytes", + "n_epochs", + "instances", + "bytes", + "wall_s", + "instances_per_s", + "MiB_per_s", + "peak_rss_MiB", + "timed_out", + "git_sha", + "host", + "started_at", ] @@ -909,8 +937,7 @@ def run_child(n_threads: int) -> None: # dataset paths were written by the parent, keyed by region length ds_paths = { - length: tmp_dir / f"dataset_rL{length}.gvl" - for length in C.REGION_LENGTHS + length: tmp_dir / f"dataset_rL{length}.gvl" for length in C.REGION_LENGTHS } cells = C.cells_for_threads(n_threads) @@ -918,8 +945,12 @@ def run_child(n_threads: int) -> None: ds_path = ds_paths[cell.region_length] try: row = C.measure_cell( - cell, ds_path, REF, - git_sha=git_sha, host=host, started_at=started_at, + cell, + ds_path, + REF, + git_sha=git_sha, + host=host, + started_at=started_at, ) except Exception as e: # noqa: BLE001 - one bad cell must not kill the run print(f"[threads={n_threads}] cell {i}/{len(cells)} FAILED: {cell} -> {e}") @@ -1047,15 +1078,30 @@ OUT_PNG = HERE / "results_plot.png" # column axis -> (csv column, fan values, {pinned other axis: midpoint}) AXES = { - "threads": ("threads", C.THREADS_FAN, - {"region_length": C.REGION_MID, "batch_size": C.BATCH_MID}), - "region_length": ("region_length", C.REGION_FAN, - {"threads": C.THREADS_MID, "batch_size": C.BATCH_MID}), - "batch_size": ("batch_size", C.BATCH_FAN, - {"threads": C.THREADS_MID, "region_length": C.REGION_MID}), - "buffer_bytes": ("buffer_bytes", C.BUFFER_FAN, - {"threads": C.THREADS_MID, "region_length": C.REGION_MID, - "batch_size": C.BATCH_MID}), + "threads": ( + "threads", + C.THREADS_FAN, + {"region_length": C.REGION_MID, "batch_size": C.BATCH_MID}, + ), + "region_length": ( + "region_length", + C.REGION_FAN, + {"threads": C.THREADS_MID, "batch_size": C.BATCH_MID}, + ), + "batch_size": ( + "batch_size", + C.BATCH_FAN, + {"threads": C.THREADS_MID, "region_length": C.REGION_MID}, + ), + "buffer_bytes": ( + "buffer_bytes", + C.BUFFER_FAN, + { + "threads": C.THREADS_MID, + "region_length": C.REGION_MID, + "batch_size": C.BATCH_MID, + }, + ), } MODE_STYLE = { @@ -1092,9 +1138,11 @@ def main() -> None: axis_names = list(AXES) fig, axs = plt.subplots( - len(C.OUTPUTS), len(axis_names), + len(C.OUTPUTS), + len(axis_names), figsize=(4 * len(axis_names), 3 * len(C.OUTPUTS)), - constrained_layout=True, squeeze=False, + constrained_layout=True, + squeeze=False, ) for r, output in enumerate(C.OUTPUTS): @@ -1118,7 +1166,8 @@ def main() -> None: fig.suptitle( "DataLoader throughput: mode comparison across knobs " "(other axes pinned at midpoint)", - fontsize=13, fontweight="bold", + fontsize=13, + fontweight="bold", ) fig.savefig(OUT_PNG, dpi=150) print(f"Saved {OUT_PNG}") diff --git a/docs/superpowers/plans/2026-05-30-dense2sparse-with-length.md b/docs/superpowers/plans/2026-05-30-dense2sparse-with-length.md index bd5bb65a..5325ddca 100644 --- a/docs/superpowers/plans/2026-05-30-dense2sparse-with-length.md +++ b/docs/superpowers/plans/2026-05-30-dense2sparse-with-length.md @@ -229,9 +229,7 @@ def _vcf_region_chunks( ends = df["chromEnd"].to_numpy() # unextended in-range variant indices, split per region v_idx, v_offsets = vcf._var_idxs(contig, starts, ends) - unextended_idxs = np.array_split( - v_idx.astype(V_IDX_TYPE), v_offsets[1:-1] - ) + unextended_idxs = np.array_split(v_idx.astype(V_IDX_TYPE), v_offsets[1:-1]) contig_desc = f"Processing genotypes for {df.height} regions on contig {contig}" first_in_contig = True @@ -392,9 +390,7 @@ def _pgen_region_chunks( region_iter = pgen._chunk_ranges_with_length(contig, starts, ends, max_mem) else: v_idx, v_offsets = pgen.var_idxs(contig, starts, ends) - unextended_idxs = np.array_split( - v_idx.astype(V_IDX_TYPE), v_offsets[1:-1] - ) + unextended_idxs = np.array_split(v_idx.astype(V_IDX_TYPE), v_offsets[1:-1]) region_iter = pgen.chunk_ranges(contig, starts, ends, max_mem) for ri, range_ in enumerate(region_iter): diff --git a/docs/superpowers/plans/2026-05-30-vcfixture-phase1.md b/docs/superpowers/plans/2026-05-30-vcfixture-phase1.md index 05f211f6..51ebad53 100644 --- a/docs/superpowers/plans/2026-05-30-vcfixture-phase1.md +++ b/docs/superpowers/plans/2026-05-30-vcfixture-phase1.md @@ -166,9 +166,10 @@ LOCI = [ def test_reference_has_expected_bases_at_loci(tmp_path: Path): ref = write_synthetic_reference(tmp_path / "synthetic.fa.bgz", seed=0) assert ref.exists() - assert ref.with_suffix(ref.suffix + ".fai").exists() or ( - ref.parent / (ref.name + ".fai") - ).exists() + assert ( + ref.with_suffix(ref.suffix + ".fai").exists() + or (ref.parent / (ref.name + ".fai")).exists() + ) with pysam.FastaFile(str(ref)) as fa: for contig, pos, expected in LOCI: got = fa.fetch(contig, pos - 1, pos - 1 + len(expected)).upper() @@ -406,20 +407,45 @@ def build_source_vcf(reference_path: str | Path) -> "object": # chr20 block — carries INFO (test_sitesonly) and IDs/FILTERs. b.record( - "chr20", 14370, ref="N", alt=["A"], ids=["rs6054257"], qual=29.0, filter=(), - gt=["0|0", "1|0", "1/1"], info={"NS": 3, "DP": 14, "AF": [0.5], "DB": True, "H2": True}, + "chr20", + 14370, + ref="N", + alt=["A"], + ids=["rs6054257"], + qual=29.0, + filter=(), + gt=["0|0", "1|0", "1/1"], + info={"NS": 3, "DP": 14, "AF": [0.5], "DB": True, "H2": True}, ) b.record( - "chr20", 17330, ref="N", alt=["A"], qual=3.0, filter=["q10"], - gt=["0|0", "0|1", "0/0"], info={"NS": 3, "DP": 11, "AF": [0.017]}, + "chr20", + 17330, + ref="N", + alt=["A"], + qual=3.0, + filter=["q10"], + gt=["0|0", "0|1", "0/0"], + info={"NS": 3, "DP": 11, "AF": [0.017]}, ) b.record( - "chr20", 1110696, ref="G", alt=["A", "T"], ids=["rs6040355"], qual=67.0, filter=(), + "chr20", + 1110696, + ref="G", + alt=["A", "T"], + ids=["rs6040355"], + qual=67.0, + filter=(), gt=["1|2", "2|1", "2/2"], info={"NS": 2, "DP": 10, "AF": [0.333, 0.667], "AA": "T", "DB": True}, ) b.record( - "chr20", 1234567, ref="A", alt=["GA", "AC"], ids=["microsat1"], qual=50.0, filter=(), + "chr20", + 1234567, + ref="A", + alt=["GA", "AC"], + ids=["microsat1"], + qual=50.0, + filter=(), gt=["0/1", "0/2", "./."], info={"NS": 3, "DP": 9, "AA": "G", "AN": 6, "AC": [3, 1]}, ) diff --git a/docs/superpowers/plans/2026-05-31-flat-buffer-getitem-pipeline.md b/docs/superpowers/plans/2026-05-31-flat-buffer-getitem-pipeline.md index b450c662..fdbf62bd 100644 --- a/docs/superpowers/plans/2026-05-31-flat-buffer-getitem-pipeline.md +++ b/docs/superpowers/plans/2026-05-31-flat-buffer-getitem-pipeline.md @@ -290,6 +290,7 @@ def test_to_padded_matches_seqpro(): off = np.array([0, 2, 5], np.int64) # rows len 2 and 3 f = _Flat.from_offsets(data, (2, None), off) from genvarloader._ragged import to_padded + expected = to_padded(_rag(data, (2, None), off), -1) np.testing.assert_array_equal(f.to_padded(-1), expected) @@ -423,6 +424,7 @@ Append to `tests/dataset/test_flat.py`: ```python def test_reverse_masked_int_matches_awkward(): import awkward as ak + data = np.arange(10, dtype=np.int32) off = np.array([0, 3, 6, 10], np.int64) # 3 rows mask = np.array([True, False, True]) @@ -437,6 +439,7 @@ def test_reverse_masked_int_matches_awkward(): def test_reverse_masked_dna_matches_existing(): from genvarloader._ragged import reverse_complement_masked, _COMP # noqa + seq = np.frombuffer(b"ACGTAACCGGTT", dtype="S1") off = np.array([0, 4, 12], np.int64) # 2 rows mask = np.array([True, False]) @@ -477,30 +480,32 @@ def _reverse_rows_masked(data, offsets, mask): # pragma: no cover - njit Add the method to `_Flat`: ```python - def reverse_masked(self, mask: NDArray[np.bool_], comp: NDArray | None = None) -> "_Flat": - """Reverse (DNA: reverse-complement) the `mask`-selected rows, in place. - - `mask` is one entry per outer query; replicate across any inner fixed - axes in C order to get one entry per flattened ragged row, matching the - awkward `ak.where` broadcast it replaces. - """ - m = np.ascontiguousarray(mask, np.bool_).reshape(-1) - if m.size != self.n_rows: - factor, rem = divmod(self.n_rows, m.size) - if rem != 0: - raise ValueError( - f"mask has {m.size} entries but {self.n_rows} rows " - "(not an integer multiple)." - ) - m = np.repeat(m, factor) - if comp is not None: - # DNA reverse-complement via the flat seqpro kernel (reuses gvl's LUT). - from ._ragged import reverse_complement_masked - - rag = reverse_complement_masked(self.to_ragged(), m) - return _Flat(np.asarray(rag.data), self.offsets, self.shape) - _reverse_rows_masked(self.data, self.offsets, m) - return self +def reverse_masked( + self, mask: NDArray[np.bool_], comp: NDArray | None = None +) -> "_Flat": + """Reverse (DNA: reverse-complement) the `mask`-selected rows, in place. + + `mask` is one entry per outer query; replicate across any inner fixed + axes in C order to get one entry per flattened ragged row, matching the + awkward `ak.where` broadcast it replaces. + """ + m = np.ascontiguousarray(mask, np.bool_).reshape(-1) + if m.size != self.n_rows: + factor, rem = divmod(self.n_rows, m.size) + if rem != 0: + raise ValueError( + f"mask has {m.size} entries but {self.n_rows} rows " + "(not an integer multiple)." + ) + m = np.repeat(m, factor) + if comp is not None: + # DNA reverse-complement via the flat seqpro kernel (reuses gvl's LUT). + from ._ragged import reverse_complement_masked + + rag = reverse_complement_masked(self.to_ragged(), m) + return _Flat(np.asarray(rag.data), self.offsets, self.shape) + _reverse_rows_masked(self.data, self.offsets, m) + return self ``` > **Implementer note:** `reverse_complement_masked` already does the mask replication; passing the already-replicated `m` is harmless (replication is idempotent when `m.size == n_rows`). If pyrefly complains about the circular import, the local import inside the method is the fix. @@ -543,7 +548,9 @@ class _FlatAnnotatedHaps: def shape(self) -> tuple[int | None, ...]: return self.haps.shape - def reverse_masked(self, mask: NDArray[np.bool_], comp: NDArray) -> "_FlatAnnotatedHaps": + def reverse_masked( + self, mask: NDArray[np.bool_], comp: NDArray + ) -> "_FlatAnnotatedHaps": self.haps = self.haps.reverse_masked(mask, comp=comp) self.var_idxs = self.var_idxs.reverse_masked(mask) self.ref_coords = self.ref_coords.reverse_masked(mask) @@ -558,7 +565,9 @@ class _FlatAnnotatedHaps: def squeeze(self, axis=None) -> "_FlatAnnotatedHaps": return _FlatAnnotatedHaps( - self.haps.squeeze(axis), self.var_idxs.squeeze(axis), self.ref_coords.squeeze(axis) + self.haps.squeeze(axis), + self.var_idxs.squeeze(axis), + self.ref_coords.squeeze(axis), ) def to_ragged(self): @@ -594,8 +603,11 @@ Append to `tests/dataset/test_flat.py`: ```python def test_flat_annotated_to_ragged(): from genvarloader._flat import _Flat, _FlatAnnotatedHaps + off = np.array([0, 2, 4], np.int64) - h = _Flat.from_offsets(np.frombuffer(b"ACGT", "S1").view(np.uint8).copy(), (2, None), off) + h = _Flat.from_offsets( + np.frombuffer(b"ACGT", "S1").view(np.uint8).copy(), (2, None), off + ) v = _Flat.from_offsets(np.array([0, 1, 2, 3], np.int32), (2, None), off) p = _Flat.from_offsets(np.array([10, 11, 12, 13], np.int32), (2, None), off) rah = _FlatAnnotatedHaps(h, v, p).to_ragged() @@ -638,37 +650,41 @@ At the top of `reverse_complement_ragged` (before the `isinstance(rag, Ragged)` In `pad` (`:357`), add before the `isinstance(rag, Ragged)` chain: ```python - if isinstance(rag, (_Flat, _FlatAnnotatedHaps)): - if isinstance(rag, _Flat): - pad_value = b"N" if rag.data.dtype.kind in "SU" else 0 - return rag.view("S1").to_padded(pad_value) if rag.data.dtype == np.uint8 else rag.to_padded(pad_value) - return rag.to_padded() +if isinstance(rag, (_Flat, _FlatAnnotatedHaps)): + if isinstance(rag, _Flat): + pad_value = b"N" if rag.data.dtype.kind in "SU" else 0 + return ( + rag.view("S1").to_padded(pad_value) + if rag.data.dtype == np.uint8 + else rag.to_padded(pad_value) + ) + return rag.to_padded() ``` In `getitem`, replace the densify block (`:94-103`) so it handles flat and routes ragged-output flats to `to_ragged`: ```python - if view.output_length == "variable": - recon = tuple( - r if isinstance(r, (RaggedVariants, RaggedIntervals)) else pad(r) - for r in recon - ) - elif isinstance(view.output_length, int): - recon = tuple( - r if isinstance(r, (RaggedVariants, RaggedIntervals)) - else r.to_fixed(view.output_length) if isinstance(r, (_Flat, _FlatAnnotatedHaps)) - else r.to_numpy() - for r in recon - ) +if view.output_length == "variable": + recon = tuple( + r if isinstance(r, (RaggedVariants, RaggedIntervals)) else pad(r) for r in recon + ) +elif isinstance(view.output_length, int): + recon = tuple( + r + if isinstance(r, (RaggedVariants, RaggedIntervals)) + else r.to_fixed(view.output_length) + if isinstance(r, (_Flat, _FlatAnnotatedHaps)) + else r.to_numpy() + for r in recon + ) ``` Then, immediately before the `out_reshape` step (`:105`), add a final wrap so any still-flat (ragged-output) element becomes its public `Ragged` type: ```python - recon = tuple( - o.to_ragged() if isinstance(o, (_Flat, _FlatAnnotatedHaps)) else o - for o in recon - ) +recon = tuple( + o.to_ragged() if isinstance(o, (_Flat, _FlatAnnotatedHaps)) else o for o in recon +) ``` > **Implementer note:** for ragged output, `reshape`/`squeeze` then run on the resulting `Ragged` (existing behavior). `_Flat.reshape`/`squeeze` exist too, so if you prefer to reshape-then-wrap, that also works — but wrapping first keeps the diff smallest and reuses the proven `Ragged` reshape. The `to_fixed`/`to_padded` for `_Flat[uint8]` must `.view("S1")` first so dtype matches the legacy `to_numpy()` output; the helper methods handle this when called via `_FlatAnnotatedHaps`, but a bare S1 `_Flat` (haplotypes mode) needs the `.view("S1")` — encode that in `_Flat.to_fixed`/`to_padded` by checking `data.dtype == np.uint8`? No — keep `_Flat` dtype-agnostic; instead the reconstructor returns the haps `_Flat` already `.view("S1")` (Task 6). Verify dtype at the boundary with the snapshot gate. @@ -756,30 +772,32 @@ from .._flat import _Flat, _FlatAnnotatedHaps Replace the non-splice body of `_reconstruct_haplotypes` (`:758-784`). Allocate the buffer/offsets directly instead of via `Ragged.from_offsets`, run the kernel, return a `_Flat` viewed as S1: ```python - if req.splice_plan is None: - data = np.empty(req.out_offsets[-1], np.uint8) - shape = (*req.shifts.shape, None) - reconstruct_haplotypes_from_sparse( - geno_offset_idx=req.geno_offset_idx, - out=data, - out_offsets=req.out_offsets, - regions=req.regions, - shifts=req.shifts, - geno_offsets=self.genotypes.offsets, - geno_v_idxs=self.genotypes.data, - v_starts=self.variants.start, - ilens=self.variants.ilen, - alt_alleles=self.variants.alt.data.view(np.uint8), - alt_offsets=self.variants.alt.offsets, - ref=self.reference.reference, - ref_offsets=self.reference.offsets, - pad_char=self.reference.pad_char, - keep=req.keep, - keep_offsets=req.keep_offsets, - annot_v_idxs=None, - annot_ref_pos=None, - ) - return cast("Ragged[np.bytes_]", _Flat.from_offsets(data, shape, req.out_offsets).view("S1")) +if req.splice_plan is None: + data = np.empty(req.out_offsets[-1], np.uint8) + shape = (*req.shifts.shape, None) + reconstruct_haplotypes_from_sparse( + geno_offset_idx=req.geno_offset_idx, + out=data, + out_offsets=req.out_offsets, + regions=req.regions, + shifts=req.shifts, + geno_offsets=self.genotypes.offsets, + geno_v_idxs=self.genotypes.data, + v_starts=self.variants.start, + ilens=self.variants.ilen, + alt_alleles=self.variants.alt.data.view(np.uint8), + alt_offsets=self.variants.alt.offsets, + ref=self.reference.reference, + ref_offsets=self.reference.offsets, + pad_char=self.reference.pad_char, + keep=req.keep, + keep_offsets=req.keep_offsets, + annot_v_idxs=None, + annot_ref_pos=None, + ) + return cast( + "Ragged[np.bytes_]", _Flat.from_offsets(data, shape, req.out_offsets).view("S1") + ) ``` > **Implementer note:** the kernel writes into `data` (uint8); `.view("S1")` makes the `_Flat` dtype `S1` so the boundary's S1/`to_fixed` path matches the old `to_numpy()` dtype, AND the RC branch must then test `data.dtype.kind == "S"` (see Task 4 Step 3 note) — pick one convention and make both consistent. Recommended: keep `_Flat` data as `S1` for haps and branch RC on `dtype.kind == "S"`; update the Task 4 RC check accordingly. Leave the splice path (`:786-824`) on `Ragged` for Task 9. @@ -789,36 +807,39 @@ Replace the non-splice body of `_reconstruct_haplotypes` (`:758-784`). Allocate Replace the non-splice body of `_reconstruct_annotated_haplotypes` (`:837-879`) to allocate three flat buffers, run the kernel, and return a tuple of `_Flat`s (keeping the method's `tuple` return contract): ```python - if req.splice_plan is None: - shape = (*req.shifts.shape, None) - haps = np.empty(req.out_offsets[-1], np.uint8) - annot_v = np.empty(req.out_offsets[-1], V_IDX_TYPE) - annot_pos = np.empty(req.out_offsets[-1], np.int32) - reconstruct_haplotypes_from_sparse( - geno_offset_idx=req.geno_offset_idx, - out=haps, - out_offsets=req.out_offsets, - regions=req.regions, - shifts=req.shifts, - geno_offsets=self.genotypes.offsets, - geno_v_idxs=self.genotypes.data, - v_starts=self.variants.start, - ilens=self.variants.ilen, - alt_alleles=self.variants.alt.data.view(np.uint8), - alt_offsets=self.variants.alt.offsets, - ref=self.reference.reference, - ref_offsets=self.reference.offsets, - pad_char=self.reference.pad_char, - keep=req.keep, - keep_offsets=req.keep_offsets, - annot_v_idxs=annot_v, - annot_ref_pos=annot_pos, - ) - return ( - cast("Ragged[np.bytes_]", _Flat.from_offsets(haps, shape, req.out_offsets).view("S1")), - cast("Ragged", _Flat.from_offsets(annot_v, shape, req.out_offsets)), - cast("Ragged", _Flat.from_offsets(annot_pos, shape, req.out_offsets)), - ) +if req.splice_plan is None: + shape = (*req.shifts.shape, None) + haps = np.empty(req.out_offsets[-1], np.uint8) + annot_v = np.empty(req.out_offsets[-1], V_IDX_TYPE) + annot_pos = np.empty(req.out_offsets[-1], np.int32) + reconstruct_haplotypes_from_sparse( + geno_offset_idx=req.geno_offset_idx, + out=haps, + out_offsets=req.out_offsets, + regions=req.regions, + shifts=req.shifts, + geno_offsets=self.genotypes.offsets, + geno_v_idxs=self.genotypes.data, + v_starts=self.variants.start, + ilens=self.variants.ilen, + alt_alleles=self.variants.alt.data.view(np.uint8), + alt_offsets=self.variants.alt.offsets, + ref=self.reference.reference, + ref_offsets=self.reference.offsets, + pad_char=self.reference.pad_char, + keep=req.keep, + keep_offsets=req.keep_offsets, + annot_v_idxs=annot_v, + annot_ref_pos=annot_pos, + ) + return ( + cast( + "Ragged[np.bytes_]", + _Flat.from_offsets(haps, shape, req.out_offsets).view("S1"), + ), + cast("Ragged", _Flat.from_offsets(annot_v, shape, req.out_offsets)), + cast("Ragged", _Flat.from_offsets(annot_pos, shape, req.out_offsets)), + ) ``` - [ ] **Step 3: Assemble `_FlatAnnotatedHaps` in `get_haps_and_shifts`** @@ -951,7 +972,7 @@ The spliced reconstruction paths (`_reconstruct_haplotypes` splice branch, `_rec Append spliced cases to `tests/dataset/test_flat_getitem_snapshot.py` CASES (regenerate snapshots for the new cases only — delete nothing existing): ```python - ("haps_spliced", dict(seqs="haplotypes"), "ragged"), # with a SpliceIndexer +(("haps_spliced", dict(seqs="haplotypes"), "ragged"),) # with a SpliceIndexer ``` > **Implementer note:** read `tests/dataset/` for an existing spliced test to copy the `subset_to`/splice setup; the snapshot harness's `_build` needs a splice branch. If splicing needs a specific BED/region setup the fixture lacks, add a minimal spliced unit test comparing flat vs a pre-refactor `Ragged` reference computed in the same test (capture by temporarily forcing the legacy path) instead of a committed snapshot. @@ -1047,6 +1068,7 @@ Create `tests/dataset/test_no_awkward_in_hotpath.py` — a regression guard that ```python """Guard: the fixed/ragged getitem hot path must not dispatch awkward kernels.""" + import numpy as np import pytest import genvarloader as gvl @@ -1064,7 +1086,12 @@ def test_tracks_fixed_no_awkward(monkeypatch, dataset_path_with_ref): return orig(*a, **k) monkeypatch.setattr(ak, "to_numpy", counting) - ds = gvl.Dataset.open(*dataset_path_with_ref).with_seqs(None).with_tracks("read-depth").with_len(64) + ds = ( + gvl.Dataset.open(*dataset_path_with_ref) + .with_seqs(None) + .with_tracks("read-depth") + .with_len(64) + ) _ = ds[[0, 1, 2, 3], [0, 0, 0, 0]] assert calls["n"] == 0 ``` diff --git a/docs/superpowers/plans/2026-05-31-gvl-phase2.md b/docs/superpowers/plans/2026-05-31-gvl-phase2.md index 764c1d10..147b9ff3 100644 --- a/docs/superpowers/plans/2026-05-31-gvl-phase2.md +++ b/docs/superpowers/plans/2026-05-31-gvl-phase2.md @@ -115,6 +115,7 @@ Create `tests/_builders/test_case.py`: ```python """Tests for the shared build_case fixture builder.""" + from __future__ import annotations from pathlib import Path @@ -182,6 +183,7 @@ Used by BOTH the property-test module (random reference-consistent draws) and the session conftest fixture / `gen` task (one fixed standardized document), so the generation logic lives in exactly one place. """ + from __future__ import annotations import subprocess @@ -240,9 +242,7 @@ def _bgzip_index(vcf_text: bytes, out_gz: Path) -> Path: return out_gz -def _derive_bed( - vcf_gz: Path, extra_regions: pl.DataFrame | None -) -> pl.DataFrame: +def _derive_bed(vcf_gz: Path, extra_regions: pl.DataFrame | None) -> pl.DataFrame: """Group variant positions into SEQ_LEN-wide regions (Phase-1 logic, keyed off the *normalized* VCF positions). Optionally append manual regions.""" df = pl.read_csv( @@ -290,8 +290,15 @@ def _write_consensus( out_fa = out_dir / f"source_{sample}_nr{row_nr}_h{hap}.fa" _run( [ - "bcftools", "consensus", "-H", str(hap + 1), - "-s", sample, "-o", str(out_fa), str(vcf_gz), + "bcftools", + "consensus", + "-H", + str(hap + 1), + "-s", + sample, + "-o", + str(out_fa), + str(vcf_gz), ], input=seq, ) @@ -350,8 +357,14 @@ def build_case( pgen_path = workdir / "filtered.pgen" _run( [ - "plink2", "--vcf", str(vcf_gz), "--make-pgen", - "--vcf-half-call", "r", "--out", str(pgen_path.with_suffix("")), + "plink2", + "--vcf", + str(vcf_gz), + "--make-pgen", + "--vcf-half-call", + "r", + "--out", + str(pgen_path.with_suffix("")), ] ) @@ -553,22 +566,46 @@ def session_document(spec): # chr2 block (relabeled from chr20) — carries INFO/IDs/FILTERs. b.record( - "chr2", 14370, ref="N", alt=["A"], ids=["rs6054257"], qual=29.0, filter=(), + "chr2", + 14370, + ref="N", + alt=["A"], + ids=["rs6054257"], + qual=29.0, + filter=(), gt=["0|0", "1|0", "1/1"], info={"NS": 3, "DP": 14, "AF": [0.5], "DB": True, "H2": True}, ) b.record( - "chr2", 17330, ref="N", alt=["A"], qual=3.0, filter=["q10"], - gt=["0|0", "0|1", "0/0"], info={"NS": 3, "DP": 11, "AF": [0.017]}, + "chr2", + 17330, + ref="N", + alt=["A"], + qual=3.0, + filter=["q10"], + gt=["0|0", "0|1", "0/0"], + info={"NS": 3, "DP": 11, "AF": [0.017]}, ) b.record( - "chr2", 1110696, ref="G", alt=["A", "T"], ids=["rs6040355"], qual=67.0, - filter=(), gt=["1|2", "2|1", "2/2"], + "chr2", + 1110696, + ref="G", + alt=["A", "T"], + ids=["rs6040355"], + qual=67.0, + filter=(), + gt=["1|2", "2|1", "2/2"], info={"NS": 2, "DP": 10, "AF": [0.333, 0.667], "AA": "T", "DB": True}, ) b.record( - "chr2", 1234567, ref="A", alt=["GA", "AC"], ids=["microsat1"], qual=50.0, - filter=(), gt=["0/1", "0/2", "./."], + "chr2", + 1234567, + ref="A", + alt=["GA", "AC"], + ids=["microsat1"], + qual=50.0, + filter=(), + gt=["0/1", "0/2", "./."], info={"NS": 3, "DP": 9, "AA": "G", "AN": 6, "AC": [3, 1]}, ) return b.build() @@ -610,6 +647,7 @@ property tests) from the same `session_document`. This task only persists the inputs that the self-contained write/track/edge tests and the FASTA-only unit tests consume. """ + from __future__ import annotations import shutil @@ -828,6 +866,7 @@ It currently imports the deleted `_synthetic` module. Repoint it at `case.py`'s ```python """Unit tests for the standardized session reference + source document.""" + from __future__ import annotations import subprocess @@ -879,7 +918,8 @@ def test_reference_n_masks_chr1_telomere(tmp_path: Path): def _norm(vcf_path: Path, ref_path: Path) -> str: out = subprocess.run( ["bcftools", "norm", "-f", str(ref_path), str(vcf_path)], - check=True, capture_output=True, + check=True, + capture_output=True, ) return out.stdout.decode() @@ -1016,7 +1056,7 @@ Expected: PASS (skipped on py3.12 per the existing `pytestmark`; otherwise green Lines 23 + 30 — relabel the self-contained DataFrame (semantically arbitrary): ```python - {"chrom": ["chr1", "chr1"]}, +({"chrom": ["chr1", "chr1"]},) ``` ```python ContigNormalizer(["chr1"]) @@ -1056,6 +1096,7 @@ Each example shells out to bcftools/plink2/samtools, so deadlines are disabled and per-example isolation uses a TemporaryDirectory created inside the test body (NOT a function-scoped fixture — that trips Hypothesis's health check). """ + from __future__ import annotations import sys @@ -1080,10 +1121,16 @@ _ALL_VIOLATIONS = frozenset({"multiallelic", "non_atomic", "non_left_aligned"}) @pytest.mark.parametrize("src", ["vcf", "pgen", "svar"]) @settings(max_examples=25, deadline=None, suppress_health_check=_SUPPRESS) -@given(case_inputs=st.reference_and_documents( - violations=_ALL_VIOLATIONS, max_samples=2, max_records=4, - max_contigs=2, max_contig_len=2000, max_repeats=3, -)) +@given( + case_inputs=st.reference_and_documents( + violations=_ALL_VIOLATIONS, + max_samples=2, + max_records=4, + max_contigs=2, + max_contig_len=2000, + max_repeats=3, + ) +) def test_haplotypes_match_consensus(src, case_inputs): spec, doc, _truth = case_inputs with tempfile.TemporaryDirectory() as tmp: @@ -1102,7 +1149,8 @@ def test_haplotypes_match_consensus(src, case_inputs): with pysam.FastaFile(str(fa_path)) as f: desired = sp.cast_seqs(f.fetch(f.references[0]).upper()) np.testing.assert_equal( - actual, desired, + actual, + desired, f"src={src} region={region} sample={sample} hap={h}", ) ``` @@ -1176,9 +1224,7 @@ def test_track1b_helpers_agree_on_session_case(synthetic_case): truth = case.truth sample_to_idx = {s: i for i, s in enumerate(case.samples)} for region in range(ds.n_regions): - contig, start0, end0 = case.regions.select( - "chrom", "start", "end" - ).row(region) + contig, start0, end0 = case.regions.select("chrom", "start", "end").row(region) expected = _truth_allele_counts(truth, contig, start0, end0) for sample in case.samples: rv = ds[region, sample] @@ -1211,10 +1257,16 @@ def _truth_alt_af(truth, v): @settings(max_examples=15, deadline=None, suppress_health_check=_SUPPRESS) -@given(case_inputs=st.reference_and_documents( - violations=frozenset(), max_samples=2, max_records=4, - max_contigs=2, max_contig_len=2000, max_repeats=3, -)) +@given( + case_inputs=st.reference_and_documents( + violations=frozenset(), + max_samples=2, + max_records=4, + max_contigs=2, + max_contig_len=2000, + max_repeats=3, + ) +) @pytest.mark.parametrize("src", ["vcf", "pgen", "svar"]) def test_allele_calls_match_truth(src, case_inputs): spec, doc, truth = case_inputs @@ -1228,9 +1280,9 @@ def test_allele_calls_match_truth(src, case_inputs): ) sample_to_idx = {s: i for i, s in enumerate(case.samples)} for region in range(ds.n_regions): - contig, start0, end0 = case.regions.select( - "chrom", "start", "end" - ).row(region) + contig, start0, end0 = case.regions.select("chrom", "start", "end").row( + region + ) expected = _truth_allele_counts(truth, contig, start0, end0) for sample in case.samples: got = _gvl_allele_counts(ds[region, sample]) @@ -1299,8 +1351,11 @@ def test_multiallelic_raw_is_rejected(data): import hypothesis.strategies as hyp spec = data - doc = st.documents(reference=spec, violations={"multiallelic"}).example() \ - if False else None + doc = ( + st.documents(reference=spec, violations={"multiallelic"}).example() + if False + else None + ) # Draw the document from the same example's reference deterministically: # use a composite strategy instead (see note below). ``` @@ -1314,8 +1369,11 @@ import hypothesis.strategies as hyp @hyp.composite def _spec_and_violating_doc(draw, violation): spec = draw(st.references(max_contigs=2, max_contig_len=2000, max_repeats=3)) - doc = draw(st.documents(reference=spec, violations={violation}, - max_samples=2, max_records=4)) + doc = draw( + st.documents( + reference=spec, violations={violation}, max_samples=2, max_records=4 + ) + ) return spec, doc diff --git a/docs/superpowers/plans/2026-06-01-flat-variants-path.md b/docs/superpowers/plans/2026-06-01-flat-variants-path.md index bb89dd5a..65c16ffb 100644 --- a/docs/superpowers/plans/2026-06-01-flat-variants-path.md +++ b/docs/superpowers/plans/2026-06-01-flat-variants-path.md @@ -83,7 +83,7 @@ Add this branch in `_flatten_output` **before** the final `isinstance(obj, np.nd Add to the `CASES` list (variants is seqs-only; tracks must be off so the return is a bare `RaggedVariants`, not a tuple): ```python - ("variants_ragged", dict(seqs="variants"), "ragged"), +(("variants_ragged", dict(seqs="variants"), "ragged"),) ``` Confirm `_build` turns tracks off for seqs-only cases (it calls `with_tracks(False)` when `tracks is None`). `with_seqs("variants")` is valid for the fixture (it has a VCF with variants). @@ -173,20 +173,18 @@ def _alt_layout_parts( Replace the layout-surgery block in `_get_alleles` (the `node = alleles.layout` … `RegularArray(...)` … `ak.Array(pvl_content)` lines) so it delegates to the helper. The gather still uses the current `ak.to_packed(... [v_idxs])` for now (Task 2 swaps it): ```python - def _get_alleles( - self, genos: Ragged[V_IDX_TYPE], kind: Literal["alt", "ref"] - ) -> ak.Array: - v_idxs = genos.data - # (b*p*v ~l) packed allele bytes for the selected variants - alleles = ak.to_packed( - cast(RaggedAlleles, getattr(self.variants, kind)[v_idxs]) - ) - return _build_allele_layout( - np.asarray(alleles.data).view(np.uint8), - np.asarray(alleles.offsets), - np.asarray(genos.offsets), - genos.shape[-2], - ) +def _get_alleles( + self, genos: Ragged[V_IDX_TYPE], kind: Literal["alt", "ref"] +) -> ak.Array: + v_idxs = genos.data + # (b*p*v ~l) packed allele bytes for the selected variants + alleles = ak.to_packed(cast(RaggedAlleles, getattr(self.variants, kind)[v_idxs])) + return _build_allele_layout( + np.asarray(alleles.data).view(np.uint8), + np.asarray(alleles.offsets), + np.asarray(genos.offsets), + genos.shape[-2], + ) ``` (`genos.shape[-2]` is ploidy `p`.) @@ -239,11 +237,15 @@ In the AF-filter branch, swap the pack but keep `ak.to_regular` (it is the only In the dosage branch: ```python - if self.dosages is not None and "dosage" in self.var_fields: - dosages = self.dosages[r, s] - if _keep is not None: - dosages = ak.to_regular(dosages[_keep], 1) - fields["dosage"] = Ragged(dosages).to_packed() if not isinstance(dosages, Ragged) else dosages.to_packed() +if self.dosages is not None and "dosage" in self.var_fields: + dosages = self.dosages[r, s] + if _keep is not None: + dosages = ak.to_regular(dosages[_keep], 1) + fields["dosage"] = ( + Ragged(dosages).to_packed() + if not isinstance(dosages, Ragged) + else dosages.to_packed() + ) ``` > Keep it simple: if `self.dosages[r, s]` already returns a `Ragged`, use `dosages.to_packed()`; the `_keep`/`to_regular` sub-branch may yield an awkward Array, so wrap with `Ragged(...)` before `.to_packed()`. Confirm the runtime types and pick the minimal correct form; the snapshot gate is the check. @@ -304,14 +306,20 @@ from genvarloader._dataset._haps import _build_allele_layout def _make_rv(alt_rows, ref_rows, starts, group_off, ploidy): """alt_rows/ref_rows: list[bytes] per variant; group_off: variant boundaries per (b*p) row.""" + def alleles(rows): data = np.frombuffer(b"".join(rows), np.uint8) off = np.concatenate([[0], np.cumsum([len(r) for r in rows])]).astype(np.int64) return _build_allele_layout(data, off, np.asarray(group_off, np.int64), ploidy) + alt = alleles(alt_rows) ref = alleles(ref_rows) n = len(starts) - start = Ragged.from_offsets(np.asarray(starts, np.int32), (len(group_off) - 1, None), np.asarray(group_off, np.int64)) + start = Ragged.from_offsets( + np.asarray(starts, np.int32), + (len(group_off) - 1, None), + np.asarray(group_off, np.int64), + ) return RaggedVariants(alt=alt, start=start, ref=ref) @@ -322,15 +330,20 @@ def _ref_rc(rv, to_rc): return alt, ref -@pytest.mark.parametrize("mask", [ - np.array([True, True]), # all - np.array([False, False]), # none (early return) - np.array([True, False]), # mixed -]) +@pytest.mark.parametrize( + "mask", + [ + np.array([True, True]), # all + np.array([False, False]), # none (early return) + np.array([True, False]), # mixed + ], +) def test_rc_matches_awkward(mask): # b=2, p=1, group_off over 2 rows: row0 has 2 variants, row1 has 1 group_off = [0, 2, 3] - rv = _make_rv([b"ACG", b"T", b"GG"], [b"A", b"CC", b"T"], [1, 5, 9], group_off, ploidy=1) + rv = _make_rv( + [b"ACG", b"T", b"GG"], [b"A", b"CC", b"T"], [1, 5, 9], group_off, ploidy=1 + ) exp_alt, exp_ref = _ref_rc(rv, mask) rv.rc_(mask) np.testing.assert_array_equal(ak.to_list(rv["alt"]), ak.to_list(exp_alt)) @@ -339,7 +352,9 @@ def test_rc_matches_awkward(mask): def test_rc_none_means_all(): group_off = [0, 2, 3] - rv = _make_rv([b"ACG", b"T", b"GG"], [b"A", b"CC", b"T"], [1, 5, 9], group_off, ploidy=1) + rv = _make_rv( + [b"ACG", b"T", b"GG"], [b"A", b"CC", b"T"], [1, 5, 9], group_off, ploidy=1 + ) exp_alt, exp_ref = _ref_rc(rv, np.ones(2, bool)) rv.rc_(None) np.testing.assert_array_equal(ak.to_list(rv["alt"]), ak.to_list(exp_alt)) @@ -357,31 +372,29 @@ Expected: FAIL — the current `rc_` uses the eager awkward path; tests fail onl In `python/genvarloader/_dataset/_rag_variants.py`, replace the body of `rc_` (keep the signature, docstring, and the `to_rc is None` / `not to_rc.any()` guards): ```python - def rc_(self, to_rc: NDArray[np.bool_] | None = None) -> Self: - if to_rc is None: - to_rc = np.ones(self.shape[0], np.bool_) - elif not to_rc.any(): - return self - - from .._ragged import _COMP, reverse_complement_masked - from seqpro.rag import Ragged - from ._haps import _alt_layout_parts - - for field in ("alt", "ref"): - if field not in self.fields: - continue - arr = self[field] - leaf, allele_off, group_off, ploidy = _alt_layout_parts(arr) - # per-allele mask: to_rc is per-batch; broadcast across ploidy then variants - per_bp = np.repeat(np.ascontiguousarray(to_rc, np.bool_), ploidy) - per_allele = np.repeat(per_bp, np.diff(group_off)) - view = Ragged.from_offsets( - leaf.view("S1"), (per_allele.size, None), allele_off - ) - # in-place: mutates `leaf`, which shares memory with `arr`'s buffer - reverse_complement_masked(view, per_allele) - +def rc_(self, to_rc: NDArray[np.bool_] | None = None) -> Self: + if to_rc is None: + to_rc = np.ones(self.shape[0], np.bool_) + elif not to_rc.any(): return self + + from .._ragged import _COMP, reverse_complement_masked + from seqpro.rag import Ragged + from ._haps import _alt_layout_parts + + for field in ("alt", "ref"): + if field not in self.fields: + continue + arr = self[field] + leaf, allele_off, group_off, ploidy = _alt_layout_parts(arr) + # per-allele mask: to_rc is per-batch; broadcast across ploidy then variants + per_bp = np.repeat(np.ascontiguousarray(to_rc, np.bool_), ploidy) + per_allele = np.repeat(per_bp, np.diff(group_off)) + view = Ragged.from_offsets(leaf.view("S1"), (per_allele.size, None), allele_off) + # in-place: mutates `leaf`, which shares memory with `arr`'s buffer + reverse_complement_masked(view, per_allele) + + return self ``` > The leaf buffer shares memory with the `ak.Array` (verified), so mutating `view` (which wraps `leaf`) reverse-complements `self["alt"]`/`self["ref"]` in place — preserving `rc_`'s in-place contract without `ak.where`/`ak.to_packed`. Confirm `NDArray`/`Self` are imported in this module (they are used in the existing signature). @@ -420,12 +433,16 @@ Append to `tests/dataset/test_flat_variants.py`: ```python def test_to_packed_matches_awkward_contiguous(): group_off = [0, 2, 3] - rv = _make_rv([b"ACG", b"T", b"GG"], [b"A", b"CC", b"T"], [1, 5, 9], group_off, ploidy=1) + rv = _make_rv( + [b"ACG", b"T", b"GG"], [b"A", b"CC", b"T"], [1, 5, 9], group_off, ploidy=1 + ) exp = ak.to_packed(ak.Array(rv)) # old behavior got = rv.to_packed() assert ak.to_list(got["alt"]) == ak.to_list(exp["alt"]) assert ak.to_list(got["ref"]) == ak.to_list(exp["ref"]) - np.testing.assert_array_equal(np.asarray(got["start"].data), np.asarray(exp["start"].data)) + np.testing.assert_array_equal( + np.asarray(got["start"].data), np.asarray(exp["start"].data) + ) def test_to_packed_matches_awkward_sliced(): @@ -435,14 +452,17 @@ def test_to_packed_matches_awkward_sliced(): [b"ACG", b"T", b"GG", b"AA", b"C"], [b"A", b"CC", b"T", b"G", b"TT"], [1, 5, 9, 12, 20], - group_off, ploidy=1, + group_off, + ploidy=1, ) - sliced = rv[1:] # drop the first (b,p) row + sliced = rv[1:] # drop the first (b,p) row exp = ak.to_packed(ak.Array(sliced)) got = sliced.to_packed() assert ak.to_list(got["alt"]) == ak.to_list(exp["alt"]) assert ak.to_list(got["ref"]) == ak.to_list(exp["ref"]) - np.testing.assert_array_equal(np.asarray(got["start"].data), np.asarray(exp["start"].data)) + np.testing.assert_array_equal( + np.asarray(got["start"].data), np.asarray(exp["start"].data) + ) ``` - [ ] **Step 2: Run to verify (guards the refactor)** @@ -455,30 +475,34 @@ Expected: PASS against the current `ak.to_packed(self)`; this pins behavior so S Replace `RaggedVariants.to_packed`: ```python - def to_packed(self) -> Self: - from seqpro.rag import Ragged - from ._haps import _alt_layout_parts, _build_allele_layout - - packed = {} - for field in self.fields: - arr = self[field] - if field in ("alt", "ref"): - leaf, allele_off, group_off, ploidy = _alt_layout_parts(arr) - # pack the allele (byte) level: contiguates bytes, zero-bases allele_off - allele_lvl = Ragged.from_offsets( - leaf.view("S1"), (allele_off.size - 1, None), allele_off - ).to_packed() - # group_off may be non-zero-based (sliced view) -> rebase - rebased_group = np.asarray(group_off, np.int64) - int(group_off[0]) - packed[field] = _build_allele_layout( - np.asarray(allele_lvl.data).view(np.uint8), - np.asarray(allele_lvl.offsets), - rebased_group, - ploidy, - ) - else: - packed[field] = Ragged(arr).to_packed() if not isinstance(arr, Ragged) else arr.to_packed() - return type(self)(**packed) +def to_packed(self) -> Self: + from seqpro.rag import Ragged + from ._haps import _alt_layout_parts, _build_allele_layout + + packed = {} + for field in self.fields: + arr = self[field] + if field in ("alt", "ref"): + leaf, allele_off, group_off, ploidy = _alt_layout_parts(arr) + # pack the allele (byte) level: contiguates bytes, zero-bases allele_off + allele_lvl = Ragged.from_offsets( + leaf.view("S1"), (allele_off.size - 1, None), allele_off + ).to_packed() + # group_off may be non-zero-based (sliced view) -> rebase + rebased_group = np.asarray(group_off, np.int64) - int(group_off[0]) + packed[field] = _build_allele_layout( + np.asarray(allele_lvl.data).view(np.uint8), + np.asarray(allele_lvl.offsets), + rebased_group, + ploidy, + ) + else: + packed[field] = ( + Ragged(arr).to_packed() + if not isinstance(arr, Ragged) + else arr.to_packed() + ) + return type(self)(**packed) ``` > `type(self)(**packed)` reuses `RaggedVariants.__init__`, which `ak.zip`s the packed fields (the one remaining, documented awkward call — cheap layout wrap). For the allele level, `allele_off.size - 1` is the number of variants; `to_packed` reorders by the existing allele order, which is `(b,p,variant)` row-major — matching `ak.to_packed`'s canonical order. Confirm at runtime whether `self[field]` for numeric fields returns a seqpro `Ragged` (has `.to_packed()`) or a bare `ak.Array`; the `isinstance(arr, Ragged)` guard handles both, and the byte-identity test is the check. @@ -516,14 +540,18 @@ Read the existing `guard_dataset` fixture + `_install_ak_counters` in `tests/dat def test_variants_ragged_minimal_awkward(monkeypatch, guard_dataset): """Variants gather + rc_ + to_packed must dispatch no awkward kernels. ak.zip (record construction) is the documented remaining awkward and is NOT patched here.""" - calls = _install_ak_counters(monkeypatch) # patches to_numpy/to_packed/flatten/where + calls = _install_ak_counters( + monkeypatch + ) # patches to_numpy/to_packed/flatten/where ds = guard_dataset.with_seqs("variants").with_tracks(False) regions = list(range(min(4, ds.shape[0]))) samples = [i % ds.shape[1] for i in range(len(regions))] rv = ds[regions, samples] - rv.rc_(np.ones(len(regions), np.bool_)) # exercise rc_ explicitly - rv.to_packed() # exercise field-wise to_packed - assert calls["n"] == 0, "variants gather/rc_/to_packed dispatched awkward (to_packed/where/flatten/to_numpy)" + rv.rc_(np.ones(len(regions), np.bool_)) # exercise rc_ explicitly + rv.to_packed() # exercise field-wise to_packed + assert calls["n"] == 0, ( + "variants gather/rc_/to_packed dispatched awkward (to_packed/where/flatten/to_numpy)" + ) ``` > If `ds[regions, samples]` for variants returns a tuple (tracks somehow on), use `with_tracks(False)` (already applied) and unwrap. If the guard fails because a path still calls a patched function, investigate — it reveals a missed swap; report DONE_WITH_CONCERNS rather than weakening the assert. Note: `ak.zip` is intentionally not in the patched set. diff --git a/docs/superpowers/plans/2026-06-03-robust-ondisk-artifacts.md b/docs/superpowers/plans/2026-06-03-robust-ondisk-artifacts.md index 30756e49..b6b56072 100644 --- a/docs/superpowers/plans/2026-06-03-robust-ondisk-artifacts.md +++ b/docs/superpowers/plans/2026-06-03-robust-ondisk-artifacts.md @@ -452,9 +452,7 @@ def _ensure_built(source_fa: Path, gvlfa_dir: Path) -> FastaCache: except Exception: pass # unreadable/corrupt -> fall through and rebuild _build_into(source_fa, tmp, gvlfa_dir) - return FastaCache.model_validate_json( - (gvlfa_dir / METADATA_FILENAME).read_text() - ) + return FastaCache.model_validate_json((gvlfa_dir / METADATA_FILENAME).read_text()) ``` In `_ensure_from_fasta` (current lines 216-246) replace the two `meta = build(source_fa, gvlfa_dir)` calls (the rebuild inside `if not valid` and the final fallback) and the migrate path with calls that go through the lock. Concretely: @@ -929,7 +927,9 @@ def test_concurrent_ensure_cache_no_corruption(tmp_path, ref_fasta): expected = np.array(np.memmap(single_data, np.uint8, "r")) # N concurrent builders against the same source - procs = [_CTX.Process(target=_build_cache_worker, args=(str(src),)) for _ in range(6)] + procs = [ + _CTX.Process(target=_build_cache_worker, args=(str(src),)) for _ in range(6) + ] for p in procs: p.start() for p in procs: @@ -972,9 +972,7 @@ def test_concurrent_gvl_write_one_valid_dataset(tmp_path, phased_vcf_gvl, refere dest = tmp_path / "shared.gvl" procs = [ - _CTX.Process( - target=_write_worker, args=(str(dest), str(src_vcf), bed_rows) - ) + _CTX.Process(target=_write_worker, args=(str(dest), str(src_vcf), bed_rows)) for _ in range(4) ] for p in procs: diff --git a/docs/superpowers/plans/2026-06-05-to_dataloader-drop_last.md b/docs/superpowers/plans/2026-06-05-to_dataloader-drop_last.md index 71c406d2..3d70cf41 100644 --- a/docs/superpowers/plans/2026-06-05-to_dataloader-drop_last.md +++ b/docs/superpowers/plans/2026-06-05-to_dataloader-drop_last.md @@ -86,30 +86,28 @@ Expected: FAIL — `ValueError: len(r_idx)=7 is not a multiple of batch_size=3. In `python/genvarloader/_chunked.py`, replace the body of `__init__` from the `n = len(r_idx)` line through the `batch_totals = per_inst.reshape(-1, batch_size).sum(-1)` line (currently lines 37-50): ```python - if len(r_idx) != len(s_idx): - raise ValueError("r_idx and s_idx must have the same length") - n = len(r_idx) - self.r_idx = np.asarray(r_idx) - self.s_idx = np.asarray(s_idx) - self.batch_size = batch_size - self.bytes_per_instance = bytes_per_instance - self.slot_bytes = int(slot_bytes) - self._n = n - - # Per-instance byte cost in epoch order, grouped into mini-batches. The - # final batch may be partial (drop_last=False); its bytes are summed into - # a trailing batch_totals entry so chunk packing and peak-byte sizing - # account for it like any other batch. - per_inst = bytes_per_instance[self.r_idx, self.s_idx].astype(np.int64) - n_full = n // batch_size - full_totals = ( - per_inst[: n_full * batch_size].reshape(n_full, batch_size).sum(-1) - ) - remainder = per_inst[n_full * batch_size :] - if remainder.size: - batch_totals = np.concatenate([full_totals, remainder.sum(keepdims=True)]) - else: - batch_totals = full_totals +if len(r_idx) != len(s_idx): + raise ValueError("r_idx and s_idx must have the same length") +n = len(r_idx) +self.r_idx = np.asarray(r_idx) +self.s_idx = np.asarray(s_idx) +self.batch_size = batch_size +self.bytes_per_instance = bytes_per_instance +self.slot_bytes = int(slot_bytes) +self._n = n + +# Per-instance byte cost in epoch order, grouped into mini-batches. The +# final batch may be partial (drop_last=False); its bytes are summed into +# a trailing batch_totals entry so chunk packing and peak-byte sizing +# account for it like any other batch. +per_inst = bytes_per_instance[self.r_idx, self.s_idx].astype(np.int64) +n_full = n // batch_size +full_totals = per_inst[: n_full * batch_size].reshape(n_full, batch_size).sum(-1) +remainder = per_inst[n_full * batch_size :] +if remainder.size: + batch_totals = np.concatenate([full_totals, remainder.sum(keepdims=True)]) +else: + batch_totals = full_totals ``` Note: the `too_big` check and `_compute_peak_chunk_bytes()` that follow already read `batch_totals`, so they work unchanged with the partial entry included. @@ -254,9 +252,7 @@ def test_buffered_modes_respect_drop_last(small_gvl_ds, mode, drop_last): bs = next((c for c in range(2, N) if N % c), 1) assert N % bs != 0, "need an indivisible batch_size to exercise drop_last" - dl = ds.to_dataloader( - batch_size=bs, shuffle=False, drop_last=drop_last, mode=mode - ) + dl = ds.to_dataloader(batch_size=bs, shuffle=False, drop_last=drop_last, mode=mode) batches = list(dl) expected = N // bs if drop_last else math.ceil(N / bs) assert len(batches) == expected diff --git a/docs/superpowers/plans/2026-06-07-open-variants-no-reference.md b/docs/superpowers/plans/2026-06-07-open-variants-no-reference.md index aa5e9865..c7e6e894 100644 --- a/docs/superpowers/plans/2026-06-07-open-variants-no-reference.md +++ b/docs/superpowers/plans/2026-06-07-open-variants-no-reference.md @@ -16,6 +16,7 @@ The crash, confirmed by repro: ```python import genvarloader as gvl + gvl.Dataset.open("tests/data/phased_dataset.vcf.gvl") # ValueError: Cannot return RaggedSeqs: no reference genome was provided. ``` diff --git a/docs/superpowers/plans/2026-06-07-ragged-variants-pack-lazy-views.md b/docs/superpowers/plans/2026-06-07-ragged-variants-pack-lazy-views.md index 9bd1cc3d..bacaa0fa 100644 --- a/docs/superpowers/plans/2026-06-07-ragged-variants-pack-lazy-views.md +++ b/docs/superpowers/plans/2026-06-07-ragged-variants-pack-lazy-views.md @@ -40,9 +40,10 @@ class TestIndexedLayouts: field = ak.zip({"x": r}, depth_limit=1)[perm]["x"] from awkward.contents import IndexedArray + assert isinstance(field.layout, IndexedArray) - rag = Ragged(field) # used to raise: Expected 1 ragged dimension, got 0 + rag = Ragged(field) # used to raise: Expected 1 ragged dimension, got 0 # accessors that route through unbox() must all work assert rag.offsets is not None assert rag.data is not None @@ -59,9 +60,10 @@ class TestIndexedLayouts: rec = ak.zip({"a": r, "b": r}, depth_limit=1)[np.array([2, 0, 1])] from awkward.contents import IndexedArray + assert isinstance(rec.layout, IndexedArray) - rag = Ragged(rec) # record-layout Ragged over an indexed layout + rag = Ragged(rec) # record-layout Ragged over an indexed layout # offsets extraction (via _extract_list_offsets) must not crash assert rag.offsets is not None assert ak.to_list(rag["a"]) == ak.to_list(rec["a"]) @@ -129,31 +131,38 @@ In `unbox`, extend the loop condition and add an Indexed* branch that projects a without recording a dimension: ```python - while isinstance( - node, - (ListArray, ListOffsetArray, RegularArray, RecordArray, IndexedArray, IndexedOptionArray), - ): - if isinstance(node, (IndexedArray, IndexedOptionArray)): - node = node.project() - continue - if isinstance(node, RecordArray): - raise ValueError( # noqa: TRY004 - "Must extract a single field before unboxing a Ragged array of records." - ) - elif isinstance(node, RegularArray): - shape.append(node.size) +while isinstance( + node, + ( + ListArray, + ListOffsetArray, + RegularArray, + RecordArray, + IndexedArray, + IndexedOptionArray, + ), +): + if isinstance(node, (IndexedArray, IndexedOptionArray)): + node = node.project() + continue + if isinstance(node, RecordArray): + raise ValueError( # noqa: TRY004 + "Must extract a single field before unboxing a Ragged array of records." + ) + elif isinstance(node, RegularArray): + shape.append(node.size) + else: + shape.append(None) + n_ragged += 1 + if isinstance(node, ListOffsetArray): + offsets = node.offsets.data else: - shape.append(None) - n_ragged += 1 - if isinstance(node, ListOffsetArray): - offsets = node.offsets.data - else: - offsets = np.stack( # pyrefly: ignore[no-matching-overload] # awkward .data is ArrayLike, not _ArrayLike - [node.starts.data, node.stops.data], # type: ignore - 0, - ) - - node = node.content + offsets = np.stack( # pyrefly: ignore[no-matching-overload] # awkward .data is ArrayLike, not _ArrayLike + [node.starts.data, node.stops.data], # type: ignore + 0, + ) + + node = node.content ``` - [ ] **Step 4: Update the `unbox` docstring (zero-copy caveat)** @@ -234,7 +243,9 @@ def test_to_packed_numeric_fields_reorder_after_fancy_index(): fancy = RaggedVariants.from_ak(rv[np.array([2, 0])]) got = fancy.to_packed() exp = ak.to_packed(ak.Array(fancy)) - np.testing.assert_array_equal(np.asarray(got["start"].data), np.asarray(exp["start"].data)) + np.testing.assert_array_equal( + np.asarray(got["start"].data), np.asarray(exp["start"].data) + ) assert ak.to_list(got["start"]) == ak.to_list(exp["start"]) ``` @@ -281,12 +292,19 @@ def test_pack_alleles_kernel_identity_and_reorder(): leaf = np.frombuffer(b"ACGTGG", np.uint8) allele_starts = np.array([0, 3, 4], np.int64) allele_stops = np.array([3, 4, 6], np.int64) - var_starts = np.array([0, 2], np.int64) # row0 -> alleles[0:2], row1 -> alleles[2:3] + var_starts = np.array( + [0, 2], np.int64 + ) # row0 -> alleles[0:2], row1 -> alleles[2:3] var_stops = np.array([2, 3], np.int64) # identity order packed, allele_off, group_off = _pack_alleles( - np.array([0, 1], np.int64), var_starts, var_stops, allele_starts, allele_stops, leaf + np.array([0, 1], np.int64), + var_starts, + var_stops, + allele_starts, + allele_stops, + leaf, ) assert bytes(packed) == b"ACGTGG" assert allele_off.tolist() == [0, 3, 4, 6] @@ -294,7 +312,12 @@ def test_pack_alleles_kernel_identity_and_reorder(): # reversed row order packed, allele_off, group_off = _pack_alleles( - np.array([1, 0], np.int64), var_starts, var_stops, allele_starts, allele_stops, leaf + np.array([1, 0], np.int64), + var_starts, + var_stops, + allele_starts, + allele_stops, + leaf, ) assert bytes(packed) == b"GGACGT" assert allele_off.tolist() == [0, 2, 5, 6] @@ -314,18 +337,22 @@ def test_decompose_alleles_reversed(): from genvarloader._dataset._rag_variants import _decompose_alleles, _pack_alleles rv = _make_rv( - [b"A", b"C", b"G", b"T", b"N"], [b"a", b"c", b"g", b"t", b"n"], - [1, 2, 3, 4, 5], [0, 2, 3, 5], ploidy=1, + [b"A", b"C", b"G", b"T", b"N"], + [b"a", b"c", b"g", b"t", b"n"], + [1, 2, 3, 4, 5], + [0, 2, 3, 5], + ploidy=1, ) fancy = RaggedVariants.from_ak(rv[np.array([2, 0])]) - row_src, var_starts, var_stops, allele_starts, allele_stops, leaf, ploidy = _decompose_alleles( - fancy["alt"] + row_src, var_starts, var_stops, allele_starts, allele_stops, leaf, ploidy = ( + _decompose_alleles(fancy["alt"]) ) assert ploidy == 1 packed, allele_off, group_off = _pack_alleles( row_src, var_starts, var_stops, allele_starts, allele_stops, leaf ) from genvarloader._dataset._haps import _build_allele_layout + rebuilt = _build_allele_layout(packed, allele_off, group_off, ploidy) assert ak.to_list(rebuilt) == ak.to_list(fancy["alt"]) ``` @@ -392,9 +419,7 @@ def _decompose_alleles(arr: ak.Array): reg = lay if not isinstance(reg, RegularArray): - raise ValueError( - f"Unsupported allele layout for packing: {arr.layout.form}" - ) + raise ValueError(f"Unsupported allele layout for packing: {arr.layout.form}") ploidy = int(reg.size) var_node = reg.content @@ -410,9 +435,9 @@ def _decompose_alleles(arr: ak.Array): n_out_rows = len(reg) * ploidy row_src = np.arange(n_out_rows, dtype=np.int64) else: - row_src = ( - index[:, None] * ploidy + np.arange(ploidy, dtype=np.int64) - ).reshape(-1) + row_src = (index[:, None] * ploidy + np.arange(ploidy, dtype=np.int64)).reshape( + -1 + ) return row_src, var_starts, var_stops, allele_starts, allele_stops, leaf, ploidy @@ -479,8 +504,11 @@ Append to `tests/dataset/test_flat_variants.py`: def test_to_packed_alt_ref_on_lazy_views(transform): group_off = [0, 2, 3, 5] rv = _make_rv( - [b"ACG", b"T", b"GG", b"AA", b"C"], [b"A", b"CC", b"T", b"G", b"TT"], - [1, 5, 9, 12, 20], group_off, ploidy=1, + [b"ACG", b"T", b"GG", b"AA", b"C"], + [b"A", b"CC", b"T", b"G", b"TT"], + [1, 5, 9, 12, 20], + group_off, + ploidy=1, ) view = rv[::-1] if transform == "reverse" else rv[np.array([2, 0, 3, 1])] view = RaggedVariants.from_ak(view) @@ -488,7 +516,9 @@ def test_to_packed_alt_ref_on_lazy_views(transform): exp = ak.to_packed(ak.Array(view)) assert ak.to_list(got["alt"]) == ak.to_list(exp["alt"]) assert ak.to_list(got["ref"]) == ak.to_list(exp["ref"]) - np.testing.assert_array_equal(np.asarray(got["start"].data), np.asarray(exp["start"].data)) + np.testing.assert_array_equal( + np.asarray(got["start"].data), np.asarray(exp["start"].data) + ) def test_to_packed_explicit_listarray_variant_level(): @@ -497,16 +527,26 @@ def test_to_packed_explicit_listarray_variant_level(): from awkward.index import Index def listarray_alleles(joined_bytes, allele_off, starts, stops): - leaf = NumpyArray(np.frombuffer(joined_bytes, np.uint8), parameters={"__array__": "byte"}) + leaf = NumpyArray( + np.frombuffer(joined_bytes, np.uint8), parameters={"__array__": "byte"} + ) allele = ListOffsetArray( - Index(np.asarray(allele_off, np.int64)), leaf, parameters={"__array__": "bytestring"} + Index(np.asarray(allele_off, np.int64)), + leaf, + parameters={"__array__": "bytestring"}, + ) + var = ListArray( + Index(np.asarray(starts, np.int64)), + Index(np.asarray(stops, np.int64)), + allele, ) - var = ListArray(Index(np.asarray(starts, np.int64)), Index(np.asarray(stops, np.int64)), allele) return ak.Array(RegularArray(var, 1)) alt = listarray_alleles(b"ACGTGG", [0, 3, 4, 6], [0, 2], [2, 3]) ref = listarray_alleles(b"ACCT", [0, 1, 3, 4], [0, 2], [2, 3]) - start = Ragged.from_offsets(np.array([1, 5, 9], np.int32), (2, None), np.array([0, 2, 3], np.int64)) + start = Ragged.from_offsets( + np.array([1, 5, 9], np.int32), (2, None), np.array([0, 2, 3], np.int64) + ) rv = RaggedVariants(alt=alt, start=start, ref=ref) got = rv.to_packed() @@ -526,39 +566,44 @@ In `to_packed`, replace the body of the `if field in ("alt", "ref"):` branch (th `_alt_layout_parts` and does the g0/a0 rebase, ~lines 213–237) with: ```python - if field in ("alt", "ref"): - if _is_canonical_alleles(arr.layout): - # fast path (unchanged): canonical (possibly sliced) layout - leaf, allele_off, group_off, ploidy = _alt_layout_parts(arr) - g0 = int(group_off[0]) - rebased_group = np.asarray(group_off, np.int64) - g0 - a0 = int(allele_off[g0]) - sliced_allele_off = np.asarray(allele_off[g0:], np.int64) - a0 - sliced_leaf = leaf[a0:] - allele_lvl = Ragged.from_offsets( - sliced_leaf.view("S1"), - (sliced_allele_off.size - 1, None), - sliced_allele_off, - ).to_packed() - packed[field] = _build_allele_layout( - np.asarray(allele_lvl.data).view(np.uint8), - np.asarray(allele_lvl.offsets), - rebased_group, - ploidy, - ) - else: - # non-canonical (IndexedArray/ListArray from slicing/reorder): - # numba gather, no ak.to_packed / awkward gather primitives. - ( - row_src, var_starts, var_stops, - allele_starts, allele_stops, leaf, ploidy, - ) = _decompose_alleles(arr) - packed_bytes, allele_off, group_off = _pack_alleles( - row_src, var_starts, var_stops, allele_starts, allele_stops, leaf - ) - packed[field] = _build_allele_layout( - packed_bytes, allele_off, group_off, ploidy - ) +if field in ("alt", "ref"): + if _is_canonical_alleles(arr.layout): + # fast path (unchanged): canonical (possibly sliced) layout + leaf, allele_off, group_off, ploidy = _alt_layout_parts(arr) + g0 = int(group_off[0]) + rebased_group = np.asarray(group_off, np.int64) - g0 + a0 = int(allele_off[g0]) + sliced_allele_off = np.asarray(allele_off[g0:], np.int64) - a0 + sliced_leaf = leaf[a0:] + allele_lvl = Ragged.from_offsets( + sliced_leaf.view("S1"), + (sliced_allele_off.size - 1, None), + sliced_allele_off, + ).to_packed() + packed[field] = _build_allele_layout( + np.asarray(allele_lvl.data).view(np.uint8), + np.asarray(allele_lvl.offsets), + rebased_group, + ploidy, + ) + else: + # non-canonical (IndexedArray/ListArray from slicing/reorder): + # numba gather, no ak.to_packed / awkward gather primitives. + ( + row_src, + var_starts, + var_stops, + allele_starts, + allele_stops, + leaf, + ploidy, + ) = _decompose_alleles(arr) + packed_bytes, allele_off, group_off = _pack_alleles( + row_src, var_starts, var_stops, allele_starts, allele_stops, leaf + ) + packed[field] = _build_allele_layout( + packed_bytes, allele_off, group_off, ploidy + ) ``` > The fast-path block is the existing code verbatim, now guarded by `_is_canonical_alleles`. The @@ -594,15 +639,20 @@ Append to `tests/dataset/test_flat_variants.py`: def test_rc_on_lazy_views_matches_reference(transform): group_off = [0, 2, 3, 5] rv = _make_rv( - [b"ACG", b"T", b"GG", b"AA", b"C"], [b"A", b"CC", b"T", b"G", b"TT"], - [1, 5, 9, 12, 20], group_off, ploidy=1, + [b"ACG", b"T", b"GG", b"AA", b"C"], + [b"A", b"CC", b"T", b"G", b"TT"], + [1, 5, 9, 12, 20], + group_off, + ploidy=1, ) view = rv[::-1] if transform == "reverse" else rv[np.array([2, 0, 3, 1])] view = RaggedVariants.from_ak(view) n = view.shape[0] mask = np.ones(n, np.bool_) - exp_alt, exp_ref = _ref_rc(view, mask) # independent awkward reference (top of file) + exp_alt, exp_ref = _ref_rc( + view, mask + ) # independent awkward reference (top of file) out = view.rc_(mask) assert ak.to_list(out["alt"]) == ak.to_list(exp_alt) @@ -664,13 +714,15 @@ def test_to_packed_ploidy2_reordered(): alt = _build_allele_layout( np.frombuffer(b"ACGTGGA", np.uint8), np.array([0, 2, 3, 4, 6, 7], np.int64), - group_off, ploidy=2, + group_off, + ploidy=2, ) # ref alleles: ["a","c","g","t","n"] -> b"acgtn" ref = _build_allele_layout( np.frombuffer(b"acgtn", np.uint8), np.array([0, 1, 2, 3, 4, 5], np.int64), - group_off, ploidy=2, + group_off, + ploidy=2, ) start = Ragged.from_offsets( np.array([1, 2, 3, 4, 5], np.int32), (2, 2, None), group_off diff --git a/docs/superpowers/plans/2026-06-13-flat-output-mode-C-flank-tokenize.md b/docs/superpowers/plans/2026-06-13-flat-output-mode-C-flank-tokenize.md index 5543c504..34115e0d 100644 --- a/docs/superpowers/plans/2026-06-13-flat-output-mode-C-flank-tokenize.md +++ b/docs/superpowers/plans/2026-06-13-flat-output-mode-C-flank-tokenize.md @@ -144,9 +144,8 @@ import numpy as np def test_with_settings_stores_flank_config(snap_dataset): # snap_dataset is the session phased VCF+reference dataset (see test_flat_getitem_snapshot.py) - ds = ( - snap_dataset.with_seqs("variants") - .with_settings(flank_length=5, token_alphabet=b"ACGT", unknown_token=4) + ds = snap_dataset.with_seqs("variants").with_settings( + flank_length=5, token_alphabet=b"ACGT", unknown_token=4 ) haps = ds._seqs assert haps.flank_length == 5 @@ -184,40 +183,32 @@ In `python/genvarloader/_dataset/_haps.py`, add to the `Haps` dataclass field bl In `python/genvarloader/_dataset/_impl.py`, extend the `with_settings` signature (after `var_filter`): ```python - flank_length: int | None = None, - token_alphabet: bytes | None = None, - unknown_token: int | None = None, +flank_length: int | None = (None,) +token_alphabet: bytes | None = (None,) +unknown_token: int | None = (None,) ``` Add to the docstring Parameters and insert this block before the `# If any source state changed` rebuild (around line 348): ```python - if ( - flank_length is not None - or token_alphabet is not None - or unknown_token is not None - ): - if not isinstance(self._seqs, Haps): - raise ValueError( - "Flank settings require a dataset with genotypes (variants)." - ) - haps = to_evolve.get("_seqs", self._seqs) - new_flank_len = haps.flank_length if flank_length is None else flank_length - lut, lut_dtype = haps.token_lut, haps.token_dtype - if token_alphabet is not None or unknown_token is not None: - if token_alphabet is None or unknown_token is None: - raise ValueError( - "token_alphabet and unknown_token must be set together." - ) - from ._flat_flanks import build_token_lut - - lut, lut_dtype = build_token_lut(token_alphabet, unknown_token) - to_evolve["_seqs"] = replace( - haps, - flank_length=new_flank_len, - token_lut=lut, - token_dtype=lut_dtype, - ) +if flank_length is not None or token_alphabet is not None or unknown_token is not None: + if not isinstance(self._seqs, Haps): + raise ValueError("Flank settings require a dataset with genotypes (variants).") + haps = to_evolve.get("_seqs", self._seqs) + new_flank_len = haps.flank_length if flank_length is None else flank_length + lut, lut_dtype = haps.token_lut, haps.token_dtype + if token_alphabet is not None or unknown_token is not None: + if token_alphabet is None or unknown_token is None: + raise ValueError("token_alphabet and unknown_token must be set together.") + from ._flat_flanks import build_token_lut + + lut, lut_dtype = build_token_lut(token_alphabet, unknown_token) + to_evolve["_seqs"] = replace( + haps, + flank_length=new_flank_len, + token_lut=lut, + token_dtype=lut_dtype, + ) ``` - [ ] **Step 5: Run test to verify it passes** @@ -281,9 +272,11 @@ def test_compute_flank_tokens_matches_oracle(snap_dataset): ```python def test_compute_flank_tokens_unit(snap_dataset): - haps = snap_dataset.with_seqs("variants").with_settings( - flank_length=3, token_alphabet=b"ACGT", unknown_token=4 - )._seqs + haps = ( + snap_dataset.with_seqs("variants") + .with_settings(flank_length=3, token_alphabet=b"ACGT", unknown_token=4) + ._seqs + ) ref = haps.reference lut = haps.token_lut # one (b=1, ploidy=1) group with two variants @@ -309,14 +302,15 @@ Expected: FAIL with `ImportError: cannot import name 'compute_flank_tokens'`. ```python # add to python/genvarloader/_dataset/_flat_flanks.py + def compute_flank_tokens( reference, - v_contigs: NDArray[np.integer], # (n_var,) contig id per variant - starts: NDArray[np.integer], # (n_var,) - ilens: NDArray[np.integer], # (n_var,) + v_contigs: NDArray[np.integer], # (n_var,) contig id per variant + starts: NDArray[np.integer], # (n_var,) + ilens: NDArray[np.integer], # (n_var,) flank_len: int, lut: NDArray, - row_offsets: NDArray[np.int64], # (b*p + 1,) per-(instance,ploid) variant offsets + row_offsets: NDArray[np.int64], # (b*p + 1,) per-(instance,ploid) variant offsets ) -> tuple[NDArray, NDArray[np.int64]]: """Ride-along flank tokens: ``[flank5 | flank3]`` (2*flank_len tokens) per variant. Returns ``(token_data, offsets)`` where ``token_data`` is flat @@ -367,8 +361,8 @@ from genvarloader._dataset._flat_variants import _FlatWindow, _FlatVariantWindow def test_flat_window_to_ragged_roundtrip(): # two groups (b*p=2), variant counts [2, 1]; window lens [3,4 | 2] token_data = np.arange(3 + 4 + 2, dtype=np.uint8) - seq_offsets = np.array([0, 3, 7, 9], dtype=np.int64) # per-variant - var_offsets = np.array([0, 2, 3], dtype=np.int64) # per group + seq_offsets = np.array([0, 3, 7, 9], dtype=np.int64) # per-variant + var_offsets = np.array([0, 2, 3], dtype=np.int64) # per group shape = (2, 1, None, None) w = _FlatWindow(token_data, seq_offsets, var_offsets, shape) rag = w.to_ragged() @@ -391,9 +385,9 @@ class _FlatWindow: """Two-level flat token buffer for ref/alt windows, shape (b, p, ~v, ~win). Mirrors _FlatAlleles but `data` holds tokens (configured dtype), not bytes.""" - data: NDArray # tokens (uint8 or int32), flat - seq_offsets: NDArray[np.int64] # per-variant window offsets, n_variants + 1 - var_offsets: NDArray[np.int64] # per (instance, ploid) offsets, b*p + 1 + data: NDArray # tokens (uint8 or int32), flat + seq_offsets: NDArray[np.int64] # per-variant window offsets, n_variants + 1 + var_offsets: NDArray[np.int64] # per (instance, ploid) offsets, b*p + 1 shape: tuple[int | None, ...] def to_ragged(self): @@ -405,8 +399,9 @@ class _FlatWindow: def reshape(self, shape) -> "_FlatWindow": if isinstance(shape, int): shape = (shape,) - return _FlatWindow(self.data, self.seq_offsets, self.var_offsets, - (*shape, None, None)) + return _FlatWindow( + self.data, self.seq_offsets, self.var_offsets, (*shape, None, None) + ) def squeeze(self, axis: int | None = None) -> "_FlatWindow": # delegate to the same outer-dim squeeze logic as _FlatAlleles @@ -418,7 +413,7 @@ class _FlatVariantWindows: """Window-mode variants output: scalar fields + ref/alt token windows. Raw alleles are intentionally absent (folded into the windows).""" - fields: dict[str, Any] # start / ilen / dosage / info -> _Flat + fields: dict[str, Any] # start / ilen / dosage / info -> _Flat ref_window: _FlatWindow alt_window: _FlatWindow @@ -476,8 +471,9 @@ rtk git commit -m "feat(flat): _FlatWindow + _FlatVariantWindows two-level token from genvarloader._dataset._flat_flanks import compute_windows -def _oracle_windows(reference, v_contigs, starts, ilens, alt_data, alt_seq_off, - flank_len, lut): +def _oracle_windows( + reference, v_contigs, starts, ilens, alt_data, alt_seq_off, flank_len, lut +): ends = starts - np.minimum(ilens, 0) + 1 # ref_window: single contiguous read [start-L, end+L) rw = reference.fetch(v_contigs, starts - flank_len, ends + flank_len) @@ -489,16 +485,18 @@ def _oracle_windows(reference, v_contigs, starts, ilens, alt_data, alt_seq_off, f3 = f3.reshape(len(starts), flank_len) alt_rows, alt_lens = [], np.diff(alt_seq_off) for i in range(len(starts)): - a = alt_data[alt_seq_off[i]:alt_seq_off[i + 1]] + a = alt_data[alt_seq_off[i] : alt_seq_off[i + 1]] alt_rows.append(np.concatenate([f5[i], a, f3[i]])) alt_tok = lut[np.concatenate(alt_rows)] if alt_rows else np.empty(0, lut.dtype) return ref_tok, np.asarray(rw.offsets), alt_tok def test_compute_windows_unit(snap_dataset): - haps = snap_dataset.with_seqs("variants").with_settings( - flank_length=3, token_alphabet=b"ACGT", unknown_token=4 - )._seqs + haps = ( + snap_dataset.with_seqs("variants") + .with_settings(flank_length=3, token_alphabet=b"ACGT", unknown_token=4) + ._seqs + ) ref, lut = haps.reference, haps.token_lut v_contigs = np.array([0, 0], dtype=np.int32) starts = np.array([10, 20], dtype=np.int32) @@ -558,9 +556,14 @@ def _assemble_alt_windows(f5, f3, alt_data, alt_seq_off, flank_len): def compute_windows( reference, - v_contigs, starts, ilens, - alt_data, alt_seq_off, - flank_len, lut, row_offsets, + v_contigs, + starts, + ilens, + alt_data, + alt_seq_off, + flank_len, + lut, + row_offsets, ) -> tuple["_FlatWindow", "_FlatWindow"]: """ref_window = tokenized [start-L, end+L) (single contiguous read); alt_window = tokenized flank5 . alt . flank3 (assembly).""" @@ -570,18 +573,26 @@ def compute_windows( rw = reference.fetch(v_contigs, starts - flank_len, ends + flank_len) ref_tok = lut[rw.data.view(np.uint8)] - ref_window = _FlatWindow(ref_tok, np.asarray(rw.offsets, np.int64), - np.asarray(row_offsets, np.int64), (None,)) # shape set by caller + ref_window = _FlatWindow( + ref_tok, + np.asarray(rw.offsets, np.int64), + np.asarray(row_offsets, np.int64), + (None,), + ) # shape set by caller f5 = reference.fetch(v_contigs, starts - flank_len, starts).data.view(np.uint8) f3 = reference.fetch(v_contigs, ends, ends + flank_len).data.view(np.uint8) alt_bytes, alt_off = _assemble_alt_windows( - np.ascontiguousarray(f5), np.ascontiguousarray(f3), - np.asarray(alt_data, np.uint8), np.asarray(alt_seq_off, np.int64), flank_len, + np.ascontiguousarray(f5), + np.ascontiguousarray(f3), + np.asarray(alt_data, np.uint8), + np.asarray(alt_seq_off, np.int64), + flank_len, ) alt_tok = lut[alt_bytes] - alt_window = _FlatWindow(alt_tok, alt_off, - np.asarray(row_offsets, np.int64), (None,)) + alt_window = _FlatWindow( + alt_tok, alt_off, np.asarray(row_offsets, np.int64), (None,) + ) return ref_window, alt_window ``` @@ -623,14 +634,14 @@ def test_flank_tokens_end_to_end_matches_oracle(snap_dataset): flat = ds[[0, 1, 2], [0, 1, 2]] assert flat.flank_tokens is not None # oracle: build from the SAME RaggedVariants the ragged path returns - rag = ( - snap_dataset.with_seqs("variants")[[0, 1, 2], [0, 1, 2]] - ) + rag = snap_dataset.with_seqs("variants")[[0, 1, 2], [0, 1, 2]] # flank tokens are (b, p, ~v, 2L); compare against an independent fetch+tokenize # over rag.start / rag.ilen. (Helper mirrors genvarformer _read_flank_seq.) expected = _oracle_from_ragged(snap_dataset, rag, flank_len=5) np.testing.assert_array_equal( - np.asarray(flat.flank_tokens.to_ragged().data).view(flat.flank_tokens.data.dtype), + np.asarray(flat.flank_tokens.to_ragged().data).view( + flat.flank_tokens.data.dtype + ), expected, ) ``` @@ -640,6 +651,7 @@ Add the `_oracle_from_ragged` helper near the top of the test file: ```python def _oracle_from_ragged(dataset, rag, flank_len): import seqpro as sp + ref = dataset._seqs.reference lut, _ = build_token_lut(b"ACGT", 4) # per-variant contig: repeat region contig by ploidy then by variant counts @@ -648,8 +660,16 @@ def _oracle_from_ragged(dataset, rag, flank_len): ilens = np.asarray(rag.ilen.data) contigs = _per_variant_contigs(dataset, rag) # see helper below ends = starts - np.minimum(ilens, 0) + 1 - f5 = ref.fetch(contigs, starts - flank_len, starts).data.view(np.uint8).reshape(-1, flank_len) - f3 = ref.fetch(contigs, ends, ends + flank_len).data.view(np.uint8).reshape(-1, flank_len) + f5 = ( + ref.fetch(contigs, starts - flank_len, starts) + .data.view(np.uint8) + .reshape(-1, flank_len) + ) + f3 = ( + ref.fetch(contigs, ends, ends + flank_len) + .data.view(np.uint8) + .reshape(-1, flank_len) + ) return lut[np.concatenate([f5, f3], axis=1)].reshape(-1) ``` @@ -697,13 +717,20 @@ def get_variants_flat(haps: "Haps", idx, regions=None): ploidy = genotypes.shape[-2] starts_v = np.asarray(haps.variants.start)[v_idxs] ilens_v = np.asarray(haps.variants.ilen)[v_idxs] - group_contigs = np.repeat(regions[:, 0], ploidy) # (b*p,) + group_contigs = np.repeat(regions[:, 0], ploidy) # (b*p,) v_contigs = np.repeat(group_contigs, np.diff(row_offsets)) # (n_var,) if issubclass(haps.kind, _FlatVariantWindows): ref_w, alt_w = compute_windows( - haps.reference, v_contigs, starts_v, ilens_v, - alt_data, alt_seq_off, L, haps.token_lut, row_offsets, + haps.reference, + v_contigs, + starts_v, + ilens_v, + alt_data, + alt_seq_off, + L, + haps.token_lut, + row_offsets, ) wshape = (b, ploidy, None, None) ref_w.shape = wshape @@ -713,7 +740,13 @@ def get_variants_flat(haps: "Haps", idx, regions=None): return _FlatVariantWindows(wfields, ref_w, alt_w) tok, off = compute_flank_tokens( - haps.reference, v_contigs, starts_v, ilens_v, L, haps.token_lut, row_offsets, + haps.reference, + v_contigs, + starts_v, + ilens_v, + L, + haps.token_lut, + row_offsets, ) flat.flank_tokens = _Flat.from_offsets(tok, (b, ploidy, None, 2 * L), off) @@ -777,6 +810,7 @@ def test_variant_windows_kind_end_to_end(snap_dataset): def test_variant_windows_requires_flank_settings(snap_dataset): import pytest + with pytest.raises(ValueError, match="flank"): snap_dataset.with_seqs("variant-windows") # no flank_length set ``` @@ -857,6 +891,7 @@ rtk git commit -m "feat(flat): variant-windows kind, dispatch, validation, bound ```python def test_public_exports(): import genvarloader as gvl + assert hasattr(gvl, "FlatVariantWindows") assert hasattr(gvl, "FlatVariants") # from A ``` @@ -871,7 +906,9 @@ Expected: FAIL with `AttributeError`. In `python/genvarloader/_dataset/_flat_variants.py`: ```python -FlatVariantWindows = _FlatVariantWindows # public alias (keep underscored alias working) +FlatVariantWindows = ( + _FlatVariantWindows # public alias (keep underscored alias working) +) ``` In `python/genvarloader/__init__.py`, import and add `"FlatVariantWindows"` to `__all__` (and confirm `FlatVariants`/`FlatAlleles`/`FlatRagged` from A are present). @@ -901,14 +938,20 @@ rtk git commit -m "feat(flat): export FlatVariantWindows" import pytest -@pytest.mark.parametrize("idx", [ - (0, 0), # scalar - ([0, 1, 2], [0, 1, 2]), # list / 2-D (region, sample) - ([0, 0], [0, 1]), # same region, two samples (dedup-relevant) -]) +@pytest.mark.parametrize( + "idx", + [ + (0, 0), # scalar + ([0, 1, 2], [0, 1, 2]), # list / 2-D (region, sample) + ([0, 0], [0, 1]), # same region, two samples (dedup-relevant) + ], +) def test_flank_tokens_index_matrix(snap_dataset, idx): - ds = (snap_dataset.with_seqs("variants").with_output_format("flat") - .with_settings(flank_length=5, token_alphabet=b"ACGT", unknown_token=4)) + ds = ( + snap_dataset.with_seqs("variants") + .with_output_format("flat") + .with_settings(flank_length=5, token_alphabet=b"ACGT", unknown_token=4) + ) flat = ds[idx] rag = snap_dataset.with_seqs("variants")[idx] expected = _oracle_from_ragged(snap_dataset, rag, flank_len=5) @@ -918,8 +961,11 @@ def test_flank_tokens_index_matrix(snap_dataset, idx): def test_oob_flank_at_contig_start(snap_dataset): # region whose variant sits within flank_len of position 0 -> N -> unknown_token - ds = (snap_dataset.with_seqs("variants").with_output_format("flat") - .with_settings(flank_length=50, token_alphabet=b"ACGT", unknown_token=4)) + ds = ( + snap_dataset.with_seqs("variants") + .with_output_format("flat") + .with_settings(flank_length=50, token_alphabet=b"ACGT", unknown_token=4) + ) flat = ds[(0, 0)] toks = np.asarray(flat.flank_tokens.to_ragged().data) assert (toks == 4).any() # some padded positions tokenized to unknown @@ -927,14 +973,20 @@ def test_oob_flank_at_contig_start(snap_dataset): def test_no_awkward_on_flank_hot_path(snap_dataset, monkeypatch): import awkward as ak + calls = {"n": 0} orig = ak.highlevel.Array.__getitem__ + def spy(self, *a, **k): calls["n"] += 1 return orig(self, *a, **k) + monkeypatch.setattr(ak.highlevel.Array, "__getitem__", spy) - ds = (snap_dataset.with_seqs("variants").with_output_format("flat") - .with_settings(flank_length=5, token_alphabet=b"ACGT", unknown_token=4)) + ds = ( + snap_dataset.with_seqs("variants") + .with_output_format("flat") + .with_settings(flank_length=5, token_alphabet=b"ACGT", unknown_token=4) + ) calls["n"] = 0 _ = ds[[0, 1, 2], [0, 1, 2]] assert calls["n"] == 0, "awkward __getitem__ called on flat flank hot path" @@ -1024,8 +1076,11 @@ rtk git commit -m "bench(flat): flank dedup decision micro-benchmark" ```python def test_dedup_invariant(snap_dataset, monkeypatch): # output identical with dedup forced on vs off - base = (snap_dataset.with_seqs("variants").with_output_format("flat") - .with_settings(flank_length=5, token_alphabet=b"ACGT", unknown_token=4)) + base = ( + snap_dataset.with_seqs("variants") + .with_output_format("flat") + .with_settings(flank_length=5, token_alphabet=b"ACGT", unknown_token=4) + ) monkeypatch.setenv("GVL_FLANK_DEDUP", "0") a = base[[0, 0, 1], [0, 1, 1]].flank_tokens.to_ragged() monkeypatch.setenv("GVL_FLANK_DEDUP", "1") diff --git a/docs/superpowers/plans/2026-06-13-flat-output-mode-a0-a.md b/docs/superpowers/plans/2026-06-13-flat-output-mode-a0-a.md index 45284b78..33b081db 100644 --- a/docs/superpowers/plans/2026-06-13-flat-output-mode-a0-a.md +++ b/docs/superpowers/plans/2026-06-13-flat-output-mode-a0-a.md @@ -182,9 +182,7 @@ class _FlatVariants: return _FlatVariants({k: v.reshape(shape) for k, v in self.fields.items()}) def squeeze(self, axis: int | None = None) -> "_FlatVariants": - return _FlatVariants( - {k: v.squeeze(axis) for k, v in self.fields.items()} - ) + return _FlatVariants({k: v.squeeze(axis) for k, v in self.fields.items()}) def reverse_masked(self, mask: NDArray[np.bool_]) -> "_FlatVariants": # Only alt/ref alleles are reverse-complemented; scalar fields unchanged @@ -257,10 +255,10 @@ from ._dataset._flat_variants import _FlatAlleles as FlatAlleles And add the four names to `__all__` (keep alphabetical ordering used in the file): ```python - "FlatAlleles", - "FlatAnnotatedHaps", - "FlatRagged", - "FlatVariants", +("FlatAlleles",) +("FlatAnnotatedHaps",) +("FlatRagged",) +("FlatVariants",) ``` - [ ] **Step 4: Run test to verify it passes** @@ -395,6 +393,7 @@ rtk git commit -m "feat(flat): output_format field + with_output_format + QueryV ```python # tests/dataset/test_flat_mode_equivalence.py """flat-mode output, re-wrapped via .to_ragged(), must be byte-identical to ragged mode.""" + from __future__ import annotations import numpy as np @@ -410,8 +409,10 @@ def _to_plain(obj): """Normalize a ragged/annot/flat object into dict of ndarrays for comparison.""" if isinstance(obj, RaggedAnnotatedHaps): return { - "haps": np.asarray(obj.haps.data), "haps_off": np.asarray(obj.haps.offsets), - "vidx": np.asarray(obj.var_idxs.data), "pos": np.asarray(obj.ref_coords.data), + "haps": np.asarray(obj.haps.data), + "haps_off": np.asarray(obj.haps.offsets), + "vidx": np.asarray(obj.var_idxs.data), + "pos": np.asarray(obj.ref_coords.data), } if isinstance(obj, Ragged): return {"data": np.asarray(obj.data), "off": np.asarray(obj.offsets)} @@ -518,8 +519,8 @@ def _rv_to_lists(rv: RaggedVariants) -> dict: @pytest.mark.parametrize("idx", IDX) def test_a_flat_variants_to_ragged_matches_ragged(snap_dataset, idx): ds = snap_dataset.with_seqs("variants").with_tracks(False) - ragged = ds[idx] # RaggedVariants (current path) - flat = ds.with_output_format("flat")[idx] # _FlatVariants + ragged = ds[idx] # RaggedVariants (current path) + flat = ds.with_output_format("flat")[idx] # _FlatVariants rewrapped = flat.to_ragged() assert _rv_to_lists(rewrapped) == _rv_to_lists(ragged) @@ -546,6 +547,7 @@ Expected: FAIL — flat mode still returns a `RaggedVariants` from the reconstru ```python # add to python/genvarloader/_dataset/_flat_variants.py + @nb.njit(nogil=True, cache=True) def _gather_v_idxs(geno_offset_idx, geno_offsets, geno_v_idxs): """Concatenate the per-(b*p)-row sparse variant-index slices into one flat @@ -555,7 +557,9 @@ def _gather_v_idxs(geno_offset_idx, geno_offsets, geno_v_idxs): out_offsets[0] = 0 for i in range(n_rows): goi = geno_offset_idx[i] - out_offsets[i + 1] = out_offsets[i] + (geno_offsets[goi + 1] - geno_offsets[goi]) + out_offsets[i + 1] = out_offsets[i] + ( + geno_offsets[goi + 1] - geno_offsets[goi] + ) total = out_offsets[n_rows] v_idxs = np.empty(total, geno_v_idxs.dtype) dst = 0 @@ -578,7 +582,9 @@ def _gather_alleles(v_idxs, allele_bytes, allele_offsets): seq_offsets[0] = 0 for i in range(n): v = v_idxs[i] - seq_offsets[i + 1] = seq_offsets[i] + (allele_offsets[v + 1] - allele_offsets[v]) + seq_offsets[i + 1] = seq_offsets[i] + ( + allele_offsets[v + 1] - allele_offsets[v] + ) data = np.empty(seq_offsets[n], np.uint8) dst = 0 for i in range(n): @@ -649,7 +655,7 @@ def get_variants_flat(haps, idx) -> "_FlatVariants": # get_haps_and_shifts (not the bare __call__). Replicate: if keep is supplied, # compact with it here. See Step 3b for how keep is threaded. - n_rows = row_offsets.shape[0] # b*p + 1 + n_rows = row_offsets.shape[0] # b*p + 1 b = goi_flat.shape[0] // ploidy shape = (b, ploidy, None) @@ -685,7 +691,9 @@ def get_variants_flat(haps, idx) -> "_FlatVariants": # if AF/exonic filtered, dosage must be compacted with the same keep mask # — keep dosage gather BEFORE compaction OR carry keep; simplest: gather # dosage parallel to the UNfiltered v_idxs then apply the same _compact_keep. - fields["dosage"] = _Flat.from_offsets(np.ascontiguousarray(dos), shape, row_offsets) + fields["dosage"] = _Flat.from_offsets( + np.ascontiguousarray(dos), shape, row_offsets + ) # remaining info fields for k in haps.var_fields: @@ -707,21 +715,39 @@ def get_variants_flat(haps, idx) -> "_FlatVariants": In `_haps.py` `Haps.__call__` (line 502), add `flat: bool = False`: ```python - def __call__(self, idx, r_idx, regions, output_length, jitter, rng, deterministic, - splice_plan=None, flat: bool = False): - if issubclass(self.kind, RaggedVariants): - if splice_plan is not None: - raise NotImplementedError("Spliced output is not supported for RaggedVariants.") - if flat: - from ._flat_variants import get_variants_flat - return cast(_H, get_variants_flat(self, idx)) - ragv = self._get_variants(idx=idx, regions=None, shifts=None) - return cast(_H, ragv) - else: - haps, *_ = self.get_haps_and_shifts(idx=idx, regions=regions, - output_length=output_length, rng=rng, deterministic=deterministic, - splice_plan=splice_plan) - return haps +def __call__( + self, + idx, + r_idx, + regions, + output_length, + jitter, + rng, + deterministic, + splice_plan=None, + flat: bool = False, +): + if issubclass(self.kind, RaggedVariants): + if splice_plan is not None: + raise NotImplementedError( + "Spliced output is not supported for RaggedVariants." + ) + if flat: + from ._flat_variants import get_variants_flat + + return cast(_H, get_variants_flat(self, idx)) + ragv = self._get_variants(idx=idx, regions=None, shifts=None) + return cast(_H, ragv) + else: + haps, *_ = self.get_haps_and_shifts( + idx=idx, + regions=regions, + output_length=output_length, + rng=rng, + deterministic=deterministic, + splice_plan=splice_plan, + ) + return haps ``` (For exonic filter parity, also branch in `get_haps_and_shifts` where `_get_variants` is called with `keep`/`keep_offsets`: when `flat`, call `get_variants_flat(self, idx, keep=req.keep, keep_offsets=req.keep_offsets)` and add those optional params to `get_variants_flat` applying `_compact_keep`. Only needed if the variants-kind path flows through `get_haps_and_shifts` — verify which call site the variants output uses; the bare `__call__` path is the primary one.) @@ -731,11 +757,16 @@ Add `flat: bool = False` to the `Reconstructor` protocol `__call__` and to each In `_query.py` `_getitem_unspliced` (line 145) and `_getitem_spliced`, pass `flat=view.flat_output`: ```python - recon = view.recon( - idx=ds_idx, r_idx=r_idx, regions=regions, output_length=view.output_length, - jitter=view.jitter, rng=view.rng, deterministic=view.deterministic, - flat=view.flat_output, - ) +recon = view.recon( + idx=ds_idx, + r_idx=r_idx, + regions=regions, + output_length=view.output_length, + jitter=view.jitter, + rng=view.rng, + deterministic=view.deterministic, + flat=view.flat_output, +) ``` - [ ] **Step 3c: Teach `_query.py` to carry `_FlatVariants` through rc / boundary** @@ -743,9 +774,10 @@ In `_query.py` `_getitem_unspliced` (line 145) and `_getitem_spliced`, pass `fla In `reverse_complement_ragged` (line 335), add a branch before `RaggedVariants`: ```python - from ._flat_variants import _FlatVariants - if isinstance(rag, _FlatVariants): - return rag.reverse_masked(to_rc) +from ._flat_variants import _FlatVariants + +if isinstance(rag, _FlatVariants): + return rag.reverse_masked(to_rc) ``` Add the overload signature too. In `getitem`, the flat branch (Task 4) already leaves non-`(_Flat, _FlatAnnotatedHaps)` objects alone — but `_FlatVariants` must also bypass `to_ragged()` in ragged mode? No: in **ragged** mode the variant recon returns `RaggedVariants` (flat=False), so nothing changes. In **flat** mode the recon returns `_FlatVariants` and the `if not view.flat_output:` guard skips conversion — so `_FlatVariants` flows straight out. `reshape`/`squeeze` apply via `_FlatVariants` methods. Confirm `squeeze` matches `RaggedVariants.squeeze` (→ `self[0]`): `_FlatVariants.squeeze(0)` delegates to `_Flat.squeeze(0)` / `_FlatAlleles` (add a `squeeze` to `_FlatAlleles` that drops the leading fixed axis). The equivalence test on a scalar `idx` (which triggers squeeze) is the gate. @@ -753,14 +785,15 @@ Add the overload signature too. In `getitem`, the flat branch (Task 4) already l - [ ] **Step 3d: Add `squeeze` to `_FlatAlleles`** ```python - def squeeze(self, axis: int | None = None) -> "_FlatAlleles": - fixed = [d for d in self.shape if d is not None] - if axis is None: - fixed = [d for d in fixed if d != 1] - else: - del fixed[axis] - return _FlatAlleles(self.byte_data, self.seq_offsets, self.var_offsets, - (*fixed, None)) +def squeeze(self, axis: int | None = None) -> "_FlatAlleles": + fixed = [d for d in self.shape if d is not None] + if axis is None: + fixed = [d for d in fixed if d != 1] + else: + del fixed[axis] + return _FlatAlleles( + self.byte_data, self.seq_offsets, self.var_offsets, (*fixed, None) + ) ``` - [ ] **Step 4: Run the variant equivalence tests** @@ -797,10 +830,14 @@ Add a test that runs `ds.with_seqs("variants").with_tracks(False).with_output_fo ```python def test_flat_variants_decode_has_no_awkward(snap_dataset, awkward_getitem_counter): - ds = snap_dataset.with_seqs("variants").with_tracks(False).with_output_format("flat") + ds = ( + snap_dataset.with_seqs("variants").with_tracks(False).with_output_format("flat") + ) with awkward_getitem_counter() as count: _ = ds[(np.arange(4),)] - assert count.value == 0, f"awkward.__getitem__ called {count.value}x in flat variant decode" + assert count.value == 0, ( + f"awkward.__getitem__ called {count.value}x in flat variant decode" + ) ``` - [ ] **Step 3: Run to verify it fails (or passes)** diff --git a/docs/superpowers/plans/2026-06-13-ref-fetch-parallel-overhead.md b/docs/superpowers/plans/2026-06-13-ref-fetch-parallel-overhead.md index 470de26c..81247f47 100644 --- a/docs/superpowers/plans/2026-06-13-ref-fetch-parallel-overhead.md +++ b/docs/superpowers/plans/2026-06-13-ref-fetch-parallel-overhead.md @@ -220,22 +220,18 @@ Replace the body of `fetch` that currently reads (lines 133-143): with: ```python - seqs = np.empty(offsets[-1], np.uint8) - kernel = ( - _fetch_impl_par - if should_parallelize(int(offsets[-1])) - else _fetch_impl_ser - ) - kernel( - c_idxs, - starts, - ends, - self.reference, - self.offsets, - self.pad_char, - seqs, - offsets, - ) +seqs = np.empty(offsets[-1], np.uint8) +kernel = _fetch_impl_par if should_parallelize(int(offsets[-1])) else _fetch_impl_ser +kernel( + c_idxs, + starts, + ends, + self.reference, + self.offsets, + self.pad_char, + seqs, + offsets, +) ``` Replace the existing `_fetch_impl` definition (lines 150-165): @@ -263,7 +259,9 @@ with a shared inner row kernel plus serial and parallel wrappers: ```python @nb.njit(nogil=True, cache=True, inline="always") -def _fetch_row(i, c_idxs, starts, ends, reference, ref_offsets, pad_char, out, out_offsets): +def _fetch_row( + i, c_idxs, starts, ends, reference, ref_offsets, pad_char, out, out_offsets +): r_s, r_e = ref_offsets[c_idxs[i]], ref_offsets[c_idxs[i] + 1] o_s, o_e = out_offsets[i], out_offsets[i + 1] padded_slice(reference[r_s:r_e], starts[i], ends[i], pad_char, out[o_s:o_e]) @@ -387,14 +385,18 @@ def _get_reference_row(i, regions, out_offsets, reference, ref_offsets, pad_char @nb.njit(parallel=True, nogil=True, cache=True) def _get_reference_par(regions, out_offsets, reference, ref_offsets, pad_char, out): for i in nb.prange(len(regions)): - _get_reference_row(i, regions, out_offsets, reference, ref_offsets, pad_char, out) + _get_reference_row( + i, regions, out_offsets, reference, ref_offsets, pad_char, out + ) return out @nb.njit(nogil=True, cache=True) def _get_reference_ser(regions, out_offsets, reference, ref_offsets, pad_char, out): for i in range(len(regions)): - _get_reference_row(i, regions, out_offsets, reference, ref_offsets, pad_char, out) + _get_reference_row( + i, regions, out_offsets, reference, ref_offsets, pad_char, out + ) return out diff --git a/docs/superpowers/plans/2026-06-13-variant-assembly-B.md b/docs/superpowers/plans/2026-06-13-variant-assembly-B.md index 96404f98..1e1ca835 100644 --- a/docs/superpowers/plans/2026-06-13-variant-assembly-B.md +++ b/docs/superpowers/plans/2026-06-13-variant-assembly-B.md @@ -47,14 +47,12 @@ Expected: PASS. This is the pre-refactor reference (esp. `test_flat_getitem_snap In `python/genvarloader/_dataset/_haps.py`, replace the variants branch of `__call__` (currently the `if flat: ... else: ragv = self._get_variants(...)` block) with an unconditional flat decode: ```python - if issubclass(self.kind, RaggedVariants): - if splice_plan is not None: - raise NotImplementedError( - "Spliced output is not supported for RaggedVariants." - ) - from ._flat_variants import get_variants_flat +if issubclass(self.kind, RaggedVariants): + if splice_plan is not None: + raise NotImplementedError("Spliced output is not supported for RaggedVariants.") + from ._flat_variants import get_variants_flat - return cast(_H, get_variants_flat(self, idx)) + return cast(_H, get_variants_flat(self, idx)) ``` (The `flat` parameter on `__call__` is retained for signature stability across reconstructors but is no longer read for variants — the `_query.py` boundary decides ragged-vs-flat via `view.flat_output`.) @@ -138,19 +136,23 @@ def test_fill_empty_seq_kernel(): # 3 rows: empty, ["AC","G"], empty data = np.frombuffer(b"ACG", np.uint8).copy() - var_off = np.array([0, 0, 2, 2], np.int64) # per-row variant boundaries - seq_off = np.array([0, 2, 3], np.int64) # per-variant byte boundaries + var_off = np.array([0, 0, 2, 2], np.int64) # per-row variant boundaries + seq_off = np.array([0, 2, 3], np.int64) # per-variant byte boundaries dummy = np.frombuffer(b"N", np.uint8).copy() nd, nvar, nseq = _fill_empty_seq(data, var_off, seq_off, dummy) - assert nvar.tolist() == [0, 1, 3, 4] # each empty row gains 1 variant - assert nseq.tolist() == [0, 1, 3, 4, 5] # dummy(1) AC(2) G(1) dummy(1) + assert nvar.tolist() == [0, 1, 3, 4] # each empty row gains 1 variant + assert nseq.tolist() == [0, 1, 3, 4, 5] # dummy(1) AC(2) G(1) dummy(1) assert bytes(nd) == b"NACGN" def test_fill_empty_groups_roundtrip(): import awkward as ak - from genvarloader._dataset._flat_variants import DummyVariant, _FlatAlleles, _FlatVariants + from genvarloader._dataset._flat_variants import ( + DummyVariant, + _FlatAlleles, + _FlatVariants, + ) from genvarloader._flat import _Flat # b*p = 3 rows: row0 empty, row1 has [b"AC", b"G"], row2 empty @@ -171,13 +173,21 @@ def test_fill_empty_groups_roundtrip(): def test_fill_empty_groups_noop_when_no_empties(): - from genvarloader._dataset._flat_variants import DummyVariant, _FlatAlleles, _FlatVariants + from genvarloader._dataset._flat_variants import ( + DummyVariant, + _FlatAlleles, + _FlatVariants, + ) from genvarloader._flat import _Flat import awkward as ak group_off = np.array([0, 1, 2], np.int64) # every row has 1 variant - alt = _FlatAlleles(np.frombuffer(b"AG", np.uint8).copy(), - np.array([0, 1, 2], np.int64), group_off.copy(), (2, None)) + alt = _FlatAlleles( + np.frombuffer(b"AG", np.uint8).copy(), + np.array([0, 1, 2], np.int64), + group_off.copy(), + (2, None), + ) start = _Flat.from_offsets(np.array([3, 7], np.int32), (2, None), group_off.copy()) fv = _FlatVariants(fields={"alt": alt, "start": start}) filled = fv.fill_empty_groups(DummyVariant()) @@ -303,23 +313,27 @@ def _fill_empty_seq(data, var_offsets, seq_offsets, dummy): # pragma: no cover Then `fill_empty_groups` on `_FlatVariants` (add as a method): ```python - def fill_empty_groups(self, dummy: "DummyVariant") -> "_FlatVariants": - """Insert one dummy variant into each empty (b*p) group; non-empty - groups are unchanged. Every field shares the same empty-row pattern, so - the rebuilt offsets stay consistent across fields.""" - from .._flat import _Flat - - new_fields: dict[str, Any] = {} - for name, f in self.fields.items(): - if isinstance(f, _FlatAlleles): - db = np.frombuffer(dummy.alt if name == "alt" else dummy.ref, np.uint8).copy() - nd, nvar, nseq = _fill_empty_seq(f.byte_data, f.var_offsets, f.seq_offsets, db) - new_fields[name] = _FlatAlleles(nd, nseq, nvar, f.shape) - else: - fill = dummy.scalar_for(name, f.data.dtype) - nd, noff = _fill_empty_scalar(f.data, f.offsets, fill) - new_fields[name] = _Flat.from_offsets(nd, f.shape, noff) - return _FlatVariants(new_fields) +def fill_empty_groups(self, dummy: "DummyVariant") -> "_FlatVariants": + """Insert one dummy variant into each empty (b*p) group; non-empty + groups are unchanged. Every field shares the same empty-row pattern, so + the rebuilt offsets stay consistent across fields.""" + from .._flat import _Flat + + new_fields: dict[str, Any] = {} + for name, f in self.fields.items(): + if isinstance(f, _FlatAlleles): + db = np.frombuffer( + dummy.alt if name == "alt" else dummy.ref, np.uint8 + ).copy() + nd, nvar, nseq = _fill_empty_seq( + f.byte_data, f.var_offsets, f.seq_offsets, db + ) + new_fields[name] = _FlatAlleles(nd, nseq, nvar, f.shape) + else: + fill = dummy.scalar_for(name, f.data.dtype) + nd, noff = _fill_empty_scalar(f.data, f.offsets, fill) + new_fields[name] = _Flat.from_offsets(nd, f.shape, noff) + return _FlatVariants(new_fields) ``` - [ ] **Step 4: Run to verify pass** @@ -491,7 +505,7 @@ Expected: FAIL (`with_settings` has no `dummy_variant` kwarg; `gvl.DummyVariant` In `python/genvarloader/_dataset/_impl.py`, add to the `with_settings` signature (after `var_filter`): ```python - dummy_variant: "DummyVariant | Literal[False] | None" = None, +dummy_variant: "DummyVariant | Literal[False] | None" = (None,) ``` Add a runtime import at the top of `_impl.py` (no cycle — `_flat_variants` imports only `_flat` at runtime): @@ -503,14 +517,12 @@ from ._flat_variants import DummyVariant Add a handling block inside `with_settings` (alongside the `min_af`/`var_filter` blocks, before the `if "_seqs" in to_evolve ...` rebuild): ```python - if dummy_variant is not None: - if not isinstance(self._seqs, Haps): - raise ValueError( - "dummy_variant requires a dataset with variants/genotypes." - ) - dv = None if dummy_variant is False else dummy_variant - haps = to_evolve.get("_seqs", self._seqs) - to_evolve["_seqs"] = replace(haps, dummy_variant=dv) +if dummy_variant is not None: + if not isinstance(self._seqs, Haps): + raise ValueError("dummy_variant requires a dataset with variants/genotypes.") + dv = None if dummy_variant is False else dummy_variant + haps = to_evolve.get("_seqs", self._seqs) + to_evolve["_seqs"] = replace(haps, dummy_variant=dv) ``` Add a `dummy_variant` entry to the `with_settings` docstring Parameters section: @@ -586,8 +598,12 @@ import genvarloader as gvl @pytest.mark.parametrize("idx", IDX) def test_b_dummy_fill_flat_to_ragged_matches_ragged(snap_dataset, idx): dv = gvl.DummyVariant(start=-1, alt=b"N", ref=b"N", ilen=0) - ds = snap_dataset.with_seqs("variants").with_tracks(False).with_settings(dummy_variant=dv) - ragged = ds[idx] # ragged mode (now flat decode + to_ragged) + ds = ( + snap_dataset.with_seqs("variants") + .with_tracks(False) + .with_settings(dummy_variant=dv) + ) + ragged = ds[idx] # ragged mode (now flat decode + to_ragged) rewrapped = ds.with_output_format("flat")[idx].to_ragged() assert _rv_to_lists(rewrapped) == _rv_to_lists(ragged) @@ -596,7 +612,11 @@ def test_b_dummy_fill_no_empty_groups(snap_dataset): import awkward as ak dv = gvl.DummyVariant(start=-1, alt=b"N", ref=b"N") - ds = snap_dataset.with_seqs("variants").with_tracks(False).with_settings(dummy_variant=dv) + ds = ( + snap_dataset.with_seqs("variants") + .with_tracks(False) + .with_settings(dummy_variant=dv) + ) idx = (np.arange(min(6, snap_dataset.shape[0])),) rv = ds[idx] for ploid_groups in ak.to_list(rv["start"]): diff --git a/docs/superpowers/plans/2026-06-16-update-and-annot-tracks.md b/docs/superpowers/plans/2026-06-16-update-and-annot-tracks.md index 7dc6762b..f7bc0324 100644 --- a/docs/superpowers/plans/2026-06-16-update-and-annot-tracks.md +++ b/docs/superpowers/plans/2026-06-16-update-and-annot-tracks.md @@ -342,9 +342,7 @@ def test_annot_intervals_from_bigwig(tmp_path): data_dir = Path(__file__).parent.parent / "data" / "bigwig" bw = data_dir / "sample_0.bw" # a region known to overlap intervals in the fixture bigwig - regions = pl.DataFrame( - {"chrom": ["chr1"], "chromStart": [0], "chromEnd": [1000]} - ) + regions = pl.DataFrame({"chrom": ["chr1"], "chromStart": [0], "chromEnd": [1000]}) itvs = _annot_intervals(regions, bw, max_mem=2**30) # shape (regions, None), one region assert itvs.values.offsets.shape == (2,) @@ -414,13 +412,9 @@ def _annot_intervals_from_bigwig( out_ends.append(np.asarray(itvs.ends[r, 0], dtype=np.int32)) out_values.append(np.asarray(itvs.values[r, 0], dtype=np.float32)) lengths.append(len(s)) - flat_starts = ( - np.concatenate(out_starts) if out_starts else np.empty(0, np.int32) - ) + flat_starts = np.concatenate(out_starts) if out_starts else np.empty(0, np.int32) flat_ends = np.concatenate(out_ends) if out_ends else np.empty(0, np.int32) - flat_values = ( - np.concatenate(out_values) if out_values else np.empty(0, np.float32) - ) + flat_values = np.concatenate(out_values) if out_values else np.empty(0, np.float32) offsets = lengths_to_offsets(np.asarray(lengths, np.int32)) shape = (regions.height, None) return RaggedIntervals( @@ -483,13 +477,13 @@ def test_write_with_annot_tracks(phased_vcf, ref_fasta, tmp_path): import polars as pl out = tmp_path / "ds" - bed = pl.DataFrame( - {"chrom": ["chr1"], "chromStart": [0], "chromEnd": [64]} - ) + bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [0], "chromEnd": [64]}) annot = bed.with_columns(chromEnd=pl.col("chromStart") + 1, score=pl.lit(1.0)) gvl.write(out, bed, variants=phased_vcf, annot_tracks={"5ss": annot}) - ds = gvl.Dataset.open(out, ref_fasta).with_seqs("annotated").with_tracks( - "5ss", "tracks" + ds = ( + gvl.Dataset.open(out, ref_fasta) + .with_seqs("annotated") + .with_tracks("5ss", "tracks") ) assert "5ss" in ds.available_tracks ``` @@ -538,15 +532,13 @@ Update the "at least one input" guard: After the existing track-writing block (~286-289), add: ```python - if annot_tracks is not None: - logger.info("Writing annotation tracks.") - annot_bed = regions_to_bed( - np.load(path / "regions.npy"), contigs - ).select("chrom", "chromStart", "chromEnd") - for name, source in annot_tracks.items(): - _write_annot_track( - path / "annot_intervals" / name, annot_bed, source, max_mem - ) +if annot_tracks is not None: + logger.info("Writing annotation tracks.") + annot_bed = regions_to_bed(np.load(path / "regions.npy"), contigs).select( + "chrom", "chromStart", "chromEnd" + ) + for name, source in annot_tracks.items(): + _write_annot_track(path / "annot_intervals" / name, annot_bed, source, max_mem) ``` Add the import at the top of `_write.py`: @@ -672,37 +664,35 @@ Replace the sequential track / annot blocks in `write` with job construction. Va still run first, serially (they finalize `gvl_bed` and write `regions.npy`): ```python - # variants already written above; regions.npy is finalized here. - _write_regions(path, gvl_bed, contigs) +# variants already written above; regions.npy is finalized here. +_write_regions(path, gvl_bed, contigs) - jobs: list[Callable[[int], None]] = [] - if tracks is not None: - _tracks = list(tracks) - _bed = gvl_bed +jobs: list[Callable[[int], None]] = [] +if tracks is not None: + _tracks = list(tracks) + _bed = gvl_bed - def _tracks_job(mm: int, _tracks=_tracks, _bed=_bed): - for tr in _tracks: - _write_track(path / "intervals" / tr.name, _bed, tr, samples, mm) + def _tracks_job(mm: int, _tracks=_tracks, _bed=_bed): + for tr in _tracks: + _write_track(path / "intervals" / tr.name, _bed, tr, samples, mm) - jobs.append(_tracks_job) + jobs.append(_tracks_job) - if annot_tracks is not None: - annot_bed = regions_to_bed( - np.load(path / "regions.npy"), contigs - ).select("chrom", "chromStart", "chromEnd") - _annots = dict(annot_tracks) +if annot_tracks is not None: + annot_bed = regions_to_bed(np.load(path / "regions.npy"), contigs).select( + "chrom", "chromStart", "chromEnd" + ) + _annots = dict(annot_tracks) - def _annot_job(mm: int, _annots=_annots, _bed=annot_bed): - for name, source in _annots.items(): - _write_annot_track( - path / "annot_intervals" / name, _bed, source, mm - ) + def _annot_job(mm: int, _annots=_annots, _bed=annot_bed): + for name, source in _annots.items(): + _write_annot_track(path / "annot_intervals" / name, _bed, source, mm) - jobs.append(_annot_job) + jobs.append(_annot_job) - if jobs: - logger.info(f"Writing {len(jobs)} track categor(ies).") - _run_jobs(jobs, max_mem) +if jobs: + logger.info(f"Writing {len(jobs)} track categor(ies).") + _run_jobs(jobs, max_mem) ``` Remove the now-superseded sequential `if tracks` / `if annot_tracks` blocks added in @@ -774,7 +764,9 @@ def test_update_accepts_dataset_object(phased_vcf, ref_fasta, bigwigs, tmp_path) assert bigwigs.name in gvl.Dataset.open(out, ref_fasta).available_tracks -def test_update_rejects_extra_or_missing_samples(phased_vcf, ref_fasta, bigwigs, tmp_path): +def test_update_rejects_extra_or_missing_samples( + phased_vcf, ref_fasta, bigwigs, tmp_path +): out = tmp_path / "ds" bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [0], "chromEnd": [64]}) gvl.write(out, bed, variants=phased_vcf, samples=bigwigs.samples[:-1]) @@ -853,7 +845,9 @@ def update( _tracks = list(tracks) if tracks is not None else [] if tracks is None and annot_tracks is None: - raise ValueError("At least one of `tracks` or `annot_tracks` must be provided.") + raise ValueError( + "At least one of `tracks` or `annot_tracks` must be provided." + ) # validate strict sample-set agreement for per-sample tracks for tr in _tracks: @@ -914,24 +908,25 @@ file in the same parent; the latter is a file and `p.iterdir()` on it raises. Gu loops: ```python - def _is_track_dir(p: Path) -> bool: - return ( - p.is_dir() - and ".tmp." not in p.name - and ".old." not in p.name - and not p.name.endswith(".lock") - ) +def _is_track_dir(p: Path) -> bool: + return ( + p.is_dir() + and ".tmp." not in p.name + and ".old." not in p.name + and not p.name.endswith(".lock") + ) - available_tracks: list[str] = [] - if strack_dir.exists(): - for p in strack_dir.iterdir(): - if not _is_track_dir(p): - continue - if len(list(p.iterdir())) == 0: - p.rmdir() - else: - available_tracks.append(p.name) - available_tracks.sort() + +available_tracks: list[str] = [] +if strack_dir.exists(): + for p in strack_dir.iterdir(): + if not _is_track_dir(p): + continue + if len(list(p.iterdir())) == 0: + p.rmdir() + else: + available_tracks.append(p.name) + available_tracks.sort() ``` Apply the same `_is_track_dir` guard to the `available_annots` loop. Define @@ -991,8 +986,10 @@ def test_annot_tracks(phased_vcf, ref_fasta, tmp_path): chromEnd=pl.col("chromStart") + 1, score=pl.lit(1.0) ) gvl.update(out, annot_tracks={"5ss": annots}) - annot_ds = gvl.Dataset.open(out, ref_fasta).with_seqs("annotated").with_tracks( - "5ss", "tracks" + annot_ds = ( + gvl.Dataset.open(out, ref_fasta) + .with_seqs("annotated") + .with_tracks("5ss", "tracks") ) haps, tracks = annot_ds[:] mask = haps.ref_coords == ak.Array( @@ -1024,9 +1021,7 @@ def test_annot_overlap_explicit(): got = annot_overlap(regions, annot) # region 0 [chr1:0-100] overlaps the 3 chr1 annots; region 1 [chr1:50-150] overlaps # the chr1 annots at 60-70 and 90-95; region 2 [chr2:0-100] overlaps the chr2 annot. - np.testing.assert_array_equal( - np.diff(got.values.offsets), np.array([3, 2, 1]) - ) + np.testing.assert_array_equal(np.diff(got.values.offsets), np.array([3, 2, 1])) # region 2's single interval is the chr2 annot 5-15 with score 4.0 np.testing.assert_array_equal(np.asarray(got.starts[2]), [5]) np.testing.assert_array_equal(np.asarray(got.ends[2]), [15]) diff --git a/docs/superpowers/plans/2026-06-17-flat-tracks-intervals-realign.md b/docs/superpowers/plans/2026-06-17-flat-tracks-intervals-realign.md index f16676a5..9dc591fc 100644 --- a/docs/superpowers/plans/2026-06-17-flat-tracks-intervals-realign.md +++ b/docs/superpowers/plans/2026-06-17-flat-tracks-intervals-realign.md @@ -150,9 +150,9 @@ from ._ragged import FlatIntervals, RaggedAnnotatedHaps, RaggedIntervals Add `"FlatIntervals"` to `__all__` (keep alphabetical order — place after `"FlatAnnotatedHaps"`, line 51): ```python - "FlatAnnotatedHaps", - "FlatIntervals", - "FlatRagged", +("FlatAnnotatedHaps",) +("FlatIntervals",) +("FlatRagged",) ``` - [ ] **Step 5: Run tests to verify they pass** @@ -283,13 +283,9 @@ def build_flat_intervals( tb_ends.append(np.asarray(itv.ends.data)[src]) tb_values.append(np.asarray(itv.values.data)[src]) - data_starts = ( - np.concatenate(tb_starts) if tb_starts else np.empty(0, np.int32) - ) + data_starts = np.concatenate(tb_starts) if tb_starts else np.empty(0, np.int32) data_ends = np.concatenate(tb_ends) if tb_ends else np.empty(0, np.int32) - data_values = ( - np.concatenate(tb_values) if tb_values else np.empty(0, np.float32) - ) + data_values = np.concatenate(tb_values) if tb_values else np.empty(0, np.float32) offsets_tb = lengths_to_offsets(lengths_tb.ravel()) # (T*B + 1) # Pass 2: reorder groups (t, b) -> (b, t). For output group (b, t) the @@ -298,9 +294,8 @@ def build_flat_intervals( final_lengths = lengths_tb.ravel()[perm] final_offsets = lengths_to_offsets(final_lengths) total = int(final_offsets[-1]) - src = ( - np.repeat(offsets_tb[perm] - final_offsets[:-1], final_lengths) - + np.arange(total, dtype=np.int64) + src = np.repeat(offsets_tb[perm] - final_offsets[:-1], final_lengths) + np.arange( + total, dtype=np.int64 ) shape = (B, T, None) @@ -592,8 +587,10 @@ def test_default_haps_tracks_realigns(): def test_realign_false_haps_tracks_uses_seqstracks_and_is_reference_coord(): ds = gvl.get_dummy_dataset() - asis = ds.with_seqs("haplotypes").with_tracks(["read-depth"]).with_settings( - realign_tracks=False + asis = ( + ds.with_seqs("haplotypes") + .with_tracks(["read-depth"]) + .with_settings(realign_tracks=False) ) assert type(asis._recon) is SeqsTracks @@ -612,10 +609,9 @@ def test_intervals_plus_haplotypes_requires_realign_false(): def test_intervals_plus_haplotypes_ok_when_realign_false(): ds = gvl.get_dummy_dataset() - out = ( - ds.with_settings(realign_tracks=False) - .with_tracks(["read-depth"], kind="intervals")[[0], [0]] - ) + out = ds.with_settings(realign_tracks=False).with_tracks( + ["read-depth"], kind="intervals" + )[[0], [0]] seqs, itvs = out assert isinstance(itvs, gvl.RaggedIntervals) @@ -710,9 +706,7 @@ Update each `_build_reconstructor(...)` call to pass `self.realign_tracks`: - `with_seqs` (line 717): ```python - new_recon = _build_reconstructor( - new_seqs, self._tracks, kind, self.realign_tracks - ) +new_recon = _build_reconstructor(new_seqs, self._tracks, kind, self.realign_tracks) ``` - `with_tracks` (line 769): @@ -775,18 +769,14 @@ In the body, after the `unphased_union` block (lines 423-429), add: Update the recon-rebuild guard (lines 432-437) to also rebuild when `realign_tracks` changes and pass it through: ```python - # If any source state changed, rebuild _recon via the factory. - if ( - "_seqs" in to_evolve - or "_tracks" in to_evolve - or "realign_tracks" in to_evolve - ): - new_seqs = to_evolve.get("_seqs", self._seqs) - new_tracks = to_evolve.get("_tracks", self._tracks) - new_realign = to_evolve.get("realign_tracks", self.realign_tracks) - to_evolve["_recon"] = _build_reconstructor( - new_seqs, new_tracks, self._seqs_kind, new_realign - ) +# If any source state changed, rebuild _recon via the factory. +if "_seqs" in to_evolve or "_tracks" in to_evolve or "realign_tracks" in to_evolve: + new_seqs = to_evolve.get("_seqs", self._seqs) + new_tracks = to_evolve.get("_tracks", self._tracks) + new_realign = to_evolve.get("realign_tracks", self.realign_tracks) + to_evolve["_recon"] = _build_reconstructor( + new_seqs, new_tracks, self._seqs_kind, new_realign + ) ``` - [ ] **Step 8: Forward `realign_tracks` (and `output_format`) through `with_len`** @@ -921,7 +911,9 @@ def test_flat_float_tracks_only_returns_flatragged(): def test_flat_haps_plus_tracks_returns_flat_pair(): ds = gvl.get_dummy_dataset() flat = ( - ds.with_seqs("haplotypes").with_tracks(["read-depth"]).with_output_format("flat") + ds.with_seqs("haplotypes") + .with_tracks(["read-depth"]) + .with_output_format("flat") ) seqs, tracks = flat[[0, 1], [0, 1]] assert type(seqs).__name__ == "_Flat" diff --git a/docs/superpowers/plans/2026-06-18-svar-custom-format-fields.md b/docs/superpowers/plans/2026-06-18-svar-custom-format-fields.md index 999073c9..17336c29 100644 --- a/docs/superpowers/plans/2026-06-18-svar-custom-format-fields.md +++ b/docs/superpowers/plans/2026-06-18-svar-custom-format-fields.md @@ -88,7 +88,10 @@ def custom_field_ds(tmp_path, filtered_svar, source_bed): # Custom int16 field, values 0..n-1. mm = np.memmap( - svar_copy / f"{FIELD_NAME}.npy", dtype=FIELD_DTYPE, mode="w+", shape=(n_records,) + svar_copy / f"{FIELD_NAME}.npy", + dtype=FIELD_DTYPE, + mode="w+", + shape=(n_records,), ) mm[:] = np.arange(n_records, dtype=FIELD_DTYPE) mm.flush() @@ -307,8 +310,9 @@ def _open_variants(gvl_path, ref_fasta, field_name, **settings): gvl.Dataset.open(gvl_path, ref_fasta, rc_neg=False) .with_len("ragged") .with_seqs("variants") - .with_settings(var_fields=["alt", "ilen", "start", "dosage", field_name], - **settings) + .with_settings( + var_fields=["alt", "ilen", "start", "dosage", field_name], **settings + ) ) @@ -321,7 +325,9 @@ def test_custom_field_present_in_ragged_variants(custom_field_ds, ref_fasta): flat = ak.to_numpy(ak.flatten(batch[field_name], axis=None)) assert flat.dtype == np.dtype(FIELD_DTYPE) # Per-cell variant counts equal `start`'s (call-aligned with the genotypes). - assert ak.num(batch[field_name], -1).to_list() == ak.num(batch["start"], -1).to_list() + assert ( + ak.num(batch[field_name], -1).to_list() == ak.num(batch["start"], -1).to_list() + ) def test_custom_field_matches_dosage_gather(custom_field_ds, ref_fasta): diff --git a/docs/superpowers/plans/2026-06-19-bigwig-streaming-write.md b/docs/superpowers/plans/2026-06-19-bigwig-streaming-write.md index f42f9efb..679bfa3e 100644 --- a/docs/superpowers/plans/2026-06-19-bigwig-streaming-write.md +++ b/docs/superpowers/plans/2026-06-19-bigwig-streaming-write.md @@ -62,7 +62,9 @@ def test_make_synthetic_bigwigs_deterministic(tmp_path): def test_make_regions_grouped_in_contig_order(tmp_path): - regions = make_regions({"chr21": 200_000, "chr22": 150_000}, n_per_contig=4, width=1000, seed=1) + regions = make_regions( + {"chr21": 200_000, "chr22": 150_000}, n_per_contig=4, width=1000, seed=1 + ) assert regions.columns == ["chrom", "chromStart", "chromEnd"] # contig-grouped in dict order (chr21 block then chr22 block) chroms = regions["chrom"].to_list() @@ -1004,7 +1006,9 @@ def main() -> None: _write_track(out, bed, track, samples, 4 << 30) dt = time.perf_counter() - t0 - print(f"impl={args.impl} regions={bed.height} samples={len(samples)} wall={dt:.3f}s") + print( + f"impl={args.impl} regions={bed.height} samples={len(samples)} wall={dt:.3f}s" + ) if __name__ == "__main__": diff --git a/docs/superpowers/plans/2026-06-19-rust-table-overlap.md b/docs/superpowers/plans/2026-06-19-rust-table-overlap.md index 35050243..42a25134 100644 --- a/docs/superpowers/plans/2026-06-19-rust-table-overlap.md +++ b/docs/superpowers/plans/2026-06-19-rust-table-overlap.md @@ -1164,9 +1164,9 @@ def _write_track_table( # c_idxs maps via auto-normalization; regions whose contig is absent from the # Table must map to -1 (zero intervals). Detect via norm(). norm = track._cnorm.norm(bed["chrom"].to_list()) - chrom_codes = np.where( - np.array([n is None for n in norm]), -1, chrom_codes - ).astype(np.int32) + chrom_codes = np.where(np.array([n is None for n in norm]), -1, chrom_codes).astype( + np.int32 + ) starts = np.ascontiguousarray(bed["chromStart"].to_numpy(), dtype=np.int32) ends = np.ascontiguousarray(bed["chromEnd"].to_numpy(), dtype=np.int32) track._rust.write_track( @@ -1261,9 +1261,7 @@ def _dense_table(n_intervals: int) -> Table: def test_write_track_table_raises_when_region_exceeds_max_mem(tmp_path): t = _dense_table(1000) - bed = pl.DataFrame( - {"chrom": ["chr1"], "chromStart": [0], "chromEnd": [10_000]} - ) + bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [0], "chromEnd": [10_000]}) # One region overlaps ~1000 intervals = ~12 KB; cap at 12 bytes -> must raise. with pytest.raises(RuntimeError, match="max_mem"): _write_track_table(tmp_path, bed, t, ["s0"], max_mem=12) @@ -1271,9 +1269,7 @@ def test_write_track_table_raises_when_region_exceeds_max_mem(tmp_path): def test_write_track_table_succeeds_within_budget(tmp_path): t = _dense_table(1000) - bed = pl.DataFrame( - {"chrom": ["chr1"], "chromStart": [0], "chromEnd": [10_000]} - ) + bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [0], "chromEnd": [10_000]}) _write_track_table(tmp_path, bed, t, ["s0"], max_mem=1 << 20) assert (tmp_path / "intervals.npy").exists() assert (tmp_path / "offsets.npy").exists() @@ -1501,9 +1497,7 @@ def test_count_and_intervals_match_oracle( np.testing.assert_array_equal(counts, exp_counts) offsets = lengths_to_offsets(counts.ravel()) - itvs = t._intervals_from_offsets( - contig, starts, ends, offsets, sample=present - ) + itvs = t._intervals_from_offsets(contig, starts, ends, offsets, sample=present) n_sel = len(present) for ri in range(n_regions): for sj in range(n_sel): diff --git a/docs/superpowers/plans/2026-06-19-track-write-window-input-floor.md b/docs/superpowers/plans/2026-06-19-track-write-window-input-floor.md index e9440b41..2a2a8783 100644 --- a/docs/superpowers/plans/2026-06-19-track-write-window-input-floor.md +++ b/docs/superpowers/plans/2026-06-19-track-write-window-input-floor.md @@ -70,9 +70,7 @@ def test_stored_window_floored_to_input( # One wide region spanning the chr1 variant cluster out to the contig end. # Its tail is variant-free, so a pre-fix writer truncates chromEnd to the # rightmost variant (< chr1_len); the fix floors it at the input end. - bed = pl.DataFrame( - {"chrom": ["chr1"], "chromStart": [100], "chromEnd": [chr1_len]} - ) + bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [100], "chromEnd": [chr1_len]}) variants = _open_variants(source, vcf_dir, pgen_dir, filtered_svar) out = tmp_path / "ds.gvl" gvl.write(out, bed, variants=variants, overwrite=True) @@ -164,9 +162,7 @@ def test_annot_track_tail_not_truncated_by_variants(vcf_dir, ref_fasta, tmp_path bw.addEntries(["chr1"], [0], ends=[chr1_len], values=[0.5]) bw.close() - bed = pl.DataFrame( - {"chrom": ["chr1"], "chromStart": [100], "chromEnd": [chr1_len]} - ) + bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [100], "chromEnd": [chr1_len]}) out = tmp_path / "ds.gvl" gvl.write( out, @@ -238,9 +234,7 @@ def test_warns_on_truncated_track_window(vcf_dir, ref_fasta, tmp_path): bw.addHeader([("chr1", chr1_len)]) bw.addEntries(["chr1"], [0], ends=[chr1_len], values=[0.5]) bw.close() - bed = pl.DataFrame( - {"chrom": ["chr1"], "chromStart": [100], "chromEnd": [chr1_len]} - ) + bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [100], "chromEnd": [chr1_len]}) out = tmp_path / "ds.gvl" gvl.write( out, diff --git a/docs/superpowers/plans/2026-06-22-migrate-gvl-off-awkward-to-core-ragged.md b/docs/superpowers/plans/2026-06-22-migrate-gvl-off-awkward-to-core-ragged.md index f9640aa2..07462bee 100644 --- a/docs/superpowers/plans/2026-06-22-migrate-gvl-off-awkward-to-core-ragged.md +++ b/docs/superpowers/plans/2026-06-22-migrate-gvl-off-awkward-to-core-ragged.md @@ -63,12 +63,15 @@ import numpy as np import seqpro.rag as r from seqpro.rag import Ragged + def test_to_ak_multi_leading_axis_record(): # (b=2, p=2, ~v) record with an opaque-string and a numeric field var_off = np.array([0, 2, 3, 3, 4], dtype=np.int64) # 4 groups char_off = np.array([0, 2, 3, 6, 7], dtype=np.int64) chars = np.frombuffer(b"ACGTTTX", dtype="S1").copy() - alt = Ragged.from_offsets(chars, (2, 2, None, None), [var_off, char_off]).to_strings() + alt = Ragged.from_offsets( + chars, (2, 2, None, None), [var_off, char_off] + ).to_strings() start = Ragged.from_offsets(np.arange(4, dtype=np.int32), (2, 2, None), alt.offsets) rv = Ragged.from_fields({"alt": alt, "start": start}) got = rv.to_ak() # must not raise @@ -115,27 +118,34 @@ import numpy as np import seqpro.rag as r from seqpro.rag import Ragged + def _record(var_off, char_off, chars): - alt = Ragged.from_offsets(chars, (len(var_off) - 1, None, None), - [var_off, char_off]).to_strings() - start = Ragged.from_offsets(np.arange(int(var_off[-1]), dtype=np.int32), - (len(var_off) - 1, None), alt.offsets) + alt = Ragged.from_offsets( + chars, (len(var_off) - 1, None, None), [var_off, char_off] + ).to_strings() + start = Ragged.from_offsets( + np.arange(int(var_off[-1]), dtype=np.int32), + (len(var_off) - 1, None), + alt.offsets, + ) return Ragged.from_fields({"alt": alt, "start": start}), alt + def test_to_packed_opaque_string_under_axis(): var_off = np.array([0, 2, 3], dtype=np.int64) char_off = np.array([0, 2, 3, 6], dtype=np.int64) rv, alt = _record(var_off, char_off, np.frombuffer(b"ACGTTT", "S1").copy()) - sl = alt[np.array([1, 0])] # produces (2,N) gather offsets - packed = sl.to_packed() # must not raise + sl = alt[np.array([1, 0])] # produces (2,N) gather offsets + packed = sl.to_packed() # must not raise assert packed.to_ak().to_list() == [[b"TTT"], [b"AC", b"G"]] + def test_to_packed_record_with_string_field(): var_off = np.array([0, 2, 3], dtype=np.int64) char_off = np.array([0, 2, 3, 6], dtype=np.int64) rv, _ = _record(var_off, char_off, np.frombuffer(b"ACGTTT", "S1").copy()) sl = rv[np.array([1, 0])] - packed = sl.to_packed() # must not raise + packed = sl.to_packed() # must not raise assert packed["alt"].to_ak().to_list() == [[b"TTT"], [b"AC", b"G"]] assert packed["start"].to_ak().to_list() == [[2], [0, 1]] ``` @@ -182,21 +192,28 @@ import numpy as np import seqpro.rag as r from seqpro.rag import Ragged + def test_concatenate_ragged_axis_prepend_regular(): # prepend a size-1 pad per group (the prepend_pad_itv use case) - base = Ragged.from_offsets(np.array([10, 11, 12], np.int32), (2, None), - np.array([0, 2, 3], np.int64)) # [[10,11],[12]] - pad = Ragged.from_offsets(np.array([-1, -1], np.int32), (2, None), - np.array([0, 1, 2], np.int64)) # [[-1],[-1]] + base = Ragged.from_offsets( + np.array([10, 11, 12], np.int32), (2, None), np.array([0, 2, 3], np.int64) + ) # [[10,11],[12]] + pad = Ragged.from_offsets( + np.array([-1, -1], np.int32), (2, None), np.array([0, 1, 2], np.int64) + ) # [[-1],[-1]] out = r.concatenate([pad, base], axis=-1) assert out.to_ak().to_list() == [[-1, 10, 11], [-1, 12]] + def test_concatenate_matches_awkward_oracle(): import awkward as ak - a = Ragged.from_offsets(np.arange(5, dtype=np.float32), (2, None), - np.array([0, 3, 5], np.int64)) - b = Ragged.from_offsets(np.arange(5, 9, dtype=np.float32), (2, None), - np.array([0, 1, 4], np.int64)) + + a = Ragged.from_offsets( + np.arange(5, dtype=np.float32), (2, None), np.array([0, 3, 5], np.int64) + ) + b = Ragged.from_offsets( + np.arange(5, 9, dtype=np.float32), (2, None), np.array([0, 1, 4], np.int64) + ) got = r.concatenate([a, b], axis=-1).to_ak().to_list() exp = ak.concatenate([a.to_ak(), b.to_ak()], axis=-1).to_list() assert got == exp @@ -218,6 +235,7 @@ In `src/ragged.rs`, add a kernel (mirroring `_ragged_pack`/`_ragged_nested_pack` def concatenate(rags, axis): """Concatenate Rageds along the ragged axis. See SKILL.md.""" from ._core import Ragged + if not rags: raise ValueError("concatenate requires at least one Ragged") rags = [r if isinstance(r, Ragged) else Ragged(r) for r in rags] @@ -225,11 +243,15 @@ def concatenate(rags, axis): ref = rags[0] ax = axis % len(ref.shape) if ax != ref.rag_dim: - raise ValueError(f"concatenate only supports the ragged axis ({ref.rag_dim}), got {axis}") + raise ValueError( + f"concatenate only supports the ragged axis ({ref.rag_dim}), got {axis}" + ) packed = [x.to_packed() for x in rags] from seqpro.seqpro import _ragged_concat # rust - data, offsets = _ragged_concat([p.data for p in packed], - [np.ascontiguousarray(p.offsets) for p in packed]) + + data, offsets = _ragged_concat( + [p.data for p in packed], [np.ascontiguousarray(p.offsets) for p in packed] + ) return Ragged.from_offsets(data, ref.shape, offsets) ``` @@ -289,11 +311,13 @@ import seqpro.rag as r from seqpro.rag import Ragged from genvarloader import RaggedVariants + def _char_alt(var_off, char_off, chars): return Ragged.from_offsets(chars, (2, 1, None, None), [var_off, char_off]) + def test_construct_from_char_and_numeric_fields(): - var_off = np.array([0, 2, 3], np.int64) # b=2,p=1 -> 2 groups + var_off = np.array([0, 2, 3], np.int64) # b=2,p=1 -> 2 groups char_off = np.array([0, 2, 3, 6], np.int64) alt = _char_alt(var_off, char_off, np.frombuffer(b"ACGTTT", "S1").copy()) start = Ragged.from_offsets(np.array([10, 20, 30], np.int32), (2, 1, None), var_off) @@ -306,13 +330,17 @@ def test_construct_from_char_and_numeric_fields(): # ilen derived from alt/ref char lengths assert rv.ilen.to_ak().to_list() == [[[0, 0]], [[0]]] + def test_getitem_returns_raggedvariants(): var_off = np.array([0, 2, 3], np.int64) char_off = np.array([0, 2, 3, 6], np.int64) alt = _char_alt(var_off, char_off, np.frombuffer(b"ACGTTT", "S1").copy()) start = Ragged.from_offsets(np.array([10, 20, 30], np.int32), (2, 1, None), var_off) - rv = RaggedVariants(alt=alt, start=start, ilen=Ragged.from_offsets( - np.zeros(3, np.int32), (2, 1, None), var_off)) + rv = RaggedVariants( + alt=alt, + start=start, + ilen=Ragged.from_offsets(np.zeros(3, np.int32), (2, 1, None), var_off), + ) sub = rv[0] assert isinstance(sub, RaggedVariants) assert sub.alt.to_ak().to_list() == [[b"AC", b"G"]] @@ -337,11 +365,13 @@ from genoray._types import DOSAGE_TYPE, POS_TYPE _ALLELE_FIELDS = ("alt", "ref") + def _as_opaque(rag: Ragged) -> Ragged: """Normalize an allele field to opaque-string (b,p,~v). Accepts an S1 char (b,p,~v,~l) Ragged (collapse via to_strings) or an already-opaque Ragged.""" return rag.to_strings() if not getattr(rag, "is_string", False) else rag + def _share_offsets(rag: Ragged, offsets) -> Ragged: """Rebuild `rag` onto the given (identical) variant-level offsets object so all record fields share it (Ragged.from_fields requires identity, not equality).""" @@ -349,10 +379,12 @@ def _share_offsets(rag: Ragged, offsets) -> Ragged: return rag if getattr(rag, "is_string", False): chars = rag.to_chars() - return Ragged.from_offsets(chars.data, rag.shape, offsets, - str_offsets=chars.offsets).to_strings() + return Ragged.from_offsets( + chars.data, rag.shape, offsets, str_offsets=chars.offsets + ).to_strings() return Ragged.from_offsets(rag.data, rag.shape, offsets) + class RaggedVariants: """Variable-length variants as a single record Ragged with shape (batch, ploidy, ~variants). `alt`/`ref` are opaque-string fields; `start` and @@ -416,14 +448,20 @@ class RaggedVariants: return self._rag["ilen"] alt_len = self._alt_chars("alt").lengths ref_len = self._alt_chars("ref").lengths - return Ragged.from_offsets((alt_len - ref_len).astype(np.int32), - self._rag["start"].shape, self._rag["start"].offsets) + return Ragged.from_offsets( + (alt_len - ref_len).astype(np.int32), + self._rag["start"].shape, + self._rag["start"].offsets, + ) @property def end(self): if "ref" in self.fields: - reflen = Ragged.from_offsets(self._alt_chars("ref").lengths.astype(POS_TYPE), - self.start.shape, self.start.offsets) + reflen = Ragged.from_offsets( + self._alt_chars("ref").lengths.astype(POS_TYPE), + self.start.shape, + self.start.offsets, + ) return self.start + reflen ilen = self.ilen return self.start - np.clip(ilen, None, 0) + 1 @@ -435,8 +473,11 @@ class RaggedVariants: return RaggedVariants.from_record(self._rag[idx]) def reshape(self, shape) -> "RaggedVariants": - return RaggedVariants.from_record(self._rag.reshape(*shape) - if isinstance(shape, tuple) else self._rag.reshape(shape)) + return RaggedVariants.from_record( + self._rag.reshape(*shape) + if isinstance(shape, tuple) + else self._rag.reshape(shape) + ) def squeeze(self, axis=None, **kw) -> "RaggedVariants": return self[0] @@ -475,14 +516,19 @@ import numpy as np from seqpro.rag import Ragged from genvarloader import RaggedVariants + def test_to_packed_after_slice_roundtrips(): - var_off = np.array([0, 2, 3, 4], np.int64) # 3 groups (b=3,p=1) + var_off = np.array([0, 2, 3, 4], np.int64) # 3 groups (b=3,p=1) char_off = np.array([0, 2, 3, 6, 7], np.int64) - alt = Ragged.from_offsets(np.frombuffer(b"ACGTTTX", "S1").copy(), - (3, 1, None, None), [var_off, char_off]) + alt = Ragged.from_offsets( + np.frombuffer(b"ACGTTTX", "S1").copy(), (3, 1, None, None), [var_off, char_off] + ) start = Ragged.from_offsets(np.array([1, 2, 3, 4], np.int32), (3, 1, None), var_off) - rv = RaggedVariants(alt=alt, start=start, - ilen=Ragged.from_offsets(np.zeros(4, np.int32), (3, 1, None), var_off)) + rv = RaggedVariants( + alt=alt, + start=start, + ilen=Ragged.from_offsets(np.zeros(4, np.int32), (3, 1, None), var_off), + ) sub = rv[np.array([2, 0])].to_packed() assert sub.alt.to_ak().to_list() == [[b"X"], [b"AC", b"G"]] assert sub.start.to_ak().to_list() == [[[4]], [[1, 2]]] @@ -535,16 +581,21 @@ import numpy as np from seqpro.rag import Ragged from genvarloader import RaggedVariants + def test_rc_all_complements_and_reverses(): - var_off = np.array([0, 1, 2], np.int64) # 2 groups, 1 variant each + var_off = np.array([0, 1, 2], np.int64) # 2 groups, 1 variant each char_off = np.array([0, 2, 5], np.int64) - alt = Ragged.from_offsets(np.frombuffer(b"ACGTA", "S1").copy(), - (2, 1, None, None), [var_off, char_off]) + alt = Ragged.from_offsets( + np.frombuffer(b"ACGTA", "S1").copy(), (2, 1, None, None), [var_off, char_off] + ) start = Ragged.from_offsets(np.array([0, 0], np.int32), (2, 1, None), var_off) - rv = RaggedVariants(alt=alt, start=start, - ilen=Ragged.from_offsets(np.zeros(2, np.int32), (2, 1, None), var_off)) + rv = RaggedVariants( + alt=alt, + start=start, + ilen=Ragged.from_offsets(np.zeros(2, np.int32), (2, 1, None), var_off), + ) out = rv.rc_(np.array([True, True])) - assert out.alt.to_ak().to_list() == [[b"GT"], [b"TAC"]] # AC->GT, GTA->TAC + assert out.alt.to_ak().to_list() == [[b"GT"], [b"TAC"]] # AC->GT, GTA->TAC ``` - [ ] **Step 2: Run to verify failure** @@ -611,14 +662,19 @@ import numpy as np from seqpro.rag import Ragged from genvarloader import RaggedVariants + def test_pad_fills_empty_groups_only(): var_off = np.array([0, 2, 2, 3], np.int64) # group1 empty char_off = np.array([0, 2, 3, 6], np.int64) - alt = Ragged.from_offsets(np.frombuffer(b"ACGTTT", "S1").copy(), - (3, 1, None, None), [var_off, char_off]) + alt = Ragged.from_offsets( + np.frombuffer(b"ACGTTT", "S1").copy(), (3, 1, None, None), [var_off, char_off] + ) start = Ragged.from_offsets(np.array([1, 2, 3], np.int32), (3, 1, None), var_off) - rv = RaggedVariants(alt=alt, start=start, - ilen=Ragged.from_offsets(np.zeros(3, np.int32), (3, 1, None), var_off)) + rv = RaggedVariants( + alt=alt, + start=start, + ilen=Ragged.from_offsets(np.zeros(3, np.int32), (3, 1, None), var_off), + ) out = rv.pad() assert out.alt.to_ak().to_list() == [[b"AC", b"G"], [b"N"], [b"TTT"]] assert out.start.to_ak().to_list() == [[[1, 2]], [[-1]], [[3]]] @@ -681,17 +737,22 @@ Co-Authored-By: Claude Opus 4.8 " import numpy as np, pytest from seqpro.rag import Ragged from genvarloader import RaggedVariants + torch = pytest.importorskip("torch") + def test_to_nested_tensor_batch_shapes(): var_off = np.array([0, 2, 3], np.int64) char_off = np.array([0, 2, 3, 6], np.int64) - alt = Ragged.from_offsets(np.frombuffer(b"ACGTTT", "S1").copy(), - (2, 1, None, None), [var_off, char_off]) + alt = Ragged.from_offsets( + np.frombuffer(b"ACGTTT", "S1").copy(), (2, 1, None, None), [var_off, char_off] + ) start = Ragged.from_offsets(np.array([1, 2, 3], np.int32), (2, 1, None), var_off) - rv = RaggedVariants(alt=alt, start=start, - ilen=Ragged.from_offsets(np.zeros(3, np.int32), (2, 1, None), var_off) - ).to_packed() + rv = RaggedVariants( + alt=alt, + start=start, + ilen=Ragged.from_offsets(np.zeros(3, np.int32), (2, 1, None), var_off), + ).to_packed() out = rv.to_nested_tensor_batch() assert out["max_n_vars"] == 2 assert out["max_alt_len"] == 3 @@ -766,20 +827,29 @@ Co-Authored-By: Claude Opus 4.8 " import numpy as np from seqpro.rag import Ragged from genvarloader import RaggedVariants -from genvarloader._shm_layout import pack_to_shm, unpack_from_shm # use the actual public entry points +from genvarloader._shm_layout import ( + pack_to_shm, + unpack_from_shm, +) # use the actual public entry points + def _rv(): var_off = np.array([0, 2, 3], np.int64) char_off = np.array([0, 2, 3, 6], np.int64) - alt = Ragged.from_offsets(np.frombuffer(b"ACGTTT", "S1").copy(), - (2, 1, None, None), [var_off, char_off]) + alt = Ragged.from_offsets( + np.frombuffer(b"ACGTTT", "S1").copy(), (2, 1, None, None), [var_off, char_off] + ) start = Ragged.from_offsets(np.array([1, 2, 3], np.int32), (2, 1, None), var_off) - return RaggedVariants(alt=alt, start=start, - ilen=Ragged.from_offsets(np.zeros(3, np.int32), (2, 1, None), var_off)) + return RaggedVariants( + alt=alt, + start=start, + ilen=Ragged.from_offsets(np.zeros(3, np.int32), (2, 1, None), var_off), + ) + def test_rag_variants_shm_roundtrip(): rv = _rv() - rt = unpack_from_shm(pack_to_shm(rv)) # adapt to real API surface + rt = unpack_from_shm(pack_to_shm(rv)) # adapt to real API surface assert rt.alt.to_ak().to_list() == rv.alt.to_ak().to_list() assert rt.start.to_ak().to_list() == rv.start.to_ak().to_list() ``` @@ -830,13 +900,20 @@ import numpy as np from seqpro.rag import Ragged from genvarloader._ragged import RaggedIntervals + def test_prepend_pad_itv_prepends_one_per_group(): def mk(vals, off): - return Ragged.from_offsets(np.array(vals, np.int32), (1, 1, None), - np.array(off, np.int64)) - ri = RaggedIntervals(mk([0, 5], [0, 2]), mk([5, 9], [0, 2]), - Ragged.from_offsets(np.array([1.0, 2.0], np.float32), - (1, 1, None), np.array([0, 2], np.int64))) + return Ragged.from_offsets( + np.array(vals, np.int32), (1, 1, None), np.array(off, np.int64) + ) + + ri = RaggedIntervals( + mk([0, 5], [0, 2]), + mk([5, 9], [0, 2]), + Ragged.from_offsets( + np.array([1.0, 2.0], np.float32), (1, 1, None), np.array([0, 2], np.int64) + ), + ) out = ri.prepend_pad_itv(start=-1, end=-1, value=0.0) assert out.starts.to_ak().to_list() == [[[-1, 0, 5]]] assert out.values.to_ak().to_list() == [[[0.0, 1.0, 2.0]]] @@ -854,13 +931,19 @@ Replace the three `ak.concatenate([pad, self.X.to_ak()], axis=2)` blocks. Build ```python # python/genvarloader/_ragged.py (prepend_pad_itv) import seqpro.rag as spr + b, t, *_ = self.values.shape n = b * t + + def _pad(value, dtype): - return Ragged.from_offsets(np.full(n, value, dtype), (b, t, None), - np.arange(n + 1, dtype=np.int64)) + return Ragged.from_offsets( + np.full(n, value, dtype), (b, t, None), np.arange(n + 1, dtype=np.int64) + ) + + new_starts = spr.concatenate([_pad(start, np.int32), self.starts], axis=-1) -new_ends = spr.concatenate([_pad(end, np.int32), self.ends], axis=-1) +new_ends = spr.concatenate([_pad(end, np.int32), self.ends], axis=-1) new_values = spr.concatenate([_pad(value, np.float32), self.values], axis=-1) return RaggedIntervals(new_starts, new_ends, new_values) ``` @@ -962,6 +1045,7 @@ Confirm consumers (the genotype gather that builds `alt`/`ref`). Each will now r # python/genvarloader/_dataset/_haps.py from seqpro.rag import Ragged + def _build_allele_layout(data, allele_offsets, group_offsets, ploidy) -> Ragged: """Flat allele bytes + two offset levels -> (b, p, ~v, ~l) S1 Ragged.""" buf = np.ascontiguousarray(data) @@ -970,7 +1054,8 @@ def _build_allele_layout(data, allele_offsets, group_offsets, ploidy) -> Ragged: n_groups = group_offsets.size - 1 b = n_groups // ploidy return Ragged.from_offsets( - buf.view("S1"), (b, ploidy, None, None), + buf.view("S1"), + (b, ploidy, None, None), [np.asarray(group_offsets, np.int64), np.asarray(allele_offsets, np.int64)], ) ``` @@ -1022,9 +1107,10 @@ Both build a `(…, ~v, ~l)` allele awkward layout from flat buffers. Replace wi ```python # python/genvarloader/_dataset/_flat_variants.py (_FlatVariants.to_ragged) from ._rag_variants import RaggedVariants + kw = {name: flat.to_ragged() for name, flat in self.fields.items()} # numeric -> Ragged # build alt/ref allele Rageds via _build_allele_layout from the flat allele buffers -return RaggedVariants(**kw) # alt/ref passed as Ragged char arrays +return RaggedVariants(**kw) # alt/ref passed as Ragged char arrays ``` Adapt to the file's actual field layout (allele fields vs numeric `_Flat`s). The constructor accepts char-`Ragged` alt/ref and numeric `Ragged` fields. diff --git a/docs/superpowers/plans/2026-06-23-raggedvariants-subclass-ragged.md b/docs/superpowers/plans/2026-06-23-raggedvariants-subclass-ragged.md index 72a5bb51..e89a2463 100644 --- a/docs/superpowers/plans/2026-06-23-raggedvariants-subclass-ragged.md +++ b/docs/superpowers/plans/2026-06-23-raggedvariants-subclass-ragged.md @@ -121,7 +121,9 @@ def test_no_rag_composition_attribute(): assert rv._layout is not None # holds the record layout directly -@pytest.mark.parametrize("key", [0, slice(0, 2), np.array([1, 0])], ids=["int", "slice", "fancy"]) +@pytest.mark.parametrize( + "key", [0, slice(0, 2), np.array([1, 0])], ids=["int", "slice", "fancy"] +) def test_positional_indexing_preserves_subclass(key): rv = _rv() out = rv[key] @@ -130,8 +132,8 @@ def test_positional_indexing_preserves_subclass(key): def test_int_index_collapses_leading_axis(): rv = _rv() # (2, 2, ~v) - assert rv[0].shape == (2, None) # int collapses batch -> (ploidy, ~v) - assert rv[0:2].shape == (2, 2, None) # slice keeps batch (ploidy preserved) + assert rv[0].shape == (2, None) # int collapses batch -> (ploidy, ~v) + assert rv[0:2].shape == (2, 2, None) # slice keeps batch (ploidy preserved) def test_string_key_returns_base_ragged(): @@ -161,11 +163,15 @@ def test_extra_field_via_getattr(): alt = _rv()["alt"] start = _rv()["start"] af = Ragged.from_offsets( - np.arange(4, dtype=np.float32), (2, 2, None), np.array([0, 1, 2, 3, 4], np.int64) + np.arange(4, dtype=np.float32), + (2, 2, None), + np.array([0, 1, 2, 3, 4], np.int64), ) rv = RaggedVariants(alt=alt.to_strings(), start=start, ilen=_rv()["ilen"], AF=af) assert "AF" in rv.fields - np.testing.assert_array_equal(np.asarray(rv.AF.data), np.arange(4, dtype=np.float32)) + np.testing.assert_array_equal( + np.asarray(rv.AF.data), np.arange(4, dtype=np.float32) + ) with pytest.raises(AttributeError): _ = rv.not_a_field ``` @@ -296,8 +302,10 @@ def test_reshape_ragged_for_chunk_leaves_raggedvariants_untouched(): from genvarloader._double_buffered_loader import _reshape_ragged_for_chunk alt = Ragged.from_offsets( - np.frombuffer(b"ACGT", dtype="S1").copy(), (2, 1, None), - np.array([0, 1, 2], np.int64), str_offsets=np.array([0, 2, 4], np.int64), + np.frombuffer(b"ACGT", dtype="S1").copy(), + (2, 1, None), + np.array([0, 1, 2], np.int64), + str_offsets=np.array([0, 2, 4], np.int64), ).to_strings() start = Ragged.from_offsets( np.arange(2, dtype=np.int32), (2, 1, None), np.array([0, 1, 2], np.int64) diff --git a/docs/superpowers/plans/2026-06-23-rust-migration-phase-0-foundation.md b/docs/superpowers/plans/2026-06-23-rust-migration-phase-0-foundation.md index 719dd558..7b4e1860 100644 --- a/docs/superpowers/plans/2026-06-23-rust-migration-phase-0-foundation.md +++ b/docs/superpowers/plans/2026-06-23-rust-migration-phase-0-foundation.md @@ -165,9 +165,7 @@ def get(name: str) -> Callable: if backend is None: backend = entry["default"] # type: ignore[assignment] elif backend not in ("numba", "rust"): - raise ValueError( - f"GVL_BACKEND must be 'numba' or 'rust', got {backend!r}" - ) + raise ValueError(f"GVL_BACKEND must be 'numba' or 'rust', got {backend!r}") return entry[backend] # type: ignore[return-value] @@ -392,9 +390,7 @@ register( ) -def splits_sum_le_value( - arr: NDArray[np.number], max_value: float -) -> NDArray[np.intp]: +def splits_sum_le_value(arr: NDArray[np.number], max_value: float) -> NDArray[np.intp]: """Greedy split offsets for groups summing to no more than ``max_value``. Dispatches to the numba or Rust backend via :mod:`genvarloader._dispatch`. @@ -522,7 +518,13 @@ Add a throwaway sanity test, run it, then delete it: def test_harness_detects_mismatch(): import numpy as np from genvarloader import _dispatch - _dispatch.register("bad", numba=lambda a: np.array([1]), rust=lambda a: np.array([2]), default="numba") + + _dispatch.register( + "bad", + numba=lambda a: np.array([1]), + rust=lambda a: np.array([2]), + default="numba", + ) with pytest.raises(AssertionError): assert_kernel_parity("bad", np.array([0])) ``` @@ -727,8 +729,8 @@ def main() -> None: # reuse those exact inputs (see tests/data/generate_1kg_ground_truth.py). from tests.data import generate_1kg_ground_truth as g1k - bcf = g1k.ONE_KG_DIR / "filtered.bcf" # confirm the actual filename in g1k - bed = g1k.ONE_KG_DIR / "regions.bed" # confirm the actual filename in g1k + bcf = g1k.ONE_KG_DIR / "filtered.bcf" # confirm the actual filename in g1k + bed = g1k.ONE_KG_DIR / "regions.bed" # confirm the actual filename in g1k if not bcf.exists(): raise SystemExit("Run `pixi run -e dev gen-1kg` first to build 1kg inputs.") diff --git a/docs/superpowers/plans/2026-06-23-seqpro-core-shared-ragged-crate.md b/docs/superpowers/plans/2026-06-23-seqpro-core-shared-ragged-crate.md index 048b6e5a..c8e96e6d 100644 --- a/docs/superpowers/plans/2026-06-23-seqpro-core-shared-ragged-crate.md +++ b/docs/superpowers/plans/2026-06-23-seqpro-core-shared-ragged-crate.md @@ -394,6 +394,7 @@ In `to_padded`, replace the kernel call (currently `_to_padded_copy(data_u1, off ```python from seqpro.seqpro import _ragged_to_padded # type: ignore[missing-import] # rust + _ragged_to_padded(data_u1, offsets, out_u1, itemsize, out_len) ``` @@ -545,6 +546,7 @@ In `reverse_complement`, replace `_reverse_complement_ragged(u1, offsets, comp_l ```python from seqpro.seqpro import _ragged_reverse_complement # type: ignore[missing-import] # rust + _ragged_reverse_complement(u1, offsets, comp_lut, mask_flat) ``` @@ -691,8 +693,11 @@ from genvarloader._ragged import to_padded as gvl_to_padded def test_gvl_to_padded_matches_seqpro(dtype, pad, rows): offsets = np.concatenate([[0], np.cumsum(rows)]).astype(np.int64) n = int(offsets[-1]) - data = (np.arange(n, dtype=np.int64) % 4).astype(dtype) if dtype != "S1" \ + data = ( + (np.arange(n, dtype=np.int64) % 4).astype(dtype) + if dtype != "S1" else np.frombuffer(b"ACGT" * (n // 4 + 1), dtype="S1")[:n] + ) rag = Ragged.from_offsets(np.ascontiguousarray(data), (len(rows),), offsets) np.testing.assert_array_equal(gvl_to_padded(rag, pad), sp_to_padded(rag, pad)) ``` @@ -716,7 +721,9 @@ def to_padded(rag: Ragged[RDTYPE], pad_value: Any) -> NDArray[RDTYPE]: the shared seqpro-core kernel (Rust->Rust, no Python-seqpro round-trip). """ if rag._is_record: - raise NotImplementedError("to_padded is not defined on record-layout Ragged arrays.") + raise NotImplementedError( + "to_padded is not defined on record-layout Ragged arrays." + ) rag_dim = rag.rag_dim if any(d is not None for d in rag.shape[rag_dim + 1 :]): raise ValueError( diff --git a/docs/superpowers/plans/2026-06-24-issue-242-intervals-track-jitter-clip.md b/docs/superpowers/plans/2026-06-24-issue-242-intervals-track-jitter-clip.md index e5fdf796..210e1c8c 100644 --- a/docs/superpowers/plans/2026-06-24-issue-242-intervals-track-jitter-clip.md +++ b/docs/superpowers/plans/2026-06-24-issue-242-intervals-track-jitter-clip.md @@ -287,9 +287,7 @@ to: Change the first-start initialization (lines ~34-36): ```python - cur = qstart + draw( - st.integers(min_value=0, max_value=10) - ) # first start >= qstart +cur = qstart + draw(st.integers(min_value=0, max_value=10)) # first start >= qstart ``` to: ```python diff --git a/docs/superpowers/plans/2026-06-24-phase-3-closeout.md b/docs/superpowers/plans/2026-06-24-phase-3-closeout.md index 4b52920a..c302bb03 100644 --- a/docs/superpowers/plans/2026-06-24-phase-3-closeout.md +++ b/docs/superpowers/plans/2026-06-24-phase-3-closeout.md @@ -199,9 +199,11 @@ pytestmark = pytest.mark.parity def test_reference_fetch_parity(reference, monkeypatch): - ref = _ref_mod.Reference.from_path_and_contigs(reference, None) \ - if hasattr(_ref_mod.Reference, "from_path_and_contigs") \ + ref = ( + _ref_mod.Reference.from_path_and_contigs(reference, None) + if hasattr(_ref_mod.Reference, "from_path_and_contigs") else _ref_mod.Reference.from_path(reference) + ) contigs = ref.contigs[:1] starts = np.array([0], dtype=np.int64) ends = np.array([50], dtype=np.int64) @@ -226,9 +228,7 @@ def test_reference_fetch_parity(reference, monkeypatch): _dispatch._REGISTRY["get_reference"] = orig assert rust_calls > 0, "rust get_reference never invoked via fetch — vacuous" - np.testing.assert_array_equal( - np.asarray(out_numba.data), np.asarray(out_rust.data) - ) + np.testing.assert_array_equal(np.asarray(out_numba.data), np.asarray(out_rust.data)) np.testing.assert_array_equal( np.asarray(out_numba.offsets, np.int64), np.asarray(out_rust.offsets, np.int64), @@ -249,21 +249,19 @@ Expected: FAIL — `rust get_reference never invoked via fetch` (fetch currently In `_reference.py`, replace the kernel-selection block inside `fetch` (currently lines 135-148) with a call to the dispatched `get_reference`, assembling a `(n,3)` regions array: ```python - lengths = ends - starts - offsets = lengths_to_offsets(lengths) - regions = np.stack( - [ - np.asarray(c_idxs, np.int32), - np.asarray(starts, np.int32), - np.asarray(ends, np.int32), - ], - axis=1, - ) - seqs = get_reference( - regions, offsets, self.reference, self.offsets, int(self.pad_char) - ) - seqs = Ragged.from_offsets(seqs.view("S1"), (len(contigs), None), offsets) - return seqs +lengths = ends - starts +offsets = lengths_to_offsets(lengths) +regions = np.stack( + [ + np.asarray(c_idxs, np.int32), + np.asarray(starts, np.int32), + np.asarray(ends, np.int32), + ], + axis=1, +) +seqs = get_reference(regions, offsets, self.reference, self.offsets, int(self.pad_char)) +seqs = Ragged.from_offsets(seqs.view("S1"), (len(contigs), None), offsets) +return seqs ``` (`get_reference` is defined later in the same module; it is module-level, so the forward reference resolves at call time.) @@ -333,24 +331,25 @@ Co-Authored-By: Claude Opus 4.8 " In `tests/parity/test_haplotypes_dataset_parity.py::test_annotated_haplotypes_mode_dataset_parity`, change the spy from the dispatched `reconstruct_haplotypes_from_sparse` to the new module-level fused entry, mirroring `test_haplotypes_mode_dataset_parity` (which spies `_haps_mod.reconstruct_haplotypes_fused`): ```python - import genvarloader._dataset._haps as _haps_mod - orig_fused = _haps_mod.reconstruct_annotated_haplotypes_fused - calls = {"n": 0} +import genvarloader._dataset._haps as _haps_mod - def _spy_fused(*a, **k): - calls["n"] += 1 - return orig_fused(*a, **k) +orig_fused = _haps_mod.reconstruct_annotated_haplotypes_fused +calls = {"n": 0} - monkeypatch.setattr( - _haps_mod, "reconstruct_annotated_haplotypes_fused", _spy_fused - ) - monkeypatch.setenv("GVL_BACKEND", "rust") - out_rust = ds[:, :] - rust_call_count = calls["n"] - monkeypatch.setenv("GVL_BACKEND", "numba") - out_numba = ds[:, :] - assert calls["n"] == rust_call_count, "fused spy fired during numba read" - assert calls["n"] > 0, "rust annotated fused entry never invoked — vacuous" + +def _spy_fused(*a, **k): + calls["n"] += 1 + return orig_fused(*a, **k) + + +monkeypatch.setattr(_haps_mod, "reconstruct_annotated_haplotypes_fused", _spy_fused) +monkeypatch.setenv("GVL_BACKEND", "rust") +out_rust = ds[:, :] +rust_call_count = calls["n"] +monkeypatch.setenv("GVL_BACKEND", "numba") +out_numba = ds[:, :] +assert calls["n"] == rust_call_count, "fused spy fired during numba read" +assert calls["n"] > 0, "rust annotated fused entry never invoked — vacuous" ``` Keep the existing three-array byte-identical comparison (`_compare_ragged_bytes` + two `_compare_ragged_int`). @@ -444,6 +443,7 @@ Create `tests/parity/test_spliced_haplotypes_parity.py`. It needs a spliced data ```python """Spliced-haplotypes dataset parity backstop (fused rust splice entry).""" + from __future__ import annotations import numpy as np import pytest @@ -470,9 +470,7 @@ def test_spliced_haplotypes_parity(spliced_gvl, reference, monkeypatch): out_numba = ds[:, :] assert calls["n"] == rc, "fused splice spy fired during numba read" assert calls["n"] > 0, "rust spliced fused entry never invoked — vacuous" - np.testing.assert_array_equal( - np.asarray(out_numba.data), np.asarray(out_rust.data) - ) + np.testing.assert_array_equal(np.asarray(out_numba.data), np.asarray(out_rust.data)) np.testing.assert_array_equal( np.asarray(out_numba.offsets, np.int64), np.asarray(out_rust.offsets, np.int64), diff --git a/docs/superpowers/plans/2026-06-24-rust-migration-phase-2-genotypes-variants.md b/docs/superpowers/plans/2026-06-24-rust-migration-phase-2-genotypes-variants.md index e736d6cd..a65ef0ba 100644 --- a/docs/superpowers/plans/2026-06-24-rust-migration-phase-2-genotypes-variants.md +++ b/docs/superpowers/plans/2026-06-24-rust-migration-phase-2-genotypes-variants.md @@ -421,16 +421,20 @@ Append to `tests/parity/strategies.py`: ```python @st.composite -def _sparse_geno(draw, max_queries=4, max_ploidy=2, max_vars_per_group=5, - max_total_unique=12): +def _sparse_geno( + draw, max_queries=4, max_ploidy=2, max_vars_per_group=5, max_total_unique=12 +): """Shared sparse-genotype layout: returns (geno_offset_idx (q,p) int64, geno_v_idxs int32, geno_offsets (n+1,) int64, v_starts int32, ilens int32, q_starts int32, q_ends int32). geno_offset_idx is arange so each (q,p) row maps to its own offset slice.""" n_unique = draw(st.integers(min_value=1, max_value=max_total_unique)) v_starts = np.sort( - draw(st.lists(st.integers(0, 1000), min_size=n_unique, max_size=n_unique) - .map(np.array)) + draw( + st.lists(st.integers(0, 1000), min_size=n_unique, max_size=n_unique).map( + np.array + ) + ) ).astype(np.int32) ilens = np.array( draw(st.lists(st.integers(-5, 5), min_size=n_unique, max_size=n_unique)), @@ -443,8 +447,9 @@ def _sparse_geno(draw, max_queries=4, max_ploidy=2, max_vars_per_group=5, v_idx_list = [] for c in counts: # sorted variant indices within a group (reconstruction assumes sorted pos) - idxs = sorted(draw(st.lists(st.integers(0, n_unique - 1), - min_size=c, max_size=c))) + idxs = sorted( + draw(st.lists(st.integers(0, n_unique - 1), min_size=c, max_size=c)) + ) v_idx_list.extend(idxs) geno_v_idxs = np.array(v_idx_list, dtype=np.int32) geno_offsets = np.concatenate([[0], np.cumsum(counts)]).astype(np.int64) @@ -453,8 +458,15 @@ def _sparse_geno(draw, max_queries=4, max_ploidy=2, max_vars_per_group=5, draw(st.lists(st.integers(0, 800), min_size=n_q, max_size=n_q)), np.int32 ) q_ends = (q_starts + draw(st.integers(1, 200))).astype(np.int32) - return (geno_offset_idx, geno_v_idxs, geno_offsets, v_starts, ilens, - q_starts, q_ends) + return ( + geno_offset_idx, + geno_v_idxs, + geno_offsets, + v_starts, + ilens, + q_starts, + q_ends, + ) @st.composite @@ -978,7 +990,9 @@ def _gather_rows_numba(geno_offset_idx, geno_offsets, geno_v_idxs): ) -register("gather_rows", numba=_gather_rows_numba, rust=_gather_rows_rust, default="rust") +register( + "gather_rows", numba=_gather_rows_numba, rust=_gather_rows_rust, default="rust" +) ``` 3. Replace the body of the existing `_gather_rows(...)` (line 538) with: @@ -1019,7 +1033,9 @@ def gather_rows_inputs(draw): np.int64, ) twod = draw(st.booleans()) - off = offsets if not twod else np.stack([offsets[:-1], offsets[1:]]).astype(np.int64) + off = ( + offsets if not twod else np.stack([offsets[:-1], offsets[1:]]).astype(np.int64) + ) return (goi, off, data) ``` @@ -1160,7 +1176,12 @@ Expected: PASS. In `_flat_variants.py`: add `from ..genvarloader import gather_alleles as _gather_alleles_rust`; rename njit to `_gather_alleles_numba`; add a thin dispatch wrapper named `_gather_alleles` (preserving the existing internal call name) + register: ```python -register("gather_alleles", numba=_gather_alleles_numba, rust=_gather_alleles_rust, default="rust") +register( + "gather_alleles", + numba=_gather_alleles_numba, + rust=_gather_alleles_rust, + default="rust", +) def _gather_alleles(v_idxs, allele_bytes, allele_offsets): @@ -1323,8 +1344,18 @@ Expected: PASS. In `_flat_variants.py`: import both rust fns; rename njit → `_compact_keep_numba`; add: ```python -register("compact_keep_i32", numba=_compact_keep_numba, rust=_compact_keep_i32_rust, default="rust") -register("compact_keep_f32", numba=_compact_keep_numba, rust=_compact_keep_f32_rust, default="rust") +register( + "compact_keep_i32", + numba=_compact_keep_numba, + rust=_compact_keep_i32_rust, + default="rust", +) +register( + "compact_keep_f32", + numba=_compact_keep_numba, + rust=_compact_keep_f32_rust, + default="rust", +) def _compact_keep(v_idxs, row_offsets, keep): @@ -1332,8 +1363,12 @@ def _compact_keep(v_idxs, row_offsets, keep): row_offsets = np.ascontiguousarray(row_offsets, np.int64) keep = np.ascontiguousarray(keep, np.bool_) if np.issubdtype(values.dtype, np.floating): - return get("compact_keep_f32")(values.astype(np.float32, copy=False), row_offsets, keep) - return get("compact_keep_i32")(values.astype(np.int32, copy=False), row_offsets, keep) + return get("compact_keep_f32")( + values.astype(np.float32, copy=False), row_offsets, keep + ) + return get("compact_keep_i32")( + values.astype(np.int32, copy=False), row_offsets, keep + ) ``` If Step 1 found a float64 dosage/ccf dtype, the `.astype(np.float32)` would lose precision and break parity — in that case add a `compact_keep_f64` core/wrapper and route float64 to it instead of down-casting. The numba reference preserves the input dtype, so the parity test (which feeds the same dtype to both) will catch any mismatch. @@ -1351,13 +1386,22 @@ def compact_keep_inputs(draw, dtype): total = int(row_offsets[-1]) if np.issubdtype(np.dtype(dtype), np.floating): values = np.array( - draw(st.lists(st.floats(width=32, allow_nan=False, allow_infinity=False), - min_size=total, max_size=total)), dtype) + draw( + st.lists( + st.floats(width=32, allow_nan=False, allow_infinity=False), + min_size=total, + max_size=total, + ) + ), + dtype, + ) else: values = np.array( - draw(st.lists(st.integers(0, 1000), min_size=total, max_size=total)), dtype) + draw(st.lists(st.integers(0, 1000), min_size=total, max_size=total)), dtype + ) keep = np.array( - draw(st.lists(st.booleans(), min_size=total, max_size=total)), np.bool_) + draw(st.lists(st.booleans(), min_size=total, max_size=total)), np.bool_ + ) return (values, row_offsets, keep) ``` @@ -1647,7 +1691,8 @@ def test_variants_getitem_parity_and_kernels_invoked(variants_dataset, monkeypat # rust + spy monkeypatch.setenv("GVL_BACKEND", "rust") monkeypatch.setattr( - _flat_variants, "get", + _flat_variants, + "get", lambda name: spy if name == "gather_rows" else _dispatch.get(name), ) out_rust = _run_variants_getitem(variants_dataset) diff --git a/docs/superpowers/plans/2026-06-24-rust-migration-phase-3.md b/docs/superpowers/plans/2026-06-24-rust-migration-phase-3.md index 831208e9..2bba9e15 100644 --- a/docs/superpowers/plans/2026-06-24-rust-migration-phase-3.md +++ b/docs/superpowers/plans/2026-06-24-rust-migration-phase-3.md @@ -209,6 +209,7 @@ def get_reference_inputs(draw): """Generate (regions, out_offsets, reference, ref_offsets, pad_char, parallel) with regions whose [start,end) windows may run off either contig edge.""" import numpy as np + n_contigs = draw(st.integers(1, 3)) contig_lens = [draw(st.integers(1, 40)) for _ in range(n_contigs)] ref_offsets = np.concatenate([[0], np.cumsum(contig_lens)]).astype(np.int64) @@ -249,7 +250,13 @@ pytestmark = pytest.mark.parity def test_get_reference_parity(inputs): regions, out_offsets, reference, ref_offsets, pad_char, parallel = inputs assert_kernel_parity( - "get_reference", regions, out_offsets, reference, ref_offsets, pad_char, parallel + "get_reference", + regions, + out_offsets, + reference, + ref_offsets, + pad_char, + parallel, ) ``` @@ -362,13 +369,17 @@ from genvarloader import _genvarloader as _gvl_rust # match existing import ali from genvarloader._dispatch import register, get -def _get_reference_numba(regions, out_offsets, reference, ref_offsets, pad_char, parallel): +def _get_reference_numba( + regions, out_offsets, reference, ref_offsets, pad_char, parallel +): out = np.empty(out_offsets[-1], np.uint8) kernel = _get_reference_par if parallel else _get_reference_ser return kernel(regions, out_offsets, reference, ref_offsets, pad_char, out) -def _get_reference_rust(regions, out_offsets, reference, ref_offsets, pad_char, parallel): +def _get_reference_rust( + regions, out_offsets, reference, ref_offsets, pad_char, parallel +): return _gvl_rust.get_reference( np.ascontiguousarray(regions, np.int32), np.ascontiguousarray(out_offsets, np.int64), @@ -379,12 +390,19 @@ def _get_reference_rust(regions, out_offsets, reference, ref_offsets, pad_char, ) -register("get_reference", numba=_get_reference_numba, rust=_get_reference_rust, default="rust") +register( + "get_reference", + numba=_get_reference_numba, + rust=_get_reference_rust, + default="rust", +) def get_reference(regions, out_offsets, reference, ref_offsets, pad_char): parallel = should_parallelize(int(out_offsets[-1])) - return get("get_reference")(regions, out_offsets, reference, ref_offsets, pad_char, parallel) + return get("get_reference")( + regions, out_offsets, reference, ref_offsets, pad_char, parallel + ) ``` Note: `parallel` is computed in the Python entry (not inside the kernels) so both backends receive the identical flag — this keeps the numba twin byte-identical to today's behavior and makes the strategy's `parallel` field meaningful. @@ -585,11 +603,16 @@ pub fn hash4(a: u64, b: u64, c: u64, d: u64) -> u64 { - [ ] **Step 2: Add a direct numba-vs-rust PRNG parity test.** Temporarily expose the rust `hash4` via a `#[pyfunction]` (e.g. `ffi::_debug_hash4`) and a numba `_hash4` accessor in `_tracks.py`, then over a hypothesis grid of `(a,b,c,d)` `uint64` quadruples assert `rust_hash4(a,b,c,d) == int(_hash4(a,b,c,d))`. This is the single most important guard for FlankSample byte-identity. ```python -@given(st.integers(0, 2**64 - 1), st.integers(0, 2**64 - 1), - st.integers(0, 2**64 - 1), st.integers(0, 2**64 - 1)) +@given( + st.integers(0, 2**64 - 1), + st.integers(0, 2**64 - 1), + st.integers(0, 2**64 - 1), + st.integers(0, 2**64 - 1), +) def test_hash4_parity(a, b, c, d): from genvarloader._dataset._tracks import _hash4 import numpy as np + exp = int(_hash4(np.uint64(a), np.uint64(b), np.uint64(c), np.uint64(d))) assert _gvl_rust._debug_hash4(a, b, c, d) == exp ``` diff --git a/docs/superpowers/plans/2026-06-25-rust-variant-rc-fold.md b/docs/superpowers/plans/2026-06-25-rust-variant-rc-fold.md index e1b20079..57f4a1dd 100644 --- a/docs/superpowers/plans/2026-06-25-rust-variant-rc-fold.md +++ b/docs/superpowers/plans/2026-06-25-rust-variant-rc-fold.md @@ -250,8 +250,11 @@ def _allele_batch(draw): lens = [draw(st.integers(0, 5)) for _ in range(n_alleles)] seq_offsets = np.concatenate([[0], np.cumsum(lens)]).astype(np.int64) total = int(seq_offsets[-1]) - data = _ACGTN[draw(st.lists(st.integers(0, 4), min_size=total, max_size=total))] \ - if total else np.zeros(0, np.uint8) + data = ( + _ACGTN[draw(st.lists(st.integers(0, 4), min_size=total, max_size=total))] + if total + else np.zeros(0, np.uint8) + ) data = np.ascontiguousarray(data, np.uint8) mask = np.array([draw(st.booleans()) for _ in range(n_rows)], np.bool_) return data, seq_offsets, var_offsets, mask @@ -596,7 +599,9 @@ def test_neg_strand_variants_rc_parity_and_kernel_invoked( ds_dir = build_strand_mixed_dataset(tmp_path, synthetic_case.svar_path) ref = gvl.Reference.from_path(synthetic_case.ref_path, in_memory=False) - ds = gvl.Dataset.open(ds_dir, reference=ref).with_tracks(False).with_seqs("variants") + ds = ( + gvl.Dataset.open(ds_dir, reference=ref).with_tracks(False).with_seqs("variants") + ) # Non-vacuity: fixture must carry −strand regions (rc_neg defaults True). assert np.any(ds._full_regions[:, 3] == -1), "fixture has no −strand regions" diff --git a/docs/superpowers/plans/2026-06-25-target7-variant-windows-rust-assembly.md b/docs/superpowers/plans/2026-06-25-target7-variant-windows-rust-assembly.md index 9353664f..5bbbbd64 100644 --- a/docs/superpowers/plans/2026-06-25-target7-variant-windows-rust-assembly.md +++ b/docs/superpowers/plans/2026-06-25-target7-variant-windows-rust-assembly.md @@ -1099,8 +1099,15 @@ def _assemble_variant_buffers_numba( out["ref"] = (rw.data, rw.seq_offsets) if alt_mode == 1: aw = compute_alt_window( - ref_shim, v_contigs, starts_v, ilens_v, alt_data, alt_seq_off, - flank_len, lut_arr, row_offsets, + ref_shim, + v_contigs, + starts_v, + ilens_v, + alt_data, + alt_seq_off, + flank_len, + lut_arr, + row_offsets, ) out["alt_window"] = (aw.data, aw.seq_offsets) elif alt_mode == 2: @@ -1157,7 +1164,9 @@ def _assemble_variant_buffers_rust( np.ascontiguousarray(alt_global, np.uint8), np.ascontiguousarray(alt_off_global, np.int64), None if ref_global is None else np.ascontiguousarray(ref_global, np.uint8), - None if ref_off_global is None else np.ascontiguousarray(ref_off_global, np.int64), + None + if ref_off_global is None + else np.ascontiguousarray(ref_off_global, np.int64), bool(want_ref_bytes), bool(want_flank), int(ref_mode), @@ -1213,8 +1222,12 @@ def assert_kernel_parity_dict(name: str, *inputs) -> None: for i, (a, b) in enumerate(zip(nt, rt)): a = np.asarray(a) b = np.asarray(b) - assert a.dtype == b.dtype, f"{name}[{key}][{i}]: dtype {a.dtype} != {b.dtype}" - assert a.shape == b.shape, f"{name}[{key}][{i}]: shape {a.shape} != {b.shape}" + assert a.dtype == b.dtype, ( + f"{name}[{key}][{i}]: dtype {a.dtype} != {b.dtype}" + ) + assert a.shape == b.shape, ( + f"{name}[{key}][{i}]: shape {a.shape} != {b.shape}" + ) np.testing.assert_array_equal(a, b) ``` @@ -1290,101 +1303,107 @@ Concretely, after the scalar/dosage/custom fields are built into `fields` (keep Replace the windows branch (`if regions is not None and issubclass(haps.kind, _FlatVariantWindows) and opt is not None:` ... `return win`) with: ```python - opt = haps.window_opt - if ( - regions is not None - and issubclass(haps.kind, _FlatVariantWindows) - and opt is not None - ): - L = opt.flank_length - ref_mode = 1 if opt.ref == "window" else 2 - alt_mode = 1 if opt.alt == "window" else 2 - bufs = get("assemble_variant_buffers")( - 1, # windows mode - v_idxs, - row_offsets, - stat.alt_alleles, - stat.alt_offsets, - ref_global, - ref_off_global, - False, # want_ref_bytes (windows mode emits tokens, not raw bytes) - False, # want_flank - ref_mode, - alt_mode, - L, - haps.token_lut, - v_contigs, - stat.v_starts, - stat.ilens, - stat.ref, # reference genome buffer - stat.ref_offsets, # contig offsets - haps.reference.pad_char, - ) - wshape = (b, eff_ploidy, None, None) - wfields = {k: v for k, v in fields.items() if k not in ("alt", "ref")} - win = _FlatVariantWindows(wfields) - for name, (data, seq_off) in bufs.items(): - fw = _FlatWindow(data, np.asarray(seq_off, np.int64), row_offsets, wshape) - setattr(win, name, fw) - if haps.dummy_variant is not None: - win = win.fill_empty_groups( - haps.dummy_variant, unk=haps.unknown_token, flank_length=L - ) - return win -``` - -- [ ] **Step 3: Build the plain-variants alt/ref + flank result from the dict** - -Replace the inline alt/ref allele gather and the flank ride-along so the plain-variants path also goes through the kernel. Where the code currently does `fields["alt"] = _FlatAlleles(...)` and `fields["ref"] = _FlatAlleles(...)`, and the later `if haps.flank_length and ...: compute_flank_tokens(...)` block, replace with a single call after the scalar fields are assembled: - -```python - want_flank = bool( - haps.flank_length and haps.token_lut is not None and regions is not None - ) - L = haps.flank_length or 0 +opt = haps.window_opt +if ( + regions is not None + and issubclass(haps.kind, _FlatVariantWindows) + and opt is not None +): + L = opt.flank_length + ref_mode = 1 if opt.ref == "window" else 2 + alt_mode = 1 if opt.alt == "window" else 2 bufs = get("assemble_variant_buffers")( - 0, # variants mode + 1, # windows mode v_idxs, row_offsets, stat.alt_alleles, stat.alt_offsets, ref_global, ref_off_global, - ref_present, # want_ref_bytes - want_flank, - 0, # ref_mode (unused in variants mode) - 0, # alt_mode (unused) + False, # want_ref_bytes (windows mode emits tokens, not raw bytes) + False, # want_flank + ref_mode, + alt_mode, L, haps.token_lut, v_contigs, stat.v_starts, stat.ilens, - stat.ref if stat.ref is not None else np.zeros(0, np.uint8), - stat.ref_offsets if stat.ref_offsets is not None else np.zeros(1, np.int64), - haps.reference.pad_char if haps.reference is not None else 0, - ) - alt_data, alt_seq_off = bufs["alt"] - fields["alt"] = _FlatAlleles( - np.asarray(alt_data, np.uint8), np.asarray(alt_seq_off, np.int64), row_offsets, shape + stat.ref, # reference genome buffer + stat.ref_offsets, # contig offsets + haps.reference.pad_char, ) - if "ref" in bufs: - ref_data, ref_seq_off = bufs["ref"] - fields["ref"] = _FlatAlleles( - np.asarray(ref_data, np.uint8), np.asarray(ref_seq_off, np.int64), row_offsets, shape + wshape = (b, eff_ploidy, None, None) + wfields = {k: v for k, v in fields.items() if k not in ("alt", "ref")} + win = _FlatVariantWindows(wfields) + for name, (data, seq_off) in bufs.items(): + fw = _FlatWindow(data, np.asarray(seq_off, np.int64), row_offsets, wshape) + setattr(win, name, fw) + if haps.dummy_variant is not None: + win = win.fill_empty_groups( + haps.dummy_variant, unk=haps.unknown_token, flank_length=L ) - flat = _FlatVariants(fields) - if "flank_tokens" in bufs: - from .._flat import _Flat + return win +``` - tok, off = bufs["flank_tokens"] - flat.flank_tokens = _Flat.from_offsets( - tok, (b, eff_ploidy, None, 2 * L), np.asarray(off, np.int64) - ) +- [ ] **Step 3: Build the plain-variants alt/ref + flank result from the dict** - if haps.dummy_variant is not None: - flat = flat.fill_empty_groups(haps.dummy_variant, unk=haps.unknown_token) +Replace the inline alt/ref allele gather and the flank ride-along so the plain-variants path also goes through the kernel. Where the code currently does `fields["alt"] = _FlatAlleles(...)` and `fields["ref"] = _FlatAlleles(...)`, and the later `if haps.flank_length and ...: compute_flank_tokens(...)` block, replace with a single call after the scalar fields are assembled: - return flat +```python +want_flank = bool( + haps.flank_length and haps.token_lut is not None and regions is not None +) +L = haps.flank_length or 0 +bufs = get("assemble_variant_buffers")( + 0, # variants mode + v_idxs, + row_offsets, + stat.alt_alleles, + stat.alt_offsets, + ref_global, + ref_off_global, + ref_present, # want_ref_bytes + want_flank, + 0, # ref_mode (unused in variants mode) + 0, # alt_mode (unused) + L, + haps.token_lut, + v_contigs, + stat.v_starts, + stat.ilens, + stat.ref if stat.ref is not None else np.zeros(0, np.uint8), + stat.ref_offsets if stat.ref_offsets is not None else np.zeros(1, np.int64), + haps.reference.pad_char if haps.reference is not None else 0, +) +alt_data, alt_seq_off = bufs["alt"] +fields["alt"] = _FlatAlleles( + np.asarray(alt_data, np.uint8), + np.asarray(alt_seq_off, np.int64), + row_offsets, + shape, +) +if "ref" in bufs: + ref_data, ref_seq_off = bufs["ref"] + fields["ref"] = _FlatAlleles( + np.asarray(ref_data, np.uint8), + np.asarray(ref_seq_off, np.int64), + row_offsets, + shape, + ) +flat = _FlatVariants(fields) +if "flank_tokens" in bufs: + from .._flat import _Flat + + tok, off = bufs["flank_tokens"] + flat.flank_tokens = _Flat.from_offsets( + tok, (b, eff_ploidy, None, 2 * L), np.asarray(off, np.int64) + ) + +if haps.dummy_variant is not None: + flat = flat.fill_empty_groups(haps.dummy_variant, unk=haps.unknown_token) + +return flat ``` > IMPORTANT ordering: the `fields` dict insertion order determines downstream wrapping; today `alt` is inserted before `start`/`ref`/etc. Preserve the existing field order — build `fields["alt"]` placeholder position by keeping the scalar block as-is and only swapping the alt/ref *values* to come from `bufs`. If the original code inserted `alt` first, keep `alt` first (move the `bufs["alt"]` assignment up to where `fields["alt"]` was originally set, not appended at the end). Verify with `RaggedVariants` field order in a parity run (Task 8). @@ -1462,9 +1481,9 @@ def _globals(): alt_bytes = np.frombuffer(b"ACGT", np.uint8) # alt alleles: v0="A", v1="CG", v2="T" alt_data = np.frombuffer(b"ACGT", np.uint8) - alt_data = np.frombuffer(b"A" b"CG" b"T", np.uint8) + alt_data = np.frombuffer(b"ACGT", np.uint8) alt_off = np.array([0, 1, 3, 4], np.int64) - ref_data = np.frombuffer(b"C" b"G" b"AA", np.uint8) + ref_data = np.frombuffer(b"CGAA", np.uint8) ref_off = np.array([0, 1, 2, 4], np.int64) v_starts = np.array([5, 12, 20], np.int32) ilens = np.array([0, -1, 1], np.int32) # SNP, 1bp del, 1bp ins @@ -1484,14 +1503,31 @@ def test_windows_mode_matrix(tok_dtype, ref_mode, alt_mode): assert_kernel_parity_dict( "assemble_variant_buffers", 1, # windows - v_idxs, row_offsets, alt_data, alt_off, ref_data, ref_off, - False, False, ref_mode, alt_mode, 2, lut, v_contigs, v_starts, ilens, - ref, ref_offsets, ord("N"), + v_idxs, + row_offsets, + alt_data, + alt_off, + ref_data, + ref_off, + False, + False, + ref_mode, + alt_mode, + 2, + lut, + v_contigs, + v_starts, + ilens, + ref, + ref_offsets, + ord("N"), ) @pytest.mark.parametrize("tok_dtype", [np.uint8, np.int32]) -@pytest.mark.parametrize("want_ref,want_flank", [(False, False), (True, False), (False, True), (True, True)]) +@pytest.mark.parametrize( + "want_ref,want_flank", [(False, False), (True, False), (False, True), (True, True)] +) def test_variants_mode_matrix(tok_dtype, want_ref, want_flank): ref, ref_offsets = _reference() alt_data, alt_off, ref_data, ref_off, v_starts, ilens = _globals() @@ -1502,9 +1538,24 @@ def test_variants_mode_matrix(tok_dtype, want_ref, want_flank): assert_kernel_parity_dict( "assemble_variant_buffers", 0, # variants - v_idxs, row_offsets, alt_data, alt_off, ref_data, ref_off, - want_ref, want_flank, 0, 0, 2, lut, v_contigs, v_starts, ilens, - ref, ref_offsets, ord("N"), + v_idxs, + row_offsets, + alt_data, + alt_off, + ref_data, + ref_off, + want_ref, + want_flank, + 0, + 0, + 2, + lut, + v_contigs, + v_starts, + ilens, + ref, + ref_offsets, + ord("N"), ) @@ -1520,9 +1571,24 @@ def test_empty_selection(mode, ref_mode, alt_mode): assert_kernel_parity_dict( "assemble_variant_buffers", mode, - v_idxs, row_offsets, alt_data, alt_off, ref_data, ref_off, - False, (mode == 0), ref_mode, alt_mode, 2, lut, v_contigs, v_starts, ilens, - ref, ref_offsets, ord("N"), + v_idxs, + row_offsets, + alt_data, + alt_off, + ref_data, + ref_off, + False, + (mode == 0), + ref_mode, + alt_mode, + 2, + lut, + v_contigs, + v_starts, + ilens, + ref, + ref_offsets, + ord("N"), ) ``` diff --git a/docs/superpowers/plans/2026-06-25-zero-copy-scale-safe-readpath.md b/docs/superpowers/plans/2026-06-25-zero-copy-scale-safe-readpath.md index 40f2eb87..a8b6e66e 100644 --- a/docs/superpowers/plans/2026-06-25-zero-copy-scale-safe-readpath.md +++ b/docs/superpowers/plans/2026-06-25-zero-copy-scale-safe-readpath.md @@ -718,7 +718,9 @@ def _downgrade_to_aos(path): rec["start"] = starts rec["end"] = ends rec["value"] = values - out = np.memmap(d / "intervals.npy", dtype=INTERVAL_DTYPE, mode="w+", shape=rec.shape) + out = np.memmap( + d / "intervals.npy", dtype=INTERVAL_DTYPE, mode="w+", shape=rec.shape + ) out[:] = rec out.flush() del starts, ends, values, out @@ -732,7 +734,9 @@ def _downgrade_to_aos(path): def test_round_trip_byte_identical(track_dataset_path, reference): - before = gvl.Dataset.open(track_dataset_path, reference=reference).with_tracks("cov")[0, 0] + before = gvl.Dataset.open(track_dataset_path, reference=reference).with_tracks( + "cov" + )[0, 0] before = np.asarray(before).copy() _downgrade_to_aos(track_dataset_path) @@ -743,9 +747,14 @@ def test_round_trip_byte_identical(track_dataset_path, reference): assert (track_dir / "ends.npy").exists() assert (track_dir / "values.npy").exists() assert not (track_dir / "intervals.npy").exists() - assert json.loads((track_dataset_path / "metadata.json").read_text())["format_version"] == "2.0.0" + assert ( + json.loads((track_dataset_path / "metadata.json").read_text())["format_version"] + == "2.0.0" + ) - after = gvl.Dataset.open(track_dataset_path, reference=reference).with_tracks("cov")[0, 0] + after = gvl.Dataset.open(track_dataset_path, reference=reference).with_tracks( + "cov" + )[0, 0] np.testing.assert_array_equal(np.asarray(after), before) @@ -771,7 +780,10 @@ def test_resumable_after_interrupt_before_metadata_bump(track_dataset_path): assert (track_dir / "intervals.npy").exists() # AoS still present gvl.migrate(track_dataset_path) # completes the migration - assert json.loads((track_dataset_path / "metadata.json").read_text())["format_version"] == "2.0.0" + assert ( + json.loads((track_dataset_path / "metadata.json").read_text())["format_version"] + == "2.0.0" + ) assert not (track_dir / "intervals.npy").exists() @@ -783,7 +795,9 @@ def test_cleans_leftover_aos_after_interrupt_before_delete(track_dataset_path): # Re-introduce a leftover AoS file (as if delete was interrupted). starts = np.memmap(track_dir / "starts.npy", dtype=np.int32, mode="r") rec = np.zeros(len(starts), dtype=INTERVAL_DTYPE) - out = np.memmap(track_dir / "intervals.npy", dtype=INTERVAL_DTYPE, mode="w+", shape=rec.shape) + out = np.memmap( + track_dir / "intervals.npy", dtype=INTERVAL_DTYPE, mode="w+", shape=rec.shape + ) out[:] = rec out.flush() del starts, out @@ -930,9 +944,9 @@ from ._dataset._migrate import migrate and insert `"migrate"` into `__all__` (alphabetically, between `"get_splice_bed"` and `"migrate_svar_link"`): ```python - "get_splice_bed", - "migrate", - "migrate_svar_link", +("get_splice_bed",) +("migrate",) +("migrate_svar_link",) ``` - [ ] **Step 5: Run the test to verify it passes** @@ -1053,24 +1067,16 @@ In `python/genvarloader/_dataset/_reconstruct.py`, add the import near the top ( `geno_v_idxs` (`:232-234`): ```python - geno_v_idxs=_ffi_array( - self.haps.genotypes.data, np.int32, "geno_v_idxs" - ), +geno_v_idxs = (_ffi_array(self.haps.genotypes.data, np.int32, "geno_v_idxs"),) ``` `itv_starts` / `itv_ends` / `itv_values` / `itv_offsets` (`:241-250`): ```python - itv_starts=_ffi_array( - intervals.starts.data, np.int32, "itv_starts" - ), - itv_ends=_ffi_array(intervals.ends.data, np.int32, "itv_ends"), - itv_values=_ffi_array( - intervals.values.data, np.float32, "itv_values" - ), - itv_offsets=_ffi_array( - intervals.starts.offsets, np.int64, "itv_offsets" - ), +itv_starts = (_ffi_array(intervals.starts.data, np.int32, "itv_starts"),) +itv_ends = (_ffi_array(intervals.ends.data, np.int32, "itv_ends"),) +itv_values = (_ffi_array(intervals.values.data, np.float32, "itv_values"),) +itv_offsets = (_ffi_array(intervals.starts.offsets, np.int64, "itv_offsets"),) ``` Leave `v_starts` and `ilens` (`:236-239`) as `np.ascontiguousarray(...)` — Task 5 converts those to the cached arrays. Leave `o_idx`, `out_ofsts_per_t`, `regions`, `shifts`, `geno_idx`, `track_ofsts_per_t`, `params`, `keep`, `keep_offsets` as `np.ascontiguousarray(...)` (batch-bounded). @@ -1082,19 +1088,19 @@ In `python/genvarloader/_dataset/_haps.py`, add `from ._utils import _ffi_array` `:796` (plain `reconstruct_haplotypes_fused`): ```python - geno_v_idxs=_ffi_array(self.genotypes.data, np.int32, "geno_v_idxs"), +geno_v_idxs = (_ffi_array(self.genotypes.data, np.int32, "geno_v_idxs"),) ``` `:869` (`reconstruct_haplotypes_spliced_fused`): ```python - geno_v_idxs=_ffi_array(self.genotypes.data, np.int32, "geno_v_idxs"), +geno_v_idxs = (_ffi_array(self.genotypes.data, np.int32, "geno_v_idxs"),) ``` `:958` (`reconstruct_annotated_haplotypes_fused`): ```python - geno_v_idxs=_ffi_array(self.genotypes.data, np.int32, "geno_v_idxs"), +geno_v_idxs = (_ffi_array(self.genotypes.data, np.int32, "geno_v_idxs"),) ``` Leave the sub-linear args (`v_starts`, `ilens`, `alt_alleles`, `alt_offsets`, `ref_`, `ref_offsets`) as `np.ascontiguousarray(...)` for now — Task 5. Leave `regions`, `shifts`, `geno_offset_idx`, `keep`, `keep_offsets`, `permuted_regions`, `flat_shifts`, `flat_geno_offset_idx`, `out_offsets` as `np.ascontiguousarray(...)` (batch-bounded). Leave `_as_starts_stops(self.genotypes.offsets)` untouched. @@ -1147,13 +1153,17 @@ def test_tracks_only_no_memmap_copy(track_dataset_path, reference, _no_memmap_co def test_haps_no_memmap_copy(track_dataset_path, reference, _no_memmap_copies): - ds = gvl.Dataset.open(track_dataset_path, reference=reference).with_seqs("haplotypes") + ds = gvl.Dataset.open(track_dataset_path, reference=reference).with_seqs( + "haplotypes" + ) _ = ds[0, 0] assert _no_memmap_copies == [], f"sample-scale memmap copies: {_no_memmap_copies}" def test_annotated_no_memmap_copy(track_dataset_path, reference, _no_memmap_copies): - ds = gvl.Dataset.open(track_dataset_path, reference=reference).with_seqs("annotated") + ds = gvl.Dataset.open(track_dataset_path, reference=reference).with_seqs( + "annotated" + ) _ = ds[0, 0] assert _no_memmap_copies == [], f"sample-scale memmap copies: {_no_memmap_copies}" ``` @@ -1225,7 +1235,9 @@ from genvarloader._dataset._haps import Haps def _haps(track_dataset_path, reference) -> Haps: - ds = gvl.Dataset.open(track_dataset_path, reference=reference).with_seqs("haplotypes") + ds = gvl.Dataset.open(track_dataset_path, reference=reference).with_seqs( + "haplotypes" + ) seqs = ds._seqs assert isinstance(seqs, Haps) return seqs @@ -1315,27 +1327,23 @@ And add the property (anywhere in the `Haps` class body, e.g. after `__post_init In `python/genvarloader/_dataset/_haps.py`, at the plain fused call (`:797-806`) replace: ```python - v_starts=np.ascontiguousarray(self.variants.start, np.int32), - ilens=np.ascontiguousarray(self.variants.ilen, np.int32), - alt_alleles=np.ascontiguousarray( - self.variants.alt.data.view(np.uint8), np.uint8 - ), - alt_offsets=np.ascontiguousarray( - self.variants.alt.offsets, np.int64 - ), - ref_=np.ascontiguousarray(self.reference.reference, np.uint8), - ref_offsets=np.ascontiguousarray(self.reference.offsets, np.int64), +v_starts = (np.ascontiguousarray(self.variants.start, np.int32),) +ilens = (np.ascontiguousarray(self.variants.ilen, np.int32),) +alt_alleles = (np.ascontiguousarray(self.variants.alt.data.view(np.uint8), np.uint8),) +alt_offsets = (np.ascontiguousarray(self.variants.alt.offsets, np.int64),) +ref_ = (np.ascontiguousarray(self.reference.reference, np.uint8),) +ref_offsets = (np.ascontiguousarray(self.reference.offsets, np.int64),) ``` with: ```python - v_starts=self.ffi_static.v_starts, - ilens=self.ffi_static.ilens, - alt_alleles=self.ffi_static.alt_alleles, - alt_offsets=self.ffi_static.alt_offsets, - ref_=self.ffi_static.ref, - ref_offsets=self.ffi_static.ref_offsets, +v_starts = (self.ffi_static.v_starts,) +ilens = (self.ffi_static.ilens,) +alt_alleles = (self.ffi_static.alt_alleles,) +alt_offsets = (self.ffi_static.alt_offsets,) +ref_ = (self.ffi_static.ref,) +ref_offsets = (self.ffi_static.ref_offsets,) ``` Apply the identical replacement at the spliced fused call (`:870-877`) and the annotated fused call (`:959-970`), matching each call's indentation. (Each of those three sites asserts `self.reference is not None` upstream, so `ffi_static.ref`/`ref_offsets` are non-`None` there.) @@ -1345,17 +1353,15 @@ Apply the identical replacement at the spliced fused call (`:870-877`) and the a In `python/genvarloader/_dataset/_reconstruct.py`, at the `intervals_and_realign_track_fused(...)` call (`:236-239`) replace: ```python - v_starts=np.ascontiguousarray( - self.haps.variants.start, np.int32 - ), - ilens=np.ascontiguousarray(self.haps.variants.ilen, np.int32), +v_starts = (np.ascontiguousarray(self.haps.variants.start, np.int32),) +ilens = (np.ascontiguousarray(self.haps.variants.ilen, np.int32),) ``` with: ```python - v_starts=self.haps.ffi_static.v_starts, - ilens=self.haps.ffi_static.ilens, +v_starts = (self.haps.ffi_static.v_starts,) +ilens = (self.haps.ffi_static.ilens,) ``` - [ ] **Step 6: Run the cache test** diff --git a/docs/superpowers/plans/2026-06-26-rust-migration-phase-5-w3.md b/docs/superpowers/plans/2026-06-26-rust-migration-phase-5-w3.md index ce763c21..25e9abab 100644 --- a/docs/superpowers/plans/2026-06-26-rust-migration-phase-5-w3.md +++ b/docs/superpowers/plans/2026-06-26-rust-migration-phase-5-w3.md @@ -308,93 +308,89 @@ In `python/genvarloader/_dataset/_haps.py`, in the extension-import block (after Replace the current splice-plan block (`python/genvarloader/_dataset/_haps.py:1100-1157`, from the `# ---- splice plan path ----` comment through the final `return haps_rag, annot_v_rag, annot_pos_rag`) with: ```python - # ---- splice plan path ---- - flat_geno_idx, flat_shifts, permuted_regions, keep_perm, keep_offsets_perm = ( - self._permute_request_for_splice(req) +# ---- splice plan path ---- +flat_geno_idx, flat_shifts, permuted_regions, keep_perm, keep_offsets_perm = ( + self._permute_request_for_splice(req) +) +splice_plan = req.splice_plan +per_elem_shape = (splice_plan.permuted_lengths.shape[0], None) +off = splice_plan.permuted_out_offsets + +_backend = os.environ.get("GVL_BACKEND", "rust") +if _backend == "rust": + # Fused path: one FFI crossing. RC is folded in-kernel (sequence bytes + # reverse-complemented, annotation rows reversed), so there is NO Python + # reverse_masked post-pass. to_rc is already in permuted per-element order + # (from _getitem_spliced), and _getitem_spliced treats the rust output as + # already-RC'd (its post-pass is numba-only). + _to_rc_spliced = None if to_rc is None else np.ascontiguousarray(to_rc, np.bool_) + out_buf, annot_v_buf, annot_pos_buf = ( + reconstruct_annotated_haplotypes_spliced_fused( + permuted_regions=np.ascontiguousarray(permuted_regions, np.int32), + flat_shifts=np.ascontiguousarray(flat_shifts.reshape(-1, 1), np.int32), + flat_geno_offset_idx=np.ascontiguousarray( + flat_geno_idx.reshape(-1, 1), np.int64 + ), + out_offsets=np.ascontiguousarray(off, np.int64), + geno_offsets=_as_starts_stops(self.genotypes.offsets), + geno_v_idxs=_ffi_array(self.genotypes.data, np.int32, "geno_v_idxs"), + v_starts=self.ffi_static.v_starts, + ilens=self.ffi_static.ilens, + alt_alleles=self.ffi_static.alt_alleles, + alt_offsets=self.ffi_static.alt_offsets, + ref_=self.ffi_static.ref, + ref_offsets=self.ffi_static.ref_offsets, + pad_char=np.uint8(self.reference.pad_char), + keep=None + if keep_perm is None + else np.ascontiguousarray(keep_perm, np.bool_), + keep_offsets=None + if keep_offsets_perm is None + else np.ascontiguousarray(keep_offsets_perm, np.int64), + to_rc=_to_rc_spliced, ) - splice_plan = req.splice_plan - per_elem_shape = (splice_plan.permuted_lengths.shape[0], None) - off = splice_plan.permuted_out_offsets - - _backend = os.environ.get("GVL_BACKEND", "rust") - if _backend == "rust": - # Fused path: one FFI crossing. RC is folded in-kernel (sequence bytes - # reverse-complemented, annotation rows reversed), so there is NO Python - # reverse_masked post-pass. to_rc is already in permuted per-element order - # (from _getitem_spliced), and _getitem_spliced treats the rust output as - # already-RC'd (its post-pass is numba-only). - _to_rc_spliced = ( - None if to_rc is None else np.ascontiguousarray(to_rc, np.bool_) - ) - out_buf, annot_v_buf, annot_pos_buf = ( - reconstruct_annotated_haplotypes_spliced_fused( - permuted_regions=np.ascontiguousarray(permuted_regions, np.int32), - flat_shifts=np.ascontiguousarray( - flat_shifts.reshape(-1, 1), np.int32 - ), - flat_geno_offset_idx=np.ascontiguousarray( - flat_geno_idx.reshape(-1, 1), np.int64 - ), - out_offsets=np.ascontiguousarray(off, np.int64), - geno_offsets=_as_starts_stops(self.genotypes.offsets), - geno_v_idxs=_ffi_array(self.genotypes.data, np.int32, "geno_v_idxs"), - v_starts=self.ffi_static.v_starts, - ilens=self.ffi_static.ilens, - alt_alleles=self.ffi_static.alt_alleles, - alt_offsets=self.ffi_static.alt_offsets, - ref_=self.ffi_static.ref, - ref_offsets=self.ffi_static.ref_offsets, - pad_char=np.uint8(self.reference.pad_char), - keep=None - if keep_perm is None - else np.ascontiguousarray(keep_perm, np.bool_), - keep_offsets=None - if keep_offsets_perm is None - else np.ascontiguousarray(keep_offsets_perm, np.int64), - to_rc=_to_rc_spliced, - ) - ) - else: - # Numba composed oracle path. RC is applied externally in - # _getitem_spliced (numba branch), so no to_rc / RC is applied here. - total = int(off[-1]) - out_buf = np.empty(total, np.uint8) - annot_v_buf = np.empty(total, V_IDX_TYPE) - annot_pos_buf = np.empty(total, np.int32) - reconstruct_haplotypes_from_sparse( - geno_offset_idx=flat_geno_idx.reshape(-1, 1), - out=out_buf, - out_offsets=off, - regions=permuted_regions, - shifts=flat_shifts.reshape(-1, 1), - geno_offsets=self.genotypes.offsets, - geno_v_idxs=self.genotypes.data, - v_starts=self.variants.start, - ilens=self.variants.ilen, - alt_alleles=self.variants.alt.data.view(np.uint8), - alt_offsets=self.variants.alt.offsets, - ref=self.reference.reference, - ref_offsets=self.reference.offsets, - pad_char=self.reference.pad_char, - keep=keep_perm, - keep_offsets=keep_offsets_perm, - annot_v_idxs=annot_v_buf, - annot_ref_pos=annot_pos_buf, - ) - - haps_rag = cast( - "Ragged[np.bytes_]", - _Flat.from_offsets(out_buf, per_elem_shape, off).view("S1"), - ) - annot_v_rag = cast( - "Ragged[V_IDX_TYPE]", - _Flat.from_offsets(annot_v_buf, per_elem_shape, off), - ) - annot_pos_rag = cast( - "Ragged[np.int32]", - _Flat.from_offsets(annot_pos_buf, per_elem_shape, off), - ) - return haps_rag, annot_v_rag, annot_pos_rag + ) +else: + # Numba composed oracle path. RC is applied externally in + # _getitem_spliced (numba branch), so no to_rc / RC is applied here. + total = int(off[-1]) + out_buf = np.empty(total, np.uint8) + annot_v_buf = np.empty(total, V_IDX_TYPE) + annot_pos_buf = np.empty(total, np.int32) + reconstruct_haplotypes_from_sparse( + geno_offset_idx=flat_geno_idx.reshape(-1, 1), + out=out_buf, + out_offsets=off, + regions=permuted_regions, + shifts=flat_shifts.reshape(-1, 1), + geno_offsets=self.genotypes.offsets, + geno_v_idxs=self.genotypes.data, + v_starts=self.variants.start, + ilens=self.variants.ilen, + alt_alleles=self.variants.alt.data.view(np.uint8), + alt_offsets=self.variants.alt.offsets, + ref=self.reference.reference, + ref_offsets=self.reference.offsets, + pad_char=self.reference.pad_char, + keep=keep_perm, + keep_offsets=keep_offsets_perm, + annot_v_idxs=annot_v_buf, + annot_ref_pos=annot_pos_buf, + ) + +haps_rag = cast( + "Ragged[np.bytes_]", + _Flat.from_offsets(out_buf, per_elem_shape, off).view("S1"), +) +annot_v_rag = cast( + "Ragged[V_IDX_TYPE]", + _Flat.from_offsets(annot_v_buf, per_elem_shape, off), +) +annot_pos_rag = cast( + "Ragged[np.int32]", + _Flat.from_offsets(annot_pos_buf, per_elem_shape, off), +) +return haps_rag, annot_v_rag, annot_pos_rag ``` This deletes the old unconditional `reconstruct_haplotypes_from_sparse` call (it now lives only in the numba `else` branch) and the `if ... == "rust" and to_rc is not None: ... reverse_masked(...)` post-pass block (RC is now in-kernel on rust). If removing that block leaves `_FlatAnnotatedHaps` and/or the local `from .._ragged import _COMP` unused in the file, the lint step in Task 2 will catch it — remove the now-dead import(s). Do NOT change `_query.py::_getitem_spliced`: its `if _active_backend() == "numba"` RC guard remains correct (rust output is already RC'd, numba is post-passed there). diff --git a/docs/superpowers/plans/2026-06-26-rust-migration-phase-5-w5.md b/docs/superpowers/plans/2026-06-26-rust-migration-phase-5-w5.md index eaa47a37..8757e065 100644 --- a/docs/superpowers/plans/2026-06-26-rust-migration-phase-5-w5.md +++ b/docs/superpowers/plans/2026-06-26-rust-migration-phase-5-w5.md @@ -61,6 +61,7 @@ ```python # tests/parity/test_golden_infra.py """Self-tests for the golden snapshot/replay infrastructure.""" + from __future__ import annotations import numpy as np @@ -114,6 +115,7 @@ the numba oracle at generation time (see generate_goldens.py). Replay imports rust callables DIRECTLY — never via _dispatch — so these tests survive the numba/dispatch deletion in Stage B. """ + from __future__ import annotations from collections.abc import Callable @@ -217,12 +219,16 @@ def replay_tuple(name: str, cases: list) -> None: got = fn(*inputs) got = got if isinstance(got, tuple) else (got,) gold = golden if isinstance(golden, tuple) else (golden,) - assert len(got) == len(gold), f"{name}#{ci}: tuple len {len(got)} != {len(gold)}" + assert len(got) == len(gold), ( + f"{name}#{ci}: tuple len {len(got)} != {len(gold)}" + ) for j, (a, b) in enumerate(zip(got, gold)): _eq(f"{name}#{ci}", j, a, b) -def replay_inplace(name: str, cases: list, out_factory: Callable, out_index: int) -> None: +def replay_inplace( + name: str, cases: list, out_factory: Callable, out_index: int +) -> None: fn = RUST_KERNELS[name] for ci, (inputs, golden) in enumerate(cases): out = out_factory(inputs) @@ -238,9 +244,18 @@ def replay_dict(name: str, cases: list) -> None: got = fn(*inputs) assert set(got) == set(golden), f"{name}#{ci}: keys {set(got)} != {set(golden)}" for k in sorted(golden): - _eq(f"{name}#{ci}:{k}.data", 0, np.asarray(got[k][0]), np.asarray(golden[k][0])) - _eq(f"{name}#{ci}:{k}.off", 1, - np.asarray(got[k][1], np.int64), np.asarray(golden[k][1], np.int64)) + _eq( + f"{name}#{ci}:{k}.data", + 0, + np.asarray(got[k][0]), + np.asarray(golden[k][0]), + ) + _eq( + f"{name}#{ci}:{k}.off", + 1, + np.asarray(got[k][1], np.int64), + np.asarray(golden[k][1], np.int64), + ) ``` Note: `replay_inplace`'s `out_factory` takes `inputs` (so it can size the out buffer from `total_out` carried in the frozen case — the in-place strategies return `(total_out, inputs)`). @@ -311,6 +326,7 @@ and assert the numba oracle agrees BEFORE saving. After numba deletion this script still regenerates from rust (the numba cross-check is skipped if the backend is gone). """ + from __future__ import annotations import numpy as np @@ -330,12 +346,27 @@ SPEC = [ # in-place kernels: strategy yields (total_out, inputs); out inserted at index 0. INPLACE_SPEC = [ - ("intervals_to_tracks", strategies.intervals_to_tracks_inputs(), 200, - lambda inp: np.zeros(int(inp[-1][-1]), np.float32), 7), # out_index per existing test - ("shift_and_realign_tracks_sparse", strategies.shift_and_realign_tracks_inputs(), 200, - lambda total_out: np.zeros(total_out, np.float32), 0), - ("reconstruct_haplotypes_from_sparse", strategies.reconstruct_haplotypes_inputs(), 200, - lambda total_out: np.zeros(total_out, np.uint8), 0), + ( + "intervals_to_tracks", + strategies.intervals_to_tracks_inputs(), + 200, + lambda inp: np.zeros(int(inp[-1][-1]), np.float32), + 7, + ), # out_index per existing test + ( + "shift_and_realign_tracks_sparse", + strategies.shift_and_realign_tracks_inputs(), + 200, + lambda total_out: np.zeros(total_out, np.float32), + 0, + ), + ( + "reconstruct_haplotypes_from_sparse", + strategies.reconstruct_haplotypes_inputs(), + 200, + lambda total_out: np.zeros(total_out, np.uint8), + 0, + ), ] @@ -357,8 +388,9 @@ def _assert_oracle(name, a, b): assert set(a) == set(b) for k in a: np.testing.assert_array_equal(a[k][0], b[k][0]) - np.testing.assert_array_equal(np.asarray(a[k][1], np.int64), - np.asarray(b[k][1], np.int64)) + np.testing.assert_array_equal( + np.asarray(a[k][1], np.int64), np.asarray(b[k][1], np.int64) + ) else: np.testing.assert_array_equal(a, b, err_msg=f"{name} oracle mismatch") @@ -402,10 +434,14 @@ def gen_inplace_kernels(): inputs = ex of = out_factory out_r = of(inputs) - args = list(inputs); args.insert(out_index, out_r); rust(*args) + args = list(inputs) + args.insert(out_index, out_r) + rust(*args) if nb is not None: out_n = of(inputs) - an = list(inputs); an.insert(out_index, out_n); nb(*an) + an = list(inputs) + an.insert(out_index, out_n) + nb(*an) np.testing.assert_array_equal(out_n, out_r, err_msg=f"{name} oracle") cases.append((inputs, np.asarray(out_r))) _golden.save_golden(name, cases) @@ -459,6 +495,7 @@ Co-Authored-By: Claude Opus 4.8 " ```python # tests/parity/test_get_diffs_sparse_parity.py """get_diffs_sparse: rust vs frozen golden (oracle frozen Phase 5 W5).""" + from __future__ import annotations import pytest @@ -525,13 +562,24 @@ def flatten_output(out): from genvarloader._ragged import RaggedAnnotatedHaps if isinstance(out, RaggedAnnotatedHaps): - return {"kind": "annot", - "haps": (np.asarray(out.haps.data), np.asarray(out.haps.offsets, np.int64)), - "var_idxs": (np.asarray(out.var_idxs.data), np.asarray(out.var_idxs.offsets, np.int64)), - "ref_coords": (np.asarray(out.ref_coords.data), np.asarray(out.ref_coords.offsets, np.int64))} + return { + "kind": "annot", + "haps": (np.asarray(out.haps.data), np.asarray(out.haps.offsets, np.int64)), + "var_idxs": ( + np.asarray(out.var_idxs.data), + np.asarray(out.var_idxs.offsets, np.int64), + ), + "ref_coords": ( + np.asarray(out.ref_coords.data), + np.asarray(out.ref_coords.offsets, np.int64), + ), + } if isinstance(out, Ragged): - return {"kind": "ragged", - "data": np.asarray(out.data), "offsets": np.asarray(out.offsets, np.int64)} + return { + "kind": "ragged", + "data": np.asarray(out.data), + "offsets": np.asarray(out.offsets, np.int64), + } if isinstance(out, tuple): return {"kind": "tuple", "items": [flatten_output(o) for o in out]} return {"kind": "array", "data": np.asarray(out)} @@ -558,15 +606,19 @@ def test_haplotypes_mode_dataset_golden(phased_svar_gvl, reference, monkeypatch) # spy guard stays — proves the fused rust kernel fires orig = _haps_mod.reconstruct_haplotypes_fused calls = {"n": 0} + def _spy(*a, **k): calls["n"] += 1 return orig(*a, **k) + monkeypatch.setattr(_haps_mod, "reconstruct_haplotypes_fused", _spy) out_rust = ds[:, :] assert calls["n"] > 0, "fused rust kernel never fired — vacuous" # non-triviality + golden compare - _golden.assert_output_matches_golden(out_rust, _golden.load_flat_golden("ds_haplotypes_mode")) + _golden.assert_output_matches_golden( + out_rust, _golden.load_flat_golden("ds_haplotypes_mode") + ) ``` (`load_flat_golden` = `load_golden` returning the single flattened dict; add a thin variant or store as a 1-element `cases` list.) @@ -766,6 +818,7 @@ NOTE: `import genvarloader` may still pull numba transitively via seqpro this guard asserts genvarloader's own source is numba-free. See the seqpro follow-up issue for the transitive import and the W6 RSS impact. """ + from __future__ import annotations import pathlib @@ -825,6 +878,7 @@ Co-Authored-By: Claude Opus 4.8 " ```python # tests/parity/test_rayon_equivalence.py """Serial vs parallel rust output must be byte-identical (and == golden).""" + from __future__ import annotations import numpy as np import pytest @@ -845,7 +899,9 @@ def test_reconstruct_haplotypes_serial_eq_parallel(): fn(*args, parallel=parallel) # signature gains keyword `parallel` outs[parallel] = out np.testing.assert_array_equal(outs[False], outs[True], err_msg=f"case {ci}") - np.testing.assert_array_equal(outs[True], golden, err_msg=f"case {ci} vs golden") + np.testing.assert_array_equal( + outs[True], golden, err_msg=f"case {ci} vs golden" + ) ``` (If the FFI signature passes `parallel` positionally, adjust the call. Decide the FFI arg convention and keep it consistent across kernels.) diff --git a/docs/superpowers/plans/2026-06-26-rust-migration-phase-5.md b/docs/superpowers/plans/2026-06-26-rust-migration-phase-5.md index 9c301c2c..98999bf6 100644 --- a/docs/superpowers/plans/2026-06-26-rust-migration-phase-5.md +++ b/docs/superpowers/plans/2026-06-26-rust-migration-phase-5.md @@ -168,18 +168,20 @@ def test_overshoot_pads_full_tail(): # out_len=8: [1,2] ref + [50] allele, then ref exhausted -> pad rest with 0. out = np.full(8, 255, dtype=np.uint8) # 0xFF sentinel: catches unwritten positions reconstruct_haplotype_from_sparse( - np.array([0], dtype=np.int32), # v_idxs - np.array([2], dtype=np.int32), # v_starts - np.array([-5], dtype=np.int32), # ilens - 0, # shift - np.array([50], dtype=np.uint8), # alt_alleles - np.array([0, 1], dtype=np.int64), # alt_offsets + np.array([0], dtype=np.int32), # v_idxs + np.array([2], dtype=np.int32), # v_starts + np.array([-5], dtype=np.int32), # ilens + 0, # shift + np.array([50], dtype=np.uint8), # alt_alleles + np.array([0, 1], dtype=np.int64), # alt_offsets np.array([1, 2, 3, 4], dtype=np.uint8), # ref - 0, # ref_start - out, # out - 0, # pad_char + 0, # ref_start + out, # out + 0, # pad_char + ) + np.testing.assert_array_equal( + out, np.array([1, 2, 50, 0, 0, 0, 0, 0], dtype=np.uint8) ) - np.testing.assert_array_equal(out, np.array([1, 2, 50, 0, 0, 0, 0, 0], dtype=np.uint8)) ``` - [ ] **Step 2: Run to verify it FAILS** diff --git a/docs/superpowers/plans/2026-06-30-rayon-multithread-verification.md b/docs/superpowers/plans/2026-06-30-rayon-multithread-verification.md index a9bae62f..7e70d605 100644 --- a/docs/superpowers/plans/2026-06-30-rayon-multithread-verification.md +++ b/docs/superpowers/plans/2026-06-30-rayon-multithread-verification.md @@ -251,7 +251,9 @@ Replace `_detect_cpus`: ```python def _detect_cpus() -> int: try: - affinity = max(1, len(os.sched_getaffinity(0))) # respects cgroup cpuset (Linux) + affinity = max( + 1, len(os.sched_getaffinity(0)) + ) # respects cgroup cpuset (Linux) except AttributeError: affinity = max(1, os.cpu_count() or 1) quota = _cgroup_cpu_quota() @@ -599,8 +601,12 @@ def _materialize(ds): @pytest.mark.parametrize("seq_kind", ["haplotypes", "variants"]) -def test_forced_parallel_matches_serial(variant_track_dataset, reference, monkeypatch, seq_kind): - open_ds = lambda: gvl.Dataset.open(variant_track_dataset, reference=reference).with_seqs(seq_kind) +def test_forced_parallel_matches_serial( + variant_track_dataset, reference, monkeypatch, seq_kind +): + open_ds = lambda: gvl.Dataset.open( + variant_track_dataset, reference=reference + ).with_seqs(seq_kind) monkeypatch.delenv("GVL_FORCE_PARALLEL", raising=False) serial = _materialize(open_ds()) @@ -680,7 +686,9 @@ def _iterate_dataset(ds_path: str, reference_path: str, iters: int) -> int: # Force the parallel path and oversubscribe: many rayon threads per worker. os.environ["GVL_FORCE_PARALLEL"] = "1" os.environ["RAYON_NUM_THREADS"] = "8" - ds = gvl.Dataset.open(Path(ds_path), reference=Path(reference_path)).with_seqs("haplotypes") + ds = gvl.Dataset.open(Path(ds_path), reference=Path(reference_path)).with_seqs( + "haplotypes" + ) total = 0 n = len(ds) for _ in range(iters): @@ -699,13 +707,20 @@ def stress_dataset(source_bed, vcf_dir, reference, tmp_path: Path) -> tuple[Path with pyBigWig.open(str(bw_path), "w") as bw: bw.addHeader(contig_sizes, maxZooms=0) v = float(i + 1) - bw.addEntries(["chr1", "chr2"], [499_990, 17_320], - ends=[500_030, 17_340], values=[v, v]) + bw.addEntries( + ["chr1", "chr2"], + [499_990, 17_320], + ends=[500_030, 17_340], + values=[v, v], + ) bw_paths[sample] = str(bw_path) out = tmp_path / "stress.gvl" gvl.write( - path=out, bed=source_bed, variants=VCF(vcf_dir / "filtered_source.vcf.gz"), - tracks=gvl.BigWigs("5ss", bw_paths), max_jitter=2, + path=out, + bed=source_bed, + variants=VCF(vcf_dir / "filtered_source.vcf.gz"), + tracks=gvl.BigWigs("5ss", bw_paths), + max_jitter=2, ) return out, Path(reference) @@ -716,7 +731,9 @@ def test_concurrent_spawn_workers_do_not_deadlock(stress_dataset): for launch in range(LAUNCHES): with ProcessPoolExecutor(max_workers=N_WORKERS, mp_context=ctx) as ex: futs = [ - ex.submit(_iterate_dataset, str(ds_path), str(ref_path), ITERS_PER_WORKER) + ex.submit( + _iterate_dataset, str(ds_path), str(ref_path), ITERS_PER_WORKER + ) for _ in range(N_WORKERS) ] try: diff --git a/docs/superpowers/plans/2026-07-03-svar2-genoray-search-gather-split.md b/docs/superpowers/plans/2026-07-03-svar2-genoray-search-gather-split.md index cab59f9a..a0980da7 100644 --- a/docs/superpowers/plans/2026-07-03-svar2-genoray-search-gather-split.md +++ b/docs/superpowers/plans/2026-07-03-svar2-genoray-search-gather-split.md @@ -732,8 +732,16 @@ def _assert_dicts_equal(a: dict, b: dict, keys): PAYLOAD_KEYS = [ - "vk_pos", "vk_key", "vk_off", "dense_pos", "dense_key", "dense_range", - "dense_present", "dense_present_off", "lut_bytes", "lut_off", + "vk_pos", + "vk_key", + "vk_off", + "dense_pos", + "dense_key", + "dense_range", + "dense_present", + "dense_present_off", + "lut_bytes", + "lut_off", ] @@ -765,8 +773,8 @@ def test_read_ranges_sample_subset(svar2_store): fh = 1 * ploidy + p sh = 0 * ploidy + p np.testing.assert_array_equal( - full["vk_pos"][full["vk_off"][fh]:full["vk_off"][fh + 1]], - sub["vk_pos"][sub["vk_off"][sh]:sub["vk_off"][sh + 1]], + full["vk_pos"][full["vk_off"][fh] : full["vk_off"][fh + 1]], + sub["vk_pos"][sub["vk_off"][sh] : sub["vk_off"][sh + 1]], ) @@ -774,8 +782,16 @@ def test_find_ranges_out_streaming(svar2_store): sv = SparseVar2(svar2_store) ranges = sv.find_ranges("chr1", [0], [40]) # Pre-allocate matching-shape buffers and stream into them. - out = {k: np.empty_like(np.asarray(ranges[k])) for k in - ("dense_range", "region_starts", "sample_cols", "vk_snp_range", "vk_indel_range")} + out = { + k: np.empty_like(np.asarray(ranges[k])) + for k in ( + "dense_range", + "region_starts", + "sample_cols", + "vk_snp_range", + "vk_indel_range", + ) + } ranges2 = sv.find_ranges("chr1", [0], [40], out=out) for k in out: np.testing.assert_array_equal(np.asarray(ranges2[k]), np.asarray(ranges[k])) @@ -793,40 +809,43 @@ Expected: FAIL — `SparseVar2` has no attribute `read_ranges`. Add to `_BatchQueryMixin` in `python/genoray/_svar2_batch.py`. Resolve `samples=` names → original integer indices via `self.samples.index(...)`; validate membership. ```python - def _sample_idxs(self, samples): - if samples is None: - return None - idxs = [] - for s in np.atleast_1d(np.asarray(samples)).tolist(): - if s not in self.samples: - raise ValueError(f"Sample {s!r} not found in the dataset.") - idxs.append(self.samples.index(s)) - return idxs - - def read_ranges(self, contig, starts, ends, samples=None): - """Fused search+gather query (byte-identical to ``overlap_batch`` for - ``samples=None``). See ``overlap_batch`` for the returned dict contract.""" - reg = self._regions(starts, ends) - return self._readers[contig].read_ranges(reg, self._sample_idxs(samples)) - - def find_ranges(self, contig, starts, ends, samples=None, out=None): - """Search-only step: returns the compact ranges bundle to be replayed by - ``gather_ranges``. When ``out`` is a dict of preallocated arrays keyed by - the bundle field names, the ranges are written into it in place.""" - reg = self._regions(starts, ends) - d = self._readers[contig].find_ranges(reg, self._sample_idxs(samples)) - if out is not None: - for k, buf in out.items(): - np.asarray(buf)[...] = np.asarray(d[k]) - d[k] = buf - return d - - def gather_ranges(self, contig, ranges, samples=None): - """Tree-free gather step: replay a ``find_ranges`` bundle into the full - ``overlap_batch`` payload dict. ``samples`` is accepted for symmetry but - the subset is already fixed by the bundle; passing a different subset is - a ValueError.""" - return self._readers[contig].gather_ranges(ranges) +def _sample_idxs(self, samples): + if samples is None: + return None + idxs = [] + for s in np.atleast_1d(np.asarray(samples)).tolist(): + if s not in self.samples: + raise ValueError(f"Sample {s!r} not found in the dataset.") + idxs.append(self.samples.index(s)) + return idxs + + +def read_ranges(self, contig, starts, ends, samples=None): + """Fused search+gather query (byte-identical to ``overlap_batch`` for + ``samples=None``). See ``overlap_batch`` for the returned dict contract.""" + reg = self._regions(starts, ends) + return self._readers[contig].read_ranges(reg, self._sample_idxs(samples)) + + +def find_ranges(self, contig, starts, ends, samples=None, out=None): + """Search-only step: returns the compact ranges bundle to be replayed by + ``gather_ranges``. When ``out`` is a dict of preallocated arrays keyed by + the bundle field names, the ranges are written into it in place.""" + reg = self._regions(starts, ends) + d = self._readers[contig].find_ranges(reg, self._sample_idxs(samples)) + if out is not None: + for k, buf in out.items(): + np.asarray(buf)[...] = np.asarray(d[k]) + d[k] = buf + return d + + +def gather_ranges(self, contig, ranges, samples=None): + """Tree-free gather step: replay a ``find_ranges`` bundle into the full + ``overlap_batch`` payload dict. ``samples`` is accepted for symmetry but + the subset is already fixed by the bundle; passing a different subset is + a ValueError.""" + return self._readers[contig].gather_ranges(ranges) ``` Add the `_regions` helper (shared with `overlap_batch`, which currently inlines `[(int(s), int(e)) for s, e in regions]`): @@ -876,6 +895,7 @@ Extend `tests/test_svar2_ranges.py`. Mirror whatever reconstruction check `tests ```python def test_split_reconstructs_like_decode_oracle(svar2_store): from tests.test_svar2_decode import decode_from_payload # reuse existing helper + sv = SparseVar2(svar2_store) starts, ends = [0], [40] ob = sv.overlap_batch("chr1", list(zip(starts, ends))) diff --git a/docs/superpowers/plans/2026-07-03-svar2-gvl-dataset-wiring.md b/docs/superpowers/plans/2026-07-03-svar2-gvl-dataset-wiring.md index d200cb9c..c612e31e 100644 --- a/docs/superpowers/plans/2026-07-03-svar2-gvl-dataset-wiring.md +++ b/docs/superpowers/plans/2026-07-03-svar2-gvl-dataset-wiring.md @@ -97,9 +97,11 @@ def test_resolve_falls_back_to_relative_then_absolute(tmp_path): gvl = tmp_path / "ds.gvl" gvl.mkdir() import os + rel = os.path.relpath(d, start=gvl).replace(os.sep, "/") link = Svar2Link( - relative_path=rel, absolute_path=str(d), + relative_path=rel, + absolute_path=str(d), fingerprint=Svar2Fingerprint(n_variants=3, store_bytes=4), ) assert _resolve_svar2(gvl, link, None) == d @@ -109,7 +111,8 @@ def test_resolve_raises_when_unfindable(tmp_path): gvl = tmp_path / "ds.gvl" gvl.mkdir() link = Svar2Link( - relative_path="missing.svar2", absolute_path=str(tmp_path / "missing.svar2"), + relative_path="missing.svar2", + absolute_path=str(tmp_path / "missing.svar2"), fingerprint=Svar2Fingerprint(n_variants=3, store_bytes=4), ) with pytest.raises(FileNotFoundError): @@ -131,6 +134,7 @@ Copy `_svar_link.py`'s structure. Fingerprint on a **stable** `.svar2` identity: Mirrors _svar_link.py. SVAR2 fingerprint identity = n_variants (from the SparseVar2 index) + byte count of a canonical store file. """ + from __future__ import annotations from pathlib import Path @@ -187,9 +191,7 @@ def _verify_fingerprint2(svar2_path: Path, link: Svar2Link | None) -> None: return store = svar2_path / _STORE_FILE if not store.exists(): - raise FileNotFoundError( - f"Expected {store}; resolved svar2 is malformed." - ) + raise FileNotFoundError(f"Expected {store}; resolved svar2 is malformed.") import polars as pl n_variants_observed = ( @@ -247,11 +249,15 @@ def test_metadata_roundtrips_svar2_link(): from genvarloader._dataset._svar2_link import Svar2Fingerprint, Svar2Link link = Svar2Link( - relative_path="c.svar2", absolute_path="/abs/c.svar2", + relative_path="c.svar2", + absolute_path="/abs/c.svar2", fingerprint=Svar2Fingerprint(n_variants=5, store_bytes=99), ) m = Metadata( - contigs=["chr1"], samples=["s0"], ploidy=2, n_regions=1, + contigs=["chr1"], + samples=["s0"], + ploidy=2, + n_regions=1, svar2_link=link, ) m2 = Metadata.model_validate_json(m.model_dump_json()) @@ -341,9 +347,7 @@ def test_write_svar2_produces_ranges_cache(svar2_store, tmp_path): from genoray import SparseVar2 sv = SparseVar2(svar2_store) - bed = pl.DataFrame( - {"chrom": ["chr1"], "chromStart": [0], "chromEnd": [40]} - ) + bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [0], "chromEnd": [40]}) out = tmp_path / "ds.gvl" gvl.write(path=out, bed=bed, variants=Path(svar2_store)) @@ -414,14 +418,24 @@ def _write_from_svar2( S = len(samples) P = svar2.ploidy - dense_range = np.memmap(out_dir / "svar2_dense_range.npy", np.int32, "w+", shape=(R, 2)) - region_starts = np.memmap(out_dir / "svar2_region_starts.npy", np.int32, "w+", shape=(R,)) - vk_snp = np.memmap(out_dir / "svar2_vk_snp_range.npy", np.int64, "w+", shape=(R, S, P, 2)) - vk_indel = np.memmap(out_dir / "svar2_vk_indel_range.npy", np.int64, "w+", shape=(R, S, P, 2)) + dense_range = np.memmap( + out_dir / "svar2_dense_range.npy", np.int32, "w+", shape=(R, 2) + ) + region_starts = np.memmap( + out_dir / "svar2_region_starts.npy", np.int32, "w+", shape=(R,) + ) + vk_snp = np.memmap( + out_dir / "svar2_vk_snp_range.npy", np.int64, "w+", shape=(R, S, P, 2) + ) + vk_indel = np.memmap( + out_dir / "svar2_vk_indel_range.npy", np.int64, "w+", shape=(R, S, P, 2) + ) sample_cols: list[int] | None = None contig_offset = 0 - for (c,), df in bed.partition_by("chrom", as_dict=True, maintain_order=True).items(): + for (c,), df in bed.partition_by( + "chrom", as_dict=True, maintain_order=True + ).items(): c = cast(str, c) rc = df.height rows = slice(contig_offset, contig_offset + rc) @@ -446,7 +460,9 @@ def _write_from_svar2( with open(out_dir / "svar2_meta.json", "w") as f: json.dump( { - "n_regions": R, "n_samples": S, "ploidy": P, + "n_regions": R, + "n_samples": S, + "ploidy": P, "sample_cols": sample_cols, "dense_range": {"shape": [R, 2], "dtype": " (region, sample); group by contig; gather+reconstruct per contig; @@ -715,25 +767,26 @@ Expected: FAIL — `_recon` is a plain `Haps`, or `_build_seqs` errors on missin 3a. In `_open.py::_build_seqs` (`:149`), branch on `metadata.svar2_link`: ```python - if self._has_genotypes(): - if metadata.ploidy is None: - raise ValueError("Malformed dataset: found genotypes but not ploidy.") - if metadata.svar2_link is not None: - from ._haps import HapsSvar2 - if reference is None: - raise ValueError( - "SVAR2 datasets require a reference genome for haplotype output." - ) - return HapsSvar2.from_path( - path=self.path, - reference=reference, - contigs=metadata.contigs, - samples=metadata.samples, - ploidy=metadata.ploidy, - svar2_link=metadata.svar2_link, - svar2_override=getattr(self, "svar2", None), - ) - seqs = Haps.from_path(...) # unchanged SVAR1 path +if self._has_genotypes(): + if metadata.ploidy is None: + raise ValueError("Malformed dataset: found genotypes but not ploidy.") + if metadata.svar2_link is not None: + from ._haps import HapsSvar2 + + if reference is None: + raise ValueError( + "SVAR2 datasets require a reference genome for haplotype output." + ) + return HapsSvar2.from_path( + path=self.path, + reference=reference, + contigs=metadata.contigs, + samples=metadata.samples, + ploidy=metadata.ploidy, + svar2_link=metadata.svar2_link, + svar2_override=getattr(self, "svar2", None), + ) + seqs = Haps.from_path(...) # unchanged SVAR1 path ``` `self._has_genotypes()` checks for `genotypes/` — confirm it does not require SVAR1-specific files (`svar_meta.json`); if it does, relax it to also accept `svar2_meta.json`. @@ -743,14 +796,15 @@ Expected: FAIL — `_recon` is a plain `Haps`, or `_build_seqs` errors on missin 3c. In `_reconstruct.py::_build_reconstructor`, accept `HapsSvar2` for the `haplotypes` kind. The simplest wiring: treat `HapsSvar2` like `Haps` in the `seqs_kind in ("haplotypes", ...)` branch but restrict to `"haplotypes"`: ```python - from ._haps import HapsSvar2 - if isinstance(seqs, HapsSvar2): - if seqs_kind not in (None, "haplotypes"): - raise NotImplementedError( - f"SVAR2 datasets support only 'haplotypes', not {seqs_kind!r}." - ) - active_seqs = seqs - # dispatch: HapsSvar2 alone -> itself; with tracks -> HapsSvar2Tracks (Task 6) +from ._haps import HapsSvar2 + +if isinstance(seqs, HapsSvar2): + if seqs_kind not in (None, "haplotypes"): + raise NotImplementedError( + f"SVAR2 datasets support only 'haplotypes', not {seqs_kind!r}." + ) + active_seqs = seqs + # dispatch: HapsSvar2 alone -> itself; with tracks -> HapsSvar2Tracks (Task 6) ``` Add `HapsSvar2` (and `HapsSvar2Tracks` from Task 6) to the `_recon` type union in `_impl.py` (`:899`) and the `match self._recon` in `__getitem__` (`:1028`). @@ -792,15 +846,17 @@ def test_svar2_tracks_match_live(svar2_store, reference_fasta, bigwig_track, tmp out = tmp_path / "ds.gvl" gvl.write(path=out, bed=bed, variants=Path(svar2_store), tracks=[bigwig_track]) - ds = gvl.Dataset.open(out, reference=reference_fasta).with_seqs("haplotypes").with_tracks(...) + ds = ( + gvl.Dataset.open(out, reference=reference_fasta) + .with_seqs("haplotypes") + .with_tracks(...) + ) _, cached_tracks = ds[0, :] # Live oracle: SparseVar2Source.realign_tracks with the same track buffer. sv = SparseVar2(svar2_store) live = SparseVar2Source(sv).realign_tracks("chr1", [(0, 40)], *_track_args(...)) - np.testing.assert_array_equal( - cached_tracks.to_packed().data, live.to_packed().data - ) + np.testing.assert_array_equal(cached_tracks.to_packed().data, live.to_packed().data) ``` Adapt fixtures/`_track_args` to whatever the existing `tests/test_svar2_realign_tracks.py` uses. @@ -851,11 +907,13 @@ Parametrize over `{SNP, INS, DEL} × {1, 2, 4} samples × {1, 2} ploidy` (reuse ```python import pytest + @pytest.mark.parametrize("variant_kind", ["snp", "ins", "del"]) @pytest.mark.parametrize("n_samples", [1, 2, 4]) @pytest.mark.parametrize("ploidy", [1, 2]) -def test_svar2_cached_matches_decode_matrix(variant_kind, n_samples, ploidy, tmp_path): - ... # synth store -> gvl.write -> open -> compare cached vs live vs decode +def test_svar2_cached_matches_decode_matrix( + variant_kind, n_samples, ploidy, tmp_path +): ... # synth store -> gvl.write -> open -> compare cached vs live vs decode ``` - [ ] **Step 2: Add the real chr21 germline + somatic parity test (slow)** diff --git a/docs/superpowers/plans/2026-07-03-svar2-gvl-mvp-validate-benchmark.md b/docs/superpowers/plans/2026-07-03-svar2-gvl-mvp-validate-benchmark.md index 5193f6b4..34207449 100644 --- a/docs/superpowers/plans/2026-07-03-svar2-gvl-mvp-validate-benchmark.md +++ b/docs/superpowers/plans/2026-07-03-svar2-gvl-mvp-validate-benchmark.md @@ -129,8 +129,15 @@ def svar2_del_store(tmp_path_factory) -> Path: out = d / "store" _core.run_conversion_pipeline( - str(bcf), str(ref), ["chr1"], str(out), ["S0", "S1"], - 25_000, 2, 1, 8 * 1024 * 1024, + 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 @@ -155,7 +162,7 @@ def test_svar2_realign_tracks_matches_svar1_oracle(svar2_del_store): rng = np.random.default_rng(0) track = rng.random(region_len).astype(np.float32) - strategy_id = 0 # irrelevant for DEL-only (insertion-fill unused) + strategy_id = 0 # irrelevant for DEL-only (insertion-fill unused) params = np.zeros(1, np.float64) base_seed = 0 @@ -164,12 +171,12 @@ def test_svar2_realign_tracks_matches_svar1_oracle(svar2_del_store): out_rag = src.realign_tracks( contig, regions, - track, # flat per-region track buffer - np.array([0, region_len], np.int64), # (R+1) offsets + track, # flat per-region track buffer + np.array([0, region_len], np.int64), # (R+1) offsets params, strategy_id, base_seed, - shifts=None, # no jitter + shifts=None, # no jitter parallel=False, ) @@ -177,7 +184,7 @@ def test_svar2_realign_tracks_matches_svar1_oracle(svar2_del_store): raw = sv._readers[contig].decode_batch([(q_start, q_end)]) R, So, Po = int(raw["n_regions"]), int(raw["n_samples"]), int(raw["ploidy"]) assert (R, So, Po) == (1, S, P) - off = np.asarray(raw["off"]) # (H+1,) per-hap variant offsets + off = np.asarray(raw["off"]) # (H+1,) per-hap variant offsets d_pos = np.asarray(raw["pos"]) d_ilen = np.asarray(raw["ilen"]) @@ -187,7 +194,7 @@ def test_svar2_realign_tracks_matches_svar1_oracle(svar2_del_store): for s in range(S): for p in range(P): - h = (0 * S + s) * P + p # region-major h=(r*S+s)*P+p + h = (0 * S + s) * P + p # region-major h=(r*S+s)*P+p gi0, gi1 = int(off[h]), int(off[h + 1]) pos_h = np.ascontiguousarray(d_pos[gi0:gi1], np.int32) ilen_h = np.ascontiguousarray(d_ilen[gi0:gi1], np.int32) @@ -223,9 +230,12 @@ def test_svar2_realign_tracks_matches_svar1_oracle(svar2_del_store): hap=h, ) np.testing.assert_allclose( - got, expected, rtol=0, atol=0, + got, + expected, + rtol=0, + atol=0, err_msg=f"(s={s},p={p}) SVAR2 track != SVAR1 oracle " - f"(pos={pos_h.tolist()}, ilen={ilen_h.tolist()})", + f"(pos={pos_h.tolist()}, ilen={ilen_h.tolist()})", ) ``` @@ -298,27 +308,37 @@ Create `tmp/svar2_mvp/build_stores.py`: ```python """Build .svar (SVAR1) and .svar2 (SVAR2) stores from a normalized biallelic BCF.""" + import sys from pathlib import Path from genoray import VCF, SparseVar, _core + def build(bcf: str, chrom: str, samples: list[str], out_prefix: str, ploidy: int): bcf = str(bcf) # SVAR 1.0 SparseVar.from_vcf(f"{out_prefix}.svar", VCF(bcf), "8g", overwrite=True) # SVAR 2.0 _core.run_conversion_pipeline( - bcf, "/carter/shared/data/gdc/resources/GRCh38.d1.vd1.fa", - [chrom], f"{out_prefix}.svar2", samples, - 25_000, ploidy, 8, 8 * 1024 * 1024, + bcf, + "/carter/shared/data/gdc/resources/GRCh38.d1.vd1.fa", + [chrom], + f"{out_prefix}.svar2", + samples, + 25_000, + ploidy, + 8, + 8 * 1024 * 1024, ) print(f"built {out_prefix}.svar and {out_prefix}.svar2") + if __name__ == "__main__": # argv: bcf, chrom, out_prefix = sys.argv[1], sys.argv[2], sys.argv[3] import subprocess + samples = subprocess.run( ["bcftools", "query", "-l", bcf], capture_output=True, text=True, check=True ).stdout.split() @@ -347,6 +367,7 @@ Create `tmp/svar2_mvp/validate.py`: both the SVAR1 (gvl Dataset over .svar) and SVAR2 (SparseVar2Source over .svar2) backends, on a handful of regions x a few samples. Correctness is already proven by the test suite; this proves the REAL-DATA plumbing works.""" + import sys from pathlib import Path @@ -357,6 +378,7 @@ from genvarloader._dataset._svar2_source import SparseVar2Source REF = "/carter/shared/data/gdc/resources/GRCh38.d1.vd1.fa" + def main(prefix: str, chrom: str): # A few small regions (0-based, half-open) in a variant-dense chr21 window. regions = [(20_000_000, 20_001_000), (30_000_000, 30_000_500)] @@ -367,36 +389,47 @@ def main(prefix: str, chrom: str): ref_bytes = _contig_ref(REF, chrom) src = SparseVar2Source(sv2) hap = src.reconstruct( - chrom, regions, + chrom, + regions, np.frombuffer(ref_bytes, np.uint8), np.array([0, len(ref_bytes)], np.int64), - pad_char=ord("N"), shifts=None, output_length=-1, + pad_char=ord("N"), + shifts=None, + output_length=-1, ) lens = np.asarray(hap.offsets) - print(f"[svar2] hap ragged rows={len(lens) - 1} " - f"min_len={int(np.diff(lens).min())} max_len={int(np.diff(lens).max())}") + print( + f"[svar2] hap ragged rows={len(lens) - 1} " + f"min_len={int(np.diff(lens).min())} max_len={int(np.diff(lens).max())}" + ) var = sv2.decode(chrom, regions) print(f"[svar2] decode variants: {var}") # --- SVAR1 backend (gvl Dataset over .svar) --- import polars as pl - bed = pl.DataFrame({ - "chrom": [chrom] * len(regions), - "chromStart": [s for s, _ in regions], - "chromEnd": [e for _, e in regions], - }) + + bed = pl.DataFrame( + { + "chrom": [chrom] * len(regions), + "chromStart": [s for s, _ in regions], + "chromEnd": [e for _, e in regions], + } + ) ds_path = f"{prefix}.gvl" gvl.write(ds_path, bed, variants=f"{prefix}.svar", overwrite=True) ds = gvl.Dataset.open(ds_path, reference=REF).with_seqs("haplotypes") - seqs = ds[:len(regions), :3] # a few regions x first 3 samples + seqs = ds[: len(regions), :3] # a few regions x first 3 samples print(f"[svar1] gvl haplotypes sample shape/type: {type(seqs)}") + def _contig_ref(fasta: str, chrom: str) -> bytes: import pysam + return pysam.FastaFile(fasta).fetch(chrom).encode() + if __name__ == "__main__": - main(sys.argv[1], sys.argv[2]) # argv: + main(sys.argv[1], sys.argv[2]) # argv: ``` ```bash @@ -452,6 +485,7 @@ Create `tmp/svar2_mvp/benchmark.py`: """Benchmark SVAR1 (gvl Dataset over .svar) vs SVAR2 (SparseVar2Source over .svar2): hap latency, variant latency, store size, for one source prefix. Fair workload: ALL samples for a fixed region set. Warm caches, median of N.""" + import sys import time import subprocess @@ -465,10 +499,13 @@ from genvarloader._dataset._svar2_source import SparseVar2Source REF = "/carter/shared/data/gdc/resources/GRCh38.d1.vd1.fa" N = 5 # repeats + def _contig_ref(fasta, chrom): import pysam + return pysam.FastaFile(fasta).fetch(chrom).encode() + def _timed(fn, warmup=1): for _ in range(warmup): fn() @@ -479,9 +516,13 @@ def _timed(fn, warmup=1): ts.append(time.perf_counter() - t0) return median(ts) + def main(prefix, chrom): - regions = [(20_000_000, 20_001_000), (30_000_000, 30_000_500), - (40_000_000, 40_001_000)] + regions = [ + (20_000_000, 20_001_000), + (30_000_000, 30_000_500), + (40_000_000, 40_001_000), + ] ref_bytes = _contig_ref(REF, chrom) ref_u8 = np.frombuffer(ref_bytes, np.uint8) ref_off = np.array([0, len(ref_bytes)], np.int64) @@ -489,37 +530,55 @@ def main(prefix, chrom): # SVAR2 backend sv2 = SparseVar2(f"{prefix}.svar2") src = SparseVar2Source(sv2) - svar2_hap = _timed(lambda: src.reconstruct( - chrom, regions, ref_u8, ref_off, pad_char=ord("N"), - shifts=None, output_length=-1)) + svar2_hap = _timed( + lambda: src.reconstruct( + chrom, + regions, + ref_u8, + ref_off, + pad_char=ord("N"), + shifts=None, + output_length=-1, + ) + ) svar2_var = _timed(lambda: sv2.decode(chrom, regions)) # SVAR1 backend (all samples, same regions) import polars as pl - bed = pl.DataFrame({"chrom": [chrom] * len(regions), - "chromStart": [s for s, _ in regions], - "chromEnd": [e for _, e in regions]}) + + bed = pl.DataFrame( + { + "chrom": [chrom] * len(regions), + "chromStart": [s for s, _ in regions], + "chromEnd": [e for _, e in regions], + } + ) ds_path = f"{prefix}.gvl" ds = gvl.Dataset.open(ds_path, reference=REF) ds_hap = ds.with_seqs("haplotypes") ds_var = ds.with_seqs("variants") n_s = sv2.n_samples - svar1_hap = _timed(lambda: ds_hap[:len(regions), :n_s]) - svar1_var = _timed(lambda: ds_var[:len(regions), :n_s]) + svar1_hap = _timed(lambda: ds_hap[: len(regions), :n_s]) + svar1_var = _timed(lambda: ds_var[: len(regions), :n_s]) def du(path): - return subprocess.run(["du", "-sb", path], capture_output=True, - text=True).stdout.split()[0] + return subprocess.run( + ["du", "-sb", path], capture_output=True, text=True + ).stdout.split()[0] - print(f"source={prefix.split('/')[-1]} chrom={chrom} n_samples={n_s} " - f"regions={len(regions)} N={N}") + print( + f"source={prefix.split('/')[-1]} chrom={chrom} n_samples={n_s} " + f"regions={len(regions)} N={N}" + ) print(f" hap_latency_s svar1={svar1_hap:.4f} svar2={svar2_hap:.4f}") print(f" var_latency_s svar1={svar1_var:.4f} svar2={svar2_var:.4f}") - print(f" store_bytes svar1={du(prefix + '.svar')} " - f"svar2={du(prefix + '.svar2')}") + print( + f" store_bytes svar1={du(prefix + '.svar')} svar2={du(prefix + '.svar2')}" + ) + if __name__ == "__main__": - main(sys.argv[1], sys.argv[2]) # argv: + main(sys.argv[1], sys.argv[2]) # argv: ``` - [ ] **Step 2: Run the benchmark for both sources** diff --git a/docs/superpowers/plans/2026-07-03-svar2-profiling-followup.md b/docs/superpowers/plans/2026-07-03-svar2-profiling-followup.md index 8ccd96a6..b562dca5 100644 --- a/docs/superpowers/plans/2026-07-03-svar2-profiling-followup.md +++ b/docs/superpowers/plans/2026-07-03-svar2-profiling-followup.md @@ -193,6 +193,7 @@ py-spy/perf attribute time to that path only. Prints: per_call_s= For svar1, the 3-region .gvl is written ONCE before the loop (we profile the query, not gvl.write).""" + import sys import time @@ -206,6 +207,7 @@ REGIONS = [(20_000_000, 20_001_000), (30_000_000, 30_000_500), (40_000_000, 40_0 def _ref(): import pysam + rb = pysam.FastaFile(REF).fetch(CHROM).encode() return np.frombuffer(rb, np.uint8), np.array([0, len(rb)], np.int64) @@ -213,12 +215,15 @@ def _ref(): def make_svar2(cohort): from genoray import SparseVar2 from genvarloader._dataset._svar2_source import SparseVar2Source + src = SparseVar2Source(SparseVar2(f"{W}/{cohort}.svar2")) ru, ro = _ref() def call(): - src.reconstruct(CHROM, REGIONS, ru, ro, pad_char=ord("N"), - shifts=None, output_length=-1) + src.reconstruct( + CHROM, REGIONS, ru, ro, pad_char=ord("N"), shifts=None, output_length=-1 + ) + return call @@ -226,16 +231,22 @@ def make_svar1(cohort): import polars as pl import genvarloader as gvl from genoray import SparseVar2 + n_s = SparseVar2(f"{W}/{cohort}.svar2").n_samples - bed = pl.DataFrame({"chrom": [CHROM] * len(REGIONS), - "chromStart": [s for s, _ in REGIONS], - "chromEnd": [e for _, e in REGIONS]}) + bed = pl.DataFrame( + { + "chrom": [CHROM] * len(REGIONS), + "chromStart": [s for s, _ in REGIONS], + "chromEnd": [e for _, e in REGIONS], + } + ) ds_path = f"{W}/{cohort}.gvl" gvl.write(ds_path, bed, variants=f"{W}/{cohort}.svar", overwrite=True) # ONCE ds_hap = gvl.Dataset.open(ds_path, reference=REF).with_seqs("haplotypes") def call(): - ds_hap[:len(REGIONS), :n_s] + ds_hap[: len(REGIONS), :n_s] + return call @@ -290,6 +301,7 @@ self-time by LEAF frame. A leaf frame is Python iff it contains '.py:'. python split_folded.py """ + import sys from collections import Counter @@ -321,8 +333,11 @@ def main(path): nat += n tot = py + nat if tot == 0: - print("no samples parsed"); return - print(f"python_pct={100 * py / tot:.1f} native_pct={100 * nat / tot:.1f} total_samples={tot}") + print("no samples parsed") + return + print( + f"python_pct={100 * py / tot:.1f} native_pct={100 * nat / tot:.1f} total_samples={tot}" + ) print("top-15 leaf frames (self-time):") for leaf, n in leaves.most_common(15): print(f" {100 * n / tot:5.1f}% [{classed[leaf]:6s}] {leaf}") @@ -585,6 +600,7 @@ Fixed dataset family (subsampled somatic), same 3 regions, warm, median N=5. python e2_bench.py > tmp/svar2_mvp/prof_out/e2_curve.tsv """ + import time from statistics import median @@ -627,6 +643,7 @@ def variants_at(S): def main(): import pysam + rb = pysam.FastaFile(REF).fetch(CHROM).encode() ru, ro = np.frombuffer(rb, np.uint8), np.array([0, len(rb)], np.int64) print("S\tvariants\tsvar1_hap_s\tsvar2_hap_s") @@ -635,14 +652,21 @@ def main(): sv2 = SparseVar2(f"{p}.svar2") src = SparseVar2Source(sv2) n_s = sv2.n_samples - svar2 = timed(lambda: src.reconstruct(CHROM, REGIONS, ru, ro, - pad_char=ord("N"), shifts=None, output_length=-1)) - bed = pl.DataFrame({"chrom": [CHROM] * len(REGIONS), - "chromStart": [s for s, _ in REGIONS], - "chromEnd": [e for _, e in REGIONS]}) + svar2 = timed( + lambda: src.reconstruct( + CHROM, REGIONS, ru, ro, pad_char=ord("N"), shifts=None, output_length=-1 + ) + ) + bed = pl.DataFrame( + { + "chrom": [CHROM] * len(REGIONS), + "chromStart": [s for s, _ in REGIONS], + "chromEnd": [e for _, e in REGIONS], + } + ) gvl.write(f"{p}.gvl", bed, variants=f"{p}.svar", overwrite=True) ds_hap = gvl.Dataset.open(f"{p}.gvl", reference=REF).with_seqs("haplotypes") - svar1 = timed(lambda: ds_hap[:len(REGIONS), :n_s]) + svar1 = timed(lambda: ds_hap[: len(REGIONS), :n_s]) print(f"{S}\t{variants_at(S)}\t{svar1:.4f}\t{svar2:.4f}", flush=True) @@ -707,6 +731,7 @@ Germline (high-AF -> large n_dense_variants) is the stress cohort. python e3_probe.py > tmp/svar2_mvp/prof_out/e3.tsv """ + import time from statistics import median @@ -727,7 +752,9 @@ def timed(fn): fn() ts = [] for _ in range(N): - t0 = time.perf_counter(); fn(); ts.append(time.perf_counter() - t0) + t0 = time.perf_counter() + fn() + ts.append(time.perf_counter() - t0) return median(ts) @@ -743,8 +770,11 @@ def main(): # dense variants actually spanned by this region (dense_range gives [lo,hi) per region) dr = np.asarray(d["dense_range"]).reshape(-1, 2) n_dense = int((dr[:, 1] - dr[:, 0]).sum()) - t = timed(lambda: src.reconstruct(CHROM, regs, ru, ro, - pad_char=ord("N"), shifts=None, output_length=-1)) + t = timed( + lambda: src.reconstruct( + CHROM, regs, ru, ro, pad_char=ord("N"), shifts=None, output_length=-1 + ) + ) print(f"{w}\t{n_dense}\t{t:.4f}", flush=True) @@ -796,6 +826,7 @@ Create `tmp/svar2_mvp/e4_convert_driver.py`: python e4_convert_driver.py Prints: build_wall_s= Set GENORAY_SAMPLE_INTERVAL in the environment to enable genoray's sampler.""" + import sys import time import subprocess @@ -806,12 +837,20 @@ REF = "/carter/shared/data/gdc/resources/GRCh38.d1.vd1.fa" def main(bcf, chrom, out_prefix, max_threads): - samples = subprocess.run(["bcftools", "query", "-l", bcf], - capture_output=True, text=True, check=True).stdout.split() + samples = subprocess.run( + ["bcftools", "query", "-l", bcf], capture_output=True, text=True, check=True + ).stdout.split() t0 = time.perf_counter() _core.run_conversion_pipeline( - bcf, REF, [chrom], f"{out_prefix}.svar2", samples, - 25_000, 2, int(max_threads), 8 * 1024 * 1024, + bcf, + REF, + [chrom], + f"{out_prefix}.svar2", + samples, + 25_000, + 2, + int(max_threads), + 8 * 1024 * 1024, ) print(f"build_wall_s={time.perf_counter() - t0:.2f}") diff --git a/docs/superpowers/plans/2026-07-04-svar2-genoray-readbound-gather.md b/docs/superpowers/plans/2026-07-04-svar2-genoray-readbound-gather.md index aff20cc1..4d94f8ae 100644 --- a/docs/superpowers/plans/2026-07-04-svar2-genoray-readbound-gather.md +++ b/docs/superpowers/plans/2026-07-04-svar2-genoray-readbound-gather.md @@ -983,6 +983,8 @@ Create `tests/test_py_ranges_readbound.py` (run via the built wheel — this is ```python import numpy as np + + # genoray._core.PyContigReader is constructed the same way the existing # py_query_ranges tests do; reuse that harness path if one exists in genoray's # python test suite. Placeholder assertion of the new keys: diff --git a/docs/superpowers/plans/2026-07-04-svar2-gvl-readbound-wiring.md b/docs/superpowers/plans/2026-07-04-svar2-gvl-readbound-wiring.md index 89f2eda7..7063fa39 100644 --- a/docs/superpowers/plans/2026-07-04-svar2-gvl-readbound-wiring.md +++ b/docs/superpowers/plans/2026-07-04-svar2-gvl-readbound-wiring.md @@ -66,15 +66,21 @@ Create `tests/unit/dataset/test_svar2_link.py`: from pathlib import Path import pytest from genvarloader._dataset._svar2_link import ( - Svar2Link, Svar2Fingerprint, _resolve_svar2, _verify_svar2_fingerprint, + Svar2Link, + Svar2Fingerprint, + _resolve_svar2, + _verify_svar2_fingerprint, ) def test_resolve_prefers_override(tmp_path: Path): real = tmp_path / "cohort.svar2" real.mkdir() - link = Svar2Link(relative_path="nope.svar2", absolute_path="/nope.svar2", - fingerprint=Svar2Fingerprint(n_variants=1, store_bytes=1)) + link = Svar2Link( + relative_path="nope.svar2", + absolute_path="/nope.svar2", + fingerprint=Svar2Fingerprint(n_variants=1, store_bytes=1), + ) assert _resolve_svar2(tmp_path, link, real) == real @@ -103,6 +109,7 @@ Mirrors _svar_link.py; the fingerprint keys on the .svar2 store's stable identit (variant count + a canonical store-file byte count) rather than SVAR1's variant_idxs.npy, which .svar2 does not have. """ + from __future__ import annotations import os @@ -134,6 +141,7 @@ def _svar2_store_bytes(svar2_path: Path) -> int: def _svar2_n_variants(svar2_path: Path) -> int: import polars as pl + # .svar2 index; confirm filename against a real store (SparseVar2().index). return pl.scan_ipc(svar2_path / "index.arrow").select(pl.len()).collect().item() @@ -176,7 +184,9 @@ def _verify_svar2_fingerprint(svar2_path: Path, link: "Svar2Link | None") -> Non if n_obs != exp.n_variants: mismatches.append(f"n_variants: expected {exp.n_variants}, observed {n_obs}") if bytes_obs != exp.store_bytes: - mismatches.append(f"store_bytes: expected {exp.store_bytes}, observed {bytes_obs}") + mismatches.append( + f"store_bytes: expected {exp.store_bytes}, observed {bytes_obs}" + ) if mismatches: raise ValueError( f"svar2 fingerprint mismatch at {svar2_path}: " + "; ".join(mismatches) @@ -186,7 +196,9 @@ def _verify_svar2_fingerprint(svar2_path: Path, link: "Svar2Link | None") -> Non def make_svar2_link(gvl_path: Path, svar2_path: Path) -> Svar2Link: svar2_resolved = svar2_path.resolve() return Svar2Link( - relative_path=os.path.relpath(svar2_resolved, start=gvl_path).replace(os.sep, "/"), + relative_path=os.path.relpath(svar2_resolved, start=gvl_path).replace( + os.sep, "/" + ), absolute_path=str(svar2_resolved), fingerprint=Svar2Fingerprint( n_variants=_svar2_n_variants(svar2_resolved), @@ -261,20 +273,27 @@ SVAR2_FIXTURE = ... # Path to a small .svar2 store (reuse existing test fixture def test_write_svar2_emits_cache(tmp_path: Path): from genoray import SparseVar2 + svar2 = SparseVar2(SVAR2_FIXTURE) - bed = pl.DataFrame({ - "chrom": ["chr1", "chr1"], - "chromStart": [0, 250], - "chromEnd": [1000, 400], - }) + bed = pl.DataFrame( + { + "chrom": ["chr1", "chr1"], + "chromStart": [0, 250], + "chromEnd": [1000, 400], + } + ) out = tmp_path / "ds.gvl" gvl.write(out, bed, variants=svar2, samples=None, overwrite=True) rd = out / "genotypes" / "svar2_ranges" meta = json.loads((rd / "svar2_meta.json").read_text()) assert set(meta) >= { - "vk_snp_range", "vk_indel_range", "dense_snp_range", - "dense_indel_range", "region_starts", "sample_cols", + "vk_snp_range", + "vk_indel_range", + "dense_snp_range", + "dense_indel_range", + "region_starts", + "sample_cols", } # metadata.json carries the link + ploidy. md = json.loads((out / "metadata.json").read_text()) @@ -336,9 +355,13 @@ def _write_from_svar2( R, S, P = bed.height, len(samples), svar2.ploidy vk_snp = np.memmap(out_dir / "vk_snp_range.npy", np.int64, "w+", shape=(R, S, P, 2)) - vk_indel = np.memmap(out_dir / "vk_indel_range.npy", np.int64, "w+", shape=(R, S, P, 2)) + vk_indel = np.memmap( + out_dir / "vk_indel_range.npy", np.int64, "w+", shape=(R, S, P, 2) + ) dense_snp = np.memmap(out_dir / "dense_snp_range.npy", np.int64, "w+", shape=(R, 2)) - dense_indel = np.memmap(out_dir / "dense_indel_range.npy", np.int64, "w+", shape=(R, 2)) + dense_indel = np.memmap( + out_dir / "dense_indel_range.npy", np.int64, "w+", shape=(R, 2) + ) region_starts = np.memmap(out_dir / "region_starts.npy", np.int64, "w+", shape=(R,)) # sample_cols: selected slot -> original sample index (same for every contig). sample_cols_full = np.asarray( @@ -366,14 +389,18 @@ def _write_from_svar2( max_ends = np.empty(R, np.int32) contig_offset = 0 pbar = tqdm(total=R, unit=" region") - for (c,), df in bed.partition_by("chrom", as_dict=True, maintain_order=True).items(): + for (c,), df in bed.partition_by( + "chrom", as_dict=True, maintain_order=True + ).items(): c = cast(str, c) pbar.set_description(f"Processing svar2 ranges for {df.height} regions on {c}") lo, hi = contig_offset, contig_offset + df.height d = svar2.find_ranges( c, df["chromStart"].to_numpy(), - df["chromEnd"].to_numpy() if not extend_to_length else df["chromEnd"].to_numpy(), + df["chromEnd"].to_numpy() + if not extend_to_length + else df["chromEnd"].to_numpy(), samples=samples, ) rc = df.height @@ -397,6 +424,7 @@ def _write_from_svar2( mm.flush() from ._svar2_link import make_svar2_link + svar2_link = make_svar2_link(path, svar2.path) return bed.with_columns( chromEnd=pl.max_horizontal(pl.Series(max_ends), pl.col("chromEnd")) @@ -680,13 +708,16 @@ Create `tests/dataset/test_svar2_readbound_haps.py`. The oracle is the existing ```python import numpy as np from genvarloader._dataset._svar2_source import SparseVar2Source -from genvarloader._dataset._svar2_store_py import build_readbound_haps # thin py wrapper (Task 7) +from genvarloader._dataset._svar2_store_py import ( + build_readbound_haps, +) # thin py wrapper (Task 7) SVAR2_FIXTURE = ... # same fixture def test_readbound_haps_match_union_oracle(): from genoray import SparseVar2 + svar2 = SparseVar2(SVAR2_FIXTURE) contig = "chr1" regions = [(0, 1000), (250, 400), (150, 250)] @@ -695,8 +726,10 @@ def test_readbound_haps_match_union_oracle(): union = SparseVar2Source(svar2).reconstruct( contig, regions, ref, ref_offsets, pad, shifts=None, output_length=-1 ) - readbound = build_readbound_haps( # opens Svar2Store, slices no cache (direct find_ranges), - svar2, contig, regions, ref, ref_offsets, pad, shifts=None, output_length=-1 + readbound = ( + build_readbound_haps( # opens Svar2Store, slices no cache (direct find_ranges), + svar2, contig, regions, ref, ref_offsets, pad, shifts=None, output_length=-1 + ) ) # Ragged equality: same offsets + same bytes. assert np.array_equal(np.asarray(union.offsets), np.asarray(readbound.offsets)) @@ -743,14 +776,18 @@ Create `tests/dataset/test_svar2_readbound_tracks.py`, oracle = `SparseVar2Sourc ```python def test_readbound_tracks_match_union_oracle(): from genoray import SparseVar2 + svar2 = SparseVar2(SVAR2_FIXTURE) contig, regions = "chr1", [(0, 1000), (250, 400)] tracks, toff, params, strat, seed = _synthetic_track_inputs(regions) union = SparseVar2Source(svar2).realign_tracks( - contig, regions, tracks, toff, params, strat, seed, shifts=None) + contig, regions, tracks, toff, params, strat, seed, shifts=None + ) rb = build_readbound_tracks( - svar2, contig, regions, tracks, toff, params, strat, seed, shifts=None) + svar2, contig, regions, tracks, toff, params, strat, seed, shifts=None + ) import numpy as np + assert np.array_equal(np.asarray(union.offsets), np.asarray(rb.offsets)) assert np.allclose(union.data, rb.data, equal_nan=True) ``` @@ -826,7 +863,9 @@ import genvarloader as gvl def _open_pair(tmp_path, bed, svar_fixture, svar2_fixture, ref): from genoray import SparseVar, SparseVar2 - d1 = tmp_path / "d1.gvl"; d2 = tmp_path / "d2.gvl" + + d1 = tmp_path / "d1.gvl" + d2 = tmp_path / "d2.gvl" gvl.write(d1, bed, variants=SparseVar(svar_fixture), overwrite=True) gvl.write(d2, bed, variants=SparseVar2(svar2_fixture), overwrite=True) return gvl.Dataset.open(d1, reference=ref), gvl.Dataset.open(d2, reference=ref) @@ -840,7 +879,9 @@ def test_svar2_haplotypes_match_svar1(tmp_path, bed, svar_fixture, svar2_fixture assert np.array_equal(a.data.view("u1"), b.data.view("u1")) -def test_svar2_tracks_match_svar1(tmp_path, bed, svar_fixture, svar2_fixture, ref, bigwig): +def test_svar2_tracks_match_svar1( + tmp_path, bed, svar_fixture, svar2_fixture, ref, bigwig +): ds1, ds2 = _open_pair(tmp_path, bed, svar_fixture, svar2_fixture, ref) a = ds1.with_tracks(bigwig)[:, :] b = ds2.with_tracks(bigwig)[:, :] diff --git a/docs/superpowers/plans/2026-07-05-svar2-readbound-getitem-perf.md b/docs/superpowers/plans/2026-07-05-svar2-readbound-getitem-perf.md index 3dc3fa28..38cf6f3b 100644 --- a/docs/superpowers/plans/2026-07-05-svar2-readbound-getitem-perf.md +++ b/docs/superpowers/plans/2026-07-05-svar2-readbound-getitem-perf.md @@ -28,10 +28,17 @@ ```python import genvarloader as gvl from genoray import SparseVar2 - gvl.write(ds_path, bed, variants=SparseVar2(f"{prefix}.svar2"), samples=None, - max_jitter=0, overwrite=True) # ONCE, before the profiled loop + + gvl.write( + ds_path, + bed, + variants=SparseVar2(f"{prefix}.svar2"), + samples=None, + max_jitter=0, + overwrite=True, + ) # ONCE, before the profiled loop ds = gvl.Dataset.open(ds_path, reference=REF) - ds.with_seqs("haplotypes")[:, :] # or with_seqs("variants") + ds.with_seqs("haplotypes")[:, :] # or with_seqs("variants") ``` --- @@ -80,6 +87,7 @@ attribute cleanly. gvl.write + Dataset.open run ONCE (we profile the READ, not the write). Prints per_call_s over K warm calls. Tracks mode is out of scope; variant-windows is guarded NotImplementedError in Svar2Haps and cannot be profiled yet.""" + import sys import time from pathlib import Path @@ -94,11 +102,13 @@ WORK = Path("tmp/svar2_mvp/prof_out/readbound") def _bed(): - return pl.DataFrame({ - "chrom": [CHROM] * len(REGIONS), - "chromStart": [s for s, _ in REGIONS], - "chromEnd": [e for _, e in REGIONS], - }) + return pl.DataFrame( + { + "chrom": [CHROM] * len(REGIONS), + "chromStart": [s for s, _ in REGIONS], + "chromEnd": [e for _, e in REGIONS], + } + ) def make_call(mode, cohort): @@ -111,10 +121,16 @@ def make_call(mode, cohort): ds_path = WORK / f"{cohort}_{mode}.gvl" WORK.mkdir(parents=True, exist_ok=True) - gvl.write(ds_path, _bed(), variants=SparseVar2(f"{prefix}.svar2"), - samples=None, max_jitter=0, overwrite=True) + gvl.write( + ds_path, + _bed(), + variants=SparseVar2(f"{prefix}.svar2"), + samples=None, + max_jitter=0, + overwrite=True, + ) ds = gvl.Dataset.open(ds_path, reference=REF) - view = ds.with_seqs(mode) # "haplotypes" or "variants" + view = ds.with_seqs(mode) # "haplotypes" or "variants" R = len(REGIONS) @@ -179,6 +195,7 @@ Create `tmp/svar2_mvp/prof_python.py`: cProfile ranks Python functions by cumulative time; pyinstrument gives a low-overhead statistical wall-clock call tree as a cross-check (cProfile's own per-call overhead can distort tiny hot loops).""" + import cProfile import io import pstats @@ -202,6 +219,7 @@ def main(mode, cohort, K): print("```\n" + s.getvalue() + "```\n") from pyinstrument import Profiler + p = Profiler(interval=0.0005) p.start() for _ in range(K): @@ -353,16 +371,18 @@ def test_deterministic_haps_read_skips_pre_reconstruct_diffs(monkeypatch): calls = {"diffs": 0} real = m.hap_diffs_from_svar2_readbound + def counting(*a, **k): calls["diffs"] += 1 return real(*a, **k) + monkeypatch.setattr(m, "hap_diffs_from_svar2_readbound", counting) # Build the same small live svar2 dataset the module parity tests use, then: # ds.with_seqs("haplotypes")[:, :] # (reuse this file's existing fixture that yields a ds2 Svar2Haps-backed view; # if none is exposed, lift the _open_pair helper from test_svar2_dataset.py.) - ds2 = _svar2_haps_dataset() # existing/lifted fixture -> haplotypes view + ds2 = _svar2_haps_dataset() # existing/lifted fixture -> haplotypes view ds2[:, :] assert calls["diffs"] == 0 ``` @@ -379,38 +399,44 @@ Expected: FAIL with `assert 1 == 0` (the diffs kernel is currently called uncond In `python/genvarloader/_dataset/_svar2_haps.py`, add `need_hap_lengths: bool = False` to `get_haps_and_shifts`'s signature, and replace the unconditional diffs block + shifts block (currently ~lines 352-384) with: ```python - groups = self._contig_groups(contig_ids) - - # diffs are needed pre-reconstruct ONLY to (a) bound randomized jitter - # shifts, or (b) return hap_lengths/diffs to a caller that uses them - # (the tracks path). A deterministic/ragged haplotypes read needs - # neither: reconstruct sizes itself internally. Avoid the redundant - # gather+split+diffs in that (common warm-read) case. - randomized = not (deterministic or isinstance(output_length, str)) - need_diffs = randomized or need_hap_lengths - - if need_diffs: - diffs = np.empty((b, P), np.int32) - for ci, qsel in groups: - gi = self._gather_inputs(r_q[qsel], si_q[qsel], regions[qsel], P) - d = hap_diffs_from_svar2_readbound( - self.store, self.ds_contigs[ci], - gi[0], gi[1], gi[2], gi[3], gi[4], gi[5], gi[6], P, - ) - diffs[qsel] = np.asarray(d, np.int32).reshape(len(qsel), P) - hap_lengths = (lengths[:, None] + diffs).astype(np.int32) - else: - diffs = np.zeros((b, P), np.int32) # placeholder (unused downstream) - hap_lengths = np.broadcast_to( - lengths[:, None].astype(np.int32), (b, P) - ).copy() - - if randomized: - max_shift = diffs.clip(min=0) - max_shift = max_shift + (lengths - output_length).clip(min=0)[:, None] - shifts = rng.integers(0, max_shift + 1, dtype=np.int32) - else: - shifts = np.zeros((b, P), np.int32) +groups = self._contig_groups(contig_ids) + +# diffs are needed pre-reconstruct ONLY to (a) bound randomized jitter +# shifts, or (b) return hap_lengths/diffs to a caller that uses them +# (the tracks path). A deterministic/ragged haplotypes read needs +# neither: reconstruct sizes itself internally. Avoid the redundant +# gather+split+diffs in that (common warm-read) case. +randomized = not (deterministic or isinstance(output_length, str)) +need_diffs = randomized or need_hap_lengths + +if need_diffs: + diffs = np.empty((b, P), np.int32) + for ci, qsel in groups: + gi = self._gather_inputs(r_q[qsel], si_q[qsel], regions[qsel], P) + d = hap_diffs_from_svar2_readbound( + self.store, + self.ds_contigs[ci], + gi[0], + gi[1], + gi[2], + gi[3], + gi[4], + gi[5], + gi[6], + P, + ) + diffs[qsel] = np.asarray(d, np.int32).reshape(len(qsel), P) + hap_lengths = (lengths[:, None] + diffs).astype(np.int32) +else: + diffs = np.zeros((b, P), np.int32) # placeholder (unused downstream) + hap_lengths = np.broadcast_to(lengths[:, None].astype(np.int32), (b, P)).copy() + +if randomized: + max_shift = diffs.clip(min=0) + max_shift = max_shift + (lengths - output_length).clip(min=0)[:, None] + shifts = rng.integers(0, max_shift + 1, dtype=np.int32) +else: + shifts = np.zeros((b, P), np.int32) ``` Then have the tracks caller pass `need_hap_lengths=True`. Find the caller: `grep -n "get_haps_and_shifts" python/genvarloader/_dataset/*.py` — it is invoked from `HapsTracks` dispatch (the tracks path) and from `Svar2Haps.__call__` (haplotypes). Update the tracks call site to `get_haps_and_shifts(..., need_hap_lengths=True)`; leave the haplotypes call site at the `False` default. diff --git a/docs/superpowers/plans/2026-07-06-svar2-variant-windows.md b/docs/superpowers/plans/2026-07-06-svar2-variant-windows.md index fdff4701..c4ce092e 100644 --- a/docs/superpowers/plans/2026-07-06-svar2-variant-windows.md +++ b/docs/superpowers/plans/2026-07-06-svar2-variant-windows.md @@ -158,18 +158,16 @@ def _reconstruct_variant_windows( for ci, qsel in groups: gi = self._gather_inputs(r_q[qsel], si_q[qsel], regions[qsel], P) - pos, ilen, alt_bytes, str_off, var_off = ( - decode_variants_from_svar2_readbound( - self.store, - self.ds_contigs[ci], - gi[0], - gi[1], - gi[2], - gi[3], - gi[4], - gi[5], - P, - ) + pos, ilen, alt_bytes, str_off, var_off = decode_variants_from_svar2_readbound( + self.store, + self.ds_contigs[ci], + gi[0], + gi[1], + gi[2], + gi[3], + gi[4], + gi[5], + P, ) pos = np.asarray(pos, np.int32) ilen = np.asarray(ilen, np.int32) @@ -333,7 +331,11 @@ def test_svar2_variant_windows_alt_window_decomposition( L = _WIN_OPT.flank_length w_win = ds2.with_output_format("flat").with_seqs("variant-windows", _WIN_OPT)[:, :] alt_opt = VarWindowOpt( - flank_length=L, token_alphabet=b"ACGT", unknown_token=4, ref="window", alt="allele" + flank_length=L, + token_alphabet=b"ACGT", + unknown_token=4, + ref="window", + alt="allele", ) w_alt = ds2.with_output_format("flat").with_seqs("variant-windows", alt_opt)[:, :] @@ -375,7 +377,11 @@ def test_svar2_variant_windows_bare_alt_tokenizes_variants_alt( _, ds2 = _open_pair(tmp_path, bed, svar_fixture, svar2_fixture, ref) L = _WIN_OPT.flank_length alt_opt = VarWindowOpt( - flank_length=L, token_alphabet=b"ACGT", unknown_token=4, ref="window", alt="allele" + flank_length=L, + token_alphabet=b"ACGT", + unknown_token=4, + ref="window", + alt="allele", ) w_alt = ds2.with_output_format("flat").with_seqs("variant-windows", alt_opt)[:, :] v = ds2.with_seqs("variants")[:, :] # RaggedVariants (validated) @@ -401,7 +407,9 @@ def test_svar2_variant_windows_bare_alt_tokenizes_variants_alt( - [ ] **Step 9: Add multi-contig parity test** ```python -def test_svar2_variant_windows_multicontig(tmp_path, svar_fixture2, svar2_fixture2, _src2): +def test_svar2_variant_windows_multicontig( + tmp_path, svar_fixture2, svar2_fixture2, _src2 +): """ref_window byte-identical to SVAR1 across an interleaved 2-contig bed (single-contig fast path bypassed -> exercises the group-stitch reorder).""" from genoray import SparseVar, SparseVar2 @@ -417,7 +425,9 @@ def test_svar2_variant_windows_multicontig(tmp_path, svar_fixture2, svar2_fixtur d1 = tmp_path / "vw_mc1.gvl" d2 = tmp_path / "vw_mc2.gvl" gvl.write(d1, bed, variants=SparseVar(svar_fixture2), samples=None, overwrite=True) - gvl.write(d2, bed, variants=SparseVar2(svar2_fixture2), samples=None, overwrite=True) + gvl.write( + d2, bed, variants=SparseVar2(svar2_fixture2), samples=None, overwrite=True + ) ds1 = gvl.Dataset.open(d1, reference=ref) ds2 = gvl.Dataset.open(d2, reference=ref) w1 = ds1.with_output_format("flat").with_seqs("variant-windows", _WIN_OPT)[:, :] @@ -471,7 +481,11 @@ def test_svar2_variant_windows_ref_allele_guard(tmp_path, bed, svar2_fixture, _s gvl.write(d, bed, variants=SparseVar2(svar2_fixture), samples=None, overwrite=True) ds = gvl.Dataset.open(d, reference=ref).with_output_format("flat") bad = VarWindowOpt( - flank_length=3, token_alphabet=b"ACGT", unknown_token=4, ref="allele", alt="window" + flank_length=3, + token_alphabet=b"ACGT", + unknown_token=4, + ref="allele", + alt="window", ) with pytest.raises(ValueError, match="REF"): ds.with_seqs("variant-windows", bad) @@ -485,7 +499,12 @@ def test_svar2_variant_windows_jitter_guard(tmp_path, svar2_fixture, _src): jbed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [5], "chromEnd": [20]}) d = tmp_path / "d.gvl" gvl.write( - d, jbed, variants=SparseVar2(svar2_fixture), samples=None, max_jitter=2, overwrite=True + d, + jbed, + variants=SparseVar2(svar2_fixture), + samples=None, + max_jitter=2, + overwrite=True, ) ds = gvl.Dataset.open(d, reference=ref).with_output_format("flat") with pytest.raises(NotImplementedError, match="right-clip"): @@ -553,10 +572,8 @@ Expected: FAIL with `NotImplementedError: unphased_union is not supported for sv In `Svar2Haps._guard_unsupported`, delete: ```python - if self.unphased_union: - raise NotImplementedError( - "unphased_union is not supported for svar2 datasets yet." - ) +if self.unphased_union: + raise NotImplementedError("unphased_union is not supported for svar2 datasets yet.") ``` (Haplotypes/annotated + union is still blocked at `with_seqs` in `_impl.py`, so the haplotypes/tracks paths can never reach here with the flag set.) @@ -636,7 +653,9 @@ def test_svar2_variant_windows_unphased_union( _assert_window_equal(w2.ref_window, w1.ref_window, "ref_window") # Union row count == sum over haplotypes: compare to the non-union var counts. nu = np.asarray(w2.ref_window.var_offsets) - w2_diploid = ds2.with_output_format("flat").with_seqs("variant-windows", _WIN_OPT)[:, :] + w2_diploid = ds2.with_output_format("flat").with_seqs("variant-windows", _WIN_OPT)[ + :, : + ] nd = np.asarray(w2_diploid.ref_window.var_offsets) P = int(ds2._seqs.genotypes.shape[-2]) # Folded per-row counts == sum of the P per-hap counts (rows q*P+p are contiguous). @@ -687,15 +706,25 @@ def test_svar2_variant_windows_union_multicontig( d1 = tmp_path / "vwu_mc1.gvl" d2 = tmp_path / "vwu_mc2.gvl" gvl.write(d1, bed, variants=SparseVar(svar_fixture2), samples=None, overwrite=True) - gvl.write(d2, bed, variants=SparseVar2(svar2_fixture2), samples=None, overwrite=True) + gvl.write( + d2, bed, variants=SparseVar2(svar2_fixture2), samples=None, overwrite=True + ) ds1 = gvl.Dataset.open(d1, reference=ref) ds2 = gvl.Dataset.open(d2, reference=ref) - w1 = (ds1.with_output_format("flat").with_settings(unphased_union=True) - .with_seqs("variant-windows", _WIN_OPT)[:, :]) - w2 = (ds2.with_output_format("flat").with_settings(unphased_union=True) - .with_seqs("variant-windows", _WIN_OPT)[:, :]) + w1 = ( + ds1.with_output_format("flat") + .with_settings(unphased_union=True) + .with_seqs("variant-windows", _WIN_OPT)[:, :] + ) + w2 = ( + ds2.with_output_format("flat") + .with_settings(unphased_union=True) + .with_seqs("variant-windows", _WIN_OPT)[:, :] + ) assert w2.ref_window.shape[-3] == 1 # window ploidy axis - _assert_window_equal(w2.ref_window, w1.ref_window, "ref_window (union, multicontig)") + _assert_window_equal( + w2.ref_window, w1.ref_window, "ref_window (union, multicontig)" + ) w2.alt_window.to_ragged() ``` diff --git a/docs/superpowers/plans/2026-07-12-svar2-info-format-field-routing.md b/docs/superpowers/plans/2026-07-12-svar2-info-format-field-routing.md index f32a7c1b..31eee84b 100644 --- a/docs/superpowers/plans/2026-07-12-svar2-info-format-field-routing.md +++ b/docs/superpowers/plans/2026-07-12-svar2-info-format-field-routing.md @@ -641,10 +641,11 @@ and pass `store_field_keys=store_field_keys, store_fields=store_fields` into the Replace the hard-coded line (`:174`): ```python - self.available_var_fields = ["alt", "ilen", "start"] + [ - k for k in self.store_field_keys - if k not in {"alt", "ilen", "start", "ref", "dosage"} - ] +self.available_var_fields = ["alt", "ilen", "start"] + [ + k + for k in self.store_field_keys + if k not in {"alt", "ilen", "start", "ref", "dosage"} +] ``` - [ ] **Step 3: Rebuild + smoke-check discovery** @@ -740,12 +741,20 @@ At the top of `_reconstruct_variants`, compute the requested extra fields (order Change the kernel call to pass `field_specs`, and unpack the two new returns per contig group: ```python - pos, ilen, alt_bytes, str_off, var_off, field_bufs, field_isizes = ( - decode_variants_from_svar2_readbound( - self.store, self.ds_contigs[ci], - gi[0], gi[1], gi[2], gi[3], gi[4], gi[5], P, field_specs, - ) - ) +pos, ilen, alt_bytes, str_off, var_off, field_bufs, field_isizes = ( + decode_variants_from_svar2_readbound( + self.store, + self.ds_contigs[ci], + gi[0], + gi[1], + gi[2], + gi[3], + gi[4], + gi[5], + P, + field_specs, + ) +) ``` Accumulate per-field buffers per group into `cat_fields: list[list[np.ndarray]]` (one inner list per requested field), asserting `field_isizes[j] == field_dtypes[j].itemsize`. @@ -755,18 +764,20 @@ Accumulate per-field buffers per group into `cat_fields: list[list[np.ndarray]]` In the `len(cat_pos) == 1` branch, build a `fields` dict parallel to the existing `alt`/`start`/`ilen`, then splat into `RaggedVariants`: ```python - extra = { - req_keys[j]: Ragged.from_offsets( - cat_fields[j][0].view(field_dtypes[j]), shape, var_off_g - ) - for j in range(len(req_keys)) - } - return RaggedVariants( - alt=Ragged.from_offsets(cat_alt[0].view("S1"), shape, var_off_g, str_offsets=str_off_g), - start=Ragged.from_offsets(cat_pos[0], shape, var_off_g), - ilen=Ragged.from_offsets(cat_ilen[0], shape, var_off_g), - **extra, - ) +extra = { + req_keys[j]: Ragged.from_offsets( + cat_fields[j][0].view(field_dtypes[j]), shape, var_off_g + ) + for j in range(len(req_keys)) +} +return RaggedVariants( + alt=Ragged.from_offsets( + cat_alt[0].view("S1"), shape, var_off_g, str_offsets=str_off_g + ), + start=Ragged.from_offsets(cat_pos[0], shape, var_off_g), + ilen=Ragged.from_offsets(cat_ilen[0], shape, var_off_g), + **extra, +) ``` - [ ] **Step 3: Multi-contig path — reorder fields by the same `src`** @@ -774,13 +785,17 @@ In the `len(cat_pos) == 1` branch, build a `fields` dict parallel to the existin In the general path, after computing `src, var_off_g = _ragged_arange_src(grouped_var_off, perm)`, each field is per-variant so it reorders exactly like `pos`: ```python - extra = {} - for j in range(len(req_keys)): - fc = np.concatenate([g[j] for g in per_group_fields]) if per_group_fields else np.zeros(0, np.uint8) - fc_typed = fc.view(field_dtypes[j]) - fg = fc_typed[:0].copy() if src.size == 0 else fc_typed[src] - extra[req_keys[j]] = Ragged.from_offsets(fg, shape, var_off_g) - return RaggedVariants(alt=alt_r, start=pos_r, ilen=ilen_r, **extra) +extra = {} +for j in range(len(req_keys)): + fc = ( + np.concatenate([g[j] for g in per_group_fields]) + if per_group_fields + else np.zeros(0, np.uint8) + ) + fc_typed = fc.view(field_dtypes[j]) + fg = fc_typed[:0].copy() if src.size == 0 else fc_typed[src] + extra[req_keys[j]] = Ragged.from_offsets(fg, shape, var_off_g) +return RaggedVariants(alt=alt_r, start=pos_r, ilen=ilen_r, **extra) ``` (`per_group_fields[g][j]` is group `g`'s buffer for field `j`, already `.view(dtype)`-ed to length `n_var_group`. Reuse the same `src` computed for `pos_g`/`ilen_g`.) @@ -873,7 +888,9 @@ from genoray import SparseVar2 from genoray._svar2_fields import InfoField, FormatField SparseVar2.from_vcf( - out=store_dir, source=vcf_gz, reference=fasta, + out=store_dir, + source=vcf_gz, + reference=fasta, info_fields=[InfoField("AF"), InfoField("NS")], format_fields=[FormatField("DP")], ) @@ -889,12 +906,15 @@ Parse the VCF with `cyvcf2` into a dict `{(contig, pos): {"AF": ..., "NS": ..., import numpy as np import pytest + @pytest.mark.parametrize("union", [False, True]) def test_svar2_ragged_variants_fields(tmp_path, union): ds_path = _write_dataset(tmp_path) # gvl.write over the svar2 source import genvarloader as gvl + ds = gvl.Dataset.open(ds_path, reference=FASTA).with_seqs( - "variants", var_fields=["alt", "start", "ilen", "AF", "NS", "DP"], + "variants", + var_fields=["alt", "start", "ilen", "AF", "NS", "DP"], ) if union: ds = ds.with_settings(unphased_union=True) # use the real gvl API name diff --git a/docs/superpowers/plans/2026-07-13-svar2-m6b-kernel-final-pass.md b/docs/superpowers/plans/2026-07-13-svar2-m6b-kernel-final-pass.md index 789dfe90..c836684e 100644 --- a/docs/superpowers/plans/2026-07-13-svar2-m6b-kernel-final-pass.md +++ b/docs/superpowers/plans/2026-07-13-svar2-m6b-kernel-final-pass.md @@ -283,6 +283,7 @@ def test_readbound_haps_noncontiguous_input_raises(): import numpy as np import pytest from genvarloader._dataset._svar2_store_py import build_readbound_haps # noqa: F401 + # Build a minimal store + regions exactly as the existing haps parity test does, # then pass a strided (non-contiguous) view of one of the int64 range arrays. # (Reuse the fixture/store construction from test_readbound_haps_* above.) @@ -372,7 +373,10 @@ def test_svar2_extend_to_length_false_raises(tmp_path, svar2_source_and_bed): must raise, not silently produce an extended dataset.""" import pytest import genvarloader as gvl - svar2, bed = svar2_source_and_bed # reuse the existing fixture used by the write tests + + svar2, bed = ( + svar2_source_and_bed # reuse the existing fixture used by the write tests + ) with pytest.raises(NotImplementedError, match="extend_to_length"): gvl.write(tmp_path / "ds", bed, variants=svar2, extend_to_length=False) ``` @@ -445,8 +449,11 @@ def test_svar2_region_max_ends_matches_reference(svar2_source_and_bed): including the pos-then-end tie-break and the empty-region default = chromEnd.""" import numpy as np from genvarloader._dataset._write import _svar2_region_max_ends + svar2, bed = svar2_source_and_bed - for (c,), df in bed.partition_by("chrom", as_dict=True, maintain_order=True).items(): + for (c,), df in bed.partition_by( + "chrom", as_dict=True, maintain_order=True + ).items(): starts = df["chromStart"].to_numpy() ends = df["chromEnd"].to_numpy() samples = list(svar2.available_samples) @@ -459,10 +466,13 @@ def test_svar2_region_max_ends_matches_reference(svar2_source_and_bed): def _reference_region_max_ends(svar2, contig, starts, ends, samples): """Byte-for-byte copy of the ORIGINAL triple-loop, kept in the test as the oracle.""" import numpy as np + R, S_all, P = len(starts), svar2.n_samples, svar2.ploidy sel = [svar2.available_samples.index(s) for s in samples] dec = svar2.decode(contig, list(zip(starts.tolist(), ends.tolist()))) - pos_arr = dec.data["pos"]; ilen_arr = dec.data["ilen"]; off = np.asarray(dec.offsets) + pos_arr = dec.data["pos"] + ilen_arr = dec.data["ilen"] + off = np.asarray(dec.offsets) out = np.asarray(ends, np.int64).copy() for r in range(R): best_pos, best_end = -1, -1 @@ -470,10 +480,13 @@ def _reference_region_max_ends(svar2, contig, starts, ends, samples): for p in range(P): h = (r * S_all + s) * P + p a, b = int(off[h]), int(off[h + 1]) - if a == b: continue - seg_pos = pos_arr[a:b]; seg_ilen = ilen_arr[a:b] + if a == b: + continue + seg_pos = pos_arr[a:b] + seg_ilen = ilen_arr[a:b] j = int(np.argmax(seg_pos)) - p_pos = int(seg_pos[j]); p_end = (p_pos + 1) - min(int(seg_ilen[j]), 0) + p_pos = int(seg_pos[j]) + p_end = (p_pos + 1) - min(int(seg_ilen[j]), 0) if p_pos > best_pos or (p_pos == best_pos and p_end > best_end): best_pos, best_end = p_pos, p_end if best_pos >= 0: @@ -489,31 +502,38 @@ Expected: PASS. - [ ] **Step 3: Vectorize the function.** Replace the triple-loop body of `_svar2_region_max_ends` (keeping the docstring's semantics but dropping the "O(...) Python iteration ... vectorize as a follow-up" caveat) with a scatter-reduce. Key idea: for each variant, its haplotype maps to a region `r = h // (S_all * P)` but only SELECTED samples count; compute `end = (pos+1) - min(ilen,0)` per variant, form a sortable composite `key = (pos << 21) | end` (end fits well under 2^21 for realistic regions; assert it) so that a plain per-region max on `key` reproduces the pos-then-end tie-break, then unpack `end`: ```python - R, S_all, P = len(starts), svar2.n_samples, svar2.ploidy - sel = np.asarray([svar2.available_samples.index(s) for s in samples], np.int64) - dec = svar2.decode(contig, list(zip(starts.tolist(), ends.tolist()))) - pos_arr = np.asarray(dec.data["pos"], np.int64) - ilen_arr = np.asarray(dec.data["ilen"], np.int64) - off = np.asarray(dec.offsets, np.int64) # length R*S_all*P + 1 - out = np.asarray(ends, np.int64).copy() # default = chromEnd - if pos_arr.size: - n_hap = R * S_all * P - counts = np.diff(off) # variants per hap - hap_of_var = np.repeat(np.arange(n_hap), counts) # region-major hap index per variant - s_of_hap = (np.arange(n_hap) // P) % S_all - keep = np.isin(s_of_hap[hap_of_var], sel) # only selected samples - region_of_var = hap_of_var // (S_all * P) - end_var = (pos_arr + 1) - np.minimum(ilen_arr, 0) # 0-based -> 1-based, extend on DEL - SHIFT = 21 - assert int(end_var.max(initial=0)) < (1 << SHIFT), "end exceeds tie-break packing width" - key = (pos_arr << SHIFT) | end_var - key_k = key[keep]; region_k = region_of_var[keep] - if key_k.size: - best = np.full(R, -1, np.int64) - np.maximum.at(best, region_k, key_k) # per-region max composite key - has = best >= 0 - out[has] = best[has] & ((1 << SHIFT) - 1) # unpack end - return out.astype(np.int32) +R, S_all, P = len(starts), svar2.n_samples, svar2.ploidy +sel = np.asarray([svar2.available_samples.index(s) for s in samples], np.int64) +dec = svar2.decode(contig, list(zip(starts.tolist(), ends.tolist()))) +pos_arr = np.asarray(dec.data["pos"], np.int64) +ilen_arr = np.asarray(dec.data["ilen"], np.int64) +off = np.asarray(dec.offsets, np.int64) # length R*S_all*P + 1 +out = np.asarray(ends, np.int64).copy() # default = chromEnd +if pos_arr.size: + n_hap = R * S_all * P + counts = np.diff(off) # variants per hap + hap_of_var = np.repeat( + np.arange(n_hap), counts + ) # region-major hap index per variant + s_of_hap = (np.arange(n_hap) // P) % S_all + keep = np.isin(s_of_hap[hap_of_var], sel) # only selected samples + region_of_var = hap_of_var // (S_all * P) + end_var = (pos_arr + 1) - np.minimum( + ilen_arr, 0 + ) # 0-based -> 1-based, extend on DEL + SHIFT = 21 + assert int(end_var.max(initial=0)) < (1 << SHIFT), ( + "end exceeds tie-break packing width" + ) + key = (pos_arr << SHIFT) | end_var + key_k = key[keep] + region_k = region_of_var[keep] + if key_k.size: + best = np.full(R, -1, np.int64) + np.maximum.at(best, region_k, key_k) # per-region max composite key + has = best >= 0 + out[has] = best[has] & ((1 << SHIFT) - 1) # unpack end +return out.astype(np.int32) ``` Update the docstring: drop the last paragraph ("O(R * len(samples) * ploidy) Python iteration ... follow-up") and replace with a one-line note that it is a vectorized per-region scatter-max preserving the pos-then-end tie-break. diff --git a/docs/superpowers/plans/2026-07-16-svar2-spliced-scatter-write.md b/docs/superpowers/plans/2026-07-16-svar2-spliced-scatter-write.md index 0dab5c25..2489098e 100644 --- a/docs/superpowers/plans/2026-07-16-svar2-spliced-scatter-write.md +++ b/docs/superpowers/plans/2026-07-16-svar2-spliced-scatter-write.md @@ -699,8 +699,8 @@ Expected: PASS (including Task 3's multi-contig test). These are the byte-identi In `python/genvarloader/_dataset/_svar2_haps.py`, add to the `from ..genvarloader import (...)` block (line ~43, keep alphabetical): ```python - reconstruct_haplotypes_from_svar2_readbound, - reconstruct_haplotypes_from_svar2_readbound_into, +(reconstruct_haplotypes_from_svar2_readbound,) +(reconstruct_haplotypes_from_svar2_readbound_into,) ``` - [ ] **Step 3: Add the `_reconstruct_spliced` method** @@ -708,94 +708,90 @@ In `python/genvarloader/_dataset/_svar2_haps.py`, add to the `from ..genvarloade Insert after `__call__` (before `haplotype_lengths_for_plan`, line ~396): ```python - def _reconstruct_spliced( - self, - idx: NDArray[np.integer], - regions: NDArray[np.int32], - splice_plan: "SplicePlan", - to_rc: "NDArray[np.bool_] | None", - ) -> _Flat[np.bytes_]: - """Reconstruct spliced haplotypes directly into spliced layout (no re-order). - - The splice plan already knows every element's final address, so instead of - reconstructing in region order and permuting the OUTPUT BYTES afterwards, we - permute the per-row METADATA (O(rows)) and let each contig group's kernel call - scatter straight into the shared buffer — the same trick SVAR1's fused spliced - entry uses (``reconstruct_haplotypes_spliced_fused``). - - The plan's k-index (``k = query * E + e`` with ``E = ploidy`` for haplotypes, - see ``_splice.build_splice_plan``) is exactly the kernel's row index - ``k = q * P + p``, so ``plan.permutation`` indexes hap rows with no translation. - - Callers reach this only via ``_getitem_spliced``, which asserts ``jitter == 0`` - and ``deterministic`` — hence zero shifts. - """ - assert self.store is not None - regions = np.asarray(regions, np.int32) - P = int(self.genotypes.shape[-2]) - b = len(idx) - R_all, S_all = int(self.genotypes.shape[0]), int(self.genotypes.shape[1]) - r_q, si_q = np.unravel_index(np.asarray(idx), (R_all, S_all)) - - perm = np.asarray(splice_plan.permutation, np.intp) - off = np.asarray(splice_plan.permuted_out_offsets, np.int64) - n_work = b * P - if len(perm) != n_work: - raise AssertionError( - f"splice permutation length {len(perm)} != n_queries*ploidy {n_work}" - ) - - # dest_rank[k] = position of kernel row k within the permuted (spliced) layout. - dest_rank = np.empty(n_work, np.intp) - dest_rank[perm] = np.arange(n_work, dtype=np.intp) - bounds_all = np.empty((n_work, 2), np.int64) - bounds_all[:, 0] = off[dest_rank] - bounds_all[:, 1] = off[dest_rank + 1] - - # to_rc arrives in permuted order (_getitem_spliced builds it as - # to_rc_flat[plan.permutation]); the kernel wants it per row. - rc_all: NDArray[np.bool_] | None = None - if to_rc is not None and bool(np.asarray(to_rc).any()): - rc_all = np.empty(n_work, np.bool_) - rc_all[perm] = np.asarray(to_rc, np.bool_) - - out = np.empty(int(off[-1]), np.uint8) - shifts_all = np.zeros((b, P), np.int32) - p_range = np.arange(P, dtype=np.intp) - - for ci, qsel in self._contig_groups(regions[:, 0].astype(np.int64)): - gi = self._gather_inputs(r_q[qsel], si_q[qsel], regions[qsel], P) - ref_, ref_offsets = self._ref_for_contig(ci) - rows = (qsel[:, None] * P + p_range).ravel() - g_bounds = np.ascontiguousarray(bounds_all[rows], np.int64) - g_rc = ( - None - if rc_all is None - else np.ascontiguousarray(rc_all[rows], np.bool_) - ) - g_total = int((g_bounds[:, 1] - g_bounds[:, 0]).sum()) - reconstruct_haplotypes_from_svar2_readbound_into( - out, - g_bounds, - self.store, - self.ds_contigs[ci], - gi[0], - gi[1], - gi[2], - gi[3], - gi[4], - gi[5], - gi[6], - np.ascontiguousarray(shifts_all[qsel], np.int32), - ref_, - ref_offsets, - np.uint8(self.reference.pad_char), # type: ignore[union-attr] # reference guaranteed for haplotypes - g_rc, - should_parallelize(g_total), - self.filter == "exonic", - ) - - return _Flat.from_offsets(out, (len(perm), None), off).view("S1") +def _reconstruct_spliced( + self, + idx: NDArray[np.integer], + regions: NDArray[np.int32], + splice_plan: "SplicePlan", + to_rc: "NDArray[np.bool_] | None", +) -> _Flat[np.bytes_]: + """Reconstruct spliced haplotypes directly into spliced layout (no re-order). + + The splice plan already knows every element's final address, so instead of + reconstructing in region order and permuting the OUTPUT BYTES afterwards, we + permute the per-row METADATA (O(rows)) and let each contig group's kernel call + scatter straight into the shared buffer — the same trick SVAR1's fused spliced + entry uses (``reconstruct_haplotypes_spliced_fused``). + + The plan's k-index (``k = query * E + e`` with ``E = ploidy`` for haplotypes, + see ``_splice.build_splice_plan``) is exactly the kernel's row index + ``k = q * P + p``, so ``plan.permutation`` indexes hap rows with no translation. + + Callers reach this only via ``_getitem_spliced``, which asserts ``jitter == 0`` + and ``deterministic`` — hence zero shifts. + """ + assert self.store is not None + regions = np.asarray(regions, np.int32) + P = int(self.genotypes.shape[-2]) + b = len(idx) + R_all, S_all = int(self.genotypes.shape[0]), int(self.genotypes.shape[1]) + r_q, si_q = np.unravel_index(np.asarray(idx), (R_all, S_all)) + + perm = np.asarray(splice_plan.permutation, np.intp) + off = np.asarray(splice_plan.permuted_out_offsets, np.int64) + n_work = b * P + if len(perm) != n_work: + raise AssertionError( + f"splice permutation length {len(perm)} != n_queries*ploidy {n_work}" + ) + + # dest_rank[k] = position of kernel row k within the permuted (spliced) layout. + dest_rank = np.empty(n_work, np.intp) + dest_rank[perm] = np.arange(n_work, dtype=np.intp) + bounds_all = np.empty((n_work, 2), np.int64) + bounds_all[:, 0] = off[dest_rank] + bounds_all[:, 1] = off[dest_rank + 1] + + # to_rc arrives in permuted order (_getitem_spliced builds it as + # to_rc_flat[plan.permutation]); the kernel wants it per row. + rc_all: NDArray[np.bool_] | None = None + if to_rc is not None and bool(np.asarray(to_rc).any()): + rc_all = np.empty(n_work, np.bool_) + rc_all[perm] = np.asarray(to_rc, np.bool_) + + out = np.empty(int(off[-1]), np.uint8) + shifts_all = np.zeros((b, P), np.int32) + p_range = np.arange(P, dtype=np.intp) + + for ci, qsel in self._contig_groups(regions[:, 0].astype(np.int64)): + gi = self._gather_inputs(r_q[qsel], si_q[qsel], regions[qsel], P) + ref_, ref_offsets = self._ref_for_contig(ci) + rows = (qsel[:, None] * P + p_range).ravel() + g_bounds = np.ascontiguousarray(bounds_all[rows], np.int64) + g_rc = None if rc_all is None else np.ascontiguousarray(rc_all[rows], np.bool_) + g_total = int((g_bounds[:, 1] - g_bounds[:, 0]).sum()) + reconstruct_haplotypes_from_svar2_readbound_into( + out, + g_bounds, + self.store, + self.ds_contigs[ci], + gi[0], + gi[1], + gi[2], + gi[3], + gi[4], + gi[5], + gi[6], + np.ascontiguousarray(shifts_all[qsel], np.int32), + ref_, + ref_offsets, + np.uint8(self.reference.pad_char), # type: ignore[union-attr] # reference guaranteed for haplotypes + g_rc, + should_parallelize(g_total), + self.filter == "exonic", + ) + + return _Flat.from_offsets(out, (len(perm), None), off).view("S1") ``` - [ ] **Step 4: Route `__call__` to it** diff --git a/docs/superpowers/plans/2026-07-17-svar2-spliced-variants-bench.md b/docs/superpowers/plans/2026-07-17-svar2-spliced-variants-bench.md index ae845b3e..e612bf79 100644 --- a/docs/superpowers/plans/2026-07-17-svar2-spliced-variants-bench.md +++ b/docs/superpowers/plans/2026-07-17-svar2-spliced-variants-bench.md @@ -111,6 +111,7 @@ import json, subprocess from dataclasses import dataclass from pathlib import Path + @dataclass(frozen=True) class CohortResult: bcf: Path @@ -119,19 +120,47 @@ class CohortResult: n_samples: int sample_names: tuple[str, ...] -def gen_cohort(samples, records, *, contig="chr1", seed=42, - profile="germline-1kgp", payload="gt-only", - vcfixture_bin, out_dir): - out_dir = Path(out_dir); out_dir.mkdir(parents=True, exist_ok=True) + +def gen_cohort( + samples, + records, + *, + contig="chr1", + seed=42, + profile="germline-1kgp", + payload="gt-only", + vcfixture_bin, + out_dir, +): + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) raw = out_dir / f"cohort_s{samples}_r{records}_seed{seed}.raw.bcf" - subprocess.run([str(vcfixture_bin), "bulk", "--profile", profile, - "--samples", str(samples), "--contigs", contig, - "--records", str(records), "--payload", payload, - "--seed", str(seed), "-o", str(raw)], check=True) + subprocess.run( + [ + str(vcfixture_bin), + "bulk", + "--profile", + profile, + "--samples", + str(samples), + "--contigs", + contig, + "--records", + str(records), + "--payload", + payload, + "--seed", + str(seed), + "-o", + str(raw), + ], + check=True, + ) # gvl requires bi-allelic, left-aligned, atomized (docs: bcftools norm). norm = out_dir / f"cohort_s{samples}_r{records}_seed{seed}.bcf" - subprocess.run(["bcftools", "norm", "-m-", "-Ob", "-o", str(norm), str(raw)], - check=True) + subprocess.run( + ["bcftools", "norm", "-m-", "-Ob", "-o", str(norm), str(raw)], check=True + ) subprocess.run(["bcftools", "index", "-f", str(norm)], check=True) summary = json.loads((raw.with_suffix(".summary.json")).read_text()) # populated span = declared contig length (vcfixture sets length = last POS). @@ -176,6 +205,7 @@ git commit -m "bench(svar2): cohort generator wrapping vcfixture bulk" import numpy as np, polars as pl, subprocess from pathlib import Path + def _write_reference(path: Path, contig: str, length: int, seed=0): rng = np.random.default_rng(seed) seq = rng.choice(np.frombuffer(b"ACGT", "S1"), size=length).tobytes().decode() @@ -201,8 +231,18 @@ def _splice_bed(contig, span, n_transcripts, exons_per_tx, exon_len): if start >= end: continue rows.append((contig, start, end, strand, f"T{t}", e + 1)) - return pl.DataFrame(rows, schema=["chrom","chromStart","chromEnd","strand", - "transcript_id","exon_number"], orient="row") + return pl.DataFrame( + rows, + schema=[ + "chrom", + "chromStart", + "chromEnd", + "strand", + "transcript_id", + "exon_number", + ], + orient="row", + ) ``` - [ ] **Step 3: Write `build_fixture`** — reference + splice BED, then the verbatim @@ -212,6 +252,7 @@ def _splice_bed(contig, span, n_transcripts, exons_per_tx, exon_len): import genvarloader as gvl from genoray import _core, SparseVar2 + def build_fixture(cohort, *, n_transcripts, exons_per_tx=3, exon_len=200, cache_dir): key = f"s{cohort.n_samples}_r_{n_transcripts}tx_{exons_per_tx}x{exon_len}" root = Path(cache_dir) / key @@ -223,9 +264,17 @@ def build_fixture(cohort, *, n_transcripts, exons_per_tx=3, exon_len=200, cache_ root.mkdir(parents=True, exist_ok=True) _write_reference(ref, cohort.contig, cohort.span + exon_len) svar2 = root / "store.svar2" - _core.run_conversion_pipeline(str(cohort.bcf), str(ref), [cohort.contig], - str(svar2), list(cohort.sample_names), - 25_000, 2, 1, 8 * 1024 * 1024) + _core.run_conversion_pipeline( + str(cohort.bcf), + str(ref), + [cohort.contig], + str(svar2), + list(cohort.sample_names), + 25_000, + 2, + 1, + 8 * 1024 * 1024, + ) gvl.write(gvl_path, bed, variants=SparseVar2(svar2), samples=None, overwrite=True) return Fixture(gvl_path, ref, bed, n_transcripts) ``` @@ -239,10 +288,12 @@ def build_fixture(cohort, *, n_transcripts, exons_per_tx=3, exon_len=200, cache_ # scratch, run manually: c = gen_cohort(50, 2000, vcfixture_bin=BIN, out_dir=TMP) f = build_fixture(c, n_transcripts=8, cache_dir=TMP) -ds = gvl.Dataset.open(f.gvl_path, reference=f.reference).with_settings( - splice_info=("transcript_id","exon_number"), var_filter="exonic" -).with_seqs("variants") -out = ds[:, :] # must return a RaggedVariants without error +ds = ( + gvl.Dataset.open(f.gvl_path, reference=f.reference) + .with_settings(splice_info=("transcript_id", "exon_number"), var_filter="exonic") + .with_seqs("variants") +) +out = ds[:, :] # must return a RaggedVariants without error print(type(out), out.alt.shape) ``` @@ -275,14 +326,16 @@ git commit -m "bench(svar2): fixture builder (reference + splice bed + svar2 wri ```python import numpy as np + + def freeze(out): frozen = {} - for name in out.fields: # alt/start/ref/ilen/dosage/... + for name in out.fields: # alt/start/ref/ilen/dosage/... f = out[name] - frozen[name] = (np.asarray(f.data).copy(), - np.asarray(f.offsets).copy()) + frozen[name] = (np.asarray(f.data).copy(), np.asarray(f.offsets).copy()) return frozen + def assert_equal(frozen, out): got = set(out.fields) assert got == set(frozen), f"field set changed: {got} vs {set(frozen)}" @@ -302,25 +355,42 @@ Task 2's smoke prints the object (read it; don't guess). import os, time, statistics as st from dataclasses import dataclass, asdict + @dataclass class BenchRow: - n_samples: int; n_transcripts: int; parallel: bool - median_ms: float; min_ms: float; spread_ms: float + n_samples: int + n_transcripts: int + parallel: bool + median_ms: float + min_ms: float + spread_ms: float + def _time_once(ds, rows, samples): - t0 = time.perf_counter(); ds[rows, samples]; return (time.perf_counter()-t0)*1e3 + t0 = time.perf_counter() + ds[rows, samples] + return (time.perf_counter() - t0) * 1e3 + def bench(fixture, *, n_query_rows, n_samples, reps=7, warmup=2, parallel=False): import genvarloader as gvl + os.environ["GVL_FORCE_PARALLEL"] = "1" if parallel else "0" - ds = gvl.Dataset.open(fixture.gvl_path, reference=fixture.reference).with_settings( - splice_info=("transcript_id","exon_number"), var_filter="exonic" - ).with_seqs("variants") - rows = slice(0, n_query_rows); samples = slice(0, n_samples) - for _ in range(warmup): ds[rows, samples] + ds = ( + gvl.Dataset.open(fixture.gvl_path, reference=fixture.reference) + .with_settings( + splice_info=("transcript_id", "exon_number"), var_filter="exonic" + ) + .with_seqs("variants") + ) + rows = slice(0, n_query_rows) + samples = slice(0, n_samples) + for _ in range(warmup): + ds[rows, samples] ts = [_time_once(ds, rows, samples) for _ in range(reps)] - return BenchRow(n_samples, n_query_rows, parallel, - st.median(ts), min(ts), max(ts)-min(ts)) + return BenchRow( + n_samples, n_query_rows, parallel, st.median(ts), min(ts), max(ts) - min(ts) + ) ``` - [ ] **Step 3: Write the sweep `main`** — sweep `n_samples ∈ {500, 3202, 25000, diff --git a/docs/superpowers/plans/2026-07-20-buffered-variant-windows.md b/docs/superpowers/plans/2026-07-20-buffered-variant-windows.md index 27d60da9..a9bf5790 100644 --- a/docs/superpowers/plans/2026-07-20-buffered-variant-windows.md +++ b/docs/superpowers/plans/2026-07-20-buffered-variant-windows.md @@ -44,6 +44,7 @@ Deliverable: `ds.to_dataloader(mode="buffered")` works for Config A and Config B ```python # tests/unit/test_output_bytes_variant_windows.py """Byte-accounting must handle variant-windows and variants+flank_tokens.""" + import numpy as np import pytest import genvarloader as gvl @@ -57,13 +58,19 @@ def _vw_ds(ref="window", alt="window"): .with_seqs( "variant-windows", gvl.VarWindowOpt( - flank_length=3, token_alphabet=b"ACGT", unknown_token=4, ref=ref, alt=alt + flank_length=3, + token_alphabet=b"ACGT", + unknown_token=4, + ref=ref, + alt=alt, ), ) ) -@pytest.mark.parametrize("ref,alt", [("window", "window"), ("window", "allele"), ("allele", "allele")]) +@pytest.mark.parametrize( + "ref,alt", [("window", "window"), ("window", "allele"), ("allele", "allele")] +) def test_variant_windows_bytes_positive(ref, alt): ds = _vw_ds(ref, alt) bpi = ds._output_bytes_per_instance(None, None) @@ -76,9 +83,14 @@ def test_variant_windows_bytes_positive(ref, alt): def test_flank_tokens_adds_bytes(): base = ( - gvl.get_dummy_dataset().with_seqs("variants").with_tracks(False).with_output_format("flat") + gvl.get_dummy_dataset() + .with_seqs("variants") + .with_tracks(False) + .with_output_format("flat") + ) + with_flank = base.with_settings( + flank_length=3, token_alphabet=b"ACGT", unknown_token=0 ) - with_flank = base.with_settings(flank_length=3, token_alphabet=b"ACGT", unknown_token=0) b0 = base._output_bytes_per_instance(None, None).sum() b1 = with_flank._output_bytes_per_instance(None, None).sum() assert b1 > b0 # flank tokens are extra payload @@ -180,6 +192,7 @@ rtk git commit -m "feat(dataloader): byte accounting for variant-windows and var ```python # tests/unit/test_flat_window_slicing.py """Instance-axis slicing of flat window / flank types matches per-item output.""" + import numpy as np import genvarloader as gvl @@ -199,14 +212,19 @@ def test_flat_variant_windows_slice_matches_per_item(): .with_output_format("flat") .with_seqs( "variant-windows", - gvl.VarWindowOpt(flank_length=2, token_alphabet=b"ACGT", unknown_token=4, - ref="window", alt="allele"), + gvl.VarWindowOpt( + flank_length=2, + token_alphabet=b"ACGT", + unknown_token=4, + ref="window", + alt="allele", + ), ) ) r = np.array([0, 0, 1], np.intp) s = np.array([0, 1, 0], np.intp) - batch = ds[r, s] # one _FlatVariantWindows over 3 instances - sliced = batch[1:3] # instances 1,2 + batch = ds[r, s] # one _FlatVariantWindows over 3 instances + sliced = batch[1:3] # instances 1,2 expected = ds[r[1:3], s[1:3]] _win_eq(sliced, expected) @@ -223,7 +241,7 @@ def test_flat_variants_flank_tokens_slice_carries_tokens(): s = np.array([0, 1, 0], np.intp) batch = ds[r, s] assert batch.flank_tokens is not None - sliced = batch[1:3] # must NOT raise, must keep flank_tokens + sliced = batch[1:3] # must NOT raise, must keep flank_tokens assert sliced.flank_tokens is not None exp = ds[r[1:3], s[1:3]] np.testing.assert_array_equal( @@ -323,7 +341,9 @@ def _iter_mode_none(ds, batch_size): yield ds[r[i : i + batch_size], s[i : i + batch_size]] -@pytest.mark.parametrize("ref,alt", [("window", "window"), ("window", "allele"), ("allele", "allele")]) +@pytest.mark.parametrize( + "ref,alt", [("window", "window"), ("window", "allele"), ("allele", "allele")] +) def test_buffered_variant_windows_matches_per_item(ref, alt): ds = ( gvl.get_dummy_dataset() @@ -331,12 +351,25 @@ def test_buffered_variant_windows_matches_per_item(ref, alt): .with_output_format("flat") .with_seqs( "variant-windows", - gvl.VarWindowOpt(flank_length=2, token_alphabet=b"ACGT", unknown_token=4, ref=ref, alt=alt), + gvl.VarWindowOpt( + flank_length=2, + token_alphabet=b"ACGT", + unknown_token=4, + ref=ref, + alt=alt, + ), ) ) bs = 2 - got = list(ds.to_dataloader(mode="buffered", batch_size=bs, shuffle=False, - drop_last=True, buffer_bytes=10 * 1024 * 1024)) + got = list( + ds.to_dataloader( + mode="buffered", + batch_size=bs, + shuffle=False, + drop_last=True, + buffer_bytes=10 * 1024 * 1024, + ) + ) exp = list(_iter_mode_none(ds, bs)) assert len(got) == len(exp) for g, e in zip(got, exp): @@ -345,13 +378,22 @@ def test_buffered_variant_windows_matches_per_item(ref, alt): def test_buffered_variants_flank_tokens_matches_per_item(): ds = ( - gvl.get_dummy_dataset().with_seqs("variants").with_tracks(False) + gvl.get_dummy_dataset() + .with_seqs("variants") + .with_tracks(False) .with_settings(flank_length=2, token_alphabet=b"ACGT", unknown_token=0) .with_output_format("flat") ) bs = 2 - got = list(ds.to_dataloader(mode="buffered", batch_size=bs, shuffle=False, - drop_last=True, buffer_bytes=10 * 1024 * 1024)) + got = list( + ds.to_dataloader( + mode="buffered", + batch_size=bs, + shuffle=False, + drop_last=True, + buffer_bytes=10 * 1024 * 1024, + ) + ) exp = list(_iter_mode_none(ds, bs)) assert len(got) == len(exp) for g, e in zip(got, exp): @@ -378,38 +420,45 @@ Expected: FAIL — `ValueError: mode='buffered' does not support 'variant-window - [ ] **Step 3b: Narrow both guards** in `_torch.py:get_dataloader` from unconditional to `mode == "double_buffered"` only. Replace the two `if` blocks (`:164-192`): ```python - # 'variant-windows' and flat variants+flank_tokens cannot yet ride the - # double_buffered transport (the producer schema / shm format do not carry - # the VarWindowOpt or the flank tokens). buffered runs in-process and does. - if mode == "double_buffered" and getattr(dataset, "sequence_type", None) == "variant-windows": - raise ValueError( - "mode='double_buffered' does not support 'variant-windows' output yet: the " - "producer schema/shared-memory format cannot carry the VarWindowOpt. Use " - "mode='buffered' (in-process) or mode=None." - ) +# 'variant-windows' and flat variants+flank_tokens cannot yet ride the +# double_buffered transport (the producer schema / shm format do not carry +# the VarWindowOpt or the flank tokens). buffered runs in-process and does. +if ( + mode == "double_buffered" + and getattr(dataset, "sequence_type", None) == "variant-windows" +): + raise ValueError( + "mode='double_buffered' does not support 'variant-windows' output yet: the " + "producer schema/shared-memory format cannot carry the VarWindowOpt. Use " + "mode='buffered' (in-process) or mode=None." + ) +if ( + mode == "double_buffered" + and getattr(dataset, "output_format", "ragged") == "flat" + and getattr(dataset, "sequence_type", None) == "variants" +): + _seqs = getattr(dataset, "_seqs", None) if ( - mode == "double_buffered" - and getattr(dataset, "output_format", "ragged") == "flat" - and getattr(dataset, "sequence_type", None) == "variants" + getattr(_seqs, "flank_length", None) + and getattr(_seqs, "token_lut", None) is not None ): - _seqs = getattr(dataset, "_seqs", None) - if getattr(_seqs, "flank_length", None) and getattr(_seqs, "token_lut", None) is not None: - raise ValueError( - "mode='double_buffered' with output_format='flat' does not support variants " - "output carrying ride-along flank tokens yet; use mode='buffered' or mode=None." - ) + raise ValueError( + "mode='double_buffered' with output_format='flat' does not support variants " + "output carrying ride-along flank tokens yet; use mode='buffered' or mode=None." + ) ``` - [ ] **Step 3c: Flip the two buffered rejection tests.** In `tests/unit/test_buffered_loader.py`, change `test_flat_buffered_rejects_variant_windows` and `test_flat_buffered_rejects_variants_flank_tokens` so the `mode="buffered"` parametrization asserts success, and only `mode="double_buffered"` still asserts the `ValueError`. Simplest: change the parametrize to `["double_buffered"]` (buffered no longer rejects) and add a note that PR 2 removes the double_buffered arm. Keep `test_flat_buffered_plain_variants_still_works` unchanged. ```python -@pytest.mark.parametrize("mode", ["double_buffered"]) # buffered now supports both (PR1); PR2 drops these -def test_flat_buffered_rejects_variants_flank_tokens(mode): - ... # body unchanged +@pytest.mark.parametrize( + "mode", ["double_buffered"] +) # buffered now supports both (PR1); PR2 drops these +def test_flat_buffered_rejects_variants_flank_tokens(mode): ... # body unchanged + @pytest.mark.parametrize("mode", ["double_buffered"]) -def test_flat_buffered_rejects_variant_windows(mode): - ... # body unchanged +def test_flat_buffered_rejects_variant_windows(mode): ... # body unchanged ``` - [ ] **Step 4: Run the buffered suite** @@ -476,6 +525,7 @@ Deliverable: `ds.to_dataloader(mode="double_buffered")` works for Config A and C ```python # tests/unit/test_producer_schema.py """_apply_schema reconstructs variant-windows and flank configs in the child.""" + import genvarloader as gvl from genvarloader._producer import _apply_schema @@ -485,8 +535,13 @@ def test_apply_schema_rebuilds_variant_windows(): schema = { "with_seqs": "variant-windows", "output_format": "flat", - "window_opt": {"flank_length": 3, "token_alphabet": b"ACGT", - "unknown_token": 4, "ref": "window", "alt": "allele"}, + "window_opt": { + "flank_length": 3, + "token_alphabet": b"ACGT", + "unknown_token": 4, + "ref": "window", + "alt": "allele", + }, } ds = _apply_schema(dummy, schema) assert ds.sequence_type == "variant-windows" @@ -499,7 +554,9 @@ def test_apply_schema_rebuilds_flank_tokens(): schema = { "with_seqs": "variants", "output_format": "flat", - "flank_length": 2, "token_alphabet": b"ACGT", "unknown_token": 0, + "flank_length": 2, + "token_alphabet": b"ACGT", + "unknown_token": 0, } ds = _apply_schema(dummy, schema) assert ds._seqs.flank_length == 2 and ds._seqs.token_lut is not None @@ -513,20 +570,22 @@ Expected: FAIL — `with_seqs('variant-windows') requires a VarWindowOpt` (schem - [ ] **Step 3a: Emit the config in `_spawn_producer`.** After the existing `if isinstance(seqs, Haps):` block (`_double_buffered_loader.py:208-217`), add: ```python - window_opt = getattr(seqs, "window_opt", None) - if window_opt is not None: - schema["window_opt"] = { - "flank_length": window_opt.flank_length, - "token_alphabet": window_opt.token_alphabet, - "unknown_token": window_opt.unknown_token, - "ref": window_opt.ref, - "alt": window_opt.alt, - } - elif getattr(seqs, "flank_length", None) and getattr(seqs, "token_lut", None) is not None: - # plain-variants ride-along flank tokens (Config B) - schema["flank_length"] = seqs.flank_length - schema["token_alphabet"] = seqs.token_alphabet - schema["unknown_token"] = seqs.unknown_token +window_opt = getattr(seqs, "window_opt", None) +if window_opt is not None: + schema["window_opt"] = { + "flank_length": window_opt.flank_length, + "token_alphabet": window_opt.token_alphabet, + "unknown_token": window_opt.unknown_token, + "ref": window_opt.ref, + "alt": window_opt.alt, + } +elif ( + getattr(seqs, "flank_length", None) and getattr(seqs, "token_lut", None) is not None +): + # plain-variants ride-along flank tokens (Config B) + schema["flank_length"] = seqs.flank_length + schema["token_alphabet"] = seqs.token_alphabet + schema["unknown_token"] = seqs.unknown_token ``` (Verify the attribute names `token_alphabet`/`unknown_token` exist on the `Haps` for the flank path; if the flank config lives under different attribute names, read them from `with_settings`'s stored fields. Test `test_apply_schema_rebuilds_flank_tokens` pins the round-trip.) @@ -534,18 +593,19 @@ Expected: FAIL — `with_seqs('variant-windows') requires a VarWindowOpt` (schem - [ ] **Step 3b: Reconstruct in `_apply_schema`.** Replace the `with_seqs` replay line and extend the settings block: ```python - if schema.get("with_seqs", "UNSET") != "UNSET": - if schema.get("window_opt") is not None: - from ._dataset._flat_variants import VarWindowOpt - ds = ds.with_seqs(schema["with_seqs"], VarWindowOpt(**schema["window_opt"])) - else: - ds = ds.with_seqs(schema["with_seqs"]) - ... - # inside the settings_kwargs block, add: - if schema.get("flank_length") is not None: - settings_kwargs["flank_length"] = schema["flank_length"] - settings_kwargs["token_alphabet"] = schema["token_alphabet"] - settings_kwargs["unknown_token"] = schema["unknown_token"] +if schema.get("with_seqs", "UNSET") != "UNSET": + if schema.get("window_opt") is not None: + from ._dataset._flat_variants import VarWindowOpt + + ds = ds.with_seqs(schema["with_seqs"], VarWindowOpt(**schema["window_opt"])) + else: + ds = ds.with_seqs(schema["with_seqs"]) +... +# inside the settings_kwargs block, add: +if schema.get("flank_length") is not None: + settings_kwargs["flank_length"] = schema["flank_length"] + settings_kwargs["token_alphabet"] = schema["token_alphabet"] + settings_kwargs["unknown_token"] = schema["unknown_token"] ``` - [ ] **Step 4: Run test to verify it passes** @@ -580,6 +640,7 @@ rtk git commit -m "feat(dataloader): replay VarWindowOpt and flank config in the ```python # tests/unit/test_shm_variant_windows.py """kind=4 round-trip for _FlatVariantWindows over the shm layout.""" + import numpy as np from genvarloader._shm_layout import write_chunk, read_chunk, HEADER_RESERVED from genvarloader._flat import _Flat @@ -588,12 +649,14 @@ from genvarloader._dataset._flat_variants import _FlatWindow, _FlatVariantWindow def _make_fvw(): # 2 instances, ploidy 1. start scalar field + ref_window + alt (bare) slots. - start = _Flat(np.array([10, 20, 30], np.int32), np.array([0, 2, 3], np.int64), (2, 1, None)) + start = _Flat( + np.array([10, 20, 30], np.int32), np.array([0, 2, 3], np.int64), (2, 1, None) + ) # ref_window: b*p=2 rows; var_offsets len 3; per-variant token runs via seq_offsets. rw = _FlatWindow( data=np.array([1, 2, 3, 4, 1, 2], np.uint8), - seq_offsets=np.array([0, 3, 4, 6], np.int64), # 3 variants - var_offsets=np.array([0, 2, 3], np.int64), # 2 rows -> 2,1 variants + seq_offsets=np.array([0, 3, 4, 6], np.int64), # 3 variants + var_offsets=np.array([0, 2, 3], np.int64), # 2 rows -> 2,1 variants shape=(2, 1, None, None), ) al = _FlatWindow( @@ -647,51 +710,79 @@ def _write_flat_variant_windows(buf: memoryview, fvw, cursor: int) -> tuple[dict outer = np.ascontiguousarray(var_off, np.int64) inner = np.ascontiguousarray(seq_off, np.int64) leaf = np.ascontiguousarray(data) - cursor = _align(cursor); outer_off = cursor + cursor = _align(cursor) + outer_off = cursor np.frombuffer(buf, np.int64, outer.size, outer_off)[...] = outer cursor += outer.nbytes - cursor = _align(cursor); inner_off = cursor + cursor = _align(cursor) + inner_off = cursor np.frombuffer(buf, np.int64, inner.size, inner_off)[...] = inner cursor += inner.nbytes - cursor = _align(cursor); data_off = cursor + cursor = _align(cursor) + data_off = cursor np.frombuffer(buf, leaf.dtype, leaf.size, data_off)[...] = leaf.ravel() cursor += leaf.nbytes - field_descs.append({ - "field_kind": 1, "dtype_str": _dtype_to_bytes(leaf.dtype), - "outer_offsets_offset": outer_off, "outer_offsets_nbytes": outer.nbytes, - "inner_offsets_offset": inner_off, "inner_offsets_nbytes": inner.nbytes, - "data_offset": data_off, "data_nbytes": leaf.nbytes, - "regular_size": regular_size, "name": name.encode("utf-8"), - }) + field_descs.append( + { + "field_kind": 1, + "dtype_str": _dtype_to_bytes(leaf.dtype), + "outer_offsets_offset": outer_off, + "outer_offsets_nbytes": outer.nbytes, + "inner_offsets_offset": inner_off, + "inner_offsets_nbytes": inner.nbytes, + "data_offset": data_off, + "data_nbytes": leaf.nbytes, + "regular_size": regular_size, + "name": name.encode("utf-8"), + } + ) # scalar .fields first (numeric, field_kind=0) — mirror _write_flat_variants for name, f in fvw.fields.items(): outer = np.ascontiguousarray(f.offsets, np.int64) leaf = np.ascontiguousarray(f.data) - cursor = _align(cursor); outer_off = cursor + cursor = _align(cursor) + outer_off = cursor np.frombuffer(buf, np.int64, outer.size, outer_off)[...] = outer cursor += outer.nbytes - cursor = _align(cursor); data_off = cursor + cursor = _align(cursor) + data_off = cursor np.frombuffer(buf, leaf.dtype, leaf.size, data_off)[...] = leaf.ravel() cursor += leaf.nbytes - field_descs.append({ - "field_kind": 0, "dtype_str": _dtype_to_bytes(leaf.dtype), - "outer_offsets_offset": outer_off, "outer_offsets_nbytes": outer.nbytes, - "inner_offsets_offset": 0, "inner_offsets_nbytes": 0, - "data_offset": data_off, "data_nbytes": leaf.nbytes, - "regular_size": _flat_ploidy(f.shape), "name": name.encode("utf-8"), - }) + field_descs.append( + { + "field_kind": 0, + "dtype_str": _dtype_to_bytes(leaf.dtype), + "outer_offsets_offset": outer_off, + "outer_offsets_nbytes": outer.nbytes, + "inner_offsets_offset": 0, + "inner_offsets_nbytes": 0, + "data_offset": data_off, + "data_nbytes": leaf.nbytes, + "regular_size": _flat_ploidy(f.shape), + "name": name.encode("utf-8"), + } + ) # present window slots (two-level, field_kind=1) for slot in _WINDOW_FIELD_NAMES: w = getattr(fvw, slot) if w is not None: - _emit_two_level(slot, w.data, w.seq_offsets, w.var_offsets, _flat_ploidy(w.shape)) + _emit_two_level( + slot, w.data, w.seq_offsets, w.var_offsets, _flat_ploidy(w.shape) + ) return { - "kind": 4, "dtype_str": b"\x00" * 4, "shape": [len(field_descs)], - "data_offset": 0, "data_nbytes": 0, "offsets_offset": 0, "offsets_nbytes": 0, - "inner_offsets_offset": 0, "inner_offsets_nbytes": 0, "name": b"", + "kind": 4, + "dtype_str": b"\x00" * 4, + "shape": [len(field_descs)], + "data_offset": 0, + "data_nbytes": 0, + "offsets_offset": 0, + "offsets_nbytes": 0, + "inner_offsets_offset": 0, + "inner_offsets_nbytes": 0, + "name": b"", "_field_descs": field_descs, }, cursor ``` @@ -703,7 +794,11 @@ def _write_flat_variant_windows(buf: memoryview, fvw, cursor: int) -> tuple[dict ```python def _read_flat_variant_windows(buf: memoryview, d: dict, copy: bool = True): from ._flat import _Flat - from ._dataset._flat_variants import _FlatWindow, _FlatVariantWindows, _WINDOW_FIELD_NAMES + from ._dataset._flat_variants import ( + _FlatWindow, + _FlatVariantWindows, + _WINDOW_FIELD_NAMES, + ) fields: dict = {} windows: dict = {} @@ -713,7 +808,9 @@ def _read_flat_variant_windows(buf: memoryview, d: dict, copy: bool = True): rs = fd["regular_size"] n_outer = fd["outer_offsets_nbytes"] // 8 var_off = np.frombuffer(buf, np.int64, n_outer, fd["outer_offsets_offset"]) - leaf = np.frombuffer(buf, leaf_dtype, fd["data_nbytes"] // leaf_dtype.itemsize, fd["data_offset"]) + leaf = np.frombuffer( + buf, leaf_dtype, fd["data_nbytes"] // leaf_dtype.itemsize, fd["data_offset"] + ) if copy: var_off, leaf = var_off.copy(), leaf.copy() n_bp = len(var_off) - 1 @@ -767,6 +864,7 @@ rtk git commit -m "feat(dataloader): shm kind=4 serialization for _FlatVariantWi ```python # tests/unit/test_shm_flank_tokens.py """kind=2 flank_tokens round-trip over the shm layout.""" + import numpy as np from genvarloader._shm_layout import write_chunk, read_chunk, HEADER_RESERVED from genvarloader._flat import _Flat @@ -775,10 +873,12 @@ from genvarloader._dataset._flat_variants import _FlatVariants def test_kind2_flank_tokens_roundtrip(): # 2 instances, ploidy 1, 2L=4. start scalar + flank_tokens (b,p,None,2L). - start = _Flat(np.array([1, 2, 3], np.int32), np.array([0, 2, 3], np.int64), (2, 1, None)) + start = _Flat( + np.array([1, 2, 3], np.int32), np.array([0, 2, 3], np.int64), (2, 1, None) + ) ft = _Flat( - np.arange(3 * 4, dtype=np.uint8), # 3 variants * 2L(=4) tokens - np.array([0, 4, 8, 12], np.int64), # n_rows = b*p*2L? verify vs builder + np.arange(3 * 4, dtype=np.uint8), # 3 variants * 2L(=4) tokens + np.array([0, 4, 8, 12], np.int64), # n_rows = b*p*2L? verify vs builder (2, 1, None, 4), ) fv = _FlatVariants({"start": start}) @@ -789,8 +889,12 @@ def test_kind2_flank_tokens_roundtrip(): out = views[0] assert out.flank_tokens is not None assert out.flank_tokens.shape == (2, 1, None, 4) - np.testing.assert_array_equal(np.asarray(out.flank_tokens.data), np.asarray(ft.data)) - np.testing.assert_array_equal(np.asarray(out.flank_tokens.offsets), np.asarray(ft.offsets)) + np.testing.assert_array_equal( + np.asarray(out.flank_tokens.data), np.asarray(ft.data) + ) + np.testing.assert_array_equal( + np.asarray(out.flank_tokens.offsets), np.asarray(ft.offsets) + ) ``` (If the real `flank_tokens.offsets` length differs from this synthetic one, adjust the fixture to match a `_FlatVariants` produced by `ds[r,s]` on a dummy variants+flank dataset — the round-trip must equal whatever the constructor produces; the shape/offset preservation is what matters.) @@ -803,24 +907,29 @@ Expected: FAIL — `out.flank_tokens is None` (dropped by the current kind=2 wri - [ ] **Step 3a: Writers.** In `_write_flat_variants`, after the field loop and before the `return`, serialize flank_tokens if present and attach a `_flank` descriptor; else `_flank=None`: ```python - flank = None - if fv.flank_tokens is not None: - ft = fv.flank_tokens - data = np.ascontiguousarray(ft.data) - off = np.ascontiguousarray(ft.offsets, np.int64) - cursor = _align(cursor); data_off = cursor - np.frombuffer(buf, data.dtype, data.size, data_off)[...] = data.ravel() - cursor += data.nbytes - cursor = _align(cursor); off_off = cursor - np.frombuffer(buf, np.int64, off.size, off_off)[...] = off - cursor += off.nbytes - flank = { - "shape": list(ft.shape), "dtype_str": _dtype_to_bytes(data.dtype), - "data_offset": data_off, "data_nbytes": data.nbytes, - "offsets_offset": off_off, "offsets_nbytes": off.nbytes, - } - # add to the returned dict: - # "_flank": flank, +flank = None +if fv.flank_tokens is not None: + ft = fv.flank_tokens + data = np.ascontiguousarray(ft.data) + off = np.ascontiguousarray(ft.offsets, np.int64) + cursor = _align(cursor) + data_off = cursor + np.frombuffer(buf, data.dtype, data.size, data_off)[...] = data.ravel() + cursor += data.nbytes + cursor = _align(cursor) + off_off = cursor + np.frombuffer(buf, np.int64, off.size, off_off)[...] = off + cursor += off.nbytes + flank = { + "shape": list(ft.shape), + "dtype_str": _dtype_to_bytes(data.dtype), + "data_offset": data_off, + "data_nbytes": data.nbytes, + "offsets_offset": off_off, + "offsets_nbytes": off.nbytes, + } +# add to the returned dict: +# "_flank": flank, ``` In `_write_rag_variants`'s returned dict add `"_flank": None`. @@ -828,14 +937,20 @@ In `_write_rag_variants`'s returned dict add `"_flank": None`. - [ ] **Step 3b: pack/unpack.** In `_pack_descriptor`'s `if kind == 2:` block, after the field loop, append: ```python - flank = d.get("_flank") - out += struct.pack("= 4096 # floor ``` @@ -135,6 +146,7 @@ every serialized array. Those are per-chunk constants (independent of instance c that must be covered by the slot's fixed slack. This module derives a true upper bound on them from the schema, replacing the historical magic 4096. """ + from __future__ import annotations _OFF = 8 # int64 offset entry / terminator @@ -155,7 +167,9 @@ def _array_counts(dataset) -> tuple[int, int]: if seq == "variant-windows": n_scalar = len(scalars) n_window_slots = 2 # exactly one ref-derived + one alt-derived slot - n_off += n_scalar * 1 + n_window_slots * 2 # scalars: outer; windows: outer+inner + n_off += ( + n_scalar * 1 + n_window_slots * 2 + ) # scalars: outer; windows: outer+inner n_arr += n_scalar * 2 + n_window_slots * 3 # +1 data array each elif seq == "variants": n_scalar = len(scalars) @@ -367,6 +381,7 @@ Co-Authored-By: Claude Opus 4.8 " schema-derived slot overhead must upper-bound the real serialized payload for every chunk. This is the invariant #315 violated; it must hold across record types and storage backends.""" + import numpy as np import pytest import seqpro as sp @@ -382,11 +397,19 @@ def _views(ds): continue for uu in (True, False): for L in (8, 128): - opt = gvl.VarWindowOpt(flank_length=L, token_alphabet=DNA, - unknown_token=len(DNA), ref=ref, alt=alt) - yield (ds.with_tracks(False).with_output_format("flat") - .with_seqs("variant-windows", opt) - .with_settings(unphased_union=uu, jitter=0)) + opt = gvl.VarWindowOpt( + flank_length=L, + token_alphabet=DNA, + unknown_token=len(DNA), + ref=ref, + alt=alt, + ) + yield ( + ds.with_tracks(False) + .with_output_format("flat") + .with_seqs("variant-windows", opt) + .with_settings(unphased_union=uu, jitter=0) + ) def _assert_upper_bound(view): @@ -397,9 +420,12 @@ def _assert_upper_bound(view): arrays = list(chunk) if isinstance(chunk, tuple) else [chunk] buf = memoryview(bytearray(64 * 1024 * 1024)) real = write_chunk(buf, arrays, n_instances=len(r)) - HEADER_RESERVED - est = int(np.asarray(view._output_bytes_per_instance(r, s, include_offsets=True)).sum()) + est = int( + np.asarray(view._output_bytes_per_instance(r, s, include_offsets=True)).sum() + ) assert est + slot_overhead_bytes(view) >= real, ( - f"slot under-sized: est={est} overhead={slot_overhead_bytes(view)} real={real}") + f"slot under-sized: est={est} overhead={slot_overhead_bytes(view)} real={real}" + ) def test_slot_fit_dummy_backend(): diff --git a/docs/superpowers/plans/2026-07-23-svar2-variant-windows-slot-fit.md b/docs/superpowers/plans/2026-07-23-svar2-variant-windows-slot-fit.md index f2b87d50..e6e1f85d 100644 --- a/docs/superpowers/plans/2026-07-23-svar2-variant-windows-slot-fit.md +++ b/docs/superpowers/plans/2026-07-23-svar2-variant-windows-slot-fit.md @@ -71,6 +71,7 @@ Create `scratch/diag_315_realcorpus.py`. It reproduces the reported config, conf """Pin the #315 estimate divergence against the real SVAR2 Hartwig corpus. Not committed under python/. Run: pixi run -e dev python scratch/diag_315_realcorpus.py """ + import sys import numpy as np import seqpro as sp @@ -90,20 +91,28 @@ def make_view(): # subset to the reported region count (all samples) ds = ds.subset_to(regions=slice(N_REGIONS)) opt = gvl.VarWindowOpt( - flank_length=128, token_alphabet=DNA, unknown_token=len(DNA), - ref="window", alt="allele", + flank_length=128, + token_alphabet=DNA, + unknown_token=len(DNA), + ref="window", + alt="allele", ) return ( - ds.with_tracks(False).with_output_format("flat") - .with_seqs("variant-windows", opt) - .with_settings(unphased_union=True, jitter=0) + ds.with_tracks(False) + .with_output_format("flat") + .with_seqs("variant-windows", opt) + .with_settings(unphased_union=True, jitter=0) ) def main() -> int: view = make_view() - print("reconstructor:", type(view._seqs).__name__, - "is Svar2Haps:", isinstance(view._seqs, Svar2Haps)) + print( + "reconstructor:", + type(view._seqs).__name__, + "is Svar2Haps:", + isinstance(view._seqs, Svar2Haps), + ) R, S = view.shape[:2] print(f"shape: {R} regions x {S} samples") @@ -118,11 +127,16 @@ def main() -> int: arrays = list(chunk) if isinstance(chunk, tuple) else [chunk] buf = memoryview(bytearray(512 * 1024 * 1024)) real = write_chunk(buf, arrays, n_instances=len(r)) - HEADER_RESERVED - est = int(np.asarray( - view._output_bytes_per_instance(r, s, include_offsets=True)).sum()) + est = int( + np.asarray( + view._output_bytes_per_instance(r, s, include_offsets=True) + ).sum() + ) ovh = slot_overhead_bytes(view) - print(f"N={len(r):>7} est={est:>12} overhead={ovh:>8} real={real:>12} " - f"est+ovh-real={est + ovh - real:>12} per_inst_gap={(real - est) / len(r):.1f}") + print( + f"N={len(r):>7} est={est:>12} overhead={ovh:>8} real={real:>12} " + f"est+ovh-real={est + ovh - real:>12} per_inst_gap={(real - est) / len(r):.1f}" + ) # Decompose one region's estimate: n_vars_total vs emitted window count W. r = np.zeros(S, np.int64) @@ -131,16 +145,21 @@ def main() -> int: n_vars = view.n_variants(r, s) n_vars_total = n_vars.reshape(-1, n_vars.shape[-1]).astype(np.int64).sum(-1) chunk = view[r, s] - ref_slot = (chunk[0] if isinstance(chunk, tuple) else chunk) + ref_slot = chunk[0] if isinstance(chunk, tuple) else chunk # emitted window count W per instance = len(ref window seq_offsets) - 1, per instance print("sum n_vars_total (estimate M):", int(n_vars_total.sum())) - print("real_ploidy:", haps.genotypes.shape[-2], - "unphased_union:", view.unphased_union) + print( + "real_ploidy:", haps.genotypes.shape[-2], "unphased_union:", view.unphased_union + ) # Dump the worst-under-counted instance for its record class. est_pi = np.asarray(view._output_bytes_per_instance(r, s, include_offsets=True)) - worst = int(np.argmin(est_pi - 0)) # smallest estimate; refine vs per-instance real if needed - print("example instance (r=0, s=%d): est_bytes=%d n_vars_total=%d" - % (worst, int(est_pi[worst]), int(n_vars_total[worst]))) + worst = int( + np.argmin(est_pi - 0) + ) # smallest estimate; refine vs per-instance real if needed + print( + "example instance (r=0, s=%d): est_bytes=%d n_vars_total=%d" + % (worst, int(est_pi[worst]), int(n_vars_total[worst])) + ) return 0 @@ -211,6 +230,7 @@ def test_slot_fit_svar2_backend(phased_svar2_gvl, reference): coverage gap that let #315 through. The estimate must upper-bound the real serialized payload here too.""" from genvarloader._dataset._svar2_haps import Svar2Haps + ds = gvl.Dataset.open(phased_svar2_gvl, reference=reference) assert isinstance(ds._seqs, Svar2Haps), "fixture must open as Svar2Haps" for view in _views(ds): diff --git a/docs/superpowers/plans/2026-07-30-svar2-write-memory.md b/docs/superpowers/plans/2026-07-30-svar2-write-memory.md index 203fde6a..d6a67d26 100644 --- a/docs/superpowers/plans/2026-07-30-svar2-write-memory.md +++ b/docs/superpowers/plans/2026-07-30-svar2-write-memory.md @@ -1212,95 +1212,93 @@ class RangesStream: Then add the method to `_BatchQueryMixin`: ```python - def _find_ranges_chunked( - self, - contig: str, - starts: "ArrayLike", - ends: "ArrayLike", - samples: "ArrayLike | None" = None, - *, - max_mem: int | None = None, - ) -> RangesStream: - """Chunked, memory-bounded ``_find_ranges``. - - ``starts``/``ends`` and ``samples`` behave as in :meth:`read_ranges`. - - The var_key payload is ``n_regions * n_samples * ploidy * 2`` int64 - pairs per channel, which is tens of GiB at cohort scale. This splits it - along the SAMPLE axis -- not the region axis -- because the search is - column-outer: chunking regions would re-sweep the whole packed store per - chunk, while chunking samples keeps a single sweep. - - Args: - contig: Contig name. - starts: 0-based start positions of the query regions. - ends: 0-based, exclusive end positions of the query regions. - samples: Sample names selecting (and reordering) a subset. - max_mem: Approximate byte budget for one chunk's payload. ``None`` - yields a single chunk covering every sample. - - Returns: - A :class:`RangesStream` whose ``chunks`` generator yields - :class:`RangesChunk` in ascending ``sample_start`` order. - - Raises: - ValueError: If ``max_mem`` cannot fit a single sample's payload, or - if the contig's largest deletion overflows the max-end key - packing width. - """ - reg = self._regions(starts, ends) - sample_idxs = self._sample_idxs(samples) - reader = self._reader(contig) - header = reader.find_ranges_header(reg, sample_idxs) - - n_regions = int(header["n_regions"]) - n_samples = int(header["n_samples"]) - ploidy = int(header["ploidy"]) - - # Both channels, 2 endpoints, int64. The 2x is slop for the transient - # the binding holds while handing the arrays back. - bytes_per_sample = n_regions * ploidy * 2 * 8 * 2 - if max_mem is None: - per = max(n_samples, 1) - else: - per = int(max_mem) // (2 * bytes_per_sample) if bytes_per_sample else n_samples - if per < 1: - raise ValueError( - f"max_mem ({int(max_mem)} bytes) is too small for even one " - f"sample of {n_regions} regions at ploidy {ploidy}: needs at " - f"least {2 * bytes_per_sample} bytes." - ) - per = min(per, max(n_samples, 1)) +def _find_ranges_chunked( + self, + contig: str, + starts: "ArrayLike", + ends: "ArrayLike", + samples: "ArrayLike | None" = None, + *, + max_mem: int | None = None, +) -> RangesStream: + """Chunked, memory-bounded ``_find_ranges``. + + ``starts``/``ends`` and ``samples`` behave as in :meth:`read_ranges`. + + The var_key payload is ``n_regions * n_samples * ploidy * 2`` int64 + pairs per channel, which is tens of GiB at cohort scale. This splits it + along the SAMPLE axis -- not the region axis -- because the search is + column-outer: chunking regions would re-sweep the whole packed store per + chunk, while chunking samples keeps a single sweep. - def _gen() -> "Iterator[RangesChunk]": - for s0 in range(0, n_samples, per): - s1 = min(s0 + per, n_samples) - d = reader.find_ranges_chunk( - reg, sample_idxs, s0 * ploidy, s1 * ploidy - ) - cs = s1 - s0 - shape = (cs, ploidy, n_regions, 2) - yield RangesChunk( - sample_start=s0, - n_samples=cs, - vk_snp_range=np.asarray(d["vk_snp_range"]).reshape(shape), - vk_indel_range=np.asarray(d["vk_indel_range"]).reshape(shape), - max_end_keys=np.asarray(d["max_end_keys"], np.int64), - ) + Args: + contig: Contig name. + starts: 0-based start positions of the query regions. + ends: 0-based, exclusive end positions of the query regions. + samples: Sample names selecting (and reordering) a subset. + max_mem: Approximate byte budget for one chunk's payload. ``None`` + yields a single chunk covering every sample. - return RangesStream( - n_regions=n_regions, - n_samples=n_samples, - ploidy=ploidy, - samples_per_chunk=per, - region_starts=np.asarray(header["region_starts"]), - dense_range=np.asarray(header["dense_range"]), - dense_snp_range=np.asarray(header["dense_snp_range"]), - dense_indel_range=np.asarray(header["dense_indel_range"]), - sample_cols=np.asarray(header["sample_cols"]), - dense_max_end_keys=np.asarray(header["dense_max_end_keys"], np.int64), - chunks=_gen(), - ) + Returns: + A :class:`RangesStream` whose ``chunks`` generator yields + :class:`RangesChunk` in ascending ``sample_start`` order. + + Raises: + ValueError: If ``max_mem`` cannot fit a single sample's payload, or + if the contig's largest deletion overflows the max-end key + packing width. + """ + reg = self._regions(starts, ends) + sample_idxs = self._sample_idxs(samples) + reader = self._reader(contig) + header = reader.find_ranges_header(reg, sample_idxs) + + n_regions = int(header["n_regions"]) + n_samples = int(header["n_samples"]) + ploidy = int(header["ploidy"]) + + # Both channels, 2 endpoints, int64. The 2x is slop for the transient + # the binding holds while handing the arrays back. + bytes_per_sample = n_regions * ploidy * 2 * 8 * 2 + if max_mem is None: + per = max(n_samples, 1) + else: + per = int(max_mem) // (2 * bytes_per_sample) if bytes_per_sample else n_samples + if per < 1: + raise ValueError( + f"max_mem ({int(max_mem)} bytes) is too small for even one " + f"sample of {n_regions} regions at ploidy {ploidy}: needs at " + f"least {2 * bytes_per_sample} bytes." + ) + per = min(per, max(n_samples, 1)) + + def _gen() -> "Iterator[RangesChunk]": + for s0 in range(0, n_samples, per): + s1 = min(s0 + per, n_samples) + d = reader.find_ranges_chunk(reg, sample_idxs, s0 * ploidy, s1 * ploidy) + cs = s1 - s0 + shape = (cs, ploidy, n_regions, 2) + yield RangesChunk( + sample_start=s0, + n_samples=cs, + vk_snp_range=np.asarray(d["vk_snp_range"]).reshape(shape), + vk_indel_range=np.asarray(d["vk_indel_range"]).reshape(shape), + max_end_keys=np.asarray(d["max_end_keys"], np.int64), + ) + + return RangesStream( + n_regions=n_regions, + n_samples=n_samples, + ploidy=ploidy, + samples_per_chunk=per, + region_starts=np.asarray(header["region_starts"]), + dense_range=np.asarray(header["dense_range"]), + dense_snp_range=np.asarray(header["dense_snp_range"]), + dense_indel_range=np.asarray(header["dense_indel_range"]), + sample_cols=np.asarray(header["sample_cols"]), + dense_max_end_keys=np.asarray(header["dense_max_end_keys"], np.int64), + chunks=_gen(), + ) ``` - [ ] **Step 4: Run the tests — expect PASS** @@ -1454,9 +1452,7 @@ def _svar2_ranges_cache_bytes(n_regions: int, n_samples: int, ploidy: int) -> in return 2 * n_regions * n_samples * ploidy * 2 * 8 -def _svar2_preflight( - out_dir: Path, n_regions: int, n_samples: int, ploidy: int -) -> int: +def _svar2_preflight(out_dir: Path, n_regions: int, n_samples: int, ploidy: int) -> int: """Log the projected ``svar2_ranges`` cache size and warn if disk is short. Warns rather than raising: free-space reporting is unreliable on some @@ -1682,54 +1678,50 @@ Remove the entire function at `python/genvarloader/_dataset/_write.py:1067-1121` Replace the loop body in `_write_from_svar2` (currently `_write.py:1178-1200`) with: ```python - max_ends = np.empty(R, np.int32) - contig_offset = 0 - pbar = tqdm(total=R, unit=" region") - for (c,), df in bed.partition_by( - "chrom", as_dict=True, maintain_order=True - ).items(): - c = cast(str, c) - pbar.set_description(f"Processing svar2 ranges for {df.height} regions on {c}") - lo, hi = contig_offset, contig_offset + df.height - rc = df.height - starts = df["chromStart"].to_numpy() - ends = df["chromEnd"].to_numpy() - # extend_to_length is validated at function entry (False raises); the - # read-bound kernel sizes haplotype output at read time. - stream = svar2._find_ranges_chunked( - c, starts, ends, samples=samples, max_mem=max_mem - ) - dense_snp[lo:hi] = np.asarray(stream.dense_snp_range, np.int64).reshape(rc, 2) - dense_indel[lo:hi] = np.asarray(stream.dense_indel_range, np.int64).reshape( - rc, 2 - ) +max_ends = np.empty(R, np.int32) +contig_offset = 0 +pbar = tqdm(total=R, unit=" region") +for (c,), df in bed.partition_by("chrom", as_dict=True, maintain_order=True).items(): + c = cast(str, c) + pbar.set_description(f"Processing svar2 ranges for {df.height} regions on {c}") + lo, hi = contig_offset, contig_offset + df.height + rc = df.height + starts = df["chromStart"].to_numpy() + ends = df["chromEnd"].to_numpy() + # extend_to_length is validated at function entry (False raises); the + # read-bound kernel sizes haplotype output at read time. + stream = svar2._find_ranges_chunked( + c, starts, ends, samples=samples, max_mem=max_mem + ) + dense_snp[lo:hi] = np.asarray(stream.dense_snp_range, np.int64).reshape(rc, 2) + dense_indel[lo:hi] = np.asarray(stream.dense_indel_range, np.int64).reshape(rc, 2) + + # Packed (pos << SHIFT) | ext keys, NOT unpacked ends: SVAR1 parity picks + # the highest-POSITION variant (ties by end), so a lower-position variant + # with a longer deletion must not win the cross-chunk reduction. + keys = stream.dense_max_end_keys.copy() + for ch in stream.chunks: + s0, s1 = ch.sample_start, ch.sample_start + ch.n_samples + # Chunks are hap-major (samples, ploidy, regions, 2); the cache is + # region-major. transpose() is a view -- numpy copies straight into + # the memmap with no intermediate array. + vk_snp[lo:hi, s0:s1] = ch.vk_snp_range.transpose(2, 0, 1, 3) + vk_indel[lo:hi, s0:s1] = ch.vk_indel_range.transpose(2, 0, 1, 3) + np.maximum(keys, ch.max_end_keys, out=keys) + # Bound the dirty page cache: at cohort scale these memmaps are tens + # of GiB and the kernel would otherwise reclaim at unpredictable times. + vk_snp.flush() + vk_indel.flush() + pbar.update(rc * ch.n_samples / S) + + mask = (1 << MAX_END_SHIFT) - 1 + region_ends = np.asarray(ends, np.int64).copy() + has = keys > 0 # 0 is the "no variant in this region" sentinel + region_ends[has] = (keys[has] >> MAX_END_SHIFT) + (keys[has] & mask) + max_ends[lo:hi] = region_ends.astype(np.int32) - # Packed (pos << SHIFT) | ext keys, NOT unpacked ends: SVAR1 parity picks - # the highest-POSITION variant (ties by end), so a lower-position variant - # with a longer deletion must not win the cross-chunk reduction. - keys = stream.dense_max_end_keys.copy() - for ch in stream.chunks: - s0, s1 = ch.sample_start, ch.sample_start + ch.n_samples - # Chunks are hap-major (samples, ploidy, regions, 2); the cache is - # region-major. transpose() is a view -- numpy copies straight into - # the memmap with no intermediate array. - vk_snp[lo:hi, s0:s1] = ch.vk_snp_range.transpose(2, 0, 1, 3) - vk_indel[lo:hi, s0:s1] = ch.vk_indel_range.transpose(2, 0, 1, 3) - np.maximum(keys, ch.max_end_keys, out=keys) - # Bound the dirty page cache: at cohort scale these memmaps are tens - # of GiB and the kernel would otherwise reclaim at unpredictable times. - vk_snp.flush() - vk_indel.flush() - pbar.update(rc * ch.n_samples / S) - - mask = (1 << MAX_END_SHIFT) - 1 - region_ends = np.asarray(ends, np.int64).copy() - has = keys > 0 # 0 is the "no variant in this region" sentinel - region_ends[has] = (keys[has] >> MAX_END_SHIFT) + (keys[has] & mask) - max_ends[lo:hi] = region_ends.astype(np.int32) - - contig_offset += df.height - pbar.close() + contig_offset += df.height +pbar.close() ``` Add the import near the other genoray imports at the top of `_write.py`: diff --git a/docs/superpowers/plans/2026-07-31-dataset-concat.md b/docs/superpowers/plans/2026-07-31-dataset-concat.md index 0f23a6dd..c016bd38 100644 --- a/docs/superpowers/plans/2026-07-31-dataset-concat.md +++ b/docs/superpowers/plans/2026-07-31-dataset-concat.md @@ -105,8 +105,12 @@ def test_provenance_regions_appends_blocks(): prov = provenance("regions", [(2, 2), (1, 2)], ploidy=1) # merged order is (r, s): A(r0s0) A(r0s1) A(r1s0) A(r1s1) B(r0s0) B(r0s1) assert prov.tolist() == [ - [0, 0], [0, 1], [0, 2], [0, 3], - [1, 0], [1, 1], + [0, 0], + [0, 1], + [0, 2], + [0, 3], + [1, 0], + [1, 1], ] @@ -115,8 +119,12 @@ def test_provenance_samples_interleaves_per_region(): prov = provenance("samples", [(2, 1), (2, 2)], ploidy=1) # merged S' = 3. Per region: A's sample, then B's two. assert prov.tolist() == [ - [0, 0], [1, 0], [1, 1], # region 0 - [0, 1], [1, 2], [1, 3], # region 1 + [0, 0], + [1, 0], + [1, 1], # region 0 + [0, 1], + [1, 2], + [1, 3], # region 1 ] @@ -397,8 +405,8 @@ def test_copy_runs_concatenates_two_ragged_sources(tmp_path): def test_copy_runs_interleaves_out_of_order_runs(tmp_path): # sample-axis shape: A slot0, B slot0, A slot1, B slot1 a, b = tmp_path / "a.npy", tmp_path / "b.npy" - _write_raw(a, np.array([10, 11, 12], np.int32)) # slots [10,11], [12] - _write_raw(b, np.array([20, 21], np.int32)) # slots [20], [21] + _write_raw(a, np.array([10, 11, 12], np.int32)) # slots [10,11], [12] + _write_raw(b, np.array([20, 21], np.int32)) # slots [20], [21] off_a = np.array([0, 2, 3], np.int64) off_b = np.array([0, 1, 2], np.int64) @@ -424,9 +432,7 @@ def test_copy_runs_handles_empty_slots(tmp_path): def test_copy_runs_spans_multiple_chunks(tmp_path, monkeypatch): """Force >1 chunk to exercise the streaming loop.""" - monkeypatch.setattr( - "genvarloader._dataset._concat_io.CONCAT_CHUNK_BYTES", 64 - ) + monkeypatch.setattr("genvarloader._dataset._concat_io.CONCAT_CHUNK_BYTES", 64) a = tmp_path / "a.npy" data = np.arange(1000, dtype=np.int32) _write_raw(a, data) @@ -690,8 +696,15 @@ def _mk(samples, n_regions, *, backend="pgen_vcf", ploidy=2, tracks=(), chroms=N ploidy=ploidy, ) return ConcatInput( - path=None, meta=meta, bed=bed, n_regions=n, n_samples=len(samples), - backend=backend, tracks=list(tracks), annot_tracks=[], has_dosages=False, + path=None, + meta=meta, + bed=bed, + n_regions=n, + n_samples=len(samples), + backend=backend, + tracks=list(tracks), + annot_tracks=[], + has_dosages=False, ) @@ -1096,12 +1109,14 @@ Append to `tests/dataset/test_concat.py`: ```python def _read_offsets(p: Path): import numpy as np + return np.fromfile(p / "genotypes" / "offsets.npy", dtype=np.int64) def _read_v_idxs(p: Path): import numpy as np from genoray._types import V_IDX_TYPE + return np.fromfile(p / "genotypes" / "variant_idxs.npy", dtype=V_IDX_TYPE) @@ -1122,6 +1137,7 @@ def test_concat_regions_metadata_matches(tmp_path, region_shards): gvl.concat(out, shards, axis="regions") import json + got = json.loads((out / "metadata.json").read_text()) exp = json.loads((whole / "metadata.json").read_text()) assert got["samples"] == exp["samples"] @@ -1136,6 +1152,7 @@ def test_concat_samples_merges_sample_list(tmp_path, sample_shards): gvl.concat(out, shards, axis="samples") import json + got = json.loads((out / "metadata.json").read_text()) exp = json.loads((whole / "metadata.json").read_text()) assert got["samples"] == exp["samples"] @@ -1162,6 +1179,7 @@ def test_concat_records_variants_fingerprint(tmp_path, region_shards): gvl.concat(out, shards, axis="regions") import json + meta = json.loads((out / "metadata.json").read_text()) fp = meta["variants_fingerprint"] assert fp["algorithm"] == "blake2b" @@ -1270,9 +1288,7 @@ def concat( """ from ._impl import Dataset as _Dataset - paths = [ - Path(d.path if isinstance(d, _Dataset) else d) for d in datasets - ] + paths = [Path(d.path if isinstance(d, _Dataset) else d) for d in datasets] dest = Path(path) if dest.exists() and not overwrite: raise FileExistsError(f"{dest} exists; pass overwrite=True to replace it") @@ -1419,6 +1435,7 @@ def svar_region_shards(tmp_path_factory, concat_case) -> tuple[list[Path], Path] def test_concat_svar_regions_matches_single_shot(tmp_path, svar_region_shards): import numpy as np + shards, whole = svar_region_shards out = tmp_path / "merged.gvl" gvl.concat(out, shards, axis="regions") @@ -1430,6 +1447,7 @@ def test_concat_svar_regions_matches_single_shot(tmp_path, svar_region_shards): def test_concat_svar_preserves_link(tmp_path, svar_region_shards): import json + shards, whole = svar_region_shards out = tmp_path / "merged.gvl" gvl.concat(out, shards, axis="regions") @@ -1442,6 +1460,7 @@ def test_concat_svar_preserves_link(tmp_path, svar_region_shards): def test_concat_regions_reads_equal_to_single_shot(tmp_path, region_shards, reference): """The real acceptance check: every merged cell reads identically.""" import numpy as np + shards, whole = region_shards out = tmp_path / "merged.gvl" gvl.concat(out, shards, axis="regions") @@ -1463,6 +1482,7 @@ def test_concat_samples_reads_equal_to_single_shot(tmp_path, sample_shards, refe full write's. Byte identity is not expected here. """ import numpy as np + shards, whole = sample_shards out = tmp_path / "merged.gvl" gvl.concat(out, shards, axis="samples") @@ -1617,8 +1637,9 @@ def _concat_svar2_ranges( for name in ("vk_snp_range", "vk_indel_range"): srcs = [ - np.fromfile(p / "genotypes" / "svar2_ranges" / f"{name}.npy", np.int64) - .reshape(-1, 2) + np.fromfile( + p / "genotypes" / "svar2_ranges" / f"{name}.npy", np.int64 + ).reshape(-1, 2) for p in paths ] out = np.empty((n_regions * n_samples * ploidy, 2), np.int64) @@ -1629,8 +1650,9 @@ def _concat_svar2_ranges( for name in ("dense_snp_range", "dense_indel_range"): srcs = [ - np.fromfile(p / "genotypes" / "svar2_ranges" / f"{name}.npy", np.int64) - .reshape(-1, 2) + np.fromfile( + p / "genotypes" / "svar2_ranges" / f"{name}.npy", np.int64 + ).reshape(-1, 2) for p in paths ] if axis == "samples": @@ -1643,8 +1665,9 @@ def _concat_svar2_ranges( np.load(p / "genotypes" / "svar2_ranges" / "sample_cols.npy") for p in paths ] if axis == "samples": - all_samples = [(s, c) for i, inp in zip(cols, inputs) - for s, c in zip(inp.meta.samples, i)] + all_samples = [ + (s, c) for i, inp in zip(cols, inputs) for s, c in zip(inp.meta.samples, i) + ] merged_cols = np.array( [c for _s, c in sorted(all_samples, key=lambda t: t[0])], np.int64 ) @@ -1737,6 +1760,7 @@ def test_open_rejects_mutated_variants_arrow(tmp_path, region_shards, reference) def test_open_accepts_absent_fingerprint(tmp_path, region_shards, reference): """Datasets written before the field exists must still open.""" import json + shards, _ = region_shards out = tmp_path / "merged.gvl" gvl.concat(out, shards, axis="regions") diff --git a/docs/superpowers/plans/2026-09-14-svar2-sparse-range-cache.md b/docs/superpowers/plans/2026-09-14-svar2-sparse-range-cache.md index 7e3d5667..c541e40f 100644 --- a/docs/superpowers/plans/2026-09-14-svar2-sparse-range-cache.md +++ b/docs/superpowers/plans/2026-09-14-svar2-sparse-range-cache.md @@ -111,7 +111,7 @@ chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1\t0|0 In the `svar2_store` fixture (`:53-71`), change the sample list argument to `run_conversion_pipeline`: ```python - ["S0", "S1", "S2"], +(["S0", "S1", "S2"],) ``` Leave `svar2_store_unsorted` (`:460-484`) at `["S1", "S0"]` — `run_conversion_pipeline` takes an explicit selection, so that fixture keeps exactly 2 samples and its `available_samples == ["S1", "S0"]` / `sample_cols == [1, 0]` assertions stay true. @@ -1331,7 +1331,9 @@ def test_iter_entries_is_sorted_and_complete(): dense = _random_dense(rng, R, S, P, fill=0.35) sparse = _sparse_from_dense(dense, R, S, P) - keys = np.concatenate([k for k, _ in sparse.iter_entries()] or [np.empty(0, np.int64)]) + keys = np.concatenate( + [k for k, _ in sparse.iter_entries()] or [np.empty(0, np.int64)] + ) ents = np.concatenate( [e for _, e in sparse.iter_entries()] or [np.empty(0, ENTRY_DTYPE)] ) @@ -1348,7 +1350,12 @@ def test_iter_entries_is_sorted_and_complete(): def test_iter_entries_empty_table(): R, S, P = 4, 3, 2 sparse = _SparseRanges( - np.zeros(R + 1, np.int64), np.empty(0, np.int32), np.empty(0, ENTRY_DTYPE), R, S, P + np.zeros(R + 1, np.int64), + np.empty(0, np.int32), + np.empty(0, ENTRY_DTYPE), + R, + S, + P, ) assert list(sparse.iter_entries()) == [] ``` @@ -1980,24 +1987,22 @@ Add the import at the top of the file: `from ._svar2_ranges import _RangeLookup, (c) Replace the cache construction in `from_path` (`:420-441`) — delete the local `_mm` helper and the `R`/`S` lines: ```python - ranges = _ranges_reader(ranges_dir) - R, S, P = ranges.n_regions, ranges.n_samples, ranges.ploidy - if P != ploidy: - raise ValueError(f"svar2 cache ploidy ({P}) != dataset ploidy ({ploidy}).") - - def _mm(name: str, shape: list[int]) -> NDArray[np.int64]: - return np.memmap( - ranges_dir / name, dtype=np.int64, mode="r", shape=tuple(shape) - ) +ranges = _ranges_reader(ranges_dir) +R, S, P = ranges.n_regions, ranges.n_samples, ranges.ploidy +if P != ploidy: + raise ValueError(f"svar2 cache ploidy ({P}) != dataset ploidy ({ploidy}).") - cache = _Svar2Cache( - ranges=ranges, - dense_snp_range=_mm("dense_snp_range.npy", meta["dense_snp_range"]["shape"]), - dense_indel_range=_mm( - "dense_indel_range.npy", meta["dense_indel_range"]["shape"] - ), - sample_cols=np.load(ranges_dir / "sample_cols.npy"), - ) + +def _mm(name: str, shape: list[int]) -> NDArray[np.int64]: + return np.memmap(ranges_dir / name, dtype=np.int64, mode="r", shape=tuple(shape)) + + +cache = _Svar2Cache( + ranges=ranges, + dense_snp_range=_mm("dense_snp_range.npy", meta["dense_snp_range"]["shape"]), + dense_indel_range=_mm("dense_indel_range.npy", meta["dense_indel_range"]["shape"]), + sample_cols=np.load(ranges_dir / "sample_cols.npy"), +) ``` (d) Replace the two `vk_*` lines in `_gather_inputs` (`:1562-1567`): @@ -3682,8 +3687,6 @@ Closes #355" Append to `tests/dataset/test_concat_svar2.py`: ```python - - def test_concat_svar2_emits_sparse_layout(svar2_shards_by_samples, tmp_path: Path): """The merged output must be sparse, whatever the inputs were.""" import json @@ -4676,7 +4679,9 @@ def main(): print(f"\nresident dense baseline (128 MB, 8192 cells): {ms:.2f} ms") print(f"\nworst shuffled: {worst:.2f}% of a {BATCH_MS} ms batch (gate: {GATE:.0%})") - print("PASS" if worst <= GATE * 100 else "FAIL -- take the spec's flat-key fallback") + print( + "PASS" if worst <= GATE * 100 else "FAIL -- take the spec's flat-key fallback" + ) if __name__ == "__main__": diff --git a/docs/superpowers/plans/2026-09-15-concat-analytic-run-planning.md b/docs/superpowers/plans/2026-09-15-concat-analytic-run-planning.md index 9dce1cb6..a8cb1322 100644 --- a/docs/superpowers/plans/2026-09-15-concat-analytic-run-planning.md +++ b/docs/superpowers/plans/2026-09-15-concat-analytic-run-planning.md @@ -219,9 +219,7 @@ class RunPlan: self.shape_per_ds = [(int(r), int(s)) for r, s in shape_per_ds] n_ds = len(self.shape_per_ds) if order is None: - counts = [ - r if axis == "regions" else s for r, s in self.shape_per_ds - ] + counts = [r if axis == "regions" else s for r, s in self.shape_per_ds] self.order = _default_order(n_ds, counts) else: self.order = np.asarray(order, dtype=np.int64).reshape(-1, 2) @@ -448,57 +446,58 @@ avoid. The samples axis is naturally bounded at `n_merged * ploidy` instead. Add to `RunPlan`: ```python - @property - def n_slots(self) -> int: - """Total merged flat slots this plan covers, computed arithmetically.""" - if not self.shape_per_ds: - return 0 - if self.axis == "regions": - return len(self.order) * self.shape_per_ds[0][1] * self.ploidy - return self.shape_per_ds[0][0] * len(self.order) * self.ploidy - - def slot_batches( - self, - ) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]": - """Yield ``(dst_start, src_ds, src_slots)`` batches 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. - - Yields: - ``(dst_start, src_ds, src_slots)``, where the two arrays are int64 - and equal in length. - """ - if self.axis == "regions": - for run in self: - pos, dst = run.src_start, run.dst_start - while pos < run.src_stop: - n = min(_SLOT_BATCH_SLOTS, run.src_stop - pos) - yield ( - dst, - np.full(n, run.src, np.int64), - np.arange(pos, pos + n, dtype=np.int64), - ) - pos += n - dst += n - return - - n_regions = self.shape_per_ds[0][0] if self.shape_per_ds else 0 - n_merged = len(self.order) - if n_regions == 0 or n_merged == 0 or self.ploidy == 0: - return - per_ds_samples = np.asarray([s for _, s in self.shape_per_ds], np.int64) - ds, w = self.order[:, 0], self.order[:, 1] - s_d = per_ds_samples[ds] - 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) - 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) +@property +def n_slots(self) -> int: + """Total merged flat slots this plan covers, computed arithmetically.""" + if not self.shape_per_ds: + return 0 + if self.axis == "regions": + return len(self.order) * self.shape_per_ds[0][1] * self.ploidy + return self.shape_per_ds[0][0] * len(self.order) * self.ploidy + + +def slot_batches( + self, +) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]": + """Yield ``(dst_start, src_ds, src_slots)`` batches 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. + + Yields: + ``(dst_start, src_ds, src_slots)``, where the two arrays are int64 + and equal in length. + """ + if self.axis == "regions": + for run in self: + pos, dst = run.src_start, run.dst_start + while pos < run.src_stop: + n = min(_SLOT_BATCH_SLOTS, run.src_stop - pos) + yield ( + dst, + np.full(n, run.src, np.int64), + np.arange(pos, pos + n, dtype=np.int64), + ) + pos += n + dst += n + return + + n_regions = self.shape_per_ds[0][0] if self.shape_per_ds else 0 + n_merged = len(self.order) + if n_regions == 0 or n_merged == 0 or self.ploidy == 0: + return + per_ds_samples = np.asarray([s for _, s in self.shape_per_ds], np.int64) + ds, w = self.order[:, 0], self.order[:, 1] + s_d = per_ds_samples[ds] + 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) + 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) ``` Then append the adapter and the normalizer: @@ -538,7 +537,9 @@ class ExplicitRunPlan: ) -def as_plan(runs: "RunPlan | ExplicitRunPlan | Sequence[Run]") -> "RunPlan | ExplicitRunPlan": +def as_plan( + runs: "RunPlan | ExplicitRunPlan | Sequence[Run]", +) -> "RunPlan | ExplicitRunPlan": """Normalize a run source so the IO layer has one consuming path. Args: @@ -837,9 +838,7 @@ def test_gather_svar_offsets_from_a_run_plan_is_byte_identical(tmp_path): out_plan.mkdir() out_list.mkdir() - _gather_svar_offsets( - paths, out_plan, RunPlan(axis, shapes, ploidy), shapes, ploidy - ) + _gather_svar_offsets(paths, out_plan, RunPlan(axis, shapes, ploidy), shapes, ploidy) _gather_svar_offsets( paths, out_list, @@ -915,18 +914,14 @@ Keep `t_runs` hoisted above `for name in ref.tracks:` — it is built once regar Replace ```python - a_prov = provenance( - "regions", [(r, 1) for r, _ in shapes], 1, order=order - ) - a_runs = coalesce(a_prov) +a_prov = provenance("regions", [(r, 1) for r, _ in shapes], 1, order=order) +a_runs = coalesce(a_prov) ``` with ```python - a_runs = RunPlan( - "regions", [(r, 1) for r, _ in shapes], 1, order=order - ) +a_runs = RunPlan("regions", [(r, 1) for r, _ in shapes], 1, order=order) ``` - [ ] **Step 7: Correct the two comments this change makes false** @@ -1138,8 +1133,15 @@ def _build_svar2(vcf_text: str, samples: list[str], d: Path, name: str) -> Path: out = d / name _core.run_conversion_pipeline( - str(bcf), str(ref), ["chr1"], str(out), samples, - 25_000, 2, 1, 8 * 1024 * 1024, + str(bcf), + str(ref), + ["chr1"], + str(out), + samples, + 25_000, + 2, + 1, + 8 * 1024 * 1024, ) assert (out / "meta.json").exists(), "conversion did not finish" return out @@ -1256,75 +1258,77 @@ Expected: FAIL — at minimum `test_fixture_has_empty_cells`, which pins the old Replace the body of the fixture guard at `tests/dataset/test_write_svar2.py:624-690` (the test asserting `sorted_samples == ["S0", "S1", "S2"]`, reshaping `nonempty` to `(3, S, P)`, and pinning `grid[0, S0, 0]`) with: ```python - svar2 = SparseVar2(svar2_store) - sorted_samples = sorted(svar2.available_samples) - assert sorted_samples == ["S0", "S1", "S2"], "fixture lost its third sample" - - d = svar2._find_ranges( - "chr1", - np.array([0, 5, 25]), - np.array([20, 15, 40]), - samples=sorted_samples, - ) - snp = np.asarray(d["vk_snp_range"], np.int64) # (R*S*P, 2) - indel = np.asarray(d["vk_indel_range"], np.int64) - w_snp = (snp[:, 1] - snp[:, 0]).reshape(3, len(sorted_samples), svar2.ploidy) - w_indel = (indel[:, 1] - indel[:, 0]).reshape(w_snp.shape) - occ = (w_snp > 0) | (w_indel > 0) - - s0, s1, s2 = (sorted_samples.index(s) for s in ("S0", "S1", "S2")) - - # Row-major (R, S, P) -- pinned by the layout oracle in - # test_write_svar2_emits_cache, which asserts this same reshape against the - # cache memmaps. Assert each structure SEPARATELY: a single `not occ.all()` - # is a disjunction that stays green when any one of them regresses alone. - assert not occ[:, s2].any(), "S2 is no longer all-reference; the empty COLUMN is gone" - assert not occ[2].any(), "region [25, 40) now holds sparse variants; the empty ROW is gone" - - # The measured grid. Every sparse variant carries exactly one call, which is - # what keeps genoray's cost model from routing it to the per-region dense - # channel; if a cost-model change pushes any of them dense this fails loudly - # rather than letting the sparse/dense parity tests pass against a thinner - # table while appearing green. - expected = np.zeros_like(occ) - expected[0, s0] = [True, True] - expected[0, s1] = [True, True] - expected[1, s0] = [False, True] - expected[1, s1] = [True, True] - np.testing.assert_array_equal(occ, expected) - assert int(occ.sum()) == 7, "fixture occupancy changed; update this pin deliberately" - - # Non-vacuity, kept from the original guard: an all-empty vk grid would make - # every sparse/dense parity test built on this fixture pass trivially. - assert occ.any(), ( - "vk channel is entirely empty: genoray routed every variant to the " - "dense channel, so all sparse-cache parity tests on this fixture are " - "now vacuous" - ) +svar2 = SparseVar2(svar2_store) +sorted_samples = sorted(svar2.available_samples) +assert sorted_samples == ["S0", "S1", "S2"], "fixture lost its third sample" + +d = svar2._find_ranges( + "chr1", + np.array([0, 5, 25]), + np.array([20, 15, 40]), + samples=sorted_samples, +) +snp = np.asarray(d["vk_snp_range"], np.int64) # (R*S*P, 2) +indel = np.asarray(d["vk_indel_range"], np.int64) +w_snp = (snp[:, 1] - snp[:, 0]).reshape(3, len(sorted_samples), svar2.ploidy) +w_indel = (indel[:, 1] - indel[:, 0]).reshape(w_snp.shape) +occ = (w_snp > 0) | (w_indel > 0) + +s0, s1, s2 = (sorted_samples.index(s) for s in ("S0", "S1", "S2")) + +# Row-major (R, S, P) -- pinned by the layout oracle in +# test_write_svar2_emits_cache, which asserts this same reshape against the +# cache memmaps. Assert each structure SEPARATELY: a single `not occ.all()` +# is a disjunction that stays green when any one of them regresses alone. +assert not occ[:, s2].any(), "S2 is no longer all-reference; the empty COLUMN is gone" +assert not occ[2].any(), ( + "region [25, 40) now holds sparse variants; the empty ROW is gone" +) + +# The measured grid. Every sparse variant carries exactly one call, which is +# what keeps genoray's cost model from routing it to the per-region dense +# channel; if a cost-model change pushes any of them dense this fails loudly +# rather than letting the sparse/dense parity tests pass against a thinner +# table while appearing green. +expected = np.zeros_like(occ) +expected[0, s0] = [True, True] +expected[0, s1] = [True, True] +expected[1, s0] = [False, True] +expected[1, s1] = [True, True] +np.testing.assert_array_equal(occ, expected) +assert int(occ.sum()) == 7, "fixture occupancy changed; update this pin deliberately" + +# Non-vacuity, kept from the original guard: an all-empty vk grid would make +# every sparse/dense parity test built on this fixture pass trivially. +assert occ.any(), ( + "vk channel is entirely empty: genoray routed every variant to the " + "dense channel, so all sparse-cache parity tests on this fixture are " + "now vacuous" +) - # Mixed per-channel emptiness inside a PRESENT cell -- the property the old - # one-cell grid could not express, and the reason two shipped oracles were - # able to mask per channel without failing. - assert (w_snp[0, s0, 0] > 0) and (w_indel[0, s0, 0] == 0) - assert (w_snp[0, s0, 1] > 0) and (w_indel[0, s0, 1] > 0) - assert (w_snp[0, s1, 0] == 0) and (w_indel[0, s1, 0] > 0) - - # An ABSENT cell inside a non-empty region: region 1 holds variants, but - # (S0, ploid 0) has none, so lookup must return (0, 0) there rather than a - # neighbour's range. - assert not occ[1, s0, 0] - - # The sample axis is now ORDERED, not just occupied: S0 and S1 differ at - # (region 1, ploid 0), so a transposed sample axis is detectable. During - # #357 this guard had to be withdrawn as unsatisfiable. - assert not np.array_equal(occ[:, s0], occ[:, s1]) - - # Both dense channels are exercised. dense_snp_range was all zeros before - # this fixture was enriched, so the dense SNP path had no coverage here. - dense_snp = np.asarray(d["dense_snp_range"], np.int64) - dense_indel = np.asarray(d["dense_indel_range"], np.int64) - assert (dense_snp[:, 1] > dense_snp[:, 0]).any(), "dense SNP channel is empty again" - assert (dense_indel[:, 1] > dense_indel[:, 0]).any(), "dense indel channel is empty" +# Mixed per-channel emptiness inside a PRESENT cell -- the property the old +# one-cell grid could not express, and the reason two shipped oracles were +# able to mask per channel without failing. +assert (w_snp[0, s0, 0] > 0) and (w_indel[0, s0, 0] == 0) +assert (w_snp[0, s0, 1] > 0) and (w_indel[0, s0, 1] > 0) +assert (w_snp[0, s1, 0] == 0) and (w_indel[0, s1, 0] > 0) + +# An ABSENT cell inside a non-empty region: region 1 holds variants, but +# (S0, ploid 0) has none, so lookup must return (0, 0) there rather than a +# neighbour's range. +assert not occ[1, s0, 0] + +# The sample axis is now ORDERED, not just occupied: S0 and S1 differ at +# (region 1, ploid 0), so a transposed sample axis is detectable. During +# #357 this guard had to be withdrawn as unsatisfiable. +assert not np.array_equal(occ[:, s0], occ[:, s1]) + +# Both dense channels are exercised. dense_snp_range was all zeros before +# this fixture was enriched, so the dense SNP path had no coverage here. +dense_snp = np.asarray(d["dense_snp_range"], np.int64) +dense_indel = np.asarray(d["dense_indel_range"], np.int64) +assert (dense_snp[:, 1] > dense_snp[:, 0]).any(), "dense SNP channel is empty again" +assert (dense_indel[:, 1] > dense_indel[:, 0]).any(), "dense indel channel is empty" ``` Rename the test to `test_fixture_grid_is_non_degenerate` and update its docstring to describe the 7-of-18 grid. Keep the module's existing import style (`from genoray import SparseVar2` inside the test, as the current version does). diff --git a/docs/superpowers/specs/2026-05-12-1kg-bcftools-parity-design.md b/docs/superpowers/specs/2026-05-12-1kg-bcftools-parity-design.md index 0181586f..a2c0502e 100644 --- a/docs/superpowers/specs/2026-05-12-1kg-bcftools-parity-design.md +++ b/docs/superpowers/specs/2026-05-12-1kg-bcftools-parity-design.md @@ -146,8 +146,7 @@ pytestmark = pytest.mark.slow def dataset_bcf(): return ( - gvl.Dataset - .open(data_dir / "1kg" / "phased_1kg.bcf.gvl", ref, rc_neg=False) + gvl.Dataset.open(data_dir / "1kg" / "phased_1kg.bcf.gvl", ref, rc_neg=False) .with_len("ragged") .with_seqs("haplotypes") .with_tracks(False) diff --git a/docs/superpowers/specs/2026-05-28-prefetching-dataloader-design.md b/docs/superpowers/specs/2026-05-28-prefetching-dataloader-design.md index d60e015a..bea45482 100644 --- a/docs/superpowers/specs/2026-05-28-prefetching-dataloader-design.md +++ b/docs/superpowers/specs/2026-05-28-prefetching-dataloader-design.md @@ -38,9 +38,9 @@ def to_dataloader( self, *, mode: Literal["buffered", "double_buffered"] | None = None, - buffer_bytes: int = 2 * 1024**3, # 2 GiB default; total footprint across all slots - copy: bool = True, # zero-copy opt-out; see §4 - heartbeat_seconds: float = 60.0, # double_buffered only; see §5 + buffer_bytes: int = 2 * 1024**3, # 2 GiB default; total footprint across all slots + copy: bool = True, # zero-copy opt-out; see §4 + heartbeat_seconds: float = 60.0, # double_buffered only; see §5 # ... existing args (batch_size, shuffle, sampler, ...) ... ) -> torch.utils.data.DataLoader: ... ``` diff --git a/docs/superpowers/specs/2026-05-29-dataloader-bench-design.md b/docs/superpowers/specs/2026-05-29-dataloader-bench-design.md index e722fae4..308801a5 100644 --- a/docs/superpowers/specs/2026-05-29-dataloader-bench-design.md +++ b/docs/superpowers/specs/2026-05-29-dataloader-bench-design.md @@ -160,9 +160,7 @@ else: "MKL_NUM_THREADS": str(n_threads), "OPENBLAS_NUM_THREADS": str(n_threads), } - subprocess.run( - [sys.executable, __file__, "--child"], env=env, check=True - ) + subprocess.run([sys.executable, __file__, "--child"], env=env, check=True) ``` The parent writes the CSV header once and child runs append. Children only diff --git a/docs/superpowers/specs/2026-05-30-dense2sparse-with-length-design.md b/docs/superpowers/specs/2026-05-30-dense2sparse-with-length-design.md index e3752877..e683edd2 100644 --- a/docs/superpowers/specs/2026-05-30-dense2sparse-with-length-design.md +++ b/docs/superpowers/specs/2026-05-30-dense2sparse-with-length-design.md @@ -36,12 +36,12 @@ input types (VCF, PGEN, SVAR) agree on the per-region `chromEnd`. ```python # genoray._svar def _dense2sparse_with_length( - genos: NDArray[np.integer], # (samples, ploidy, variants) — full window - var_idxs: NDArray[V_IDX_TYPE], # (variants,) — global, window-aligned - q_start: int, # 0-based, original (unextended) query start - q_end: int, # 0-based exclusive, original query end - v_starts: NDArray[np.int32], # POS - 1, window-aligned with var_idxs - ilens: NDArray[np.int32], # ILEN (ALT - REF length), window-aligned + genos: NDArray[np.integer], # (samples, ploidy, variants) — full window + var_idxs: NDArray[V_IDX_TYPE], # (variants,) — global, window-aligned + q_start: int, # 0-based, original (unextended) query start + q_end: int, # 0-based exclusive, original query end + v_starts: NDArray[np.int32], # POS - 1, window-aligned with var_idxs + ilens: NDArray[np.int32], # ILEN (ALT - REF length), window-aligned dosages: NDArray | None = None, ) -> Ragged[V_IDX_TYPE] | tuple[Ragged, Ragged]: """Convert a dense with_length window into per-haplotype-minimal sparse output, @@ -68,12 +68,12 @@ its current contract — it consumes `(list[Ragged], region_end, desc)` per regi ```python def _window_to_sparse( - genos: NDArray[np.integer], # (samples, ploidy, variants) — full region window - var_idxs: NDArray[V_IDX_TYPE], # (variants,) — global, window-aligned + genos: NDArray[np.integer], # (samples, ploidy, variants) — full region window + var_idxs: NDArray[V_IDX_TYPE], # (variants,) — global, window-aligned q_start: int, q_end: int, - v_starts: NDArray[np.int32], # POS - 1, window-aligned - ilens: NDArray[np.int32], # ILEN first-elem, window-aligned + v_starts: NDArray[np.int32], # POS - 1, window-aligned + ilens: NDArray[np.int32], # ILEN first-elem, window-aligned extend_to_length: bool, ) -> Ragged: if extend_to_length: @@ -111,7 +111,7 @@ Keeping the yielded type as `list[Ragged]` for both paths leaves Each generator computes, once from its index: ```python -v_ends = POS - ILEN_first.clip(upper_bound=0) # = POS + deletion length +v_ends = POS - ILEN_first.clip(upper_bound=0) # = POS + deletion length ``` For each region, `region_end = v_ends[max retained global v_idx]`, falling back to diff --git a/docs/superpowers/specs/2026-06-01-flat-buffer-getitem-followups-design.md b/docs/superpowers/specs/2026-06-01-flat-buffer-getitem-followups-design.md index 095636db..435b3d1b 100644 --- a/docs/superpowers/specs/2026-06-01-flat-buffer-getitem-followups-design.md +++ b/docs/superpowers/specs/2026-06-01-flat-buffer-getitem-followups-design.md @@ -50,10 +50,11 @@ Add two cases to the guard test, both using the **default track kind** (no ```python def test_tracks_ragged_no_awkward(monkeypatch, guard_dataset): calls = _install_ak_counters(monkeypatch) - ds = guard_dataset.with_seqs(None).with_tracks("5ss") # ragged output + ds = guard_dataset.with_seqs(None).with_tracks("5ss") # ragged output _ = ds[regions, samples] assert calls["n"] == 0 + def test_haps_tracks_ragged_no_awkward(monkeypatch, guard_dataset): calls = _install_ak_counters(monkeypatch) ds = guard_dataset.with_seqs("haplotypes").with_tracks("5ss") # ragged output @@ -121,18 +122,20 @@ elsewhere, for users who request raw intervals. ```python @dataclass(slots=True) class _FlatIntervals: - starts: _Flat # int32 - ends: _Flat # int32 - values: _Flat # float32 + starts: _Flat # int32 + ends: _Flat # int32 + values: _Flat # float32 @property - def shape(self): return self.starts.shape - def reshape(self, shape): ... # delegate to each _Flat + def shape(self): + return self.starts.shape + + def reshape(self, shape): ... # delegate to each _Flat def squeeze(self, axis=None): ... - def to_ragged(self): # boundary import RaggedIntervals - return RaggedIntervals(self.starts.to_ragged(), - self.ends.to_ragged(), - self.values.to_ragged()) + def to_ragged(self): # boundary import RaggedIntervals + return RaggedIntervals( + self.starts.to_ragged(), self.ends.to_ragged(), self.values.to_ragged() + ) ``` (Three component dtypes differ, so it cannot be a single `_Flat`. No `reverse_masked`/`to_fixed`/`to_padded` — intervals bypass RC and densify.) diff --git a/docs/superpowers/specs/2026-06-05-to_dataloader-drop_last-design.md b/docs/superpowers/specs/2026-06-05-to_dataloader-drop_last-design.md index 3cc8327d..7ca22edf 100644 --- a/docs/superpowers/specs/2026-06-05-to_dataloader-drop_last-design.md +++ b/docs/superpowers/specs/2026-06-05-to_dataloader-drop_last-design.md @@ -35,9 +35,9 @@ divisibility: ```python if n % batch_size != 0: - raise ValueError(...) # line 38-42 + raise ValueError(...) # line 38-42 ... -batch_totals = per_inst.reshape(-1, batch_size).sum(-1) # line 50 +batch_totals = per_inst.reshape(-1, batch_size).sum(-1) # line 50 ``` ### Bug 2 — default mode crashes on `drop_last=True` diff --git a/docs/superpowers/specs/2026-06-07-open-variants-no-reference-design.md b/docs/superpowers/specs/2026-06-07-open-variants-no-reference-design.md index c31fa16a..fff53ee6 100644 --- a/docs/superpowers/specs/2026-06-07-open-variants-no-reference-design.md +++ b/docs/superpowers/specs/2026-06-07-open-variants-no-reference-design.md @@ -10,6 +10,7 @@ crashes during open: ```python import genvarloader as gvl + gvl.Dataset.open("tests/data/phased_dataset.vcf.gvl") # ValueError: Cannot return RaggedSeqs: no reference genome was provided. ``` diff --git a/docs/superpowers/specs/2026-06-13-ref-fetch-parallel-overhead-design.md b/docs/superpowers/specs/2026-06-13-ref-fetch-parallel-overhead-design.md index f3e8ea49..2d4da7e0 100644 --- a/docs/superpowers/specs/2026-06-13-ref-fetch-parallel-overhead-design.md +++ b/docs/superpowers/specs/2026-06-13-ref-fetch-parallel-overhead-design.md @@ -56,9 +56,9 @@ def _resolve_num_threads() -> int: if env: return int(env) try: - real = len(os.sched_getaffinity(0)) # respects cgroup cpuset (Linux) + real = len(os.sched_getaffinity(0)) # respects cgroup cpuset (Linux) except AttributeError: - real = os.cpu_count() or 1 # non-Linux fallback + real = os.cpu_count() or 1 # non-Linux fallback return min(numba.get_num_threads(), real) ``` diff --git a/docs/superpowers/specs/2026-06-13-variant-assembly-B-design.md b/docs/superpowers/specs/2026-06-13-variant-assembly-B-design.md index 1837e4ff..bcfc19f5 100644 --- a/docs/superpowers/specs/2026-06-13-variant-assembly-B-design.md +++ b/docs/superpowers/specs/2026-06-13-variant-assembly-B-design.md @@ -85,8 +85,12 @@ transitions like `with_seqs` / `with_output_format`; type-preserving knobs live ```python from genvarloader import DummyVariant -ds.with_settings(dummy_variant=DummyVariant(start=-1, ref=b"N", alt=b"N", ilen=0, dosage=0.0)) -ds.with_settings(dummy_variant=False) # disable; mirrors the min_af/max_af `False` convention +ds.with_settings( + dummy_variant=DummyVariant(start=-1, ref=b"N", alt=b"N", ilen=0, dosage=0.0) +) +ds.with_settings( + dummy_variant=False +) # disable; mirrors the min_af/max_af `False` convention ``` - New parameter `dummy_variant: DummyVariant | Literal[False] | None = None` on `with_settings` diff --git a/docs/superpowers/specs/2026-06-18-svar-custom-format-fields-design.md b/docs/superpowers/specs/2026-06-18-svar-custom-format-fields-design.md index 4869c410..10685e6d 100644 --- a/docs/superpowers/specs/2026-06-18-svar-custom-format-fields-design.md +++ b/docs/superpowers/specs/2026-06-18-svar-custom-format-fields-design.md @@ -114,7 +114,9 @@ fields. for name in var_fields: if name in custom_fmt: mm = np.memmap(svar_path / f"{name}.npy", dtype=custom_fmt[name], mode="r") - var_field_data[name] = Ragged.from_offsets(mm, rag_shape, offsets.reshape(2, -1)) + var_field_data[name] = Ragged.from_offsets( + mm, rag_shape, offsets.reshape(2, -1) + ) ``` - Pass `var_field_data=var_field_data` to the `cls(...)` constructor. For the legacy (non-SVAR) branch it stays empty. diff --git a/docs/superpowers/specs/2026-06-30-rayon-multithread-verification-design.md b/docs/superpowers/specs/2026-06-30-rayon-multithread-verification-design.md index edc53776..e5f2c4d9 100644 --- a/docs/superpowers/specs/2026-06-30-rayon-multithread-verification-design.md +++ b/docs/superpowers/specs/2026-06-30-rayon-multithread-verification-design.md @@ -43,7 +43,7 @@ Single chokepoint. Add a `GVL_FORCE_PARALLEL` env var read by `should_paralleliz ```python def should_parallelize(total_bytes: int) -> bool: - if _force_parallel(): # GVL_FORCE_PARALLEL truthy → always parallel + if _force_parallel(): # GVL_FORCE_PARALLEL truthy → always parallel return True return total_bytes >= num_threads() * _MIN_BYTES_PER_THREAD ``` @@ -59,7 +59,7 @@ Two bugs, both named in #263: **(a) `setdefault` → overwrite.** An ambient `RAYON_NUM_THREADS=16` (base image) currently wins in spawn workers, so `cap_threads()` never caps. Change to assign directly: ```python -os.environ["RAYON_NUM_THREADS"] = str(_NUM_THREADS) # GVL's resolved count wins +os.environ["RAYON_NUM_THREADS"] = str(_NUM_THREADS) # GVL's resolved count wins ``` GVL's resolved count wins; users steer explicitly via `GVL_NUM_THREADS`. Still runs before the first rust parallel call, so it takes effect at global-pool init. (Overwriting *after* pool init is a no-op — rayon reads the env var once — so the existing "must run first" contract is unchanged and still documented.) diff --git a/docs/superpowers/specs/2026-07-27-seqpro-0.22-bump-design.md b/docs/superpowers/specs/2026-07-27-seqpro-0.22-bump-design.md index cf938449..012ab844 100644 --- a/docs/superpowers/specs/2026-07-27-seqpro-0.22-bump-design.md +++ b/docs/superpowers/specs/2026-07-27-seqpro-0.22-bump-design.md @@ -14,8 +14,8 @@ per item, so the two were not shape-consistent: ```python rv.start[0][h] # -> array([10, 20], dtype=int32) 2 elements, per-variant -rv.ilen[0][h] # -> array([0, 1], dtype=int32) 2 elements, per-variant -rv.alt[0][h] # -> b'AGG' ONE bytes: 'A'+'GG' or 'AG'+'G'? +rv.ilen[0][h] # -> array([0, 1], dtype=int32) 2 elements, per-variant +rv.alt[0][h] # -> b'AGG' ONE bytes: 'A'+'GG' or 'AG'+'G'? ``` This surfaces through gvl's public API on `RaggedVariants`, where `alt`, `ref`, `start`, diff --git a/docs/superpowers/specs/2026-07-30-svar2-write-memory-design.md b/docs/superpowers/specs/2026-07-30-svar2-write-memory-design.md index 7c3cb413..a4ccd4f6 100644 --- a/docs/superpowers/specs/2026-07-30-svar2-write-memory-design.md +++ b/docs/superpowers/specs/2026-07-30-svar2-write-memory-design.md @@ -142,25 +142,25 @@ alongside the existing `_find_ranges`: ```python @dataclass(frozen=True) class RangesChunk: - sample_start: int # into the SELECTED sample axis + sample_start: int # into the SELECTED sample axis n_samples: int - vk_snp_range: NDArray[np.int64] # (n_samples, ploidy, R, 2), hap-major + vk_snp_range: NDArray[np.int64] # (n_samples, ploidy, R, 2), hap-major vk_indel_range: NDArray[np.int64] - max_end_keys: NDArray[np.int64] # (R,), packed key; 0 = no variant + max_end_keys: NDArray[np.int64] # (R,), packed key; 0 = no variant @dataclass(frozen=True) class RangesStream: n_regions: int - n_samples: int # progress denominator + n_samples: int # progress denominator ploidy: int - samples_per_chunk: int # derived; exposed for observability - region_starts: NDArray[np.int32] # eager, R-sized - dense_range: NDArray[np.int32] # (R, 2) - dense_snp_range: NDArray[np.int32] # (R, 2) - dense_indel_range: NDArray[np.int32] # (R, 2) - sample_cols: NDArray[np.int64] # (S,) - dense_max_end_keys: NDArray[np.int64] # (R,), dense-channel contribution + samples_per_chunk: int # derived; exposed for observability + region_starts: NDArray[np.int32] # eager, R-sized + dense_range: NDArray[np.int32] # (R, 2) + dense_snp_range: NDArray[np.int32] # (R, 2) + dense_indel_range: NDArray[np.int32] # (R, 2) + sample_cols: NDArray[np.int64] # (S,) + dense_max_end_keys: NDArray[np.int64] # (R,), dense-channel contribution chunks: Iterator[RangesChunk] @@ -268,7 +268,7 @@ the parity test. Before creating any memmap: ```python -cache_bytes = 2 * R * S * P * 2 * 8 # both vk channels +cache_bytes = 2 * R * S * P * 2 * 8 # both vk channels ``` Log it through the writer's logger and compare against @@ -293,7 +293,7 @@ for ch in stream.chunks: pbar.update(rc * ch.n_samples / S) mask = (1 << 21) - 1 region_ends = np.asarray(ends, np.int64).copy() -has = keys > 0 # 0 = no variant; keep the original chromEnd +has = keys > 0 # 0 = no variant; keep the original chromEnd region_ends[has] = (keys[has] >> 21) + (keys[has] & mask) max_ends[lo:hi] = region_ends.astype(np.int32) ``` diff --git a/docs/superpowers/specs/2026-09-15-concat-analytic-run-planning-design.md b/docs/superpowers/specs/2026-09-15-concat-analytic-run-planning-design.md index 69ba589c..d0a96532 100644 --- a/docs/superpowers/specs/2026-09-15-concat-analytic-run-planning-design.md +++ b/docs/superpowers/specs/2026-09-15-concat-analytic-run-planning-design.md @@ -11,8 +11,8 @@ source slot it comes from. Today it does that by materializing the whole mapping and then compressing it: ```python -prov = provenance(axis, shapes, ploidy, order=order) # (n_slots, 2) int64 -runs = coalesce(prov) # list[Run] +prov = provenance(axis, shapes, ploidy, order=order) # (n_slots, 2) int64 +runs = coalesce(prov) # list[Run] ``` Both intermediates are sized by the `(R, S[, P])` grid, which is exactly the @@ -98,7 +98,9 @@ re-iterable object in `_concat_plan.py`: class RunPlan: def __init__(self, axis, shape_per_ds, ploidy, *, order=None): ... def __iter__(self) -> Iterator[Run]: ... - def slot_batches(self) -> Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]: ... + def slot_batches( + self, + ) -> Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]: ... @property def n_slots(self) -> int: ... ``` diff --git a/skills/genvarloader/SKILL.md b/skills/genvarloader/SKILL.md index c0232e3f..2bf2e687 100644 --- a/skills/genvarloader/SKILL.md +++ b/skills/genvarloader/SKILL.md @@ -26,8 +26,7 @@ gvl.write( # 3. Open and configure (chainable fluent API) ds = ( - gvl.Dataset - .open("ds.gvl", reference="ref.fa") + gvl.Dataset.open("ds.gvl", reference="ref.fa") .with_seqs("haplotypes") .with_tracks(["signal"]) .with_insertion_fill(gvl.Repeat5pNormalized()) @@ -74,7 +73,9 @@ Create an SVAR from a normalized VCF/PGEN with `genoray`: ```python from genoray import VCF, SparseVar -SparseVar.from_vcf("normed.svar", VCF("normed.bcf"), max_mem="4g") # writes a .svar/ directory +SparseVar.from_vcf( + "normed.svar", VCF("normed.bcf"), max_mem="4g" +) # writes a .svar/ directory ``` SVARs are resolved at `Dataset.open` time via `metadata.json` → caller `svar=` arg → recorded relative path → recorded absolute path → sibling `*.svar`. See `docs/source/format.md` ("SVAR resolution at open time") and `_dataset/_svar_link.py`. Legacy symlink-based SVAR layouts: run `gvl.migrate_svar_link(path)` once to upgrade. @@ -283,9 +284,10 @@ In `"flat"` mode the hot path is zero-awkward; the returned containers carry `.d ```python ds_flat = ds.with_output_format("flat") -result = ds_flat[0:8, :] # FlatRagged or FlatAnnotatedHaps or FlatVariants +result = ds_flat[0:8, :] # FlatRagged or FlatAnnotatedHaps or FlatVariants # direct tensorization — no awkward round-trip import torch + t = torch.from_numpy(result.data) # or convert back ragged = result.to_ragged() @@ -305,19 +307,34 @@ import genvarloader as gvl # produced when tracks are active (see gotchas). # (a) ride-along flank tokens on the "variants" output -fv = (ds.with_tracks(False).with_seqs("variants").with_output_format("flat") - .with_settings(flank_length=128, token_alphabet=sp.DNA.alphabet, - unknown_token=len(sp.DNA)))[0:8] -fv.flank_tokens # FlatRagged, shape (b, p, ~v, 2*128), or None if not configured +fv = ( + ds.with_tracks(False) + .with_seqs("variants") + .with_output_format("flat") + .with_settings( + flank_length=128, token_alphabet=sp.DNA.alphabet, unknown_token=len(sp.DNA) + ) +)[0:8] +fv.flank_tokens # FlatRagged, shape (b, p, ~v, 2*128), or None if not configured # (b) per-allele windows: ref as a flanked window, alt as a bare tokenized allele -fw = (ds.with_tracks(False).with_output_format("flat") - .with_seqs("variant-windows", - gvl.VarWindowOpt(flank_length=128, token_alphabet=sp.DNA.alphabet, - unknown_token=len(sp.DNA), ref="window", alt="allele")))[0:8] -fw.ref_window # flanked ref window tokens (two-level token buffer) -fw.alt # bare alt allele tokens (no flanks); fw.alt_window is None -fw.ref_window.shape # the window buffer's own shape: (b, p, ~v, ~len) +fw = ( + ds.with_tracks(False) + .with_output_format("flat") + .with_seqs( + "variant-windows", + gvl.VarWindowOpt( + flank_length=128, + token_alphabet=sp.DNA.alphabet, + unknown_token=len(sp.DNA), + ref="window", + alt="allele", + ), + ) +)[0:8] +fw.ref_window # flanked ref window tokens (two-level token buffer) +fw.alt # bare alt allele tokens (no flanks); fw.alt_window is None +fw.ref_window.shape # the window buffer's own shape: (b, p, ~v, ~len) ``` **Ride-along `FlatVariants.flank_tokens`** (`with_seqs("variants")` + `with_settings(flank_length=L, token_alphabet=..., unknown_token=...)`): appends a `FlatRagged` of shape `(b, p, ~v, 2L)` to the returned `FlatVariants`. Per variant the buffer holds `[flank5 | flank3]` reference-context tokens (each `L` long). Coordinate rule: `flank5 = [start-L, start)`, `flank3 = [end, end+L)` where `end = start - min(ilen, 0) + 1`. `token_alphabet` (`str`, `bytes`, or `seqpro.NucleotideAlphabet` — e.g. `sp.alphabets.DNA` / `sp.DNA.alphabet`; normalized to `bytes` at the `with_settings`/`build_token_lut` boundary) and `unknown_token` (int) together build a 256-entry byte→token LUT (seqpro-style): each alphabet byte → its 0-based index; every other byte (including `N` and out-of-bounds padding) → `unknown_token`. `flank_length=0`/`None` disables; both `token_alphabet` and `unknown_token` must be set together. Token dtype is `uint8` when max token id ≤ 255, else `int32`; offsets are `int64`. When `with_settings(dummy_variant=...)` is set, each empty `(region, sample, ploid)` group's `flank_tokens` row is a `2L`-long run of `unknown_token`. @@ -386,10 +403,10 @@ Use `gvl.sites_vcf_to_table(vcf)` → `pl.DataFrame` (bi-allelic SNPs only), the ```python loader = ds.to_dataloader( batch_size=32, - mode="double_buffered", # or "buffered", or None - buffer_bytes=2 * 1024**3, # total RAM budget; split across slots in double mode - copy=True, # zero-copy opt-out (default True = safe) - heartbeat_seconds=60.0, # double_buffered: max wait per chunk before liveness check + mode="double_buffered", # or "buffered", or None + buffer_bytes=2 * 1024**3, # total RAM budget; split across slots in double mode + copy=True, # zero-copy opt-out (default True = safe) + heartbeat_seconds=60.0, # double_buffered: max wait per chunk before liveness check ) ```