Add AF and AMN estimators with a shared control-function correction

[hmgaudecker]

Summary

skillmodels now hosts three estimators behind one ModelSpec and one parameter index, reorganised into per-estimator subpackages over a shared core:

• chs — the existing Cunha–Heckman–Schennach Kalman-filter MLE.
• af — new: the Antweiler & Freyberger (2025) sequential Halton-quadrature estimator.
• amn — new: the Attanasio, Meghir & Nix (2020) three-stage mixture-of-normals estimator.

All three get a uniform turnkey surface — estimate_chs / estimate_af / estimate_amn (model_spec, data, options) -> …Result carrying inference — while get_maximization_inputs stays as the CHS power-user hatch. A single declarative CorrectionSpec drives the control-function endogeneity correction for both CHS and AMN. The user-facing API is guarded by a typed beartype perimeter, and an external correctness audit was ingested and its findings fixed (see Correctness pass).

The diff is large because it also folds in the package reorg (af/amn/chs/common), the beartype perimeter (previously stacked PR #90), and the audit pass.

AF estimator (src/skillmodels/af/)

• Core estimation: estimate_af(model_spec, data, af_options, start_params) → AFEstimationResult.
• Initial period: mixture-of-normals + measurement system via 1D / KD Halton quadrature.
• Transition periods: triple integral (state nodes × investment shocks × production shocks) with previous-period conditioning.
• Transition constraints: ProbabilityConstraint for log_ces gammas, satisfied at start values.
• Investment equation: I = β₀ + β₁θ + β₂Y + σ_I ε for endogenous factors.
• State propagation: quadrature-based moment matching between periods.
• start_params support: user-supplied starting values override heuristic defaults.
• Posterior states: get_af_posterior_states computes quadrature-based posterior means per individual / period (the unified get_individual_states dispatches on result type).
• Score-based bootstrap for standard errors (paper Sec. 4.2) — compute_af_standard_errors.

AMN estimator (src/skillmodels/amn/)

• Core estimation: estimate_amn(model_spec, data, options) → AMNEstimationResult — the three-stage Attanasio–Meghir–Nix (2020) estimator.
• Stage 1 — mixture EM over the augmented measurement vector. Two methods on AMNEstimationOptions.mixture_em_method: complete_case (sklearn GaussianMixture on listwise-complete rows; raises InsufficientCompleteCasesError under an unbalanced panel) and missing_data (masked-covariance EM that marginalises over each row's missing entries, MAR — valid with no complete cases at all).
• Stage 2 — minimum distance: solve_minimum_distance recovers the structural parameters with a jitted analytical JAX gradient (value_and_grad), so the optimisation scales to large factor-period blocks instead of finite-differencing thousands of parameters.
• Stage 3 — simulate & regress: simulate_and_regress draws a synthetic latent panel and recovers the per-period transition / investment parameters by (N)LS.
• Inference: all-stage, case-level cluster bootstrap (compute_amn_standard_errors) that re-runs every stage per replicate. Posterior states: get_amn_posterior_states.

CHS turnkey driver + AMN seeding

• estimate_chs(model_spec, data, options) → CHSEstimationResult — a one-call wrapper over get_maximization_inputs + estimagic.estimate_ml, giving CHS the same estimate_* surface with full ML inference (result.likelihood_result → .se() / .cov() / .summary()). get_maximization_inputs remains the escape hatch for driving the optimiser by hand.
• CHSEstimationOptions.start_params_strategy="amn" seeds the CHS optimiser by running AMN and translating its calendar/cf coordinates onto the CHS augmented-period/kappa index; "spearman" and "none" are also available.

Control-function endogeneity correction — CorrectionSpec

• A single CorrectionSpec, declared on the endogenous investment FactorSpec.correction, configures the investment control function (AF Sec. 3.5 / AMN eq. 7–8) in one place, read by both CHS and AMN.
• Fields: instruments (excluded observed factors entering only the first stage, identifying kappa), state_predictors, targets, and kappa_degree / kappa_terms. generate_kappa_terms(factors, max_degree) builds the cf-interaction monomial basis.
• Support: CHS honours the full polynomial kappa basis; AMN implements the linear cf term; AF has no correction and raises if one is declared. ModelSpec.with_correction / without_correction builders register / strip it.

Package layout

Flat top-level modules were reorganised into per-estimator subpackages src/skillmodels/{af,amn,chs}/ over a shared src/skillmodels/common/ (model spec, params index, constraints, transition functions, data/simulation, control function, diagnostics & plotting). Public entry points are re-exported from the top-level skillmodels package.

Optimizer

Each period's MLE runs via optimagic.minimize with the algorithm in AFEstimationOptions.optimizer_algorithm (default "fides"; "scipy_lbfgsb" for MC sweeps). A jaxopt LBFGSB backend was explored and dropped — on the translog n=500 h=10k MC sweep it had ~14% period-0 maxiter timeouts and occasional silent overflow to ±1e17, while AF optimagic had 0% non-convergence on the same data. The per-step API accepts the full set of optimagic constraint kinds (FixedConstraintWithValue, ProbabilityConstraint, EqualityConstraint).

AFEstimationResult.to_numpy()

estimate_af returns on-device JAX arrays so consecutive calls reuse the XLA compilation cache (essential for MC sweeps — without this, every sim recompiles every per-period likelihood + gradient). Callers that need host residency (pickling, plotting, sending across processes) explicitly invoke result.to_numpy(), which drops samples_per_component and clears caches as a side effect.

Beartype perimeter on user-facing API

Mirrors the pattern in pylcm PR #355: a per-exception BeartypeConf plus a beartype_init class decorator routes parameter-type violations at every documented entry point through a skillmodels-specific exception class, so callers can write narrowly-scoped except clauses against a stable hierarchy rather than catching beartype's framework exception.

Exceptions (src/skillmodels/exceptions.py)

Six TypeError subclasses of a common SkillmodelsInputError, organised by perimeter:

• ModelSpecInitializationError — FactorSpec, AnchoringSpec, ModelSpec, Normalizations, CorrectionSpec
• OptionsInitializationError — CHSEstimationOptions, AFEstimationOptions, AMNEstimationOptions
• EstimationCallError — get_maximization_inputs, estimate_chs, estimate_af, estimate_amn, get_individual_states, get_af_posterior_states, get_amn_posterior_states
• InferenceCallError — compute_af_standard_errors, compute_amn_standard_errors
• SimulationCallError — simulate_dataset, simulate_policy_effect
• DiagnosticsCallError — the diagnostics / plotting helpers (decompose_measurement_variance, plot_residual_boxplots, get_transition_plots, …)

Decorator + config (src/skillmodels/_beartype_conf.py)

• _conf(exc) — BeartypeConf with violation_param_type=exc, strategy=BeartypeStrategy.On (full O(n) container scan), is_pep484_tower=True.
• beartype_init(conf) — class decorator that wraps only __init__.

Whole-package beartype.claw activation in tests (tests/conftest.py)

beartype.claw.beartype_package("skillmodels", conf=...) turns annotation-drift on internal helpers into BeartypeCallHintParamViolation during the test run. skillmodels.chs.qr is excluded because JAX's @custom_jvp .defjvp attribute doesn't survive beartype's wrap. Activating the claw surfaced ~80 internal annotation drifts, all fixed.

Correctness pass (external audit + adversarial re-verification)

A ChatGPT-Pro correctness audit of the math ↔ code — with the skane-struct-bw / health-cognition call-sites bundled in — was ingested, and every serious finding was independently re-verified against the source before fixing (several turned out to be over-escalated; only one affected results):

• estimate_chs now enforces user FixedConstraintWithValue values before optimising. They were appended after seeding and never written into the start vector, so optimagic's plain fixed constraint pinned them at the AMN/Spearman seed rather than the requested value — silently mis-restricting any model that fixes parameters to specific values (e.g. no-feedback / identity restrictions). Equality / pairwise reconciliation is now fixed-aware.
• Missing-data EM raises on a never-observed augmented-measure column (with an explicit opt-in for the seeding path) instead of feeding arbitrary ridge moments into Stage 2; the co-observation connectivity diagnostic is documented as necessary but not sufficient for cross-covariance identification and renamed accordingly.
• AMN corrected production keeps non-instrument observed controls — only CorrectionSpec.instruments are excluded from the production design (with a regression test covering a registered transition that names a factor the simulated panel omits).
• get_transition_plots accepts numpy-array periods (the beartype Sequence[int] boundary rejected ndarrays passed by applications).
• Docs rewritten around CorrectionSpec (the old period-augmentation / is_correction decorator interface is gone); stale estimate_* how-to examples and the default-CHS-standard-error description (OPG / inverse-score, not the misspecification-robust sandwich) corrected; estimate_chs, CorrectionSpec and generate_kappa_terms documented.

Test plan

• pixi run -e tests-cpu pytest tests/ -q -k "not long_running" — all green with beartype.claw enabled (634 tests)
• pixi run ty — clean
• prek run --all-files — clean
• End-to-end: health-cognition CHS pipeline runs green on this branch (estimation → tables → figures)
• pytest -m long_running — MODEL2 AF vs CHS vs AMN comparison

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 90.80275% with 401 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.71%. Comparing base (2d56c8e) to head (6775a5e).

Files with missing lines	Patch %	Lines
src/skillmodels/af/transition_period.py	65.61%	131 Missing ⚠️
src/skillmodels/af/initial_period.py	80.84%	50 Missing ⚠️
src/skillmodels/af/inference.py	88.56%	39 Missing ⚠️
src/skillmodels/amn/start_values.py	89.90%	32 Missing ⚠️
src/skillmodels/amn/simulate_and_regress.py	92.76%	22 Missing ⚠️
src/skillmodels/amn/moments.py	83.49%	17 Missing ⚠️
src/skillmodels/common/constraints.py	94.21%	17 Missing ⚠️
src/skillmodels/amn/minimum_distance.py	96.34%	13 Missing ⚠️
src/skillmodels/amn/posterior_states.py	86.45%	13 Missing ⚠️
src/skillmodels/amn/mixture_em.py	91.33%	11 Missing ⚠️
... and 14 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #89      +/-   ##
==========================================
- Coverage   96.91%   95.71%   -1.20%     
==========================================
  Files          57      122      +65     
  Lines        4952    12579    +7627     
==========================================
+ Hits         4799    12040    +7241     
- Misses        153      539     +386     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[hmgaudecker]

Code review

Found 2 issues:

1. _extract_period_measurement_info in posterior_states.py only reads the "constant" control coefficient, ignoring all other control variables. For models with non-constant controls (e.g. MODEL2 with x1), posterior state means will be biased because the control contribution to measurement residuals is incomplete. The test test_af_get_filtered_states uses a model without controls, so this is not caught.

skillmodels/src/skillmodels/af/posterior_states.py

Lines 151 to 158 in 766ad09

	ctrl_list = [
	float(period_params.loc[loc, "value"]) # ty: ignore[invalid-argument-type]
	if (loc := ("controls", period, meas, "constant")) in period_params.index
	else 0.0
	for meas in all_measures
	]

2. Distribution propagation in estimate_transition_period uses obs_factor_values[0] (the first individual's observed factor values) when constructing the state_only_transition wrapper for moment matching. For models with individual-specific observed factors, this uses one person's values for the population-level distribution update.

skillmodels/src/skillmodels/af/transition_period.py

Lines 246 to 250 in 766ad09

	def state_only_transition(state_factors_val: Array, params: Array) -> Array:
	"""Transition wrapper that fills in mean investment + observed."""
	full = jnp.concatenate([state_factors_val, mean_inv, obs_factor_values[0]])
	return combined_transition(full, params)

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[hmgaudecker]

Update — income-conditional initial draws, translog, and time-invariant latents

Three rounds of improvements since the last review, ending at commit e5b9176.

What changed

1. Income-conditional initial draws (Schur complement) — when observed_factors is non-empty, Step 0 now models the joint (latent, observed) distribution and per-individual conditions Halton draws on observed values. Parses into marginal × conditional via Σ_{θY} Σ_{YY}⁻¹, concentrating nodes where the likelihood has mass. Matches the variance-reduction trick in Antweiler-Freyberger's MATLAB (L804-812, L1185). Back-compat preserved via fast path when n_obs_factors == 0.

2. Translog transition — added a smoke test confirming the existing getattr(transition_functions, name) dispatch just works for "translog". No core changes needed.

3. fixed_params argument — new optional DataFrame that clamps value + bounds for specified parameters, so the optimizer skips them via the free-mask. Primary use case: pin time-invariant latent factors to identity linear transitions with zero shock SDs (same convention CHS uses for augmented periods).

4. MATLAB reproduction scaffolding — loaded NLSY complete_7_9_11.xls via libreoffice CSV conversion, built a ModelSpec matching the MATLAB structure (theta dynamic, MC/MN time-invariant latents, investment endogenous, log_income observed with Schur conditioning). AF now runs end-to-end on 1403 cases across 3 periods with this setup — the full structural pipeline works on real data.

Remaining gap for full MATLAB reproduction

MATLAB's CES production is 2-dim in (theta, investment); our log_ces takes ALL state factors with a ProbabilityConstraint on gammas. Pinning cross-factor gammas (mc, mn, log_income) to 0 via fixed_params breaks the constraint selector in optimagic (selected params must remain free). To fully match the MATLAB CES, skillmodels would need either (a) an "input factors" concept per transition, or (b) a custom CES-on-two-inputs transition function. Left as a follow-up.

Validation

• All 401 unit tests pass (pixi run -e tests-cpu tests).
• pixi run ty clean.
• prek run --all-files clean.
• Three new end-to-end tests added:
  • test_af_estimate_with_translog — translog runs and recovers linear coefficient from linear DGP.
  • test_af_joint_initial_distribution_with_observed_factor — verifies initial_states includes observed factors and recovers positive skill-income cross-covariance.
  • test_af_fixed_params_pins_time_invariant_latent — verifies pinned MC-style factors keep identity transitions and near-zero shock SDs after optimization.

Files touched

src/skillmodels/af/{params,initial_period,likelihood,estimate,transition_period}.py, tests/test_af_estimate.py.

🤖 Generated with Claude Code

[hmgaudecker]

Code review

No issues found. Checked for bugs and CLAUDE.md compliance in the two standard-error commits (6fd7502 Phase 1 block-diagonal sandwich, ab87767 Phase 2 full cross-period sandwich).

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[hmgaudecker]

Code review (full, including low-confidence items)

Below is the full list of issues surfaced across five review agents on the Phase 1 + Phase 2 standard-error commits (6fd7502, ab87767), with confidence scores (0-100). Items below the usual 80-threshold are still shown for transparency.

Real potential issue (85) — shock_sds shape mismatch for models with n_shock_factors < n_state_factors

The JAX-pure propagator does + jnp.diag(shock_sds**2) directly. _parse_transition_params returns shock_sds with shape (n_shock_factors,), so the result is (n_shock, n_shock) — it cannot be added to a (n_state, n_state) covariance when they differ. The existing _update_conditional_distribution has the same pattern, so this is a pre-existing bug the mirror replicates rather than a new regression; still worth flagging since the mirror is a new call site.

skillmodels/src/skillmodels/af/inference.py

Lines 803 to 808 in ab87767

	new_cov = jnp.einsum(
	"q,qi,qj->ij", state_weights, centered, centered
	) + jnp.diag(shock_sds**2)
	new_chol = jnp.linalg.cholesky(new_cov + 1e-8 * jnp.eye(n_state))
	return new_mean, new_chol

Pre-existing sibling:

skillmodels/src/skillmodels/af/transition_period.py

Lines 730 to 734 in ab87767

	new_cov = jnp.einsum(
	"q,qi,qj->ij", state_weights, centered, centered
	) + jnp.diag(shock_sds**2)
	# Cholesky factorization of new covariance

CLAUDE.md: # type: ignore[arg-type] instead of # ty: ignore[...] (75)

AGENTS.md says: "Suppress errors with # ty: ignore[rule-name] (not # type: ignore)".

https://github.com/OpenSourceEconomics/skillmodels/blob/ab877673637a59c87520b20e27ff0a5dc1faa5b2/tests/test_af_inference.py#L315-L318

CLAUDE.md: internal dataclass uses Mapping, should be MappingProxyType (75)

The repo CLAUDE.md (Immutability Conventions) says internal dataclass dict fields use MappingProxyType, with MappingProxyType(...) wrapping at the call site. _PeriodMeta is internal (underscore-prefixed, not in __all__) but declares three Mapping[str, Any] fields and is constructed with plain dicts.

skillmodels/src/skillmodels/af/inference.py

Lines 258 to 286 in ab87767

	params_df: pd.DataFrame
	loglike_kwargs: Mapping[str, Any]
	"""Keyword arguments forwarded to ``af_per_obs_loglike_initial`` (if
	``is_initial``) or ``af_per_obs_loglike_transition`` otherwise.
	"""
	parse_kwargs: Mapping[str, Any]
	"""Keyword arguments forwarded to ``_parse_initial_params`` or
	``_parse_transition_params`` respectively. Used by the Phase 2 chain.
	"""
	n_components: int
	n_factors_joint: int
	"""Joint factor count in the initial mixture (state_latent + observed).
	Only meaningful for the initial period; zero otherwise.
	"""
	n_state: int
	"""State-factor count (``n_state_latent`` in the initial period;
	``n_state_factors`` in transition periods).
	"""
	n_endog: int
	n_shock: int
	n_observed_factors: int
	state_factor_indices_in_joint: tuple[int, ...]
	"""Integer positions within the joint factor vector at which state
	factors live (the complement is observed factors). Used to marginalise
	the joint cond-dist to its state-factor sub-block.
	"""
	propagation: Mapping[str, Any] = field(default_factory=dict)
	"""Extra JAX-pure bits for propagation of the conditional distribution
	through this period's transition. Only populated for transition

model_spec: Any / processed_model: Any with ANN401 suppressions (unscored)

ModelSpec and ProcessedModel are concrete types already in use in this file's imports (indirectly via AFEstimationResult). Using Any + # noqa: ANN401 sidesteps type-safety; TYPE_CHECKING imports would avoid circular-import concerns if that is the motivation.

skillmodels/src/skillmodels/af/inference.py

Lines 293 to 299 in ab87767

	*,
	result: AFEstimationResult,
	period_data: dict[int, dict[str, Array]],
	model_spec: Any, # noqa: ANN401
	processed_model: Any, # noqa: ANN401
	af_options: AFEstimationOptions,
	observed_factors: tuple[str, ...],

CLAUDE.md: multiple assertions per test (unscored)

AGENTS.md says "One assertion per test". Several tests pack 2-4 independent assertions, e.g.:

skillmodels/tests/test_af_inference.py

Lines 109 to 123 in ab87767

	def test_af_inference_fixed_entries_have_zero_se(
	fitted_result: tuple[AFInferenceResult, pd.DataFrame],
	) -> None:
	"""Normalization pins (e.g. loadings[m1, skill] == 1) must have SE = 0."""
	inference, all_params = fitted_result
	se = inference.standard_errors
	pinned_loading = ("loadings", 0, "m1", "skill")
	assert pinned_loading in all_params.index
	assert se.loc[pinned_loading] == 0.0
	pinned_intercept = ("controls", 0, "m1", "constant")
	assert pinned_intercept in all_params.index
	assert se.loc[pinned_intercept] == 0.0

Performance note: jax.hessian on the full flat_super bypasses n_obs_per_batch (unscored)

The n_obs_per_batch memory-control contract in likelihood.py applies only to single-direction reverse mode. jax.hessian materialises a full tape over the gradient, which can scale with O(n_params × n_obs) regardless of n_obs_per_batch. For large models this may OOM at inference time where estimation did not.

skillmodels/src/skillmodels/af/inference.py

Lines 677 to 680 in ab87767

	jac_full = jax.jacfwd(per_obs_loglike_full)(flat_values)
	hess_full = jax.hessian(neg_mean_loglike_full)(flat_values)

skillmodels/src/skillmodels/af/inference.py

Lines 937 to 940 in ab87767

	score_matrices_full.append(jax.jacfwd(_per_obs_t)(flat_super))
	hessian_blocks_full.append(jax.hessian(_neg_mean_t)(flat_super))

Latent inconsistency (25) — conditional_weights never propagated in the JAX chain

_build_prev_dist_arrays always broadcasts uniform mixture_weights. The estimation-time _prepare_transition_inputs instead honours prev_distribution.conditional_weights when it is non-None. In current AF code every estimation path sets conditional_weights=None, so this is latent/defensive only — but it is an asymmetry that will break silently if per-individual posterior weights are ever introduced.

skillmodels/src/skillmodels/af/inference.py

Lines 868 to 875 in ab87767

	meta_target = metas[target_t]
	n_obs = int(meta_target.loglike_kwargs["measurements"].shape[0])
	n_components = metas[0].n_components
	cond_weights = jnp.broadcast_to(mixture_weights[None, :], (n_obs, n_components))
	return {
	"cond_weights": cond_weights,
	"means": state_means,
	"chol_covs": state_chols,

Flagged but confirmed false positives (0):

• prior_mean = prev_means[0] used for all components' mean-investment — faithfully mirrors transition_period._compute_mean_investment.
• mixture_weights carried unchanged across propagation — matches _update_conditional_distribution's intentional design (docstring: "compute the new mean and covariance").
• prev_loading_mask not overridden in the Phase 2 kwargs — structural (boolean mask derived from model spec, not from estimated values), so correct.

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}