rfmix-reader 0.2.1__tar.gz → 0.3.1__tar.gz

rfmix_reader-0.3.1/PKG-INFO

@@ -0,0 +1,512 @@
+Metadata-Version: 2.4
+Name: rfmix-reader
+Version: 0.3.1
+Summary: RFMix-reader is a Python package designed to efficiently read and process output files generated by RFMix, a popular tool for estimating local ancestry in admixed populations. The package employs a lazy loading approach, which minimizes memory consumption by reading only the loci that are accessed by the user, rather than loading the entire dataset into memory at once.
+License: GPL-3.0-or-later
+License-File: LICENSE
+Keywords: file parser,rfmix,gpu acceleration,local ancestry
+Author: Kynon J.M. Benjamin
+Author-email: kj.benjamin90@gmail.com
+Maintainer: Kynon J.M. Benjamin
+Maintainer-email: kj.benjamin90@gmail.com
+Requires-Python: >=3.11,<3.15
+Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Provides-Extra: all
+Provides-Extra: docs
+Provides-Extra: gpu
+Provides-Extra: io
+Provides-Extra: tests
+Provides-Extra: viz
+Requires-Dist: cairosvg (>=2.7,<3.0) ; extra == "viz" or extra == "all"
+Requires-Dist: cudf-cu12 ; extra == "all"
+Requires-Dist: cupy-cuda12x ; extra == "all"
+Requires-Dist: cyvcf2 (>=0.31)
+Requires-Dist: dask (>=2025.1,<2026.0)
+Requires-Dist: dask-cudf-cu12 ; extra == "all"
+Requires-Dist: fonttools (>=4.61.0,<5.0.0)
+Requires-Dist: matplotlib (>=3.10,<4.0) ; extra == "viz" or extra == "all"
+Requires-Dist: numpy (>=1.23,<3)
+Requires-Dist: pandas (>=2.0)
+Requires-Dist: psutil (>=7,<8)
+Requires-Dist: pyarrow ; extra == "io" or extra == "all"
+Requires-Dist: seaborn (>=0.13,<0.14) ; extra == "viz" or extra == "all"
+Requires-Dist: sphinx (>=7,<9) ; extra == "docs"
+Requires-Dist: sphinx-autodoc-typehints (>=2,<3) ; extra == "docs"
+Requires-Dist: sphinx-copybutton (>=0.5,<0.6) ; extra == "docs"
+Requires-Dist: sphinx-rtd-theme (>=2,<3) ; extra == "docs"
+Requires-Dist: torch (>=2.8) ; extra == "gpu" or extra == "all"
+Requires-Dist: tqdm (>=4.66)
+Project-URL: Bug Tracker, https://github.com/heart-gen/rfmix_reader/issues
+Project-URL: homepage, https://rfmix-reader.readthedocs.io/en/latest/
+Project-URL: repository, https://github.com/heart-gen/rfmix_reader.git
+Description-Content-Type: text/markdown
+
+# RFMix-reader
+`RFMix-reader` is a Python package for efficiently reading and processing output
+files generated by [`RFMix`](https://github.com/slowkoni/rfmix), a widely used tool
+for estimating local ancestry in admixed populations.  
+It employs a **lazy loading approach** to minimize memory usage, and leverages **GPU acceleration**
+for major speedups when available.
+
+---
+
+## Installation
+
+`rfmix-reader` requires **Python 3.11+**. Install from PyPI:
+
+```bash
+pip install rfmix-reader
+````
+
+### Installation Options
+
+* **Basic install** (CPU only):
+
+  ```bash
+  pip install rfmix-reader
+  ```
+
+* **With GPU acceleration** (`cupy`, `cudf`, `dask-cudf`, `torch`):
+
+  ```bash
+  pip install rfmix-reader[gpu]
+  ```
+
+  The default GPU extra targets **CUDA 12** wheels (`cupy-cuda12x`,
+  `cudf-cu12`, `dask-cudf-cu12`). Use the appropriate CUDA build strings if
+  your environment requires a different version.
+  
+* **With documentation tools** (`sphinx`, `sphinx-rtd-theme`):
+
+  ```bash
+  pip install rfmix-reader[docs]
+  ```
+
+* **With testing tools** (`pytest`):
+
+  ```bash
+  pip install rfmix-reader[tests]
+  ```
+
+### GPU Notes
+
+* PyTorch is **not** installed in the base package. It is pulled in only when you install the `[gpu]` extra (`pip install rfmix-reader[gpu]`) or add it yourself.
+* Install a PyTorch wheel that matches your platform and CUDA version using the [official selector](https://pytorch.org/get-started/locally/). For example:
+  * **CUDA 12 (Linux/Windows):** `pip install torch --index-url https://download.pytorch.org/whl/cu121`
+  * **CUDA 11 (Linux/Windows):** `pip install torch --index-url https://download.pytorch.org/whl/cu118`
+  * **CPU-only:** `pip install torch --index-url https://download.pytorch.org/whl/cpu`
+* RAPIDS (`cudf`, `cupy`) wheels are version- and CUDA-specific. See the [RAPIDS install guide](https://docs.rapids.ai/install).
+* CPU-only installations will still run efficiently, just without GPU acceleration.
+
+---
+
+## Quickstart
+
+```python
+from rfmix_reader import read_rfmix
+
+# Load RFMix outputs (two-population admixture example)
+file_path = "examples/two_populations/out/"
+loci_df, g_anc, local_array = read_rfmix(file_path)
+
+print(loci_df.head())
+print(g_anc.head())
+print(local_array.shape)
+"""See the phasing section below for how to phase per-chromosome outputs and
+write them to Zarr with `phase_rfmix_chromosome_to_zarr`."""
+```
+
+---
+
+## Key Features
+
+* **Lazy Loading**: Reads data on-the-fly, reducing memory footprint.
+* **Efficient Access**: Query specific loci or regions of interest.
+* **Seamless Integration**: Works smoothly with `pandas`, `dask`, and other analysis tools.
+* **Loci Imputation**: Impute local ancestry loci to dense genotype variant sites.
+* **GPU Acceleration**: Automatic CUDA acceleration via PyTorch/CuPy when available.
+
+---
+
+## Simulation Data
+
+Test datasets for two- and three-population admixture are available on Synapse:
+[Synapse Project syn61691659](https://www.synapse.org/Synapse:syn61691659).
+
+---
+
+## Usage
+
+### Binary Conversion
+
+RFMix does not generate binary files directly.
+Use `create_binaries` to generate them (also available as a CLI):
+
+```bash
+create-binaries two_pops/out/
+```
+
+```python
+from rfmix_reader import create_binaries
+
+create_binaries("two_pops/out/", binary_dir="./binary_files")
+```
+
+### Preparing Reference Data for Phasing
+
+Use `prepare-reference` to convert bgzipped, indexed reference VCF/BCF files
+into per-chromosome VCF-Zarr stores that the phasing pipeline consumes.
+The command writes one `<chrom>.zarr` directory per input file.
+
+```
+prepare-reference -h
+```
+
+```
+usage: prepare-reference [-h] [--chunk-length CHUNK_LENGTH]
+                         [--samples-chunk-size SAMPLES_CHUNK_SIZE]
+                         [--worker-processes WORKER_PROCESSES]
+                         [--verbose | --no-verbose] [--version]
+                         output_dir vcf_paths [vcf_paths ...]
+
+Convert one or more bgzipped reference VCF/BCF files into Zarr stores.
+
+positional arguments:
+  output_dir            Directory where the Zarr outputs will be written.
+  vcf_paths             Paths to reference VCF/BCF files (bgzipped and
+                        indexed).
+
+options:
+  -h, --help            show this help message and exit
+  --chunk-length CHUNK_LENGTH
+                        Genomic chunk size for the output Zarr stores
+                        (default: 100000).
+  --samples-chunk-size SAMPLES_CHUNK_SIZE
+                        Chunk size for samples in the output Zarr stores
+                        (default: library chosen).
+  --worker-processes WORKER_PROCESSES
+                        Number of worker processes to use for conversion
+                        (default: 0, use library default).
+  --verbose, --no-verbose
+                        Print progress messages (default: enabled).
+  --version             Show the version of the program and exit.
+```
+
+Example data preparation:
+
+```bash
+# Sample annotations: two columns (no header): sample_id<TAB>group
+cat > sample_annotations.tsv <<'EOF'
+NA19700	AFR
+NA19701	AFR
+NA20847	EUR
+EOF
+
+# Convert chromosome VCFs into a reference store directory
+prepare-reference refs/ 1kg_chr20.vcf.gz 1kg_chr21.vcf.gz \
+  --chunk-length 50000 --samples-chunk-size 512
+```
+
+### Main Function
+
+Once binaries are available, process RFMix results:
+
+```python
+from rfmix_reader import read_rfmix
+
+loci, g_anc, admix = read_rfmix("two_pops/out/")
+```
+
+### Three Population Example
+
+Binaries can also be generated on-the-fly within `read_rfmix` with
+`generate_binary` set to `True`.
+
+```python
+loci, g_anc, admix = read_rfmix("examples/three_populations/out/",
+                                binary_dir="./binary_files",
+                                generate_binary=True)
+```
+
+### Phasing data
+
+For optimal memory and computational speed, phasing is done per
+chromosome.
+
+```python
+from rfmix_reader import phase_rfmix_chromosome_to_zarr
+
+# Use the reference store + annotations during phasing
+admix = phase_rfmix_chromosome_to_zarr(
+    file_prefix="two_pops/out/",
+    ref_zarr_root="refs",
+    sample_annot_path="sample_annotations.tsv",
+    output_path="./phased_chr21.zarr",
+    chrom="21",
+)
+```
+
+The chunking is suboptimal for phasing, so remember to
+rechunk before using for optimal processing. This is only needed
+when loading individual chromosomes. The merged data is already
+optimized.
+
+```python
+local_array = admix["local_ancestry"].chunk({"variant": 20000, "sample": 100})
+# Compute into memory, if needed
+local_array = local_array.compute()
+```
+
+This also saves the data to Zarr for later merging or data processing.
+
+```bash
+merge-phased-zarrs ./phased_all.zarr ./phased_chr21.zarr ./phased_chr22.zarr
+```
+### Loci Imputation
+
+The imputation workflow now lives in ``rfmix_reader.processing.imputation`` and
+is exported as ``interpolate_array``. It interpolates the local ancestry matrix
+onto a denser variant grid and writes the result to ``<zarr_outdir>/local-ancestry.zarr``
+as a Zarr array shaped ``(variants, samples, ancestries)``.
+
+**Inputs**
+
+* ``variant_loci_df``: a pandas DataFrame defining the variant grid. Provide at
+  least ``chrom``/``pos`` and an ``i`` column that points to the source RFMix
+  row index; rows with ``i`` set to ``NaN`` are treated as missing loci to
+  interpolate. Sort the frame by genomic coordinate, and include ``pos`` if you
+  plan to interpolate in base-pair space.
+* ``admix``: the local ancestry Dask array returned by ``read_rfmix`` (shape
+  ``(loci, samples, ancestries)``).
+* ``zarr_outdir``: an output directory where the new ``local-ancestry.zarr``
+  store will be created.
+
+**Key options**
+
+* ``interpolation``: ``"linear"`` (default), ``"nearest"``, or ``"stepwise"``.
+* ``use_bp_positions``: set to ``True`` to interpolate along ``variant_loci_df['pos']``
+  rather than treating loci as equally spaced indices.
+* ``chunk_size``/``batch_size``: tune how many rows are materialized at a time
+  when filling and interpolating the Zarr array.
+
+**Workflow example**
+
+```python
+import pandas as pd
+from pathlib import Path
+from rfmix_reader import interpolate_array, read_rfmix
+
+# Load RFMix loci and local ancestry
+loci_df, _, admix = read_rfmix("two_pops/out/", binary_dir="./binary_files")
+
+# Build the variant grid by merging genotype sites with the RFMix loci index
+variants = pd.read_parquet("genotypes/variants.parquet")  # must include chrom/pos
+variants = variants.drop_duplicates(subset=["chrom", "pos"]).sort_values("pos")
+variant_loci_df = (
+    variants.merge(loci_df.to_pandas(), on=["chrom", "pos"], how="outer", indicator=True)
+            .loc[:, ["chrom", "pos", "i", "_merge"]]
+)
+
+z = interpolate_array(
+    variant_loci_df,
+    admix,
+    zarr_outdir=Path("./imputed_local_ancestry"),
+    interpolation="linear",
+    use_bp_positions=True,
+    chunk_size=50_000,
+)
+print(z)
+```
+
+The interpolator uses GPU acceleration transparently when ``cupy`` and a CUDA
+PyTorch build are available; otherwise it falls back to NumPy. All methods other
+than ``"hmm"`` operate on diploid-summed trajectories and preserve the original
+ancestry dimension.
+
+### Phasing workflow (`rfmix_reader.processing.phase`)
+
+`rfmix_reader.processing.phase` implements gnomix-style tail-flip corrections
+for local ancestry haplotypes. Phasing now lives outside `read_rfmix` so you can
+process each chromosome independently and write outputs directly to Zarr.
+
+#### Required reference inputs
+
+* **VCF-Zarr reference** (`ref_zarr_root`): either a single `*.zarr` store or a
+  directory containing per-chromosome stores (e.g., `1.zarr`, `chr1.zarr`).
+* **Sample annotations** (`sample_annot_path`): two-column file mapping
+  `sample_id` to `group` (ancestry label). One representative sample per group
+  is pulled to build reference haplotypes.
+
+#### Per-chromosome pipeline
+
+1. (Optional) generate RFMix binary caches with `create_binaries`.
+2. Call `phase_rfmix_chromosome_to_zarr` for each chromosome you want to
+   process.
+3. Optionally concatenate those per-chromosome Zarr stores with
+   `merge_phased_zarrs`.
+
+```python
+from rfmix_reader.processing.phase import (
+    PhasingConfig,
+    merge_phased_zarrs,
+    phase_rfmix_chromosome_to_zarr,
+)
+
+# Phase chromosome 21 and write to Zarr
+dataset = phase_rfmix_chromosome_to_zarr(
+    file_prefix="examples/two_populations/out/",
+    ref_zarr_root="/refs/1kg_chr_zarr/",
+    sample_annot_path="/refs/1kg_annotations.tsv",
+    output_path="/tmp/phased_chr21.zarr",
+    chrom="21",
+)
+
+# Merge multiple per-chromosome Zarr stores
+merged = merge_phased_zarrs(
+    ["/tmp/phased_chr21.zarr", "/tmp/phased_chr22.zarr"],
+    output_path="/tmp/phased_all.zarr",
+)
+```
+
+If you need fine-grained control, you can still start from unphased outputs and
+call `phase_admix_dask_with_index` directly:
+
+```python
+from rfmix_reader import read_rfmix
+from rfmix_reader.processing.phase import PhasingConfig, phase_admix_dask_with_index
+
+loci_df, g_anc, admix, X_raw = read_rfmix(
+    "examples/two_populations/out/",
+    return_original=True,
+    chrom="21",
+)
+
+config = PhasingConfig(window_size=100, min_block_len=10, max_mismatch_frac=0.3)
+phased = phase_admix_dask_with_index(
+    admix=admix,
+    X_raw=X_raw,
+    positions=loci_df.physical_position.to_numpy(),
+    chrom=str(loci_df.chromosome.iloc[0]),
+    ref_zarr_root="/refs/1kg_chr_zarr/",
+    sample_annot_path="/refs/1kg_annotations.tsv",
+    config=config,
+)
+```
+
+### Reading Haptools simulations
+
+Use `read_simu` to load BGZF-compressed VCF files created by
+`haptools simgenotype --pop_field`:
+
+```python
+from rfmix_reader import read_simu
+
+loci_df, g_anc, admix = read_simu("/path/to/simulations/")
+```
+
+Haptools does **not** include the chromosome length in the `##contig`
+header lines, but `read_simu` requires that metadata to index each VCF.
+Copy the `contigs.txt` file Haptools generates from the FASTA you used
+for simulation and reheader every file with the appropriate contig entry
+before calling `read_simu`. The following snippet shows one approach
+using `bcftools` and `tabix`:
+
+```bash
+CONTIGS="../../three_populations/_m/contigs.txt"
+VCFDIR="gt-files"
+CHR="chr${SLURM_ARRAY_TASK_ID}"
+OUT="${VCFDIR}/${CHR}.vcf.gz"
+IN="${VCFDIR}/back/${CHR}.vcf.gz"
+
+CONTIG_LINE=$(grep -w "ID=${CHR}" "$CONTIGS")
+if [[ -z "$CONTIG_LINE" ]]; then
+    echo "ERROR: No contig line found for ${CHR} in $CONTIGS"
+    exit 1
+fi
+
+bcftools view -h "$IN" \
+    | sed "s/^##contig=<ID=${CHR}>.*/${CONTIG_LINE}/" > header.${CHR}.tmp
+bcftools reheader -h header.${CHR}.tmp -o "$OUT" "$IN"
+tabix -p vcf "$OUT"
+```
+
+### Visualization
+
+`read_rfmix`, `read_flare`, and `read_simu` all return the same
+`(loci_df, g_anc, admix)` tuple, so the plotting utilities in
+`rfmix_reader._visualization` work identically for RFMix, FLARE, and
+Haptools-simulated inputs. The snippet below shows the typical workflow
+for each reader:
+
+```python
+from rfmix_reader import (
+    plot_ancestry_by_chromosome,
+    plot_global_ancestry,
+    read_flare,
+    read_rfmix,
+    read_simu,
+)
+
+# RFMix run directory
+loci_df, g_anc, admix = read_rfmix("two_pops/out/")
+plot_global_ancestry(g_anc, save_path="rfmix_global.png")
+plot_ancestry_by_chromosome(loci_df, admix, save_path="rfmix_local.png")
+
+# FLARE output directory (contains *.anc.vcf.gz + global.anc.gz)
+loci_df, g_anc, admix = read_flare("flare_runs/chr1/")
+plot_global_ancestry(g_anc, save_path="flare_global.png")
+plot_ancestry_by_chromosome(loci_df, admix, save_path="flare_local.png")
+
+# Haptools simulations (after reheadering contigs)
+loci_df, g_anc, admix = read_simu("/path/to/simulations/")
+plot_global_ancestry(g_anc, save_path="simu_global.png")
+plot_ancestry_by_chromosome(loci_df, admix, save_path="simu_local.png")
+```
+
+`plot_global_ancestry` builds per-individual stacked bars of global
+ancestry while `plot_ancestry_by_chromosome` summarizes local ancestry
+along each chromosome, giving you quick visual QC for every supported
+input format.
+
+---
+
+## Development Install
+
+For contributors:
+
+```bash
+git clone https://github.com/heart-gen/rfmix_reader.git
+cd rfmix_reader
+pip install -e ".[gpu,docs,tests]"
+```
+
+---
+
+## Citation
+
+If you use this software, please cite:
+
+[![DOI](https://zenodo.org/badge/807052842.svg)](https://zenodo.org/doi/10.5281/zenodo.12629787)
+
+Benjamin, K. J. M. (2024). **RFMix-reader (Version 0.2.0)** \[Computer software].
+[https://github.com/heart-gen/rfmix\_reader](https://github.com/heart-gen/rfmix_reader)
+
+Kynon JM Benjamin. *"RFMix-reader: Accelerated reading and processing for local ancestry studies."*
+**bioRxiv** (2024).
+DOI: [10.1101/2024.07.13.603370](https://www.biorxiv.org/content/10.1101/2024.07.13.603370v2).
+
+---
+
+## Funding
+
+This work was supported by the National Institutes of Health,
+National Institute on Minority Health and Health Disparities (NIMHD)
+K99MD016964 / R00MD016964.
+
+