Skip to content

refactor: narrow the mode-specific inputs in compute_lmm_chunk_numpy #125

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michael-denyer merged 1 commit into from
Jul 26, 2026
Merged

refactor: narrow the mode-specific inputs in compute_lmm_chunk_numpy #125

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 0 additions & 72 deletions pyrefly-baseline.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,78 +144,6 @@
"concise_description": "Argument `int | ndarray` is not assignable to parameter `total` with type `int` in function `jamma.core.progress.progress_iterator`",
"severity": "error"
},
{
"line": 1740,
"column": 22,
"stop_line": 1752,
"stop_column": 10,
"path": "src/jamma/lmm/compute_numpy.py",
"code": -2,
"name": "no-matching-overload",
"description": "No matching overload found for function `typing.MutableMapping.update` called with arguments: (WaldResult)\n Possible overloads:\n (m: SupportsKeysAndGetItem[str, ndarray | None], /) -> None [closest match]\n (m: SupportsKeysAndGetItem[str, ndarray | None], /, ...) -> None\n (m: Iterable[tuple[str, ndarray | None]], /) -> None\n (m: Iterable[tuple[str, ndarray | None]], /, ...) -> None\n (**kwargs: ndarray | None) -> None\n Argument `WaldResult` is not assignable to parameter `m` with type `SupportsKeysAndGetItem[str, ndarray | None]` in function `typing.MutableMapping.update`",
"concise_description": "No matching overload found for function `typing.MutableMapping.update` called with arguments: (WaldResult)",
"severity": "error"
},
{
"line": 1764,
"column": 17,
"stop_line": 1764,
"stop_column": 24,
"path": "src/jamma/lmm/compute_numpy.py",
"code": -2,
"name": "bad-argument-type",
"description": "Argument `float | None` is not assignable to parameter `logl_H0` with type `float` in function `_compute_lrt_numpy`\n The declared type does not allow `None`. Consider narrowing the value with an `is not None` check.",
"concise_description": "Argument `float | None` is not assignable to parameter `logl_H0` with type `float` in function `_compute_lrt_numpy`",
"severity": "error"
},
{
"line": 1774,
"column": 17,
"stop_line": 1774,
"stop_column": 29,
"path": "src/jamma/lmm/compute_numpy.py",
"code": -2,
"name": "bad-argument-type",
"description": "Argument `ndarray | None` is not assignable to parameter `Hi_eval_null` with type `ndarray` in function `_compute_score_numpy`\n The declared type does not allow `None`. Consider narrowing the value with an `is not None` check.",
"concise_description": "Argument `ndarray | None` is not assignable to parameter `Hi_eval_null` with type `ndarray` in function `_compute_score_numpy`",
"severity": "error"
},
{
"line": 1787,
"column": 13,
"stop_line": 1787,
"stop_column": 25,
"path": "src/jamma/lmm/compute_numpy.py",
"code": -2,
"name": "bad-argument-type",
"description": "Argument `ndarray | None` is not assignable to parameter `Hi_eval_null` with type `ndarray` in function `_compute_score_numpy`\n The declared type does not allow `None`. Consider narrowing the value with an `is not None` check.",
"concise_description": "Argument `ndarray | None` is not assignable to parameter `Hi_eval_null` with type `ndarray` in function `_compute_score_numpy`",
"severity": "error"
},
{
"line": 1802,
"column": 17,
"stop_line": 1802,
"stop_column": 24,
"path": "src/jamma/lmm/compute_numpy.py",
"code": -2,
"name": "bad-argument-type",
"description": "Argument `float | None` is not assignable to parameter `logl_H0` with type `float` in function `_compute_lrt_numpy`\n The declared type does not allow `None`. Consider narrowing the value with an `is not None` check.",
"concise_description": "Argument `float | None` is not assignable to parameter `logl_H0` with type `float` in function `_compute_lrt_numpy`",
"severity": "error"
},
{
"line": 1808,
"column": 22,
"stop_line": 1821,
"stop_column": 10,
"path": "src/jamma/lmm/compute_numpy.py",
"code": -2,
"name": "no-matching-overload",
"description": "No matching overload found for function `typing.MutableMapping.update` called with arguments: (WaldResult)\n Possible overloads:\n (m: SupportsKeysAndGetItem[str, ndarray | None], /) -> None [closest match]\n (m: SupportsKeysAndGetItem[str, ndarray | None], /, ...) -> None\n (m: Iterable[tuple[str, ndarray | None]], /) -> None\n (m: Iterable[tuple[str, ndarray | None]], /, ...) -> None\n (**kwargs: ndarray | None) -> None\n Argument `WaldResult` is not assignable to parameter `m` with type `SupportsKeysAndGetItem[str, ndarray | None]` in function `typing.MutableMapping.update`",
"concise_description": "No matching overload found for function `typing.MutableMapping.update` called with arguments: (WaldResult)",
"severity": "error"
},
{
"line": 178,
"column": 23,
Expand Down
79 changes: 65 additions & 14 deletions src/jamma/lmm/compute_numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
from __future__ import annotations

from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypedDict
from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypedDict, TypeVar

import numpy as np

Expand Down Expand Up @@ -1679,6 +1679,54 @@ def _compute_lrt_split_numpy(
)


_LOGL_H0_REQUIRED = "logl_H0 is required for LRT (mode 2) and All (mode 4)"
_HI_EVAL_NULL_REQUIRED = "Hi_eval_null is required for Score (mode 3) and All (mode 4)"

_ModeInput = TypeVar("_ModeInput")


def _require_mode_input(value: _ModeInput | None, message: str) -> _ModeInput:
"""Return a mode-specific input, or raise naming which mode needed it.

Only some LMM modes take logl_H0 / Hi_eval_null, so both arrive optional.
Checking them at the point of use ties the guard to the branch that reads
the value, which a mode check made up front cannot express.

Args:
value: The optional input.
message: Error text naming the input and the modes that require it.

Returns:
The value, known to be present.

Raises:
ValueError: If the value is None.
"""
if value is None:
raise ValueError(message)
return value


def _store_wald(result: dict[str, np.ndarray | None], wald: WaldResult) -> None:
"""Copy a WaldResult's five arrays into the mode-agnostic result dict.

Spelled out per key rather than ``result.update(wald)`` because a TypedDict
is not a ``Mapping[str, ndarray | None]`` — its value types are per-key, so
the update overloads reject it.

Args:
result: The chunk result dict to populate.
wald: Wald statistics for the chunk.
"""
result.update(
lambdas=wald["lambdas"],
logls=wald["logls"],
betas=wald["betas"],
ses=wald["ses"],
pwalds=wald["pwalds"],
)


def compute_lmm_chunk_numpy(
lmm_mode: LmmMode,
n_cvt: int,
Expand Down Expand Up @@ -1720,11 +1768,6 @@ def compute_lmm_chunk_numpy(
"""
n_refine = max(n_refine, 20)

if lmm_mode in (2, 4) and logl_H0 is None:
raise ValueError("logl_H0 is required for LRT (mode 2) and All (mode 4)")
if lmm_mode in (3, 4) and Hi_eval_null is None:
raise ValueError("Hi_eval_null is required for Score (mode 3) and All (mode 4)")

result: dict[str, np.ndarray | None] = {
"lambdas": None,
"logls": None,
Expand All @@ -1737,7 +1780,8 @@ def compute_lmm_chunk_numpy(
}

if lmm_mode == 1:
result.update(
_store_wald(
result,
_compute_wald_numpy(
n_cvt,
eigenvalues,
Expand All @@ -1748,7 +1792,7 @@ def compute_lmm_chunk_numpy(
n_grid,
n_refine,
n_threads=n_threads,
)
),
)

elif lmm_mode == 2:
Expand All @@ -1761,7 +1805,7 @@ def compute_lmm_chunk_numpy(
l_max,
n_grid,
n_refine,
logl_H0,
_require_mode_input(logl_H0, _LOGL_H0_REQUIRED),
n_threads=n_threads,
)
)
Expand All @@ -1771,20 +1815,26 @@ def compute_lmm_chunk_numpy(
_compute_score_numpy(
n_cvt,
eigenvalues,
Hi_eval_null,
_require_mode_input(Hi_eval_null, _HI_EVAL_NULL_REQUIRED),
Uab_batch,
n_samples,
n_threads=n_threads,
)
)

elif lmm_mode == 4:
# Both mode-4 inputs are checked before any compute runs, so an omitted
# one still fails before the Score step rather than partway through.
# logl_H0 first: with both absent that is the one reported, which is
# the order the previous up-front guards established.
null_logl = _require_mode_input(logl_H0, _LOGL_H0_REQUIRED)
hi_eval_null = _require_mode_input(Hi_eval_null, _HI_EVAL_NULL_REQUIRED)
# Compose all three tests; only take p_scores from Score —
# Wald provides REML-optimized beta/SE below
score_result = _compute_score_numpy(
n_cvt,
eigenvalues,
Hi_eval_null,
hi_eval_null,
Uab_batch,
n_samples,
n_threads=n_threads,
Expand All @@ -1799,13 +1849,14 @@ def compute_lmm_chunk_numpy(
l_max,
n_grid,
n_refine,
logl_H0,
null_logl,
n_threads=n_threads,
)
)
# Pre-compute Iab once for Wald (lambda-independent)
Iab_batch = batch_compute_iab_numpy(n_cvt, Uab_batch)
result.update(
_store_wald(
result,
_compute_wald_numpy(
n_cvt,
eigenvalues,
Expand All @@ -1817,7 +1868,7 @@ def compute_lmm_chunk_numpy(
n_refine,
Iab_batch=Iab_batch,
n_threads=n_threads,
)
),
)

else:
Expand Down
Loading