Add Antweiler-Freyberger (2025) iterative quadrature estimator

[hmgaudecker]

Summary

• New af/ subpackage implementing the Antweiler & Freyberger (2025) estimator as an alternative to the CHS Kalman filter
• Same ModelSpec interface — users switch estimator by calling estimate_af() instead of get_maximization_inputs() + om.maximize()
• Period-by-period MLE with Halton quadrature, JAX AD for gradients, LogSumExp for numerical stability
• Supports arbitrary factor counts, log_ces/linear/translog transitions, endogenous factors via explicit investment equation
• AF and CHS agree closely on both measurement and transition parameters (tested on synthetic data and MODEL2)
• Common get_filtered_states() interface: pass af_result= for AF posterior states, omit for CHS filtered states

What's done

• Core estimation: estimate_af(model_spec, data, af_options, start_params) → AFEstimationResult
• Initial period: mixture-of-normals + measurement system via 1D/KD 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_filtered_states(model_spec, data, params, af_result=result) computes quadrature-based posterior means per individual/period
• 10 tests (9 regular + 1 long_running MODEL2 comparison)

Still to do

• Score-based bootstrap for standard errors (AF paper Section 4.2) — postponed

Test plan

• pixi run -e tests-cpu tests — 399 passed, 1 deselected (long_running)
• pixi run ty — all checks passed
• prek run --all-files — all passed
• pytest -m long_running — MODEL2 AF vs CHS comparison (both estimators optimised from same naive start values)

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 96.79715% with 36 lines in your changes missing coverage. Please review.
✅ Project coverage is 96.95%. Comparing base (2d56c8e) to head (765d1c6).

Files with missing lines	Patch %	Lines
src/skillmodels/af/params.py	92.68%	9 Missing ⚠️
src/skillmodels/af/transition_period.py	95.97%	8 Missing ⚠️
src/skillmodels/af/validate.py	66.66%	7 Missing ⚠️
src/skillmodels/af/posterior_states.py	95.14%	5 Missing ⚠️
src/skillmodels/af/estimate.py	94.59%	4 Missing ⚠️
src/skillmodels/af/initial_period.py	98.42%	2 Missing ⚠️
src/skillmodels/af/halton.py	94.11%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #89      +/-   ##
==========================================
+ Coverage   96.91%   96.95%   +0.04%     
==========================================
  Files          57       68      +11     
  Lines        4952     6075    +1123     
==========================================
+ Hits         4799     5890    +1091     
- Misses        153      185      +32     

☔ View full report in Codecov by Sentry.
📢 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 👎.}