statgen 0.2.2__tar.gz → 0.2.3__tar.gz

{statgen-0.2.2 → statgen-0.2.3}/PKG-INFO

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

{statgen-0.2.2 → statgen-0.2.3}/pyproject.toml

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

statgen-0.2.3/statgen/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.2.3"

statgen-0.2.3/statgen/_bfile_utils.py

@@ -0,0 +1,271 @@
+import re
+import warnings
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from pandas.errors import EmptyDataError, ParserError, ParserWarning
+
+from ._utils import CANONICAL_CHR_ORDER, CHR_RANK, allele_hash64
+from ._variant_match import check_unique_sorted_variant_keys, numeric_variant_keys
+
+IGNORED_CHR = {"Y", "MT"}
+FAM_PUBLIC_COLUMNS = ["fid", "iid", "father_id", "mother_id", "sex"]
+
+_BIM_NAMES = ["chr", "snp", "cm", "bp", "a1", "a2"]
+_FAM_NAMES = ["fid", "iid", "father_id", "mother_id", "sex", "pheno"]
+_PLOIDY_NAMES = ["ploidy_male", "ploidy_female"]
+_DNA_ALLELE_RE = re.compile(r"^[ACGT]+$")
+_BED_LOOKUP_INT8 = None
+
+
+def _read_csv_strict(path: Path, *, sep, names, description: str | None = None) -> pd.DataFrame:
+    expected = description or f"{len(names)} delimited columns"
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", ParserWarning)
+            df = pd.read_csv(
+                path,
+                sep=sep,
+                header=None,
+                dtype=str,
+                index_col=False,
+                keep_default_na=False,
+                na_filter=False,
+            )
+    except (EmptyDataError, ParserError, ParserWarning, ValueError) as exc:
+        raise ValueError(f"{path}: expected {expected}") from exc
+    if df.shape[1] != len(names):
+        raise ValueError(f"{path}: expected {expected}, got {df.shape[1]}")
+    df.columns = names
+    return df
+
+
+def _validate_source_sort_order(source_df: pd.DataFrame, path: Path) -> None:
+    if source_df.empty:
+        return
+    line = (
+        source_df["line"].to_numpy(dtype=np.int64)
+        if "line" in source_df.columns
+        else source_df.index.to_numpy(dtype=np.int64) + 1
+    )
+    chr_rank = source_df["chr"].map(CHR_RANK)
+    bad_chr = chr_rank.isna()
+    if bad_chr.any():
+        idx = int(bad_chr.idxmax())
+        raise ValueError(f"{path}:{idx + 1}: chr must use canonical labels 1-22 or X")
+
+    keys = numeric_variant_keys(
+        source_df["bp"].to_numpy(dtype=np.int64),
+        allele_hash64(source_df["a1"].to_numpy(dtype=object)),
+        allele_hash64(source_df["a2"].to_numpy(dtype=object)),
+    )
+    order_df = pd.DataFrame(
+        {
+            "chr_rank": chr_rank.to_numpy(dtype=np.int64),
+            "bp": source_df["bp"].to_numpy(dtype=np.int64),
+            "line": line,
+        }
+    )
+
+    for rank in np.unique(order_df["chr_rank"].to_numpy(dtype=np.int64)):
+        mask = order_df["chr_rank"].to_numpy(dtype=np.int64) == rank
+        shard_keys = keys[mask]
+        sorted_keys = np.sort(shard_keys, order=("bp", "a1_hash64", "a2_hash64"))
+        check_unique_sorted_variant_keys(sorted_keys, f"BIM {path}", "variant")
+
+    prev = order_df.shift(1)
+    bad_order = (
+        (order_df["chr_rank"] < prev["chr_rank"])
+        | (
+            (order_df["chr_rank"] == prev["chr_rank"])
+            & (order_df["bp"] < prev["bp"])
+        )
+    ).fillna(False)
+    if bad_order.any():
+        line = int(order_df.loc[bad_order.idxmax(), "line"])
+        raise ValueError(
+            f"{path}:{line}: rows must be sorted by (chr_rank, bp) in canonical contig order"
+        )
+
+
+def _parse_bim(path: Path) -> pd.DataFrame:
+    path = Path(path)
+    df = _read_csv_strict(
+        path,
+        sep=r"\s+",
+        names=_BIM_NAMES,
+        description="6 whitespace-delimited columns",
+    )
+    df = df.copy()
+
+    chr_col = df["chr"]
+    bad_chr = chr_col.eq("")
+    if bad_chr.any():
+        idx = int(bad_chr.idxmax())
+        raise ValueError(f"{path}:{idx + 1}: chr must be non-empty")
+
+    chr_style = chr_col.str.lower().str.startswith("chr")
+    if chr_style.any():
+        idx = int(chr_style.idxmax())
+        raise ValueError(f"{path}:{idx + 1}: chr-style labels (e.g., chr1/chrX) are not allowed")
+
+    known_chr = chr_col.isin(CANONICAL_CHR_ORDER) | chr_col.isin(IGNORED_CHR)
+    if not known_chr.all():
+        idx = int((~known_chr).idxmax())
+        raise ValueError(
+            f"{path}:{idx + 1}: unsupported chr label {chr_col.iat[idx]!r}; expected 1-22, X (Y/MT are ignored)"
+        )
+
+    bad_allele = df["a1"].eq("") | df["a2"].eq("")
+    if bad_allele.any():
+        idx = int(bad_allele.idxmax())
+        raise ValueError(f"{path}:{idx + 1}: a1 and a2 must be non-empty")
+
+    bad_a1_syntax = ~df["a1"].str.fullmatch(_DNA_ALLELE_RE.pattern)
+    if bad_a1_syntax.any():
+        idx = int(bad_a1_syntax.idxmax())
+        raise ValueError(
+            f"{path}:{idx + 1}: a1 must be uppercase DNA bases (A/C/G/T): {df['a1'].iat[idx]!r}"
+        )
+    bad_a2_syntax = ~df["a2"].str.fullmatch(_DNA_ALLELE_RE.pattern)
+    if bad_a2_syntax.any():
+        idx = int(bad_a2_syntax.idxmax())
+        raise ValueError(
+            f"{path}:{idx + 1}: a2 must be uppercase DNA bases (A/C/G/T): {df['a2'].iat[idx]!r}"
+        )
+
+    cm_num = pd.to_numeric(df["cm"], errors="coerce")
+    bad_cm = cm_num.isna()
+    if bad_cm.any():
+        idx = int(bad_cm.idxmax())
+        raise ValueError(f"{path}:{idx + 1}: cm is not a number: {df['cm'].iat[idx]!r}")
+
+    bp_num = pd.to_numeric(df["bp"], errors="coerce")
+    bad_bp = bp_num.isna() | (np.floor(bp_num) != bp_num)
+    if bad_bp.any():
+        idx = int(bad_bp.idxmax())
+        raise ValueError(f"{path}:{idx + 1}: bp is not an integer: {df['bp'].iat[idx]!r}")
+    df["bp"] = bp_num.astype(np.int64)
+
+    _validate_source_sort_order(df.loc[df["chr"].isin(CANONICAL_CHR_ORDER)], path)
+    return df.reset_index(drop=True)
+
+
+def _parse_fam(path: Path) -> pd.DataFrame:
+    path = Path(path)
+    df = _read_csv_strict(
+        path,
+        sep=r"\s+",
+        names=_FAM_NAMES,
+        description="6 whitespace-delimited columns",
+    )
+    required = df[_FAM_NAMES].eq("")
+    if required.any().any():
+        row = int(np.flatnonzero(required.any(axis=1).to_numpy())[0] + 1)
+        raise ValueError(f"{path}:{row}: FAM must contain exactly 6 non-empty columns")
+
+    sex_num = pd.to_numeric(df["sex"], errors="coerce")
+    bad_sex = sex_num.isna() | (np.floor(sex_num) != sex_num) | ~sex_num.isin([0, 1, 2])
+    if bad_sex.any():
+        idx = int(bad_sex.idxmax())
+        raise ValueError(f"{path}:{idx + 1}: FAM sex must be one of PLINK values 1, 2, or 0")
+
+    out = df[["fid", "iid", "father_id", "mother_id"]].copy()
+    out["sex"] = sex_num.astype(np.int8)
+    dup = out.duplicated(subset=["fid", "iid"], keep="first")
+    if dup.any():
+        idx = int(dup.idxmax())
+        raise ValueError(f"{path}:{idx + 1}: duplicate FAM subject pair (fid, iid)")
+    return out.reset_index(drop=True)
+
+
+def _parse_ploidy(path: Path | None, source_num_snp: int) -> tuple[np.ndarray, np.ndarray]:
+    if path is None:
+        return (
+            np.full(source_num_snp, 2.0, dtype=float),
+            np.full(source_num_snp, 2.0, dtype=float),
+        )
+
+    path = Path(path)
+    df = _read_csv_strict(
+        path,
+        sep="\t",
+        names=_PLOIDY_NAMES,
+        description="2 tab-delimited columns",
+    )
+    if df.shape[0] != source_num_snp:
+        raise ValueError(
+            f"{path}: PLOIDY row count mismatch: expected {source_num_snp}, got {df.shape[0]}"
+        )
+    numeric = df[_PLOIDY_NAMES].apply(pd.to_numeric, errors="coerce")
+    bad_numeric = numeric.isna()
+    if bad_numeric.any().any():
+        row = int(np.flatnonzero(bad_numeric.any(axis=1).to_numpy())[0] + 1)
+        raise ValueError(f"{path}:{row}: ploidy values must be integers 0, 1, or 2")
+    values = numeric.to_numpy(dtype=np.int64)
+    bad = ~np.isin(values, [0, 1, 2])
+    if bad.any():
+        row = int(np.flatnonzero(bad.any(axis=1))[0] + 1)
+        raise ValueError(f"{path}:{row}: ploidy values must be 0, 1, or 2")
+    return (
+        numeric["ploidy_male"].to_numpy(dtype=float),
+        numeric["ploidy_female"].to_numpy(dtype=float),
+    )
+
+
+def _expected_bed_size(source_num_sample: int, source_num_snp: int) -> int:
+    return 3 + ((int(source_num_sample) + 3) // 4) * int(source_num_snp)
+
+
+def _validate_bed(path: Path, source_num_sample: int, source_num_snp: int) -> int:
+    path = Path(path)
+    with open(path, "rb") as f:
+        magic = f.read(3)
+    if len(magic) != 3 or magic[:2] != b"\x6c\x1b":
+        raise ValueError(f"{path}: invalid PLINK BED magic bytes")
+    if magic[2] != 1:
+        raise ValueError(f"{path}: PLINK BED must be SNP-major mode")
+    size = path.stat().st_size
+    expected = _expected_bed_size(source_num_sample, source_num_snp)
+    if size != expected:
+        raise ValueError(f"{path}: BED file size mismatch: expected {expected}, got {size}")
+    return int(size)
+
+
+def _bed_lookup_int8() -> np.ndarray:
+    global _BED_LOOKUP_INT8
+    if _BED_LOOKUP_INT8 is None:
+        values = np.arange(256, dtype=np.uint16)[:, None]
+        shifts = (np.arange(4, dtype=np.uint16) * 2)[None, :]
+        two_bit = ((values >> shifts) & 0b11).astype(np.uint8)
+        lookup = np.empty((256, 4), dtype=np.int8)
+        lookup[two_bit == 0b00] = 2
+        lookup[two_bit == 0b01] = -1
+        lookup[two_bit == 0b10] = 1
+        lookup[two_bit == 0b11] = 0
+        lookup.setflags(write=False)
+        _BED_LOOKUP_INT8 = lookup
+    return _BED_LOOKUP_INT8
+
+
+def _read_bed_rows_int8(path: Path, source_rows: np.ndarray, source_num_sample: int) -> np.ndarray:
+    path = Path(path)
+    source_rows = np.asarray(source_rows, dtype=np.int64).reshape(-1)
+    bytes_per_snp = (int(source_num_sample) + 3) // 4
+    decoded = np.empty((int(source_num_sample), source_rows.size), dtype=np.int8)
+    lookup = _bed_lookup_int8()
+    with open(path, "rb") as f:
+        # PERF: loop over requested SNPs retained; each SNP is a separate BED row
+        #       requiring an individual seek. Inner decode is vectorized across subjects.
+        for j, source_row0 in enumerate(source_rows):
+            f.seek(3 + int(source_row0) * bytes_per_snp)
+            packed = f.read(bytes_per_snp)
+            if len(packed) != bytes_per_snp:
+                raise ValueError(
+                    f"{path}: short BED read for source SNP row {int(source_row0)}: "
+                    f"expected {bytes_per_snp} bytes, got {len(packed)}"
+                )
+            byte_values = np.frombuffer(packed, dtype=np.uint8)
+            decoded[:, j] = lookup[byte_values].reshape(-1)[:source_num_sample]
+    return decoded

statgen-0.2.3/statgen/_variant_match.py

@@ -0,0 +1,58 @@
+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",
+) -> np.ndarray:
+    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)
+
+    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]]
+    return out

