statgen 0.3.4__tar.gz → 0.4.0__tar.gz

{statgen-0.3.4 → statgen-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: statgen
-Version: 0.3.4
+Version: 0.4.0
 Summary: Statistical genetics data objects and loaders for Python and MATLAB/Octave.
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

{statgen-0.3.4 → statgen-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "statgen"
-version = "0.3.4"
+version = "0.4.0"
 description = "Statistical genetics data objects and loaders for Python and MATLAB/Octave."
 readme = "README.md"
 requires-python = ">=3.10"

{statgen-0.3.4 → statgen-0.4.0}/statgen/__init__.py

@@ -1,5 +1,5 @@
 from ._verbosity import get_verbosity, set_verbosity
 
-__version__ = "0.3.4"
+__version__ = "0.4.0"
 
 __all__ = ["__version__", "get_verbosity", "set_verbosity"]

{statgen-0.3.4 → statgen-0.4.0}/statgen/_bfile_utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import re
 import warnings
 from pathlib import Path
@@ -275,7 +277,7 @@ def _read_bed_rows_int8(path: Path, source_rows: np.ndarray, source_num_sample:
         # PERF: one seek/read per contiguous BED row block; requested order and
         #       repeats are restored after block decoding. Blocks are capped to
         #       bound transient packed/decode allocations for large requests.
-        for start, stop in zip(group_starts, group_stops, strict=True):
+        for start, stop in zip(group_starts, group_stops):
             for chunk_start in range(start, stop, rows_per_read):
                 chunk_stop = min(chunk_start + rows_per_read, stop)
                 block_rows = unique_rows[chunk_start:chunk_stop]

{statgen-0.3.4 → statgen-0.4.0}/statgen/_ld_schema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import hashlib
 import json
 from pathlib import Path

{statgen-0.3.4 → statgen-0.4.0}/statgen/_utils.py

@@ -114,3 +114,31 @@ def _allele_hash_lanes_from_utf8(encoded) -> tuple[np.ndarray, np.ndarray]:
         h2[active] = (h2[active] * _ALLELE_HASH_BASE2 + x) % _ALLELE_HASH_P
 
     return h1, h2
+
+
+_A_HASH64, _C_HASH64, _G_HASH64, _T_HASH64 = _SINGLE_BASE_HASH64
+
+
+def _is_single_base(h) -> np.ndarray:
+    return np.isin(h, _SINGLE_BASE_HASH64)
+
+
+def _is_single_nucleotide_variant(a1_hash64, a2_hash64) -> np.ndarray:
+    return _is_single_base(a1_hash64) & _is_single_base(a2_hash64)
+
+
+def _is_strand_ambiguous(a1_hash64, a2_hash64) -> np.ndarray:
+    a1 = np.asarray(a1_hash64, dtype=np.uint64).reshape(-1)
+    a2 = np.asarray(a2_hash64, dtype=np.uint64).reshape(-1)
+    return (
+        ((a1 == _A_HASH64) & (a2 == _T_HASH64))
+        | ((a1 == _T_HASH64) & (a2 == _A_HASH64))
+        | ((a1 == _C_HASH64) & (a2 == _G_HASH64))
+        | ((a1 == _G_HASH64) & (a2 == _C_HASH64))
+    )
+
+
+def _validate_allow_indel_swaps(value) -> bool:
+    if not isinstance(value, (bool, np.bool_)):
+        raise ValueError("allow_indel_swaps must be a boolean scalar")
+    return bool(value)

statgen-0.4.0/statgen/_variant_match.py

@@ -0,0 +1,105 @@
+from dataclasses import dataclass
+
+import numpy as np
+
+from ._utils import _is_single_base, _is_single_nucleotide_variant, _is_strand_ambiguous
+
+_KEY_DTYPE = np.dtype([("bp", "<i8"), ("a1_hash64", "<u8"), ("a2_hash64", "<u8")])
+
+
+def numeric_variant_keys(bp_arr, a1_hash64, a2_hash64) -> np.ndarray:
+    bp = np.asarray(bp_arr, dtype=np.int64).reshape(-1)
+    a1h = np.asarray(a1_hash64, dtype=np.uint64).reshape(-1)
+    a2h = np.asarray(a2_hash64, dtype=np.uint64).reshape(-1)
+    if not (bp.size == a1h.size == a2h.size):
+        raise ValueError("numeric variant key vector lengths mismatch")
+    keys = np.empty(bp.size, dtype=_KEY_DTYPE)
+    keys["bp"] = bp
+    keys["a1_hash64"] = a1h
+    keys["a2_hash64"] = a2h
+    return keys
+
+
+def check_unique_sorted_variant_keys(sorted_keys: np.ndarray, where: str, object_name: str = "source") -> None:
+    if sorted_keys.size < 2:
+        return
+    dup = sorted_keys[1:] == sorted_keys[:-1]
+    if np.any(dup):
+        pos = int(np.flatnonzero(dup)[0] + 1)
+        key = sorted_keys[pos]
+        raise ValueError(
+            f"Ambiguous duplicate {object_name}/reference matching key in {where}: "
+            f"bp={int(key['bp'])}, a1_hash64={int(key['a1_hash64'])}, "
+            f"a2_hash64={int(key['a2_hash64'])}"
+        )
+
+
+@dataclass(frozen=True)
+class VariantMatch:
+    source_rows: np.ndarray
+    is_allele_swapped: np.ndarray
+    num_strand_ambiguous: int
+    num_disabled_indels: int
+
+
+def _lookup_candidates(ref_keys, src_keys, source_rows, shard_label, source_name, candidate_eligible=None):
+    """Return unique eligible source rows and count hits from ineligible rows."""
+    order = np.argsort(src_keys, kind="mergesort", order=("bp", "a1_hash64", "a2_hash64"))
+    sorted_keys = src_keys[order]
+    left = np.searchsorted(sorted_keys, ref_keys, side="left")
+    right = np.searchsorted(sorted_keys, ref_keys, side="right")
+    present = right > left
+    accepted = present.copy()
+    if candidate_eligible is not None:
+        # A key fixes both allele hashes, so all candidates for a slot have
+        # the same allele class and eligibility, even when they are duplicates.
+        accepted[present] = candidate_eligible[order[left[present]]]
+    excluded_hits = int(np.sum((right - left)[present & ~accepted]))
+    ambiguous = (right - left > 1) & accepted
+    if np.any(ambiguous):
+        key = ref_keys[np.flatnonzero(ambiguous)[0]]
+        raise ValueError(
+            f"Ambiguous duplicate {source_name}/reference matching key in shard {shard_label}: "
+            f"bp={int(key['bp'])}, a1_hash64={int(key['a1_hash64'])}, "
+            f"a2_hash64={int(key['a2_hash64'])}"
+        )
+    out = np.full(ref_keys.size, -1, dtype=np.int64)
+    out[accepted] = source_rows[order[left[accepted]]]
+    return out, excluded_hits
+
+
+def match_shard_numeric(
+    ref_keys: np.ndarray,
+    src_keys: np.ndarray,
+    shard_label: str,
+    *,
+    source_name: str = "source",
+    allow_indel_swaps: bool = False,
+) -> VariantMatch:
+    # Reference validation belongs to load_reference; cached references are trusted.
+    rows, _ = _lookup_candidates(
+        ref_keys, src_keys, np.arange(src_keys.size), shard_label, source_name
+    )
+    swapped = np.zeros(ref_keys.size, dtype=bool)
+    unused = np.ones(src_keys.size, dtype=bool)
+    unused[rows[rows >= 0]] = False
+    a1h = src_keys["a1_hash64"]
+    a2h = src_keys["a2_hash64"]
+    snv = _is_single_nucleotide_variant(a1h, a2h)
+    indel = _is_single_base(a1h) ^ _is_single_base(a2h)
+    eligible = snv | (indel & allow_indel_swaps)
+    source_rows = np.flatnonzero(unused & (snv | indel))
+    reversed_keys = numeric_variant_keys(
+        src_keys["bp"][source_rows], a2h[source_rows], a1h[source_rows]
+    )
+    missing = np.flatnonzero(rows < 0)
+    fallback, disabled = _lookup_candidates(
+        ref_keys[missing], reversed_keys, source_rows, shard_label, source_name,
+        candidate_eligible=eligible[source_rows],
+    )
+    accepted = missing[fallback >= 0]
+    rows[accepted] = fallback[fallback >= 0]
+    swapped[accepted] = True
+    selected = rows[swapped]
+    ambiguous = _is_strand_ambiguous(a1h[selected], a2h[selected])
+    return VariantMatch(rows, swapped, int(np.count_nonzero(ambiguous)), disabled)

statgen-0.4.0/statgen/_verbosity.py

@@ -0,0 +1,30 @@
+import sys
+
+_LEVELS = ("quiet", "info")
+_VERBOSITY = "info"
+
+
+def set_verbosity(level: str) -> None:
+    global _VERBOSITY
+    if level not in _LEVELS:
+        raise ValueError("verbosity must be one of: quiet, info")
+    _VERBOSITY = level
+
+
+def get_verbosity() -> str:
+    return _VERBOSITY
+
+
+def info(message: str) -> None:
+    if _VERBOSITY == "info":
+        print(message, file=sys.stderr)
+
+
+def report_variant_swaps(source_path, shard_label, num_swaps, num_ambiguous, num_disabled):
+    context = f"{source_path}: shard {shard_label}"
+    if num_swaps:
+        info(f"{context}: {num_swaps} successful allele swap(s)")
+    if num_ambiguous:
+        info(f"{context}: {num_ambiguous} strand-ambiguous A/T or C/G swap(s)")
+    if num_disabled:
+        info(f"{context}: {num_disabled} indel row(s) unmatched because indel swaps are disabled")

{statgen-0.3.4 → statgen-0.4.0}/statgen/annotations.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 from pathlib import Path
 

{statgen-0.3.4 → statgen-0.4.0}/statgen/genotype.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 import warnings
 from dataclasses import dataclass
@@ -15,11 +17,13 @@ from ._bfile_utils import (
     _read_bed_rows_int8,
     _validate_bed,
 )
-from ._utils import allele_hash64, validate_requested_shards
+from ._utils import _validate_allow_indel_swaps, allele_hash64, validate_requested_shards
+from ._verbosity import report_variant_swaps
 from ._variant_match import match_shard_numeric, numeric_variant_keys
 
-_CACHE_SCHEMA = "genotype_cache/0.1"
+_CACHE_SCHEMA = "genotype_cache/0.2"
 _BED_MISSING_INT8 = np.int8(-1)
+_SWAP_LUT = np.array([-1, 2, 1, 0], dtype=np.int8)
 _HAPLOID_MODES = {"raw", "ploidy_scaled"}
 _CACHE_ARRAY_FIELDS = (
     "is_present",
@@ -150,6 +154,7 @@ class GenotypeShard:
         subject_present,
         source_subject_row0,
         is_present,
+        is_allele_swapped,
         ploidy_male,
         ploidy_female,
         reference_checksum,
@@ -170,10 +175,12 @@ class GenotypeShard:
 
         self._source_row0 = np.asarray(source_row0, dtype=np.int64).reshape(-1)
         self._is_present = np.asarray(is_present, dtype=bool).reshape(-1)
+        self._is_allele_swapped = np.asarray(is_allele_swapped, dtype=bool).reshape(-1)
         self._ploidy_male = np.asarray(ploidy_male, dtype=float).reshape(-1)
         self._ploidy_female = np.asarray(ploidy_female, dtype=float).reshape(-1)
         if not (
             self._source_row0.size
+            == self._is_allele_swapped.size
             == self._is_present.size
             == self._ploidy_male.size
             == self._ploidy_female.size
@@ -184,6 +191,8 @@ class GenotypeShard:
         bad_source = (self._source_row0 < -1) | (self._source_row0 >= self._source_num_snp)
         if bad_source.any():
             raise ValueError(f"GenotypeShard {self._label}: source_row0 out of source BIM bounds")
+        if np.any(self._is_allele_swapped & ~self._is_present):
+            raise ValueError("is_allele_swapped requires is_present")
         absent = ~self._is_present
         if np.any(~np.isnan(self._ploidy_male[absent])) or np.any(~np.isnan(self._ploidy_female[absent])):
             raise ValueError(f"GenotypeShard {self._label}: absent SNPs must have NaN ploidy")
@@ -207,6 +216,7 @@ class GenotypeShard:
         for name in (
             "_source_row0",
             "_is_present",
+            "_is_allele_swapped",
             "_ploidy_male",
             "_ploidy_female",
             "_subject_present",
@@ -258,6 +268,10 @@ class GenotypeShard:
     def is_present(self) -> np.ndarray:
         return self._is_present
 
+    @property
+    def is_allele_swapped(self) -> np.ndarray:
+        return self._is_allele_swapped
+
     @property
     def ploidy_male(self) -> np.ndarray:
         return self._ploidy_male
@@ -340,6 +354,12 @@ class GenotypePanel:
             return np.array([], dtype=bool)
         return np.concatenate([s.is_present for s in self._shards])
 
+    @property
+    def is_allele_swapped(self) -> np.ndarray:
+        if not self._shards:
+            return np.array([], dtype=bool)
+        return np.concatenate([s.is_allele_swapped for s in self._shards])
+
     @property
     def ploidy_male(self) -> np.ndarray:
         if not self._shards:
@@ -413,7 +433,7 @@ class GenotypePanel:
 
     def _addressed_shards(self, snp_indices: np.ndarray) -> list[tuple[GenotypeShard, np.ndarray, np.ndarray]]:
         addressed = []
-        for shard, offset in zip(self._shards, self._shard_offsets, strict=True):
+        for shard, offset in zip(self._shards, self._shard_offsets):
             start = int(offset["start0"])
             stop = int(offset["stop0"])
             cols = np.flatnonzero((snp_indices >= start) & (snp_indices < stop))
@@ -475,6 +495,9 @@ class GenotypePanel:
                 source_rows,
                 shard.source_num_sample,
             )
+            sc = np.flatnonzero(shard.is_allele_swapped[local_indices])
+            if sc.size:
+                source_geno[:, sc] = _SWAP_LUT[source_geno[:, sc] + 1]
             panel_rows = np.flatnonzero(shard.subject_present)
             if panel_rows.size:
                 source_subject_rows = shard.source_subject_row0[panel_rows]
@@ -572,8 +595,9 @@ def _map_chrx_subjects(chrx_fam: pd.DataFrame, panel_fam: pd.DataFrame, where: s
     return subject_present, source_subject_row0
 
 
-def _build_shard(ref_shard, source: _SourceRecord, subject_present, source_subject_row0) -> GenotypeShard:
+def _build_shard(ref_shard, source: _SourceRecord, subject_present, source_subject_row0, allow_indel_swaps) -> GenotypeShard:
     source_bim = source.bim.loc[source.bim["chr"] == ref_shard.label]
+    is_allele_swapped = np.zeros(int(ref_shard.num_snp), dtype=bool)
     if source_bim.empty:
         local_match = np.full(int(ref_shard.num_snp), -1, dtype=np.int64)
     else:
@@ -583,15 +607,14 @@ def _build_shard(ref_shard, source: _SourceRecord, subject_present, source_subje
             source_bim["a1_hash64"].to_numpy(dtype=np.uint64),
             source_bim["a2_hash64"].to_numpy(dtype=np.uint64),
         )
-        local_match, n_swapped = match_shard_numeric(ref_keys, src_keys, ref_shard.label, source_name="genotype")
-        if n_swapped > 0:
-            warnings.warn(
-                f"{source.bim_path}: shard {ref_shard.label}: {n_swapped} unmatched genotype "
-                "variant(s) would match the reference if a1/a2 were swapped; "
-                "variants remain unmatched",
-                RuntimeWarning,
-                stacklevel=3,
-            )
+        match = match_shard_numeric(
+            ref_keys, src_keys, ref_shard.label, source_name="genotype",
+            allow_indel_swaps=allow_indel_swaps,
+        )
+        local_match = match.source_rows
+        is_allele_swapped = match.is_allele_swapped
+        report_variant_swaps(source.bim_path, ref_shard.label, np.count_nonzero(is_allele_swapped),
+                             match.num_strand_ambiguous, match.num_disabled_indels)
 
     is_present = local_match >= 0
     source_row0 = np.full(int(ref_shard.num_snp), -1, dtype=np.int64)
@@ -613,6 +636,7 @@ def _build_shard(ref_shard, source: _SourceRecord, subject_present, source_subje
         subject_present=subject_present,
         source_subject_row0=source_subject_row0,
         is_present=is_present,
+        is_allele_swapped=is_allele_swapped,
         ploidy_male=ploidy_male,
         ploidy_female=ploidy_female,
         reference_checksum=ref_shard.checksum,
@@ -628,7 +652,8 @@ def _validate_reference_for_genotype(reference) -> list:
     return ref_shards
 
 
-def load_genotype(bfile_prefix, reference) -> GenotypePanel:
+def load_genotype(bfile_prefix, reference, allow_indel_swaps=False) -> GenotypePanel:
+    allow_indel_swaps = _validate_allow_indel_swaps(allow_indel_swaps)
     ref_shards = _validate_reference_for_genotype(reference)
     prefix = str(bfile_prefix)
     ref_labels = [s.label for s in ref_shards]
@@ -682,7 +707,7 @@ def load_genotype(bfile_prefix, reference) -> GenotypePanel:
             subject_maps["X"] = _all_subjects_present(sources["X"].source_num_sample)
 
     shards = [
-        _build_shard(ref_shard, sources[ref_shard.label], *subject_maps[ref_shard.label])
+        _build_shard(ref_shard, sources[ref_shard.label], *subject_maps[ref_shard.label], allow_indel_swaps)
         for ref_shard in ref_shards
     ]
     return GenotypePanel(
@@ -714,6 +739,7 @@ def save_genotype_cache(panel: GenotypePanel, path) -> None:
     arrays = {
         "_meta": np.frombuffer(json.dumps(meta).encode(), dtype=np.uint8),
         "is_present": panel.is_present.astype(bool, copy=False),
+        "is_allele_swapped": panel.is_allele_swapped,
         "ploidy_male": panel.ploidy_male.astype(float, copy=False),
         "ploidy_female": panel.ploidy_female.astype(float, copy=False),
         "source_row0": panel.source_row0.astype(np.int64, copy=False),
@@ -744,9 +770,10 @@ def _load_cache_meta(data) -> dict:
 
 def _validate_cache_meta(meta: dict, data) -> None:
     schema = meta.get("schema")
-    if schema != _CACHE_SCHEMA:
+    if schema not in {_CACHE_SCHEMA, "genotype_cache/0.1"}:
         raise ValueError(f"Unsupported genotype cache schema: {schema!r}")
-    missing = [name for name in _CACHE_ARRAY_FIELDS if name not in data.files]
+    required = _CACHE_ARRAY_FIELDS + (("is_allele_swapped",) if schema == _CACHE_SCHEMA else ())
+    missing = [name for name in required if name not in data.files]
     if missing:
         raise ValueError(f"Invalid genotype cache: missing array field {missing[0]!r}")
     n_shards = int(meta.get("n_shards", -1))
@@ -773,6 +800,12 @@ def _validate_cache_meta(meta: dict, data) -> None:
         raise ValueError("Invalid genotype cache: source_layout must be 'non_sharded' or 'sharded'")
 
     is_present = data["is_present"]
+    if schema == _CACHE_SCHEMA:
+        swapped = data["is_allele_swapped"]
+        if swapped.dtype != np.dtype(bool) or swapped.shape != is_present.shape:
+            raise ValueError("Invalid genotype cache: is_allele_swapped must be a logical SNP-axis vector")
+        if np.any(swapped & ~is_present.astype(bool)):
+            raise ValueError("Invalid genotype cache: is_allele_swapped requires is_present")
     source_row0 = data["source_row0"]
     if is_present.ndim != 1 or source_row0.ndim != 1 or is_present.shape[0] != source_row0.shape[0]:
         raise ValueError("Invalid genotype cache: SNP-axis vector lengths mismatch")
@@ -817,6 +850,8 @@ def load_genotype_cache(path, shards=None) -> GenotypePanel:
         selected = validate_requested_shards(shards, labels, "load_genotype_cache")
         label_to_index = {label: i for i, label in enumerate(labels)}
 
+        swap_flags = (data["is_allele_swapped"] if meta["schema"] == _CACHE_SCHEMA
+                      else np.zeros(data["is_present"].size, dtype=bool))
         shard_objs = []
         for label in selected:
             i = label_to_index[label]
@@ -833,6 +868,7 @@ def load_genotype_cache(path, shards=None) -> GenotypePanel:
                     subject_present=data["subject_present"][:, i],
                     source_subject_row0=data["source_subject_row0"][:, i],
                     is_present=data["is_present"][start:stop],
+                    is_allele_swapped=swap_flags[start:stop],
                     ploidy_male=data["ploidy_male"][start:stop],
                     ploidy_female=data["ploidy_female"][start:stop],
                     reference_checksum=meta["shard_checksums"][i],

{statgen-0.3.4 → statgen-0.4.0}/statgen/reference.py

@@ -9,34 +9,32 @@ import numpy as np
 import pandas as pd
 
 from ._bfile_utils import _parse_bim
-from ._utils import CANONICAL_CHR_ORDER, allele_hash64, validate_requested_shards
+from ._utils import (
+    CANONICAL_CHR_ORDER,
+    _is_single_nucleotide_variant,
+    _is_strand_ambiguous,
+    allele_hash64,
+    validate_requested_shards,
+)
 
 logger = logging.getLogger(__name__)
 
 _CACHE_SCHEMA = "reference_cache/0.1"
-_A_HASH64, _C_HASH64, _G_HASH64, _T_HASH64 = allele_hash64(
-    np.array(["A", "C", "G", "T"], dtype=object)
-)
-_SINGLE_BASE_HASH64 = np.array(
-    [_A_HASH64, _C_HASH64, _G_HASH64, _T_HASH64], dtype=np.uint64
-)
 
 
-def _is_single_nucleotide_variant(a1_hash64, a2_hash64) -> np.ndarray:
-    return np.isin(a1_hash64, _SINGLE_BASE_HASH64) & np.isin(
-        a2_hash64, _SINGLE_BASE_HASH64
-    )
-
-
-def _is_strand_ambiguous(a1_hash64, a2_hash64) -> np.ndarray:
-    a1 = np.asarray(a1_hash64, dtype=np.uint64).reshape(-1)
-    a2 = np.asarray(a2_hash64, dtype=np.uint64).reshape(-1)
-    return (
-        ((a1 == _A_HASH64) & (a2 == _T_HASH64))
-        | ((a1 == _T_HASH64) & (a2 == _A_HASH64))
-        | ((a1 == _C_HASH64) & (a2 == _G_HASH64))
-        | ((a1 == _G_HASH64) & (a2 == _C_HASH64))
-    )
+def _validate_reference_keys(shard) -> None:
+    # Ordered-key duplicates have already been rejected by the BIM parser.
+    snv = shard.is_single_nucleotide_variant
+    bp = shard.bp[snv]
+    a1 = shard.a1_hash64[snv]
+    a2 = shard.a2_hash64[snv]
+    lo, hi = np.minimum(a1, a2), np.maximum(a1, a2)
+    order = np.lexsort((hi, lo, bp))
+    left, right = order[:-1], order[1:]
+    duplicate = (bp[left] == bp[right]) & (lo[left] == lo[right]) & (hi[left] == hi[right])
+    if np.any(duplicate):
+        pos = right[np.flatnonzero(duplicate)[0]]
+        raise ValueError(f"reference shard {shard.label}: reciprocal SNV keys are not allowed at bp={int(bp[pos])}")
 
 
 def _checksum_from_arrays(
@@ -410,6 +408,8 @@ def load_reference(path, shards=None) -> ReferencePanel:
                     bim_df["a2"].to_numpy(dtype=object),
                 )
             )
+        for shard in shards:
+            _validate_reference_keys(shard)
         return ReferencePanel(shards)
 
     bim_df = _parse_bim(path_obj)
@@ -431,6 +431,8 @@ def load_reference(path, shards=None) -> ReferencePanel:
                 shard_df["a2"].to_numpy(dtype=object),
             )
         )
+    for shard in out_shards:
+        _validate_reference_keys(shard)
     return ReferencePanel(out_shards)
 
 

{statgen-0.3.4 → statgen-0.4.0}/statgen/sumstats.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import json
 import re
 import warnings
@@ -8,7 +10,8 @@ import numpy as np
 import pandas as pd
 from pandas.errors import EmptyDataError, ParserError
 
-from ._utils import CANONICAL_CHR_ORDER, allele_hash64, validate_requested_shards
+from ._utils import CANONICAL_CHR_ORDER, _validate_allow_indel_swaps, allele_hash64, validate_requested_shards
+from ._verbosity import info, report_variant_swaps
 from ._variant_match import match_shard_numeric, numeric_variant_keys
 
 _CACHE_SCHEMA = "sumstats_cache/0.1"
@@ -136,15 +139,19 @@ def _parse_sumstats(path: Path) -> pd.DataFrame:
     out["bp"] = out["bp"].astype(np.int64)
 
     p_arr = out["p"].to_numpy(dtype=float)
-    bad_p = out["p"].isna() | ~np.isfinite(p_arr) | (p_arr < 0.0) | (p_arr > 1.0)
+    bad_p = ~np.isnan(p_arr) & (~np.isfinite(p_arr) | (p_arr < 0.0) | (p_arr > 1.0))
     if bad_p.any():
-        idx = int(bad_p.idxmax())
+        idx = int(out.index[np.flatnonzero(bad_p)[0]])
         raise ValueError(f"{path}: row {idx + 2}: p must be finite numeric in [0, 1]: {out.loc[idx, 'p']!r}")
 
     for opt in _OPTIONAL_COLS:
         if opt in out.columns:
             out[opt] = out[opt].mask(~np.isfinite(out[opt].to_numpy(dtype=float)))
 
+    missing_p = out["p"].isna()
+    if missing_p.any():
+        info(f"{path}: dropped {int(missing_p.sum())} row(s) with missing p")
+        return out.loc[~missing_p].copy()
     return out
 
 
@@ -336,7 +343,7 @@ def _derive_logp(p_vals: np.ndarray) -> np.ndarray:
 
 
 def _validate_pvec(p_vals: np.ndarray, name: str) -> None:
-    bad = ~np.isfinite(p_vals) | (p_vals < 0.0) | (p_vals > 1.0)
+    bad = ~np.isnan(p_vals) & (~np.isfinite(p_vals) | (p_vals < 0.0) | (p_vals > 1.0))
     if bad.any():
         idx = int(np.flatnonzero(bad)[0])
         raise ValueError(f"{name}[{idx}] must be finite numeric in [0, 1]")
@@ -396,18 +403,7 @@ def _build_sumstats_from_aligned(
     return Sumstats(shards)
 
 
-def _warn_swapped_allele_matches(source_path, shard_label: str, count: int, context: str) -> None:
-    if count > 0:
-        warnings.warn(
-            f"{source_path}: shard {shard_label}: {count} unmatched {context} "
-            "variant(s) would match the reference if a1/a2 were swapped; "
-            "variants remain unmatched",
-            RuntimeWarning,
-            stacklevel=3,
-        )
-
-
-def _build_sumstats_panel(df: pd.DataFrame, reference, source_path) -> Sumstats:
+def _build_sumstats_panel(df: pd.DataFrame, reference, source_path, allow_indel_swaps=False) -> Sumstats:
     n = int(reference.num_snp)
     aligned_optional = {
         col: (np.full(n, np.nan, dtype=float) if col in df.columns else None)
@@ -427,8 +423,13 @@ def _build_sumstats_panel(df: pd.DataFrame, reference, source_path) -> Sumstats:
         src_idx = np.flatnonzero(src_mask)
         ref_keys = numeric_variant_keys(ref_shard.bp, ref_shard.a1_hash64, ref_shard.a2_hash64)
         src_keys = numeric_variant_keys(src_bp[src_idx], src_a1_hash64[src_idx], src_a2_hash64[src_idx])
-        local_match, n_swapped = match_shard_numeric(ref_keys, src_keys, ref_shard.label, source_name="sumstats")
-        _warn_swapped_allele_matches(source_path, ref_shard.label, n_swapped, "sumstats")
+        match = match_shard_numeric(
+            ref_keys, src_keys, ref_shard.label, source_name="sumstats",
+            allow_indel_swaps=allow_indel_swaps,
+        )
+        local_match = match.source_rows
+        report_variant_swaps(source_path, ref_shard.label, np.count_nonzero(match.is_allele_swapped),
+                             match.num_strand_ambiguous, match.num_disabled_indels)
         has_match = local_match >= 0
         matched_src = src_idx[local_match[has_match]]
         aligned_ix = np.arange(start, stop, dtype=np.int64)[has_match]
@@ -437,6 +438,14 @@ def _build_sumstats_panel(df: pd.DataFrame, reference, source_path) -> Sumstats:
         for col, out in aligned_optional.items():
             if out is not None:
                 out[aligned_ix] = df[col].to_numpy(dtype=float)[matched_src]
+        swapped_ix = np.arange(start, stop)[match.is_allele_swapped]
+        for col in ("z", "beta", "eaf"):
+            out = aligned_optional[col]
+            if out is not None:
+                out[swapped_ix] = 1 - out[swapped_ix] if col == "eaf" else -out[swapped_ix]
+        if not has_match.any():
+            warnings.warn(f"{source_path}: no matching source rows for reference shard {ref_shard.label!r}",
+                          RuntimeWarning, stacklevel=3)
 
     aligned_logp = _derive_logp(aligned_p)
     _warn_optional_zn_completeness(
@@ -454,9 +463,10 @@ def _build_sumstats_panel(df: pd.DataFrame, reference, source_path) -> Sumstats:
     )
 
 
-def load_sumstats(path, reference) -> Sumstats:
+def load_sumstats(path, reference, allow_indel_swaps=False) -> Sumstats:
+    allow_indel_swaps = _validate_allow_indel_swaps(allow_indel_swaps)
     path = Path(path)
-    return _build_sumstats_panel(_parse_sumstats(path), reference, path)
+    return _build_sumstats_panel(_parse_sumstats(path), reference, path, allow_indel_swaps)
 
 
 def create_sumstats(
@@ -483,6 +493,10 @@ def create_sumstats(
         "info": None if info is None else _coerce_aligned_vec("info", info, n=num_snp_value, allow_inf=False),
     }
 
+    for name, vec in aligned_optional.items():
+        if vec is not None:
+            aligned_optional[name] = np.where(np.isnan(aligned_p), np.nan, vec)
+
     return _build_sumstats_from_aligned(
         reference,
         aligned_logp=aligned_logp,

{statgen-0.3.4 → statgen-0.4.0}/statgen.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: statgen
-Version: 0.3.4
+Version: 0.4.0
 Summary: Statistical genetics data objects and loaders for Python and MATLAB/Octave.
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown

statgen-0.3.4/statgen/_variant_match.py

@@ -1,76 +0,0 @@
-import numpy as np
-
-_KEY_DTYPE = np.dtype([("bp", "<i8"), ("a1_hash64", "<u8"), ("a2_hash64", "<u8")])
-
-
-def numeric_variant_keys(bp_arr, a1_hash64, a2_hash64) -> np.ndarray:
-    bp = np.asarray(bp_arr, dtype=np.int64).reshape(-1)
-    a1h = np.asarray(a1_hash64, dtype=np.uint64).reshape(-1)
-    a2h = np.asarray(a2_hash64, dtype=np.uint64).reshape(-1)
-    if not (bp.size == a1h.size == a2h.size):
-        raise ValueError("numeric variant key vector lengths mismatch")
-    keys = np.empty(bp.size, dtype=_KEY_DTYPE)
-    keys["bp"] = bp
-    keys["a1_hash64"] = a1h
-    keys["a2_hash64"] = a2h
-    return keys
-
-
-def check_unique_sorted_variant_keys(sorted_keys: np.ndarray, where: str, object_name: str = "source") -> None:
-    if sorted_keys.size < 2:
-        return
-    dup = sorted_keys[1:] == sorted_keys[:-1]
-    if np.any(dup):
-        pos = int(np.flatnonzero(dup)[0] + 1)
-        key = sorted_keys[pos]
-        raise ValueError(
-            f"Ambiguous duplicate {object_name}/reference matching key in {where}: "
-            f"bp={int(key['bp'])}, a1_hash64={int(key['a1_hash64'])}, "
-            f"a2_hash64={int(key['a2_hash64'])}"
-        )
-
-
-def match_shard_numeric(
-    ref_keys: np.ndarray,
-    src_keys: np.ndarray,
-    shard_label: str,
-    *,
-    source_name: str = "source",
-) -> tuple[np.ndarray, int]:
-    ref_order = np.argsort(ref_keys, kind="mergesort", order=("bp", "a1_hash64", "a2_hash64"))
-    ref_sorted = ref_keys[ref_order]
-    check_unique_sorted_variant_keys(ref_sorted, f"reference shard {shard_label}", source_name)
-
-    if src_keys.size == 0:
-        return np.full(ref_keys.size, -1, dtype=np.int64), 0
-
-    src_order = np.argsort(src_keys, kind="mergesort", order=("bp", "a1_hash64", "a2_hash64"))
-    src_sorted = src_keys[src_order]
-    check_unique_sorted_variant_keys(src_sorted, f"{source_name} shard {shard_label}", source_name)
-
-    pos = np.searchsorted(src_sorted, ref_keys)
-    in_range = pos < src_sorted.size
-    matched = np.zeros(ref_keys.size, dtype=bool)
-    matched[in_range] = src_sorted[pos[in_range]] == ref_keys[in_range]
-
-    out = np.full(ref_keys.size, -1, dtype=np.int64)
-    out[matched] = src_order[pos[matched]]
-
-    matched_src = np.zeros(src_keys.size, dtype=bool)
-    matched_src[out[matched]] = True
-    unmatched_src = src_keys[~matched_src]
-    if unmatched_src.size == 0 or ref_sorted.size == 0:
-        return out, 0
-
-    swapped_src = np.empty(unmatched_src.size, dtype=_KEY_DTYPE)
-    swapped_src["bp"] = unmatched_src["bp"]
-    swapped_src["a1_hash64"] = unmatched_src["a2_hash64"]
-    swapped_src["a2_hash64"] = unmatched_src["a1_hash64"]
-    swapped_pos = np.searchsorted(ref_sorted, swapped_src)
-    swapped_in_range = swapped_pos < ref_sorted.size
-    swapped_match = np.zeros(swapped_src.size, dtype=bool)
-    swapped_match[swapped_in_range] = (
-        ref_sorted[swapped_pos[swapped_in_range]]
-        == swapped_src[swapped_in_range]
-    )
-    return out, int(np.count_nonzero(swapped_match))

statgen-0.3.4/statgen/_verbosity.py

@@ -1,20 +0,0 @@
-import sys
-
-_LEVELS = ("quiet", "info")
-_VERBOSITY = "info"
-
-
-def set_verbosity(level: str) -> None:
-    global _VERBOSITY
-    if level not in _LEVELS:
-        raise ValueError("verbosity must be one of: quiet, info")
-    _VERBOSITY = level
-
-
-def get_verbosity() -> str:
-    return _VERBOSITY
-
-
-def info(message: str) -> None:
-    if _VERBOSITY == "info":
-        print(message, file=sys.stderr)