pg_gpu

GPU-accelerated population genetics statistics using CuPy.

Installation

pg_gpu requires a Linux x86_64 machine with an NVIDIA GPU and a CUDA 12 driver. Nothing else is needed -- the full GPU runtime, including the CUDA toolkit headers cupy uses to JIT-compile its kernels, is pulled from PyPI via the cupy-cuda12x[ctk] dependency.

With pixi (recommended)

The pinned, reproducible environment is managed with pixi and is the recommended way to install pg_gpu:

pixi install
pixi shell

Into an existing conda / venv environment

To use pg_gpu from your own workflow (Snakemake, Jupyter, an existing conda env), install it with pip:

pip install "git+https://github.com/kr-colab/pg_gpu"

This pulls the full runtime stack (cupy-cuda12x with toolkit headers, bio2zarr, kvikio, nvcomp) as declared in pyproject.toml. For development against a local checkout, use an editable install:

pip install -e ".[dev]"

Quick Start

from pg_gpu import HaplotypeMatrix, diversity, divergence, selection, sfs

# Load from VCF
hm = HaplotypeMatrix.from_vcf("data.vcf.gz", region="chr1:1-1000000")
hm.load_pop_file("populations.txt")

# Diversity
diversity.pi(hm, population="pop1")
diversity.tajimas_d(hm, population="pop1")

# Divergence
divergence.fst_hudson(hm, "pop1", "pop2")
divergence.dxy(hm, "pop1", "pop2")

# Selection scans
selection.ihs(hm)
selection.nsl(hm)

# Windowed statistics (fused CUDA kernels)
from pg_gpu import windowed_analysis
results = windowed_analysis(hm, statistics=["pi", "theta_w", "tajimas_d"],
                            window_size=50000)

Documentation

Full documentation at https://pg-gpu.readthedocs.io/.

Interactive walkthrough: examples/pg_gpu_tour.ipynb.

Statistics

Development

pixi run pytest tests/
pixi run -e lint ruff check pg_gpu/