diff --git a/.github/scripts/wheel_smoke_test.py b/.github/scripts/wheel_smoke_test.py new file mode 100644 index 00000000..0dfd8a3c --- /dev/null +++ b/.github/scripts/wheel_smoke_test.py @@ -0,0 +1,54 @@ +"""Smoke-test an installed snputils wheel without using the source checkout.""" + +from importlib import import_module +from pathlib import Path +import platform + +import snputils + + +project_root = Path(__file__).resolve().parents[2] +installed_package = Path(snputils.__file__).resolve() +try: + installed_package.relative_to(project_root) +except ValueError: + pass +else: + raise RuntimeError(f"snputils was imported from the source checkout: {installed_package}") + +if snputils.__version__ == "unknown": + raise RuntimeError("The installed wheel does not expose package version metadata.") + +for module_name in ( + "snputils.snp.io.read._bcf", + "snputils.snp.io.write._bcf", + "snputils.snp.io._bgen", +): + import_module(module_name) + +for public_name in ( + "SNPObject", + "VCFReader", + "BCFReader", + "BGENReader", + "PGENReader", + "VCFWriter", + "BCFWriter", + "BGENWriter", + "PGENWriter", +): + getattr(snputils, public_name) + +snpobj = snputils.build_synthetic_snp_dataset(n_samples=4, n_snps=8, seed=1) +if snpobj.n_samples != 4 or snpobj.n_snps != 8: + raise RuntimeError("The installed wheel failed the synthetic SNPObject smoke test.") + +from snputils.tools.cli import main + +if main(["version"]) != 0: + raise RuntimeError("The installed wheel's CLI smoke test failed.") + +print( + f"Wheel smoke test passed: snputils {snputils.__version__}, " + f"{platform.system()} {platform.machine()}" +) diff --git a/.github/workflows/ci-cd.yml b/.github/workflows/ci-cd.yml index 64a01a13..14f9e8df 100644 --- a/.github/workflows/ci-cd.yml +++ b/.github/workflows/ci-cd.yml @@ -3,10 +3,11 @@ # 1. For Pull Requests: # - Runs all tests # - Verifies documentation builds +# - Builds and smoke-tests wheels on every supported platform # 2. For Pre-releases: # - Runs all tests # - Builds documentation -# - Publishes package to PyPI +# - Publishes the tested wheels and source distribution to PyPI # - Updates the release to a full release (removes pre-release status) and marks latest # - Deploys documentation to GitHub Pages (at the end; non-blocking) @@ -104,7 +105,21 @@ jobs: # Build documentation and verify it builds correctly - name: Build documentation - run: sphinx-build -b html docs docs/_build/html + run: sphinx-build -W --keep-going -b html docs docs/_build/html + + - name: Verify documentation discovery files + run: | + test -s docs/_build/html/robots.txt + test -s docs/_build/html/sitemap.xml + grep -Fxq 'Sitemap: https://docs.snputils.org/sitemap.xml' docs/_build/html/robots.txt + grep -Fq 'https://docs.snputils.org/' docs/_build/html/sitemap.xml + if grep -Eq '[^<]*/(search|genindex|py-modindex)\.html|[^<]*/_modules/' docs/_build/html/sitemap.xml; then + echo 'Sitemap contains an excluded utility page.' >&2 + exit 1 + fi + + - name: Check documentation links + run: sphinx-build -b linkcheck docs docs/_build/linkcheck # Upload Pages artifact if this is a release - name: Upload Pages artifact @@ -113,12 +128,51 @@ jobs: with: path: docs/_build/html + build-wheels: + needs: test + if: github.event_name == 'pull_request' || (github.event_name == 'release' && github.event.release.prerelease == true) + strategy: + fail-fast: false + matrix: + include: + - runner: ubuntu-latest + arch: x86_64 + artifact: wheels-linux-x86_64 + - runner: ubuntu-24.04-arm + arch: aarch64 + artifact: wheels-linux-aarch64 + - runner: macos-14 + arch: arm64 + artifact: wheels-macos-arm64 + - runner: macos-15-intel + arch: x86_64 + artifact: wheels-macos-x86_64 + runs-on: ${{ matrix.runner }} + steps: + - name: Checkout repository + uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: Build and smoke-test wheels + uses: pypa/cibuildwheel@v4.1.0 + with: + output-dir: wheelhouse + env: + CIBW_ARCHS: ${{ matrix.arch }} + + - name: Upload wheels + uses: actions/upload-artifact@v4 + with: + name: ${{ matrix.artifact }} + path: wheelhouse/*.whl + if-no-files-found: error + publish: - needs: [test, build-docs] - if: needs.build-docs.result == 'success' && github.event_name == 'release' && github.event.release.prerelease == true + needs: [test, build-docs, build-wheels] + if: needs.build-docs.result == 'success' && needs.build-wheels.result == 'success' && github.event_name == 'release' && github.event.release.prerelease == true runs-on: ubuntu-latest steps: - # Build and publish to PyPI - name: Checkout repository uses: actions/checkout@v6 with: @@ -136,23 +190,20 @@ jobs: - name: Install build dependencies run: | python -m pip install --upgrade pip - python -m pip install build - - - name: Build sdist - run: | - rm -rf dist wheelhouse - python -m build --sdist + python -m pip install build twine - - name: Build wheels - uses: pypa/cibuildwheel@v4.1.0 + - name: Download tested wheels + uses: actions/download-artifact@v4 with: - output-dir: wheelhouse - env: - CIBW_BUILD: "cp39-* cp310-* cp311-* cp312-* cp313-* cp314-*" - CIBW_SKIP: "*-musllinux_*" + pattern: wheels-* + path: dist + merge-multiple: true + + - name: Build sdist + run: python -m build --sdist --outdir dist - - name: Move wheels into dist - run: cp wheelhouse/*.whl dist/ + - name: Verify distributions + run: python -m twine check dist/* - name: Publish package to PyPI uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.gitignore b/.gitignore index 22f1cbd9..47931ea6 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,7 @@ __pycache__/ *$py.class *.lock *.txt +!docs/_extra/robots.txt # C extensions *.so @@ -161,4 +162,4 @@ snputils/tools/__test__/test_admixture_mapping_internal_large.py *.md !README.md !docs/**/*.md -.benchmarks/ \ No newline at end of file +.benchmarks/ diff --git a/README.md b/README.md index e885f471..24757df1 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,29 @@ Developed in collaboration between Stanford University's Department of Biomedica ## Quickstart +Start with a fully self-contained synthetic dataset—no external files or downloads required: + +```python +import snputils as su + +snp = su.build_synthetic_snp_dataset( + n_samples=30, + n_snps=100, + seed=42, +) + +af = snp.allele_freq() +pcs = su.PCA(n_components=2).fit_transform(snp) + +print(snp.n_samples, snp.n_snps) +print(af[:5]) +print(pcs.shape) +``` + +### Working with your own files + +For file-backed workflows, `read_snp` detects the genotype format from its extension and the other readers load ancestry, phenotype, and IBD data: + ```python import snputils as su @@ -182,7 +205,7 @@ The Python API remains the full surface for low-level readers/writers, object ma ## Documentation and Examples - **Documentation**: [docs.snputils.org](https://docs.snputils.org) -- **Quickstart**: [Quickstart guide](https://docs.snputils.org/en/latest/quickstart.html) +- **Quickstart**: [Quickstart guide](https://docs.snputils.org/quickstart.html) - **Tutorials**: PCA, mdPCA, maasMDS, SNP objects, allele frequency, local ancestry visualization, admixture mapping, and GRG workflows - **API Reference**: Readers, writers, data objects, processing classes, statistics, datasets, and visualization helpers - **Issues and feature requests**: [GitHub Issues](https://github.com/AI-sandbox/snputils/issues) diff --git a/benchmark/README.md b/benchmark/README.md index cdc86af4..90cce6b5 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -8,7 +8,16 @@ On the chromosome 22 of the 1000 Genomes Project dataset: **Legend: `genotype_mode`**: - 🟦 `"dosage"`: ALT dosage (`0|0 -> 0`, `0|1 -> 1`, `1|0 -> 1`, `1|1 -> 2`) -- 🟧 `"phased"`: allele codes (`0|0 -> [0, 0]`, `0|1 -> [0, 1]`, `1|0 -> [1, 0]`, `1|1 -> [1, 1]`) +- 🟧 `"phased"`: allele codes for PGEN, VCF, and BCF (`0|1 -> [0, 1]`); the BGEN panel instead uses `"probabilities"` to materialize the complete `float32` genotype-probability (`GP`) tensor + +The BGEN GP results are: + +| Reader | Time (mean ± SD) | Peak memory (mean ± SD) | +| --- | ---: | ---: | +| **snputils** | 13.31 ± 1.11 s | 28.5082 ± 0.0002 GiB | +| bgen | 16.49 ± 0.23 s | 28.3963 ± 0.0002 GiB | +| pysnptools | 51.49 ± 1.33 s | 28.5433 ± 0.0003 GiB | +| sgkit | 58.05 ± 1.94 s | 56.8980 ± 0.2273 GiB | ## Methodology @@ -17,6 +26,8 @@ The reader benchmark measures wall-clock read time and peak memory on chromosome Each reader is run in an Slurm allocation with 8 AMD EPYC 9684X CPU cores and 200 GB of RAM and a Python 3.12 environment. Python 3.12 is used because some libraries are not yet compatible with Python 3.13 or 3.14. +For BGEN, the dosage condition materializes expected alternate-allele counts from biallelic probabilities, while the orange condition materializes the complete `float32` genotype-probability tensor. The input contains 993,881 variants and 2,548 samples. This permits a like-for-like comparison across BGEN readers, including libraries that do not preserve phase. Every completed BGEN reader is checked against the corresponding output from the independent `bgen` package outside the timed and memory-profiled region. The correctness comparison is performed in variant chunks so that it does not create a second full-size comparison mask. Hail is omitted because its first GP materialization exhausted the fixed 200 GB allocation before returning the tensor. + After running the benchmark, the plotting utility `benchmark/plot_time_memory.py` writes both `benchmark/readers_benchmark.png` and `benchmark/readers_benchmark.pdf` from existing benchmark JSON files. ## Contributing diff --git a/benchmark/conftest.py b/benchmark/conftest.py index 0c5bc521..355dd4dd 100644 --- a/benchmark/conftest.py +++ b/benchmark/conftest.py @@ -27,8 +27,11 @@ def pytest_addoption(parser): "--genotype-mode", action="store", default="dosage", - choices=("dosage", "phased"), - help="Whether readers should return genotype dosages or phased allele calls." + choices=("dosage", "probabilities", "phased"), + help=( + "Whether readers should return genotype dosages, genotype probabilities, " + "or phased allele calls. Probability output is currently benchmarked for BGEN." + ), ) @@ -55,5 +58,8 @@ def reader_name(request): @pytest.fixture def genotype_mode(request): - """Fixture selecting dosage or phased genotype output.""" - return request.config.getoption("--genotype-mode") + """Fixture selecting the materialized genotype representation.""" + mode = request.config.getoption("--genotype-mode") + if mode == "probabilities" and request.node.path.name != "read_bgen.py": + pytest.skip('genotype_mode="probabilities" is only supported by the BGEN benchmark.') + return mode diff --git a/benchmark/plot_time_memory.py b/benchmark/plot_time_memory.py index a68c885c..d981e4f8 100644 --- a/benchmark/plot_time_memory.py +++ b/benchmark/plot_time_memory.py @@ -129,6 +129,7 @@ def _draw_grouped_bars( y_cap: float | None = None, axis_top: float | None = None, show_phased: bool = True, + secondary_label: str = 'genotype_mode="phased"', ) -> None: x = np.arange(len(names)) width = 0.36 if show_phased else 0.42 @@ -170,7 +171,7 @@ def _draw_grouped_bars( series = ( (-width / 2, dosage_values, 'genotype_mode="dosage"', colors["dosage"]), - (width / 2, phased_values, 'genotype_mode="phased"', colors["phased"]), + (width / 2, phased_values, secondary_label, colors["phased"]), ) if show_phased else ( (0.0, dosage_values, 'genotype_mode="dosage"', colors["dosage"]), ) @@ -278,9 +279,14 @@ def plot_time_memory( time_dosage, time_phased, "time", - f"{title} time", + f"{title} time" + (" (orange: GP)" if fmt == "bgen" else ""), axis_top=600 if fmt == "vcf" else None, show_phased=fmt != "bed", + secondary_label=( + 'genotype_mode="probabilities"' + if fmt == "bgen" + else 'genotype_mode="phased"' + ), ) _draw_grouped_bars( axs[row, 1], @@ -288,9 +294,14 @@ def plot_time_memory( memory_dosage, memory_phased, "memory", - f"{title} peak memory", + f"{title} peak memory" + (" (orange: GP)" if fmt == "bgen" else ""), y_cap=memory_y_cap, show_phased=fmt != "bed", + secondary_label=( + 'genotype_mode="probabilities"' + if fmt == "bgen" + else 'genotype_mode="phased"' + ), ) axs[row, 0].set_ylabel("Time (seconds)") axs[row, 1].set_ylabel("Peak memory (GiB)") diff --git a/benchmark/read_bgen.py b/benchmark/read_bgen.py index 23df68a1..00286f4c 100644 --- a/benchmark/read_bgen.py +++ b/benchmark/read_bgen.py @@ -50,6 +50,53 @@ def _probabilities_to_dosage(probabilities): return dosage +def _probabilities_to_phased_calls(probabilities): + """Hard-call diploid haplotypes from phased BGEN probabilities.""" + probabilities = np.asarray(probabilities, dtype=np.float32) + if probabilities.ndim not in (2, 3): + raise ValueError("Phased BGEN probabilities must be a 2D or 3D array.") + + width = probabilities.shape[-1] + if width < 4 or width % 2 != 0: + raise ValueError( + "Phased diploid BGEN probabilities must contain two equally sized " + "haplotype probability vectors." + ) + + n_alleles = width // 2 + haplotype_probabilities = probabilities.reshape( + probabilities.shape[:-1] + (2, n_alleles) + ) + finite = np.isfinite(haplotype_probabilities) + any_finite = np.any(finite, axis=-1) + fully_finite = np.all(finite, axis=-1) + if np.any(any_finite & ~fully_finite): + raise ValueError("Each phased haplotype probability vector must be complete or missing.") + if np.any(np.sum(fully_finite, axis=-1) == 1): + raise ValueError("The phased-call benchmark requires diploid BGEN records.") + + probability_sums = np.sum( + np.where(finite, haplotype_probabilities, 0.0), + axis=-1, + dtype=np.float32, + ) + if np.any(fully_finite & ~np.isclose(probability_sums, 1.0, atol=1e-4, rtol=0)): + raise ValueError("The probability vectors do not look like phased BGEN probabilities.") + + safe_probabilities = np.where(finite, haplotype_probabilities, -np.inf) + calls = np.argmax(safe_probabilities, axis=-1).astype(np.int8) + calls[~fully_finite] = -1 + return calls + + +def _phased_calls_or_skip(probabilities, reader_name): + """Return phased calls, omitting readers that do not preserve BGEN phase.""" + try: + return _probabilities_to_phased_calls(probabilities) + except ValueError as exc: + pytest.skip(f"{reader_name} does not expose phased BGEN probabilities: {exc}") + + def read_bgen_snputils(path, genotype_mode="dosage"): """Read BGEN file using snputils""" from snputils.snp.io.read.bgen import BGENReader @@ -57,8 +104,11 @@ def read_bgen_snputils(path, genotype_mode="dosage"): reader = BGENReader(path) if genotype_mode == "dosage": return reader.read_dosage() - snpobj = reader.read(fields=["GP"]) - return snpobj.calldata_gp.astype(np.float32, copy=False) + if genotype_mode == "probabilities": + return reader.read(fields=["GP"]).calldata_gp.astype(np.float32, copy=False) + if genotype_mode == "phased": + return reader.read(fields=["GT"], genotype_mode="phased").genotypes + raise ValueError(f"Unsupported BGEN benchmark genotype mode: {genotype_mode}") def read_bgen_bgen(path, genotype_mode="dosage"): @@ -66,25 +116,43 @@ def read_bgen_bgen(path, genotype_mode="dosage"): from bgen import BgenReader sample_path = str(Path(path).with_suffix(".sample")) if Path(path).with_suffix(".sample").exists() else "" with BgenReader(str(path), sample_path, delay_parsing=True) as bfile: - first_probabilities = np.asarray(bfile[0].probabilities, dtype=np.float32) n_variants = len(bfile) - n_samples, width = first_probabilities.shape + variants = iter(bfile) + first_variant = next(variants) + first_probabilities = np.asarray(first_variant.probabilities, dtype=np.float32) + n_samples = first_probabilities.shape[0] if genotype_mode == "dosage": out = np.empty((n_variants, n_samples), dtype=np.float32) + out[0] = _probabilities_to_dosage(first_probabilities) + elif genotype_mode == "probabilities": + out = np.empty( + (n_variants, n_samples, first_probabilities.shape[1]), + dtype=np.float32, + ) + out[0] = first_probabilities + elif genotype_mode == "phased": + out = np.empty((n_variants, n_samples, 2), dtype=np.int8) + out[0] = _probabilities_to_phased_calls(first_probabilities) else: - out = np.empty((n_variants, n_samples, width), dtype=np.float32) + raise ValueError(f"Unsupported BGEN benchmark genotype mode: {genotype_mode}") - for i, variant in enumerate(bfile): + for i, variant in enumerate(variants, start=1): probabilities = np.asarray(variant.probabilities, dtype=np.float32) if genotype_mode == "dosage": out[i] = _probabilities_to_dosage(probabilities) - else: + elif genotype_mode == "probabilities": if probabilities.shape[1] > out.shape[2]: - expanded = np.full((n_variants, n_samples, probabilities.shape[1]), np.nan, dtype=np.float32) + expanded = np.full( + (n_variants, n_samples, probabilities.shape[1]), + np.nan, + dtype=np.float32, + ) expanded[:, :, : out.shape[2]] = out out = expanded out[i].fill(np.nan) out[i, :, : probabilities.shape[1]] = probabilities + else: + out[i] = _probabilities_to_phased_calls(probabilities) return out @@ -93,8 +161,12 @@ def read_bgen_pysnptools(path, genotype_mode="dosage"): from pysnptools.distreader import Bgen probabilities = Bgen(str(path)).read(order="C", dtype=np.float32).val.transpose(1, 0, 2) if genotype_mode == "dosage": - return probabilities @ np.array([0.0, 1.0, 2.0], dtype=np.float32) - return probabilities + return _probabilities_to_dosage(probabilities) + if genotype_mode == "probabilities": + return probabilities + if genotype_mode == "phased": + return _phased_calls_or_skip(probabilities, "pysnptools") + raise ValueError(f"Unsupported BGEN benchmark genotype mode: {genotype_mode}") def read_bgen_sgkit(path, genotype_mode="dosage"): @@ -112,11 +184,18 @@ def read_bgen_sgkit(path, genotype_mode="dosage"): probabilities = ds["call_genotype_probability"] if "call_genotype_probability_mask" in ds: probabilities = probabilities.where(~ds["call_genotype_probability_mask"]) - return probabilities.compute().values.astype(np.float32, copy=False) + probabilities = probabilities.compute().values.astype(np.float32, copy=False) + if genotype_mode == "probabilities": + return probabilities + if genotype_mode == "phased": + return _phased_calls_or_skip(probabilities, "sgkit") + raise ValueError(f"Unsupported BGEN benchmark genotype mode: {genotype_mode}") def read_bgen_hail(path, genotype_mode="dosage"): """Read BGEN file using hail""" + if genotype_mode == "phased": + pytest.skip("Hail does not preserve phase when importing BGEN files.") import os import hail as hl spark_memory = os.environ.get("HAIL_SPARK_MEMORY", "192g") @@ -144,11 +223,19 @@ def read_bgen_hail(path, genotype_mode="dosage"): ) n_samples = mt.count_cols() if genotype_mode == "dosage": - return np.array(hl.or_else(mt.dosage, float("nan")).collect(), dtype=np.float32).reshape((-1, n_samples)) - return np.array( - hl.or_else(mt.GP, hl.array([float("nan"), float("nan"), float("nan")])).collect(), - dtype=np.float32, - ).reshape((-1, n_samples, 3)) + return np.array( + hl.or_else(mt.dosage, float("nan")).collect(), + dtype=np.float32, + ).reshape((-1, n_samples)) + if genotype_mode == "probabilities": + return np.array( + hl.or_else( + mt.GP, + hl.array([float("nan"), float("nan"), float("nan")]), + ).collect(), + dtype=np.float32, + ).reshape((-1, n_samples, 3)) + raise ValueError(f"Unsupported BGEN benchmark genotype mode: {genotype_mode}") finally: hl.stop() @@ -181,9 +268,9 @@ def test_bgen_readers(benchmark, reader, name, path, memory_profile, reader_name ref_array, memory_profile, genotype_mode=genotype_mode, - ref_reader_func=read_bgen_snputils, - assert_allclose=True, - atol=1 / 255 + 1e-6, - equal_nan=True, - verify=not (name == "snputils" and not memory_profile), + ref_reader_func=read_bgen_bgen, + assert_allclose=genotype_mode in {"dosage", "probabilities"}, + atol=(1 / 255 + 1e-6) if genotype_mode in {"dosage", "probabilities"} else 0.0, + equal_nan=genotype_mode in {"dosage", "probabilities"}, + comparison_chunk_size=(1024 if genotype_mode in {"dosage", "probabilities"} else None), ) diff --git a/benchmark/readers_benchmark.pdf b/benchmark/readers_benchmark.pdf index 3e714b0e..ab3967a5 100644 Binary files a/benchmark/readers_benchmark.pdf and b/benchmark/readers_benchmark.pdf differ diff --git a/benchmark/readers_benchmark.png b/benchmark/readers_benchmark.png index d2f2c795..903dd02f 100644 Binary files a/benchmark/readers_benchmark.png and b/benchmark/readers_benchmark.png differ diff --git a/benchmark/test_read_bgen.py b/benchmark/test_read_bgen.py new file mode 100644 index 00000000..f731fd37 --- /dev/null +++ b/benchmark/test_read_bgen.py @@ -0,0 +1,121 @@ +import numpy as np +import pytest + +from .read_bgen import ( + _probabilities_to_dosage, + _probabilities_to_phased_calls, + read_bgen_bgen, + read_bgen_snputils, +) + + +def test_probabilities_to_dosage_handles_unphased_probabilities(): + probabilities = np.array( + [ + [[0.8, 0.1, 0.1], [0.1, 0.2, 0.7]], + [[np.nan, np.nan, np.nan], [0.0, 1.0, 0.0]], + ], + dtype=np.float32, + ) + + observed = _probabilities_to_dosage(probabilities) + + expected = np.array([[0.3, 1.6], [np.nan, 1.0]], dtype=np.float32) + np.testing.assert_allclose(observed, expected, equal_nan=True) + + +def test_probabilities_to_dosage_handles_phased_probabilities(): + probabilities = np.array( + [ + [[0.9, 0.1, 0.2, 0.8], [0.4, 0.6, 0.7, 0.3]], + [[np.nan, np.nan, np.nan, np.nan], [1.0, 0.0, 0.0, 1.0]], + ], + dtype=np.float32, + ) + + observed = _probabilities_to_dosage(probabilities) + + expected = np.array([[0.9, 0.9], [np.nan, 1.0]], dtype=np.float32) + np.testing.assert_allclose(observed, expected, equal_nan=True) + + +def test_probabilities_to_phased_calls_materializes_allele_codes(): + probabilities = np.array( + [ + [ + [0.9, 0.1, 0.2, 0.8], + [np.nan, np.nan, np.nan, np.nan], + ], + [ + [0.2, 0.8, 0.7, 0.3], + [0.4, 0.6, 0.9, 0.1], + ], + ], + dtype=np.float32, + ) + + observed = _probabilities_to_phased_calls(probabilities) + + expected = np.array( + [ + [[0, 1], [-1, -1]], + [[1, 0], [1, 0]], + ], + dtype=np.int8, + ) + np.testing.assert_array_equal(observed, expected) + assert observed.shape == (2, 2, 2) + assert observed.dtype == np.int8 + + +def test_probabilities_to_phased_calls_rejects_unphased_probabilities(): + probabilities = np.array([[0.1, 0.8, 0.1]], dtype=np.float32) + + with pytest.raises(ValueError, match="two equally sized haplotype"): + _probabilities_to_phased_calls(probabilities) + + +def test_probabilities_to_phased_calls_rejects_partial_ploidy(): + probabilities = np.array([[0.1, 0.9, np.nan, np.nan]], dtype=np.float32) + + with pytest.raises(ValueError, match="requires diploid"): + _probabilities_to_phased_calls(probabilities) + + +def test_bgen_package_reference_matches_snputils(tmp_path): + pytest.importorskip("bgen") + from snputils import BGENWriter, SNPObject + + genotypes = np.array( + [ + [[0, 0], [0, 1], [1, 0]], + [[1, 1], [1, 0], [0, 1]], + ], + dtype=np.int8, + ) + snpobj = SNPObject( + genotypes=genotypes, + samples=np.array(["s1", "s2", "s3"], dtype=object), + variants_ref=np.array(["A", "C"], dtype=object), + variants_alt=np.array(["G", "T"], dtype=object), + variants_chrom=np.array(["1", "1"], dtype=object), + variants_id=np.array(["rs1", "rs2"], dtype=object), + variants_pos=np.array([100, 200]), + ) + path = tmp_path / "phased.bgen" + BGENWriter(snpobj, path).write(phased=True, bit_depth=16) + + np.testing.assert_allclose( + read_bgen_snputils(path, genotype_mode="dosage"), + read_bgen_bgen(path, genotype_mode="dosage"), + atol=1 / 65535 + 1e-6, + ) + np.testing.assert_allclose( + read_bgen_snputils(path, genotype_mode="probabilities"), + read_bgen_bgen(path, genotype_mode="probabilities"), + atol=1 / 65535 + 1e-6, + ) + np.testing.assert_array_equal( + read_bgen_snputils(path, genotype_mode="phased"), + read_bgen_bgen(path, genotype_mode="phased"), + ) diff --git a/benchmark/utils.py b/benchmark/utils.py index 050b351a..407c92b2 100644 --- a/benchmark/utils.py +++ b/benchmark/utils.py @@ -16,6 +16,7 @@ def create_benchmark_test( atol=0.0, equal_nan=False, verify=True, + comparison_chunk_size=None, ): timed_reader = functools.partial(reader_func, file_path, genotype_mode=genotype_mode) @@ -68,13 +69,29 @@ def read_file(): raise ValueError("ref_reader_func is required when ref_array is None.") ref_array = ref_reader_func(file_path, genotype_mode=genotype_mode) if assert_allclose: - np.testing.assert_allclose( - result, - ref_array, - rtol=rtol, - atol=atol, - equal_nan=equal_nan, - err_msg=f"Output does not match reference for {name}", - ) + if comparison_chunk_size is None: + chunks = ((result, ref_array),) + else: + if result.shape != ref_array.shape: + raise AssertionError( + f"Output shape {result.shape} does not match reference shape " + f"{ref_array.shape} for {name}" + ) + chunks = ( + ( + result[start : start + comparison_chunk_size], + ref_array[start : start + comparison_chunk_size], + ) + for start in range(0, result.shape[0], comparison_chunk_size) + ) + for result_chunk, reference_chunk in chunks: + np.testing.assert_allclose( + result_chunk, + reference_chunk, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + err_msg=f"Output does not match reference for {name}", + ) else: assert np.array_equal(result, ref_array), f"Output does not match reference for {name}" diff --git a/benchmark/write_timeout_result.py b/benchmark/write_timeout_result.py index 8fa0867e..0e853639 100644 --- a/benchmark/write_timeout_result.py +++ b/benchmark/write_timeout_result.py @@ -72,7 +72,11 @@ def main() -> None: parser.add_argument("--format", dest="benchmark_format", required=True) parser.add_argument("--reader", required=True) parser.add_argument("--path", required=True) - parser.add_argument("--genotype-mode", default="dosage", choices=("dosage", "phased")) + parser.add_argument( + "--genotype-mode", + default="dosage", + choices=("dosage", "probabilities", "phased"), + ) parser.add_argument("--timeout-seconds", type=float, required=True) args = parser.parse_args() diff --git a/docs/_extra/robots.txt b/docs/_extra/robots.txt new file mode 100644 index 00000000..62b5e81c --- /dev/null +++ b/docs/_extra/robots.txt @@ -0,0 +1,4 @@ +User-agent: * +Allow: / + +Sitemap: https://docs.snputils.org/sitemap.xml diff --git a/docs/_templates/base.html b/docs/_templates/base.html new file mode 100644 index 00000000..990c5fa9 --- /dev/null +++ b/docs/_templates/base.html @@ -0,0 +1,11 @@ +{% extends "!base.html" %} + +{% block htmltitle %} + {% if not docstitle %} + {{ title|striptags|e }} + {% elif pagename == master_doc %} + {{ title|striptags|e }} + {% else %} + {{ title|striptags|e }} - {{ docstitle|striptags|e }} + {% endif %} +{% endblock %} diff --git a/docs/api/datasets.md b/docs/api/datasets.md index 9475da70..0d30817c 100644 --- a/docs/api/datasets.md +++ b/docs/api/datasets.md @@ -1,6 +1,8 @@ # Datasets -Helpers for discovering bundled example datasets and generating synthetic data for tests, tutorials, and quick experiments. Synthetic builders return fully populated objects matching the shapes described in {doc}`../user_guide/data-model`. +Helpers for discovering registered downloadable datasets and generating synthetic data for tests, tutorials, and quick experiments. Synthetic builders return fully populated objects matching the shapes described in {doc}`../user_guide/data-model`. + +{func}`snputils.load_dataset` downloads and caches the selected remote resources on demand. ## Example Datasets @@ -12,6 +14,10 @@ expanded 1000 Genomes Project cohort including 602 trios". Because these VCFs ar large, `load_dataset("1kgp", resource="high_coverage_2022")` defaults to chromosome 1; pass `chromosomes=` explicitly to load additional chromosomes. +All 1KGP resources are chromosome-scale downloads. Always choose `resource=` and +`chromosomes=` deliberately. A bare `load_dataset("1kgp")` uses the default resource +and its full default chromosome set, which can download a very large amount of data. + ```{eval-rst} .. autofunction:: snputils.available_datasets_list ``` diff --git a/docs/conf.py b/docs/conf.py index 619d1e27..a5936962 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -37,6 +37,7 @@ "sphinx.ext.viewcode", "sphinx_copybutton", "sphinx_design", + "sphinx_sitemap", ] source_suffix = { @@ -108,12 +109,21 @@ templates_path = ["_templates"] html_theme = "furo" html_title = "snputils" +html_baseurl = "https://docs.snputils.org/" html_logo = "../assets/logo.png" html_favicon = "https://snputils.org/su_favicon.png" html_static_path = ["_static"] html_css_files = ["custom.css"] html_extra_path = ["_extra"] html_show_sourcelink = True +sitemap_url_scheme = "{link}" +sitemap_excludes = [ + "search.html", + "genindex.html", + "py-modindex.html", + "_modules/*", +] +sitemap_indent = 2 html_theme_options = { "source_repository": "https://github.com/AI-sandbox/snputils/", "source_branch": "main", diff --git a/docs/contributing.md b/docs/contributing.md index 3198f4fb..c631ae9a 100644 --- a/docs/contributing.md +++ b/docs/contributing.md @@ -27,7 +27,8 @@ The CI workflow runs the test suite on supported Python versions before releases ## Building documentation ```bash -sphinx-build -b html docs docs/_build/html +sphinx-build -W --keep-going -b html docs docs/_build/html +sphinx-build -b linkcheck docs docs/_build/linkcheck ``` See {doc}`installation` for details. User-facing docs live in `docs/`; API pages use Sphinx autodoc against the Python sources. diff --git a/docs/index.md b/docs/index.md index 95d9f480..f6d60539 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,6 +1,12 @@ -# snputils +--- +myst: + html_meta: + description: snputils is a Python toolkit for reading, processing, analyzing, and visualizing genotype, ancestry, phenotype, and population-genetics data. +--- -**snputils** is a Python library for processing genetic variation, ancestry, phenotype, and relatedness data. It focuses on fast file I/O, ergonomic data containers, population-genetic statistics, dimensionality reduction, and visualization workflows for genomic analyses. +# snputils: Python toolkit for genotype, ancestry, and population-genetics workflows + +**snputils** is a Python toolkit for reading, processing, analyzing, and visualizing genotype, ancestry, phenotype, and relatedness data. It focuses on fast file I/O, ergonomic data containers, population-genetic statistics, dimensionality reduction, and visualization workflows for genomic analyses. Developed in collaboration between Stanford University's Department of Biomedical Data Science, UC Santa Cruz Genomics Institute, and collaborators worldwide. diff --git a/docs/quickstart.md b/docs/quickstart.md index 65843a36..5f0e10e9 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -1,6 +1,27 @@ # Quickstart -The top-level package exposes the most common readers, data containers, and analysis helpers. A typical workflow loads genotypes, optional ancestry and phenotypes, runs analyses, and plots results: +Start with a fully self-contained synthetic dataset—no external files or downloads required: + +```python +import snputils as su + +snp = su.build_synthetic_snp_dataset( + n_samples=30, + n_snps=100, + seed=42, +) + +af = snp.allele_freq() +pcs = su.PCA(n_components=2).fit_transform(snp) + +print(snp.n_samples, snp.n_snps) +print(af[:5]) +print(pcs.shape) +``` + +## Working with your own files + +The top-level package exposes the most common readers, data containers, and analysis helpers. A typical file-backed workflow loads genotypes, optional ancestry and phenotypes, runs analyses, and plots results: ```python import snputils as su diff --git a/docs/tutorials/admixture-mapping.ipynb b/docs/tutorials/admixture-mapping.ipynb index 477285ba..69e243a2 100644 --- a/docs/tutorials/admixture-mapping.ipynb +++ b/docs/tutorials/admixture-mapping.ipynb @@ -5,9 +5,9 @@ "id": "e0370902", "metadata": {}, "source": [ - "# Admixture Mapping\n", + "# Run admixture mapping from local ancestry\n", "\n", - "Admixture mapping tests whether local ancestry at a genomic region is associated with a phenotype in recently admixed individuals. Instead of asking whether a specific SNP allele is associated with a trait, we ask whether carrying ancestry from a source population at a local genomic window changes trait values after accounting for background ancestry and other covariates.\n", + "Use snputils to run admixture mapping from local ancestry and phenotype data. Admixture mapping tests whether local ancestry at a genomic region is associated with a phenotype in recently admixed individuals after accounting for background ancestry and other covariates.\n", "\n", "In this tutorial, we build a synthetic local ancestry dataset with a quantitative trait, run `snputils.tools.run_admixture_mapping`, and make a Manhattan plot from the resulting association table." ] @@ -1140,6 +1140,11 @@ } ], "metadata": { + "myst": { + "html_meta": { + "description": "Run admixture mapping from local ancestry and phenotype data in Python with snputils." + } + }, "kernelspec": { "display_name": ".venv", "language": "python", diff --git a/docs/tutorials/f-statistics.ipynb b/docs/tutorials/f-statistics.ipynb index 4f7e41e0..a00a6a49 100644 --- a/docs/tutorials/f-statistics.ipynb +++ b/docs/tutorials/f-statistics.ipynb @@ -5,9 +5,9 @@ "id": "c836f712", "metadata": {}, "source": [ - "# F-statistics and population statistics\n", + "# Calculate FST and f-statistics in Python\n", "\n", - "F-statistics summarize allele-frequency covariance among populations. They are useful for checking population separation, testing whether a group is consistent with a mixture of two references, and measuring whether two population pairs share excess genetic drift.\n", + "Use snputils to calculate FST and f-statistics in Python. These statistics summarize allele-frequency covariance among populations and help assess population separation, admixture, and excess shared genetic drift.\n", "\n", "This tutorial shows how to:\n", "\n", @@ -1353,6 +1353,11 @@ } ], "metadata": { + "myst": { + "html_meta": { + "description": "Calculate FST, f2, f3, f4, f4-ratios, and D-statistics in Python with snputils." + } + }, "kernelspec": { "display_name": "Python 3", "language": "python", diff --git a/docs/tutorials/index.md b/docs/tutorials/index.md index b7ce3864..0caf4376 100644 --- a/docs/tutorials/index.md +++ b/docs/tutorials/index.md @@ -1,6 +1,6 @@ # Tutorials -End-to-end notebooks for data objects, population structure, ancestry visualization, association testing, and GRG workflows. Examples use synthetic or bundled datasets unless a notebook states otherwise. +End-to-end notebooks for data objects, population structure, ancestry visualization, association testing, and GRG workflows. Examples use synthetic or registered downloadable datasets unless a notebook states otherwise. ::::{grid} 1 1 2 2 :gutter: 2 @@ -72,13 +72,13 @@ Work with genotype representation graph data. :maxdepth: 1 snp-object -pca +PCA snp-allele-frequency -f-statistics -local-ancestry-visualization +F-statistics +Local ancestry visualization mdpca maasmds -admixture-mapping +Admixture mapping phenotype-object grg ``` diff --git a/docs/tutorials/local-ancestry-visualization.ipynb b/docs/tutorials/local-ancestry-visualization.ipynb index 609af30d..5fec90aa 100644 --- a/docs/tutorials/local-ancestry-visualization.ipynb +++ b/docs/tutorials/local-ancestry-visualization.ipynb @@ -5,9 +5,9 @@ "id": "66e392a8", "metadata": {}, "source": [ - "# Local Ancestry Visualization\n", + "# Local ancestry analysis and visualization in Python\n", "\n", - "Local ancestry can be viewed at two levels: chromosome paintings for individual genomes, and stacked cohort plots for many samples. This tutorial starts with full-chromosome paintings, then uses `plot_lai` for a dataset-level view." + "Use snputils for local ancestry analysis and visualization in Python at two levels: chromosome paintings for individual genomes and stacked cohort plots for many samples. This tutorial starts with full-chromosome paintings, then uses `plot_lai` for a dataset-level view." ] }, { @@ -271,6 +271,11 @@ } ], "metadata": { + "myst": { + "html_meta": { + "description": "Analyze and visualize local ancestry in Python with snputils using chromosome paintings and cohort-level plots." + } + }, "kernelspec": { "display_name": ".venv", "language": "python", diff --git a/docs/tutorials/pca.ipynb b/docs/tutorials/pca.ipynb index 3033d039..502c704c 100644 --- a/docs/tutorials/pca.ipynb +++ b/docs/tutorials/pca.ipynb @@ -5,9 +5,9 @@ "id": "a21f983a", "metadata": {}, "source": [ - "# Principal Component Analysis (PCA)\n", + "# Run population-genetics PCA in Python\n", "\n", - "Principal Component Analysis (PCA) summarizes genome-wide genotype variation into a small number of axes called principal components (PCs). In population genetics, the first few PCs often capture broad ancestry structure, sample outliers, or batch effects.\n", + "Use snputils to run population-genetics principal component analysis (PCA) in Python. PCA summarizes genome-wide genotype variation into a small number of axes called principal components (PCs); the first few PCs often capture broad ancestry structure, sample outliers, or batch effects.\n", "\n", "This tutorial shows how to:\n", "\n", @@ -1276,6 +1276,11 @@ } ], "metadata": { + "myst": { + "html_meta": { + "description": "Run population-genetics principal component analysis (PCA) in Python with snputils using CPU or GPU implementations." + } + }, "kernelspec": { "display_name": "Python (snputils)", "language": "python", diff --git a/docs/user_guide/analysis.md b/docs/user_guide/analysis.md index 7d431b57..e015c9a3 100644 --- a/docs/user_guide/analysis.md +++ b/docs/user_guide/analysis.md @@ -1,4 +1,12 @@ -# Analysis +--- +myst: + html_meta: + description: Run PCA, calculate allele frequencies and f-statistics, perform GWAS and admixture mapping, and simulate genomic data in Python with snputils. +--- + +# Population-genetics analysis in Python + +Use snputils for population-genetics workflows including dimensionality reduction, allele frequencies, f-statistics, association analysis, simulation, and command-line execution. ## Dimensionality Reduction diff --git a/docs/user_guide/file-io.md b/docs/user_guide/file-io.md index fcf4b992..06458676 100644 --- a/docs/user_guide/file-io.md +++ b/docs/user_guide/file-io.md @@ -1,6 +1,12 @@ -# File I/O +--- +myst: + html_meta: + description: Read and write VCF, BCF, BGEN, PLINK BED, and PLINK PGEN genomic files in Python with snputils. +--- -High-level readers dispatch to the right implementation based on file extension. +# Read and write VCF, BCF, BGEN, BED, and PGEN files in Python + +Use the high-level snputils readers and writers to work with common genomic file formats in Python. Readers dispatch to the right implementation based on file extension. ```python import snputils as su @@ -180,14 +186,11 @@ sample2 sample2 49 2 In Python you can pass the same data as a {class}`~snputils.CovariateObject` instead of a path. -## Datasets +## Dataset registry and synthetic data -Built-in datasets and synthetic data builders for quick experimentation: +Synthetic data builders provide fully populated objects for quick experimentation without external files or downloads: ```python -su.available_datasets_list() # list bundled datasets -ds = su.load_dataset("1kgp") # returns a SNPObject - snpobj = su.build_synthetic_snp_dataset() laiobj = su.build_synthetic_chromosome_painting_dataset() admobj = su.build_synthetic_admixture_dataset() @@ -196,3 +199,5 @@ grg = su.build_synthetic_grg() mdpca_ds = su.build_synthetic_mdpca_dataset() maas_ds = su.build_synthetic_maasmds_dataset() ``` + +`su.available_datasets_list()` lists registered downloadable datasets. Their files are fetched from remote sources and cached on demand; they are not included in the `snputils` package. The registered 1KGP resources are large, so choose `resource=` and `chromosomes=` explicitly when calling `su.load_dataset("1kgp", ...)`. See {doc}`../api/datasets` for details. diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index 3e123634..dae67266 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -5,8 +5,8 @@ The user guide groups the library by workflow rather than by module path. Each p ## Recommended reading order 1. **{doc}`data-model`** — Object shapes, attributes, and filtering before you load real data. -2. **{doc}`file-io`** — Format dispatch, reader/writer options, and synthetic datasets for experiments. -3. **{doc}`analysis`** — Dimensionality reduction, allele frequencies, *f*-statistics, GWAS, admixture mapping, simulation, and the CLI. +2. **{doc}`File I/O `** — Format dispatch, reader/writer options, and synthetic datasets for experiments. +3. **{doc}`Analysis `** — Dimensionality reduction, allele frequencies, *f*-statistics, GWAS, admixture mapping, simulation, and the CLI. 4. **{doc}`visualization`** — Scatter, local ancestry, admixture, Manhattan, and Q–Q plots. After the quickstart, open the {doc}`../tutorials/index` notebooks for end-to-end examples with saved outputs. @@ -44,7 +44,7 @@ Embedding scatter plots, chromosome painting, admixture bars, Manhattan, and Q :maxdepth: 1 data-model -file-io -analysis +File I/O +Analysis visualization ``` diff --git a/pyproject.toml b/pyproject.toml index ee340016..e2b77e43 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,15 +5,17 @@ build-backend = "setuptools.build_meta" [project] name = "snputils" dynamic = ["version"] # Version is set by setuptools_scm and determined by the latest git tag -description = "Process genomes with ease" +description = "Python toolkit for fast genotype I/O, ancestry analysis, population-genetics statistics, genotype QC, GWAS, and visualization" readme = "README.md" requires-python = ">=3.9" license = {text = "BSD 3-Clause License"} +keywords = ["bioinformatics", "genomics", "population genetics", "genotype", "ancestry", "VCF", "BCF", "BGEN", "PLINK", "PGEN", "GWAS"] classifiers = [ "Development Status :: 5 - Production/Stable", "Intended Audience :: Science/Research", "Intended Audience :: Developers", "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Bio-Informatics", "License :: OSI Approved :: BSD License", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", @@ -52,6 +54,7 @@ docs = [ "myst-nb>=1.0", "sphinx-copybutton>=0.5", "sphinx-design>=0.6", + "sphinx-sitemap>=2.8", ] benchmark = ["pytest", "pytest-benchmark", "memory_profiler", "pandas-plink", "pysam", "scikit-allel", "sgkit[plink,bgen]", "zarr<3", "hail>=0.2.130", "pysnptools[bgen]", "bgen>=1.9.9", "Pgenlib", "cyvcf2", "plinkio", "PyVCF3"] @@ -60,6 +63,9 @@ grg = ["pygrgl>=2.9"] [project.urls] "Homepage" = "https://snputils.org" "Documentation" = "https://docs.snputils.org" +"Tutorials" = "https://docs.snputils.org/tutorials/index.html" +"Changelog" = "https://docs.snputils.org/changelog.html" +"Paper" = "https://doi.org/10.64898/2026.02.28.708618" "Source Code" = "https://github.com/AI-sandbox/snputils" "Issue Tracker" = "https://github.com/AI-sandbox/snputils/issues" @@ -89,6 +95,11 @@ libraries = ["z"] [tool.setuptools_scm] +[tool.cibuildwheel] +build = "cp39-* cp310-* cp311-* cp312-* cp313-* cp314-*" +skip = "*-musllinux_* *-win*" +test-command = "python {project}/.github/scripts/wheel_smoke_test.py" + [tool.mypy] python_version = "3.9" ignore_missing_imports = true diff --git a/snputils/datasets/__test__/test_load_dataset.py b/snputils/datasets/__test__/test_load_dataset.py index 09f5b591..47ef1296 100644 --- a/snputils/datasets/__test__/test_load_dataset.py +++ b/snputils/datasets/__test__/test_load_dataset.py @@ -54,6 +54,51 @@ def _maf_chunk() -> SNPObject: ) +def test_load_dataset_streaming_dosage_with_require_biallelic_skips_multiallelic(tmp_path): + """Dosage parsing rejects multiallelic sites; require_biallelic must filter first.""" + load_dataset_module = importlib.import_module("snputils.datasets.load_dataset") + vcf_path = tmp_path / "mixed.vcf" + vcf_path.write_text( + "##fileformat=VCFv4.2\n" + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ts1\ts2\n" + "1\t100\trs_multi\tA\tG,T\t.\tPASS\t.\tGT\t0|1\t1|2\n" + "1\t200\trs_bi\tC\tT\t.\tPASS\t.\tGT\t0|1\t1|1\n" + "1\t300\trs_indel\tA\tAT\t.\tPASS\t.\tGT\t0|0\t0|1\n" + ) + + snpobj = load_dataset_module.load_dataset( + "1kgp", + genotype_sources=[vcf_path], + download_genotypes=False, + max_variants=2, + require_biallelic=True, + verbose=False, + ) + + assert snpobj.variants_id.tolist() == ["rs_bi", "rs_indel"] + assert snpobj.is_dosage is True + np.testing.assert_array_equal(snpobj.genotypes, np.array([[1, 2], [0, 1]], dtype=np.int8)) + + +def test_load_dataset_streaming_dosage_without_biallelic_filter_still_rejects_multiallelic(tmp_path): + load_dataset_module = importlib.import_module("snputils.datasets.load_dataset") + vcf_path = tmp_path / "multi.vcf" + vcf_path.write_text( + "##fileformat=VCFv4.2\n" + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\ts1\n" + "1\t100\trs_multi\tA\tG,T\t.\tPASS\t.\tGT\t0|1\n" + ) + + with pytest.raises(ValueError, match="dosage.*biallelic"): + load_dataset_module.load_dataset( + "1kgp", + genotype_sources=[vcf_path], + download_genotypes=False, + max_variants=1, + verbose=False, + ) + + def test_load_dataset_streaming_honors_variant_ids_with_max_variants(tmp_path, monkeypatch): load_dataset_module = importlib.import_module("snputils.datasets.load_dataset") diff --git a/snputils/datasets/load_dataset.py b/snputils/datasets/load_dataset.py index 91a16f30..43a31410 100644 --- a/snputils/datasets/load_dataset.py +++ b/snputils/datasets/load_dataset.py @@ -169,7 +169,8 @@ def load_dataset( maf: Optional minor allele frequency threshold. In streaming reads, this is applied before ``max_variants`` truncation. require_biallelic: When ``max_variants`` is set and source files are streamed, keep only variants with - exactly one REF allele and one ALT allele. + exactly one REF allele and one ALT allele. With the streaming default ``genotype_mode="dosage"``, + multiallelic sites are filtered before dosage conversion. require_complete: When ``max_variants`` is set and source files are streamed, keep only variants with no missing genotype calls across the selected samples. require_variable_genotypes: When ``max_variants`` is set and source files are streamed, keep only variants that @@ -369,7 +370,12 @@ def _read_snp_subset( chunks: list[SNPObject] = [] n_variants = 0 reader = SNPReader(source) - iter_kwargs = {"genotype_mode": genotype_mode, "chunk_size": 50_000} + stream_mode = genotype_mode + convert_to_dosage = False + if genotype_mode == "dosage" and (require_biallelic or snv_only): + stream_mode = "phased" + convert_to_dosage = True + iter_kwargs = {"genotype_mode": stream_mode, "chunk_size": 50_000} if selected is not None: iter_kwargs["sample_ids"] = np.asarray(selected, dtype=object) if selected_variants is not None: @@ -387,6 +393,8 @@ def _read_snp_subset( chunk = chunk.filter_maf(maf=maf) if chunk.n_snps == 0: continue + if convert_to_dosage: + chunk = chunk.to_dosage() remaining = max_variants - n_variants if chunk.n_snps > remaining: chunk = chunk.filter_variants(indexes=np.arange(remaining), include=True) diff --git a/snputils/simulation/__test__/test_simulator.py b/snputils/simulation/__test__/test_simulator.py new file mode 100644 index 00000000..75e2809a --- /dev/null +++ b/snputils/simulation/__test__/test_simulator.py @@ -0,0 +1,29 @@ +import numpy as np +import pytest + +torch = pytest.importorskip("torch") +if not hasattr(torch, "Tensor"): + pytest.skip("PyTorch is required for simulator tests.", allow_module_level=True) + +from snputils.simulation.simulator import OnlineSimulator + + +def test_map_breakpoints_use_poisson_crossover_counts(monkeypatch): + simulator = OnlineSimulator.__new__(OnlineSimulator) + simulator.rate_per_snp = np.array([0.001, 0.002, 0.003]) + observed = {} + + def poisson(lam): + observed["lam"] = lam + return np.array([1, 2, 3]) + + monkeypatch.setattr(np.random, "poisson", poisson) + + split_points = simulator._draw_split_points( + n_snps=4, + num_generation_max=10, + num_generations=10, + ) + + np.testing.assert_allclose(observed["lam"], 10 * simulator.rate_per_snp) + np.testing.assert_array_equal(split_points, [1, 3]) diff --git a/snputils/simulation/__test__/test_simulator_cli.py b/snputils/simulation/__test__/test_simulator_cli.py new file mode 100644 index 00000000..3d65fe46 --- /dev/null +++ b/snputils/simulation/__test__/test_simulator_cli.py @@ -0,0 +1,186 @@ +import argparse +import pickle +import sys +import types + +import numpy as np +import pytest + +import snputils.simulation.simulator_cli as simulator_cli +from snputils.simulation.simulator_cli import ( + _batch_genotypes_to_snp_layout, + _infer_output_format, + _output_genotype_path, + _write_batch_truth, + _write_genotypes_like_input, +) +from snputils.snp.genobj.snpobj import SNPObject +from snputils.snp.io.read.bcf import BCFReader + + +def _source_snpobject() -> SNPObject: + return SNPObject( + genotypes=np.array( + [ + [[0, 1], [1, 1]], + [[1, 0], [0, 0]], + ], + dtype=np.int8, + ), + samples=np.array(["founder1", "founder2"], dtype=object), + variants_ref=np.array(["A", "C"], dtype=object), + variants_alt=np.array(["G", "T"], dtype=object), + variants_chrom=np.array(["1", "1"], dtype=object), + variants_id=np.array(["rs1", "rs2"], dtype=object), + variants_pos=np.array([10, 20]), + ) + + +def test_simulator_output_paths_include_bcf_pickle_and_npz(tmp_path): + prefix = tmp_path / "simulated" + + assert _output_genotype_path(prefix, "bcf") == tmp_path / "simulated.bcf" + assert _output_genotype_path(prefix, "pkl") == tmp_path / "simulated.pkl" + assert _output_genotype_path(prefix, "npz") == tmp_path / "simulated.npz" + assert _infer_output_format("founders.bcf", "same") == "bcf" + + +def test_simulator_writes_phased_bcf_and_snpobject_pickle(tmp_path): + source = _source_snpobject() + simulated = np.array( + [ + [[0, 1], [1, 0]], + [[1, 1], [0, 0]], + ], + dtype=np.int8, + ) + samples = np.array(["SIM1", "SIM2"], dtype=object) + + bcf_path = _write_genotypes_like_input( + source, + simulated, + samples, + tmp_path / "cohort", + "bcf", + ) + observed_bcf = BCFReader(bcf_path).read(genotype_mode="phased") + np.testing.assert_array_equal(observed_bcf.genotypes, simulated) + np.testing.assert_array_equal(observed_bcf.samples, samples) + + pickle_path = _write_genotypes_like_input( + source, + simulated, + samples, + tmp_path / "cohort_object", + "pkl", + ) + with pickle_path.open("rb") as handle: + observed_pickle = pickle.load(handle) + np.testing.assert_array_equal(observed_pickle.genotypes, simulated) + np.testing.assert_array_equal(observed_pickle.samples, samples) + + +def test_batch_output_layout_and_truth_sidecar(tmp_path): + batch = np.array( + [ + [[0, 1, 0], [1, 1, 0]], + [[1, 0, 1], [0, 0, 1]], + ], + dtype=np.float32, + ) + + observed = _batch_genotypes_to_snp_layout(batch) + expected = np.transpose(batch, (2, 0, 1)).astype(np.int8) + np.testing.assert_array_equal(observed, expected) + + truth_path = tmp_path / "batch_0001.labels.npz" + _write_batch_truth( + truth_path, + labels_d=np.array([[[0, 1]], [[1, 0]]]), + labels_c=None, + cp=np.array([[[False, True]], [[True, False]]]), + population_names=np.array(["AFR", "EUR"]), + output_ploidy=2, + ) + + with np.load(truth_path) as truth: + np.testing.assert_array_equal(truth["population_names"], ["AFR", "EUR"]) + assert int(truth["output_ploidy"]) == 2 + assert truth["labels_c"].size == 0 + + +def test_batch_genotype_formats_require_diploid_layout(): + with pytest.raises(ValueError, match="Diploid batch output"): + _batch_genotypes_to_snp_layout(np.zeros((2, 3), dtype=np.int8)) + + +def test_batch_pickle_output_keeps_npz_truth_sidecar(tmp_path, monkeypatch): + source = _source_snpobject() + metadata_path = tmp_path / "founders.tsv" + metadata_path.write_text("Sample\tPopulation\nfounder1\tAFR\nfounder2\tEUR\n") + + class FakeReader: + def read(self, genotype_mode): + assert genotype_mode == "phased" + return source + + class FakeSimulator: + population_names = np.array(["AFR", "EUR"]) + + def __init__(self, **kwargs): + assert kwargs["snp_data"] is source + + def simulate(self, *, batch_size, **kwargs): + assert batch_size == 4 + return ( + np.array( + [ + [0, 1], + [1, 1], + [1, 0], + [0, 0], + ], + dtype=np.float32, + ), + None, + None, + None, + ) + + fake_simulator_module = types.ModuleType("snputils.simulation.simulator") + fake_simulator_module.OnlineSimulator = FakeSimulator + monkeypatch.setitem(sys.modules, "snputils.simulation.simulator", fake_simulator_module) + monkeypatch.setattr(simulator_cli, "SNPReader", lambda path: FakeReader()) + + args = argparse.Namespace( + ancestry_proportions=None, + batch_size=2, + device="cpu", + diploid_output=True, + expand_haplotypes=False, + fixed_generations=False, + genetic_map=None, + metadata=str(metadata_path), + n_batches=1, + n_individuals=None, + num_generations=10, + output_dir=str(tmp_path / "output"), + output_format="pkl", + output_prefix=None, + sample_prefix="SIM", + snp="founders.pgen", + store_latlon_as_nvec=False, + verbose=False, + window_size=1000, + ) + + assert simulator_cli.run_simulator_command(args) == 0 + + pickle_path = tmp_path / "output" / "batch_0001.pkl" + truth_path = tmp_path / "output" / "batch_0001.labels.npz" + assert pickle_path.exists() + assert truth_path.exists() + with pickle_path.open("rb") as handle: + observed = pickle.load(handle) + assert observed.genotypes.shape == (2, 2, 2) + np.testing.assert_array_equal(observed.samples, ["SIMB0001_000001", "SIMB0001_000002"]) diff --git a/snputils/simulation/simulator/simulator.py b/snputils/simulation/simulator/simulator.py index 251a74d8..39184bb5 100644 --- a/snputils/simulation/simulator/simulator.py +++ b/snputils/simulation/simulator/simulator.py @@ -388,7 +388,7 @@ def _intersect_snp_metadata(self): if self.genetic_map is not None: cm_interp = np.interp(self.snp_data.variants_pos, self.genetic_map['pos'], self.genetic_map['cM']) self.cm_per_snp = cm_interp - self.rate_per_snp = np.gradient(cm_interp/100.0) + self.rate_per_snp = np.diff(cm_interp/100.0) log.info(f"rate/snp shape = {self.rate_per_snp.shape}") else: self.cm_per_snp = getattr(self.snp_data, "variants_cm", None) @@ -524,7 +524,7 @@ def _simulate_from_pool( if batch_labels_continuous is not None: batch_labels_continuous = batch_labels_continuous.to(device) - B, D = batch_snps.shape + B, D = batch_snps.shape split_points = self._draw_split_points(D, num_generation_max, num_generations) for sp in split_points: @@ -552,8 +552,8 @@ def _draw_split_points(self, n_snps, num_generation_max, num_generations=None): return np.empty(0, dtype=np.int64) if self.rate_per_snp is not None: - switch = np.random.binomial(G, self.rate_per_snp) % 2 - split_points = np.flatnonzero(switch) + switch = np.random.poisson(G * self.rate_per_snp) % 2 + split_points = np.flatnonzero(switch) + 1 else: if n_snps <= 1: return np.empty(0, dtype=np.int64) diff --git a/snputils/simulation/simulator_cli.py b/snputils/simulation/simulator_cli.py index ffa0bbed..b368a857 100644 --- a/snputils/simulation/simulator_cli.py +++ b/snputils/simulation/simulator_cli.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import sys import argparse import logging @@ -7,17 +9,18 @@ from snputils.snp.io.read import SNPReader from snputils.simulation._validation import validate_phased_simulation_input -from snputils.simulation.simulator import OnlineSimulator logging.basicConfig(level=logging.INFO, format="%(asctime)s │ %(levelname)-8s │ %(message)s", datefmt="%Y-%m-%d %H:%M:%S") log = logging.getLogger("simulator_cli") +_OUTPUT_FORMAT_CHOICES = ("same", "npz", "pkl", "pgen", "vcf", "vcf.gz", "bcf", "bgen") + def add_simulator_arguments(p: argparse.ArgumentParser) -> None: p.add_argument("--snp", required=True, - help="Path to phased SNP input (VCF, PGEN, or BGEN fileset). PLINK1 BED is not supported because it cannot store phase.") + help="Path to phased SNP input (VCF, BCF, PGEN, or BGEN fileset). PLINK1 BED is not supported because it cannot store phase.") p.add_argument("--metadata", required=True, help="TSV/CSV file with at least Sample/IID and Population columns.") p.add_argument("--output-dir", required=True, @@ -25,8 +28,8 @@ def add_simulator_arguments(p: argparse.ArgumentParser) -> None: p.add_argument("--output-prefix", default=None, help="Output prefix for cohort mode. Defaults to /simulated when --n-individuals is used.") p.add_argument("--output-format", default="same", - choices=("same", "pgen", "vcf", "vcf.gz", "bgen"), - help="Genotype output format for --n-individuals cohort output.") + choices=_OUTPUT_FORMAT_CHOICES, + help="Output format. 'same' preserves the cohort input format and keeps NPZ as the batch-mode default.") p.add_argument("--genetic-map", default=None, help="Genetic map table with columns: chrom, pos, cM.") p.add_argument("--chromosome", type=int, default=None, @@ -132,6 +135,8 @@ def _infer_output_format(input_path: str, requested_format: str) -> str: return "vcf" if suffix == ".bgen": return "bgen" + if suffix == ".bcf": + return "bcf" raise ValueError(f"Cannot infer cohort output format from input path: {input_path}") @@ -144,6 +149,12 @@ def _output_genotype_path(output_prefix: Path, output_format: str) -> Path: return output_prefix.with_suffix(".vcf.gz") if output_format == "bgen": return output_prefix.with_suffix(".bgen") + if output_format == "bcf": + return output_prefix.with_suffix(".bcf") + if output_format == "pkl": + return output_prefix.with_suffix(".pkl") + if output_format == "npz": + return output_prefix.with_suffix(".npz") raise ValueError(f"Unsupported output format: {output_format}") @@ -191,6 +202,32 @@ def _write_genotypes_like_input( from snputils.snp.io.write.bgen import BGENWriter BGENWriter(snpobj=snpobj, filename=output_path).write(phased=True) + elif output_format == "bcf": + from snputils.snp.io.write.bcf import BCFWriter + + BCFWriter(snpobj=snpobj, filename=output_path, phased=True).write() + elif output_format == "pkl": + snpobj.save_pickle(output_path) + elif output_format == "npz": + payload = { + "genotypes": np.asarray(snpobj.genotypes), + "samples": np.asarray(snpobj.samples), + } + for attr in ( + "variants_ref", + "variants_alt", + "variants_chrom", + "variants_cm", + "variants_filter_pass", + "variants_id", + "variants_pos", + "variants_qual", + "variants_info", + ): + value = getattr(snpobj, attr, None) + if value is not None: + payload[attr] = np.asarray(value) + np.savez_compressed(output_path, **payload) else: raise ValueError(f"Unsupported output format: {output_format}") @@ -291,6 +328,50 @@ def _run_cohort_output(args, simulator, snp_data, out_dir: Path) -> int: return 0 +def _to_numpy(value): + if value is None: + return None + if hasattr(value, "detach"): + value = value.detach() + if hasattr(value, "cpu"): + value = value.cpu() + if hasattr(value, "numpy"): + value = value.numpy() + return np.asarray(value) + + +def _batch_genotypes_to_snp_layout(snps) -> np.ndarray: + values = _to_numpy(snps) + if values.ndim != 3 or values.shape[1] != 2: + raise ValueError( + "Diploid batch output must have shape (n_samples, 2, n_variants)." + ) + return np.transpose(values, (2, 0, 1)).astype(np.int8, copy=False) + + +def _write_batch_truth( + path: Path, + *, + labels_d, + labels_c, + cp, + population_names, + output_ploidy: int, +) -> None: + np.savez_compressed( + path, + labels_d=(_to_numpy(labels_d) if labels_d is not None else np.empty(0)), + labels_c=(_to_numpy(labels_c) if labels_c is not None else np.empty(0)), + cp=(_to_numpy(cp) if cp is not None else np.empty(0)), + population_names=( + np.asarray(population_names, dtype=str) + if population_names is not None + else np.empty(0, dtype=str) + ), + output_ploidy=np.asarray(output_ploidy, dtype=np.int8), + ) + + def run_simulator_command(args: argparse.Namespace) -> int: if getattr(args, "verbose", False): log.setLevel(logging.DEBUG) @@ -329,6 +410,8 @@ def run_simulator_command(args: argparse.Namespace) -> int: ancestry_proportions = _parse_ancestry_proportions(args.ancestry_proportions) log.info("Initialising OnlineSimulator...") + from snputils.simulation.simulator import OnlineSimulator + simulator = OnlineSimulator( snp_data = snp_data, meta = meta, @@ -364,20 +447,49 @@ def run_simulator_command(args: argparse.Namespace) -> int: if cp is not None: cp = cp.reshape(args.batch_size, output_ploidy, cp.shape[-1]) - out_path = out_dir / f"batch_{b:04d}.npz" - np.savez_compressed( - out_path, - snps = snps.cpu().numpy(), - labels_d = (labels_d.cpu().numpy() - if labels_d is not None else np.empty(0)), - labels_c = (labels_c.cpu().numpy() - if labels_c is not None else np.empty(0)), - cp = (cp.cpu().numpy() - if cp is not None else np.empty(0)), - population_names = (np.asarray(simulator.population_names, dtype=str) - if simulator.population_names is not None else np.empty(0, dtype=str)), - output_ploidy = np.asarray(output_ploidy, dtype=np.int8), - ) + output_format = "npz" if args.output_format == "same" else args.output_format + output_prefix = out_dir / f"batch_{b:04d}" + if output_format == "npz": + out_path = output_prefix.with_suffix(".npz") + np.savez_compressed( + out_path, + snps=_to_numpy(snps), + labels_d=(_to_numpy(labels_d) if labels_d is not None else np.empty(0)), + labels_c=(_to_numpy(labels_c) if labels_c is not None else np.empty(0)), + cp=(_to_numpy(cp) if cp is not None else np.empty(0)), + population_names=( + np.asarray(simulator.population_names, dtype=str) + if simulator.population_names is not None + else np.empty(0, dtype=str) + ), + output_ploidy=np.asarray(output_ploidy, dtype=np.int8), + ) + else: + if not args.diploid_output: + raise ValueError( + "Batch genotype, BCF, BGEN, and SNPObject outputs require --diploid-output; " + "use NPZ for haplotype batches." + ) + batch_genotypes = _batch_genotypes_to_snp_layout(snps) + batch_samples = np.asarray( + [f"{args.sample_prefix}B{b:04d}_{i + 1:06d}" for i in range(batch_genotypes.shape[1])], + dtype=str, + ) + out_path = _write_genotypes_like_input( + snp_data, + batch_genotypes, + batch_samples, + output_prefix, + output_format, + ) + _write_batch_truth( + output_prefix.with_suffix(".labels.npz"), + labels_d=labels_d, + labels_c=labels_c, + cp=cp, + population_names=simulator.population_names, + output_ploidy=output_ploidy, + ) log.info("Saved %s", out_path.name) log.info("[✓] All done. %d files written to %s", args.n_batches, out_dir) diff --git a/snputils/snp/genobj/snpobj.py b/snputils/snp/genobj/snpobj.py index 433ffb1f..53a0af8d 100644 --- a/snputils/snp/genobj/snpobj.py +++ b/snputils/snp/genobj/snpobj.py @@ -5270,21 +5270,18 @@ def correct_flipped_variants( Correct flipped variants between between `self` and a reference `snpobj`, where reference (`variants_ref`) and alternate (`variants_alt`) alleles are swapped. - Flip Detection Based on `check_complement`: - - - If `check_complement=False`, only direct allele swaps are considered: - 1. Direct Swap: `self.variants_ref == snpobj.variants_alt` and `self.variants_alt == snpobj.variants_ref`. - - - If `check_complement=True`, a swap is accepted when either both original - alleles or both complemented alleles match the swapped reference pair. - Partial complements and strand-ambiguous orientations are not changed. - - Note: Variants where `self.variants_ref == self.variants_alt` are ignored as they are ambiguous. - - Correction Process: - - Swaps `variants_ref` and `variants_alt` alleles in `self` to align with `snpobj`. - - Flips called 2D diploid dosages as `2 - dosage` and called 3D allele indexes as - `1 - allele`, while preserving missing values. + If `check_complement=False`, only direct allele swaps are considered: + `self.variants_ref == snpobj.variants_alt` and + `self.variants_alt == snpobj.variants_ref`. + + If `check_complement=True`, a swap is accepted when either both original + alleles or both complemented alleles match the swapped reference pair. + Partial complements, strand-ambiguous orientations, and variants where + `self.variants_ref == self.variants_alt` are not changed. + + The correction swaps `variants_ref` and `variants_alt` to align with + `snpobj`. It flips called 2D diploid dosages as `2 - dosage` and called 3D + allele indexes as `1 - allele`, while preserving missing values. Args: snpobj (SNPObject): diff --git a/snputils/snp/io/read/bgen.py b/snputils/snp/io/read/bgen.py index b91dba54..38523536 100644 --- a/snputils/snp/io/read/bgen.py +++ b/snputils/snp/io/read/bgen.py @@ -68,6 +68,7 @@ class _BGENRecord: pos: int alleles: tuple[str, ...] probabilities: Optional[np.ndarray] + phased: Optional[bool] def _as_field_list(fields: Optional[Union[str, Sequence[str]]]) -> Optional[List[str]]: @@ -222,8 +223,9 @@ def _read_record(self, index: int, read_probabilities: bool) -> _BGENRecord: block = _read_exact(self.handle, block_len, "genotype block") probabilities = None + phased = None if read_probabilities: - probabilities = self._decode_probabilities(block, n_alleles) + probabilities, phased = self._decode_probabilities(block, n_alleles) return _BGENRecord( index=index, varid=varid, @@ -232,9 +234,10 @@ def _read_record(self, index: int, read_probabilities: bool) -> _BGENRecord: pos=pos, alleles=alleles, probabilities=probabilities, + phased=phased, ) - def _decode_probabilities(self, block: bytes, n_alleles: int) -> np.ndarray: + def _decode_probabilities(self, block: bytes, n_alleles: int) -> tuple[np.ndarray, bool]: assert self.header is not None if _native_bgen is None: raise ImportError("Native BGEN support requires the compiled snputils.snp.io._bgen extension.") @@ -252,12 +255,13 @@ def _decode_probabilities(self, block: bytes, n_alleles: int) -> np.ndarray: if len(payload) != expected_len: raise ValueError("BGEN genotype block decompressed to the wrong size.") - buffer, n_samples, width, _phased, _bit_depth = _native_bgen.decode_layout2( + buffer, n_samples, width, phased, _bit_depth = _native_bgen.decode_layout2( payload, self.header.n_samples, n_alleles, ) - return np.frombuffer(buffer, dtype=np.float32).reshape(n_samples, width) + probabilities = np.frombuffer(buffer, dtype=np.float32).reshape(n_samples, width) + return probabilities, bool(phased) @SNPBaseReader.register @@ -271,16 +275,22 @@ def read( sample_idxs: Optional[Sequence[int]] = None, variant_ids: Optional[Sequence[str]] = None, variant_idxs: Optional[Sequence[int]] = None, + genotype_mode: Optional[str] = None, ) -> SNPObject: """ Read a BGEN file into a SNPObject. Args: - fields: Fields to include. Available fields are ``GP``, ``IID``, ``REF``, - ``ALT``, ``#CHROM``, ``ID``, and ``POS``. ``GT`` is intentionally - unsupported because this reader preserves BGEN genotype probabilities - instead of converting them to hard calls. + fields: Fields to include. Available fields are ``GP``, ``GT``, ``IID``, + ``REF``, ``ALT``, ``#CHROM``, ``ID``, and ``POS``. ``GT`` requires + ``genotype_mode="phased"``. exclude_fields: Fields to exclude from the returned SNPObject. + genotype_mode: By default, preserve BGEN genotype probabilities in + ``calldata_gp``. ``"dosage"`` returns expected biallelic alternate- + allele counts with shape ``(n_variants, n_samples)``. ``"phased"`` + hard-calls each haplotype from phased genotype probabilities and + returns a diploid allele array with shape ``(n_variants, n_samples, 2)``. + Unphased or non-diploid records are rejected in phased mode. sample_path: Optional Oxford ``.sample`` file for BGEN files without embedded sample identifiers. sample_ids: Sample IDs to read. If None and sample_idxs is None, all samples are read. @@ -289,9 +299,15 @@ def read( variant_idxs: Variant indices to read. If None and variant_ids is None, all variants are read. Returns: - SNPObject: A SNPObject with genotype probabilities in ``calldata_gp``. - Mixed probability widths are padded with NaN columns. + SNPObject: A SNPObject with genotype probabilities in ``calldata_gp`` by + default, or dosages/phased hard calls in ``genotypes`` when requested. + Mixed probability widths are padded with NaN columns in probability mode. """ + if genotype_mode is not None: + genotype_mode = str(genotype_mode).strip().lower() + if genotype_mode not in {"dosage", "phased"}: + raise ValueError("BGENReader genotype_mode must be 'dosage' or 'phased' when provided.") + if sample_idxs is not None and sample_ids is not None: raise ValueError("Only one of sample_idxs and sample_ids can be specified.") if variant_idxs is not None and variant_ids is not None: @@ -301,12 +317,18 @@ def read( exclude = set(_as_field_list(exclude_fields) or []) fields_set = {field for field in fields_list if field not in exclude} - if "GT" in fields_set: + if "GT" in fields_set and genotype_mode != "phased": raise NotImplementedError( - "BGENReader preserves genotype probabilities in `calldata_gp` and does not hard-call GT." + "BGENReader does not hard-call GT by default; request " + "genotype_mode='phased' to hard-call phased probabilities." ) + if genotype_mode == "phased": + fields_set.discard("GT") + fields_set.add("GP") + elif genotype_mode == "dosage": + fields_set.discard("GP") - if self._can_use_native_bulk_gp( + if genotype_mode is None and self._can_use_native_bulk_gp( fields_set=fields_set, sample_path=sample_path, sample_ids=sample_ids, @@ -345,8 +367,20 @@ def read( samples = bfile.samples[sample_indices] if "IID" in fields_set else None metadata, calldata_gp = self._records_to_arrays(records, sample_indices, fields_set) + genotypes = None + if genotype_mode == "phased": + genotypes = self._records_to_phased_hardcalls(records, len(sample_indices)) + calldata_gp = None + elif genotype_mode == "dosage": + genotypes = self.read_dosage( + sample_path=sample_path, + sample_ids=sample_ids, + sample_idxs=sample_idxs, + variant_ids=variant_ids, + variant_idxs=variant_idxs, + ) return SNPObject( - genotypes=None, + genotypes=genotypes, calldata_gp=calldata_gp, samples=samples, variants_ref=metadata["variants_ref"], @@ -489,6 +523,7 @@ def _load_records( pos=record.pos, alleles=record.alleles, probabilities=record.probabilities[sample_indices, :], + phased=record.phased, ) if requested_idx_set is not None: by_index[record.index] = record @@ -510,6 +545,58 @@ def _load_records( raise ValueError(f"The following specified variants were not found: {missing}") return selected + @staticmethod + def _records_to_phased_hardcalls( + records: Sequence[_BGENRecord], + n_samples: int, + ) -> np.ndarray: + genotypes = np.empty((len(records), n_samples, 2), dtype=np.int8) + + for variant_index, record in enumerate(records): + variant_id = _variant_identifier(record.varid, record.rsid) or f"index {record.index}" + if record.phased is not True: + raise ValueError( + "BGEN genotype_mode='phased' requires phased probabilities; " + f"variant {variant_id!r} is unphased." + ) + if record.probabilities is None: + raise ValueError(f"BGEN probabilities were not loaded for variant {variant_id!r}.") + + n_alleles = len(record.alleles) + if n_alleles < 2 or n_alleles > np.iinfo(np.int8).max: + raise ValueError( + f"Variant {variant_id!r} has unsupported allele count {n_alleles}." + ) + + probabilities = np.asarray(record.probabilities, dtype=np.float32) + expected_width = 2 * n_alleles + if probabilities.shape != (n_samples, expected_width): + raise ValueError( + "BGEN genotype_mode='phased' requires diploid phased probabilities; " + f"variant {variant_id!r} has probability width {probabilities.shape[1]}, " + f"expected {expected_width}." + ) + + haplotype_probabilities = probabilities.reshape(n_samples, 2, n_alleles) + called = np.all(np.isfinite(haplotype_probabilities), axis=2) + partial_ploidy = np.sum(called, axis=1) == 1 + if np.any(partial_ploidy): + sample_index = int(np.flatnonzero(partial_ploidy)[0]) + raise ValueError( + "BGEN genotype_mode='phased' requires diploid phased probabilities; " + f"variant {variant_id!r} has a non-diploid call for sample index {sample_index}." + ) + safe_probabilities = np.where( + np.isfinite(haplotype_probabilities), + haplotype_probabilities, + -np.inf, + ) + variant_genotypes = np.argmax(safe_probabilities, axis=2).astype(np.int8) + variant_genotypes[~called] = -1 + genotypes[variant_index] = variant_genotypes + + return genotypes + @classmethod def _records_to_arrays( cls, @@ -574,4 +661,4 @@ def _resolve_sample_indices( missing = [str(sample) for sample in requested if str(sample) not in sample_lookup] if missing: raise ValueError(f"The following specified samples were not found: {missing}") - return np.asarray([sample_lookup[str(sample)] for sample in requested], dtype=int) \ No newline at end of file + return np.asarray([sample_lookup[str(sample)] for sample in requested], dtype=int) diff --git a/snputils/snp/io/write/__test__/test_bgen_writer.py b/snputils/snp/io/write/__test__/test_bgen_writer.py index 925eabf9..f41152e2 100644 --- a/snputils/snp/io/write/__test__/test_bgen_writer.py +++ b/snputils/snp/io/write/__test__/test_bgen_writer.py @@ -67,6 +67,36 @@ def test_bgen_reader_native_bulk_paths_match_general_reader(tmp_path): np.testing.assert_allclose(fast_dosage, general_dosage, atol=1 / 65535, equal_nan=True) +def test_bgen_reader_dosage_mode_returns_metadata_complete_snpobject(tmp_path): + path = tmp_path / "dosage_mode.bgen" + gp = np.array( + [ + [[1.0, 0.0, 0.0], [0.2, 0.3, 0.5], [0.0, 1.0, 0.0]], + [[0.0, 0.0, 1.0], [0.25, 0.5, 0.25], [1.0, 0.0, 0.0]], + ], + dtype=np.float32, + ) + snpobj = SNPObject( + calldata_gp=gp, + samples=np.array(["s1", "s2", "s3"], dtype=object), + variants_ref=np.array(["A", "C"], dtype=object), + variants_alt=np.array(["G", "T"], dtype=object), + variants_chrom=np.array(["1", "1"], dtype=object), + variants_id=np.array(["rs1", "rs2"], dtype=object), + variants_pos=np.array([10, 20]), + ) + + BGENWriter(snpobj, path).write(phased=False, bit_depth=16) + observed = BGENReader(path).read(genotype_mode="dosage") + + expected = gp[:, :, 1] + 2 * gp[:, :, 2] + np.testing.assert_allclose(observed.genotypes, expected, atol=1 / 65535) + np.testing.assert_array_equal(observed.samples, snpobj.samples) + np.testing.assert_array_equal(observed.variants_id, snpobj.variants_id) + np.testing.assert_array_equal(observed.variants_pos, snpobj.variants_pos) + assert observed.calldata_gp is None + + def test_bgen_writer_rejects_incompatible_phased_width(tmp_path): snpobj = SNPObject( calldata_gp=np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32), @@ -188,3 +218,58 @@ def test_bgen_writer_roundtrips_variable_ploidy_phased(tmp_path): observed = BGENReader(path).read() np.testing.assert_allclose(observed.calldata_gp, gp, atol=1 / 65535, equal_nan=True) + with pytest.raises(ValueError, match="non-diploid call"): + BGENReader(path).read(genotype_mode="phased") + + +def test_bgen_reader_hardcalls_phased_probabilities(tmp_path): + path = tmp_path / "phased_probabilities.bgen" + gp = np.array( + [ + [[0.9, 0.1, 0.2, 0.8], [0.1, 0.9, 0.8, 0.2], [np.nan] * 4], + [[0.4, 0.6, 0.7, 0.3], [0.8, 0.2, 0.1, 0.9], [0.2, 0.8, 0.3, 0.7]], + ], + dtype=np.float32, + ) + snpobj = SNPObject( + calldata_gp=gp, + samples=np.array(["s1", "s2", "missing"], dtype=object), + variants_ref=np.array(["A", "C"], dtype=object), + variants_alt=np.array(["G", "T"], dtype=object), + variants_chrom=np.array(["1", "1"], dtype=object), + variants_id=np.array(["rs1", "rs2"], dtype=object), + variants_pos=np.array([10, 20]), + ) + + BGENWriter(snpobj, path).write(phased=True) + observed = BGENReader(path).read(genotype_mode="phased") + + expected = np.array( + [ + [[0, 1], [1, 0], [-1, -1]], + [[1, 0], [0, 1], [1, 1]], + ], + dtype=np.int8, + ) + np.testing.assert_array_equal(observed.genotypes, expected) + np.testing.assert_array_equal(observed.samples, snpobj.samples) + np.testing.assert_array_equal(observed.variants_id, snpobj.variants_id) + assert observed.calldata_gp is None + + +def test_bgen_reader_rejects_unphased_probabilities_in_phased_mode(tmp_path): + path = tmp_path / "unphased_probabilities.bgen" + snpobj = SNPObject( + calldata_gp=np.array([[[0.2, 0.3, 0.5]]], dtype=np.float32), + samples=np.array(["s1"], dtype=object), + variants_ref=np.array(["A"], dtype=object), + variants_alt=np.array(["G"], dtype=object), + variants_chrom=np.array(["1"], dtype=object), + variants_id=np.array(["rs1"], dtype=object), + variants_pos=np.array([10]), + ) + + BGENWriter(snpobj, path).write(phased=False) + + with pytest.raises(ValueError, match="requires phased probabilities"): + BGENReader(path).read(genotype_mode="phased") diff --git a/snputils/tools/__test__/test_main_cli.py b/snputils/tools/__test__/test_main_cli.py index 9797dff9..c335081b 100644 --- a/snputils/tools/__test__/test_main_cli.py +++ b/snputils/tools/__test__/test_main_cli.py @@ -9,6 +9,8 @@ from snputils.ancestry.io.local.read.__test__.fixtures import make_synthetic_dataset, write_msp from snputils.simulation._validation import validate_phased_simulation_input +from snputils.snp.genobj.snpobj import SNPObject +from snputils.snp.io.write.bgen import BGENWriter from snputils.tools.cli import main @@ -24,6 +26,30 @@ def _write_tiny_vcf(path: Path) -> None: ) +def _write_tiny_bgen(path: Path) -> None: + dosages = np.array( + [ + [0, 1, 2, 1], + [2, 1, 0, 1], + [0, 0, 1, 2], + ], + dtype=np.int8, + ) + gp = np.eye(3, dtype=np.float32)[dosages] + BGENWriter( + SNPObject( + calldata_gp=gp, + samples=np.array(["s0", "s1", "s2", "s3"], dtype=object), + variants_ref=np.array(["A", "C", "G"], dtype=object), + variants_alt=np.array(["G", "T", "A"], dtype=object), + variants_chrom=np.array(["1", "1", "1"], dtype=object), + variants_id=np.array(["rs1", "rs2", "rs3"], dtype=object), + variants_pos=np.array([100, 200, 300]), + ), + path, + ).write(phased=False) + + def _write_gwas_vcf(path: Path, sample_ids) -> None: dosage_rows = [ [0, 1, 2, 0, 1, 2, 1, 0], @@ -179,6 +205,40 @@ def test_main_pca_sklearn_smoke_with_pgen_auto_reader( assert components.shape[1] == 2 +def test_main_pca_sklearn_smoke_with_bgen_default_dosage( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + matplotlib.use("Agg", force=True) + monkeypatch.setenv("MPLBACKEND", "Agg") + + bgen_path = tmp_path / "tiny.bgen" + fig_path = tmp_path / "pca_bgen.png" + npy_path = tmp_path / "components_bgen.npy" + _write_tiny_bgen(bgen_path) + + monkeypatch.setattr( + sys, + "argv", + [ + "snputils", + "pca", + "--snp-path", + str(bgen_path), + "--plot", + str(fig_path), + "--components", + str(npy_path), + "--backend", + "sklearn", + ], + ) + + assert main() == 0 + assert fig_path.exists() + assert npy_path.exists() + assert np.load(npy_path).shape == (4, 2) + + def test_main_admixture_map_smoke(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: sample_ids, y, lai, chromosomes, starts, ends, ancestry_map = make_synthetic_dataset( n_samples=24, diff --git a/snputils/tools/cli.py b/snputils/tools/cli.py index 785cb2b3..e828ad21 100644 --- a/snputils/tools/cli.py +++ b/snputils/tools/cli.py @@ -610,14 +610,19 @@ def _add_simulate_arguments(parser: argparse.ArgumentParser) -> None: "--snp", required=True, help=( - "Path to phased SNP input (VCF, PGEN, or BGEN fileset). " + "Path to phased SNP input (VCF, BCF, PGEN, or BGEN fileset). " "PLINK1 BED is not supported because it cannot store phase." ), ) parser.add_argument("--metadata", required=True, help="TSV/CSV file with at least Sample/IID and Population columns.") parser.add_argument("--output-dir", required=True, help="Directory in which to save the simulated batches.") parser.add_argument("--output-prefix", default=None, help="Output prefix for cohort mode. Defaults to /simulated when --n-individuals is used.") - parser.add_argument("--output-format", default="same", choices=("same", "pgen", "vcf", "vcf.gz", "bgen"), help="Genotype output format for --n-individuals cohort output.") + parser.add_argument( + "--output-format", + default="same", + choices=("same", "npz", "pkl", "pgen", "vcf", "vcf.gz", "bcf", "bgen"), + help="Output format. 'same' preserves the cohort input format and keeps NPZ as the batch-mode default.", + ) parser.add_argument("--genetic-map", default=None, help="Genetic map table with columns: chrom, pos, cM.") parser.add_argument("--chromosome", type=int, default=None, help="If provided, restrict genetic map rows to this chromosome id.") parser.add_argument("--window-size", type=int, default=1000, help="#SNPs per window.") diff --git a/snputils/tools/pca.py b/snputils/tools/pca.py index f2164850..995ff978 100644 --- a/snputils/tools/pca.py +++ b/snputils/tools/pca.py @@ -28,7 +28,7 @@ def add_pca_arguments(parser: argparse.ArgumentParser) -> None: dest="snp_path", required=True, type=str, - help="Path to genotype input (VCF, BED, or PGEN fileset).", + help="Path to genotype input (VCF, BCF, BGEN, BED, or PGEN fileset).", ) parser.add_argument( "--plot",