genoray 0.1.0__tar.gz

genoray-0.1.0/.gitattributes

@@ -0,0 +1,2 @@
+# SCM syntax highlighting & preventing 3-way merges
+pixi.lock merge=binary linguist-language=YAML linguist-generated=true

genoray-0.1.0/.gitignore

@@ -0,0 +1,8 @@
+.pytest_cache/
+__pycache__/
+notebooks/
+.venv/
+
+# pixi environments
+.pixi
+*.egg-info

genoray-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,19 @@
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v5.0.0
+  hooks:
+  - id: check-merge-conflict
+  - id: debug-statements
+  - id: mixed-line-ending
+  - id: check-case-conflict
+  - id: check-yaml
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.9.4
+  hooks:
+  - id: ruff
+  - id: ruff-format
+- repo: https://github.com/commitizen-tools/commitizen
+  rev: v4.4.1
+  hooks:
+  - id: commitizen
+    stages: [commit-msg]

genoray-0.1.0/.python-version

@@ -0,0 +1 @@
+3.9

genoray-0.1.0/CHANGELOG.md

@@ -0,0 +1,16 @@
+## 0.1.0 (2025-04-12)
+
+### Feat
+
+- sketching out support for PGEN dosages
+- refactor readers to be type safe. pass all tests.
+- **wip**: reasonable output from PGEN in notebook
+- initial PGEN support
+- rename package to genoray
+- rename package to genoray
+- **wip**: initial prototype of VCF reader
+- **wip**: VCF support
+
+### Fix
+
+- use future annotations for union types

genoray-0.1.0/LICENSE.txt

@@ -0,0 +1,7 @@
+Copyright (c) 2025 David Laub
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

genoray-0.1.0/PKG-INFO

@@ -0,0 +1,26 @@
+Metadata-Version: 2.4
+Name: genoray
+Version: 0.1.0
+Summary: Add your description here
+Author-email: David Laub <dlaub@ucsd.edu>
+License: Copyright (c) 2025 David Laub
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+License-File: LICENSE.txt
+Requires-Python: >=3.9
+Requires-Dist: cyvcf2>=0.31.1
+Requires-Dist: hirola>=0.3.0
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: pgenlib>=0.92.0
+Requires-Dist: phantom-types>=3.0.2
+Requires-Dist: polars>=1.27.1
+Requires-Dist: pyarrow>=19.0.1
+Requires-Dist: pyranges==0.1.3
+Requires-Dist: setuptools>=78.1.0
+Requires-Dist: tqdm>=4.67.1
+Requires-Dist: typing-extensions>=4.13.2

genoray-0.1.0/genoray/__init__.py

@@ -0,0 +1,5 @@
+from ._pgen import PGEN
+from ._types import Reader
+from ._vcf import VCF
+
+__all__ = ["Reader", "VCF", "PGEN"]

genoray-0.1.0/genoray/_pgen.py

