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)
• 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 91.12330% with 437 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.00%. Comparing base (2d56c8e) to head (852b095).

Files with missing lines	Patch %	Lines
src/skillmodels/af/transition_period.py	65.17%	132 Missing ⚠️
src/skillmodels/af/inference.py	68.93%	96 Missing ⚠️
src/skillmodels/af/initial_period.py	79.92%	51 Missing ⚠️
src/skillmodels/amn/start_values.py	87.28%	37 Missing ⚠️
src/skillmodels/amn/moments.py	82.52%	18 Missing ⚠️
src/skillmodels/amn/simulate_and_regress.py	90.95%	17 Missing ⚠️
src/skillmodels/common/transition_functions.py	44.00%	14 Missing ⚠️
src/skillmodels/amn/posterior_states.py	86.02%	13 Missing ⚠️
src/skillmodels/amn/mixture_em.py	90.74%	10 Missing ⚠️
src/skillmodels/af/validate.py	72.72%	9 Missing ⚠️
... and 11 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #89      +/-   ##
==========================================
- Coverage   96.91%   94.00%   -2.91%     
==========================================
  Files          57       94      +37     
  Lines        4952     9707    +4755     
==========================================
+ Hits         4799     9125    +4326     
- Misses        153      582     +429     

☔ 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 👎.}

[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 👎.}

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB