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

{rfmix_reader-0.3.0 → rfmix_reader-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rfmix-reader
-Version: 0.3.0
+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
@@ -17,25 +17,30 @@ 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"
+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"
+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: seaborn (>=0.13,<0.14) ; extra == "viz"
+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"
+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/
@@ -53,7 +58,7 @@ for major speedups when available.
 
 ## Installation
 
-`rfmix-reader` requires **Python 3.10+**. Install from PyPI:
+`rfmix-reader` requires **Python 3.11+**. Install from PyPI:
 
 ```bash
 pip install rfmix-reader
@@ -67,12 +72,16 @@ pip install rfmix-reader
   pip install rfmix-reader
   ```
 
-* **With GPU acceleration** (`cupy`, `cudf`, `dask-cudf`):
+* **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
@@ -87,8 +96,11 @@ pip install rfmix-reader
 
 ### GPU Notes
 
-* `torch` is installed automatically.
-* For CUDA builds, install a matching GPU-enabled wheel for your system following the [PyTorch guide](https://pytorch.org/get-started/locally/).
+* 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.
 
@@ -218,8 +230,8 @@ Binaries can also be generated on-the-fly within `read_rfmix` with
 
 ```python
 loci, g_anc, admix = read_rfmix("examples/three_populations/out/",
-                               binary_dir="./binary_files",
-                               generate_binary=True)
+                                binary_dir="./binary_files",
+                                generate_binary=True)
 ```
 
 ### Phasing data
@@ -241,10 +253,12 @@ admix = phase_rfmix_chromosome_to_zarr(
 ```
 
 The chunking is suboptimal for phasing, so remember to
-rechunk before using for optimal processing.
+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": 50})
+local_array = admix["local_ancestry"].chunk({"variant": 20000, "sample": 100})
 # Compute into memory, if needed
 local_array = local_array.compute()
 ```

{rfmix_reader-0.3.0 → rfmix_reader-0.3.1}/README.md

@@ -9,7 +9,7 @@ for major speedups when available.
 
 ## Installation
 
-`rfmix-reader` requires **Python 3.10+**. Install from PyPI:
+`rfmix-reader` requires **Python 3.11+**. Install from PyPI:
 
 ```bash
 pip install rfmix-reader
@@ -23,12 +23,16 @@ pip install rfmix-reader
   pip install rfmix-reader
   ```
 
-* **With GPU acceleration** (`cupy`, `cudf`, `dask-cudf`):
+* **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
@@ -43,8 +47,11 @@ pip install rfmix-reader
 
 ### GPU Notes
 
-* `torch` is installed automatically.
-* For CUDA builds, install a matching GPU-enabled wheel for your system following the [PyTorch guide](https://pytorch.org/get-started/locally/).
+* 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.
 
@@ -174,8 +181,8 @@ Binaries can also be generated on-the-fly within `read_rfmix` with
 
 ```python
 loci, g_anc, admix = read_rfmix("examples/three_populations/out/",
-                               binary_dir="./binary_files",
-                               generate_binary=True)
+                                binary_dir="./binary_files",
+                                generate_binary=True)
 ```
 
 ### Phasing data
@@ -197,10 +204,12 @@ admix = phase_rfmix_chromosome_to_zarr(
 ```
 
 The chunking is suboptimal for phasing, so remember to
-rechunk before using for optimal processing.
+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": 50})
+local_array = admix["local_ancestry"].chunk({"variant": 20000, "sample": 100})
 # Compute into memory, if needed
 local_array = local_array.compute()
 ```

{rfmix_reader-0.3.0 → rfmix_reader-0.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rfmix-reader"
-version = "0.3.0"  # This will be overridden if dynamic versioning is enabled
+version = "0.3.1"  # This will be overridden if dynamic versioning is enabled
 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."
 authors = ["Kynon J.M. Benjamin <kj.benjamin90@gmail.com>"]
 maintainers = ["Kynon J.M. Benjamin <kj.benjamin90@gmail.com>"]
@@ -28,6 +28,10 @@ seaborn = { version = "^0.13", optional = true }
 cairosvg = { version = "^2.7", optional = true }
 # GPU
 torch = { version = ">=2.8", optional = true }
+cupy-cuda12x = { version = "*", optional = true }
+cudf-cu12 = { version = "*", optional = true }
+dask-cudf-cu12 = { version = "*", optional = true }
+pyarrow = { version = "*", optional = true }
 # Docs
 sphinx                    = { version = ">=7,<9", optional = true }
 sphinx-rtd-theme          = { version = "^2", optional = true }
@@ -39,6 +43,7 @@ fonttools = "^4.61.0"
 viz = ["matplotlib", "seaborn", "cairosvg"]
 gpu = ["cupy", "cudf", "dask-cudf", "torch"]
 io  = ["pyarrow"]  # parquet export
+all = ["matplotlib", "seaborn", "cairosvg", "torch", "cupy-cuda12x", "cudf-cu12", "dask-cudf-cu12", "pyarrow"]
 docs = ["sphinx", "sphinx-rtd-theme", "sphinx-autodoc-typehints", "sphinx-copybutton"]
 tests = ["pytest"]
 

{rfmix_reader-0.3.0 → rfmix_reader-0.3.1}/rfmix_reader/processing/imputation.py

@@ -275,9 +275,6 @@ def interpolate_array(
         Batch size for processing local ancestry data. Default is 10,000.
     interpolation : {"linear","nearest","stepwise"}, default "linear"
         Interpolation scheme.
-        ``"hmm"`` is available for experimentation but requires
-        `allow_hmm=True` and haplotype-level inputs created with
-        `split_to_haplotypes`.
     use_bp_positions : bool, default False
         If True, use `variant_loci_df['pos']` as the x-axis for interpolation.
         If False, loci are treated as equally spaced (index-based).

{rfmix_reader-0.3.0 → rfmix_reader-0.3.1}/rfmix_reader/processing/phase.py

@@ -1240,10 +1240,22 @@ def merge_phased_zarrs(
     if sort:
         combined = combined.sortby(["chromosome", "variant_position"])
 
+    for v in combined.variables:
+        var = combined[v]
+        var.encoding.pop("chunks", None)
+        if var.dtype == object:
+            combined[v] = var.astype("S50")
+
+    # Ensure string coordinates are encoded consistently when writing to Zarr.
+    for coord in ("chromosome", "sample_id", "ancestry"):
+        if coord in combined.coords:
+            combined = combined.assign_coords({coord: combined[coord].astype(str)})
+
     logger.info(
         "[merge_phased_zarrs] Writing merged dataset with %d variants to %s",
-        combined.dims.get("variant", 0),
+        combined.sizes.get("variant", 0),
         output_path,
     )
+    combined = combined.chunk({"variant": 50000})
     combined.to_zarr(output_path, mode="w")
     return combined