refactor: tighten five annotations that were narrower than the values

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,

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,

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
 

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:

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

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).

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