rfmix-reader 0.3.0__tar.gz → 0.7.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (61) hide show
  1. rfmix_reader-0.7.0/PKG-INFO +288 -0
  2. rfmix_reader-0.7.0/README.md +236 -0
  3. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/pyproject.toml +23 -6
  4. rfmix_reader-0.7.0/rfmix_reader/__init__.py +114 -0
  5. rfmix_reader-0.7.0/rfmix_reader/backends.py +73 -0
  6. rfmix_reader-0.7.0/rfmix_reader/cli/__init__.py +1 -0
  7. rfmix_reader-0.7.0/rfmix_reader/cli/main.py +78 -0
  8. rfmix_reader-0.7.0/rfmix_reader/core/__init__.py +54 -0
  9. rfmix_reader-0.7.0/rfmix_reader/core/accessor.py +230 -0
  10. rfmix_reader-0.7.0/rfmix_reader/core/api.py +252 -0
  11. rfmix_reader-0.7.0/rfmix_reader/core/codes.py +92 -0
  12. rfmix_reader-0.7.0/rfmix_reader/core/legacy.py +91 -0
  13. rfmix_reader-0.7.0/rfmix_reader/core/schema.py +207 -0
  14. rfmix_reader-0.7.0/rfmix_reader/core/zarr_io.py +152 -0
  15. rfmix_reader-0.7.0/rfmix_reader/formats/__init__.py +7 -0
  16. rfmix_reader-0.7.0/rfmix_reader/formats/base.py +80 -0
  17. rfmix_reader-0.7.0/rfmix_reader/formats/common.py +437 -0
  18. rfmix_reader-0.7.0/rfmix_reader/formats/discover.py +56 -0
  19. rfmix_reader-0.7.0/rfmix_reader/formats/flare_vcf.py +127 -0
  20. rfmix_reader-0.7.0/rfmix_reader/formats/global_ancestry.py +53 -0
  21. rfmix_reader-0.7.0/rfmix_reader/formats/haptools_vcf.py +358 -0
  22. rfmix_reader-0.7.0/rfmix_reader/formats/rfmix_fb.py +125 -0
  23. rfmix_reader-0.7.0/rfmix_reader/formats/rfmix_msp.py +137 -0
  24. rfmix_reader-0.7.0/rfmix_reader/io/__init__.py +28 -0
  25. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/io/prepare_reference.py +17 -4
  26. rfmix_reader-0.7.0/rfmix_reader/ops/__init__.py +10 -0
  27. rfmix_reader-0.7.0/rfmix_reader/ops/bed.py +105 -0
  28. rfmix_reader-0.7.0/rfmix_reader/ops/interpolate.py +104 -0
  29. rfmix_reader-0.7.0/rfmix_reader/ops/parquet.py +73 -0
  30. rfmix_reader-0.7.0/rfmix_reader/ops/positions.py +208 -0
  31. rfmix_reader-0.7.0/rfmix_reader/ops/tagore.py +24 -0
  32. rfmix_reader-0.7.0/rfmix_reader/processing/__init__.py +62 -0
  33. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/processing/imputation.py +182 -89
  34. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/processing/phase.py +335 -403
  35. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/viz/__init__.py +0 -2
  36. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/viz/tagore.py +13 -19
  37. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/viz/visualization.py +44 -101
  38. rfmix_reader-0.3.0/PKG-INFO +0 -498
  39. rfmix_reader-0.3.0/README.md +0 -453
  40. rfmix_reader-0.3.0/rfmix_reader/__init__.py +0 -103
  41. rfmix_reader-0.3.0/rfmix_reader/_testing.py +0 -69
  42. rfmix_reader-0.3.0/rfmix_reader/cli/__init__.py +0 -4
  43. rfmix_reader-0.3.0/rfmix_reader/cli/create_binaries.py +0 -33
  44. rfmix_reader-0.3.0/rfmix_reader/io/__init__.py +0 -14
  45. rfmix_reader-0.3.0/rfmix_reader/io/chunk.py +0 -34
  46. rfmix_reader-0.3.0/rfmix_reader/io/errors.py +0 -53
  47. rfmix_reader-0.3.0/rfmix_reader/io/loci_bed.py +0 -431
  48. rfmix_reader-0.3.0/rfmix_reader/io/write_data.py +0 -446
  49. rfmix_reader-0.3.0/rfmix_reader/processing/__init__.py +0 -20
  50. rfmix_reader-0.3.0/rfmix_reader/processing/_hmm_lai.py +0 -333
  51. rfmix_reader-0.3.0/rfmix_reader/readers/__init__.py +0 -8
  52. rfmix_reader-0.3.0/rfmix_reader/readers/fb_read.py +0 -106
  53. rfmix_reader-0.3.0/rfmix_reader/readers/read_flare.py +0 -458
  54. rfmix_reader-0.3.0/rfmix_reader/readers/read_rfmix.py +0 -416
  55. rfmix_reader-0.3.0/rfmix_reader/readers/read_simu.py +0 -423
  56. rfmix_reader-0.3.0/rfmix_reader/utils.py +0 -656
  57. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/LICENSE +0 -0
  58. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/base.svg.p +0 -0
  59. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/cli/merge_phased_zarrs.py +0 -0
  60. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/cli/prepare_reference.py +0 -0
  61. {rfmix_reader-0.3.0 → rfmix_reader-0.7.0}/rfmix_reader/processing/constants.py +0 -0
