From 6524a6d1b0e9e8a3366f3cf5bffcea9e84b709d1 Mon Sep 17 00:00:00 2001 From: Kemal Yaylali Date: Mon, 28 Sep 2026 21:28:55 +0100 Subject: [PATCH] P9 (WIP): compiled kernels (rust/simurg-core, PyO3) with bit-identical pure-Python fallback Rust ports of the hot paths found by profiling: fused per-locus GenCall calling and scoring, BAF/LRR, libm maps with CPython math semantics, both nearest-neighbour searches of the normalization fit, and the BPM LocusEntry parser. Address coverage is memoised per IDAT layout. SIMURG_PURE_PYTHON=1 forces the pure paths. On the Pi: genotype_sample 3.4 s -> 1.1 s, 36-sample GSA run 131 s -> 46 s, BPM read 4x faster; 36 GSA + 6 Omni2.5 GTCs and QC tables byte-identical to the pure path; tests/unit/test_core_kernels.py checks every kernel on edge cases. bench/profile_genotype.py, bench/scaling.py, cloud task p9 (thread scaling). --- .github/workflows/ci.yml | 6 + bench/profile_genotype.py | 83 ++++++ bench/scaling.py | 109 ++++++++ rust/simurg-core/.gitignore | 1 + rust/simurg-core/Cargo.lock | 228 ++++++++++++++++ rust/simurg-core/Cargo.toml | 20 ++ rust/simurg-core/pyproject.toml | 14 + rust/simurg-core/src/bpm.rs | 275 +++++++++++++++++++ rust/simurg-core/src/gencall.rs | 357 +++++++++++++++++++++++++ rust/simurg-core/src/lib.rs | 226 ++++++++++++++++ scripts/cloud/tasks/p9.sh | 53 ++++ src/simurg/_core.py | 24 ++ src/simurg/formats/bpm.py | 52 +++- src/simurg/genotype/baf_lrr.py | 12 + src/simurg/genotype/gencall.py | 30 ++- src/simurg/pipeline/genotype.py | 2 +- src/simurg/signal/normalization/fit.py | 9 + src/simurg/signal/raw_xy.py | 22 +- tests/unit/test_core_kernels.py | 238 +++++++++++++++++ 19 files changed, 1755 insertions(+), 6 deletions(-) create mode 100644 bench/profile_genotype.py create mode 100644 bench/scaling.py create mode 100644 rust/simurg-core/.gitignore create mode 100644 rust/simurg-core/Cargo.lock create mode 100644 rust/simurg-core/Cargo.toml create mode 100644 rust/simurg-core/pyproject.toml create mode 100644 rust/simurg-core/src/bpm.rs create mode 100644 rust/simurg-core/src/gencall.rs create mode 100644 rust/simurg-core/src/lib.rs create mode 100644 scripts/cloud/tasks/p9.sh create mode 100644 src/simurg/_core.py create mode 100644 tests/unit/test_core_kernels.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 125afeb..cc4d69a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -23,3 +23,9 @@ jobs: run: uv run mypy - name: Unit tests (no downloaded data needed) run: uv run pytest tests/unit -q + - name: Build the compiled kernels (rust/simurg-core) + run: | + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal + PATH="$HOME/.cargo/bin:$PATH" uv pip install ./rust/simurg-core + - name: Unit tests with the compiled kernels (must equal the pure-Python results) + run: uv run pytest tests/unit -q diff --git a/bench/profile_genotype.py b/bench/profile_genotype.py new file mode 100644 index 0000000..bacdc09 --- /dev/null +++ b/bench/profile_genotype.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Per-stage timing of `simurg genotype` for one sample (plan §19 measurements). + + python bench/profile_genotype.py BPM EGT GRN_IDAT [--repeat N] +""" + +from __future__ import annotations + +import argparse +import datetime +import time +from dataclasses import replace +from pathlib import Path + +import numpy as np + +from simurg.formats.bpm import read_bpm +from simurg.formats.egt import read_egt +from simurg.formats.gtc import write_gtc +from simurg.formats.idat import read_idat +from simurg.genotype.baf_lrr import baf_lrr, logr_dev +from simurg.genotype.clusters import align_clusters +from simurg.genotype.gencall import call_genotypes +from simurg.genotype.sample import genotype_sample, manifest_context, params_for_preset +from simurg.genotype.sex import estimate_sex +from simurg.signal.normalization.apply import normalize, polar +from simurg.signal.normalization.fit import fit_transforms +from simurg.signal.raw_xy import LocusMapCache, assemble_raw_xy + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("bpm") + ap.add_argument("egt") + ap.add_argument("grn") + ap.add_argument("--repeat", type=int, default=3) + ap.add_argument("--out", default="/tmp/claude-1000/profile.gtc") + a = ap.parse_args() + m = read_bpm(a.bpm) + ctx = manifest_context(m, align_clusters(read_egt(a.egt), m)) + lite = m.lite() + ctx = replace(ctx, manifest=lite, cache=LocusMapCache(lite)) + params = params_for_preset("current") + red_path = a.grn.replace("_Grn", "_Red") + t: dict[str, list[float]] = {} + + def lap(name: str, t0: float) -> float: + now = time.perf_counter() + t.setdefault(name, []).append(now - t0) + return now + + for _ in range(a.repeat): + t0 = time.perf_counter() + grn, red = read_idat(a.grn), read_idat(red_path) + t0 = lap("read_idat x2", t0) + raw = assemble_raw_xy(grn, red, lite, cache=ctx.cache) + t0 = lap("assemble_raw_xy", t0) + xf = fit_transforms(raw.raw_x, raw.raw_y, ctx.norm_bins, params.norm_version) + t0 = lap("fit_transforms", t0) + nx, ny = normalize(raw.raw_x, raw.raw_y, ctx.lookup, xf) + th, r = polar(nx, ny) + t0 = lap("normalize+polar", t0) + z = (raw.raw_x == 0) & (raw.raw_y == 0) + thc = np.where(z, np.float32(np.nan), th).astype(np.float32) + rc = np.where(z, np.float32(np.nan), r).astype(np.float32) + calls = call_genotypes(thc, rc, z, ctx.clusters, lite.intensity_only, params.cutoff) + t0 = lap("call_genotypes", t0) + estimate_sex(calls.genotypes, r, ctx.clusters.gentrain, ctx.clusters.present, ctx.sex_classes, params.sex) + t0 = lap("estimate_sex", t0) + _, lrr = baf_lrr(thc, rc, z, ctx.clusters, lite.intensity_only) + logr_dev(lrr, ctx.autosomal, lite.intensity_only, ctx.clusters.present) + t0 = lap("baf_lrr+logr_dev", t0) + g = genotype_sample(grn, red, ctx, params, datetime.datetime(2026, 1, 1)) + t0 = lap("genotype_sample (whole)", t0) + write_gtc(g, Path(a.out)) + t0 = lap("write_gtc", t0) + print(f"{'stage':28s} {'median s':>9s}") + for k, v in t.items(): + print(f"{k:28s} {sorted(v)[len(v) // 2]:9.3f}") + + +if __name__ == "__main__": + main() diff --git a/bench/scaling.py b/bench/scaling.py new file mode 100644 index 0000000..b554663 --- /dev/null +++ b/bench/scaling.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +"""Thread scaling of ``simurg genotype`` (plan §18 P9 gate: ≥ 20 samples/min/core on GSA-size arrays, linear +scaling to 32 threads). + +For every worker count the same N samples are genotyped (``--no-cohort-qc``); the per-sample phase is timed from +the run log ("genotyping … sample(s)" → "done"), so manifest loading is reported separately and does not blur +the scaling. Writes a JSON summary. + + python bench/scaling.py --bpm M.bpm --egt C.egt --idats DIR --out OUT --threads 1,2,4,8 [--samples N] +""" + +from __future__ import annotations + +import argparse +import datetime as dt +import json +import os +import subprocess +import sys +import time +from pathlib import Path + + +def phase_seconds(log: Path) -> tuple[float, float]: + """(manifest+setup seconds, per-sample phase seconds) from a genotype run log.""" + ts: dict[str, dt.datetime] = {} + for line in log.read_text().splitlines(): + rec = json.loads(line) + msg = rec.get("msg") or "" + t = dt.datetime.fromisoformat(rec["ts"].replace("Z", "+00:00")) + ts.setdefault("first", t) + if msg.startswith("genotyping "): + ts["start"] = t + elif msg.startswith("done:"): + ts["done"] = t + return (ts["start"] - ts["first"]).total_seconds(), (ts["done"] - ts["start"]).total_seconds() + + +def subset(idats: Path, n: int, dest: Path) -> Path: + grn = sorted(p for p in idats.rglob("*_Grn.idat*")) + if n <= 0 or n >= len(grn): + return idats + dest.mkdir(parents=True, exist_ok=True) + for g in grn[:n]: + d = dest / g.parent.name + d.mkdir(exist_ok=True) + for p in (g, g.with_name(g.name.replace("_Grn.", "_Red."))): + link = d / p.name + if not link.exists(): + link.symlink_to(p.resolve()) + return dest + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--bpm", required=True) + ap.add_argument("--egt", required=True) + ap.add_argument("--idats", required=True, type=Path) + ap.add_argument("--out", required=True, type=Path) + ap.add_argument("--threads", default="1,2,4") + ap.add_argument("--samples", type=int, default=0, help="use the first N samples (0 = all)") + ap.add_argument("--simurg", default=str(Path(sys.executable).with_name("simurg"))) + ap.add_argument("--pure-python", action="store_true", help="run with SIMURG_PURE_PYTHON=1") + a = ap.parse_args() + a.out.mkdir(parents=True, exist_ok=True) + idats = subset(a.idats, a.samples, a.out / "idats") + n = len(list(idats.rglob("*_Grn.idat*"))) + env = {**os.environ, "SIMURG_PURE_PYTHON": "1" if a.pure_python else "0"} + runs = [] + for t in (int(x) for x in a.threads.split(",")): + out = a.out / f"t{t}" + t0 = time.time() + cmd = [a.simurg, "genotype", "--bpm", a.bpm, "--egt", a.egt, "--idats", str(idats), "--out", str(out), + "--threads", str(t), "--no-cohort-qc", "--autocall-date", "1/1/2026 12:00:00 AM", + "--quiet", "--force"] # fmt: skip + rc = subprocess.run(cmd, env=env).returncode + wall = time.time() - t0 + setup, phase = phase_seconds(out / "logs" / "genotype.log") + runs.append( + { + "threads": t, + "rc": rc, + "wall_s": round(wall, 1), + "setup_s": round(setup, 1), + "per_sample_phase_s": round(phase, 2), + "samples_per_min": round(60 * n / phase, 1) if phase else None, + "samples_per_min_per_core": round(60 * n / phase / t, 1) if phase else None, + } + ) + print(json.dumps(runs[-1]), flush=True) + base = runs[0]["per_sample_phase_s"] * runs[0]["threads"] + for r in runs: + r["speedup_vs_first"] = round(base / r["per_sample_phase_s"], 2) if r["per_sample_phase_s"] else None + r["efficiency"] = round(r["speedup_vs_first"] / r["threads"], 3) if r["speedup_vs_first"] else None + from simurg._core import backend + + summary = { + "samples": n, + "backend": "python" if a.pure_python else backend(), + "cpu_count": os.cpu_count(), + "machine": os.uname().machine, + "runs": runs, + } + (a.out / "scaling.json").write_text(json.dumps(summary, indent=1)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/rust/simurg-core/.gitignore b/rust/simurg-core/.gitignore new file mode 100644 index 0000000..2f7896d --- /dev/null +++ b/rust/simurg-core/.gitignore @@ -0,0 +1 @@ +target/ diff --git a/rust/simurg-core/Cargo.lock b/rust/simurg-core/Cargo.lock new file mode 100644 index 0000000..aa8d78f --- /dev/null +++ b/rust/simurg-core/Cargo.lock @@ -0,0 +1,228 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "numpy" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "portable-atomic" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" + +[[package]] +name = "portable-atomic-util" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "simurg-core" +version = "1.0.0" +dependencies = [ + "numpy", + "pyo3", +] + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "unicode-ident" +version = "1.0.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954" diff --git a/rust/simurg-core/Cargo.toml b/rust/simurg-core/Cargo.toml new file mode 100644 index 0000000..4d6c684 --- /dev/null +++ b/rust/simurg-core/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "simurg-core" +version = "1.0.0" +edition = "2021" +license = "MIT" +description = "Compiled kernels for SimurgArray (plan §18 P9)" +publish = false + +[lib] +name = "simurg_core" +crate-type = ["cdylib"] + +[dependencies] +numpy = "0.29.0" +pyo3 = { version = "0.29", features = ["extension-module", "abi3-py312"] } + +[profile.release] +opt-level = 3 +lto = true +codegen-units = 1 diff --git a/rust/simurg-core/pyproject.toml b/rust/simurg-core/pyproject.toml new file mode 100644 index 0000000..f9148eb --- /dev/null +++ b/rust/simurg-core/pyproject.toml @@ -0,0 +1,14 @@ +[build-system] +requires = ["maturin>=1.7,<2"] +build-backend = "maturin" + +[project] +name = "simurg-core" +description = "Compiled kernels for SimurgArray (optional; the pure-Python paths give identical results)" +requires-python = ">=3.12" +license = { text = "MIT" } +dependencies = ["numpy>=1.26"] +dynamic = ["version"] + +[tool.maturin] +features = ["pyo3/extension-module"] diff --git a/rust/simurg-core/src/bpm.rs b/rust/simurg-core/src/bpm.rs new file mode 100644 index 0000000..4253724 --- /dev/null +++ b/rust/simurg-core/src/bpm.rs @@ -0,0 +1,275 @@ +//! BPM LocusEntry records — port of `simurg.formats.bpm._parse_locus_entries` (format: Illumina BeadArrayFiles +//! `BeadPoolManifest.py`, BSD-2). Same field order, checks and error messages; strings of ≤ 64 bytes are shared +//! Python objects, as in the Python parser. + +use std::collections::HashMap; + +use numpy::{IntoPyArray, PyArray1, PyArray2}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::{PyList, PyString}; + +pub const STR_COLS: usize = 16; +// output column order (matches the Python `str_cols` list) +const ILMN: usize = 0; +const SNP: usize = 1; +const CHR: usize = 2; +const ISTR: usize = 3; +const SSTR: usize = 4; +const RSTR: usize = 5; +const PLOIDY: usize = 6; +const SPECIES: usize = 7; +const BUILD: usize = 8; +const SOURCE: usize = 9; +const SVER: usize = 10; +const CSTR: usize = 11; +const SSEQ: usize = 12; +const TOP: usize = 13; +const PA: usize = 14; +const PB: usize = 15; + +struct Cursor<'a> { + b: &'a [u8], + pos: usize, +} + +fn err(msg: impl Into) -> PyErr { + PyValueError::new_err(msg.into()) +} + +impl<'a> Cursor<'a> { + fn string(&mut self) -> PyResult<&'a [u8]> { + let size = self.b.len(); + let mut p = self.pos; + if p >= size { + return Err(err("unexpected end of file in locus entry")); + } + let mut byte = self.b[p]; + p += 1; + let mut ln = (byte & 0x7f) as usize; + let mut shift = 7; + while byte & 0x80 != 0 { + if p >= size || shift > 28 { + return Err(err("bad string length in locus entry")); + } + byte = self.b[p]; + p += 1; + ln |= ((byte & 0x7f) as usize) << shift; + shift += 7; + } + let end = p + ln; + if end > size { + return Err(err("string runs past end of file in locus entry")); + } + self.pos = end; + Ok(&self.b[p..end]) + } + fn i32(&mut self, what: &str) -> PyResult { + if self.pos + 4 > self.b.len() { + return Err(err(format!("unexpected end of file at {what}"))); + } + let v = i32::from_le_bytes(self.b[self.pos..self.pos + 4].try_into().unwrap()); + self.pos += 4; + Ok(v) + } +} + +type Parsed<'py> = ( + Vec>, + Bound<'py, PyArray1>, + Bound<'py, PyArray1>, + Bound<'py, PyArray1>, + Bound<'py, PyArray1>, + Bound<'py, PyArray1>, + Bound<'py, PyArray1>, + Bound<'py, PyArray1>, + Option>>, + usize, +); + +/// Parse `n` LocusEntry records starting at byte 0 of `buf`. Returns the 16 string columns, map_infos, +/// address_a, address_b, assay types (−1 = infer from sequences, entry version < 6), expected clusters, +/// intensity-only flags, entry versions, base fractions (None unless some entry is v7+) and the bytes consumed. +/// Errors are ValueError with the Python parser's message. +#[pyfunction] +pub fn bpm_locus_entries<'py>(py: Python<'py>, buf: &[u8], n: usize, names: Vec) -> PyResult> { + let name_to_idx: HashMap<&str, usize> = names.iter().enumerate().map(|(i, s)| (s.as_str(), i)).collect(); + let mut cols: Vec>>> = (0..STR_COLS).map(|_| (0..n).map(|_| None).collect()).collect(); + let mut intern: HashMap<&[u8], Py> = HashMap::new(); + let mut map_infos = vec![0i64; n]; + let mut address_a = vec![0i32; n]; + let mut address_b = vec![0i32; n]; + let mut assay = vec![0i16; n]; + let mut exp_clusters = vec![0i8; n]; + let mut int_only = vec![false; n]; + let mut versions = vec![0i8; n]; + let mut fracs = vec![0f32; 4 * n]; + let mut seen = vec![false; n]; + let mut any_v7 = false; + let mut c = Cursor { b: buf, pos: 0 }; + + let to_py = |raw: &'_ [u8], intern: &mut HashMap<&[u8], Py>| -> Py { + if let Some(s) = intern.get(raw) { + return s.clone_ref(py); + } + let s = PyString::new(py, &String::from_utf8_lossy(raw)).unbind(); + s + }; + + for _ in 0..n { + let ver = c.i32("locus entry")?; + if ![4, 6, 7, 8].contains(&ver) { + return Err(err(format!("unsupported locus entry version {ver}"))); + } + let ilmn_id = c.string()?; + let name = c.string()?; + for _ in 0..3 { + c.string()?; + } + c.i32("locus index")?; + c.string()?; + let ilmn_strand = c.string()?; + let snp = c.string()?; + let chrom = c.string()?; + let ploidy = c.string()?; + let species = c.string()?; + let map_info = c.string()?; + let top_seq = c.string()?; + let cust_strand = c.string()?; + if c.pos + 8 > buf.len() { + return Err(err("unexpected end of file at addresses")); + } + let a_addr = c.i32("addresses")?; + let b_addr = c.i32("addresses")?; + let probe_a = c.string()?; + let probe_b = c.string()?; + let build = c.string()?; + let source = c.string()?; + let source_ver = c.string()?; + let source_strand = c.string()?; + let source_seq = c.string()?; + let (mut at, mut exp_cl, mut io) = (-1i16, 0i8, 0u8); + if ver >= 6 { + if c.pos + 4 > buf.len() { + return Err(err("unexpected end of file at assay type")); + } + exp_cl = buf[c.pos + 1] as i8; + io = buf[c.pos + 2]; + at = buf[c.pos + 3] as i16; + c.pos += 4; + if at > 2 { + return Err(err(format!("invalid assay type {at}"))); + } + if (b_addr == 0) != (at == 0) { + return Err(err("assay type is inconsistent with address B")); + } + } + let mut f4: Option<[f32; 4]> = None; + if ver >= 7 { + if c.pos + 16 > buf.len() { + return Err(err("unexpected end of file at base fractions")); + } + let mut f = [0f32; 4]; + for (k, v) in f.iter_mut().enumerate() { + *v = f32::from_le_bytes(buf[c.pos + 4 * k..c.pos + 4 * k + 4].try_into().unwrap()); + } + f4 = Some(f); + c.pos += 16; + } + let ref_strand: &[u8] = if ver >= 8 { c.string()? } else { b"" }; + + let name_s = String::from_utf8_lossy(name); + let idx = match name_to_idx.get(name_s.as_ref()) { + Some(&i) => i, + None => return Err(err(format!("locus entry {} not in the names table", py_repr(&name_s)))), + }; + if seen[idx] { + return Err(err(format!("locus entry {} appears twice", py_repr(&name_s)))); + } + seen[idx] = true; + for (col, raw) in [ + (ILMN, ilmn_id), + (SNP, snp), + (CHR, chrom), + (ISTR, ilmn_strand), + (SSTR, source_strand), + (RSTR, ref_strand), + (PLOIDY, ploidy), + (SPECIES, species), + (BUILD, build), + (SOURCE, source), + (SVER, source_ver), + (CSTR, cust_strand), + (SSEQ, source_seq), + (TOP, top_seq), + (PA, probe_a), + (PB, probe_b), + ] { + let s = to_py(raw, &mut intern); + if raw.len() <= 64 && !intern.contains_key(raw) { + intern.insert(raw, s.clone_ref(py)); + } + cols[col][idx] = Some(s); + } + map_infos[idx] = if map_info.is_empty() { + 0 + } else { + let t = String::from_utf8_lossy(map_info); + match t.trim().parse::() { + Ok(v) => v, + Err(_) => { + return Err(err(format!( + "non-integer MapInfo {} for {}", + py_repr(&t), + py_repr(&name_s) + ))) + } + } + }; + address_a[idx] = a_addr; + address_b[idx] = b_addr; + assay[idx] = at; + exp_clusters[idx] = exp_cl; + int_only[idx] = io != 0; + versions[idx] = ver as i8; + if let Some(f) = f4 { + fracs[4 * idx..4 * idx + 4].copy_from_slice(&f); + any_v7 = true; + } + } + let missing = seen.iter().filter(|s| !**s).count(); + if missing > 0 { + return Err(err(format!("{missing} loci in the names table have no locus entry"))); + } + let empty = PyString::new(py, "").unbind(); + let lists = cols + .into_iter() + .map(|col| PyList::new(py, col.into_iter().map(|s| s.unwrap_or_else(|| empty.clone_ref(py))))) + .collect::>>()?; + let fr = if any_v7 { + Some(numpy::ndarray::Array2::from_shape_vec((n, 4), fracs).map_err(|e| err(e.to_string()))?.into_pyarray(py)) + } else { + None + }; + Ok(( + lists, + map_infos.into_pyarray(py), + address_a.into_pyarray(py), + address_b.into_pyarray(py), + assay.into_pyarray(py), + exp_clusters.into_pyarray(py), + int_only.into_pyarray(py), + versions.into_pyarray(py), + fr, + c.pos, + )) +} + +/// Python `repr()` of a str for error messages (single quotes unless the text contains one). +fn py_repr(s: &str) -> String { + if s.contains('\'') && !s.contains('"') { + format!("\"{s}\"") + } else { + format!("'{}'", s.replace('\\', "\\\\").replace('\'', "\\'")) + } +} diff --git a/rust/simurg-core/src/gencall.rs b/rust/simurg-core/src/gencall.rs new file mode 100644 index 0000000..0b45a7d --- /dev/null +++ b/rust/simurg-core/src/gencall.rs @@ -0,0 +1,357 @@ +//! Fused per-locus ports of `simurg.genotype.gencall.call_genotypes` and `simurg.genotype.baf_lrr.baf_lrr` +//! (which port gtc2vcf `gen_std_flair`/`modilik`/`compute_score_call_prelim`/`gencall_score_map`/`make_calls` +//! and `compute_baf_lrr`, MIT, © Giulio Genovese). Every intermediate keeps the NumPy reference's type (f32 or +//! f64) and evaluation order; NumPy's `fmin`/`fmax` are `f32::min`/`max`; libm calls use CPython `math` +//! semantics. One pass per locus instead of dozens of full-array temporaries. + +use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +use crate::{math_exp, math_log, math_pow}; + +const M_2_SQRTPI: f64 = 1.12837916709551257390; +const M_SQRT1_2: f64 = 0.707106781186547524401; +const NEG_NAN32: f32 = f32::from_bits(0xffc0_0000); + +#[inline] +fn smf(x: f32, a: f32, b: f32) -> f32 { + let mid = (a + b) / 2.0; + if a >= b { + return if x >= mid { 1.0 } else { 0.0 }; + } + let u = (x - a) / (b - a); + let v = (x - b) / (b - a); + if x <= a { + 0.0 + } else if a < x && x <= mid { + 2.0 * (u * u) + } else if mid < x && x <= b { + 1.0 - 2.0 * (v * v) + } else { + 1.0 + } +} + +#[inline] +fn zmf(x: f32, a: f32, b: f32) -> f32 { + let mid = (a + b) / 2.0; + let u = (x - a) / (b - a); + let v = (x - b) / (b - a); + if x <= a { + 1.0 + } else if a < x && x <= mid { + 1.0 - 2.0 * (u * u) + } else if mid < x && x <= b { + 2.0 * (v * v) + } else { + 0.0 + } +} + +/// Normal pdf pieces of `_pdf`: (full, peak, bad) for value x, mean mu, sd sigma. +#[inline] +fn pdf(x: f32, mu: f32, sigma: f32) -> (f64, f64, bool) { + let z = ((x - mu) / sigma) as f64; + let e = math_exp(-0.5 * (z * z)); + let sd = sigma as f64; + let full = e * 0.5 * M_2_SQRTPI * M_SQRT1_2 / sd; + let peak = (0.5 * M_2_SQRTPI * M_SQRT1_2) / sd; + (full, peak, sigma <= 0.0) +} + +#[derive(Clone, Copy)] +enum Side { + Left, + Right, + Full, +} + +#[inline] +fn get(p: (f64, f64, bool), x: f32, mu: f32, side: Side) -> f64 { + let (full, peak, bad) = p; + if bad { + return f64::NAN; + } + match side { + Side::Left => { + if x < mu { + peak + } else { + full + } + } + Side::Right => { + if x > mu { + peak + } else { + full + } + } + Side::Full => full, + } +} + +fn check(len: usize, n: usize, what: &str) -> PyResult<()> { + if len != n { + return Err(PyValueError::new_err(format!("gencall: {what} has length {len}, expected {n}"))); + } + Ok(()) +} + +/// Genotypes (0 NC, 1 AA, 2 AB, 3 BB) and GenCall scores; (n, 3) cluster arrays are passed flattened (C order). +#[pyfunction] +#[allow(clippy::too_many_arguments)] +pub fn call_genotypes<'py>( + py: Python<'py>, + theta: PyReadonlyArray1<'py, f32>, + r: PyReadonlyArray1<'py, f32>, + raw_zero: PyReadonlyArray1<'py, bool>, + present: PyReadonlyArray1<'py, bool>, + n_cl: PyReadonlyArray1<'py, i32>, + theta_mean: PyReadonlyArray1<'py, f32>, + theta_dev: PyReadonlyArray1<'py, f32>, + r_mean: PyReadonlyArray1<'py, f32>, + r_dev: PyReadonlyArray1<'py, f32>, + intensity_threshold: PyReadonlyArray1<'py, f32>, + gentrain: PyReadonlyArray1<'py, f32>, + intensity_only: PyReadonlyArray1<'py, bool>, + cutoff: f32, +) -> PyResult<(Bound<'py, PyArray1>, Bound<'py, PyArray1>)> { + let (th, rr, rz, pres) = (theta.as_slice()?, r.as_slice()?, raw_zero.as_slice()?, present.as_slice()?); + let (ncl, tmv, tdv, rmv, rdv) = + (n_cl.as_slice()?, theta_mean.as_slice()?, theta_dev.as_slice()?, r_mean.as_slice()?, r_dev.as_slice()?); + let (ithr, gt, io) = (intensity_threshold.as_slice()?, gentrain.as_slice()?, intensity_only.as_slice()?); + let n = th.len(); + for (l, w) in [(rr.len(), "r"), (rz.len(), "raw_zero"), (pres.len(), "present"), (ithr.len(), "threshold"), + (gt.len(), "gentrain"), (io.len(), "intensity_only")] { + check(l, n, w)?; + } + for (l, w) in [(ncl.len(), "n"), (tmv.len(), "theta_mean"), (tdv.len(), "theta_dev"), (rmv.len(), "r_mean"), + (rdv.len(), "r_dev")] { + check(l, 3 * n, w)?; + } + let tiny_d = f64::MIN_POSITIVE * f64::EPSILON; + let third: f32 = 1.0 / 3.0; + let mut geno = vec![0u8; n]; + let mut scores = vec![0f32; n]; + + for i in 0..n { + let t = th[i]; + let r = rr[i]; + let tm = [tmv[3 * i], tmv[3 * i + 1], tmv[3 * i + 2]]; + let td = [tdv[3 * i], tdv[3 * i + 1], tdv[3 * i + 2]]; + let rm = [rmv[3 * i], rmv[3 * i + 1], rmv[3 * i + 2]]; + let rd = [rdv[3 * i], rdv[3 * i + 1], rdv[3 * i + 2]]; + let nc = [ncl[3 * i], ncl[3 * i + 1], ncl[3 * i + 2]]; + + // gen_std_flair + let z1 = 0.5f32 * (tm[1] - tm[0]) / (td[1] + td[0]); + let z2 = 0.5f32 * (tm[2] - tm[1]) / (td[2] + td[1]); + let mz = z1.min(z2); + let alpha = (third * mz).max(1.0); + let beta = (third * mz).min(1.0); + let eta = alpha * beta; + let mut ftd = [0f32; 3]; + for k in 0..3 { + let v = td[k] * eta; + ftd[k] = if v < 0.02 { 0.02 } else { v }; + } + let med = td[0].min(td[1]).max(td[0].max(td[1]).min(td[2])); + let min_r = if 0.2f32 < med { med } else { 0.2 }; + let mut frd = [0f32; 3]; + for k in 0..3 { + let mrk0 = rm[k] / 7.0; + let mrk = if mrk0 < min_r { min_r } else { mrk0 }; + frd[k] = if rd[k] < mrk { mrk } else { rd[k] }; + } + + // modilik + let d: [f64; 3] = std::array::from_fn(|k| 100.0 * ((t - tm[k]).abs() as f64) + ((r - rm[k]).abs() as f64)); + let (daa, dab, dbb) = (d[0], d[1], d[2]); + let ok_t = !t.is_nan(); + let pt: [(f64, f64, bool); 3] = std::array::from_fn(|k| pdf(t, tm[k], ftd[k])); + let pr: [(f64, f64, bool); 3] = std::array::from_fn(|k| pdf(r, rm[k], frd[k])); + let pdf2 = |k: usize, st: Side, sr: Side| get(pt[k], t, tm[k], st) * get(pr[k], r, rm[k], sr); + use Side::{Full, Left, Right}; + let mut lik = [0f64; 3]; + let mut covered = false; + if ok_t { + let cases: [(bool, [(usize, Side, Side); 2]); 6] = [ + (daa <= dbb && dab <= dbb && rm[0] <= rm[1], [(0, Left, Left), (1, Right, Right)]), + (daa <= dab && dbb <= dab && rm[0] <= rm[2], [(0, Left, Left), (2, Right, Right)]), + (dab <= daa && dbb <= daa && rm[1] <= rm[2], [(1, Left, Left), (2, Right, Right)]), + (daa <= dbb && dab <= dbb && rm[0] > rm[1], [(0, Left, Right), (1, Right, Left)]), + (dab <= daa && dbb <= daa && rm[1] > rm[2], [(1, Left, Right), (2, Right, Left)]), + (daa <= dab && dbb <= dab && rm[0] > rm[2], [(0, Left, Right), (2, Right, Left)]), + ]; + for (cond, assigns) in cases { + if cond { + for (k, st, sr) in assigns { + lik[k] = pdf2(k, st, sr); + } + covered = true; + } + } + } + if !covered { + for (k, l) in lik.iter_mut().enumerate() { + *l = pdf2(k, Full, Full); + } + } + for l in lik.iter_mut() { + if l.is_nan() { + *l = 0.0; + } + } + if !ok_t { + lik = [f64::NAN; 3]; + } + + // priors and posteriors + let ntot = nc[0] as i64 + nc[1] as i64 + nc[2] as i64; + let p: [f64; 3] = std::array::from_fn(|k| (((nc[k] as f32) / (ntot as f32)) as f64 + 1.0) / 4.0); + let lp: [f64; 3] = std::array::from_fn(|k| lik[k] * p[k]); + let ev = (lp[0] + lp[1]) + lp[2]; + let (aa, ab, bb) = (lp[0] / ev, lp[1] / ev, lp[2] / ev); + let mut iap: u8 = if aa >= ab && aa >= bb { + 1 + } else if ab >= aa && ab >= bb { + 2 + } else if bb >= aa && bb >= ab { + 3 + } else { + 0 + }; + let mut mx = if aa > ab { aa } else { ab }; + let mut scnd = if aa > ab { ab } else { aa }; + let b_gt_mx = bb > mx; + scnd = if b_gt_mx { mx } else if bb > scnd { bb } else { scnd }; + mx = if b_gt_mx { bb } else { mx }; + let ap_lod = math_log(mx / (tiny_d + scnd), true); + let score_ap = smf(ap_lod as f32, 0.0, 2.0); + + // sub-scores + let score_r1 = smf(r, 0.0, 0.1); + let nclus = nc.iter().filter(|&&v| v > 0).count(); + let misclust: f32 = if nclus == 1 && nc[1] == 0 { + 0.7 + } else if nclus != 3 { + 0.95 + } else { + 1.0 + }; + let g = (iap as i64 - 1).clamp(0, 2) as usize; + let (gm_t, gd_t, gm_r, gd_r) = (tm[g], ftd[g], rm[g], frd[g]); + let st_aa = zmf(t, gm_t + 2.0 * gd_t, gm_t + 6.0 * gd_t); + let st_ab = zmf(((t - gm_t) / gd_t).abs(), 2.0, 6.0); + let st_bb = smf(t, gm_t - 6.0 * gd_t, gm_t - 2.0 * gd_t); + let mut score_t = match iap { + 1 => st_aa, + 2 => st_ab, + 3 => st_bb, + _ => 1.0, + }; + let r2 = smf(r, 0.0, gm_r / 10.0); + let r3 = smf(r, gm_r - 6.0 * gd_r, gm_r - 2.0 * gd_r); + let r4 = zmf(r, gm_r + 6.0 * gd_r, gm_r + 12.0 * gd_r); + let mut score_r = if iap > 0 { ((score_r1 * r2) * r3) * r4 } else { 0.0 }; + let mut score_n: f32 = 1.0; + if gt[i].is_nan() { + score_t = f32::NAN; + score_r = f32::NAN; + } + if !ok_t { + score_t = f32::NAN; + score_r = f32::NAN; + score_n = f32::NAN; + } + let mut prelim = ((((score_ap * gt[i]) * score_t) * score_r) * score_n) * misclust; + if r < ithr[i] { + prelim = f32::NAN; + iap = 0; + } + + // gencall_score_map + let x = prelim as f64; + let q = (x - 1.08) / 0.504; + let gbell = 1.0 / (1.0 + math_pow(q * q, 1.71)); + let mut score = (math_pow(x, 0.35) * gbell) as f32; + if score.is_nan() { + score = 0.0; + } + if rz[i] || !pres[i] { + score = 0.0; + iap = 0; + } + geno[i] = if score < cutoff || io[i] { 0 } else { iap }; + scores[i] = score; + } + Ok((geno.into_pyarray(py), scores.into_pyarray(py))) +} + +/// BAF and LRR (float32, canonical NaN 0xFFC00000 where undefined). +#[pyfunction] +#[allow(clippy::too_many_arguments)] +pub fn baf_lrr<'py>( + py: Python<'py>, + theta: PyReadonlyArray1<'py, f32>, + r: PyReadonlyArray1<'py, f32>, + raw_zero: PyReadonlyArray1<'py, bool>, + present: PyReadonlyArray1<'py, bool>, + theta_mean: PyReadonlyArray1<'py, f32>, + r_mean: PyReadonlyArray1<'py, f32>, + mean_r: PyReadonlyArray1<'py, f32>, + intensity_only: PyReadonlyArray1<'py, bool>, +) -> PyResult<(Bound<'py, PyArray1>, Bound<'py, PyArray1>)> { + let (th, rr, rz, pres) = (theta.as_slice()?, r.as_slice()?, raw_zero.as_slice()?, present.as_slice()?); + let (tmv, rmv, mr, io) = (theta_mean.as_slice()?, r_mean.as_slice()?, mean_r.as_slice()?, intensity_only.as_slice()?); + let n = th.len(); + for (l, w) in [(rr.len(), "r"), (rz.len(), "raw_zero"), (pres.len(), "present"), (mr.len(), "mean_r"), (io.len(), "intensity_only")] { + check(l, n, w)?; + } + check(tmv.len(), 3 * n, "theta_mean")?; + check(rmv.len(), 3 * n, "r_mean")?; + let mut baf = vec![0f32; n]; + let mut lrr = vec![0f32; n]; + for i in 0..n { + let (t, r) = (th[i], rr[i]); + let (at, bt, ct) = (tmv[3 * i], tmv[3 * i + 1], tmv[3 * i + 2]); + let (ar, br, cr) = (rmv[3 * i], rmv[3 * i + 1], rmv[3 * i + 2]); + let (eq, lo, hi) = (t == bt, t < bt, t > bt); + let (r_ref, mut b) = if eq { + (br, 0.5f32) + } else if lo { + (ar + ((t - at) * (ar - br)) / (at - bt), ((t - at) / (bt - at)) * 0.5) + } else if hi { + (br + ((t - bt) * (cr - br)) / (ct - bt), 0.5 + ((t - bt) / (ct - bt)) * 0.5) + } else { + (f32::NAN, NEG_NAN32) + }; + let denom = if io[i] { mr[i] } else { r_ref }; + let ratio = (r / denom) as f64; + let mut l = math_log(ratio, false) as f32; + if r == 0.0 { + l = -f32::MAX; + } + if !(eq || lo || hi) { + l = NEG_NAN32; + } + b = if b < 0.0 { + 0.0 + } else if b > 1.0 { + 1.0 + } else { + b + }; + if rz[i] || !pres[i] || b.is_nan() { + b = NEG_NAN32; + } + if rz[i] || !pres[i] || l.is_nan() { + l = NEG_NAN32; + } + baf[i] = b; + lrr[i] = l; + } + Ok((baf.into_pyarray(py), lrr.into_pyarray(py))) +} diff --git a/rust/simurg-core/src/lib.rs b/rust/simurg-core/src/lib.rs new file mode 100644 index 0000000..809a828 --- /dev/null +++ b/rust/simurg-core/src/lib.rs @@ -0,0 +1,226 @@ +//! Compiled kernels for SimurgArray (plan §18 P9, §19). +//! +//! Every kernel is a drop-in replacement for a pure-Python/NumPy reference in the `simurg` package and returns +//! bit-identical results (checked by `tests/unit/test_core_kernels.py`): +//! +//! * `exp`, `log10`, `log2`, `pow` — element-wise glibc libm calls with the error semantics of CPython's `math` +//! module as used by `simurg.genotype.gencall._c_call` (NaN input returned as is, a domain error gives the +//! canonical +NaN, `log10(0)` = −inf, overflow = +inf). The libm symbols are the ones CPython's `math` calls, so +//! results equal the C reference implementations (gtc2vcf, MIT) that the Python port follows. +//! * `closest_binned` — the per-bin candidate scan of the vendor's binned nearest-neighbour search +//! (`simurg.signal.normalization.fit._closest_binned`): float32 squared distances, first minimum wins. +//! +//! * `call_genotypes`, `baf_lrr` — fused per-locus GenCall calling/scoring and BAF/LRR (`gencall.rs`). +//! * `bpm_locus_entries` — the BPM LocusEntry parser (`simurg.formats.bpm._parse_locus_entries`). +//! +//! No floating-point contraction (FMA) happens: Rust never fuses `a*b+c` implicitly. + +mod bpm; +mod gencall; + +use numpy::{IntoPyArray, PyArray1, PyReadonlyArray1}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +#[link(name = "m")] +unsafe extern "C" { + #[link_name = "exp"] + fn c_exp(x: f64) -> f64; + #[link_name = "log10"] + fn c_log10(x: f64) -> f64; + #[link_name = "log2"] + fn c_log2(x: f64) -> f64; + #[link_name = "pow"] + fn c_pow(x: f64, y: f64) -> f64; +} + +/// CPython's `float('nan')`. +const PY_NAN: f64 = f64::from_bits(0x7ff8_0000_0000_0000); + +#[inline] +pub(crate) fn math_exp(v: f64) -> f64 { + if v.is_nan() { + return v; + } + let r = unsafe { c_exp(v) }; + if r.is_nan() { PY_NAN } else if r.is_infinite() && v.is_finite() { f64::INFINITY } else { r } +} + +/// CPython `m_log10`/`m_log2` + `_c_call`: x > 0 → libm; 0 → −inf (log10) or NaN (log2); x < 0, −inf → NaN. +#[inline] +pub(crate) fn math_log(v: f64, ten: bool) -> f64 { + if v.is_nan() { + return v; + } + if v.is_infinite() { + return if v > 0.0 { v } else { PY_NAN }; + } + if v > 0.0 { + return unsafe { if ten { c_log10(v) } else { c_log2(v) } }; + } + if v == 0.0 && ten { f64::NEG_INFINITY } else { PY_NAN } +} + +/// CPython `math.pow(v, e)` for a finite exponent e > 0 that is not an integer, + `_c_call`. +#[inline] +pub(crate) fn math_pow(v: f64, e: f64) -> f64 { + if v.is_nan() { + return v; + } + if v.is_infinite() { + return f64::INFINITY; + } + if v == 0.0 { + return 0.0; + } + let r = unsafe { c_pow(v, e) }; + if r.is_nan() { PY_NAN } else if r.is_infinite() { f64::INFINITY } else { r } +} + +fn map_f64<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>, f: impl Fn(f64) -> f64) -> PyResult>> { + let s = x.as_slice()?; + let out: Vec = s.iter().map(|&v| f(v)).collect(); + Ok(out.into_pyarray(py)) +} + +#[pyfunction] +fn exp<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>) -> PyResult>> { + map_f64(py, x, math_exp) +} + +#[pyfunction] +fn log10<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>) -> PyResult>> { + map_f64(py, x, |v| math_log(v, true)) +} + +#[pyfunction] +fn log2<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>) -> PyResult>> { + map_f64(py, x, |v| math_log(v, false)) +} + +#[pyfunction] +fn pow<'py>(py: Python<'py>, x: PyReadonlyArray1<'py, f64>, e: f64) -> PyResult>> { + if !(e.is_finite() && e > 0.0 && e.fract() != 0.0) { + return Err(PyValueError::new_err("pow kernel: exponent must be finite, positive and not an integer")); + } + map_f64(py, x, |v| math_pow(v, e)) +} + +/// For each sweep point `p` searching bin `choice[p]` (0..11 regular, 12..22 shifted; any other value: none), the +/// index of the closest candidate among `eligible[i]` with `b1[i] == choice` (regular) or `b2[i] == choice − 12` +/// (shifted), in ascending order; squared distance in float32; −1 if no candidate or the best is not < 1e20. +#[pyfunction] +#[allow(clippy::too_many_arguments)] +fn closest_binned<'py>( + py: Python<'py>, + raw_x: PyReadonlyArray1<'py, f32>, + raw_y: PyReadonlyArray1<'py, f32>, + axis_x: PyReadonlyArray1<'py, f32>, + axis_y: PyReadonlyArray1<'py, f32>, + choice: PyReadonlyArray1<'py, i64>, + eligible: PyReadonlyArray1<'py, i64>, + b1: PyReadonlyArray1<'py, i64>, + b2: PyReadonlyArray1<'py, i64>, +) -> PyResult>> { + let (rx, ry) = (raw_x.as_slice()?, raw_y.as_slice()?); + let (ax, ay, ch) = (axis_x.as_slice()?, axis_y.as_slice()?, choice.as_slice()?); + let (el, b1, b2) = (eligible.as_slice()?, b1.as_slice()?, b2.as_slice()?); + if ax.len() != ay.len() || ax.len() != ch.len() || el.len() != b1.len() || el.len() != b2.len() || rx.len() != ry.len() { + return Err(PyValueError::new_err("closest_binned: length mismatch")); + } + let mut groups: Vec> = vec![Vec::new(); 23]; + for (i, &e) in el.iter().enumerate() { + let e = usize::try_from(e).map_err(|_| PyValueError::new_err("closest_binned: negative index"))?; + if e >= rx.len() { + return Err(PyValueError::new_err("closest_binned: index out of range")); + } + if (0..12).contains(&b1[i]) { + groups[b1[i] as usize].push(e); + } + if (0..11).contains(&b2[i]) { + groups[12 + b2[i] as usize].push(e); + } + } + let out: Vec = (0..ax.len()) + .map(|p| { + let c = ch[p]; + if !(0..23).contains(&c) { + return -1; + } + let cand = &groups[c as usize]; + let mut best_j = usize::MAX; + let mut best = f32::INFINITY; + for (j, &e) in cand.iter().enumerate() { + let dx = rx[e] - ax[p]; + let dy = ry[e] - ay[p]; + let d = dx * dx + dy * dy; + // numpy argmin: the first NaN wins, otherwise the first strict minimum + if d.is_nan() { + best_j = j; + best = d; + break; + } + if best_j == usize::MAX || d < best { + best_j = j; + best = d; + } + } + if best_j == usize::MAX || !((best as f64) < 1e20) { -1 } else { cand[best_j] as i64 } + }) + .collect(); + Ok(out.into_pyarray(py)) +} + +/// Exact nearest neighbour (`simurg.signal.normalization.fit._closest_slow`): for each sweep point (x, y) the +/// index of the reference point with the smallest float32 squared distance `(x-xr)²+(y-yr)²`; first minimum wins, +/// a NaN distance wins as in `numpy.argmin`. `xref` must not be empty. +#[pyfunction] +fn closest_exact<'py>( + py: Python<'py>, + xref: PyReadonlyArray1<'py, f32>, + yref: PyReadonlyArray1<'py, f32>, + x: PyReadonlyArray1<'py, f32>, + y: PyReadonlyArray1<'py, f32>, +) -> PyResult>> { + let (xr, yr, xs, ys) = (xref.as_slice()?, yref.as_slice()?, x.as_slice()?, y.as_slice()?); + if xr.is_empty() || xr.len() != yr.len() || xs.len() != ys.len() { + return Err(PyValueError::new_err("closest_exact: empty reference or length mismatch")); + } + let out: Vec = xs + .iter() + .zip(ys) + .map(|(&px, &py_)| { + let mut best_j = 0usize; + let mut best = f32::INFINITY; + for (j, (&rx, &ry)) in xr.iter().zip(yr).enumerate() { + let dx = px - rx; + let dy = py_ - ry; + let d = dx * dx + dy * dy; + if d.is_nan() { + return j as i64; + } + if j == 0 || d < best { + best_j = j; + best = d; + } + } + best_j as i64 + }) + .collect(); + Ok(out.into_pyarray(py)) +} + +#[pymodule] +fn simurg_core(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + m.add_function(wrap_pyfunction!(exp, m)?)?; + m.add_function(wrap_pyfunction!(log10, m)?)?; + m.add_function(wrap_pyfunction!(log2, m)?)?; + m.add_function(wrap_pyfunction!(pow, m)?)?; + m.add_function(wrap_pyfunction!(closest_binned, m)?)?; + m.add_function(wrap_pyfunction!(closest_exact, m)?)?; + m.add_function(wrap_pyfunction!(bpm::bpm_locus_entries, m)?)?; + m.add_function(wrap_pyfunction!(gencall::call_genotypes, m)?)?; + m.add_function(wrap_pyfunction!(gencall::baf_lrr, m)?)?; + Ok(()) +} diff --git a/scripts/cloud/tasks/p9.sh b/scripts/cloud/tasks/p9.sh new file mode 100644 index 0000000..d0f2790 --- /dev/null +++ b/scripts/cloud/tasks/p9.sh @@ -0,0 +1,53 @@ +#!/bin/bash +# Phase P9 performance gate on a 32-vCPU Graviton instance, GSE165845 GSA arrays (360 samples): +# 1. build the compiled kernels (rust/simurg-core) into the venv; +# 2. thread scaling 1 → 32 with the kernels (bench/scaling.py, per-sample phase from the run log); +# 3. pure-Python reference at 1 and 32 threads on a subset, and byte-identity of all 360 GTCs (kernels vs pure). +# Args: . Run with OMNI=0. +set -uo pipefail +OUT=$2 +W=/work/p9; R=$OUT/p9; mkdir -p $W $R +N=$(nproc) + +echo "=== build simurg-core $(date -u +%T) ===" +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal > /dev/null 2>&1 +export PATH="$HOME/.cargo/bin:$PATH" +rustc --version +VIRTUAL_ENV=$PWD/.venv uv pip install -q ./rust/simurg-core || { echo "!! simurg-core build failed"; exit 1; } +.venv/bin/python -c "from simurg._core import backend; print('backend', backend())" +.venv/bin/python -m pytest tests/unit/test_core_kernels.py -q -p no:cacheprovider 2>&1 | tail -2 + +echo "=== data $(date -u +%T) ===" +GSA=https://webdata.illumina.com/downloads/productfiles/global-screening-array/v1-0 +mkdir -p $W/gsa && cd $W/gsa +for z in infinium-global-screening-array-24-v1-0-c2-manifest-file-bpm-build38.zip \ + infinium-global-screening-array-v1-0-c1-cluster-file.zip; do + curl -sSfL --retry 5 -O "$GSA/$z" && unzip -oq $z +done +cd - > /dev/null +BPM=$W/gsa/GSA-24v1-0_C2.bpm; EGT=$W/gsa/GSA-24v1-0_C1_ClusterFile.egt +.venv/bin/python scripts/p8/prepare.py --tsv scripts/p8/gse165845_samples.tsv --dl $W/geo --out $W/idats --workers 24 \ + || exit 1 + +echo "=== scaling with kernels $(date -u +%T) ===" +.venv/bin/python bench/scaling.py --bpm $BPM --egt $EGT --idats $W/idats --out $W/rust \ + --threads 1,2,4,8,16,$N 2>&1 | grep -v WARNING +cp $W/rust/scaling.json $R/scaling_rust.json + +echo "=== pure Python reference $(date -u +%T) ===" +.venv/bin/python bench/scaling.py --bpm $BPM --egt $EGT --idats $W/idats --out $W/python --threads 1 --samples 64 \ + --pure-python 2>&1 | grep -v WARNING +cp $W/python/scaling.json $R/scaling_python_t1.json +.venv/bin/python bench/scaling.py --bpm $BPM --egt $EGT --idats $W/idats --out $W/python_all --threads $N \ + --pure-python 2>&1 | grep -v WARNING +cp $W/python_all/scaling.json $R/scaling_python_t$N.json + +echo "=== identity: 360 GTCs, kernels vs pure Python $(date -u +%T) ===" +n=0; d=0 +for f in $(cd $W/rust/t$N && find samples -name '*.gtc'); do + n=$((n+1)); cmp -s $W/rust/t$N/$f $W/python_all/t$N/$f || { d=$((d+1)); echo "differs: $f"; } +done +s=$(cmp -s $W/rust/t$N/sample_summary.tsv $W/python_all/t$N/sample_summary.tsv && echo identical || echo DIFFERS) +echo "{\"gtcs\": $n, \"differ\": $d, \"sample_summary\": \"$s\"}" | tee $R/identity.json +lscpu | grep -i "model name\|^CPU(s)" > $R/cpu.txt; free -g >> $R/cpu.txt +echo "=== task done $(date -u +%T) ===" diff --git a/src/simurg/_core.py b/src/simurg/_core.py new file mode 100644 index 0000000..c6db901 --- /dev/null +++ b/src/simurg/_core.py @@ -0,0 +1,24 @@ +"""Optional compiled kernels (``rust/simurg-core``, plan §18 P9). + +``core`` is the ``simurg_core`` extension module when it is installed, else ``None``; every caller keeps a pure +Python/NumPy path with bit-identical results. ``SIMURG_PURE_PYTHON=1`` forces the pure paths (tests, debugging). +""" + +from __future__ import annotations + +import os +from types import ModuleType + +core: ModuleType | None +if os.environ.get("SIMURG_PURE_PYTHON", "") not in ("", "0"): + core = None +else: + try: + import simurg_core as core # type: ignore[no-redef] + except ImportError: # pragma: no cover - depends on the installation + core = None + + +def backend() -> str: + """``"rust "`` or ``"python"`` (recorded in run manifests).""" + return f"rust {core.__version__}" if core is not None else "python" diff --git a/src/simurg/formats/bpm.py b/src/simurg/formats/bpm.py index a445fd9..d250e86 100644 --- a/src/simurg/formats/bpm.py +++ b/src/simurg/formats/bpm.py @@ -32,6 +32,7 @@ from typing import Any import numpy as np import numpy.typing as npt +from simurg._core import core from simurg.core.errors import FormatError from simurg.formats.binutil import BinaryReader, load_bytes @@ -223,7 +224,9 @@ def read_bpm(path: str | Path) -> Manifest: def _parse_locus_entries(r: BinaryReader, n: int, name_to_idx: dict[str, int]) -> dict[str, Any]: - """Hot loop: parse ``n`` LocusEntry records (optimised pure Python; Rust in phase P9).""" + """Hot loop: parse ``n`` LocusEntry records (compiled in ``simurg_core`` when installed; same results).""" + if core is not None: + return _parse_locus_entries_core(r, n, name_to_idx) view = memoryview(r._buf) size = r.size pos = r.pos @@ -416,6 +419,53 @@ def _parse_locus_entries(r: BinaryReader, n: int, name_to_idx: dict[str, int]) - } +_STR_COLS = ( + "ilmn_ids", + "snps", + "chroms", + "ilmn_strands", + "source_strands", + "ref_strands", + "ploidies", + "species", + "genome_builds", + "sources", + "source_versions", + "customer_strands", + "source_seqs", + "top_genomic_seqs", + "allele_a_probe_seqs", + "allele_b_probe_seqs", +) + + +def _parse_locus_entries_core(r: BinaryReader, n: int, name_to_idx: dict[str, int]) -> dict[str, Any]: + """``_parse_locus_entries`` through ``simurg_core.bpm_locus_entries`` (identical columns and error messages).""" + assert core is not None + pos = r.pos + try: + lists, map_infos, a, b, assay, exp_cl, int_only, versions, fracs, used = core.bpm_locus_entries( + bytes(r._view[pos:]), n, list(name_to_idx) + ) + except ValueError as e: + raise r.fail(str(e)) from None + out: dict[str, Any] = dict(zip(_STR_COLS, lists, strict=True)) + for i in np.flatnonzero(assay < 0): + assay[i] = infer_assay_type(out["allele_a_probe_seqs"][i], out["allele_b_probe_seqs"][i], out["source_seqs"][i]) + r.pos = pos + used + return { + **out, + "map_infos": map_infos, + "address_a": a, + "address_b": b, + "assay_types": assay.astype(np.uint8), + "exp_clusters": exp_cl, + "intensity_only": int_only, + "entry_versions": versions, + "base_fractions": fracs, + } + + _IUPAC = { "A": 1, "C": 2, diff --git a/src/simurg/genotype/baf_lrr.py b/src/simurg/genotype/baf_lrr.py index f0646e1..e45267f 100644 --- a/src/simurg/genotype/baf_lrr.py +++ b/src/simurg/genotype/baf_lrr.py @@ -11,6 +11,7 @@ import math import numpy as np import numpy.typing as npt +from simurg._core import core from simurg.genotype.clusters import ClusterArrays from simurg.genotype.gencall import _libm @@ -22,6 +23,17 @@ NEG_NAN = np.frombuffer(np.uint32(0xFFC00000).tobytes(), dtype=np.float32)[0] # def baf_lrr( theta: F32Arr, r: F32Arr, raw_zero: npt.NDArray[np.bool_], c: ClusterArrays, intensity_only: npt.NDArray[np.bool_] ) -> tuple[F32Arr, F32Arr]: + if core is not None: + return core.baf_lrr( # type: ignore[no-any-return] + np.ascontiguousarray(theta, np.float32), + np.ascontiguousarray(r, np.float32), + np.ascontiguousarray(raw_zero, np.bool_), + np.ascontiguousarray(c.present, np.bool_), + np.ascontiguousarray(c.theta_mean, np.float32).ravel(), + np.ascontiguousarray(c.r_mean, np.float32).ravel(), + np.ascontiguousarray(c.mean_r, np.float32), + np.ascontiguousarray(intensity_only, np.bool_), + ) t = np.asarray(theta, np.float32) rr = np.asarray(r, np.float32) at, bt, ct = c.theta_mean[:, 0], c.theta_mean[:, 1], c.theta_mean[:, 2] diff --git a/src/simurg/genotype/gencall.py b/src/simurg/genotype/gencall.py index ba96833..1ba8718 100644 --- a/src/simurg/genotype/gencall.py +++ b/src/simurg/genotype/gencall.py @@ -15,6 +15,7 @@ from dataclasses import dataclass import numpy as np import numpy.typing as npt +from simurg._core import core from simurg.genotype.clusters import ClusterArrays F32 = np.float32 @@ -42,8 +43,17 @@ def _c_call(fn, v: float) -> float: # type: ignore[no-untyped-def] return math.inf +_CORE_FN = {math.exp: "exp", math.log10: "log10", math.log2: "log2"} + + def _libm(fn, x: F64Arr) -> F64Arr: # type: ignore[no-untyped-def] - """Element-wise libm call (bit-identical to C).""" + """Element-wise libm call (bit-identical to C); compiled when ``simurg_core`` is installed.""" + if core is not None: + xa = np.ascontiguousarray(x, dtype=np.float64).ravel() + if (e := getattr(fn, "exponent", None)) is not None: + return core.pow(xa, e).reshape(np.shape(x)) # type: ignore[no-any-return] + if (name := _CORE_FN.get(fn)) is not None: + return getattr(core, name)(xa).reshape(np.shape(x)) # type: ignore[no-any-return] return np.fromiter((_c_call(fn, v) for v in x.tolist()), dtype=np.float64, count=len(x)) @@ -51,6 +61,7 @@ def _pow(e: float): # type: ignore[no-untyped-def] def f(v: float) -> float: return math.pow(v, e) + f.exponent = e # type: ignore[attr-defined] return f @@ -137,6 +148,23 @@ def call_genotypes( cutoff: float = DEFAULT_CUTOFF, ) -> Calls: """Bayesian calls and GenCall scores for one sample (plan §7.5).""" + if core is not None: + f32 = [np.ascontiguousarray(a, np.float32).ravel() for a in (c.theta_mean, c.theta_dev, c.r_mean, c.r_dev)] + geno, score = core.call_genotypes( + np.ascontiguousarray(theta, np.float32), + np.ascontiguousarray(r, np.float32), + np.ascontiguousarray(raw_zero, np.bool_), + np.ascontiguousarray(c.present, np.bool_), + np.ascontiguousarray(c.n, np.int32).ravel(), + *f32, + np.ascontiguousarray(c.intensity_threshold, np.float32), + np.ascontiguousarray(c.gentrain, np.float32), + np.ascontiguousarray(intensity_only, np.bool_), + F32(cutoff), + ) + return Calls( + genotypes=geno, scores=score, num_calls=int((geno > 0).sum()), num_intensity_only=int(intensity_only.sum()) + ) n = len(theta) t = np.asarray(theta, np.float32) rr = np.asarray(r, np.float32) diff --git a/src/simurg/pipeline/genotype.py b/src/simurg/pipeline/genotype.py index 8668e02..f9b2a4c 100644 --- a/src/simurg/pipeline/genotype.py +++ b/src/simurg/pipeline/genotype.py @@ -160,7 +160,7 @@ def _process(job: Job) -> dict[str, Any]: try: grn, red = read_idat(job.grn), read_idat(job.red) check_pair(grn, red) - _, notes = check_coverage(grn, ctx.manifest) + _, notes = check_coverage(grn, ctx.manifest, ctx.cache) row["notes"].extend(notes) g = genotype_sample(grn, red, ctx, params, _W["started"]) write_gtc(g, Path(job.gtc_path)) diff --git a/src/simurg/signal/normalization/fit.py b/src/simurg/signal/normalization/fit.py index 1a1cf1c..21c1393 100644 --- a/src/simurg/signal/normalization/fit.py +++ b/src/simurg/signal/normalization/fit.py @@ -14,6 +14,7 @@ from dataclasses import dataclass import numpy as np import numpy.typing as npt +from simurg._core import core from simurg.signal.normalization.numerics import ( F32, TINY, @@ -86,6 +87,8 @@ def closest_points(xref: FloatArr, yref: FloatArr, x: FloatArr, y: FloatArr) -> def _closest_slow(xref: FloatArr, yref: FloatArr, x: FloatArr, y: FloatArr) -> IdxArr: """Exact search: squared distances in float32; the first minimum wins (C strict ``<``).""" + if core is not None and len(xref) and all(a.dtype == np.float32 for a in (xref, yref, x, y)): + return core.closest_exact(*(np.ascontiguousarray(a) for a in (xref, yref, x, y))) # type: ignore[no-any-return] dx = x[:, None] - xref[None, :] dy = y[:, None] - yref[None, :] return np.argmin(dx * dx + dy * dy, axis=1).astype(np.int64) @@ -118,6 +121,12 @@ def _closest_binned(raw_x: FloatArr, raw_y: FloatArr, axis_x: FloatArr, axis_y: low = rem < F32(0.25) choice = np.where(mid, k, np.where(low, np.where(k == 0, k, 12 + k - 1), np.where(k == 11, k, 12 + k))) + if core is not None and all(a.dtype == np.float32 for a in (raw_x, raw_y, axis_x, axis_y)): + i64 = np.int64 + return core.closest_binned( # type: ignore[no-any-return] + *(np.ascontiguousarray(a, np.float32) for a in (raw_x, raw_y, axis_x, axis_y)), + *(np.ascontiguousarray(a, i64) for a in (choice, eligible, b1, b2)), + ) out = np.full(len(axis_a), -1, dtype=np.int64) for c in np.unique(choice): cand = eligible[b1 == c] if c < 12 else eligible[b2 == c - 12] diff --git a/src/simurg/signal/raw_xy.py b/src/simurg/signal/raw_xy.py index 73611d4..460d57e 100644 --- a/src/simurg/signal/raw_xy.py +++ b/src/simurg/signal/raw_xy.py @@ -90,9 +90,11 @@ def address_coverage(idat: Idat, manifest: Manifest) -> float: return float(found.mean()) -def check_coverage(idat: Idat, manifest: Manifest) -> tuple[float, list[tuple[str, str]]]: +def check_coverage( + idat: Idat, manifest: Manifest, cache: LocusMapCache | None = None +) -> tuple[float, list[tuple[str, str]]]: """IN-027: < 0.99 → E_IDAT_MANIFEST_MISMATCH (raised); < 1 → W_IDAT_LOW_ADDRESS_COVERAGE (returned).""" - cov = address_coverage(idat, manifest) + cov = cache.coverage(idat) if cache is not None and cache.manifest is manifest else address_coverage(idat, manifest) notes: list[tuple[str, str]] = [] if cov < COVERAGE_ERROR: raise SimurgError( @@ -149,10 +151,24 @@ class LocusMapCache: self.manifest = manifest self.caddr = control_addresses(manifest) self._maps: dict[bytes, LocusMap] = {} + self._coverage: dict[bytes, float] = {} self.max_entries = max_entries + @staticmethod + def _key(idat: Idat) -> bytes: + return hashlib.blake2b(np.ascontiguousarray(idat.address_ids).tobytes(), digest_size=16).digest() + + def coverage(self, idat: Idat) -> float: + """``address_coverage`` memoised per address layout.""" + key = self._key(idat) + if (cov := self._coverage.get(key)) is None: + if len(self._coverage) >= self.max_entries: + self._coverage.pop(next(iter(self._coverage))) + cov = self._coverage[key] = address_coverage(idat, self.manifest) + return cov + def get(self, idat: Idat) -> LocusMap: - key = hashlib.blake2b(np.ascontiguousarray(idat.address_ids).tobytes(), digest_size=16).digest() + key = self._key(idat) lm = self._maps.get(key) if lm is None: if len(self._maps) >= self.max_entries: diff --git a/tests/unit/test_core_kernels.py b/tests/unit/test_core_kernels.py new file mode 100644 index 0000000..9370d7d --- /dev/null +++ b/tests/unit/test_core_kernels.py @@ -0,0 +1,238 @@ +"""The compiled kernels (``rust/simurg-core``, plan §18 P9) must equal the pure-Python paths bit for bit.""" + +from __future__ import annotations + +import datetime as dt +import math +from dataclasses import fields, replace + +import numpy as np +import pytest +from tests.conftest import GSA_BPM, GSA_EGT, GSA_IDATS + +import simurg.formats.bpm as bpm_mod +import simurg.genotype.baf_lrr as baf_mod +import simurg.genotype.gencall as gencall_mod +import simurg.signal.normalization.fit as fit_mod +from simurg._core import core +from simurg.core.errors import FormatError +from simurg.formats.bpm import read_bpm +from simurg.synth.files import SynthLocus, encode_bpm + +pytestmark = pytest.mark.skipif(core is None, reason="simurg_core extension not installed") + +MODULES = (bpm_mod, gencall_mod, fit_mod, baf_mod) + + +@pytest.fixture +def pure(monkeypatch): # type: ignore[no-untyped-def] + """Context switch: ``with pure(): ...`` runs the pure-Python paths.""" + + class _Pure: + def __enter__(self) -> None: + for m in MODULES: + monkeypatch.setattr(m, "core", None) + + def __exit__(self, *exc: object) -> None: + for m in MODULES: + monkeypatch.setattr(m, "core", core) + + return _Pure + + +def _bits(a: np.ndarray) -> bytes: + return np.ascontiguousarray(a).tobytes() + + +def _edge_values() -> np.ndarray: + neg_nan = np.frombuffer(np.uint64(0xFFF8000000000000).tobytes(), np.float64)[0] + payload_nan = np.frombuffer(np.uint64(0x7FF8000000001234).tobytes(), np.float64)[0] + tiny = np.finfo(np.float64).tiny + special = [0.0, -0.0, 1.0, -1.0, math.inf, -math.inf, math.nan, neg_nan, payload_nan, tiny, tiny / 4, -tiny, + 1e-320, 709.78, 709.79, -745.2, -746.0, 1e300, -1e300, 2.0**-1074, 0.5, 10.0, 1e-5] # fmt: skip + rng = np.random.default_rng(7) + rand = np.concatenate( + [rng.normal(0, 3, 20_000), rng.uniform(-800, 800, 5_000), np.exp(rng.uniform(-700, 700, 5_000))] + ) + return np.concatenate([np.array(special, np.float64), rand, -rand[:1000]]) + + +@pytest.mark.parametrize("fn", [math.exp, math.log10, math.log2, gencall_mod._pow(1.71), gencall_mod._pow(0.35)]) +def test_libm_kernels_bit_identical(fn, pure) -> None: # type: ignore[no-untyped-def] + x = _edge_values() + fast = gencall_mod._libm(fn, x) + with pure(): + ref = gencall_mod._libm(fn, x) + assert fast.dtype == ref.dtype == np.float64 + assert _bits(fast) == _bits(ref) + + +def test_pow_kernel_rejects_integer_exponent() -> None: + assert core is not None + with pytest.raises(ValueError): + core.pow(np.ones(3), 2.0) + + +def test_closest_binned_identical(pure) -> None: # type: ignore[no-untyped-def] + rng = np.random.default_rng(3) + for trial in range(20): + n = int(rng.integers(0, 3000)) + raw_x = rng.gamma(2.0, 400.0, n).astype(np.float32) + raw_y = rng.gamma(1.0, 60.0 if trial % 2 else 3000.0, n).astype(np.float32) + if trial == 5 and n: + raw_x[: n // 3] = raw_x[0] # ties: first minimum must win + raw_y[: n // 3] = raw_y[0] + if trial == 7 and n: + raw_x[n // 2] = np.nan + m = int(rng.integers(1, 400)) + axis = np.linspace(0, float(np.float32(rng.uniform(100, 5000))), m).astype(np.float32) + zeros = np.zeros(m, np.float32) + ax, ay = (axis, zeros) if trial % 3 else (zeros, axis) + fast = fit_mod._closest_binned(raw_x, raw_y, ax, ay) + with pure(): + ref = fit_mod._closest_binned(raw_x, raw_y, ax, ay) + assert fast.dtype == ref.dtype + assert np.array_equal(fast, ref), trial + + +def _loci() -> list[SynthLocus]: + return [ + SynthLocus("rs1", "1", 100, "[A/G]", 11, 0, 0, norm_id=1, source_seq="ACGT[A/G]TTGA", probe_a="ACGTA"), + SynthLocus("rs2", "2", 200, "[A/T]", 12, 13, 1, ilmn_strand="BOT", ref_strand="-", probe_a="AC", probe_b="AT"), + SynthLocus("rs3", "X", 300, "[C/G]", 14, 15, 2, norm_id=99), + SynthLocus("rsü", "Y", 400, "[T/C]", 16, 0, 0, norm_id=2, intensity_only=1), + SynthLocus("rs5", "MT", 500, "[A/C]", 17, 0, 0, norm_id=99, source_seq="A" * 200 + "[A/C]"), + ] + + +def _manifest_dict(m) -> dict: # type: ignore[no-untyped-def] + out = {} + for f in fields(m): + v = getattr(m, f.name) + out[f.name] = (v.dtype.str, v.shape, v.tobytes()) if isinstance(v, np.ndarray) else v + return out + + +@pytest.mark.parametrize("entry_version", [4, 6, 7, 8]) +@pytest.mark.parametrize("order", [None, [4, 3, 2, 1, 0]]) +def test_bpm_parser_identical(tmp_path, pure, entry_version, order) -> None: # type: ignore[no-untyped-def] + p = tmp_path / "m.bpm" + p.write_bytes(encode_bpm(_loci(), entry_version=entry_version, entry_order=order)) + fast = _manifest_dict(read_bpm(p)) + with pure(): + ref = _manifest_dict(read_bpm(p)) + assert fast == ref + + +def test_bpm_parser_errors_identical(tmp_path, pure) -> None: # type: ignore[no-untyped-def] + good = encode_bpm(_loci()) + cases = { + "truncated": good[: len(good) - 40], + "bad_assay": encode_bpm([SynthLocus("rs1", address_b=0, assay_type=1)]), + "bad_map": encode_bpm([SynthLocus("rs1", pos=12345)]).replace(b"\x0512345", b"\x0512x45"), + } + for label, data in cases.items(): + p = tmp_path / f"{label}.bpm" + p.write_bytes(data) + with pytest.raises(FormatError) as fast: + read_bpm(p) + with pure(), pytest.raises(FormatError) as ref: + read_bpm(p) + assert (fast.value.code, str(fast.value)) == (ref.value.code, str(ref.value)), label + + +@pytest.mark.data +@pytest.mark.skipif(not GSA_BPM.exists(), reason="GSA demo data not fetched") +def test_genotype_sample_identical(pure) -> None: # type: ignore[no-untyped-def] + from simurg.formats.egt import read_egt + from simurg.formats.gtc import encode_gtc + from simurg.formats.idat import read_idat + from simurg.genotype.clusters import align_clusters + from simurg.genotype.sample import genotype_sample, manifest_context, params_for_preset + from simurg.signal.raw_xy import LocusMapCache + + def run() -> bytes: + m = read_bpm(GSA_BPM) + ctx = manifest_context(m, align_clusters(read_egt(GSA_EGT), m)) + ctx = replace(ctx, manifest=m.lite(), cache=LocusMapCache(m.lite())) + stem = next(iter(sorted(GSA_IDATS.rglob("*_Grn.idat")))).as_posix()[: -len("_Grn.idat")] + g = genotype_sample(read_idat(stem + "_Grn.idat"), read_idat(stem + "_Red.idat"), ctx, + params_for_preset("current"), dt.datetime(2026, 1, 1)) # fmt: skip + return encode_gtc(g) + + fast = run() + with pure(): + ref = run() + assert fast == ref + + +def test_closest_exact_identical(pure) -> None: # type: ignore[no-untyped-def] + rng = np.random.default_rng(11) + for trial in range(10): + n = int(rng.integers(1, 1500)) + xr = rng.normal(500, 200, n).astype(np.float32) + yr = rng.normal(300, 100, n).astype(np.float32) + if trial == 2: + xr[n // 2 :] = xr[0] + yr[n // 2 :] = yr[0] + if trial == 3: + xr[n // 3] = np.nan + m = int(rng.integers(1, 300)) + x = rng.normal(500, 200, m).astype(np.float32) + y = rng.normal(300, 100, m).astype(np.float32) + fast = fit_mod._closest_slow(xr, yr, x, y) + with pure(): + ref = fit_mod._closest_slow(xr, yr, x, y) + assert np.array_equal(fast, ref), trial + + +def _synthetic_clusters(n: int, rng: np.random.Generator): # type: ignore[no-untyped-def] + from simurg.genotype.clusters import ClusterArrays + + tm = np.sort(rng.uniform(0, 1, (n, 3)), axis=1).astype(np.float32) + td = rng.uniform(0.0, 0.08, (n, 3)).astype(np.float32) + rm = rng.uniform(0.2, 3, (n, 3)).astype(np.float32) + rd = rng.uniform(0.0, 0.3, (n, 3)).astype(np.float32) + counts = rng.integers(0, 200, (n, 3)).astype(np.int32) + counts[rng.random((n, 3)) < 0.2] = 0 + counts[: n // 50] = 0 # no samples at all + td[rng.random((n, 3)) < 0.02] = 0.0 + tm[rng.random((n, 3)) < 0.01] = np.nan + rm[rng.random((n, 3)) < 0.01] = np.nan + thr = rng.uniform(0, 0.3, n).astype(np.float32) + thr[rng.random(n) < 0.3] = np.nan + gt = rng.uniform(0, 1, n).astype(np.float32) + gt[rng.random(n) < 0.02] = np.nan + return ClusterArrays( + present=rng.random(n) > 0.01, n=counts, theta_mean=tm, theta_dev=td, r_mean=rm, r_dev=rd, + intensity_threshold=thr, gentrain=gt, mean_r=rng.uniform(0.2, 3, n).astype(np.float32), + ) # fmt: skip + + +def test_call_genotypes_and_baf_lrr_identical(pure) -> None: # type: ignore[no-untyped-def] + from simurg.genotype.baf_lrr import baf_lrr + from simurg.genotype.gencall import call_genotypes + + rng = np.random.default_rng(5) + n = 60_000 + c = _synthetic_clusters(n, rng) + theta = rng.uniform(-0.05, 1.05, n).astype(np.float32) + theta[: n // 3] = c.theta_mean[: n // 3, rng.integers(0, 3)] # exactly on a cluster centre + r = rng.gamma(2.0, 0.6, n).astype(np.float32) + raw_zero = rng.random(n) < 0.01 + theta[raw_zero] = np.nan + r[raw_zero] = np.nan + r[rng.random(n) < 0.01] = 0.0 + intensity_only = rng.random(n) < 0.02 + for cutoff in (0.15, 0.0): + fast = call_genotypes(theta, r, raw_zero, c, intensity_only, cutoff) + with pure(): + ref = call_genotypes(theta, r, raw_zero, c, intensity_only, cutoff) + assert _bits(fast.genotypes) == _bits(ref.genotypes) + assert _bits(fast.scores) == _bits(ref.scores) + assert (fast.num_calls, fast.num_intensity_only) == (ref.num_calls, ref.num_intensity_only) + fb, fl = baf_lrr(theta, r, raw_zero, c, intensity_only) + with pure(): + rb, rl = baf_lrr(theta, r, raw_zero, c, intensity_only) + assert _bits(fb) == _bits(rb) + assert _bits(fl) == _bits(rl)