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.


