Type-hint ratchet, an ArrayBox stub, and the API bugs turning mypy back on found

[derrynknife]

27 commits, 77 files, +2538/−591. It began as the docs extra from #141 and grew each time validating one thing exposed another: the extra revealed a broken docs build, the docs fix silently broke lint, the broken lint hid mypy for ten commits, and turning mypy back on found real bugs in the distributions.

⚠️ This PR contains a breaking change. Bernoulli.from_params(p=...) and ExactEventTime.from_params(T=...) now raise TypeError. See §6.

---

1. docs extra (part of #141)

pip install -e ".[docs]" installs the documentation toolchain. docs/requirements.txt is removed rather than kept beside it — two copies of the same pinned toolchain is the arrangement that drifts, and the repo already settled this pattern for tests. Read the Docs installs the extra via extra_requirements.

Pins carried over unchanged, including ipykernel==6.31.0 and the reason for it.

2. The docs build was broken on develop

Validating the extra meant running a real build, which failed:

ValueError: Gamma distribution does not work with probability plot fitting

A jupyter-execute cell looped over ['MPP', 'MOM', 'MSE', 'MPS', 'MLE'] for a shifted Gamma; #338 made Gamma.fit(how="MPP") raise. Documentation cells execute during the build, so the build failed. The surrounding prose had gone stale the same way and now explains why the Gamma has no probability plot at all.

3. What CI runs, where

Nothing caught the above, because CI never built the documentation and Read the Docs builds only master. Separately, the full suite on every develop PR was the slowest part of working on the package.

Event	Jobs	Time
PR → develop	lint	~1 min
PR → master (release)	lint + suite ×3 + docs build	~10 min
push to master / tag	lint + suite	~9 min

Both conditions are now proven in both directions — PR #340 was opened into master purely to confirm the jobs fire, and they did. (The original version of this description carried a caveat that only the skip half was proven; that caveat is now resolved.)

scripts/check_all_pythons.py is the other half of the trade: one command runs the suite and both doctest passes on 3.11/3.12/3.13 and refuses to say "passed" unless all nine did. Contributing.rst carries the table and the timings.

4. Documentation: API pages, then -W

New reference pages for the previously undocumented public surfaces (multivariate, beta, frailty, Buckley-James; extended degradation). All 18 build warnings cleared — mostly autosectionlabel duplicates against the changelog's repeated headings, fixed with autosectionlabel_maxdepth = 1 rather than by suppression. -W is now on in both CI and .readthedocs.yaml; enabling one without the other lets a build go green while publishing broken pages.

5. #287 — the survival tree's log-rank split statistic was wrong

kind="non-parametric" trees, and any RandomSurvivalForest built from them, selected splits on a statistic off by factors of several — and reordered, not merely inflated: of two candidate splits, 1.8564 → 0.2765 and 0.2765 → 1.2247. The at-risk count is now computed as Y(τ) = #{x ≥ τ} − #{tl ≥ τ} via two suffix sums and a searchsorted, which is also O(N log N) rather than the previous O(N×G). Verified against a naive reference over 200 random partitions plus left-truncation, censored-tail and tie cases.

6. ⚠️ BREAKING: from_params argument names

ParametricFitter.from_params names its first argument params. Bernoulli called it p and ExactEventTime called it T, so positional calls worked and keyword calls raised:

Bernoulli.from_params(0.5)          # OK
Bernoulli.from_params(params=0.5)   # TypeError

That is the shape of bug a test suite never catches, because every internal call and every docstring example passes positionally. It only broke for someone writing generic code, on two of twenty-five distributions.

Bernoulli's was worse than a rename. The base's p is the proportion that never fails. So p=0.5 meant the never-fails fraction on twenty-four distributions and the event probability on Bernoulli — the same keyword, sibling classes, unrelated meanings, and no error either way.

What breaks: Bernoulli.from_params(p=...) and ExactEventTime.from_params(T=...) now raise TypeError. Positional calls are unaffected and unchanged in value. Nothing changes meaning silently — params has no default, so the old keyword forms fail loudly rather than being reinterpreted as the never-fails proportion. There are no keyword callers in the repository or the documentation.

All three closed-form distributions also accept gamma, p and f0 now and reject them with a ValueError naming the distribution.

7. OptimisedFitMixin — the estimation machinery left the base class

ParametricFitter.fit takes 18 named arguments. Three distributions cannot honour any of them:

Bernoulli.fit([0, 1], c=[0, 0])              TypeError
Bernoulli.fit([0, 1], how="MLE")             TypeError
Binomial.fit([1, 2], n_trials=3, how="MLE")  TypeError

An audit of all 25 subclasses found these three and no others. fit and the twelve methods it needs moved to OptimisedFitMixin, which the 21 distributions that have them inherit alongside ParametricFitter.

Every distribution is still a ParametricFitter — that is what the isinstance gates in parametric, mixture_model, parametric_regression_model, frailty_model and renewal_model check. The mixin only adds the estimation methods on top.

The point is that the wrong thing is now unwriteable rather than merely undocumented:

fit_best's candidates: list[OptimisedFitMixin]
  -> adding Bernoulli is a type error, not a runtime one

Annotate a parameter OptimisedFitMixin when it must be fittable by a chosen method, ParametricFitter when only the distribution functions are needed.

8. Every distribution exports its own type

17 distributions read Weibull: ParametricFitter = Weibull_("Weibull"), which erases the concrete class. Since the base declares none of sf/ff/df/hf/Hf/qf/mean, the example in each distribution's own docstring did not type check for anyone whose checker honours py.typed:

Weibull.sf(x, 3, 4)
error: "ParametricFitter" has no attribute "sf"

The annotation cannot simply be dropped and inferred — the regression subpackages and fit_best import these names, and mypy cannot resolve them through that cycle. They name the concrete class instead.

9. Type-hint ratchet (#143) — 17 modules enforced

disallow_untyped_defs is set per-module, so an unannotated function in a listed module is an error and coverage cannot slip back. A module joins once it is clean; the remaining 1206 unannotated functions do not have to be finished first.

Enforced: distribution, serialisation, metrics.*, univariate.information_criteria, datasets, all of univariate.nonparametric.*, recurrent.nonparametric.*, univariate.regression.frailty.*, and 9 of the 25 parametric distributions (Weibull plus all 8 discrete).

Annotating is what makes mypy check a function body, and that found real defects rather than paperwork:

• SerialisableMixin.to_json/from_json called to_dict/from_dict, which the mixin never declared. Now declared under TYPE_CHECKING — a real stub would be inherited, and copula_model probes with hasattr.
• ESTIMATOR_FUNCS was built from nonp.nelson_aalen and siblings while the package __init__ was still running, so the names resolved to the submodules, not the functions. It worked by load order, not by guarantee.
• success_run tested confidence and alpha for truthiness, so passing both with either set to zero skipped the "only one of" raise, and confidence=0 fell through every branch leaving alpha as None.
• turnbull, rank_adjust and NonParametricCounting.from_xrd declared array-like parameters and then indexed and divided them. _logrank_z_v declared c/n as arrays while handling None for both. NonParametricCounting.fit declared windows as array-like when it is the dictionary its own docstring describes.
• _parameter_initialiser disagrees on its return type across distributions — Weibull returns a tuple, the discrete ones return arrays. Harmless today; the signatures now say which is which.

The distributions needed a type vocabulary first

A distribution deals in two kinds of value and only one can be an autograd box. Instrumenting a live fit showed the runtime types directly:

(x, alpha) = ('ndarray', 'float64')     evaluating the likelihood
(x, alpha) = ('ndarray', 'ArrayBox')    differentiating it

So Numeric = npt.NDArray | float for what a function is evaluated at, and Boxable = npt.NDArray | float | ArrayBox for a parameter or anything derived from one.

ArrayBox needed a stub, because autograd ships no type information and mypy otherwise rejects it in a union. stubs/autograd/ describes that one class; permissive __getattr__ stubs at the package level keep the rest of autograd exactly as untyped as it was (supplying any stub makes mypy consider the whole package described, which otherwise produced 328 spurious errors).

This matters beyond tidiness. The np.asarray fix applied to turnbull and mcf above would be a silent numerical bug here. np.asarray on a box does not reject it — it wraps it in a 0-d object array, the forward value stays correct, and only the derivative is damaged:

alpha * x + beta      plain [6. 3.]   asarray [0. 3.]   <- no exception
(x / alpha) ** beta   plain [-1.613 -0.094]   asarray TypeError

A zero gradient is not an error to an optimiser; it means "this parameter does not affect the likelihood". The fit leaves the parameter at its initial guess and reports success.

10. The lint job had been red for ten commits

Fix the last two documentation build warnings turned a plain reference into a :class: cross-reference and pushed one line to 80 characters. flake8 runs before mypy in the lint job, so from that commit on mypy was skipped, not run — including through three ratchet batches. Bisecting:

6ef57a1  Success: no issues found in 305 source files
1c6d563  Success: no issues found in 305 source files
dd50376  Found 78 errors in 5 files
49077ef  Found 154 errors in 6 files

All fixed. A related finding: mypy's result is interpreter-dependent — d_ghosts, an array in one branch and a scalar in the other, is rejected on 3.12 and accepted on 3.11. A local mypy run on one interpreter is not a proxy for the lint job.

Verification

• scripts/check_all_pythons.py: 9/9 green on 3.11/3.12/3.13 — suite + both doctest passes — re-run after every batch. Final: 2166 passed, 271 skipped, 5 xfailed, 229 doctests each.
• mypy clean on all three interpreters (not just one, for the reason in §10).
• sphinx -b html -W --keep-going exits 0.
• Each ratchet entry verified to bite: adding an unannotated probe function to a listed module produces exactly one error.
• Behavioural diffs against the pre-change tree: 73 values for the OptimisedFitMixin split (every distribution, four estimation methods, aic/neg_ll/mean, fit_best, a mixture fit) and a separate diff for the discrete batch — no differences in either.
• Every # type: ignore[override] introduced during this work has been removed by fixing the underlying divergence. None remain as suppressions.

Known-imperfect, stated rather than glossed

• Four files unrelated to their commit (three tests, mixture_model.py) were swept in by git add -A; the changes are one isort import-reordering line each, zero behaviour, and were covered by the full test run. The commit messages do not mention them.
• isort non-conformance is repo-wide and pre-existing — 28 files at the branch point, 15 now. It runs in pre-commit, not CI.

Follow-ups, not included

• The remaining 15 continuous distributions for the ratchet, and the 15 exports still carrying the : ParametricFitter erasure (§8 fixed 16 of 31; the rest go as each file is annotated).
• Adding mypy to check_all_pythons.py — its docstring reasons that "lint runs on every pull request anyway", which §10 disproves twice over.
• Declaring sf/ff/df/hf/Hf/qf on ParametricFitter. Its docstring states the contract; the class does not express it, which is why Discretize holds its wrapped distribution as Any.
• A scheduled build on develop to close the detection window from a release cycle to a day.

🤖 Generated with Claude Code

https://claude.ai/code/session_01TRhKL2fJBiAAfNo9rihwts