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.
• Now also bundles: the beartype perimeter on the user-facing API (previously stacked PR Beartype perimeter on user-facing API #90), whole-package beartype.claw activation in tests, and JAX 0.10 / cuda13 workarounds. See sections below.

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_filtered_states(model_spec, data, params, af_result=result) computes quadrature-based posterior means per individual / period.
• Score-based bootstrap for standard errors (paper Sec. 4.2) — compute_af_standard_errors.

Optimizer backends

AFEstimationOptions.optimizer_backend chooses how each period's MLE is solved:

• "optimagic" (default fallback): om.minimize(algorithm="scipy_lbfgsb", ...). Required when user equality / probability constraints exist (jaxopt can't fold those).
• "jaxopt": jaxopt.LBFGSB, run directly on device. Avoids the host↔device transfer optimagic incurs once per likelihood call.
• "auto": pick "jaxopt" iff a JAX GPU is visible and the model is jaxopt-compatible (no log_ces* transitions, no user-supplied constraints); otherwise "optimagic".

The jaxopt wrapper now matches scipy_lbfgsb's gtol OR ftol stopping rule (drives the solver through an explicit init_state + update loop and checks ‖projected_grad‖∞ < gtol_abs OR (f_k − f_{k+1}) / max(|f_k|, |f_{k+1}|, 1) < ftol_rel after each step). The same optimizer_options keys (convergence_gtol_abs, convergence_ftol_rel, stopping_maxiter) work for either backend, so per-sim Monte Carlo timing benchmarks are byte-identical apart from internal LBFGSB mechanics.

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
• OptionsInitializationError — CHSEstimationOptions, AFEstimationOptions, AMNEstimationOptions
• EstimationCallError — get_maximization_inputs, get_filtered_states, estimate_af, estimate_amn, get_af_posterior_states, get_amn_posterior_states
• InferenceCallError — compute_af_standard_errors, compute_amn_standard_errors
• SimulationCallError — simulate_dataset, simulate_policy_effect
• DiagnosticsCallError — decompose_measurement_variance, summarize_measurement_reliability, plot_residual_boxplots, plot_likelihood_contributions, create_state_ranges, plot_correlation_heatmap, get_measurements_corr, get_quasi_scores_corr, get_scores_corr, univariate_densities, bivariate_density_contours, bivariate_density_surfaces, combine_distribution_plots, get_transition_plots, combine_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__. Avoids surfacing non-public annotation drift on instance methods that has nothing to do with parameter validation at construction time.

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 decorator's secondary .defjvp attribute doesn't survive beartype's wrap. Activating the claw surfaced ~80 internal annotation drifts (Array / np.ndarray / int / Mapping / dict-vs-MappingProxyType / TYPE_CHECKING-only forward refs), all fixed in the commits that follow.

Side effects

• _check_measurements's type-shape arm in common/check_model.py is now dead code: the tuple[tuple[str, ...], ...] annotation on FactorSpec.measurements makes beartype reject every malformed measurement structure at construction time. The function is kept for non-type issues the container scan can't see; the two tests that previously asserted soft errors are rewritten to assert ModelSpecInitializationError at FactorSpec(...) time.
• tests/test_af_jaxopt_backend.py::test_optimizer_backend_rejects_unknown_value now asserts OptionsInitializationError from beartype's Literal check (which fires before AFEstimationOptions.__post_init__'s manual ValueError).
• tests/test_amn_plot_harmonization.py::test_get_filtered_states_rejects_both_af_and_amn_results now asserts EstimationCallError.
• chs/filtered_states.py imports AFEstimationResult / AMNEstimationResult at runtime rather than under TYPE_CHECKING so beartype can resolve the annotation; ruff's TC003 autofix had been silently unforwarding the string forward refs.
• FixedConstraintWithValue moved to its own module common/fixed_constraint.py to break a circular import (constraints.py imports transition_functions.py; the leaf type now lives where both modules can pull it without a cycle).

JAX 0.10 / cuda13 workarounds

• XLA_FLAGS=--xla_disable_hlo_passes=permutation_sort_simplifier is set at package import to bypass a JAX 0.10 XLA pass that mis-lowers the argsort inside jaxopt.LBFGSB.update (emits an s32 reduction accumulator into an s64 scatter operand). No-op on JAX < 0.10.
• JAX_ENABLE_X64=1 set at package import time so transitive import jaxopt sees x64 as the default integer width.

Test plan

• pixi run -e tests-cpu pytest tests/ -q -k "not long_running" — all green with beartype.claw enabled
• pixi run ty — clean
• prek run --all-files — clean
• 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 93.48227% with 353 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.26%. Comparing base (2d56c8e) to head (2d5525a).

Files with missing lines	Patch %	Lines
src/skillmodels/af/inference.py	70.83%	91 Missing ⚠️
src/skillmodels/af/transition_period.py	84.11%	61 Missing ⚠️
src/skillmodels/amn/start_values.py	87.97%	35 Missing ⚠️
src/skillmodels/amn/moments.py	83.49%	17 Missing ⚠️
src/skillmodels/amn/simulate_and_regress.py	90.95%	17 Missing ⚠️
src/skillmodels/af/estimate.py	91.66%	15 Missing ⚠️
src/skillmodels/amn/posterior_states.py	86.45%	13 Missing ⚠️
src/skillmodels/common/constraints.py	83.33%	13 Missing ⚠️
src/skillmodels/af/initial_period.py	95.45%	12 Missing ⚠️
src/skillmodels/common/transition_functions.py	61.53%	10 Missing ⚠️
... and 19 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #89      +/-   ##
==========================================
- Coverage   96.91%   95.26%   -1.65%     
==========================================
  Files          57      105      +48     
  Lines        4952    10217    +5265     
==========================================
+ Hits         4799     9733    +4934     
- Misses        153      484     +331     

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