From c48305a78dabf60bac8d557b141bb846b62ecbb5 Mon Sep 17 00:00:00 2001 From: Michael Denyer Date: Mon, 27 Jul 2026 10:09:26 +0100 Subject: [PATCH] refactor: tighten five annotations that were narrower than the values Clears 7 pyrefly baseline entries across five src modules (77 -> 70). Each is a declaration that did not describe what the code actually passes or holds. `core/progress.progress_iterator(iterable: Iterator)` only ever does `for i, item in enumerate(iterable)`, so `Iterable` is what it requires; `Iterator` rejected the `range` that `io/plink.py` passes it. Widening a parameter cannot break a caller. `io/plink.py` also fed `total=` an `int | ndarray`, because `bed.sid_count` carries that type through `n_chunks`. It is an int at runtime, so it says so. `lmm/likelihood._frozen(data: list[int])` is called with tuples on two of its five call sites. It passes `data` straight to `np.array`, which takes any sequence, so the parameter is `Sequence[int]`. `lmm/loco.py` annotated `snps_global_mask: np.ndarray | None` on the assignment *inside* the branch that builds it, which made the declared type a union at the point of the very next item assignment. Declared before the branch instead, with the array built through a local. `lmm/special.chi2_sf_batch` called `.astype` on the result of a `np.frompyfunc` ufunc, declared as returning a scalar. `np.asarray(..., dtype=)` is the same cast and types correctly. `gwas.GWASResult.timing = field(default_factory=dict)` infers `dict[Unknown, Unknown]` rather than `GWASTiming`; the TypedDict is callable and returns the same `{}`. Same fix as `PipelineResult.timing` in #123. No behaviour change, and checked rather than argued, because `special.py` feeds p-value computation and `loco.py` gates SNP selection: - `chi2_sf_batch` is bit-identical across 13 edge cases (NaN, negative, zero, 1e-300, 1e-12, 700, 1e6, inf) against the previous expression evaluated side by side. - All four LMM modes on both backends produce byte-identical .assoc.txt. - The LOCO path is identical: same .assoc.txt, pve/pve_se, per-chromosome kinship sums and stats cache. `pytest tests/` -- 2243 passed, 10 skipped, 3 xfailed. Baseline regenerated on Python 3.11 / numpy 2.4.6 per #121; 0 errors on both 3.11.13/2.4.6 and 3.13.5/2.5.1. --- pyrefly-baseline.json | 84 ------------------------------------- src/jamma/core/progress.py | 4 +- src/jamma/gwas.py | 2 +- src/jamma/io/plink.py | 2 +- src/jamma/lmm/likelihood.py | 4 +- src/jamma/lmm/loco.py | 8 ++-- src/jamma/lmm/special.py | 5 ++- 7 files changed, 14 insertions(+), 95 deletions(-) diff --git a/pyrefly-baseline.json b/pyrefly-baseline.json index e00ce2c7..e6508c2c 100644 --- a/pyrefly-baseline.json +++ b/pyrefly-baseline.json @@ -96,18 +96,6 @@ "concise_description": "Argument `str` is not assignable to parameter `backend` with type `Literal['auto', 'numpy', 'numpy-streaming']` in function `jamma.pipeline_config.PipelineConfig.__init__`", "severity": "error" }, - { - "line": 52, - "column": 26, - "stop_line": 52, - "stop_column": 53, - "path": "src/jamma/gwas.py", - "code": -2, - "name": "bad-assignment", - "description": "`dict[Unknown, Unknown]` is not assignable to `GWASTiming`", - "concise_description": "`dict[Unknown, Unknown]` is not assignable to `GWASTiming`", - "severity": "error" - }, { "line": 189, "column": 17, @@ -120,66 +108,6 @@ "concise_description": "Argument `str` is not assignable to parameter `backend` with type `Literal['auto', 'numpy', 'numpy-streaming']` in function `jamma.pipeline_config.PipelineConfig.__init__`", "severity": "error" }, - { - "line": 386, - "column": 17, - "stop_line": 386, - "stop_column": 25, - "path": "src/jamma/io/plink.py", - "code": -2, - "name": "bad-argument-type", - "description": "Argument `range` is not assignable to parameter `iterable` with type `Iterator[Unknown]` in function `jamma.core.progress.progress_iterator`\n Protocol `Iterator` requires attribute `__next__`", - "concise_description": "Argument `range` is not assignable to parameter `iterable` with type `Iterator[Unknown]` in function `jamma.core.progress.progress_iterator`", - "severity": "error" - }, - { - "line": 386, - "column": 33, - "stop_line": 386, - "stop_column": 41, - "path": "src/jamma/io/plink.py", - "code": -2, - "name": "bad-argument-type", - "description": "Argument `int | ndarray` is not assignable to parameter `total` with type `int` in function `jamma.core.progress.progress_iterator`", - "concise_description": "Argument `int | ndarray` is not assignable to parameter `total` with type `int` in function `jamma.core.progress.progress_iterator`", - "severity": "error" - }, - { - "line": 1068, - "column": 38, - "stop_line": 1068, - "stop_column": 49, - "path": "src/jamma/lmm/likelihood.py", - "code": -2, - "name": "bad-argument-type", - "description": "Argument `tuple[int, ...]` is not assignable to parameter `data` with type `list[int]` in function `_frozen`", - "concise_description": "Argument `tuple[int, ...]` is not assignable to parameter `data` with type `list[int]` in function `_frozen`", - "severity": "error" - }, - { - "line": 1069, - "column": 36, - "stop_line": 1069, - "stop_column": 47, - "path": "src/jamma/lmm/likelihood.py", - "code": -2, - "name": "bad-argument-type", - "description": "Argument `tuple[int, ...]` is not assignable to parameter `data` with type `list[int]` in function `_frozen`", - "concise_description": "Argument `tuple[int, ...]` is not assignable to parameter `data` with type `list[int]` in function `_frozen`", - "severity": "error" - }, - { - "line": 327, - "column": 13, - "stop_line": 327, - "stop_column": 43, - "path": "src/jamma/lmm/loco.py", - "code": -2, - "name": "unsupported-operation", - "description": "Cannot set item in `None`\n Object of class `NoneType` has no attribute `__setitem__`", - "concise_description": "Cannot set item in `None`", - "severity": "error" - }, { "line": 60, "column": 12, @@ -228,18 +156,6 @@ "concise_description": "Object of class `RunnerTiming` has no attribute `clear`", "severity": "error" }, - { - "line": 216, - "column": 26, - "stop_line": 216, - "stop_column": 50, - "path": "src/jamma/lmm/special.py", - "code": -2, - "name": "missing-attribute", - "description": "Object of class `float` has no attribute `astype`", - "concise_description": "Object of class `float` has no attribute `astype`", - "severity": "error" - }, { "line": 192, "column": 40, diff --git a/src/jamma/core/progress.py b/src/jamma/core/progress.py index f8fa3de7..8ff8d50b 100644 --- a/src/jamma/core/progress.py +++ b/src/jamma/core/progress.py @@ -8,7 +8,7 @@ import sys import threading import time -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterable, Iterator from typing import TypeVar import progressbar @@ -92,7 +92,7 @@ def create_progress_bar( def progress_iterator( - iterable: Iterator, + iterable: Iterable, total: int, desc: str = "", initial_eta_seconds: float | None = None, diff --git a/src/jamma/gwas.py b/src/jamma/gwas.py index 5fc791bc..d5c9d5f3 100644 --- a/src/jamma/gwas.py +++ b/src/jamma/gwas.py @@ -49,7 +49,7 @@ class GWASResult: associations: list[AssocResult] n_samples: int n_snps_tested: int - timing: GWASTiming = field(default_factory=dict) + timing: GWASTiming = field(default_factory=GWASTiming) pve_estimate: float | None = None pve_se: float | None = None diff --git a/src/jamma/io/plink.py b/src/jamma/io/plink.py index a252b6a1..14528b00 100644 --- a/src/jamma/io/plink.py +++ b/src/jamma/io/plink.py @@ -368,7 +368,7 @@ def read_chunk(start: int, end: int) -> np.ndarray: return bed.read(index=(np.s_[:], snp_indices[start:end]), dtype=dtype) else: # Unfiltered mode: read all columns sequentially - n_total = bed.sid_count + n_total = int(bed.sid_count) label = "SNPs" def read_chunk(start: int, end: int) -> np.ndarray: diff --git a/src/jamma/lmm/likelihood.py b/src/jamma/lmm/likelihood.py index 5a7b44d7..1ac1e65c 100644 --- a/src/jamma/lmm/likelihood.py +++ b/src/jamma/lmm/likelihood.py @@ -24,7 +24,7 @@ import functools import threading -from collections.abc import Callable +from collections.abc import Callable, Sequence import numpy as np from loguru import logger @@ -1040,7 +1040,7 @@ def build_pab_table_for_c(n_cvt: int) -> dict: diag_rows = [r for r, _ in table["logdet_diag_indices"]] diag_cols = [c for _, c in table["logdet_diag_indices"]] - def _frozen(data: list[int]) -> np.ndarray: + def _frozen(data: Sequence[int]) -> np.ndarray: arr = np.array(data, dtype=np.int32) arr.flags.writeable = False return arr diff --git a/src/jamma/lmm/loco.py b/src/jamma/lmm/loco.py index 0f4aee7e..cf27e769 100644 --- a/src/jamma/lmm/loco.py +++ b/src/jamma/lmm/loco.py @@ -322,11 +322,11 @@ def run_lmm_loco( # Precompute global SNP membership mask for -snps restriction. # Avoids per-chromosome np.isin on every iteration. + snps_global_mask: np.ndarray | None = None if snps_indices is not None: - snps_global_mask: np.ndarray | None = np.zeros(n_snps_total, dtype=bool) - snps_global_mask[snps_indices] = True - else: - snps_global_mask = None + mask = np.zeros(n_snps_total, dtype=bool) + mask[snps_indices] = True + snps_global_mask = mask # Content + parameter key over every determinant of the eigendecomposition # (genotype files, filter thresholds, -ksnps set, analysed-sample mask). diff --git a/src/jamma/lmm/special.py b/src/jamma/lmm/special.py index 99cc9ad6..79bff2e4 100644 --- a/src/jamma/lmm/special.py +++ b/src/jamma/lmm/special.py @@ -213,7 +213,10 @@ def chi2_sf_batch(x: np.ndarray) -> np.ndarray: if np.any(normal): args = np.sqrt(x[normal] / 2.0) - result[normal] = _erfc_ufunc(args).astype(np.float64) + # np.frompyfunc yields an object array; asarray is the cast, and it + # types correctly where .astype does not (frompyfunc is declared as + # returning a scalar). + result[normal] = np.asarray(_erfc_ufunc(args), dtype=np.float64) return result