{statgen-0.2.2 → statgen-0.2.3}/statgen/annotations.py

@@ -203,9 +203,9 @@ def _paint_annotations(bed_paths: list[Path], reference) -> tuple[sparse.csr_mat
 
 
 class AnnotationShard:
-    def __init__(self, label: str, checksum: str, annomat):
+    def __init__(self, label: str, reference_checksum: str, annomat):
         self._label = str(label)
-        self._checksum = str(checksum)
+        self._reference_checksum = str(reference_checksum)
         self._annomat = _as_sparse_binary(annomat)
 
     @property
@@ -213,8 +213,8 @@ class AnnotationShard:
         return self._label
 
     @property
-    def checksum(self) -> str:
-        return self._checksum
+    def reference_checksum(self) -> str:
+        return self._reference_checksum
 
     @property
     def num_snp(self) -> int:
@@ -229,13 +229,15 @@ class AnnotationShard:
         return self._annomat
 
     @classmethod
-    def _from_arrays(cls, label, checksum, annomat):
-        return cls(label=label, checksum=checksum, annomat=annomat)
+    def _from_arrays(cls, label, reference_checksum, annomat):
+        return cls(label=label, reference_checksum=reference_checksum, annomat=annomat)
 
 
 class AnnotationPanel:
     def __init__(self, shards: list[AnnotationShard], annonames):
         self._shards = list(shards)
+        if not self._shards:
+            raise ValueError("AnnotationPanel requires at least one shard")
         self._annonames = np.asarray(_coerce_annonames(annonames), dtype=object)
         self._num_snp = int(sum(s.num_snp for s in self._shards))
 
@@ -297,7 +299,7 @@ class AnnotationPanel:
 
         idx = np.asarray([idx_map[name] for name in names_list], dtype=np.int64)
         out_shards = [
-            AnnotationShard._from_arrays(s.label, s.checksum, s.annomat[:, idx])
+            AnnotationShard._from_arrays(s.label, s.reference_checksum, s.annomat[:, idx])
             for s in self._shards
         ]
         return AnnotationPanel(out_shards, names_list)
@@ -322,18 +324,21 @@ class AnnotationPanel:
                 raise ValueError("union_annotations requires matching shard labels")
             if a.num_snp != getattr(b, "num_snp", None):
                 raise ValueError("union_annotations requires matching shard row counts")
-            b_checksum = getattr(b, "checksum", None)
-            if b_checksum is not None and b_checksum != a.checksum:
+            b_checksum = getattr(b, "reference_checksum", None)
+            if b_checksum is not None and b_checksum != a.reference_checksum:
                 raise ValueError("union_annotations requires checksum-compatible reference alignment")
 
             b_anno = getattr(b, "annomat", None)
             if b_anno is None:
                 raise ValueError("union_annotations requires other shards to expose annomat")
             union_mat = sparse.hstack([a.annomat, _as_sparse_binary(b_anno)], format="csr")
-            out_shards.append(AnnotationShard._from_arrays(a.label, a.checksum, union_mat))
+            out_shards.append(AnnotationShard._from_arrays(a.label, a.reference_checksum, union_mat))
 
         return AnnotationPanel(out_shards, lhs_names + rhs_names)
 
+    def save_cache(self, path, format=None) -> None:
+        save_annotations_cache(self, path, format=format)
+
 
 def create_annotations(reference, annomat, annonames) -> AnnotationPanel:
     names = _coerce_annonames(annonames)
@@ -351,7 +356,7 @@ def create_annotations(reference, annomat, annonames) -> AnnotationPanel:
         out_shards.append(
             AnnotationShard(
                 label=ref_shard.label,
-                checksum=ref_shard.checksum,
+                reference_checksum=ref_shard.checksum,
                 annomat=mat[start:stop, :],
             )
         )
@@ -382,11 +387,11 @@ def load_annotations(bed_paths, reference) -> AnnotationPanel:
     return create_annotations(reference, annomat, annonames)
 
 
-def save_annotations_cache(panel: AnnotationPanel, path) -> None:
+def save_annotations_cache(panel: AnnotationPanel, path, format=None) -> None:
     meta = {
         "schema": _CACHE_SCHEMA,
         "shard_labels": [s.label for s in panel.shards],
-        "shard_checksums": [s.checksum for s in panel.shards],
+        "shard_checksums": [s.reference_checksum for s in panel.shards],
         "annonames": panel.annonames.tolist(),
     }
     arrays = {
@@ -437,7 +442,7 @@ def load_annotations_cache(path, shards=None) -> AnnotationPanel:
             shard_objs.append(
                 AnnotationShard._from_arrays(
                     label=label,
-                    checksum=checksums[i],
+                    reference_checksum=checksums[i],
                     annomat=mat,
                 )
             )