@@ -0,0 +1,414 @@
+from __future__ import annotations
+
+from functools import partial
+from pathlib import Path
+from typing import Generator, TypeVar, cast
+
+import numpy as np
+import pgenlib
+import polars as pl
+import pyranges as pr
+from hirola import HashTable
+from numpy.typing import ArrayLike, NDArray
+from phantom import Phantom
+from typing_extensions import Self, TypeGuard, assert_never
+
+from ._types import Reader
+from ._utils import (
+    ContigNormalizer,
+    format_memory,
+    is_dtype,
+    lengths_to_offsets,
+    parse_memory,
+)
+
+
+def _is_genos_dosages(obj) -> TypeGuard[tuple[Genos, Dosages]]:
+    """Check if the object is a tuple of genotypes and dosages.
+
+    Parameters
+    ----------
+    obj
+        Object to check.
+
+    Returns
+    -------
+    bool
+        True if the object is a tuple of genotypes and dosages, False otherwise.
+    """
+    return (
+        isinstance(obj, tuple)
+        and len(obj) == 2
+        and isinstance(obj[0], Genos)
+        and isinstance(obj[1], Dosages)
+    )
+
+
+class Genos(
+    NDArray[np.int32], Phantom, predicate=partial(is_dtype, dtype=np.int32)
+): ...
+
+
+class Dosages(
+    NDArray[np.float32], Phantom, predicate=partial(is_dtype, dtype=np.float32)
+): ...
+
+
+class GenosDosages(tuple[Genos, Dosages], Phantom, predicate=_is_genos_dosages): ...
+
+
+T = TypeVar(
+    "T",
+    Genos,
+    Dosages,
+    GenosDosages,
+)
+
+
+class PGEN(Reader[T]):
+    available_samples: list[str]
+    filter: pl.Expr | None
+    ploidy = 2
+    contigs: list[str]
+    _index: pr.PyRanges
+    _geno_pgen: pgenlib.PgenReader
+    _dose_pgen: pgenlib.PgenReader
+    _s_idx: NDArray[np.uint32]
+    _read_as: type[T]
+
+    Genos = Genos
+    Dosages = Dosages
+    GenosDosages = GenosDosages
+
+    def __init__(
+        self,
+        geno_path: str | Path,
+        filter: pl.Expr | None = None,
+        read_as: type[T] = Genos,
+        dosage_path: str | Path | None = None,
+    ):
+        # TODO: support dosages and allow user to either provide a second PGEN file for dosages
+        # or else use the same PGEN file for both genotypes and dosages.
+        # That being said, there's probably not much point for a user to use the same PGEN file
+        # for genos and dosages since PLINK2 defines hardcalls as a simple threshold on the dosages
+        # when dosages are available.
+        if read_as is Dosages or read_as is GenosDosages:
+            raise NotImplementedError("PGEN dosages are not yet supported.")
+
+        geno_path = Path(geno_path)
+        samples = _read_psam(geno_path.with_suffix(".psam"))
+
+        self.filter = filter
+        self.available_samples = samples.tolist()
+        self._s2i = HashTable(
+            max=len(samples) * 2,  # type: ignore
+            dtype=samples.dtype,
+        )
+        self._s2i.add(samples)
+        self._s_idx = np.arange(len(samples), dtype=np.uint32)
+        self._geno_pgen = pgenlib.PgenReader(bytes(geno_path))
+
+        if dosage_path is not None:
+            dosage_path = Path(dosage_path)
+            dose_samples = _read_psam(dosage_path.with_suffix(".psam"))
+            if (samples != dose_samples).any():
+                raise ValueError(
+                    "Samples in dosage file do not match those in genotype file."
+                )
+            self._dose_pgen = pgenlib.PgenReader(bytes(Path(dosage_path)))
+        else:
+            self._dose_pgen = self._geno_pgen
+
+        if not geno_path.with_suffix(".gvi").exists():
+            _write_index(geno_path.with_suffix(".pvar"))
+        self._index = _read_index(geno_path.with_suffix(".gvi"), self.filter)
+        self.contigs = self._index.chromosomes
+        self._c_norm = ContigNormalizer(self._index.chromosomes)
+        self._read_as = read_as
+
+    @property
+    def current_samples(self) -> list[str]:
+        return self._s2i.keys[self._s_idx].tolist()
+
+    def set_samples(self, samples: list[str]) -> Self:
+        _samples = np.atleast_1d(samples)
+        s_idx = self._s2i.get(_samples).astype(np.uint32)
+        if (missing := _samples[s_idx == -1]).any():
+            raise ValueError(f"Samples {missing} not found in the file.")
+        self._s_idx = s_idx
+        self._geno_pgen.change_sample_subset(np.sort(s_idx))
+        return self
+
+    def __del__(self):
+        self._geno_pgen.close()
+        if self._dose_pgen is not None:
+            self._dose_pgen.close()
+
+    def n_vars_in_ranges(
+        self,
+        contig: str,
+        starts: ArrayLike = 0,
+        ends: ArrayLike | None = None,
+    ) -> NDArray[np.uint32]:
+        c = self._c_norm.norm(contig)
+        if c is None:
+            return np.zeros_like(np.atleast_1d(starts), dtype=np.uint32)
+
+        starts = np.atleast_1d(starts)
+        if ends is None:
+            ends = np.full_like(starts, np.iinfo(np.int32).max)
+        queries = pr.PyRanges(
+            pl.DataFrame(
+                {
+                    "Chromosome": np.full_like(starts, contig),
+                    "Start": starts,
+                    "End": ends,
+                }
+            ).to_pandas(use_pyarrow_extension_array=True)
+        )
+        return (
+            queries.count_overlaps(self._index)
+            .df["NumberOverlaps"]
+            .to_numpy()
+            .astype(np.uint32)
+        )
+
+    def _var_idxs(
+        self, contig: str, starts: ArrayLike = 0, ends: ArrayLike | None = None
+    ) -> tuple[NDArray[np.uint32], NDArray[np.uint64]]:
+        """Get variant indices and the number of indices per region.
+
+        Parameters
+        ----------
+        contig
+            Contig name.
+        starts
+            0-based start positions of the regions.
+        ends
+            0-based, exclusive end positions of the regions.
+
+        Returns
+        -------
+        idxs
+            Shape: (tot_variants). Variant indices for the given ranges.
+        offsets
+            Shape: (regions+1). Offsets to get variant indices for each region.
+        """
+        starts = np.atleast_1d(starts)
+
+        c = self._c_norm.norm(contig)
+        if c is None:
+            return np.empty(0, np.uint32), np.zeros_like(
+                np.atleast_1d(starts), np.uint64
+            )
+
+        starts = np.atleast_1d(starts)
+        if ends is None:
+            ends = np.full_like(starts, np.iinfo(np.int32).max)
+        queries = pr.PyRanges(
+            pl.DataFrame(
+                {
+                    "Chromosome": np.full_like(starts, contig),
+                    "Start": starts,
+                    "End": ends,
+                }
+            )
+            .with_row_index("query")
+            .to_pandas(use_pyarrow_extension_array=True)
+        )
+        join = pl.from_pandas(queries.join(self._index).df)
+        if join.height == 0:
+            return np.empty(0, np.uint32), np.zeros_like(
+                np.atleast_1d(starts), np.uint64
+            )
+        join = join.sort("query", "index")
+        idxs = join["index"].to_numpy()
+        lens = (
+            join.group_by("query", maintain_order=True).agg(pl.len())["len"].to_numpy()
+        )
+        offsets = lengths_to_offsets(lens)
+        return idxs, offsets
+
+    def read(
+        self,
+        contig: str,
+        start: int = 0,
+        end: int | None = None,
+        out: T | None = None,
+    ) -> T | None:
+        c = self._c_norm.norm(contig)
+        if c is None:
+            return
+
+        if end is None:
+            end = np.iinfo(np.int64).max
+
+        var_idxs, _ = self._var_idxs(c, start, end)
+        n_variants = len(var_idxs)
+        if n_variants == 0:
+            return
+
+        # TODO: support dosages
+
+        if out is None:
+            data = np.empty((n_variants, self.n_samples * self.ploidy), dtype=np.int32)
+        else:
+            if not isinstance(out, Genos):
+                raise ValueError(f"Expected a np.int32 array, got {type(out)}.")
+            data = out
+
+        self._geno_pgen.read_alleles_list(var_idxs, data)
+        data = data.reshape(n_variants, self.n_samples, self.ploidy).transpose(1, 2, 0)[
+            self._s_idx
+        ]
+        data[data == -9] = -1
+
+        data = cast(T, data)
+
+        return data
+
+    def read_chunks(
+        self,
+        contig: str,
+        start: int = 0,
+        end: int | None = None,
+        max_mem: int | str = "4g",
+    ) -> Generator[T]:
+        # TODO: support dosages
+
+        max_mem = parse_memory(max_mem)
+
+        c = self._c_norm.norm(contig)
+        if c is None:
+            return
+
+        if end is None:
+            end = np.iinfo(np.int64).max
+
+        var_idxs, _ = self._var_idxs(c, start, end)
+        n_variants = len(var_idxs)
+        if n_variants == 0:
+            return
+
+        mem_per_v = self._mem_per_variant()
+        vars_per_chunk = min(max_mem // mem_per_v, n_variants)
+        if vars_per_chunk == 0:
+            raise ValueError(
+                f"Maximum memory {format_memory(max_mem)} insufficient to read a single variant."
+                f" Memory per variant: {format_memory(mem_per_v)}."
+            )
+
+        n_chunks = -(-n_variants // vars_per_chunk)
+        v_chunks = np.array_split(var_idxs, n_chunks)
+        for var_idx in v_chunks:
+            chunk_size = len(var_idx)
+            out = np.empty((chunk_size, self.n_samples * self.ploidy), dtype=np.int32)
+            self._geno_pgen.read_alleles_list(var_idx, out)
+            out = out.reshape(chunk_size, self.n_samples, self.ploidy).transpose(
+                1, 2, 0
+            )[self._s_idx]
+            out[out == -9] = -1
+            yield cast(T, out)
+
+    def read_ranges(
+        self,
+        contig: str,
+        starts: ArrayLike = 0,
+        ends: ArrayLike | None = None,
+    ) -> tuple[T, NDArray[np.uint32]] | None:
+        # TODO: support dosages
+
+        starts = np.atleast_1d(starts)
+
+        c = self._c_norm.norm(contig)
+        if c is None:
+            return
+
+        var_idxs, offsets = self._var_idxs(c, starts, ends)
+        n_variants = len(var_idxs)
+        if n_variants == 0:
+            return
+
+        out = np.empty((n_variants, self.n_samples * self.ploidy), dtype=np.int32)
+
+        self._geno_pgen.read_alleles_list(var_idxs, out)
+        out = out.reshape(n_variants, self.n_samples, self.ploidy).transpose(1, 2, 0)[
+            self._s_idx
+        ]
+        out[out == -9] = -1
+
+        return cast(T, out), np.diff(offsets).astype(np.uint32)
+
+    def _mem_per_variant(self) -> int:
+        if issubclass(self._read_as, Genos):
+            return self.n_samples * self.ploidy * np.int32().itemsize
+        elif issubclass(self._read_as, (Dosages, GenosDosages)):
+            raise NotImplementedError("Dosages are not yet supported.")
+        else:
+            assert_never(self._read_as)
+
+
+def _read_psam(path: Path) -> NDArray[np.str_]:
+    with open(path.with_suffix(".psam")) as f:
+        cols = [c.strip("#") for c in f.readline().strip().split()]
+
+    psam = pl.read_csv(
+        path.with_suffix(".psam"),
+        separator="\t",
+        has_header=False,
+        skip_rows=1,
+        new_columns=cols,
+        schema_overrides={
+            "FID": pl.Utf8,
+            "IID": pl.Utf8,
+            "SID": pl.Utf8,
+            "PAT": pl.Utf8,
+            "MAT": pl.Utf8,
+            "SEX": pl.Utf8,
+        },
+    )
+    samples = psam["IID"].to_numpy().astype(str)
+    return samples
+
+
+RLEN = pl.col("REF").str.len_bytes()
+ALEN = pl.col("ALT").str.len_bytes()
+ILEN = ALEN - RLEN
+KIND = (
+    pl.when(ILEN != 0)
+    .then(pl.lit("INDEL"))
+    .when(RLEN == 1)
+    .then(pl.lit("SNP"))
+    .otherwise(pl.lit("MNP"))
+    .cast(pl.Categorical)
+)
+
+
+# TODO: index can likely be implemented using the NCLS lib underlying PyRanges and then we can
+# pass np.memmap arrays directly instead of having to futz with DataFrames. This will likely make
+# filtering less ergonomic/harder to make ergonomic though, but a memmap approach will be scalable
+# to ultra-large datasets (100k+ individuals).
+def _write_index(path: Path):
+    (
+        pl.scan_csv(
+            path.with_suffix(".pvar"),
+            separator="\t",
+            comment_prefix="##",
+            schema_overrides={"#CHROM": pl.Utf8, "POS": pl.Int32},
+        )
+        .select(
+            Chromosome="#CHROM",
+            Start=pl.col("POS") - 1,
+            End=pl.col("POS") + RLEN - 1,
+            kind=KIND,
+        )
+        .sink_ipc(path.with_suffix(".gvi"))
+    )
+
+
+def _read_index(path: Path, filter: pl.Expr | None) -> pr.PyRanges:
+    index = pl.read_ipc(path, row_index_name="index", memory_map=False)
+    if filter is not None:
+        index = index.filter(filter)
+    pyr = pr.PyRanges(index.drop("kind").to_pandas(use_pyarrow_extension_array=True))
+    return pyr

genoray-0.1.0/genoray/_types.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from typing import Any, Generator, Generic, Protocol, TypeVar
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+from typing_extensions import Self
+
+T = TypeVar("T")
+
+
+class Reader(Protocol, Generic[T]):
+    available_samples: list[str]
+    """All samples in the file, in the order they exist on-disk."""
+    ploidy: int
+    filter: Any | None
+    contigs: list[str]
+
+    @property
+    def current_samples(self) -> list[str]:
+        """The samples this reader will return, in order along the sample axis."""
+        ...
+
+    def set_samples(self, samples: list[str]) -> Self:
+        """Set the samples this reader will return, in order along the sample axis."""
+        ...
+
+    @property
+    def n_samples(self) -> int:
+        return len(self.current_samples)
+
+    def n_vars_in_ranges(
+        self, contig: str, starts: ArrayLike = 0, ends: ArrayLike | None = None
+    ) -> NDArray[np.uint32]:
+        """Return the start and end indices of the variants in the given ranges.
+
+        Parameters
+        ----------
+        contig
+            Contig name.
+        starts
+            0-based start positions of the regions.
+        ends
+            0-based, exclusive end positions of the regions.
+
+        Returns
+        -------
+        n_variants
+            Shape: (regions). Number of variants in the given ranges.
+        """
+        ...
+
+    def read(
+        self,
+        contig: str,
+        start: int = 0,
+        end: int | None = None,
+        out: T | None = None,
+    ) -> T | None:
+        """Read genotypes and/or dosages for a region.
+
+        Parameters
+        ----------
+        contig
+            Contig name.
+        start
+            0-based start position of the region.
+        end
+            0-based, exclusive end position of the region.
+        samples
+            Samples to read. If None, all samples are read.
+        ploids
+            Ploids to read. If None, all ploids are read.
+        dosage_field
+            Dosage field to read. If True, use the default dosage field for the format.
+
+        Returns
+        -------
+        data
+            Genotypes and/or dosages. Genotypes have shape (samples ploidy variants) and
+            dosages have shape (samples variants). Missing genotypes have value -1 and missing dosages
+            have value np.nan. If just using genotypes or dosages, will be a single array, otherwise
+            will be a tuple of arrays.
+        dosage
+            Shape: (samples variants)
+        """
+        ...
+
+    def read_chunks(
+        self,
+        contig: str,
+        start: int = 0,
+        end: int | None = None,
+        max_mem: int | str = "4g",
+    ) -> Generator[T]:
+        """Iterate over genotypes and/or dosages for a region in chunks limited by max_mem.
+
+        Parameters
+        ----------
+        contig
+            Contig name.
+        start
+            0-based start position.
+        end
+            0-based, exclusive end position of the region.
+        samples
+            Samples to read. If None, all samples are read.
+        ploids
+            Ploids to read. If None, all ploids are read.
+
+        Returns
+        -------
+        data
+            Generator of genotypes and/or dosages. Genotypes have shape (samples ploidy variants) and
+            dosages have shape (samples variants). Missing genotypes have value -1 and missing dosages
+            have value np.nan. If just using genotypes or dosages, will be a single array, otherwise
+            will be a tuple of arrays.
+        """
+        ...
+
+    def read_ranges(
+        self,
+        contig: str,
+        starts: ArrayLike = 0,
+        ends: ArrayLike | None = None,
+    ) -> tuple[T, NDArray[np.uint32]] | None:
+        """Read genotypes and/or dosages for multiple regions.
+
+        Parameters
+        ----------
+        contig
+            Contig name.
+        start
+            0-based start position of the region.
+        end
+            0-based, exclusive end position of the region.
+        samples
+            Samples to read. If None, all samples are read.
+        ploids
+            Ploids to read. If None, all ploids are read.
+        dosage_field
+            Dosage field to read. If True, use the default dosage field for the format.
+
+        Returns
+        -------
+        data
+            Genotypes and/or dosages. Genotypes have shape (samples ploidy variants) and
+            dosages have shape (samples variants). Missing genotypes have value -1 and missing dosages
+            have value np.nan. If just using genotypes or dosages, will be a single array, otherwise
+            will be a tuple of arrays.
+        n_variants_per_region
+            Shape: (regions). Number of variants in the given ranges.
+        """
+        ...