localqtl 0.1.1__tar.gz

localqtl-0.1.1/PKG-INFO

@@ -0,0 +1,221 @@
+Metadata-Version: 2.4
+Name: localqtl
+Version: 0.1.1
+Summary: GPU-accelerated ancestry-aware cis-QTL mapping library with tensorQTL-compatible APIs
+License: GPL-3.0-or-later
+Keywords: genomics,gpu,xQTL,qtl
+Author: Kynon J Benjamin
+Author-email: kj.benjamin90@gmail.com
+Maintainer: Kynon J 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: 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
+Requires-Dist: cudf-cu12 (>=25.10,<26.0)
+Requires-Dist: cupy-cuda12x (>=13.6,<14.0)
+Requires-Dist: dask-cudf-cu12 (>=25.10.0,<26.0.0)
+Requires-Dist: numpy (>=2.0,<3.0)
+Requires-Dist: pandas (>=2.2,<3.0)
+Requires-Dist: pandas-plink-ng (>=2.3.2.post1,<3.0.0)
+Requires-Dist: pgenlib (>=0.93,<0.94)
+Requires-Dist: py-qvalue (>=0.1,<0.2)
+Requires-Dist: pyarrow (>=22.0,<23.0)
+Requires-Dist: rfmix-reader (>=0.2,<0.3)
+Requires-Dist: scipy (>=1.16,<2.0)
+Requires-Dist: torch (>=2.8,<3.0)
+Project-URL: Bug Tracker, https://github.com/heart-gen/localQTL/issues
+Project-URL: homepage, https://github.com/heart-gen/localQTL
+Project-URL: repository, https://github.com/heart-gen/localQTL
+Description-Content-Type: text/markdown
+
+# localQTL
+
+localQTL is a **pure-Python** library for local-ancestry-aware xQTL mapping that
+lets researchers run end-to-end analyses on large cohorts **without any R/rpy2
+dependencies**. It preserves the familiar 
+[tensorQTL](https://github.com/broadinstitute/tensorqtl) data model with **GPU-first
+execution paths**, flexible genotype loaders, and streaming outputs for large-scale
+workflows, while adding ancestry-aware use cases.
+
+## Features
+
+- **GPU-accelerated workflows** powered by PyTorch, with automatic fallbacks to CPU
+  execution when CUDA is unavailable.
+- **Modular cis-QTL mapping API** exposed through functional helpers such as
+  `map_nominal`, `map_permutations`, and `map_independent`, or via the convenience
+  wrapper `CisMapper`.
+- **Multiple genotype backends** including PLINK (BED/BIM/FAM), PLINK 2 (PGEN/PSAM/PVAR),
+  and BED/Parquet inputs, with helpers to stream data in manageable windows.
+- **Local ancestry integration** by pairing genotype windows with haplotype panels
+  produced by RFMix through `RFMixReader` and `InputGeneratorCisWithHaps`.
+- **Parquet streaming sinks** that make it easy to materialise association statistics
+  without loading the entire result set in memory.
+- **Pure-Python statistics (no R/rpy2 required)**: tensorQTL's `rfunc` calls have been refactored to scipy.stats for p-values and q-values are computed with the Python port of Storey's `qvalue` (`py-qvalue`).
+
+## Installation
+
+The project uses [Poetry](https://python-poetry.org/) for dependency management.
+Clone the repository and install the package into a virtual environment:
+
+```bash
+poetry install
+```
+
+If you prefer `pip`, you can install the library in editable mode after exporting
+Poetry's dependency specification:
+
+```bash
+pip install -e .
+```
+
+> **Note:** GPU acceleration relies on PyTorch, CuPy, and cuDF. Make sure you use
+> versions that match the CUDA toolkit available on your system. The versions in
+> `pyproject.toml` target CUDA 12.
+
+## Quickstart
+
+Below is a minimal example that runs a nominal cis-QTL scan against PLINK-formatted
+genotypes and BED-formatted phenotypes. The example mirrors the data layout
+expected by tensorQTL, so **existing preprocessing pipelines can be reused**.
+
+### Standard mapping (tensorQTL equivalent)
+```python
+from localqtl import PlinkReader, read_phenotype_bed
+from localqtl.cis import map_nominal
+
+# Load genotypes and variant metadata
+plink = PlinkReader("data/genotypes")
+genotype_df = plink.load_genotypes()
+variant_df = plink.bim.set_index("snp")[["chrom", "pos"]]
+
+# Load phenotypes (BED-style) and their genomic coordinates
+phenotype_df, phenotype_pos_df = read_phenotype_bed("data/phenotypes.bed")
+
+# Optional: load covariates as a DataFrame indexed by sample IDs
+covariates_df = None
+
+results = map_nominal(
+    genotype_df=genotype_df,
+    variant_df=variant_df,
+    phenotype_df=phenotype_df,
+    phenotype_pos_df=phenotype_pos_df,
+    covariates_df=covariates_df,
+    window=1_000_000,         # ±1 Mb cis window
+    maf_threshold=0.01,       # filter on in-sample MAF
+    device="auto",            # picks CUDA when available, otherwise CPU
+    out_prefix="cis_nominal", # default prefix
+    return_df=True            # default is False, parquet streamed sink
+)
+
+print(results.head())
+```
+
+For analyses that combine nominal scans, permutations, and independent signal
+calling, the `CisMapper` class offers a thin object-oriented façade:
+
+```python
+from localqtl.cis import CisMapper
+
+mapper = CisMapper(
+    genotype_df=genotype_df,
+    variant_df=variant_df,
+    phenotype_df=phenotype_df,
+    phenotype_pos_df=phenotype_pos_df,
+    covariates_df=covariates_df,
+    window=500_000,
+)
+
+mapper.map_nominal(nperm=1_000)
+perm_df = mapper.map_permutations(nperm=1_000, beta_approx=True)
+perm_df = mapper.calculate_qvalues(perm_df, fdr=0.05)
+lead_df = mapper.map_independent(cis_df=perm_df, fdr=0.05)
+```
+
+### Local ancestry-aware mapping
+
+localQTL can incorporate local ancestry (e.g., from RFMix) so that cis-xQTL tests are performed with ancestry-aware genotype inputs. To enable this, pass **`haplotypes`** (per-variant × per-sample × per-ancestry dosages) and **`loci_df`** (variant positions corresponding to the haplotype tensor) to `CisMapper`. When these are provided, `CisMapper` switches to the ancestry-aware input generator under the hood.
+
+```python
+from localqtl import PlinkReader, read_phenotype_bed
+from localqtl.haplotypeio import RFMixReader
+from localqtl.cis import CisMapper
+
+plink = PlinkReader("data/genotypes")
+genotype_df = plink.load_genotypes()  # samples as columns, variants as rows
+variant_df = plink.bim.set_index("snp")[["chrom", "pos"]]
+
+phenotype_df, phenotype_pos_df = read_phenotype_bed("data/phenotypes.bed")
+covariates_df = None  # optional
+
+# Local ancestry from RFMix
+rfmix = RFMixReader(
+    prefix_path="data/rfmix/prefix",   # directory with per-chrom outputs + fb.tsv
+    binary_path="data/rfmix",          # where prebuilt binaries live (if used)
+    verbose=True
+)
+
+# Materialize ancestry haplotypes into memory (NumPy array)
+# Shape: (variants, samples, ancestries) for >2 ancestries.
+# For 2 ancestries, the reader exposes the ancestry channel selected internally.
+H = rfmix.load_haplotypes()
+loci_df = rfmix.loci_df  # DataFrame with ['chrom','pos', 'ancestry', 'hap', 'index'] (indexed by 'hap')
+
+# Align samples to the RFMix order (recommended)
+# Keep only intersection and put everything in the same order as rfmix.sample_ids.
+keep = [sid for sid in rfmix.sample_ids if sid in genotype_df.columns]
+genotype_df = genotype_df[keep]
+phenotype_df = phenotype_df[keep]
+if covariates_df is not None:
+    covariates_df = covariates_df.loc[keep]
+
+# (Optional) Ensure chromosome dtype matches between variant_df and loci_df
+variant_df["chrom"] = variant_df["chrom"].astype(str)
+loci_df = loci_df.copy()
+loci_df["chrom"] = loci_df["chrom"].astype(str)
+
+# Ancestry-aware mapping
+mapper = CisMapper(
+    genotype_df=genotype_df,
+    variant_df=variant_df,
+    phenotype_df=phenotype_df,
+    phenotype_pos_df=phenotype_pos_df,
+    covariates_df=covariates_df,
+    haplotypes=H,          # <-- enable local ancestry-aware mode
+    loci_df=loci_df,       # <-- positions that align to H
+    window=1_000_000,
+    device="auto",
+)
+
+# Run nominal scans and permutations as usual; ancestry-awareness is automatic
+mapper.map_nominal(nperm=0)
+perm_df = mapper.map_permutations(nperm=1_000, beta_approx=True)
+
+# FDR using the pure-Python q-value port (no R/rpy2 needed)
+perm_df = mapper.calculate_qvalues(perm_df, fdr=0.05)
+
+# Identify conditionally independent signals
+lead_df = mapper.map_independent(cis_df=perm_df, fdr=0.05)
+
+print(lead_df.head())
+```
+
+**Input expectations**
+
+* `H` (from `RFMixReader.load_haplotypes()`) is ancestry dosages with shape `(variants, samples, ancestries)` for ≥3 ancestries; for 2 ancestries the reader exposes the configured channel.
+* `loci_df` rows correspond 1:1 to `H`'s and be joinable (by `chrom`/`pos`) to `variant_df` used for genotypes.
+* Sample order in `genotype_df`, `phenotype_df`, and `covariates_df` should match `rfmix.sample_ids` (reindex as shown).
+
+## Testing
+
+Run the test suite with:
+
+```bash
+poetry run pytest
+```
+
+This exercises the core cis-QTL mapping routines using small synthetic datasets.
+

localqtl-0.1.1/README.md

@@ -0,0 +1,186 @@
+# localQTL
+
+localQTL is a **pure-Python** library for local-ancestry-aware xQTL mapping that
+lets researchers run end-to-end analyses on large cohorts **without any R/rpy2
+dependencies**. It preserves the familiar 
+[tensorQTL](https://github.com/broadinstitute/tensorqtl) data model with **GPU-first
+execution paths**, flexible genotype loaders, and streaming outputs for large-scale
+workflows, while adding ancestry-aware use cases.
+
+## Features
+
+- **GPU-accelerated workflows** powered by PyTorch, with automatic fallbacks to CPU
+  execution when CUDA is unavailable.
+- **Modular cis-QTL mapping API** exposed through functional helpers such as
+  `map_nominal`, `map_permutations`, and `map_independent`, or via the convenience
+  wrapper `CisMapper`.
+- **Multiple genotype backends** including PLINK (BED/BIM/FAM), PLINK 2 (PGEN/PSAM/PVAR),
+  and BED/Parquet inputs, with helpers to stream data in manageable windows.
+- **Local ancestry integration** by pairing genotype windows with haplotype panels
+  produced by RFMix through `RFMixReader` and `InputGeneratorCisWithHaps`.
+- **Parquet streaming sinks** that make it easy to materialise association statistics
+  without loading the entire result set in memory.
+- **Pure-Python statistics (no R/rpy2 required)**: tensorQTL's `rfunc` calls have been refactored to scipy.stats for p-values and q-values are computed with the Python port of Storey's `qvalue` (`py-qvalue`).
+
+## Installation
+
+The project uses [Poetry](https://python-poetry.org/) for dependency management.
+Clone the repository and install the package into a virtual environment:
+
+```bash
+poetry install
+```
+
+If you prefer `pip`, you can install the library in editable mode after exporting
+Poetry's dependency specification:
+
+```bash
+pip install -e .
+```
+
+> **Note:** GPU acceleration relies on PyTorch, CuPy, and cuDF. Make sure you use
+> versions that match the CUDA toolkit available on your system. The versions in
+> `pyproject.toml` target CUDA 12.
+
+## Quickstart
+
+Below is a minimal example that runs a nominal cis-QTL scan against PLINK-formatted
+genotypes and BED-formatted phenotypes. The example mirrors the data layout
+expected by tensorQTL, so **existing preprocessing pipelines can be reused**.
+
+### Standard mapping (tensorQTL equivalent)
+```python
+from localqtl import PlinkReader, read_phenotype_bed
+from localqtl.cis import map_nominal
+
+# Load genotypes and variant metadata
+plink = PlinkReader("data/genotypes")
+genotype_df = plink.load_genotypes()
+variant_df = plink.bim.set_index("snp")[["chrom", "pos"]]
+
+# Load phenotypes (BED-style) and their genomic coordinates
+phenotype_df, phenotype_pos_df = read_phenotype_bed("data/phenotypes.bed")
+
+# Optional: load covariates as a DataFrame indexed by sample IDs
+covariates_df = None
+
+results = map_nominal(
+    genotype_df=genotype_df,
+    variant_df=variant_df,
+    phenotype_df=phenotype_df,
+    phenotype_pos_df=phenotype_pos_df,
+    covariates_df=covariates_df,
+    window=1_000_000,         # ±1 Mb cis window
+    maf_threshold=0.01,       # filter on in-sample MAF
+    device="auto",            # picks CUDA when available, otherwise CPU
+    out_prefix="cis_nominal", # default prefix
+    return_df=True            # default is False, parquet streamed sink
+)
+
+print(results.head())
+```
+
+For analyses that combine nominal scans, permutations, and independent signal
+calling, the `CisMapper` class offers a thin object-oriented façade:
+
+```python
+from localqtl.cis import CisMapper
+
+mapper = CisMapper(
+    genotype_df=genotype_df,
+    variant_df=variant_df,
+    phenotype_df=phenotype_df,
+    phenotype_pos_df=phenotype_pos_df,
+    covariates_df=covariates_df,
+    window=500_000,
+)
+
+mapper.map_nominal(nperm=1_000)
+perm_df = mapper.map_permutations(nperm=1_000, beta_approx=True)
+perm_df = mapper.calculate_qvalues(perm_df, fdr=0.05)
+lead_df = mapper.map_independent(cis_df=perm_df, fdr=0.05)
+```
+
+### Local ancestry-aware mapping
+
+localQTL can incorporate local ancestry (e.g., from RFMix) so that cis-xQTL tests are performed with ancestry-aware genotype inputs. To enable this, pass **`haplotypes`** (per-variant × per-sample × per-ancestry dosages) and **`loci_df`** (variant positions corresponding to the haplotype tensor) to `CisMapper`. When these are provided, `CisMapper` switches to the ancestry-aware input generator under the hood.
+
+```python
+from localqtl import PlinkReader, read_phenotype_bed
+from localqtl.haplotypeio import RFMixReader
+from localqtl.cis import CisMapper
+
+plink = PlinkReader("data/genotypes")
+genotype_df = plink.load_genotypes()  # samples as columns, variants as rows
+variant_df = plink.bim.set_index("snp")[["chrom", "pos"]]
+
+phenotype_df, phenotype_pos_df = read_phenotype_bed("data/phenotypes.bed")
+covariates_df = None  # optional
+
+# Local ancestry from RFMix
+rfmix = RFMixReader(
+    prefix_path="data/rfmix/prefix",   # directory with per-chrom outputs + fb.tsv
+    binary_path="data/rfmix",          # where prebuilt binaries live (if used)
+    verbose=True
+)
+
+# Materialize ancestry haplotypes into memory (NumPy array)
+# Shape: (variants, samples, ancestries) for >2 ancestries.
+# For 2 ancestries, the reader exposes the ancestry channel selected internally.
+H = rfmix.load_haplotypes()
+loci_df = rfmix.loci_df  # DataFrame with ['chrom','pos', 'ancestry', 'hap', 'index'] (indexed by 'hap')
+
+# Align samples to the RFMix order (recommended)
+# Keep only intersection and put everything in the same order as rfmix.sample_ids.
+keep = [sid for sid in rfmix.sample_ids if sid in genotype_df.columns]
+genotype_df = genotype_df[keep]
+phenotype_df = phenotype_df[keep]
+if covariates_df is not None:
+    covariates_df = covariates_df.loc[keep]
+
+# (Optional) Ensure chromosome dtype matches between variant_df and loci_df
+variant_df["chrom"] = variant_df["chrom"].astype(str)
+loci_df = loci_df.copy()
+loci_df["chrom"] = loci_df["chrom"].astype(str)
+
+# Ancestry-aware mapping
+mapper = CisMapper(
+    genotype_df=genotype_df,
+    variant_df=variant_df,
+    phenotype_df=phenotype_df,
+    phenotype_pos_df=phenotype_pos_df,
+    covariates_df=covariates_df,
+    haplotypes=H,          # <-- enable local ancestry-aware mode
+    loci_df=loci_df,       # <-- positions that align to H
+    window=1_000_000,
+    device="auto",
+)
+
+# Run nominal scans and permutations as usual; ancestry-awareness is automatic
+mapper.map_nominal(nperm=0)
+perm_df = mapper.map_permutations(nperm=1_000, beta_approx=True)
+
+# FDR using the pure-Python q-value port (no R/rpy2 needed)
+perm_df = mapper.calculate_qvalues(perm_df, fdr=0.05)
+
+# Identify conditionally independent signals
+lead_df = mapper.map_independent(cis_df=perm_df, fdr=0.05)
+
+print(lead_df.head())
+```
+
+**Input expectations**
+
+* `H` (from `RFMixReader.load_haplotypes()`) is ancestry dosages with shape `(variants, samples, ancestries)` for ≥3 ancestries; for 2 ancestries the reader exposes the configured channel.
+* `loci_df` rows correspond 1:1 to `H`'s and be joinable (by `chrom`/`pos`) to `variant_df` used for genotypes.
+* Sample order in `genotype_df`, `phenotype_df`, and `covariates_df` should match `rfmix.sample_ids` (reindex as shown).
+
+## Testing
+
+Run the test suite with:
+
+```bash
+poetry run pytest
+```
+
+This exercises the core cis-QTL mapping routines using small synthetic datasets.

localqtl-0.1.1/pyproject.toml

@@ -0,0 +1,51 @@
+[tool.poetry]
+name = "localqtl"
+version = "0.1.1" # placeholder, replaced dynamically
+description = "GPU-accelerated ancestry-aware cis-QTL mapping library with tensorQTL-compatible APIs"
+authors = ["Kynon J Benjamin <kj.benjamin90@gmail.com>"]
+maintainers = ["Kynon J Benjamin <kj.benjamin90@gmail.com>"]
+readme = "README.md"
+license = "GPL-3.0-or-later"
+packages = [{include = "localqtl", from = "src"}]
+keywords = [
+    "genomics", "gpu", "xQTL", "qtl"
+]
+
+[tool.poetry.dependencies]
+python = ">=3.11,<3.15"
+torch = "^2.8"
+numpy = "^2.0"
+scipy = "^1.16"
+pandas = "^2.2"
+pgenlib = "^0.93"
+rfmix-reader = "^0.2"
+py-qvalue = "^0.1"
+cupy-cuda12x = "^13.6"
+cudf-cu12 = "^25.10"
+pyarrow = "^22.0"
+dask-cudf-cu12 = "^25.10.0"
+pandas-plink-ng = "^2.3.2.post1"
+
+[tool.poetry.urls]
+homepage = "https://github.com/heart-gen/localQTL"
+repository = "https://github.com/heart-gen/localQTL"
+"Bug Tracker" = "https://github.com/heart-gen/localQTL/issues"
+
+[tool.poetry-dynamic-versioning]
+enable = false
+style = "pep440"
+dirty = true
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[dependency-groups]
+dev = [
+    "pytest (>=8.4.2,<9.0.0)",
+    "ruff (>=0.14.2,<0.15.0)",
+    "black (>=25.9.0,<26.0.0)",
+    "mypy (>=1.18.2,<2.0.0)",
+    "pytest-cov (>=7.0.0,<8.0.0)",
+    "pytest-xdist (>=3.8.0,<4.0.0)"
+]

localqtl-0.1.1/src/localqtl/__init__.py

@@ -0,0 +1,50 @@
+"""Top-level package exports for :mod:`localqtl`.
+
+This module re-exports the most commonly used classes/functions so that they
+remain accessible from ``localqtl`` after the project was reorganized into a
+package with dedicated submodules (``cis``, ``preproc``, ``finemap`` …).
+"""
+
+from . import cis
+from . import finemap
+from . import genotypeio
+from . import haplotypeio
+from . import iosinks
+from . import pgen
+from . import phenotypeio
+from . import preproc
+from . import regression_kernels
+from . import stats
+from . import utils
+
+from .cis import CisMapper, map_independent, map_nominal, map_permutations
+from .genotypeio import InputGeneratorCis, PlinkReader
+from .haplotypeio import InputGeneratorCisWithHaps, RFMixReader
+from .phenotypeio import read_phenotype_bed
+from .pgen import PgenReader
+from .utils import gpu_available
+
+__all__ = [
+    "map_nominal",
+    "map_permutations",
+    "map_independent",
+    "CisMapper",
+    "read_phenotype_bed",
+    "gpu_available",
+    "PlinkReader",
+    "RFMixReader",
+    "PgenReader",
+    "InputGeneratorCis",
+    "InputGeneratorCisWithHaps",
+    "cis",
+    "finemap",
+    "phenotypeio",
+    "genotypeio",
+    "haplotypeio",
+    "iosinks",
+    "pgen",
+    "preproc",
+    "regression_kernels",
+    "stats",
+    "utils",
+]

localqtl-0.1.1/src/localqtl/_torch_utils.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import Optional, Union
+
+import torch
+
+TensorLike = Union[torch.Tensor, torch.nn.Parameter]
+DeviceLike = Union[str, torch.device]
+
+
+def resolve_device(device: DeviceLike) -> torch.device:
+    """Normalize device specifications to a :class:`torch.device`."""
+    return device if isinstance(device, torch.device) else torch.device(device)
+
+
+def move_to_device(
+    tensor: TensorLike,
+    device: torch.device,
+    *,
+    non_blocking: Optional[bool] = None,
+) -> torch.Tensor:
+    """Move ``tensor`` to ``device`` avoiding redundant copies.
+
+    Parameters
+    ----------
+    tensor : torch.Tensor
+        Tensor to move.
+    device : torch.device
+        Destination device.
+    non_blocking : bool, optional
+        When ``True`` and ``device`` is CUDA, allow asynchronous copies if the
+        source tensor is in pinned memory. If ``None`` (default), the flag is
+        enabled automatically for CUDA destinations.
+    """
+    if tensor.device == device:
+        return tensor
+
+    if non_blocking is None:
+        non_blocking = device.type == "cuda"
+
+    return tensor.to(device=device, non_blocking=non_blocking)
+
+
+__all__ = ["resolve_device", "move_to_device"]

localqtl-0.1.1/src/localqtl/cis/__init__.py

@@ -0,0 +1,11 @@
+from .nominal import map_nominal
+from .permutations import map_permutations
+from .independent import map_independent
+from .mapper import CisMapper
+
+__all__ = [
+    "map_nominal",
+    "map_permutations",
+    "map_independent",
+    "CisMapper"
+]

localqtl-0.1.1/src/localqtl/cis/common.py

@@ -0,0 +1,124 @@
+import torch
+import numpy as np
+import pandas as pd
+from typing import Optional, Tuple, List
+
+from .._torch_utils import move_to_device, resolve_device
+from ..regression_kernels import Residualizer
+
+__all__ = [
+    "dosage_vector_for_covariate",
+    "residualize_matrix_with_covariates",
+    "residualize_batch",
+]
+
+
+def _coerce_tensor(value, device: torch.device) -> torch.Tensor:
+    return move_to_device(value, device) if torch.is_tensor(value) else torch.as_tensor(
+        value, dtype=torch.float32, device=device
+    )
+
+
+def _make_gen(seed: int | None, device: str) -> torch.Generator | None:
+    if seed is None:
+        return None
+    g = torch.Generator(device=device)
+    g.manual_seed(int(seed) & ((1 << 63) - 1))
+    return g
+
+def dosage_vector_for_covariate(genotype_df: pd.DataFrame, variant_id: str,
+                                sample_order: pd.Index, missing: float | int | None) -> np.ndarray:
+    """Fetch a dosage row aligned to samples; impute 'missing' to mean of observed."""
+    if not genotype_df.columns.equals(sample_order):
+        row = genotype_df.loc[variant_id, sample_order].to_numpy(dtype=np.float32, copy=True)
+    else:
+        row = genotype_df.loc[variant_id].to_numpy(dtype=np.float32, copy=True)
+    if missing is not None:
+        mm = (row == missing)
+        if mm.any():
+            if (~mm).any():
+                row[mm] = row[~mm].mean(dtype=np.float32)
+            else:
+                row[mm] = 0.0
+    return row
+
+
+def residualize_matrix_with_covariates(
+        Y: torch.Tensor, C: Optional[pd.DataFrame], device: str
+) -> Tuple[torch.Tensor, Optional[Residualizer]]:
+    """
+    Residualize (features x samples) matrix Y against covariates C across samples.
+    Returns residualized Y and a Residualizer (or None if no covariates).
+    """
+    if C is None:
+        return Y, None
+    dev = resolve_device(device)
+    C_t = torch.tensor(C.values, dtype=torch.float32, device=dev)
+    rez = Residualizer(C_t)
+    (Y_resid,) = rez.transform(move_to_device(Y, dev), center=True)
+    return Y_resid, rez
+
+
+def residualize_batch(
+        y, G: torch.Tensor, H: Optional[torch.Tensor],
+        rez: Optional[Residualizer], center: bool = True, group: bool = False,
+) -> Tuple[object, torch.Tensor, Optional[torch.Tensor]]:
+    """
+    Residualize y, G, and optional H with the same Residualizer.
+    - If group=False: y is a single vector (n,), returns (y_resid 1D, G_resid, H_resid).
+    - If group=True:  y is a stack (k x n) or list of length k, returns (list_of_k 1D tensors, G_resid, H_resid).
+    """
+    if rez is not None and hasattr(rez, "Q_t") and getattr(rez.Q_t, "numel", lambda: 0)() > 0:
+        dev = resolve_device(rez.Q_t.device)
+    else:
+        dev = resolve_device(G.device)
+    G = move_to_device(G, dev)
+    if H is not None:
+        H = move_to_device(H, dev)
+
+    if rez is None:
+        # Pass-through; normalize return type for group mode
+        if group:
+            if isinstance(y, (list, tuple)):
+                y_list = [torch.as_tensor(yi) if not torch.is_tensor(yi) else yi
+                          for yi in y]
+            else:
+                Y = y if torch.is_tensor(y) else torch.as_tensor(y)
+                y_list = [Y] if Y.ndim == 1 else [Y[i, :] for i in range(Y.shape[0])]
+            return y_list, G, H
+        return (y if torch.is_tensor(y) else torch.as_tensor(y, dtype=torch.float32)), G, H
+
+    # Build Y on the same device as G/rez
+    if group:
+        if isinstance(y, (list, tuple)):
+            Y = torch.stack([_coerce_tensor(yi, dev) for yi in y], dim=0)
+        else:
+            Y = _coerce_tensor(y, dev)
+            if Y.ndim == 1:
+                Y = Y.unsqueeze(0)
+    else:
+        Y = _coerce_tensor(y, dev)
+        if Y.ndim == 1:
+            Y = Y.unsqueeze(0)
+    
+    # Prepare matrices for one-pass residualization
+    mats: List[torch.Tensor] = [G]
+    H_shape = None
+    if H is not None:
+        m, n, pH = H.shape
+        H_shape = (m, n, pH)
+        mats.append(H.reshape(m * pH, n))
+
+    mats_resid = rez.transform(*mats, Y, center=center)
+
+    G_resid = mats_resid[0]
+    idx, H_resid = 1, None
+    if H is not None:
+        H_resid = mats_resid[idx].reshape(H_shape)
+        idx += 1
+    Y_resid = mats_resid[idx]
+
+    if group:
+        return [Y_resid[i, :] for i in range(Y_resid.shape[0])], G_resid, H_resid
+    else:
+        return Y_resid.squeeze(0), G_resid, H_resid