@@ -0,0 +1,288 @@
1
+ Metadata-Version: 2.4
2
+ Name: rfmix-reader
3
+ Version: 0.7.0
4
+ 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.
5
+ License: GPL-3.0-or-later
6
+ License-File: LICENSE
7
+ Keywords: file parser,rfmix,gpu acceleration,local ancestry
8
+ Author: Kynon J.M. Benjamin
9
+ Author-email: kj.benjamin90@gmail.com
10
+ Maintainer: Kynon J.M. Benjamin
11
+ Maintainer-email: kj.benjamin90@gmail.com
12
+ Requires-Python: >=3.11,<3.15
13
+ Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)
14
+ Classifier: Operating System :: OS Independent
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Programming Language :: Python :: 3.14
20
+ Provides-Extra: all
21
+ Provides-Extra: docs
22
+ Provides-Extra: gpu
23
+ Provides-Extra: io
24
+ Provides-Extra: reference
25
+ Provides-Extra: viz
26
+ Requires-Dist: bio2zarr ; extra == "reference" or extra == "all"
27
+ Requires-Dist: cairosvg (>=2.7,<3.0) ; extra == "viz" or extra == "all"
28
+ Requires-Dist: cudf-cu12 ; extra == "gpu" or extra == "all"
29
+ Requires-Dist: cupy-cuda12x ; extra == "gpu" or extra == "all"
30
+ Requires-Dist: cyvcf2 (>=0.31)
31
+ Requires-Dist: dask (>=2025.1,<2026.0)
32
+ Requires-Dist: dask-cudf-cu12 ; extra == "gpu" or extra == "all"
33
+ Requires-Dist: matplotlib (>=3.10,<4.0) ; extra == "viz" or extra == "all"
34
+ Requires-Dist: numpy (>=1.23,<3)
35
+ Requires-Dist: pandas (>=2.0)
36
+ Requires-Dist: psutil (>=7,<8)
37
+ Requires-Dist: pyarrow ; extra == "io" or extra == "all"
38
+ Requires-Dist: seaborn (>=0.13,<0.14) ; extra == "viz" or extra == "all"
39
+ Requires-Dist: sphinx (>=7,<9) ; extra == "docs"
40
+ Requires-Dist: sphinx-autodoc-typehints (>=2,<3) ; extra == "docs"
41
+ Requires-Dist: sphinx-copybutton (>=0.5,<0.6) ; extra == "docs"
42
+ Requires-Dist: sphinx-rtd-theme (>=2,<3) ; extra == "docs"
43
+ Requires-Dist: torch (>=2.8) ; extra == "gpu" or extra == "all"
44
+ Requires-Dist: tqdm (>=4.66)
45
+ Requires-Dist: xarray (>=2024.1)
46
+ Requires-Dist: zarr (>=3,<4)
47
+ Project-URL: Bug Tracker, https://github.com/heart-gen/rfmix_reader/issues
48
+ Project-URL: homepage, https://rfmix-reader.readthedocs.io/en/latest/
49
+ Project-URL: repository, https://github.com/heart-gen/rfmix_reader.git
50
+ Description-Content-Type: text/markdown
51
+
52
+ # rfmix_reader
53
+
54
+ [![Tests](https://github.com/heart-gen/rfmix_reader/actions/workflows/tests.yml/badge.svg)](https://github.com/heart-gen/rfmix_reader/actions/workflows/tests.yml)
55
+ [![Documentation](https://readthedocs.org/projects/rfmix-reader/badge/?version=latest)](https://rfmix-reader.readthedocs.io/en/latest/)
56
+ [![DOI](https://zenodo.org/badge/807052842.svg)](https://zenodo.org/doi/10.5281/zenodo.12629787)
57
+
58
+ Fast, memory-light access to local-ancestry output — **RFMix** (`.msp.tsv`,
59
+ `.fb.tsv`), **FLARE** (`.anc.vcf.gz`) and **haptools** simulations — as one
60
+ lazily-evaluated [`xarray.Dataset`](https://docs.xarray.dev) backed by a
61
+ per-chromosome Zarr cache.
62
+
63
+ - One streaming pass over the source (constant memory), then instant reopen.
64
+ - Haplotype ancestry codes stored as int8 (2 bytes per sample-locus, a few MB
65
+ per chromosome on disk); posteriors kept on request.
66
+ - The same Dataset for every format, with position queries, BED intervals,
67
+ Parquet export, interpolation onto a variant grid, and gnomix-style phasing.
68
+
69
+ ## Installation
70
+
71
+ ```bash
72
+ pip install rfmix-reader # core: numpy, pandas, dask, xarray, zarr, cyvcf2
73
+ pip install "rfmix-reader[viz,io]" # + matplotlib/seaborn/cairosvg, pyarrow
74
+ ```
75
+
76
+ | Extra | Adds | Needed for |
77
+ |---|---|---|
78
+ | `viz` | matplotlib, seaborn, cairosvg | `plot_*`, `ds.la.to_tagore` |
79
+ | `io` | pyarrow | `ds.la.to_parquet` |
80
+ | `reference` | bio2zarr | `prepare-reference` (VCF to VCF-Zarr for `method="reference"` phasing) |
81
+ | `gpu` | torch, cupy-cuda12x, cudf-cu12, dask-cudf-cu12 | optional CuPy compute backend |
82
+ | `all` | everything above | |
83
+
84
+ Parsing is CPU-only by design (it is I/O bound). When CuPy is importable the
85
+ imputation and plotting helpers use it as the array backend
86
+ (`rfmix_reader.backends.use_gpu()`).
87
+
88
+ ## Quickstart
89
+
90
+ ```python
91
+ from rfmix_reader import open_rfmix, open_local_ancestry
92
+
93
+ # parse once (one streaming pass), cache as la_cache/<chrom>.zarr
94
+ ds = open_rfmix("two_pops/out/", cache_dir="la_cache/")
95
+
96
+ # later sessions: instant, lazy
97
+ ds = open_local_ancestry("la_cache/")
98
+ ds = open_local_ancestry("la_cache/", chrom="21")
99
+
100
+ ds.la.counts # (variant, sample, ancestry) int8 diploid counts, lazy dask
101
+ ds.la.haplotypes # (variant, sample, ploidy) int8 ancestry codes
102
+ ds.la.global_ancestry # DataFrame: sample_id, <ancestries>, chrom
103
+ ds.la.samples, ds.la.ancestries, ds.la.chromosomes
104
+ ```
105
+
106
+ ```bash
107
+ rfmix-reader convert fb two_pops/out/ la_cache/ --keep-posteriors
108
+ rfmix-reader info la_cache/
109
+ ```
110
+
111
+ ## Readers
112
+
113
+ | Source | Call | Notes |
114
+ |---|---|---|
115
+ | RFMix `.msp.tsv` (+ `.rfmix.Q`) | `open_rfmix(path)` | default; segment-level hard calls, `segment_end` coordinate |
116
+ | RFMix `.fb.tsv` (+ `.rfmix.Q`) | `open_rfmix(path, source="fb", keep_posteriors=True)` | one pass over the text; hard calls are the per-haplotype argmax, posteriors optional |
117
+ | FLARE `.anc.vcf.gz` (+ `.global.anc.gz`) | `open_flare(path)` | `AN1`/`AN2` codes in `##ANCESTRY` order |
118
+ | haptools `simgenotype --pop_field` | `open_simu(path)` | tabix regions in parallel; sorted labels (haptools defines no order) |
119
+ | existing cache | `open_local_ancestry(cache_dir, chrom=None)` | concatenates `<chrom>.zarr` stores |
120
+ | write cache only | `convert(path, fmt, cache_dir)` | `fmt` in `msp`, `fb`, `flare`, `haptools` |
121
+
122
+ `path` may be a directory, a single file, or a path prefix
123
+ (`/out/run_` matches `/out/run_chr1.msp.tsv`, ...). `chrom="21"` restricts to
124
+ one chromosome. Without `cache_dir` the Dataset is built in memory, which is
125
+ fine for `.msp.tsv`; for `.fb.tsv` with posteriors use a cache (a chr1-sized
126
+ cohort of 500 samples is ~11 GB of float32 posteriors).
127
+
128
+ ## The Dataset
129
+
130
+ ```text
131
+ dims: variant, sample, ploidy (=2), ancestry, contig
132
+ vars: haplotype_ancestry (variant, sample, ploidy) int8 code into `ancestry`, -1 missing
133
+ posterior (variant, sample, ploidy, ancestry) float32 optional
134
+ global_ancestry (contig, sample, ancestry) float32
135
+ coords: chromosome, variant_position, segment_end (variant); sample_id; ancestry; contig
136
+ attrs: source_format, source_files, rfmix_reader_version
137
+ ```
138
+
139
+ Conventions that hold for every reader:
140
+
141
+ - Ancestries are in the **tool's own order** (RFMix `#reference_panel_population`
142
+ / `#Subpopulation order/codes`, FLARE `##ANCESTRY`), and `global_ancestry`
143
+ uses the same order.
144
+ - A sample/locus with no call (e.g. an all-zero RFMix posterior) is `-1` in the
145
+ codes and in every count of that row.
146
+ - Counts are derived lazily from the codes (`ds.la.counts`), so they are never
147
+ stored twice.
148
+
149
+ `ds.la` is the accessor (registered on import of `rfmix_reader.core`, which
150
+ every `open_*` call does):
151
+
152
+ | Accessor | Returns |
153
+ |---|---|
154
+ | `ds.la.counts` / `.haplotypes` / `.posterior` | lazy DataArrays (posterior may be `None`) |
155
+ | `ds.la.global_ancestry` | long DataFrame (`sample_id`, ancestries, `chrom`) |
156
+ | `ds.la.samples`, `.ancestries`, `.chromosomes`, `.n_variants`, `.n_samples` | labels / sizes |
157
+ | `ds.la.sel_chrom("chr21")`, `ds.la.sel_region("chr21", start, end)` | contiguous Dataset slices (binary search on the sorted positions) |
158
+ | `ds.la.to_legacy()` | the legacy (0.5 and earlier) `(loci_df, g_anc, local_array)` triple |
159
+
160
+ ## Operations
161
+
162
+ All operations are lazy or streaming; `sample` is an index or a sample ID.
163
+
164
+ | Operation | Result |
165
+ |---|---|
166
+ | `ds.la.at_positions(loci_df, method="stepwise"\|"nearest", samples=None, aggregate=True)` | ancestry at listed `chrom`/`pos` (haplotype counts and fractions, or per-sample copies with `aggregate=False`) |
167
+ | `ds.la.locus_index(chrom, positions, method="stepwise"\|"nearest", tolerance=None)` | index of the variant/segment covering each position (`-1` = none) |
168
+ | `ds.la.counts_at(chrom, positions, ...)` | `(n, sample, ancestry)` int8 counts at those positions |
169
+ | `ds.la.to_bed(sample, min_segment=1)` | constant-ancestry intervals for one sample |
170
+ | `ds.la.to_tagore(sample, palette="tab10")` | the BED annotated for TAGORE; plot with `plot_local_ancestry_tagore` |
171
+ | `ds.la.to_parquet(outdir, prefix=..., rows_per_file=...)` | `<prefix>.<chrom>-<k>.parquet` (`chrom`, `pos`, `hap`, then `<sample>_<ancestry>` int8 columns), one dask block at a time |
172
+ | `ds.la.interpolate(variants_df, zarr_outdir, method="linear"\|"nearest"\|"stepwise")` | counts on a denser variant grid, Zarr-backed per chromosome; missing calls are filled |
173
+ | `ds.la.phase(config=PhasingConfig())` | phase-corrected haplotype codes (single chromosome) |
174
+
175
+ ```python
176
+ import pandas as pd
177
+
178
+ loci = pd.DataFrame({"chrom": ["chr21", "chr21"], "pos": [15_000_000, 30_000_000]})
179
+ ds.la.at_positions(loci) # AFR_haplotypes, AFR_fraction, ...
180
+ ds.la.at_positions(loci, samples=["NA19700"], aggregate=False)
181
+
182
+ # QTL-style access: one chromosome in memory, then index by genotype position
183
+ chrom = ds.la.sel_chrom("chr21") # contiguous slice, still lazy
184
+ counts = chrom.la.counts.values # (variants, samples, ancestries) int8
185
+ idx = chrom.la.locus_index("chr21", genotype_positions) # segment per SNP, -1 = uncovered
186
+ H = counts[idx[idx >= 0]]
187
+
188
+ bed = ds.la.to_bed("NA19700", min_segment=3)
189
+ ds.la.to_parquet("out/", prefix="la", rows_per_file=100_000)
190
+
191
+ variants = pd.read_parquet("genotypes/variants.parquet") # chrom, pos
192
+ dense = ds.la.interpolate(variants, "imputed/", method="stepwise")
193
+ ```
194
+
195
+ ## Phasing
196
+
197
+ `ds.la.phase()` corrects switch errors between the two haplotypes per sample
198
+ the way gnomix does: heterozygous blocks are found from the haplotype codes,
199
+ each window inside a block is scored by how well the two posterior tracks match
200
+ the block-start orientation versus the swapped one, uninformative windows
201
+ inherit the previous state, and the two haplotypes (codes and posteriors) are
202
+ exchanged wherever the track says "switched". It needs no reference panel; open
203
+ `.fb.tsv` output with `keep_posteriors=True` so the posteriors are available
204
+ (without them the hard calls are used). The result adds a
205
+ `phase_swapped (variant, sample)` mask.
206
+
207
+ ```python
208
+ from rfmix_reader import open_rfmix
209
+ from rfmix_reader.processing.phase import PhasingConfig, merge_phased_zarrs
210
+
211
+ ds = open_rfmix("two_pops/out/", source="fb", keep_posteriors=True,
212
+ cache_dir="la_cache/", chrom="21")
213
+ phased = ds.la.phase(config=PhasingConfig(window_size=50, min_block_len=20, posterior_margin=0.2))
214
+ phased["phase_swapped"].sum("variant") # exchanged loci per sample
215
+ ```
216
+
217
+ Whole chromosomes phase in seconds (chr21, 175k loci by 500 samples: about
218
+ 4 s). `phase_rfmix_chromosome_to_zarr(prefix, None, None, "phased_chr21.zarr", chrom="21")`
219
+ does read, phase and write in one call, and `merge_phased_zarrs` concatenates
220
+ per-chromosome stores; both outputs reopen with `open_local_ancestry`.
221
+
222
+ The previous reference-panel matcher remains as `method="reference"` (with
223
+ `ref_zarr_root` / `sample_annot_path`) for comparison only: it compares
224
+ ancestry labels against reference *allele* codes, which is not a sound test
225
+ of phase. To build its inputs, convert bgzipped, indexed reference VCFs with
226
+ `prepare-reference refs/ 1kg_chr21.vcf.gz` (needs the `reference` extra) and
227
+ give a two-column `sample_id<TAB>group` annotation table.
228
+
229
+ ## Haptools simulations
230
+
231
+ haptools does **not** write the chromosome length into the `##contig` header,
232
+ but the tabix region pulls need it. Reheader each file with the contig entry
233
+ from the `contigs.txt` haptools produces, for example:
234
+
235
+ ```bash
236
+ CONTIG_LINE=$(grep -w "ID=${CHR}" contigs.txt)
237
+ bcftools view -h "$IN" | sed "s/^##contig=<ID=${CHR}>.*/${CONTIG_LINE}/" > header.${CHR}.tmp
238
+ bcftools reheader -h header.${CHR}.tmp -o "$OUT" "$IN"
239
+ tabix -p vcf "$OUT"
240
+ ```
241
+
242
+ ## Visualization
243
+
244
+ ```python
245
+ from rfmix_reader import plot_global_ancestry, plot_ancestry_by_chromosome, plot_local_ancestry_tagore
246
+
247
+ g_anc = ds.la.global_ancestry
248
+ plot_global_ancestry(g_anc, save_path="global") # writes global.png and global.pdf
249
+ plot_ancestry_by_chromosome(g_anc, save_path="by_chrom")
250
+ plot_local_ancestry_tagore(ds.la.to_tagore("NA19700"), prefix="NA19700", build="hg38", oformat="png")
251
+ ```
252
+
253
+ ## Migrating from 0.x
254
+
255
+ See [MIGRATION.md](MIGRATION.md). Every legacy (0.5 and earlier) name maps to one call on the
256
+ Dataset; `ds.la.to_legacy()` returns the old triple and
257
+ `rfmix_reader.from_legacy(loci_df, g_anc, admix)` builds a Dataset from one.
258
+
259
+ ## Development
260
+
261
+ ```bash
262
+ git clone https://github.com/heart-gen/rfmix_reader.git
263
+ cd rfmix_reader
264
+ poetry install --with test --extras "viz io"
265
+ poetry run pytest # fast suite (seconds)
266
+ poetry run pytest --run-slow # also the chr21 tests (needs git-LFS data)
267
+ ```
268
+
269
+ Test fixtures for every format live under `tests/data/` (regenerate with
270
+ `python tests/data/make_fixtures.py`).
271
+
272
+ ## Citation
273
+
274
+ If you use this software, please cite:
275
+
276
+ Benjamin, K. J. M. (2024). **RFMix-reader** \[Computer software].
277
+ [https://github.com/heart-gen/rfmix\_reader](https://github.com/heart-gen/rfmix_reader)
278
+
279
+ Kynon JM Benjamin. *"RFMix-reader: Accelerated reading and processing for local ancestry studies."*
280
+ **bioRxiv** (2024).
281
+ DOI: [10.1101/2024.07.13.603370](https://www.biorxiv.org/content/10.1101/2024.07.13.603370v2).
282
+
283
+ ## Funding
284
+
285
+ This work was supported by the National Institutes of Health,
286
+ National Institute on Minority Health and Health Disparities (NIMHD)
287
+ K99MD016964 / R00MD016964.
288
+
@@ -0,0 +1,236 @@
1
+ # rfmix_reader
2
+
3
+ [![Tests](https://github.com/heart-gen/rfmix_reader/actions/workflows/tests.yml/badge.svg)](https://github.com/heart-gen/rfmix_reader/actions/workflows/tests.yml)
4
+ [![Documentation](https://readthedocs.org/projects/rfmix-reader/badge/?version=latest)](https://rfmix-reader.readthedocs.io/en/latest/)
5
+ [![DOI](https://zenodo.org/badge/807052842.svg)](https://zenodo.org/doi/10.5281/zenodo.12629787)
6
+
7
+ Fast, memory-light access to local-ancestry output — **RFMix** (`.msp.tsv`,
8
+ `.fb.tsv`), **FLARE** (`.anc.vcf.gz`) and **haptools** simulations — as one
9
+ lazily-evaluated [`xarray.Dataset`](https://docs.xarray.dev) backed by a
10
+ per-chromosome Zarr cache.
11
+
12
+ - One streaming pass over the source (constant memory), then instant reopen.
13
+ - Haplotype ancestry codes stored as int8 (2 bytes per sample-locus, a few MB
14
+ per chromosome on disk); posteriors kept on request.
15
+ - The same Dataset for every format, with position queries, BED intervals,
16
+ Parquet export, interpolation onto a variant grid, and gnomix-style phasing.
17
+
18
+ ## Installation
19
+
20
+ ```bash
21
+ pip install rfmix-reader # core: numpy, pandas, dask, xarray, zarr, cyvcf2
22
+ pip install "rfmix-reader[viz,io]" # + matplotlib/seaborn/cairosvg, pyarrow
23
+ ```
24
+
25
+ | Extra | Adds | Needed for |
26
+ |---|---|---|
27
+ | `viz` | matplotlib, seaborn, cairosvg | `plot_*`, `ds.la.to_tagore` |
28
+ | `io` | pyarrow | `ds.la.to_parquet` |
29
+ | `reference` | bio2zarr | `prepare-reference` (VCF to VCF-Zarr for `method="reference"` phasing) |
30
+ | `gpu` | torch, cupy-cuda12x, cudf-cu12, dask-cudf-cu12 | optional CuPy compute backend |
31
+ | `all` | everything above | |
32
+
33
+ Parsing is CPU-only by design (it is I/O bound). When CuPy is importable the
34
+ imputation and plotting helpers use it as the array backend
35
+ (`rfmix_reader.backends.use_gpu()`).
36
+
37
+ ## Quickstart
38
+
39
+ ```python
40
+ from rfmix_reader import open_rfmix, open_local_ancestry
41
+
42
+ # parse once (one streaming pass), cache as la_cache/<chrom>.zarr
43
+ ds = open_rfmix("two_pops/out/", cache_dir="la_cache/")
44
+
45
+ # later sessions: instant, lazy
46
+ ds = open_local_ancestry("la_cache/")
47
+ ds = open_local_ancestry("la_cache/", chrom="21")
48
+
49
+ ds.la.counts # (variant, sample, ancestry) int8 diploid counts, lazy dask
50
+ ds.la.haplotypes # (variant, sample, ploidy) int8 ancestry codes
51
+ ds.la.global_ancestry # DataFrame: sample_id, <ancestries>, chrom
52
+ ds.la.samples, ds.la.ancestries, ds.la.chromosomes
53
+ ```
54
+
55
+ ```bash
56
+ rfmix-reader convert fb two_pops/out/ la_cache/ --keep-posteriors
57
+ rfmix-reader info la_cache/
58
+ ```
59
+
60
+ ## Readers
61
+
62
+ | Source | Call | Notes |
63
+ |---|---|---|
64
+ | RFMix `.msp.tsv` (+ `.rfmix.Q`) | `open_rfmix(path)` | default; segment-level hard calls, `segment_end` coordinate |
65
+ | RFMix `.fb.tsv` (+ `.rfmix.Q`) | `open_rfmix(path, source="fb", keep_posteriors=True)` | one pass over the text; hard calls are the per-haplotype argmax, posteriors optional |
66
+ | FLARE `.anc.vcf.gz` (+ `.global.anc.gz`) | `open_flare(path)` | `AN1`/`AN2` codes in `##ANCESTRY` order |
67
+ | haptools `simgenotype --pop_field` | `open_simu(path)` | tabix regions in parallel; sorted labels (haptools defines no order) |
68
+ | existing cache | `open_local_ancestry(cache_dir, chrom=None)` | concatenates `<chrom>.zarr` stores |
69
+ | write cache only | `convert(path, fmt, cache_dir)` | `fmt` in `msp`, `fb`, `flare`, `haptools` |
70
+
71
+ `path` may be a directory, a single file, or a path prefix
72
+ (`/out/run_` matches `/out/run_chr1.msp.tsv`, ...). `chrom="21"` restricts to
73
+ one chromosome. Without `cache_dir` the Dataset is built in memory, which is
74
+ fine for `.msp.tsv`; for `.fb.tsv` with posteriors use a cache (a chr1-sized
75
+ cohort of 500 samples is ~11 GB of float32 posteriors).
76
+
77
+ ## The Dataset
78
+
79
+ ```text
80
+ dims: variant, sample, ploidy (=2), ancestry, contig
81
+ vars: haplotype_ancestry (variant, sample, ploidy) int8 code into `ancestry`, -1 missing
82
+ posterior (variant, sample, ploidy, ancestry) float32 optional
83
+ global_ancestry (contig, sample, ancestry) float32
84
+ coords: chromosome, variant_position, segment_end (variant); sample_id; ancestry; contig
85
+ attrs: source_format, source_files, rfmix_reader_version
86
+ ```
87
+
88
+ Conventions that hold for every reader:
89
+
90
+ - Ancestries are in the **tool's own order** (RFMix `#reference_panel_population`
91
+ / `#Subpopulation order/codes`, FLARE `##ANCESTRY`), and `global_ancestry`
92
+ uses the same order.
93
+ - A sample/locus with no call (e.g. an all-zero RFMix posterior) is `-1` in the
94
+ codes and in every count of that row.
95
+ - Counts are derived lazily from the codes (`ds.la.counts`), so they are never
96
+ stored twice.
97
+
98
+ `ds.la` is the accessor (registered on import of `rfmix_reader.core`, which
99
+ every `open_*` call does):
100
+
101
+ | Accessor | Returns |
102
+ |---|---|
103
+ | `ds.la.counts` / `.haplotypes` / `.posterior` | lazy DataArrays (posterior may be `None`) |
104
+ | `ds.la.global_ancestry` | long DataFrame (`sample_id`, ancestries, `chrom`) |
105
+ | `ds.la.samples`, `.ancestries`, `.chromosomes`, `.n_variants`, `.n_samples` | labels / sizes |
106
+ | `ds.la.sel_chrom("chr21")`, `ds.la.sel_region("chr21", start, end)` | contiguous Dataset slices (binary search on the sorted positions) |
107
+ | `ds.la.to_legacy()` | the legacy (0.5 and earlier) `(loci_df, g_anc, local_array)` triple |
108
+
109
+ ## Operations
110
+
111
+ All operations are lazy or streaming; `sample` is an index or a sample ID.
112
+
113
+ | Operation | Result |
114
+ |---|---|
115
+ | `ds.la.at_positions(loci_df, method="stepwise"\|"nearest", samples=None, aggregate=True)` | ancestry at listed `chrom`/`pos` (haplotype counts and fractions, or per-sample copies with `aggregate=False`) |
116
+ | `ds.la.locus_index(chrom, positions, method="stepwise"\|"nearest", tolerance=None)` | index of the variant/segment covering each position (`-1` = none) |
117
+ | `ds.la.counts_at(chrom, positions, ...)` | `(n, sample, ancestry)` int8 counts at those positions |
118
+ | `ds.la.to_bed(sample, min_segment=1)` | constant-ancestry intervals for one sample |
119
+ | `ds.la.to_tagore(sample, palette="tab10")` | the BED annotated for TAGORE; plot with `plot_local_ancestry_tagore` |
120
+ | `ds.la.to_parquet(outdir, prefix=..., rows_per_file=...)` | `<prefix>.<chrom>-<k>.parquet` (`chrom`, `pos`, `hap`, then `<sample>_<ancestry>` int8 columns), one dask block at a time |
121
+ | `ds.la.interpolate(variants_df, zarr_outdir, method="linear"\|"nearest"\|"stepwise")` | counts on a denser variant grid, Zarr-backed per chromosome; missing calls are filled |
122
+ | `ds.la.phase(config=PhasingConfig())` | phase-corrected haplotype codes (single chromosome) |
123
+
124
+ ```python
125
+ import pandas as pd
126
+
127
+ loci = pd.DataFrame({"chrom": ["chr21", "chr21"], "pos": [15_000_000, 30_000_000]})
128
+ ds.la.at_positions(loci) # AFR_haplotypes, AFR_fraction, ...
129
+ ds.la.at_positions(loci, samples=["NA19700"], aggregate=False)
130
+
131
+ # QTL-style access: one chromosome in memory, then index by genotype position
132
+ chrom = ds.la.sel_chrom("chr21") # contiguous slice, still lazy
133
+ counts = chrom.la.counts.values # (variants, samples, ancestries) int8
134
+ idx = chrom.la.locus_index("chr21", genotype_positions) # segment per SNP, -1 = uncovered
135
+ H = counts[idx[idx >= 0]]
136
+
137
+ bed = ds.la.to_bed("NA19700", min_segment=3)
138
+ ds.la.to_parquet("out/", prefix="la", rows_per_file=100_000)
139
+
140
+ variants = pd.read_parquet("genotypes/variants.parquet") # chrom, pos
141
+ dense = ds.la.interpolate(variants, "imputed/", method="stepwise")
142
+ ```
143
+
144
+ ## Phasing
145
+
146
+ `ds.la.phase()` corrects switch errors between the two haplotypes per sample
147
+ the way gnomix does: heterozygous blocks are found from the haplotype codes,
148
+ each window inside a block is scored by how well the two posterior tracks match
149
+ the block-start orientation versus the swapped one, uninformative windows
150
+ inherit the previous state, and the two haplotypes (codes and posteriors) are
151
+ exchanged wherever the track says "switched". It needs no reference panel; open
152
+ `.fb.tsv` output with `keep_posteriors=True` so the posteriors are available
153
+ (without them the hard calls are used). The result adds a
154
+ `phase_swapped (variant, sample)` mask.
155
+
156
+ ```python
157
+ from rfmix_reader import open_rfmix
158
+ from rfmix_reader.processing.phase import PhasingConfig, merge_phased_zarrs
159
+
160
+ ds = open_rfmix("two_pops/out/", source="fb", keep_posteriors=True,
161
+ cache_dir="la_cache/", chrom="21")
162
+ phased = ds.la.phase(config=PhasingConfig(window_size=50, min_block_len=20, posterior_margin=0.2))
163
+ phased["phase_swapped"].sum("variant") # exchanged loci per sample
164
+ ```
165
+
166
+ Whole chromosomes phase in seconds (chr21, 175k loci by 500 samples: about
167
+ 4 s). `phase_rfmix_chromosome_to_zarr(prefix, None, None, "phased_chr21.zarr", chrom="21")`
168
+ does read, phase and write in one call, and `merge_phased_zarrs` concatenates
169
+ per-chromosome stores; both outputs reopen with `open_local_ancestry`.
170
+
171
+ The previous reference-panel matcher remains as `method="reference"` (with
172
+ `ref_zarr_root` / `sample_annot_path`) for comparison only: it compares
173
+ ancestry labels against reference *allele* codes, which is not a sound test
174
+ of phase. To build its inputs, convert bgzipped, indexed reference VCFs with
175
+ `prepare-reference refs/ 1kg_chr21.vcf.gz` (needs the `reference` extra) and
176
+ give a two-column `sample_id<TAB>group` annotation table.
177
+
178
+ ## Haptools simulations
179
+
180
+ haptools does **not** write the chromosome length into the `##contig` header,
181
+ but the tabix region pulls need it. Reheader each file with the contig entry
182
+ from the `contigs.txt` haptools produces, for example:
183
+
184
+ ```bash
185
+ CONTIG_LINE=$(grep -w "ID=${CHR}" contigs.txt)
186
+ bcftools view -h "$IN" | sed "s/^##contig=<ID=${CHR}>.*/${CONTIG_LINE}/" > header.${CHR}.tmp
187
+ bcftools reheader -h header.${CHR}.tmp -o "$OUT" "$IN"
188
+ tabix -p vcf "$OUT"
189
+ ```
190
+
191
+ ## Visualization
192
+
193
+ ```python
194
+ from rfmix_reader import plot_global_ancestry, plot_ancestry_by_chromosome, plot_local_ancestry_tagore
195
+
196
+ g_anc = ds.la.global_ancestry
197
+ plot_global_ancestry(g_anc, save_path="global") # writes global.png and global.pdf
198
+ plot_ancestry_by_chromosome(g_anc, save_path="by_chrom")
199
+ plot_local_ancestry_tagore(ds.la.to_tagore("NA19700"), prefix="NA19700", build="hg38", oformat="png")
200
+ ```
201
+
202
+ ## Migrating from 0.x
203
+
204
+ See [MIGRATION.md](MIGRATION.md). Every legacy (0.5 and earlier) name maps to one call on the
205
+ Dataset; `ds.la.to_legacy()` returns the old triple and
206
+ `rfmix_reader.from_legacy(loci_df, g_anc, admix)` builds a Dataset from one.
207
+
208
+ ## Development
209
+
210
+ ```bash
211
+ git clone https://github.com/heart-gen/rfmix_reader.git
212
+ cd rfmix_reader
213
+ poetry install --with test --extras "viz io"
214
+ poetry run pytest # fast suite (seconds)
215
+ poetry run pytest --run-slow # also the chr21 tests (needs git-LFS data)
216
+ ```
217
+
218
+ Test fixtures for every format live under `tests/data/` (regenerate with
219
+ `python tests/data/make_fixtures.py`).
220
+
221
+ ## Citation
222
+
223
+ If you use this software, please cite:
224
+
225
+ Benjamin, K. J. M. (2024). **RFMix-reader** \[Computer software].
226
+ [https://github.com/heart-gen/rfmix\_reader](https://github.com/heart-gen/rfmix_reader)
227
+
228
+ Kynon JM Benjamin. *"RFMix-reader: Accelerated reading and processing for local ancestry studies."*
229
+ **bioRxiv** (2024).
230
+ DOI: [10.1101/2024.07.13.603370](https://www.biorxiv.org/content/10.1101/2024.07.13.603370v2).
231
+
232
+ ## Funding
233
+
234
+ This work was supported by the National Institutes of Health,
235
+ National Institute on Minority Health and Health Disparities (NIMHD)
236
+ K99MD016964 / R00MD016964.
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "rfmix-reader"
3
- version = "0.3.0" # This will be overridden if dynamic versioning is enabled
3
+ version = "0.7.0" # This will be overridden if dynamic versioning is enabled
4
4
  description = "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."
5
5
  authors = ["Kynon J.M. Benjamin <kj.benjamin90@gmail.com>"]
6
6
  maintainers = ["Kynon J.M. Benjamin <kj.benjamin90@gmail.com>"]
@@ -22,25 +22,33 @@ tqdm = ">=4.66"
22
22
  psutil = "^7"
23
23
  cyvcf2 = ">=0.31"
24
24
  dask = "^2025.1"
25
+ zarr = ">=3,<4"
26
+ xarray = ">=2024.1"
25
27
  # Plotting
26
28
  matplotlib = { version = "^3.10", optional = true }
27
29
  seaborn = { version = "^0.13", optional = true }
28
30
  cairosvg = { version = "^2.7", optional = true }
29
31
  # GPU
30
32
  torch = { version = ">=2.8", optional = true }
33
+ cupy-cuda12x = { version = "*", optional = true }
34
+ cudf-cu12 = { version = "*", optional = true }
35
+ dask-cudf-cu12 = { version = "*", optional = true }
36
+ pyarrow = { version = "*", optional = true }
37
+ # Reference preparation
38
+ bio2zarr = { version = "*", optional = true }
31
39
  # Docs
32
40
  sphinx = { version = ">=7,<9", optional = true }
33
41
  sphinx-rtd-theme = { version = "^2", optional = true }
34
42
  sphinx-autodoc-typehints = { version = "^2", optional = true }
35
43
  sphinx-copybutton = { version = "^0.5", optional = true }
36
- fonttools = "^4.61.0"
37
44
 
38
45
  [tool.poetry.extras]
39
46
  viz = ["matplotlib", "seaborn", "cairosvg"]
40
- gpu = ["cupy", "cudf", "dask-cudf", "torch"]
47
+ gpu = ["torch", "cupy-cuda12x", "cudf-cu12", "dask-cudf-cu12"]
41
48
  io = ["pyarrow"] # parquet export
49
+ reference = ["bio2zarr"] # prepare-reference CLI
50
+ all = ["matplotlib", "seaborn", "cairosvg", "torch", "cupy-cuda12x", "cudf-cu12", "dask-cudf-cu12", "pyarrow", "bio2zarr"]
42
51
  docs = ["sphinx", "sphinx-rtd-theme", "sphinx-autodoc-typehints", "sphinx-copybutton"]
43
- tests = ["pytest"]
44
52
 
45
53
  [tool.poetry.urls]
46
54
  homepage = "https://rfmix-reader.readthedocs.io/en/latest/"
@@ -48,13 +56,22 @@ repository = "https://github.com/heart-gen/rfmix_reader.git"
48
56
  "Bug Tracker" = "https://github.com/heart-gen/rfmix_reader/issues"
49
57
 
50
58
  [tool.poetry.scripts]
51
- create-binaries = "rfmix_reader.cli.create_binaries:main"
59
+ rfmix-reader = "rfmix_reader.cli.main:main"
52
60
  prepare-reference = "rfmix_reader.cli.prepare_reference:main"
53
61
  merge-phased-zarrs = "rfmix_reader.cli.merge_phased_zarrs:main"
54
62
 
55
63
  [tool.poetry.group.test.dependencies]
56
- pytest = "^8"
64
+ pytest = ">=8,<10"
57
65
  pytest-cov = "^7"
66
+ pysam = ">=0.22"
67
+ pyarrow = "*"
68
+
69
+ [tool.pytest.ini_options]
70
+ testpaths = ["tests"]
71
+ markers = [
72
+ "slow: long-running tests that need the LFS chr21 data (run with --run-slow)",
73
+ "cuda: tests that need a CUDA GPU (run with --run-cuda-tests)",
74
+ ]
58
75
 
59
76
  [tool.poetry.group.docs.dependencies]
60
77
  sphinx = ">=7,<9"