From 9c4b5c590e1dd26807f59f5dcb07f964c6b95798 Mon Sep 17 00:00:00 2001 From: Ethan Buckley Date: Mon, 27 Jul 2026 11:00:11 +0100 Subject: [PATCH] refactor(aligned): typed run plan + recorded estimand/population metadata (#394 pillar 4) Extends the ITT / gain-factor / level-factor / DiD / concurrent run-plan pattern to the per-protocol onset-aligned family. A new aligned.py provides AlignedModelSettings (fail-fast on a misspelled spec.extra key; validates the 4 legacy keys) and a resolved AlignedRunPlan that drives data preparation, factory construction and the config.json / model_recipe.md audit trail, recording the design, estimand, causal status (per-protocol association -- the cohort term is NOT the randomised ITT effect) and analysis population alongside every fit. fit_aligned now resolves and validates the plan before the context resets an output directory or the loader reads data, then loads via plan.prepare_kwargs() (the onset-aligned loader; include_dose requests the session-dose covariate only for the sensitivity variant) and builds via plan.factory_kwargs(). reporting._resolved_run_plan gains an aligned branch so write_model_recipe and config.json's resolved_run_plan pick it up. Byte-identical: the resolved prepare + factory kwargs reproduce the former inline fit_aligned logic exactly for all 9 registered aligned models (verified programmatically). al-001 dev fit is clean end-to-end and now emits model_recipe.md + a populated resolved_run_plan. 12 new run-plan tests; drift suites pass; ruff / format:check / spellcheck clean. Part of #394 (pillar 4: typed settings + resolved plans). No estimand, likelihood, prior, population, fitted equation, sampling preset or artefact schema changes. Co-Authored-By: Claude Opus 4.8 --- .../statistical_models/aligned.py | 245 ++++++++++++++++++ .../statistical_models/pipeline.py | 36 ++- .../statistical_models/reporting.py | 16 +- .../test_aligned_run_plan.py | 180 +++++++++++++ 4 files changed, 455 insertions(+), 22 deletions(-) create mode 100644 src/language_reading_predictors/statistical_models/aligned.py create mode 100644 tests/statistical_models/test_aligned_run_plan.py diff --git a/src/language_reading_predictors/statistical_models/aligned.py b/src/language_reading_predictors/statistical_models/aligned.py new file mode 100644 index 00000000..02f78a2c --- /dev/null +++ b/src/language_reading_predictors/statistical_models/aligned.py @@ -0,0 +1,245 @@ +# Copyright (c) 2026 Down Syndrome Education International and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Typed settings and a resolved run plan for the onset-aligned family (#394 pillar 4). + +Mirrors the ITT / gain-factor / level-factor / DiD / concurrent run-plan pattern for +the per-protocol onset-aligned (``kind="aligned"``) models. A model module declares +its settings; the plan is resolved and **validated before any data are loaded or an +output directory is reset**, then drives data preparation, factory construction and +the ``config.json`` / ``model_recipe.md`` audit trail. This removes the untyped +``spec.extra`` boundary (where a misspelled key silently defaulted) and records the +resolved design, estimand, causal status, analysis population and missing-data +assumption alongside every fit. + +The aligned design is a per-protocol onset-aligned single-gain ANCOVA (LRPAL): a +cross-sectional Beta-Binomial regression of the aligned post-score on its own onset +baseline, age-at-onset and cognitive ability, optionally with a cohort indicator and +the cumulative session dose. One row per child, so there is **no child random +intercept**. The cohort contrast is a **per-protocol association**, not the +randomised ITT effect (it is confounded by age-at-onset and cohort/timing), and the +dose term is a collider descendant of group and ability -- a sensitivity variant. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import asdict, dataclass +from typing import Any + +from language_reading_predictors.statistical_models.context import ModelSpec + +# The complete, closed set of legacy ``spec.extra`` keys the aligned family +# understands. Anything else is a typo and must fail before a fit starts. +_LEGACY_KEYS = frozenset( + { + "ability_covariate", + "use_cohort", + "use_dose", + "likelihood", + } +) + +_LIKELIHOODS = frozenset({"beta_binomial", "bernoulli_offfloor"}) + + +@dataclass(frozen=True, slots=True) +class AlignedModelSettings: + """Immutable settings declared by a single onset-aligned model module. + + Defaults encode the primary per-protocol ANCOVA: the cohort contrast on, without + the collider-descendant session-dose covariate, and the Beta-Binomial working + likelihood. + """ + + ability_covariate: str | None = None + use_cohort: bool = True + use_dose: bool = False + likelihood: str = "beta_binomial" + + def __post_init__(self) -> None: + if self.ability_covariate is not None and ( + not isinstance(self.ability_covariate, str) or not self.ability_covariate + ): + raise TypeError("ability_covariate must be a non-empty string or None") + for flag in ("use_cohort", "use_dose"): + if not isinstance(getattr(self, flag), bool): + raise TypeError(f"{flag} must be bool") + if self.likelihood not in _LIKELIHOODS: + raise ValueError( + f"likelihood must be one of {sorted(_LIKELIHOODS)}, got {self.likelihood!r}" + ) + + @classmethod + def from_legacy_extra( + cls, extra: Mapping[str, Any], *, model_id: str + ) -> AlignedModelSettings: + """Strictly translate the former ``spec.extra`` dictionary boundary. + + Rejects unknown keys so a misspelling fails before data loading rather than + silently taking a default.""" + unknown = sorted(set(extra) - _LEGACY_KEYS) + if unknown: + raise ValueError( + f"{model_id}: unknown aligned setting(s): {', '.join(unknown)}. " + "Declare AlignedModelSettings so misspellings fail fast." + ) + # Pass raw values through so __post_init__ is the single validation/coercion + # point; pre-coercing here would silently reshape misshaped legacy settings. + return cls( + ability_covariate=extra.get("ability_covariate"), + use_cohort=extra.get("use_cohort", True), + use_dose=extra.get("use_dose", False), + likelihood=extra.get("likelihood", "beta_binomial"), + ) + + +@dataclass(frozen=True, slots=True) +class AlignedRunPlan: + """Concrete, validated instructions consumed by preparation and modelling.""" + + model_id: str + outcome_symbol: str + settings_source: str + ability_covariate: str | None + use_cohort: bool + use_dose: bool + likelihood: str + off_floor: bool + # Recorded audit metadata (#394 pillar 4). + design: str + estimand: str + causal_status: str + analysis_population: str + missing_data_assumption: str + + @property + def obs_node(self) -> str: + return "y_offfloor" if self.off_floor else "y_post" + + def as_dict(self) -> dict[str, Any]: + """Return the JSON-ready run-plan contract for ``config.json``.""" + return asdict(self) + + def prepare_kwargs(self) -> dict[str, Any]: + """Arguments for ``load_and_prepare_aligned`` from the resolved plan. + + The aligned family uses its own onset-aligned loader (one row per child at the + per-protocol onset window); ``include_dose`` requests the cumulative-session + covariate only when the dose sensitivity variant is fit.""" + return { + "outcomes": (self.outcome_symbol,), + "ability_covariate": self.ability_covariate, + "include_dose": self.use_dose, + } + + def factory_kwargs(self) -> dict[str, Any]: + """Arguments for ``build_aligned_model`` for this plan.""" + return { + "outcome_symbol": self.outcome_symbol, + "ability_covariate": self.ability_covariate, + "use_cohort": self.use_cohort, + "use_dose": self.use_dose, + "likelihood": self.likelihood, + } + + def recipe_markdown(self, *, title: str) -> str: + """Undergraduate-friendly explanation generated from the resolved plan.""" + return ( + "Note: Generated from the validated onset-aligned run plan; template " + "drafted by an LLM-based AI tool (Claude Code/Opus 4.8).\n\n" + f"# Model recipe: {title}\n\n" + f"Model ID: `{self.model_id}`.\n\n" + f"## Design\n\n{self.design}\n\n" + f"## Estimand\n\n{self.estimand}\n\n" + f"## Causal status\n\n{self.causal_status}\n\n" + f"## Analysis population\n\n{self.analysis_population}\n\n" + f"## Missing data\n\n{self.missing_data_assumption}\n\n" + "## Terms\n\n" + f"Outcome: `{self.outcome_symbol}`. Ability covariate: " + f"{self.ability_covariate or 'none'}. Cohort contrast: {self.use_cohort}. " + f"Cumulative session dose (sensitivity): {self.use_dose}.\n\n" + "## Uncertainty and checks\n\n" + "The fit reports a posterior distribution; interpret it only after the " + "convergence gate and posterior-predictive checks pass. The saved " + "`config.json` contains the same resolved run plan in machine-readable " + "form.\n" + ) + + +def declared_aligned_settings(spec: ModelSpec) -> tuple[AlignedModelSettings, str]: + """Return typed settings and their source, rejecting mixed declarations.""" + settings = spec.model_settings + if settings is not None: + if spec.extra: + raise ValueError( + f"{spec.model_id}: aligned settings cannot be split between " + "model_settings and extra" + ) + if not isinstance(settings, AlignedModelSettings): + raise TypeError( + f"{spec.model_id}: kind='aligned' requires AlignedModelSettings, got " + f"{type(settings).__name__}" + ) + return settings, "typed" + return ( + AlignedModelSettings.from_legacy_extra(spec.extra, model_id=spec.model_id), + "legacy_extra", + ) + + +def resolve_aligned_run_plan(spec: ModelSpec) -> AlignedRunPlan: + """Resolve and validate an onset-aligned spec before any data are loaded.""" + if spec.kind != "aligned": + raise ValueError(f"{spec.model_id}: expected kind 'aligned', got {spec.kind!r}") + if not spec.outcome_symbol: + raise ValueError( + f"{spec.model_id}: outcome_symbol is required for an aligned model" + ) + + settings, source = declared_aligned_settings(spec) + own = spec.outcome_symbol + off_floor = settings.likelihood == "bernoulli_offfloor" + + design = ( + "Per-protocol onset-aligned single-gain ANCOVA: a cross-sectional " + "Beta-Binomial regression of the aligned post-score on its own onset baseline, " + "age-at-onset and cognitive ability, optionally with a cohort indicator and " + "the cumulative session dose. One row per child, so no child random intercept." + ) + estimand = ( + "The cohort contrast at the two arms' own onset-aligned endpoints -- a " + "per-protocol association, NOT the randomised ITT effect (it is confounded by " + "age-at-onset and cohort/timing). With the dose variant, the cumulative-session " + "covariate is a collider descendant of group and ability, a sensitivity variant." + ) + causal_status = ( + "Associational / per-protocol: no randomised contrast is estimated. The cohort " + "term is confounded by age-at-onset and cohort/timing and the dose term is a " + "collider descendant; neither is the ITT treatment effect." + ) + analysis_population = ( + "Available-case children with onset-aligned pre and post scores (the " + "per-protocol onset window)." + ) + missing_data_assumption = ( + "Available-case analysis under ignorable missingness: children without an " + "onset-aligned pre+post pair are dropped and assumed ignorable given the " + "modelled covariates." + ) + + return AlignedRunPlan( + model_id=spec.model_id, + outcome_symbol=own, + settings_source=source, + ability_covariate=settings.ability_covariate, + use_cohort=settings.use_cohort, + use_dose=settings.use_dose, + likelihood=settings.likelihood, + off_floor=off_floor, + design=design, + estimand=estimand, + causal_status=causal_status, + analysis_population=analysis_population, + missing_data_assumption=missing_data_assumption, + ) diff --git a/src/language_reading_predictors/statistical_models/pipeline.py b/src/language_reading_predictors/statistical_models/pipeline.py index aab60fe5..46dc86a5 100644 --- a/src/language_reading_predictors/statistical_models/pipeline.py +++ b/src/language_reading_predictors/statistical_models/pipeline.py @@ -73,6 +73,9 @@ save_plotcollection, save_styled_figure, ) +from language_reading_predictors.statistical_models.aligned import ( + resolve_aligned_run_plan, +) from language_reading_predictors.statistical_models.context import ( ModelSpec, StatisticalFitContext, @@ -6126,33 +6129,24 @@ def _al_diag_vars(spec: ModelSpec) -> list[str]: def fit_aligned(spec: ModelSpec, config: str = "dev") -> StatisticalFitContext: _require_spec(spec, "aligned", outcome=True) + # Resolve and validate the family contract before the context resets an output + # directory or the loader reads any data (#394 pillar 4). One plan drives + # preparation, factory arguments, the teaching recipe and config.json. + plan = resolve_aligned_run_plan(spec) ctx = make_context(spec, config) - extra = spec.extra + ctx.resolved_plan = plan + _report.write_model_recipe(ctx) + + off_floor = plan.off_floor + obs_node = plan.obs_node section_header("Prepare data") - ability_covariate = extra.get("ability_covariate") - use_cohort = bool(extra.get("use_cohort", True)) - use_dose = bool(extra.get("use_dose", False)) - likelihood = extra.get("likelihood", "beta_binomial") - off_floor = likelihood == "bernoulli_offfloor" - obs_node = "y_offfloor" if off_floor else "y_post" - prepared = load_and_prepare_aligned( - outcomes=(spec.outcome_symbol,), - ability_covariate=ability_covariate, - include_dose=use_dose, - ) + prepared = load_and_prepare_aligned(**plan.prepare_kwargs()) ctx.prepared = prepared _print_header(ctx) section_header("Build model") - built = _factories.build_aligned_model( - prepared, - outcome_symbol=spec.outcome_symbol, - ability_covariate=ability_covariate, - use_cohort=use_cohort, - use_dose=use_dose, - likelihood=likelihood, - ) + built = _factories.build_aligned_model(prepared, **plan.factory_kwargs()) _attach_built(ctx, built) _render_model_graph(ctx) @@ -6203,7 +6197,7 @@ def fit_aligned(spec: ModelSpec, config: str = "dev") -> StatisticalFitContext: # Items-scale cohort contrast (immediate vs wait-list at aligned endpoints). # This is a PER-PROTOCOL association, NOT a randomised treatment effect -- # confounded by age-at-onset and cohort/timing (see the LRPAL design note). - if use_cohort: + if plan.use_cohort: cohort = built.prepared.G.astype(float) n_marg = 1 if off_floor else built.prepared.n_trials[spec.outcome_symbol] cme = _report.treatment_marginal_effect( diff --git a/src/language_reading_predictors/statistical_models/reporting.py b/src/language_reading_predictors/statistical_models/reporting.py index 8fac06fe..abde0c77 100644 --- a/src/language_reading_predictors/statistical_models/reporting.py +++ b/src/language_reading_predictors/statistical_models/reporting.py @@ -27,6 +27,10 @@ from language_reading_predictors.statistical_models.context import ( StatisticalFitContext, ) +from language_reading_predictors.statistical_models.aligned import ( + AlignedRunPlan, + resolve_aligned_run_plan, +) from language_reading_predictors.statistical_models.gain_factors import ( GainFactorsRunPlan, resolve_gain_factors_run_plan, @@ -2273,12 +2277,22 @@ def _gain_factors_run_plan(context: StatisticalFitContext) -> GainFactorsRunPlan return resolve_gain_factors_run_plan(context.spec) +def _aligned_run_plan(context: StatisticalFitContext) -> AlignedRunPlan: + """Return the aligned plan resolved before loading, or reconstruct it.""" + resolved_plan = getattr(context, "resolved_plan", None) + if isinstance(resolved_plan, AlignedRunPlan): + return resolved_plan + return resolve_aligned_run_plan(context.spec) + + def _resolved_run_plan(context: StatisticalFitContext): - """The typed run plan for families that have one (ITT, gain factors), else None.""" + """The typed run plan for families that have one (ITT, gain, aligned), else None.""" if context.spec.kind == "itt": return _itt_run_plan(context) if context.spec.kind == "gain_factors": return _gain_factors_run_plan(context) + if context.spec.kind == "aligned": + return _aligned_run_plan(context) return None diff --git a/tests/statistical_models/test_aligned_run_plan.py b/tests/statistical_models/test_aligned_run_plan.py new file mode 100644 index 00000000..b9f5bab5 --- /dev/null +++ b/tests/statistical_models/test_aligned_run_plan.py @@ -0,0 +1,180 @@ +# Copyright (c) 2026 Down Syndrome Education International and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Tests for the typed onset-aligned settings and run plan (#394 pillar 4).""" + +from __future__ import annotations + +import glob +import importlib +import os + +import pytest + +from language_reading_predictors.statistical_models.aligned import ( + AlignedModelSettings, + AlignedRunPlan, + resolve_aligned_run_plan, +) +from language_reading_predictors.statistical_models.context import ModelSpec + +_META_FIELDS = ( + "design", + "estimand", + "causal_status", + "analysis_population", + "missing_data_assumption", +) + + +def _aligned_specs() -> list[ModelSpec]: + """Every registered onset-aligned model's SPEC.""" + root = os.path.dirname( + importlib.import_module( + "language_reading_predictors.statistical_models.aligned" + ).__file__ + ) + specs: list[ModelSpec] = [] + for path in sorted(glob.glob(os.path.join(root, "lrp_rli_al_*.py"))): + mod = importlib.import_module( + "language_reading_predictors.statistical_models." + os.path.basename(path)[:-3] + ) + spec = getattr(mod, "SPEC", None) + if spec is not None and spec.kind == "aligned": + specs.append(spec) + return specs + + +# --- settings validation ------------------------------------------------------ + + +def test_settings_reject_unknown_likelihood(): + with pytest.raises(ValueError, match="likelihood"): + AlignedModelSettings(likelihood="poisson") + + +def test_settings_reject_non_bool_use_cohort(): + with pytest.raises(TypeError, match="use_cohort"): + AlignedModelSettings(use_cohort=1) # type: ignore[arg-type] + + +def test_settings_reject_empty_ability_covariate(): + with pytest.raises(TypeError, match="ability_covariate"): + AlignedModelSettings(ability_covariate="") + + +def test_from_legacy_extra_rejects_unknown_key(): + with pytest.raises(ValueError, match="unknown aligned setting"): + AlignedModelSettings.from_legacy_extra( + {"use_dose": True, "use_doze": True}, # typo + model_id="lrp-rli-al-999", + ) + + +def test_from_legacy_extra_round_trips_known_keys(): + settings = AlignedModelSettings.from_legacy_extra( + { + "ability_covariate": "blocks", + "use_cohort": False, + "use_dose": True, + "likelihood": "bernoulli_offfloor", + }, + model_id="lrp-rli-al-999", + ) + assert settings.ability_covariate == "blocks" + assert settings.use_cohort is False + assert settings.use_dose is True + assert settings.likelihood == "bernoulli_offfloor" + + +# --- resolve ------------------------------------------------------------------ + + +def _spec(**extra) -> ModelSpec: + return ModelSpec( + model_id="lrp-rli-al-000", + kind="aligned", + title="test", + outcome_symbol="W", + extra=extra, + ) + + +def test_resolve_rejects_wrong_kind(): + spec = ModelSpec(model_id="x", kind="itt", title="t", outcome_symbol="W") + with pytest.raises(ValueError, match="expected kind 'aligned'"): + resolve_aligned_run_plan(spec) + + +def test_resolve_primary_is_per_protocol_association(): + plan = resolve_aligned_run_plan(_spec(ability_covariate="blocks")) + assert plan.settings_source == "legacy_extra" + assert not plan.off_floor and plan.obs_node == "y_post" + prep = plan.prepare_kwargs() + assert prep == { + "outcomes": ("W",), + "ability_covariate": "blocks", + "include_dose": False, + } + assert plan.factory_kwargs()["use_cohort"] is True + assert "per-protocol" in plan.causal_status.lower() + assert "not the randomised itt" in plan.estimand.lower() + + +def test_resolve_dose_variant_requests_include_dose(): + plan = resolve_aligned_run_plan(_spec(use_dose=True)) + assert plan.prepare_kwargs()["include_dose"] is True + assert plan.factory_kwargs()["use_dose"] is True + + +def test_resolve_off_floor_sets_bernoulli_node(): + plan = resolve_aligned_run_plan(_spec(likelihood="bernoulli_offfloor")) + assert plan.off_floor + assert plan.obs_node == "y_offfloor" + assert plan.factory_kwargs()["likelihood"] == "bernoulli_offfloor" + + +def test_typed_settings_are_accepted_and_sourced(): + spec = ModelSpec( + model_id="lrp-rli-al-000", + kind="aligned", + title="test", + outcome_symbol="W", + model_settings=AlignedModelSettings(use_dose=True), + ) + plan = resolve_aligned_run_plan(spec) + assert plan.settings_source == "typed" + assert plan.use_dose is True + + +def test_split_settings_between_typed_and_extra_is_rejected(): + spec = ModelSpec( + model_id="lrp-rli-al-000", + kind="aligned", + title="test", + outcome_symbol="W", + model_settings=AlignedModelSettings(), + extra={"use_dose": True}, + ) + with pytest.raises(ValueError, match="cannot be split"): + resolve_aligned_run_plan(spec) + + +# --- registered-specification coverage (acceptance criterion) ----------------- + + +def test_every_registered_aligned_model_resolves_with_metadata(): + """Every registered onset-aligned model resolves to a validated plan that records + the design, estimand, causal status, analysis population and missing-data + assumption (#394 pillar 4).""" + specs = _aligned_specs() + assert len(specs) >= 9, f"expected the full aligned suite, found {len(specs)}" + for spec in specs: + plan = resolve_aligned_run_plan(spec) + assert isinstance(plan, AlignedRunPlan) + recorded = plan.as_dict() + for field in _META_FIELDS: + assert isinstance(recorded[field], str) and recorded[field], ( + f"{spec.model_id}: {field} not recorded" + ) + assert plan.prepare_kwargs()["outcomes"] == (spec.outcome_symbol,)