From 99821f83178ccc1f5081da0580873eaa00cbb3d5 Mon Sep 17 00:00:00 2001 From: Hans-Martin Gaudecker Date: Thu, 8 Oct 2026 21:26:53 +0200 Subject: [PATCH 01/18] Snapshot gpu-01 working tree used for health-cognition estimation runs Commits the uncommitted state of the gpu-01 checkout that produced the health-cognition July 2026 results (translog warm-start, 2-mixture linear): QR-JVP diagonal floor for rank-deficient covariances (chs/qr.py), the criterion sanitizer (chs/maximization_inputs.py), AF posterior-state / transition-period and constraints changes, and accompanying docs and tests. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/skillmodels/chs/maximization_inputs.py | 26 ++++++++++++++++++- src/skillmodels/chs/qr.py | 30 +++++++++++++++++++++- 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index 566bb052..70784834 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -193,7 +193,7 @@ def loglike_and_gradient( start += step stop += step grad = _to_numpy(_grad) - return crit, grad + return _sanitize_criterion_for_optimizer(crit, grad) def debug_loglike(params: pd.DataFrame) -> dict[str, Any]: params_vec = partialed_get_jnp_params_vec(params) @@ -449,6 +449,30 @@ def _partial_some_log_likelihood( ) +def _sanitize_criterion_for_optimizer( + crit: float, + grad: NDArray[np.floating], +) -> tuple[float, NDArray[np.floating]]: + """Replace a non-finite criterion/gradient with a finite, unattractive one. + + A descent optimiser can probe numerically degenerate parameter regions -- e.g. + an explosive nonlinear (translog) transition whose latent-state sensitivities + overflow -- where the log-likelihood value stays finite but its gradient turns + NaN/inf. Handed that verbatim, `optimagic` aborts the entire estimation. Since + such a point is always worse than the incumbent, return it with non-finite + gradient entries zeroed (and a non-finite value floored to a very unattractive + finite number) so a line-search / trust-region step rejects it and the + optimiser backtracks instead of crashing. At well-behaved points every value is + finite and this is a no-op. + """ + grad_is_finite = np.isfinite(grad) + if np.isfinite(crit) and bool(grad_is_finite.all()): + return crit, grad + safe_grad = np.where(grad_is_finite, grad, 0.0) + safe_crit = crit if np.isfinite(crit) else -1e300 + return safe_crit, safe_grad + + def _to_numpy(obj: Any) -> Any: # noqa: ANN401 if isinstance(obj, dict): res = {} diff --git a/src/skillmodels/chs/qr.py b/src/skillmodels/chs/qr.py index d42862f7..d9ac11ee 100644 --- a/src/skillmodels/chs/qr.py +++ b/src/skillmodels/chs/qr.py @@ -53,6 +53,28 @@ def _tril(m: Array, k: int = 0) -> Array: return jax.lax.select(jax.lax.broadcast(mask, m.shape[:-2]), m, jnp.zeros_like(m)) +# Floor for the magnitude of R's diagonal in the tangent solve below. The QR +# derivative solves against R, so a zero on R's diagonal -- a rank-deficient +# input, which the square-root UKF produces when a predicted covariance collapses +# as the optimiser inflates a nonlinear transition -- would send the tangent to +# +/-inf and poison the entire gradient with NaN, aborting the optimisation. +_QR_JVP_DIAGONAL_FLOOR = 1e-6 + + +def _floor_diagonal_magnitude(r: Array, floor: float) -> Array: + """Return `r` with each diagonal entry pushed to magnitude at least `floor`. + + Off-diagonal entries and the sign of each diagonal entry are preserved; a + diagonal entry already larger than `floor` in magnitude is untouched, so a + well-conditioned R (and its gradient) is unchanged. + """ + n = r.shape[-1] + diag = jnp.diagonal(r, axis1=-2, axis2=-1) + sign = jnp.where(diag < 0.0, -1.0, 1.0) + safe = jnp.where(jnp.abs(diag) < floor, sign * floor, diag) + return r + (safe - diag)[..., :, None] * jnp.eye(n, dtype=r.dtype) + + @qr_gpu.defjvp def qr_jvp_rule( primals: tuple[Array], @@ -63,7 +85,13 @@ def qr_jvp_rule( (x,) = primals (dx,) = tangents q, r = qr_gpu(x) - dx_rinv = jax.lax.linalg.triangular_solve(r, dx) # Right side solve by default + # Solve against a diagonal-floored R so a (near-)singular decomposition yields + # a large but finite tangent instead of NaN. The primal `q, r` and every other + # term below use the unmodified `r`. + r_solve = _floor_diagonal_magnitude(r, _QR_JVP_DIAGONAL_FLOOR) + dx_rinv = jax.lax.linalg.triangular_solve( + r_solve, dx + ) # Right side solve by default qt_dx_rinv = _h(q) @ dx_rinv qt_dx_rinv_lower = _tril(m=qt_dx_rinv, k=-1) do = qt_dx_rinv_lower - _h(qt_dx_rinv_lower) # This is skew-symmetric From cadcaaa20fc0f23eda41eab446687cdc5d06cd88 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 9 Oct 2026 08:40:58 +0200 Subject: [PATCH 02/18] Drop the CHS criterion sanitizer Masking a non-finite loglikelihood as -1e300 with zeroed gradient entries produced fake convergence (health-cognition 2-mixture translog, one iteration at -1e300) and hides the degenerate steps that shrink the trust region. Restore the release behaviour: non-finite values reach optimagic and abort. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/skillmodels/chs/maximization_inputs.py | 26 +--------------------- 1 file changed, 1 insertion(+), 25 deletions(-) diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index 70784834..566bb052 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -193,7 +193,7 @@ def loglike_and_gradient( start += step stop += step grad = _to_numpy(_grad) - return _sanitize_criterion_for_optimizer(crit, grad) + return crit, grad def debug_loglike(params: pd.DataFrame) -> dict[str, Any]: params_vec = partialed_get_jnp_params_vec(params) @@ -449,30 +449,6 @@ def _partial_some_log_likelihood( ) -def _sanitize_criterion_for_optimizer( - crit: float, - grad: NDArray[np.floating], -) -> tuple[float, NDArray[np.floating]]: - """Replace a non-finite criterion/gradient with a finite, unattractive one. - - A descent optimiser can probe numerically degenerate parameter regions -- e.g. - an explosive nonlinear (translog) transition whose latent-state sensitivities - overflow -- where the log-likelihood value stays finite but its gradient turns - NaN/inf. Handed that verbatim, `optimagic` aborts the entire estimation. Since - such a point is always worse than the incumbent, return it with non-finite - gradient entries zeroed (and a non-finite value floored to a very unattractive - finite number) so a line-search / trust-region step rejects it and the - optimiser backtracks instead of crashing. At well-behaved points every value is - finite and this is a no-op. - """ - grad_is_finite = np.isfinite(grad) - if np.isfinite(crit) and bool(grad_is_finite.all()): - return crit, grad - safe_grad = np.where(grad_is_finite, grad, 0.0) - safe_crit = crit if np.isfinite(crit) else -1e300 - return safe_crit, safe_grad - - def _to_numpy(obj: Any) -> Any: # noqa: ANN401 if isinstance(obj, dict): res = {} From 6c3188c8d1d1cc3922297c6bdd47d16f629939a0 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sat, 10 Oct 2026 22:12:45 +0200 Subject: [PATCH 03/18] Reset underflowed mixture components to the dominant one in the CHS update A component whose weight underflows the smallest normal float contributes nothing to the likelihood, but its state was still propagated. Under a translog transition it exploded to ~1e281 and then inf over later periods, so the gradient through its zero weight became 0 * inf = NaN for almost every parameter (health-cognition 2-mixture translog). Overwrite such components' state and covariance with the dominant component's after each update: the likelihood is unchanged and the gradient exact up to the underflowed weight. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/kalman_filters.py | 44 ++++++++++++++++++++ src/skillmodels/chs/kalman_filters_debug.py | 9 ++++ tests/test_kalman_filters.py | 46 +++++++++++++++++++++ 3 files changed, 99 insertions(+) diff --git a/src/skillmodels/chs/kalman_filters.py b/src/skillmodels/chs/kalman_filters.py index 72360cd9..bcb67d88 100644 --- a/src/skillmodels/chs/kalman_filters.py +++ b/src/skillmodels/chs/kalman_filters.py @@ -131,6 +131,13 @@ def kalman_update( log_mixture_weights, ) + if n_mixtures >= 2: + new_states, new_upper_chols = reset_negligible_mixture_components( + states=new_states, + upper_chols=new_upper_chols, + log_mixture_weights=new_log_mixture_weights, + ) + return ( new_states, new_upper_chols, @@ -139,6 +146,43 @@ def kalman_update( ) +def reset_negligible_mixture_components( + states: Array, + upper_chols: Array, + log_mixture_weights: Array, +) -> tuple[Array, Array]: + """Give components whose weight underflows the dominant component's state. + + A component whose weight is below the smallest normal float contributes exactly + nothing to the likelihood, but its state is still propagated. Under a nonlinear + (e.g. translog) transition it can then explode to inf over later periods, and + the gradient through its zero weight becomes 0 * inf = NaN for every parameter. + Overwriting its state and covariance with the dominant component's keeps all + values finite; the likelihood is unchanged and the gradient is exact up to the + underflowed weight. The weights themselves are left untouched. + + Args: + states: Array of shape (n_obs, n_mixtures, n_states). + upper_chols: Array of shape (n_obs, n_mixtures, n_states, n_states). + log_mixture_weights: Array of shape (n_obs, n_mixtures). + + Returns: + states and upper_chols with negligible components replaced. + + """ + log_weights = jax.lax.stop_gradient(log_mixture_weights) + negligible = log_weights < jnp.log(jnp.finfo(log_weights.dtype).tiny) + dominant = jnp.argmax(log_weights, axis=1) + dominant_states = jnp.take_along_axis(states, dominant[:, None, None], axis=1) + dominant_chols = jnp.take_along_axis( + upper_chols, dominant[:, None, None, None], axis=1 + ) + return ( + jnp.where(negligible[..., None], dominant_states, states), + jnp.where(negligible[..., None, None], dominant_chols, upper_chols), + ) + + def calculate_sigma_scaling_factor_and_weights( n_states: int, kappa: float = 2, diff --git a/src/skillmodels/chs/kalman_filters_debug.py b/src/skillmodels/chs/kalman_filters_debug.py index c7dd20a8..27d49665 100644 --- a/src/skillmodels/chs/kalman_filters_debug.py +++ b/src/skillmodels/chs/kalman_filters_debug.py @@ -6,6 +6,8 @@ import jax.numpy as jnp from jax import Array +from skillmodels.chs.kalman_filters import reset_negligible_mixture_components + array_qr_jax = jax.vmap(jax.vmap(jnp.linalg.qr)) @@ -118,6 +120,13 @@ def kalman_update( log_mixture_weights, ) + if n_mixtures >= 2: + new_states, new_upper_chols = reset_negligible_mixture_components( + states=new_states, + upper_chols=new_upper_chols, + log_mixture_weights=new_log_mixture_weights, + ) + debug_info = {} residuals = jnp.where(not_missing.reshape(n_obs, 1), _residuals, jnp.nan) debug_info["residuals"] = residuals diff --git a/tests/test_kalman_filters.py b/tests/test_kalman_filters.py index 73b04040..2774b8e4 100644 --- a/tests/test_kalman_filters.py +++ b/tests/test_kalman_filters.py @@ -570,3 +570,49 @@ def test_sigma_points_multiple_mixtures() -> None: (n_sigma, n_observed), ) aaae(observed_slice, expected) + + +@pytest.mark.parametrize("update_func", UPDATE_FUNCS) +def test_kalman_update_resets_negligible_mixture_component(update_func) -> None: + """A component with underflowing weight takes the dominant component's state. + + Otherwise a nonlinear transition can blow its (irrelevant) state up to inf over + later periods, and the gradient through its zero weight becomes 0 * inf = NaN. + """ + n_states = 2 + states = jnp.array([[[1e100, -1e100], [0.3, -0.2]]]) + chols = jnp.broadcast_to(jnp.eye(n_states), (1, 2, n_states, n_states)) + log_weights = jnp.array([[-1e6, 0.0]]) + kwargs = { + "loadings": jnp.array([1.0, 0.5]), + "control_params": jnp.ones(1), + "meas_sd": 1.0, + "controls": jnp.ones((1, 1)), + } + + def two_updates(s: jax.Array) -> jax.Array: + s1, c1, w1, ll1 = update_func( + states=s, + upper_chols=chols, + measurements=jnp.array([0.4]), + log_mixture_weights=log_weights, + **kwargs, + )[:4] + _, _, _, ll2 = update_func( + states=s1**4, + upper_chols=c1, + measurements=jnp.array([0.1]), + log_mixture_weights=w1, + **kwargs, + )[:4] + return (ll1 + ll2).sum() + + new_states = update_func( + states=states, + upper_chols=chols, + measurements=jnp.array([0.4]), + log_mixture_weights=log_weights, + **kwargs, + )[0] + aaae(new_states[0, 0], new_states[0, 1]) + assert jnp.isfinite(jax.grad(two_updates)(states)).all() From 06f8ab927f5cd0d0a72c0e66f369afc1f9962cee Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sat, 10 Oct 2026 22:40:35 +0200 Subject: [PATCH 04/18] Raise on a non-finite CHS criterion; forward error_handling to optimagic loglike_and_gradient now raises when the value or gradient is non-finite instead of handing NaN to the optimiser. With the new CHSEstimationOptions.error_handling (default "raise") set to "continue", optimagic replaces such a point by a penalty worse than the start, so a trust-region step into an explosive region is rejected and the radius shrinks -- what the removed sanitizer did for the value, without its zeroed gradients being accepted at finite points. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/estimate.py | 1 + src/skillmodels/chs/maximization_inputs.py | 8 ++++++++ src/skillmodels/chs/options.py | 8 ++++++++ tests/test_chs_estimate.py | 24 ++++++++++++++++++++++ tests/test_likelihood_regression.py | 12 +++++++++++ 5 files changed, 53 insertions(+) diff --git a/src/skillmodels/chs/estimate.py b/src/skillmodels/chs/estimate.py index a7f22b12..c3a22d6b 100644 --- a/src/skillmodels/chs/estimate.py +++ b/src/skillmodels/chs/estimate.py @@ -123,6 +123,7 @@ def estimate_chs( "algorithm": options.optimizer_algorithm, "algo_options": to_plain_dict(options.optimizer_options) or None, "fun_and_jac": max_inputs["loglike_and_gradient"], + "error_handling": options.error_handling, } # Default to OPG/jacobian-based inference (`hessian=False`): the numerical diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index 566bb052..2be29155 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -193,6 +193,14 @@ def loglike_and_gradient( start += step stop += step grad = _to_numpy(_grad) + if not (np.isfinite(crit) and np.isfinite(grad).all()): + # Raise rather than return, so `error_handling="continue"` can replace + # the point by a penalty that rejects the step. + msg = ( + f"The log-likelihood ({crit}) or its gradient " + f"({int((~np.isfinite(grad)).sum())} entries) is non-finite." + ) + raise FloatingPointError(msg) return crit, grad def debug_loglike(params: pd.DataFrame) -> dict[str, Any]: diff --git a/src/skillmodels/chs/options.py b/src/skillmodels/chs/options.py index 48ad358b..0b3bbb24 100644 --- a/src/skillmodels/chs/options.py +++ b/src/skillmodels/chs/options.py @@ -59,6 +59,14 @@ class CHSEstimationOptions: Forwarded by `estimate_chs` to `estimate_ml` (e.g. convergence tolerances, trust-region settings for `fides`).""" + error_handling: Literal["raise", "continue"] = "raise" + """What `optimagic` does when the criterion raises (`optimize_options`). + + The criterion raises whenever the log-likelihood or its gradient is + non-finite. `"raise"` aborts the estimation. `"continue"` replaces that + point by a penalty worse than the start value, so a trust-region step into + a numerically explosive region is rejected and the radius shrinks.""" + estimate_ml_options: Mapping[str, Any] = field( default_factory=lambda: MappingProxyType({}) ) diff --git a/tests/test_chs_estimate.py b/tests/test_chs_estimate.py index 1f2946be..c2d71230 100644 --- a/tests/test_chs_estimate.py +++ b/tests/test_chs_estimate.py @@ -125,3 +125,27 @@ def test_estimate_chs_enforces_user_fixed_constraint_value( ) assert result.params.loc[loc, "value"] == pytest.approx(target) + + +def test_estimate_chs_forwards_error_handling( + model2_data, anchoring_start_params, monkeypatch +): + """`error_handling` reaches optimagic so failed evaluations can be penalised.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs["optimize_options"]) + raise StopIteration + + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions( + start_params_strategy="none", error_handling="continue" + ), + start_params=anchoring_start_params, + require_initial_anchors=False, + ) + assert captured["error_handling"] == "continue" diff --git a/tests/test_likelihood_regression.py b/tests/test_likelihood_regression.py index cc540be4..791d3c02 100644 --- a/tests/test_likelihood_regression.py +++ b/tests/test_likelihood_regression.py @@ -115,6 +115,18 @@ def test_splitting_does_not_change_gradient(model2, model2_data) -> None: aaae(gradient, gradient_split) +def test_loglike_and_gradient_raises_on_non_finite_value(model2, model2_data) -> None: + """A non-finite value must raise, so optimagic can penalise and reject the step.""" + inputs = get_maximization_inputs( + model2, model2_data, chs_options=MODEL2_CHS_OPTIONS + ) + params = inputs["params_template"] + params["value"] = 0.1 + params.loc[params.index[0], "value"] = np.nan + with pytest.raises(FloatingPointError, match="non-finite"): + inputs["loglike_and_gradient"](params) + + @pytest.mark.parametrize( ("model_name", "fun_key"), product(MODEL_NAMES, ["loglikeobs"]) ) From 409350a8c1f3f9e6802b5e841c1292401413ad7a Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sun, 11 Oct 2026 22:50:29 +0200 Subject: [PATCH 05/18] Let estimate_chs tighten lower bounds on chosen parameters Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/estimate.py | 16 ++++++++++++++++ tests/test_chs_estimate.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/src/skillmodels/chs/estimate.py b/src/skillmodels/chs/estimate.py index c3a22d6b..c14f63ca 100644 --- a/src/skillmodels/chs/estimate.py +++ b/src/skillmodels/chs/estimate.py @@ -11,6 +11,7 @@ hatch for callers who want to drive the optimiser themselves. """ +import numpy as np import optimagic as om import pandas as pd from beartype import beartype @@ -40,6 +41,7 @@ def estimate_chs( start_params: pd.DataFrame | None = None, fixed_params: pd.DataFrame | None = None, constraints: list[om.constraints.Constraint] | None = None, + lower_bounds: pd.Series | None = None, *, require_initial_anchors: bool = True, ) -> CHSEstimationResult: @@ -64,6 +66,10 @@ def estimate_chs( fixed. constraints: Optional extra `optimagic` constraints, appended to the model-implied constraints from `get_maximization_inputs`. + lower_bounds: Optional lower bounds, indexed like the `params_template`, + that tighten the template's bounds (e.g. a floor on a measurement + SD that would otherwise collapse). Start values below them are + raised onto them. require_initial_anchors: When True (the default), run the estimator-agnostic period-0 anchor check (`fail_if_initial_state_unanchored`) and raise `ValueError` if the @@ -101,6 +107,16 @@ def estimate_chs( overlay = start_params["value"].reindex(start.index) start.loc[overlay.notna(), "value"] = overlay[overlay.notna()] + if lower_bounds is not None: + tighter = lower_bounds.reindex(start.index) + locs = tighter.dropna().index + start.loc[locs, "lower_bound"] = np.maximum( + start.loc[locs, "lower_bound"], tighter[locs] + ) + start.loc[locs, "value"] = np.maximum( + start.loc[locs, "value"], start.loc[locs, "lower_bound"] + ) + all_constraints = [*max_inputs["constraints"], *(constraints or [])] # Write every `FixedConstraintWithValue`'s target into the start vector. diff --git a/tests/test_chs_estimate.py b/tests/test_chs_estimate.py index c2d71230..f83f221d 100644 --- a/tests/test_chs_estimate.py +++ b/tests/test_chs_estimate.py @@ -149,3 +149,34 @@ def fake_estimate_ml(**kwargs): require_initial_anchors=False, ) assert captured["error_handling"] == "continue" + + +def test_estimate_chs_applies_lower_bounds_and_clips_start( + model2_data, anchoring_start_params, monkeypatch +): + """`lower_bounds` tightens the template bounds; the start moves onto them.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + raise StopIteration + + loc = anchoring_start_params.query("category == 'meas_sds'").index[:1] + start = anchoring_start_params.copy() + start.loc[loc, "value"] = 0.05 + lower_bounds = pd.Series(0.1, index=loc) + + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions(start_params_strategy="none"), + start_params=start, + lower_bounds=lower_bounds, + require_initial_anchors=False, + ) + assert captured["bounds"].lower.loc[loc].tolist() == [0.1] + assert captured["params"].loc[loc, "value"].tolist() == [0.1] + other = captured["bounds"].lower.drop(loc) + assert other[np.isfinite(other)].max() < 0.1 From af53b2c52edfe7f5503176fb483e096134047918 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 12 Oct 2026 07:51:02 +0200 Subject: [PATCH 06/18] Reject probit/Tobit measurements in CHS instead of treating them as Gaussian Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- .../estimator_prerequisites.md | 2 +- src/skillmodels/chs/maximization_inputs.py | 23 ++++++ src/skillmodels/common/measurement_models.py | 7 +- .../test_chs_limited_measurement_rejection.py | 71 +++++++++++++++++++ 4 files changed, 99 insertions(+), 4 deletions(-) create mode 100644 tests/test_chs_limited_measurement_rejection.py diff --git a/docs/reference_guides/estimator_prerequisites.md b/docs/reference_guides/estimator_prerequisites.md index 7a91d36f..a119eab7 100644 --- a/docs/reference_guides/estimator_prerequisites.md +++ b/docs/reference_guides/estimator_prerequisites.md @@ -17,7 +17,7 @@ three estimators. | --------------------- | ------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | Estimator | Joint maximum likelihood with square-root Gaussian-component filtering. | Sequential period-by-period likelihood with Halton integration. | Three stages: Gaussian-mixture EM → minimum distance → simulate-and-regress. | | Latent distribution | Finite mixture of Gaussian initial states via `ModelSpec.n_mixtures`; Gaussian component filtering thereafter. | Initial finite mixture plus sequentially carried distributions; Halton nodes approximate the integrals. | Stage-1 Gaussian mixture over the augmented measurement vector; `n_mixtures` sets the component count. | -| Measurement families | **Gaussian only** (standard Kalman update). Probit/Tobit measurements are not consumed by the CHS path. | **Initial period only**: probit/Tobit measurement families are honoured at the period-0 measurement system; transition periods ($t \geq 1$) fall back to an all-Gaussian measurement kernel. | **Gaussian only**: `estimate_amn` raises `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | +| Measurement families | **Gaussian only** (standard Kalman update): `get_maximization_inputs` / `estimate_chs` raise `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | **Initial period only**: probit/Tobit measurement families are honoured at the period-0 measurement system; transition periods ($t \geq 1$) fall back to an all-Gaussian measurement kernel. | **Gaussian only**: `estimate_amn` raises `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | | Missing data | Gaussian measurement updates skip individually missing measurements. | Measurement masks skip missing rows in the per-step likelihood contributions. | `mixture_em_method="complete_case"` (default) or `"missing_data"` (marginalises over missing entries under MAR); never-observed columns require `allow_never_observed_measurements=True`. | | Endogenous investment | Reconstructed endogenous factors plus a full `CorrectionSpec` control-function basis. | Reconstructed endogenous investment with independent shocks (source/destination calendar adapter); rejects `CorrectionSpec` / nonzero `kappa`. | Endogenous investment in Stage 3; correction support is a linearised control-function term, narrower than CHS. | | Corrections | Full `CorrectionSpec` / `kappa` polynomial basis in the processed transition DAG. | Not implemented — the validator raises if a `CorrectionSpec` is attached. | Linear `cf` term only; a higher-order `kappa_terms` request raises. | diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index 2be29155..9063f55f 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -40,6 +40,7 @@ get_constraints, project_to_probability_constraints, ) +from skillmodels.common.measurement_models import GaussianMeasurement from skillmodels.common.model_spec import ModelSpec from skillmodels.common.params_index import get_params_index from skillmodels.common.parse_params import create_parsing_info @@ -101,6 +102,7 @@ def get_maximization_inputs( # noqa: C901, PLR0915 endogenous factors, we double up the number of periods in order to add """ + _fail_if_non_gaussian_measurements(model_spec) chs_options = chs_options or CHSEstimationOptions() processed_model = process_model(model_spec) p_index = get_params_index( @@ -270,6 +272,27 @@ def debug_loglike(params: pd.DataFrame) -> dict[str, Any]: } +def _fail_if_non_gaussian_measurements(model_spec: ModelSpec) -> None: + """Reject probit/Tobit measures: the CHS update is a Gaussian Kalman update. + + Such a measure would otherwise be silently estimated as a continuous Gaussian one. + """ + non_gaussian = sorted( + name + for name, model in model_spec.measurement_models.items() + if not isinstance(model, GaussianMeasurement) + ) + if non_gaussian: + msg = ( + f"CHS supports only Gaussian measurements, but {non_gaussian} are " + "non-Gaussian (probit/Tobit). The CHS likelihood runs a Gaussian Kalman " + "update for every measurement. Declare them Gaussian (a linear " + "probability / linear working model) or use estimate_af, which honours " + "probit/Tobit families in the initial period." + ) + raise NotImplementedError(msg) + + def _estimate_amn_for_chs_seeding( model_spec: ModelSpec, data: pd.DataFrame, diff --git a/src/skillmodels/common/measurement_models.py b/src/skillmodels/common/measurement_models.py index 9a500c95..b955d900 100644 --- a/src/skillmodels/common/measurement_models.py +++ b/src/skillmodels/common/measurement_models.py @@ -2,9 +2,10 @@ One JAX-compatible kernel evaluates the log density/probability contribution of a single measurement given its linear predictor ``eta = c + x'beta + lambda'theta``, -its scale, and a `MeasurementFamily` code. AF estimation, CHS, simulation and -posterior-state reweighting all route their measurement contributions through this -kernel, so the measurement system is defined in exactly one place. +its scale, and a `MeasurementFamily` code. AF estimation, simulation and +posterior-state reweighting route their measurement contributions through this +kernel. CHS (a Gaussian Kalman update) and AMN do not, and reject non-Gaussian +measurements. The three families: diff --git a/tests/test_chs_limited_measurement_rejection.py b/tests/test_chs_limited_measurement_rejection.py new file mode 100644 index 00000000..37ca86de --- /dev/null +++ b/tests/test_chs_limited_measurement_rejection.py @@ -0,0 +1,71 @@ +"""CHS must reject non-Gaussian measurements rather than treat them as continuous. + +The CHS likelihood runs a Gaussian Kalman update for every measurement, so a probit +or Tobit measure would be silently estimated as a continuous Gaussian one. +""" + +import numpy as np +import pandas as pd +import pytest + +from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.chs.options import CHSEstimationOptions +from skillmodels.common.measurement_models import ( + ProbitMeasurement, + TobitMeasurement, +) +from skillmodels.common.model_spec import ( + FactorSpec, + ModelSpec, + Normalizations, +) + + +def _model(measurement_models) -> ModelSpec: + return ModelSpec( + factors={ + "skills": FactorSpec( + measurements=(("y1", "y2", "y3"),) * 2, + normalizations=Normalizations( + loadings=({"y1": 1},) * 2, + intercepts=({"y1": 0},) * 2, + ), + transition_function="linear", + ), + }, + measurement_models=measurement_models, + ) + + +def _tiny_data() -> pd.DataFrame: + rng = np.random.default_rng(0) + rows = [ + { + "caseid": caseid, + "period": period, + **dict(zip(("y1", "y2", "y3"), rng.normal(size=3), strict=True)), + } + for caseid in range(20) + for period in (0, 1) + ] + return pd.DataFrame(rows).set_index(["caseid", "period"]) + + +@pytest.mark.parametrize("strategy", ["amn", "none"]) +@pytest.mark.parametrize( + "models", + [ + {"y2": ProbitMeasurement()}, + {"y3": TobitMeasurement(lower=0.0)}, + ], +) +def test_get_maximization_inputs_rejects_non_gaussian_measurements( + models, strategy +) -> None: + """Rejected whatever the start-value strategy, not only via AMN seeding.""" + with pytest.raises(NotImplementedError, match="CHS supports only Gaussian"): + get_maximization_inputs( + model_spec=_model(models), + data=_tiny_data(), + chs_options=CHSEstimationOptions(start_params_strategy=strategy), + ) From f96a086d7c28d68491b55a97a06aa71bcac03b54 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 12 Oct 2026 10:10:02 +0200 Subject: [PATCH 07/18] Let estimate_chs replace bounds and start on linear equality constraints lower_bounds now replaces the template bound (so -inf drops it, needed for parameters in a LinearConstraint), and the start point is projected onto all LinearConstraint equalities. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/estimate.py | 18 ++++---- src/skillmodels/common/constraints.py | 37 ++++++++++++++++ tests/test_chs_estimate.py | 63 +++++++++++++++++++++++++++ tests/test_constraints.py | 24 ++++++++++ 4 files changed, 134 insertions(+), 8 deletions(-) diff --git a/src/skillmodels/chs/estimate.py b/src/skillmodels/chs/estimate.py index c14f63ca..d68cd586 100644 --- a/src/skillmodels/chs/estimate.py +++ b/src/skillmodels/chs/estimate.py @@ -23,6 +23,7 @@ from skillmodels.chs.types import CHSEstimationResult from skillmodels.common.constraints import ( enforce_fixed_constraints, + project_start_to_linear_equalities, reconcile_start_to_equality, ) from skillmodels.common.identification import ( @@ -67,9 +68,11 @@ def estimate_chs( constraints: Optional extra `optimagic` constraints, appended to the model-implied constraints from `get_maximization_inputs`. lower_bounds: Optional lower bounds, indexed like the `params_template`, - that tighten the template's bounds (e.g. a floor on a measurement - SD that would otherwise collapse). Start values below them are - raised onto them. + that replace the template's bounds for those parameters: a floor on + a measurement SD that would otherwise collapse, or `-inf` to drop a + bound (optimagic counts each bound on a parameter in a + `LinearConstraint` as one more linear constraint). Start values + below them are raised onto them. require_initial_anchors: When True (the default), run the estimator-agnostic period-0 anchor check (`fail_if_initial_state_unanchored`) and raise `ValueError` if the @@ -108,11 +111,9 @@ def estimate_chs( start.loc[overlay.notna(), "value"] = overlay[overlay.notna()] if lower_bounds is not None: - tighter = lower_bounds.reindex(start.index) - locs = tighter.dropna().index - start.loc[locs, "lower_bound"] = np.maximum( - start.loc[locs, "lower_bound"], tighter[locs] - ) + replacement = lower_bounds.reindex(start.index) + locs = replacement.dropna().index + start.loc[locs, "lower_bound"] = replacement[locs] start.loc[locs, "value"] = np.maximum( start.loc[locs, "value"], start.loc[locs, "lower_bound"] ) @@ -134,6 +135,7 @@ def estimate_chs( # group's seeded value onto the constraint surface -- honouring any fixed # member's enforced value so the line above is not averaged away. start = reconcile_start_to_equality(start, all_constraints) + start = project_start_to_linear_equalities(start, all_constraints) optimize_options = { "algorithm": options.optimizer_algorithm, diff --git a/src/skillmodels/common/constraints.py b/src/skillmodels/common/constraints.py index 285aac29..6a78e390 100644 --- a/src/skillmodels/common/constraints.py +++ b/src/skillmodels/common/constraints.py @@ -196,6 +196,43 @@ def reconcile_start_to_equality( return out +def project_start_to_linear_equalities( + params: pd.DataFrame, + constraints: Iterable[om.constraints.Constraint], +) -> pd.DataFrame: + """Move the start point onto every linear equality constraint. + + `om.minimize` raises `InvalidParamsError` when a constraint is violated at the + starting point. All `om.LinearConstraint`s with a `value` are stacked into one + system `A x = b` over the parameters they involve, and those parameters' `value` + is replaced by its orthogonal projection onto the solution set (e.g. the + least-squares line for second-difference constraints). Other parameters are + untouched. Returns a copy; `params` is not modified. + """ + rows = [] + for c in constraints: + if isinstance(c, om.LinearConstraint) and c.value is not None: + selected = pd.Series(c.selector(params)).index + weights = np.broadcast_to( + np.asarray(c.weights, dtype=float), (len(selected),) + ) + rows.append((selected, weights, float(c.value))) + if not rows: + return params + involved = pd.Index( + sorted({tup for idx, _, _ in rows for tup in idx}), tupleize_cols=False + ) + a = np.zeros((len(rows), len(involved))) + b = np.array([value for _, _, value in rows]) + for i, (idx, weights, _) in enumerate(rows): + a[i, involved.get_indexer(idx)] = weights + x = params.loc[list(involved), "value"].to_numpy(dtype=float) + projected = x - np.linalg.pinv(a) @ (a @ x - b) + out = params.copy() + out.loc[list(involved), "value"] = projected + return out + + def collect_fixed_locs( constraints: Iterable[om.constraints.Constraint], ) -> set[tuple[Any, ...]]: diff --git a/tests/test_chs_estimate.py b/tests/test_chs_estimate.py index f83f221d..24462255 100644 --- a/tests/test_chs_estimate.py +++ b/tests/test_chs_estimate.py @@ -1,8 +1,10 @@ """End-to-end smoke test for `skillmodels.chs.estimate.estimate_chs`.""" +import functools from pathlib import Path import numpy as np +import optimagic as om import pandas as pd import pytest @@ -11,6 +13,7 @@ from skillmodels.chs.types import CHSEstimationResult from skillmodels.common.config import TEST_DATA_DIR from skillmodels.common.fixed_constraint import FixedConstraintWithValue +from skillmodels.common.selector import select_by_loc from skillmodels.test_data.model2 import MODEL2 REGRESSION_VAULT = Path(__file__).parent / "regression_vault" @@ -180,3 +183,63 @@ def fake_estimate_ml(**kwargs): assert captured["params"].loc[loc, "value"].tolist() == [0.1] other = captured["bounds"].lower.drop(loc) assert other[np.isfinite(other)].max() < 0.1 + + +def test_estimate_chs_lower_bounds_can_remove_a_bound( + model2_data, anchoring_start_params, monkeypatch +): + """A `-inf` entry in `lower_bounds` replaces the template's bound.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + raise StopIteration + + loc = anchoring_start_params.query("category == 'meas_sds'").index[:1] + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions(start_params_strategy="none"), + start_params=anchoring_start_params, + lower_bounds=pd.Series(-np.inf, index=loc), + require_initial_anchors=False, + ) + assert captured["bounds"].lower.loc[loc].tolist() == [-np.inf] + + +def test_estimate_chs_moves_start_onto_linear_equalities( + model2_data, anchoring_start_params, monkeypatch +): + """Start values satisfy user `LinearConstraint` equalities (a linear trend).""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + raise StopIteration + + sds = anchoring_start_params.query("category == 'meas_sds' & name1 == 'y1'").index + start = anchoring_start_params.copy() + start.loc[sds, "value"] = np.linspace(0.5, 1.5, len(sds)) ** 2 + constraints: list[om.constraints.Constraint] = [ + om.LinearConstraint( + selector=functools.partial(select_by_loc, loc=list(sds[t : t + 3])), + weights=np.array([1.0, -2.0, 1.0]), + value=0.0, + ) + for t in range(len(sds) - 2) + ] + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions(start_params_strategy="none"), + start_params=start, + constraints=constraints, + lower_bounds=pd.Series(-np.inf, index=sds), + require_initial_anchors=False, + ) + values = captured["params"].loc[sds, "value"].to_numpy() + np.testing.assert_allclose(np.diff(values, n=2), 0.0, atol=1e-12) diff --git a/tests/test_constraints.py b/tests/test_constraints.py index 5fa52218..1076ca4a 100644 --- a/tests/test_constraints.py +++ b/tests/test_constraints.py @@ -22,6 +22,7 @@ _get_transition_constraints, add_bounds, get_constraints, + project_start_to_linear_equalities, reconcile_start_to_equality, ) from skillmodels.common.process_model import process_model @@ -787,3 +788,26 @@ def test_get_constraints_for_augmented_periods(simplest_augmented_model) -> None {"loc": ("shock_sds", 1, "fac2", "-"), "type": "fixed", "value": 0.00000001}, ] assert_list_equal_except_for_order(as_dicts, expected) + + +def test_project_start_to_linear_equalities_fits_a_line() -> None: + """Second-difference equalities make the start the least-squares line.""" + index = pd.MultiIndex.from_tuples( + [("meas_sds", t, "alive", "-") for t in range(5)] + [("loadings", 0, "y", "f")], + names=["category", "aug_period", "name1", "name2"], + ) + values = np.array([0.2, 0.25, 0.2, 0.35, 0.5, 7.0]) + params = pd.DataFrame({"value": values}, index=index) + sds = list(index[:5]) + constraints = [ + om.LinearConstraint( + selector=functools.partial(select_by_loc, loc=sds[t : t + 3]), + weights=np.array([1.0, -2.0, 1.0]), + value=0.0, + ) + for t in range(3) + ] + slope, intercept = np.polyfit(np.arange(5), values[:5], deg=1) + expected = np.append(intercept + slope * np.arange(5), 7.0) + projected = project_start_to_linear_equalities(params, constraints) + np.testing.assert_allclose(projected["value"].to_numpy(), expected, atol=1e-12) From 1c5bb8ed41bfd6e703e2f1bb42a77bde2eb295e9 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 12 Oct 2026 22:14:47 +0200 Subject: [PATCH 08/18] Add upper_bounds to estimate_chs The counterpart of lower_bounds, e.g. 0 to restrict a loading's sign. Start values above a bound are lowered onto it. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/estimate.py | 12 ++++++++++++ tests/test_chs_estimate.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/src/skillmodels/chs/estimate.py b/src/skillmodels/chs/estimate.py index d68cd586..d9fc21e0 100644 --- a/src/skillmodels/chs/estimate.py +++ b/src/skillmodels/chs/estimate.py @@ -43,6 +43,7 @@ def estimate_chs( fixed_params: pd.DataFrame | None = None, constraints: list[om.constraints.Constraint] | None = None, lower_bounds: pd.Series | None = None, + upper_bounds: pd.Series | None = None, *, require_initial_anchors: bool = True, ) -> CHSEstimationResult: @@ -73,6 +74,9 @@ def estimate_chs( bound (optimagic counts each bound on a parameter in a `LinearConstraint` as one more linear constraint). Start values below them are raised onto them. + upper_bounds: Optional upper bounds, the counterpart of `lower_bounds` + (e.g. `0` to restrict a loading's sign). Start values above them + are lowered onto them. require_initial_anchors: When True (the default), run the estimator-agnostic period-0 anchor check (`fail_if_initial_state_unanchored`) and raise `ValueError` if the @@ -118,6 +122,14 @@ def estimate_chs( start.loc[locs, "value"], start.loc[locs, "lower_bound"] ) + if upper_bounds is not None: + replacement = upper_bounds.reindex(start.index) + locs = replacement.dropna().index + start.loc[locs, "upper_bound"] = replacement[locs] + start.loc[locs, "value"] = np.minimum( + start.loc[locs, "value"], start.loc[locs, "upper_bound"] + ) + all_constraints = [*max_inputs["constraints"], *(constraints or [])] # Write every `FixedConstraintWithValue`'s target into the start vector. diff --git a/tests/test_chs_estimate.py b/tests/test_chs_estimate.py index 24462255..4061eaf3 100644 --- a/tests/test_chs_estimate.py +++ b/tests/test_chs_estimate.py @@ -209,6 +209,35 @@ def fake_estimate_ml(**kwargs): assert captured["bounds"].lower.loc[loc].tolist() == [-np.inf] +def test_estimate_chs_applies_upper_bounds_and_clips_start( + model2_data, anchoring_start_params, monkeypatch +): + """`upper_bounds` replaces the template bounds; the start moves onto them.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + raise StopIteration + + loc = pd.MultiIndex.from_tuples( + [("controls", 0, "y1", "x1")], names=anchoring_start_params.index.names + ) + start = anchoring_start_params.copy() + start.loc[loc, "value"] = 0.5 + monkeypatch.setattr("skillmodels.chs.estimate.estimate_ml", fake_estimate_ml) + with pytest.raises(StopIteration): + estimate_chs( + MODEL2, + model2_data, + CHSEstimationOptions(start_params_strategy="none"), + start_params=start, + upper_bounds=pd.Series(0.0, index=loc), + require_initial_anchors=False, + ) + assert captured["bounds"].upper.loc[loc].tolist() == [0.0] + assert captured["params"].loc[loc, "value"].tolist() == [0.0] + + def test_estimate_chs_moves_start_onto_linear_equalities( model2_data, anchoring_start_params, monkeypatch ): From f6ba3f5196e9f1f485d098f8ac3f062b81e397f4 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Wed, 14 Oct 2026 21:16:18 +0200 Subject: [PATCH 09/18] CHS: integrate out missing observed factors using data-supplied moments `ModelSpec.missing_observed_factor_moments` maps an observed factor to the `(mean, variance)` data columns of its distribution where it is missing. In such cells the CHS filter treats the factor as N(m, v), independent of the latent states, and integrates it out in the predict step instead of raising. - process_data: validate that missing cells have a finite fill mean and a non-negative fill variance (errors name variable and period); substitute the mean into the observed-factor array and build a per-cell variance array for the declared factors that are actually missing somewhere. Undeclared factors and non-estimation purposes keep the "Variable has missings" check. - linear_kalman_predict: exact; the square-root stack gains rows sqrt(v_j) * B[:, j] / s_out, adding B diag(v) B' (with cross-equation terms). - kalman_predict: augment the sigma-point input by the possibly-missing factors (variance 0 where observed); sigma weights use n_latent + n_missing. The augmented sigma values are written into a promoted copy of the float32 observed-factor array so they are not rounded. - Models without declared moments, or whose declared factors are complete, take the unchanged code path (bit-identical likelihood). Tests compare the filter likelihood against a Monte-Carlo integral over the missing cells (1 and 2 mixtures), the predict steps against the analytic moments, the unscented path against the linear path for a linear translog, and cover the validation errors. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/kalman_filters.py | 111 ++++- src/skillmodels/chs/likelihood.py | 12 + src/skillmodels/chs/likelihood_debug.py | 7 + src/skillmodels/chs/maximization_inputs.py | 22 +- src/skillmodels/common/model_spec.py | 36 ++ src/skillmodels/common/process_data.py | 83 +++- src/skillmodels/common/process_model.py | 1 + src/skillmodels/common/types.py | 5 + tests/test_missing_observed_factors.py | 446 +++++++++++++++++++++ 9 files changed, 711 insertions(+), 12 deletions(-) create mode 100644 tests/test_missing_observed_factors.py diff --git a/src/skillmodels/chs/kalman_filters.py b/src/skillmodels/chs/kalman_filters.py index bcb67d88..3ba5d91e 100644 --- a/src/skillmodels/chs/kalman_filters.py +++ b/src/skillmodels/chs/kalman_filters.py @@ -219,9 +219,19 @@ def kalman_predict( anchoring_scaling_factors: Array, anchoring_constants: Array, observed_factors: Array, + observed_factor_vars: Array | None = None, + *, + missing_observed_indices: tuple[int, ...] = (), ) -> tuple[Array, Array]: """Make a unscented Kalman predict. + Observed factors at `missing_observed_indices` may be missing. They augment + the sigma-point input as extra Gaussian dimensions, independent of the latent + states, with mean `observed_factors` (the fill mean where missing) and + variance `observed_factor_vars` (0 where observed, so observed values pass + through unchanged). The number of sigma points, and hence `sigma_weights`, + is based on `n_latent + len(missing_observed_indices)` for every observation. + Args: transition_func: The transition function. states: Array of shape (n_obs, n_mixtures, n_states) with @@ -245,17 +255,31 @@ def kalman_predict( period, the second to the output period (i.e. input period + 1). observed_factors: Array of shape (n_obs, n_observed_factors) with data on the observed factors in period t. + observed_factor_vars: Array of shape + (n_obs, len(missing_observed_indices)) with the variance of each + possibly-missing observed factor (0 where it is observed). + missing_observed_indices: Positions (among the observed factors) of the + observed factors that may be missing. Returns: jax.numpy.array: Predicted states, same shape as states. jax.numpy.array: Predicted upper_chols, same shape as upper_chols. """ + if missing_observed_indices: + states, upper_chols = _augment_by_missing_observed_factors( + states=states, + upper_chols=upper_chols, + observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, + missing_observed_indices=missing_observed_indices, + ) sigma_points = _calculate_sigma_points( states=states, upper_chols=upper_chols, scaling_factor=sigma_scaling_factor, observed_factors=observed_factors, + missing_observed_indices=missing_observed_indices, ) transformed = transform_sigma_points( sigma_points=sigma_points, @@ -292,10 +316,12 @@ def linear_kalman_predict( anchoring_scaling_factors: Array, anchoring_constants: Array, observed_factors: Array, + observed_factor_vars: Array | None = None, *, latent_factors: tuple[str, ...], constant_factor_indices: frozenset[int], n_all_factors: int, + missing_observed_indices: tuple[int, ...] = (), ) -> tuple[Array, Array]: """Make a linear Kalman predict (square-root form). @@ -307,6 +333,13 @@ def linear_kalman_predict( `sigma_weights` are accepted for signature compatibility with `kalman_predict` but are ignored. + Observed factors at `missing_observed_indices` may be missing, in which case + they are `N(m, v)` independent of the latent states, with `m` given in + `observed_factors` and `v` in `observed_factor_vars`. The prediction is exact: + the mean uses `m`, and the covariance adds `B diag(v) B'` with `B` the + unanchored coefficients of those factors, so the cross-covariance induced by + the same factor entering several equations is kept. + Args: transition_func: Ignored (kept for signature compatibility). states: Array of shape (n_obs, n_mixtures, n_states). @@ -318,9 +351,14 @@ def linear_kalman_predict( anchoring_scaling_factors: Array of shape (2, n_states). anchoring_constants: Array of shape (2, n_states). observed_factors: Array of shape (n_obs, n_observed_factors). + observed_factor_vars: Array of shape + (n_obs, len(missing_observed_indices)) with the variance of each + possibly-missing observed factor (0 where it is observed). latent_factors: Tuple of latent factor names. constant_factor_indices: Indices of factors with `constant` transition. n_all_factors: Total number of factors (latent + observed). + missing_observed_indices: Positions (among the observed factors) of the + observed factors that may be missing. Returns: Predicted states, same shape as states. @@ -360,9 +398,20 @@ def linear_kalman_predict( chol_g = upper_chols @ g_mat.T # (n_obs, n_mix, n_latent, n_latent) shock_diag = jnp.diag(shock_sds / s_out) # (n_latent, n_latent) - stack = jnp.concatenate( - [chol_g, jnp.broadcast_to(shock_diag, chol_g.shape)], axis=-2 - ) # (n_obs, n_mix, 2*n_latent, n_latent) + blocks = [chol_g, jnp.broadcast_to(shock_diag, chol_g.shape)] + if missing_observed_indices: + if observed_factor_vars is None: + msg = "observed_factor_vars is required with missing_observed_indices." + raise ValueError(msg) + # Row j is sqrt(v_j) * B[:, j], so the rows add B diag(v) B' to the cov. + b_missing = f_mat[:, n_latent + jnp.array(missing_observed_indices)] + b_missing = b_missing / s_out[:, jnp.newaxis] # (n_latent, n_missing) + rows = jnp.sqrt(observed_factor_vars)[..., jnp.newaxis] * b_missing.T + blocks.append( + jnp.broadcast_to(rows[:, jnp.newaxis], (n_obs, n_mix, *b_missing.T.shape)) + ) + # Stack shape is (n_obs, n_mix, 2 * n_latent + n_missing, n_latent). + stack = jnp.concatenate(blocks, axis=-2) predicted_covs = array_qr_jax(stack)[1][:, :, :n_latent] @@ -421,6 +470,7 @@ def _calculate_sigma_points( upper_chols: Array, scaling_factor: float | Array, observed_factors: Array, + missing_observed_indices: tuple[int, ...] = (), ) -> Array: """Calculate the array of sigma_points for the unscented transform. @@ -435,6 +485,9 @@ def _calculate_sigma_points( the sigma_point algorithm chosen. observed_factors: Array of shape (n_obs, n_observed_factors) with data on the observed factors in period t. + missing_observed_indices: Positions (among the observed factors) of the + observed factors held in the trailing `len(missing_observed_indices)` + entries of `states`. Their sigma-point values replace the data. Returns: jax.numpy.array: Array of shape n_obs, n_mixtures, n_sigma, n_fac (where n_sigma @@ -444,6 +497,7 @@ def _calculate_sigma_points( n_obs, n_mixtures, n_fac = states.shape n_sigma = 2 * n_fac + 1 n_observed = observed_factors.shape[1] + n_missing = len(missing_observed_indices) scaled_upper_chols = upper_chols * scaling_factor sigma_points = jnp.repeat(states, n_sigma, axis=1).reshape( @@ -459,10 +513,61 @@ def _calculate_sigma_points( observed_factors[:, jnp.newaxis, jnp.newaxis, :], (n_obs, n_mixtures, n_sigma, n_observed), ) + if n_missing: + # Promote first: observed factors are stored as float32, and writing the + # float64 sigma-point values into them would round those values. + observed_part = observed_part.astype(sigma_points.dtype) + observed_part = observed_part.at[..., jnp.array(missing_observed_indices)].set( + sigma_points[..., n_fac - n_missing :] + ) + sigma_points = sigma_points[..., : n_fac - n_missing] return jnp.concatenate([sigma_points, observed_part], axis=-1) +def _augment_by_missing_observed_factors( + states: Array, + upper_chols: Array, + observed_factors: Array, + observed_factor_vars: Array | None, + missing_observed_indices: tuple[int, ...], +) -> tuple[Array, Array]: + """Append possibly-missing observed factors as independent Gaussian states. + + Args: + states: Array of shape (n_obs, n_mixtures, n_latent). + upper_chols: Array of shape (n_obs, n_mixtures, n_latent, n_latent). + observed_factors: Array of shape (n_obs, n_observed_factors) with the fill + mean where a factor is missing. + observed_factor_vars: Array of shape (n_obs, n_missing) with the fill + variance where a factor is missing and 0 where it is observed. + missing_observed_indices: Positions of the possibly-missing factors among + the observed factors. + + Returns: + Augmented states of shape (n_obs, n_mixtures, n_latent + n_missing) and + block-diagonal upper Cholesky factors of matching shape. + + """ + if observed_factor_vars is None: + msg = "observed_factor_vars is required with missing_observed_indices." + raise ValueError(msg) + n_obs, n_mixtures, n_latent = states.shape + n_missing = len(missing_observed_indices) + n_aug = n_latent + n_missing + means = observed_factors[:, jnp.array(missing_observed_indices)] + aug_states = jnp.concatenate( + [states, jnp.broadcast_to(means[:, None], (n_obs, n_mixtures, n_missing))], + axis=-1, + ) + sds = jnp.sqrt(observed_factor_vars) # (n_obs, n_missing) + aug_chols = jnp.zeros((n_obs, n_mixtures, n_aug, n_aug), dtype=upper_chols.dtype) + aug_chols = aug_chols.at[..., :n_latent, :n_latent].set(upper_chols) + diag = jnp.arange(n_latent, n_aug) + aug_chols = aug_chols.at[..., diag, diag].set(sds[:, None, :]) + return aug_states, aug_chols + + def transform_sigma_points( sigma_points: Array, transition_func: Callable, diff --git a/src/skillmodels/chs/likelihood.py b/src/skillmodels/chs/likelihood.py index 2f9a5225..463bf7b9 100644 --- a/src/skillmodels/chs/likelihood.py +++ b/src/skillmodels/chs/likelihood.py @@ -36,6 +36,7 @@ def log_likelihood( is_predict_iteration: Array | np.ndarray, iteration_to_period: Array | np.ndarray, observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> Array: """Aggregated log likelihood of a skill formation model. @@ -66,6 +67,9 @@ def log_likelihood( iteration_to_period: Array mapping iteration index to period. observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with data on the observed factors. + observed_factor_vars: Array of shape (n_periods, n_obs, n_missing) with + the variance of each possibly-missing observed factor (0 where it is + observed). Returns: Scalar aggregated log likelihood. @@ -86,6 +90,7 @@ def log_likelihood( is_predict_iteration=is_predict_iteration, iteration_to_period=iteration_to_period, observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, ).sum() @@ -104,6 +109,7 @@ def log_likelihood_obs( is_predict_iteration: Array | np.ndarray, iteration_to_period: Array | np.ndarray, observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> Array: """Log likelihood of a skill formation model. @@ -143,6 +149,9 @@ def log_likelihood_obs( iteration_to_period: Array mapping iteration index to period. observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with data on the observed factors. + observed_factor_vars: Array of shape (n_periods, n_obs, n_missing) with + the variance of each possibly-missing observed factor (0 where it is + observed). Returns: jnp.array: 1d array of length N, the aggregated log likelihood. @@ -181,6 +190,7 @@ def log_likelihood_obs( sigma_weights=sigma_weights, predict_func=predict_func, observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, ) _body = jax.checkpoint(_body, prevent_cse=False) static_out = jax.lax.scan(_body, carry, loop_args, unroll=False)[1] @@ -205,6 +215,7 @@ def _scan_body( sigma_weights: Array, predict_func: Callable, observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> tuple[dict[str, Array], dict[str, Array]]: # ================================================================================== # create arguments needed for update @@ -250,6 +261,7 @@ def _scan_body( ], "anchoring_constants": parsed_params.anchoring_constants[jnp.array([t, t + 1])], "observed_factors": observed_factors[t], + "observed_factor_vars": observed_factor_vars[t], } fixed_kwargs = {"predict_func": predict_func} diff --git a/src/skillmodels/chs/likelihood_debug.py b/src/skillmodels/chs/likelihood_debug.py index 6a77bc27..88a80481 100644 --- a/src/skillmodels/chs/likelihood_debug.py +++ b/src/skillmodels/chs/likelihood_debug.py @@ -36,6 +36,7 @@ def log_likelihood( is_predict_iteration: Array | np.ndarray, iteration_to_period: Array | np.ndarray, observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> dict[str, Any]: """Log likelihood of a skill formation model, returning debug data on top. @@ -65,6 +66,9 @@ def log_likelihood( iteration_to_period: Array mapping iteration index to period. observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with data on the observed factors. + observed_factor_vars: Array of shape (n_periods, n_obs, n_missing) with + the variance of each possibly-missing observed factor (0 where it is + observed). Returns: All data relevant for debugging, e.g. the log likelihood contribution of @@ -104,6 +108,7 @@ def log_likelihood( sigma_weights=sigma_weights, predict_func=predict_func, observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, ) static_out = jax.lax.scan(_body, carry, loop_args)[1] @@ -156,6 +161,7 @@ def _scan_body( sigma_weights: Array, predict_func: Callable[..., tuple[Array, Array]], observed_factors: Array | np.ndarray, + observed_factor_vars: Array | np.ndarray, ) -> tuple[dict[str, Array], dict[str, Any]]: # ================================================================================== # create arguments needed for update @@ -201,6 +207,7 @@ def _scan_body( ], "anchoring_constants": parsed_params.anchoring_constants[jnp.array([t, t + 1])], "observed_factors": observed_factors[t], + "observed_factor_vars": observed_factor_vars[t], } fixed_kwargs = {"predict_func": predict_func} diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index 9063f55f..e46fa9af 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -129,8 +129,10 @@ def get_maximization_inputs( # noqa: C901, PLR0915 purpose="estimation", ) + # Possibly-missing observed factors augment the sigma-point input. sigma_scaling_factor, sigma_weights = calculate_sigma_scaling_factor_and_weights( - n_states=processed_model.dimensions.n_latent_factors, + n_states=processed_model.dimensions.n_latent_factors + + len(processed_data["missing_observed_factors"]), kappa=chs_options.sigma_points_scale, ) @@ -151,6 +153,8 @@ def get_maximization_inputs( # noqa: C901, PLR0915 measurements=processed_data["measurements"], controls=processed_data["controls"], observed_factors=processed_data["observed_factors"], + observed_factor_vars=processed_data["observed_factor_vars"], + missing_observed_factors=processed_data["missing_observed_factors"], model=processed_model, sigma_weights=sigma_weights, sigma_scaling_factor=sigma_scaling_factor, @@ -191,6 +195,9 @@ def loglike_and_gradient( measurements=measurements_slice, controls=controls_slice, observed_factors=observed_factors_slice, + observed_factor_vars=processed_data["observed_factor_vars"][ + :, start:stop, : + ], ) start += step stop += step @@ -409,6 +416,8 @@ def _partial_some_log_likelihood( measurements: Array, controls: Array, observed_factors: Array, + observed_factor_vars: Array, + missing_observed_factors: tuple[str, ...], model: ProcessedModel, sigma_weights: Array, sigma_scaling_factor: Array, @@ -446,6 +455,9 @@ def _partial_some_log_likelihood( # states (cf is a non-linear function of the latent states), so the linear # fast-path would silently bypass the cf DAG node. Force the unscented path. has_control_function = model.endogenous_factors_info.control_function is not None + missing_observed_indices = tuple( + model.labels.observed_factors.index(f) for f in missing_observed_factors + ) if is_all_linear(model.transition_info.function_names) and not has_control_function: constant_factor_indices = frozenset( i @@ -458,9 +470,14 @@ def _partial_some_log_likelihood( latent_factors=model.labels.latent_factors, constant_factor_indices=constant_factor_indices, n_all_factors=model.dimensions.n_all_factors, + missing_observed_indices=missing_observed_indices, ) else: - predict_func = functools.partial(kalman_predict, model.transition_info.func) + predict_func = functools.partial( + kalman_predict, + model.transition_info.func, + missing_observed_indices=missing_observed_indices, + ) return functools.partial( fun, @@ -477,6 +494,7 @@ def _partial_some_log_likelihood( is_predict_iteration=is_predict_iteration, iteration_to_period=iteration_to_period, observed_factors=observed_factors, + observed_factor_vars=observed_factor_vars, ) diff --git a/src/skillmodels/common/model_spec.py b/src/skillmodels/common/model_spec.py index b9f57d9a..b90cd789 100644 --- a/src/skillmodels/common/model_spec.py +++ b/src/skillmodels/common/model_spec.py @@ -211,6 +211,19 @@ class ModelSpec: ) """Per-variable observation model (probit / Tobit / Gaussian). A measurement omitted from this mapping is Gaussian, so existing specs are unchanged.""" + missing_observed_factor_moments: MappingProxyType[str, tuple[str, str]] = field( + default_factory=lambda: MappingProxyType({}) + ) + """Map an observed factor that may be missing to its `(mean, variance)` columns. + + Where the factor is missing for a person-period, the CHS filter treats it as + `N(mean, variance)`, independent of the latent states, with both moments read + from that person-period's row of the data, and integrates it out in the + predict step. Where the factor is observed, the moment columns are ignored. + An observed factor not listed here must never be missing. Only the CHS + likelihood (and the filtered states built on it) integrates missing values + out; simulation and transition plots still require complete observed + factors.""" def __init__( self, @@ -221,6 +234,7 @@ def __init__( anchoring: AnchoringSpec | None = None, n_mixtures: int = 1, measurement_models: Mapping[str, MeasurementModel] | None = None, + missing_observed_factor_moments: Mapping[str, tuple[str, str]] | None = None, ) -> None: """Create ModelSpec, wrapping factors dict in MappingProxyType.""" object.__setattr__(self, "_factors", ensure_containers_are_immutable(factors)) @@ -230,6 +244,17 @@ def __init__( object.__setattr__(self, "anchoring", anchoring) object.__setattr__(self, "n_mixtures", n_mixtures) self._set_measurement_models(measurement_models) + moments = dict(missing_observed_factor_moments or {}) + unknown = sorted(name for name in moments if name not in observed_factors) + if unknown: + msg = ( + f"missing_observed_factor_moments references {unknown}, which is not " + f"an observed factor. Observed factors: {list(observed_factors)}." + ) + raise ValueError(msg) + object.__setattr__( + self, "missing_observed_factor_moments", MappingProxyType(moments) + ) def _set_measurement_models( self, measurement_models: Mapping[str, MeasurementModel] | None @@ -324,6 +349,12 @@ def from_dict(cls, d: dict[str, Any]) -> Self: observed += tuple(i for i in auto_instruments if i not in observed) measurement_models = _parse_measurement_models(d.get("measurement_models")) + moments = { + name: (mean_col, var_col) + for name, (mean_col, var_col) in d.get( + "missing_observed_factor_moments", {} + ).items() + } return cls( factors=factors, @@ -333,6 +364,7 @@ def from_dict(cls, d: dict[str, Any]) -> Self: anchoring=anchoring, n_mixtures=d.get("n_mixtures", 1), measurement_models=measurement_models, + missing_observed_factor_moments=moments, ) @property @@ -352,6 +384,10 @@ def _replace(self, **changes: Any) -> Self: # noqa: ANN401 measurement_models=changes.get( "measurement_models", self.measurement_models ), + missing_observed_factor_moments=changes.get( + "missing_observed_factor_moments", + self.missing_observed_factor_moments, + ), ) def with_transition_functions( diff --git a/src/skillmodels/common/process_data.py b/src/skillmodels/common/process_data.py index 8e3f48de..f1c4a075 100644 --- a/src/skillmodels/common/process_data.py +++ b/src/skillmodels/common/process_data.py @@ -42,7 +42,13 @@ def process_data( observed control variables for the measurement equations. observed_factors: Array of shape (n_periods, n_obs, n_observed_factors) with data on the observed factors. - Only returned if estimation==True + Where a factor listed in `labels.missing_observed_factor_moments` is + missing, it holds the fill mean. + missing_observed_factors: Names of the observed factors with fill moments + that are missing somewhere in the data (estimation only; else empty). + observed_factor_vars: Array of shape + (n_periods, n_obs, len(missing_observed_factors)) with the fill + variance where the factor is missing and 0 where it is observed. """ df = pre_process_data(df=df, periods=labels.periods) @@ -63,6 +69,20 @@ def process_data( df=df, controls=labels.controls, update_info=update_info ) out["controls"] = _generate_controls_array(df=df, labels=labels, n_obs=n_obs) + missing = tuple( + f + for f in labels.observed_factors + if purpose == "estimation" + and f in labels.missing_observed_factor_moments + and df[f].isna().any() + ) + out["missing_observed_factors"] = missing + out["observed_factor_vars"] = _generate_observed_factor_vars_array( + df=df, labels=labels, n_obs=n_obs, factors=missing + ) + for factor in missing: + mean_col, _ = labels.missing_observed_factor_moments[factor] + df[factor] = df[factor].fillna(df[mean_col]) out["observed_factors"] = _generate_observed_factor_array( df=df, labels=labels, n_obs=n_obs ) @@ -119,6 +139,9 @@ def _get_period_data_for_endogenous_factors( meas = _get_period_measurements(update_info=update_info, aug_period=aug_period) controls = labels.controls observed = labels.observed_factors + moment_cols = [ + col for cols in labels.missing_observed_factor_moments.values() for col in cols + ] out = df.query(f"period == {period}")[ [ @@ -126,6 +149,7 @@ def _get_period_data_for_endogenous_factors( *meas, *controls, *observed, + *moment_cols, "period", "__old_id__", "__old_period__", @@ -209,12 +233,14 @@ def _check_data( # noqa: C901 ) for factor in labels.observed_factors: - if factor not in period_data.columns: - var_report.loc[(aug_period, factor), "problem"] = "Variable is missing" - elif period_data[factor].isna().any(): - var_report.loc[(aug_period, factor), "problem"] = ( - "Variable has missings" - ) + moments = ( + labels.missing_observed_factor_moments.get(factor) + if purpose == "estimation" + else None + ) + problems = _observed_factor_problems(period_data, factor, moments) + for var, problem in problems.items(): + var_report.loc[(aug_period, var), "problem"] = problem var_report = var_report.to_string() if len(var_report) > 0 else "" @@ -222,6 +248,32 @@ def _check_data( # noqa: C901 raise ValueError(var_report) +def _observed_factor_problems( + period_data: pd.DataFrame, factor: str, moments: tuple[str, str] | None +) -> dict[str, str]: + """Map variables to problems with an observed factor and its fill moments. + + Without `moments`, any missing value is a problem. With `moments`, missing + values need a non-missing fill mean and a non-negative fill variance. + """ + if factor not in period_data.columns: + return {factor: "Variable is missing"} + missing = period_data[factor].isna() + if moments is None: + return {factor: "Variable has missings"} if missing.any() else {} + absent = [col for col in moments if col not in period_data.columns] + if absent: + return dict.fromkeys(absent, "Variable is missing") + mean, var = (period_data.loc[missing, col] for col in moments) + if mean.isna().any() or var.isna().any(): + return {factor: f"Variable has missings without fill moments {moments}"} + if (var < 0).any(): + return { + factor: f"Variable has missings with negative fill variance in {moments[1]}" + } + return {} + + def _handle_controls_with_missings( df: pd.DataFrame, controls: tuple[str, ...], @@ -282,6 +334,23 @@ def _generate_controls_array( return jnp.array(arr, dtype="float32") +def _generate_observed_factor_vars_array( + df: pd.DataFrame, + labels: Labels, + n_obs: int, + factors: tuple[str, ...], +) -> Array: + arr = np.zeros((len(labels.aug_periods), n_obs, len(factors))) + for aug_period in labels.aug_periods: + period_data = df.query(f"aug_period == {aug_period}") + for k, factor in enumerate(factors): + var_col = labels.missing_observed_factor_moments[factor][1] + arr[aug_period, :, k] = np.where( + period_data[factor].isna(), period_data[var_col], 0.0 + ) + return jnp.array(arr, dtype="float32") + + def _generate_observed_factor_array( df: pd.DataFrame, labels: Labels, diff --git a/src/skillmodels/common/process_model.py b/src/skillmodels/common/process_model.py index 71c504c2..31f48438 100644 --- a/src/skillmodels/common/process_model.py +++ b/src/skillmodels/common/process_model.py @@ -237,6 +237,7 @@ def _get_labels( aug_stagemap=tuple(aug_stagemap), aug_stages=tuple(sorted(int(v) for v in np.unique(aug_stagemap))), aug_stages_to_stages=MappingProxyType(aug_stages_to_stages), + missing_observed_factor_moments=model_spec.missing_observed_factor_moments, ) diff --git a/src/skillmodels/common/types.py b/src/skillmodels/common/types.py index 2cf2de35..ded3f768 100644 --- a/src/skillmodels/common/types.py +++ b/src/skillmodels/common/types.py @@ -153,6 +153,11 @@ class Labels: """Map each augmented stage to its original stage.""" transition_names: tuple[str, ...] = () """Names of the transition functions per factor.""" + missing_observed_factor_moments: MappingProxyType[str, tuple[str, str]] = field( + default_factory=lambda: MappingProxyType({}) + ) + """Map an observed factor that may be missing to its `(mean, variance)` data + columns, see `ModelSpec.missing_observed_factor_moments`.""" @property def all_factors(self) -> tuple[str, ...]: diff --git a/tests/test_missing_observed_factors.py b/tests/test_missing_observed_factors.py new file mode 100644 index 00000000..8f627a5d --- /dev/null +++ b/tests/test_missing_observed_factors.py @@ -0,0 +1,446 @@ +"""Tests for integrating out missing observed factors in the CHS Kalman filter. + +Where an observed factor `x` is missing for a person-period, it is treated as +`x ~ N(m, v)`, independent of the latent states, with `(m, v)` taken from data +columns named in `ModelSpec.missing_observed_factor_moments`. +""" + +import jax +import jax.numpy as jnp +import numpy as np +import pandas as pd +import pytest +from numpy.testing import assert_allclose + +from skillmodels.chs.kalman_filters import ( + calculate_sigma_scaling_factor_and_weights, + kalman_predict, + linear_kalman_predict, +) +from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.chs.options import CHSEstimationOptions +from skillmodels.common.model_spec import FactorSpec, ModelSpec + +jax.config.update("jax_enable_x64", True) + +LATENT = ("fac1", "fac2") +OBSERVED = ("ob1", "ob2") +N_PERIODS = 3 +N_PERSONS = 6 +MOMENTS = {"ob1": ("ob1_mean", "ob1_var"), "ob2": ("ob2_mean", "ob2_var")} +OPTIONS = CHSEstimationOptions(start_params_strategy="none") + +# (person, period) cells where each observed factor is missing. Person 3 misses +# both in period 1 (cross-equation covariance), person 5 misses ob1 only in the +# last period (never used by a transition), person 4 misses nothing. +MISSING_CELLS = { + "ob1": ((0, 0), (1, 0), (1, 1), (3, 1), (5, 2)), + "ob2": ((1, 0), (2, 1), (3, 1)), +} + +# Log-likelihood of `_model(n_mixtures=2, transition=...)` on `_data()` without +# missings at `_params(..., nonlinear=0.05)`: reference values pinning the +# complete-data filter, which supports observed factors without fill moments. +BASELINE_LOGLIKE = { + "linear": -198.6803696097199, + "translog": -206.7069287270014, +} + + +def _model( + n_mixtures: int = 1, + transition: str = "linear", + moments: dict | None = None, +) -> ModelSpec: + factors = { + f: FactorSpec( + measurements=((f"{f}_m1", f"{f}_m2"),) * N_PERIODS, + transition_function=transition, + ) + for f in LATENT + } + return ModelSpec( + factors=factors, + observed_factors=OBSERVED, + n_mixtures=n_mixtures, + missing_observed_factor_moments=moments, + ) + + +def _round(arr: np.ndarray) -> np.ndarray: + """Round to a float32-representable grid so dtype casts are exact.""" + return np.round(arr * 64) / 64 + + +def _data(*, with_missings: bool = True) -> pd.DataFrame: + rng = np.random.default_rng(1234) + idx = pd.MultiIndex.from_product( + [range(N_PERSONS), range(N_PERIODS)], names=["id", "period"] + ) + n_rows = len(idx) + meas = [f"{f}_m{j}" for f in LATENT for j in (1, 2)] + df = pd.DataFrame(_round(rng.normal(size=(n_rows, 4))), index=idx, columns=meas) + df.loc[(2, 1), "fac1_m1"] = np.nan + for ob, (mean_col, var_col) in MOMENTS.items(): + df[ob] = _round(rng.normal(size=n_rows)) + df[mean_col] = _round(rng.normal(scale=0.5, size=n_rows)) + df[var_col] = _round(rng.uniform(0.3, 1.2, size=n_rows)) + if with_missings: + for ob, cells in MISSING_CELLS.items(): + for cell in cells: + df.loc[cell, ob] = np.nan + return df + + +def _params(template: pd.DataFrame, nonlinear: float = 0.0) -> pd.DataFrame: + """Fill a CHS params template with fixed, well-conditioned values.""" + params = template.copy() + lin = { + "fac1": {"fac1": 0.7, "fac2": 0.2, "ob1": 0.9, "ob2": -0.6, "constant": 0.1}, + "fac2": {"fac1": -0.3, "fac2": 0.6, "ob1": 0.5, "ob2": 0.8, "constant": -0.2}, + } + values = [] + for category, period, name1, name2 in params.index: + if category == "controls": + val = 0.1 + elif category == "loadings": + val = 1.0 if name1.endswith("m1") else 0.8 + elif category == "meas_sds": + val = 0.5 + elif category == "shock_sds": + val = 0.4 + 0.1 * period + elif category == "initial_states": + val = 0.0 if name1 == "mixture_0" else 1.0 + val += 0.3 if name2 == "fac2" else 0.0 + elif category == "mixture_weights": + val = 0.4 if name1 == "mixture_0" else 0.6 + elif category == "initial_cholcovs": + a, b = name2.split("-") + val = (1.0 if name1 == "mixture_0" else 0.7) if a == b else 0.3 + elif category == "transition": + val = nonlinear if "*" in name2 else lin[name1][name2] * (1 + 0.1 * period) + else: + raise ValueError(category) + values.append(val) + params["value"] = values + return params + + +def _loglikeobs(model: ModelSpec, data: pd.DataFrame, nonlinear: float = 0.0): + inputs = get_maximization_inputs(model, data, chs_options=OPTIONS) + params = _params(inputs["params_template"], nonlinear=nonlinear) + return inputs["loglikeobs"](params) + + +def _monte_carlo_likelihood( + n_mixtures: int, n_draws: int, seed: int +) -> tuple[np.ndarray, np.ndarray]: + """Average the per-person likelihood over draws of the missing cells. + + Each draw fills every missing observed-factor cell with `N(m, v)` and runs the + ordinary filter (all observed factors present). The draws are stacked as + replicas of the panel so one likelihood evaluation covers all of them. + + Returns: + Per-person Monte-Carlo likelihood and its standard error. + + """ + rng = np.random.default_rng(seed) + base = _data().reset_index() + rep = pd.concat([base] * n_draws, ignore_index=True) + replica = np.repeat(np.arange(n_draws), len(base)) + rep["id"] = np.tile(base["id"].to_numpy(), n_draws) + N_PERSONS * replica + for ob, (mean_col, var_col) in MOMENTS.items(): + values = rep[ob].to_numpy(copy=True) + miss = np.isnan(values) + mean = rep[mean_col].to_numpy()[miss] + sd = np.sqrt(rep[var_col].to_numpy()[miss]) + values[miss] = rng.normal(mean, sd) + rep[ob] = values + rep = rep.set_index(["id", "period"]) + loglikes = _loglikeobs(_model(n_mixtures), rep).reshape(n_draws, N_PERSONS) + likes = np.exp(loglikes) + return likes.mean(axis=0), likes.std(axis=0, ddof=1) / np.sqrt(n_draws) + + +@pytest.fixture(scope="module", params=[1, 2], ids=["1mix", "2mix"]) +def mc_reference(request): + n_mixtures = request.param + mc_like, mc_se = _monte_carlo_likelihood(n_mixtures, n_draws=40_000, seed=7) + return n_mixtures, mc_like, mc_se + + +def test_linear_filter_likelihood_matches_monte_carlo_integral(mc_reference) -> None: + """The exact marginal likelihood equals the MC average of per-draw likelihoods.""" + n_mixtures, mc_like, mc_se = mc_reference + model = _model(n_mixtures, moments=MOMENTS) + exact = np.exp(_loglikeobs(model, _data())) + assert np.all(np.abs(exact - mc_like) <= 4 * mc_se + 1e-12 * mc_like) + + +def test_monte_carlo_reference_rejects_mean_imputation(mc_reference) -> None: + """Plugging in `m` without its variance is far outside the MC error band. + + This shows the MC comparison has power to detect a missing variance term. + """ + n_mixtures, mc_like, mc_se = mc_reference + data = _data() + for ob, (mean_col, _) in MOMENTS.items(): + data[ob] = data[ob].fillna(data[mean_col]) + naive = np.exp(_loglikeobs(_model(n_mixtures), data)) + affected = mc_se > 0 + assert np.all(np.abs(naive - mc_like)[affected] > 10 * mc_se[affected]) + + +@pytest.mark.parametrize("n_mixtures", [1, 2]) +def test_unscented_path_matches_linear_path_for_linear_translog(n_mixtures) -> None: + """A translog with all squares/interactions at 0 reproduces the linear filter.""" + data = _data() + linear = _loglikeobs(_model(n_mixtures, "linear", MOMENTS), data) + unscented = _loglikeobs(_model(n_mixtures, "translog", MOMENTS), data) + assert_allclose(unscented, linear, rtol=1e-10, atol=1e-10) + + +@pytest.mark.parametrize("transition", ["linear", "translog"]) +def test_declared_moments_without_missings_leave_loglike_unchanged( + transition, +) -> None: + """Declaring moments for fully observed factors does not change the filter.""" + data = _data(with_missings=False) + plain = _loglikeobs(_model(2, transition), data, nonlinear=0.05) + declared = _loglikeobs(_model(2, transition, MOMENTS), data, nonlinear=0.05) + assert np.array_equal(plain, declared) + + +@pytest.mark.parametrize("transition", ["linear", "translog"]) +def test_loglike_without_missings_matches_baseline(transition) -> None: + """Models with complete observed factors keep their pre-feature likelihood.""" + data = _data(with_missings=False) + loglike = _loglikeobs(_model(2, transition), data, nonlinear=0.05).sum() + assert_allclose(loglike, BASELINE_LOGLIKE[transition], rtol=1e-12) + + +def test_moments_where_factor_is_observed_are_ignored() -> None: + """Fill moments only matter in cells where the observed factor is missing.""" + data = _data() + garbled = data.copy() + for ob, (mean_col, var_col) in MOMENTS.items(): + observed = garbled[ob].notna() + garbled.loc[observed, mean_col] = np.nan + garbled.loc[observed, var_col] = -1.0 + model = _model(1, moments=MOMENTS) + assert np.array_equal(_loglikeobs(model, data), _loglikeobs(model, garbled)) + + +@pytest.mark.parametrize("bad_col", ["ob1_mean", "ob1_var"]) +def test_missing_factor_without_moments_raises(bad_col) -> None: + data = _data() + data.loc[(1, 1), bad_col] = np.nan + with pytest.raises(ValueError, match=r"ob1[\s\S]*fill moments") as excinfo: + get_maximization_inputs(_model(1, moments=MOMENTS), data, chs_options=OPTIONS) + assert "1" in str(excinfo.value) + + +def test_negative_fill_variance_raises() -> None: + data = _data() + data.loc[(1, 1), "ob1_var"] = -0.5 + with pytest.raises(ValueError, match=r"ob1[\s\S]*negative fill variance"): + get_maximization_inputs(_model(1, moments=MOMENTS), data, chs_options=OPTIONS) + + +def test_missing_moment_column_raises() -> None: + data = _data().drop(columns=["ob2_var"]) + with pytest.raises(ValueError, match=r"ob2_var[\s\S]*Variable is missing"): + get_maximization_inputs(_model(1, moments=MOMENTS), data, chs_options=OPTIONS) + + +def test_missing_factor_without_declared_moments_raises() -> None: + moments = {"ob1": MOMENTS["ob1"]} + with pytest.raises(ValueError, match=r"ob2[\s\S]*Variable has missings"): + get_maximization_inputs( + _model(1, moments=moments), _data(), chs_options=OPTIONS + ) + + +def test_model_spec_rejects_moments_for_unknown_observed_factor() -> None: + with pytest.raises(ValueError, match="not_an_observed_factor"): + _model(1, moments={"not_an_observed_factor": ("a", "b")}) + + +def test_model_spec_from_dict_reads_moments() -> None: + spec = ModelSpec.from_dict( + { + "factors": {"fac1": {"measurements": [["y1"], ["y1"]]}}, + "observed_factors": ["ob1"], + "missing_observed_factor_moments": {"ob1": ["ob1_mean", "ob1_var"]}, + } + ) + assert dict(spec.missing_observed_factor_moments) == { + "ob1": ("ob1_mean", "ob1_var") + } + + +def _predict_inputs(): + """Inputs for one predict step: 3 persons, 2 mixtures, 2 latent, 2 observed. + + Person 0 misses ob1, person 1 misses both, person 2 misses none. + """ + rng = np.random.default_rng(0) + n_obs, n_mix, n_lat = 3, 2, 2 + states = rng.normal(size=(n_obs, n_mix, n_lat)) + covs = np.empty((n_obs, n_mix, n_lat, n_lat)) + chols = np.empty_like(covs) + for i in range(n_obs): + for k in range(n_mix): + a = rng.normal(size=(n_lat, n_lat)) + covs[i, k] = a @ a.T + np.eye(n_lat) + chols[i, k] = np.linalg.cholesky(covs[i, k]).T + # Means are already substituted where the factor is missing. + observed = rng.normal(size=(n_obs, 2)) + obs_vars = np.array([[0.7, 0.0], [0.4, 1.3], [0.0, 0.0]]) + f_obs = np.array([[0.9, -0.6], [0.5, 0.8]]) + f_lat = np.array([[0.7, 0.2], [-0.3, 0.6]]) + const = np.array([0.1, -0.2]) + shock_sds = np.array([0.4, 0.3]) + scaling = np.array([[1.5, 0.8, 1.0, 1.0], [1.2, 0.9, 1.0, 1.0]]) + constants = np.array([[0.3, -0.1, 0.0, 0.0], [0.2, 0.4, 0.0, 0.0]]) + return { + "states": states, + "covs": covs, + "chols": chols, + "observed": observed, + "obs_vars": obs_vars, + "f_lat": f_lat, + "f_obs": f_obs, + "const": const, + "shock_sds": shock_sds, + "scaling": scaling, + "constants": constants, + } + + +def _analytic_prediction(inp): + """Predicted mean and covariance with `x ~ N(m, diag(v))` integrated out.""" + s_in, s_out = inp["scaling"][0, :2], inp["scaling"][1, :2] + c_in, c_out = inp["constants"][0, :2], inp["constants"][1, :2] + n_obs, n_mix, _ = inp["states"].shape + means = np.empty_like(inp["states"]) + covs = np.empty_like(inp["covs"]) + for i in range(n_obs): + for k in range(n_mix): + anchored = inp["states"][i, k] * s_in + c_in + pred = inp["f_lat"] @ anchored + inp["f_obs"] @ inp["observed"][i] + means[i, k] = (pred + inp["const"] - c_out) / s_out + g_mat = np.diag(1 / s_out) @ inp["f_lat"] @ np.diag(s_in) + b_mat = np.diag(1 / s_out) @ inp["f_obs"] + covs[i, k] = ( + g_mat @ inp["covs"][i, k] @ g_mat.T + + np.diag((inp["shock_sds"] / s_out) ** 2) + + b_mat @ np.diag(inp["obs_vars"][i]) @ b_mat.T + ) + return means, covs + + +def _linear_trans_coeffs(inp): + return { + f: jnp.array([*inp["f_lat"][j], *inp["f_obs"][j], inp["const"][j]]) + for j, f in enumerate(LATENT) + } + + +def test_linear_predict_integrates_out_missing_observed_factors() -> None: + inp = _predict_inputs() + exp_means, exp_covs = _analytic_prediction(inp) + means, chols = linear_kalman_predict( + None, + jnp.array(inp["states"]), + jnp.array(inp["chols"]), + 1.0, + jnp.ones(1), + _linear_trans_coeffs(inp), + jnp.array(inp["shock_sds"]), + jnp.array(inp["scaling"]), + jnp.array(inp["constants"]), + jnp.array(inp["observed"]), + observed_factor_vars=jnp.array(inp["obs_vars"]), + latent_factors=LATENT, + constant_factor_indices=frozenset(), + n_all_factors=4, + missing_observed_indices=(0, 1), + ) + covs = np.einsum("nkja,nkjb->nkab", chols, chols) + assert_allclose(means, exp_means, rtol=1e-12) + assert_allclose(covs, exp_covs, rtol=1e-12) + + +def _linear_transition(params, states): + return jnp.column_stack([states @ params[f][:-1] + params[f][-1] for f in LATENT]) + + +def test_unscented_predict_integrates_out_missing_observed_factors() -> None: + inp = _predict_inputs() + exp_means, exp_covs = _analytic_prediction(inp) + scaling_factor, weights = calculate_sigma_scaling_factor_and_weights(4, 2) + # The unscented predict adds shocks in the unanchored output scale. + shock_sds = inp["shock_sds"] / inp["scaling"][1, :2] + means, chols = kalman_predict( + _linear_transition, + jnp.array(inp["states"]), + jnp.array(inp["chols"]), + scaling_factor, + weights, + _linear_trans_coeffs(inp), + jnp.array(shock_sds), + jnp.array(inp["scaling"]), + jnp.array(inp["constants"]), + jnp.array(inp["observed"]), + observed_factor_vars=jnp.array(inp["obs_vars"]), + missing_observed_indices=(0, 1), + ) + covs = np.einsum("nkja,nkjb->nkab", chols, chols) + assert_allclose(means, exp_means, rtol=1e-10) + assert_allclose(covs, exp_covs, rtol=1e-10) + + +def _nonlinear_transition(params, states): + lat = states[:, :2] + out = _linear_transition(params, states) + return out + 0.2 * lat**2 - 0.1 * lat[:, ::-1] * states[:, 2:3] + + +def test_unscented_predict_augmented_by_observed_cells_equals_larger_kappa() -> None: + """With zero fill variance, the augmented sigma points act like `kappa + n_aug`. + + The latent marginal of the augmented sigma-point set is the plain set built + with `kappa` raised by the number of augmented dimensions. + """ + inp = _predict_inputs() + common = ( + _nonlinear_transition, + jnp.array(inp["states"]), + jnp.array(inp["chols"]), + ) + rest = ( + _linear_trans_coeffs(inp), + jnp.array(inp["shock_sds"]), + jnp.array(inp["scaling"]), + jnp.array(inp["constants"]), + jnp.array(inp["observed"]), + ) + aug_scaling, aug_weights = calculate_sigma_scaling_factor_and_weights(4, 2) + aug_means, aug_chols = kalman_predict( + *common, + aug_scaling, + aug_weights, + *rest, + observed_factor_vars=jnp.zeros((3, 2)), + missing_observed_indices=(0, 1), + ) + plain_scaling, plain_weights = calculate_sigma_scaling_factor_and_weights(2, 4) + plain_means, plain_chols = kalman_predict( + *common, plain_scaling, plain_weights, *rest + ) + aug_covs = np.einsum("nkja,nkjb->nkab", aug_chols, aug_chols) + plain_covs = np.einsum("nkja,nkjb->nkab", plain_chols, plain_chols) + assert_allclose(aug_means, plain_means, rtol=1e-12) + assert_allclose(aug_covs, plain_covs, rtol=1e-12) From ef557a751b1dabe175539e6539903618dde252e1 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Wed, 14 Oct 2026 21:52:53 +0200 Subject: [PATCH 10/18] Require chs_options in get_maximization_inputs and the CHS state filters `get_maximization_inputs` no longer falls back to `CHSEstimationOptions()`. That default seeds the template with a full AMN estimation, so callers that only wanted the likelihood functions silently ran AMN; the CHS state filter spent hours in AMN mixture EM to compute filtered states. - `get_maximization_inputs(..., *, chs_options)` is now required. - `get_individual_states_from_params(..., *, chs_options)` is required and forwarded with `start_params_strategy="none"`, so the filter uses the estimation's likelihood settings (e.g. the sigma-point scale) and never seeds start values. - `get_individual_states` takes `chs_options` and requires it for a `CHSEstimationResult` (unused for AF and AMN results). - Tests and docs notebooks pass options explicitly; new tests check that the filter never seeds, forwards the sigma-point scale, and rejects a CHS result without options. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- docs/getting_started/tutorial.ipynb | 2 +- ...sualize_pairwise_factor_distribution.ipynb | 4 +- ...ow_to_visualize_transition_equations.ipynb | 4 +- src/skillmodels/chs/maximization_inputs.py | 10 ++- src/skillmodels/common/individual_states.py | 25 +++++- tests/test_filtered_states.py | 76 ++++++++++++++++++- tests/test_visualize_factor_distributions.py | 2 +- tests/test_visualize_transition_equations.py | 7 +- 8 files changed, 115 insertions(+), 15 deletions(-) diff --git a/docs/getting_started/tutorial.ipynb b/docs/getting_started/tutorial.ipynb index 46082e77..9a9ef30c 100644 --- a/docs/getting_started/tutorial.ipynb +++ b/docs/getting_started/tutorial.ipynb @@ -534,7 +534,7 @@ ], "source": [ "chs_filtered = get_individual_states_from_params(\n", - " model_spec=model, data=data, params=chs_params\n", + " model_spec=model, data=data, params=chs_params, chs_options=chs_options\n", ")\n", "chs_states = chs_filtered[\"unanchored_states\"][\"states\"]\n", "chs_decomp = decompose_measurement_variance(\n", diff --git a/docs/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb b/docs/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb index 56fc9a95..e950d71e 100644 --- a/docs/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb +++ b/docs/how_to_guides/how_to_visualize_pairwise_factor_distribution.ipynb @@ -33,7 +33,7 @@ "id": "2", "metadata": {}, "outputs": [], - "source": "model = MODEL2\nparams = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\nparams = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n\ndata = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\ndata = data.set_index([\"caseid\", \"period\"])\n\nfiltered = get_individual_states_from_params(model_spec=model, data=data, params=params)\nstates = filtered[\"anchored_states\"][\"states\"]" + "source": "model = MODEL2\nparams = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\nparams = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n\ndata = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\ndata = data.set_index([\"caseid\", \"period\"])\n\nfiltered = get_individual_states_from_params(\n model_spec=model, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS\n)\nstates = filtered[\"anchored_states\"][\"states\"]" }, { "cell_type": "markdown", @@ -254,7 +254,7 @@ "id": "20", "metadata": {}, "outputs": [], - "source": "params = get_maximization_inputs(\n model_spec=model, data=data, chs_options=MODEL2_CHS_OPTIONS\n)[\"params_template\"]\nparams[\"value\"] = 0.1\nfiltered = get_individual_states_from_params(model_spec=model, data=data, params=params)\nstates = filtered[\"anchored_states\"][\"states\"]" + "source": "params = get_maximization_inputs(\n model_spec=model, data=data, chs_options=MODEL2_CHS_OPTIONS\n)[\"params_template\"]\nparams[\"value\"] = 0.1\nfiltered = get_individual_states_from_params(\n model_spec=model, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS\n)\nstates = filtered[\"anchored_states\"][\"states\"]" }, { "cell_type": "code", diff --git a/docs/how_to_guides/how_to_visualize_transition_equations.ipynb b/docs/how_to_guides/how_to_visualize_transition_equations.ipynb index 1225f49c..96ee44c7 100644 --- a/docs/how_to_guides/how_to_visualize_transition_equations.ipynb +++ b/docs/how_to_guides/how_to_visualize_transition_equations.ipynb @@ -6,7 +6,7 @@ "id": "0", "metadata": {}, "outputs": [], - "source": "import pandas as pd\n\nfrom skillmodels.common.config import REGRESSION_VAULT, TEST_DATA_DIR\nfrom skillmodels.common.individual_states import get_individual_states_from_params\nfrom skillmodels.common.visualize_transition_equations import (\n combine_transition_plots,\n get_transition_plots,\n)\nfrom skillmodels.test_data.model2 import MODEL2" + "source": "import pandas as pd\n\nfrom skillmodels.common.config import REGRESSION_VAULT, TEST_DATA_DIR\nfrom skillmodels.common.individual_states import get_individual_states_from_params\nfrom skillmodels.common.visualize_transition_equations import (\n combine_transition_plots,\n get_transition_plots,\n)\nfrom skillmodels.test_data.model2 import MODEL2, MODEL2_CHS_OPTIONS" }, { "cell_type": "markdown", @@ -37,7 +37,7 @@ "id": "2", "metadata": {}, "outputs": [], - "source": "model = MODEL2\n\nparams = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\nparams = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n\ndata = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\ndata = data.set_index([\"caseid\", \"period\"])\n\nfiltered = get_individual_states_from_params(model_spec=model, data=data, params=params)\nstates = filtered[\"anchored_states\"][\"states\"]" + "source": "model = MODEL2\n\nparams = pd.read_csv(REGRESSION_VAULT / \"one_stage_anchoring.csv\")\nparams = params.set_index([\"category\", \"period\", \"name1\", \"name2\"])\n\ndata = pd.read_stata(TEST_DATA_DIR / \"model2_simulated_data.dta\")\ndata = data.set_index([\"caseid\", \"period\"])\n\nfiltered = get_individual_states_from_params(\n model_spec=model, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS\n)\nstates = filtered[\"anchored_states\"][\"states\"]" }, { "cell_type": "markdown", diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index e46fa9af..5e34b227 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -57,7 +57,7 @@ def get_maximization_inputs( # noqa: C901, PLR0915 data: pd.DataFrame, split_dataset: int = 1, *, - chs_options: CHSEstimationOptions | None = None, + chs_options: CHSEstimationOptions, fixed_params: pd.DataFrame | None = None, ) -> dict[str, Any]: """Create inputs for optimagic's maximize function. @@ -67,8 +67,11 @@ def get_maximization_inputs( # noqa: C901, PLR0915 data: Dataset in long format. split_dataset: Controls into how many slices to split the dataset during the gradient computation. - chs_options: CHS-specific tuning parameters. Defaults to - ``CHSEstimationOptions()`` when not provided. + chs_options: CHS-specific tuning parameters. Its `start_params_strategy` + decides how the `params_template` values are seeded; pass + `CHSEstimationOptions(start_params_strategy="none")` when only the + likelihood functions or the bare template are needed, since the + `"amn"` strategy runs a full AMN estimation. fixed_params: Optional DataFrame with a ``"value"`` column pinning specified parameters to fixed values. Uses the same 4-level MultiIndex as the returned ``params_template``. Each matching @@ -103,7 +106,6 @@ def get_maximization_inputs( # noqa: C901, PLR0915 """ _fail_if_non_gaussian_measurements(model_spec) - chs_options = chs_options or CHSEstimationOptions() processed_model = process_model(model_spec) p_index = get_params_index( update_info=processed_model.update_info, diff --git a/src/skillmodels/common/individual_states.py b/src/skillmodels/common/individual_states.py index 9d073c04..ee884fff 100644 --- a/src/skillmodels/common/individual_states.py +++ b/src/skillmodels/common/individual_states.py @@ -10,6 +10,7 @@ result machinery, so there is no raw-params path for them). """ +from dataclasses import replace from typing import Any import pandas as pd @@ -23,6 +24,7 @@ from skillmodels.af.types import AFEstimationResult from skillmodels.amn.types import AMNEstimationResult from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.chs.options import CHSEstimationOptions from skillmodels.chs.types import CHSEstimationResult from skillmodels.common.anchoring import anchor_states_df from skillmodels.common.model_spec import ModelSpec @@ -34,6 +36,8 @@ def get_individual_states( data: pd.DataFrame, result: CHSEstimationResult | AFEstimationResult | AMNEstimationResult, + *, + chs_options: CHSEstimationOptions | None = None, ) -> dict[str, dict[str, Any]]: """Compute per-individual latent state estimates from a fitted result. @@ -47,6 +51,9 @@ def get_individual_states( data: Dataset in long format with MultiIndex (id, period). result: A `CHSEstimationResult`, `AFEstimationResult`, or `AMNEstimationResult`. + chs_options: The options the CHS result was estimated with, so the + filter uses the same likelihood settings. Required for a + `CHSEstimationResult`; unused for AF and AMN. Return: Nested dict. The return shape is **asymmetric** across estimators: @@ -78,10 +85,17 @@ def get_individual_states( data=data, ) + if chs_options is None: + msg = ( + "get_individual_states needs `chs_options` for a CHSEstimationResult: " + "pass the CHSEstimationOptions used for the estimation." + ) + raise TypeError(msg) return get_individual_states_from_params( model_spec=result.model_spec, data=data, params=result.params, + chs_options=chs_options, ) @@ -90,6 +104,8 @@ def get_individual_states_from_params( model_spec: ModelSpec, data: pd.DataFrame, params: pd.DataFrame, + *, + chs_options: CHSEstimationOptions, ) -> dict[str, dict[str, Any]]: """Compute CHS filtered states directly from raw parameters. @@ -102,13 +118,20 @@ def get_individual_states_from_params( model_spec: Model specification. data: Dataset in long format with MultiIndex (id, period). params: Estimated parameter DataFrame (standard 4-level MultiIndex). + chs_options: The options used for the estimation, so the filter uses the + same likelihood settings (e.g. the sigma-point scale). Its + `start_params_strategy` is ignored: no start values are seeded. Return: Dict with both `"anchored_states"` and `"unanchored_states"`, each `{"states": DataFrame, "state_ranges": ...}`. """ - max_inputs = get_maximization_inputs(model_spec=model_spec, data=data) + max_inputs = get_maximization_inputs( + model_spec=model_spec, + data=data, + chs_options=replace(chs_options, start_params_strategy="none"), + ) params = params.loc[max_inputs["params_template"].index] debug_loglike = max_inputs["debug_loglike"] debug_data = debug_loglike(params) diff --git a/tests/test_filtered_states.py b/tests/test_filtered_states.py index f8024f44..58369711 100644 --- a/tests/test_filtered_states.py +++ b/tests/test_filtered_states.py @@ -6,9 +6,15 @@ import pandas as pd import pytest +import skillmodels.chs.maximization_inputs as mi from skillmodels.chs.maximization_inputs import get_maximization_inputs +from skillmodels.chs.options import CHSEstimationOptions +from skillmodels.chs.types import CHSEstimationResult from skillmodels.common.config import TEST_DATA_DIR -from skillmodels.common.individual_states import get_individual_states_from_params +from skillmodels.common.individual_states import ( + get_individual_states, + get_individual_states_from_params, +) from skillmodels.test_data.model2 import MODEL2, MODEL2_CHS_OPTIONS REGRESSION_VAULT = Path(__file__).parent / "regression_vault" @@ -35,7 +41,10 @@ def test_get_individual_states_from_params(model2, model2_data) -> None: params = params.loc[max_inputs["params_template"].index] calculated = get_individual_states_from_params( - model_spec=model2, data=model2_data, params=params + model_spec=model2, + data=model2_data, + params=params, + chs_options=MODEL2_CHS_OPTIONS, ) factors = ["fac1", "fac2", "fac3"] @@ -45,3 +54,66 @@ def test_get_individual_states_from_params(model2, model2_data) -> None: unanch_ranges = calculated["unanchored_states"]["state_ranges"][factor] ratio = (anch_ranges / unanch_ranges).to_numpy() assert np.allclose(ratio, expected_ratio) + + +@pytest.fixture +def model2_params(model2, model2_data): + params = pd.read_csv(REGRESSION_VAULT / "one_stage_anchoring.csv") + params = params.set_index(["category", "period", "name1", "name2"]) + template = get_maximization_inputs( + model2, + model2_data, + chs_options=CHSEstimationOptions(start_params_strategy="none"), + )["params_template"] + return params.loc[template.index] + + +def _fail(*args, **kwargs): + raise AssertionError("start values must not be seeded") + + +def test_individual_states_from_params_never_seed_start_values( + model2, model2_data, model2_params, monkeypatch +) -> None: + """Filtering at given params skips AMN / Spearman start-value seeding.""" + monkeypatch.setattr(mi, "_estimate_amn_for_chs_seeding", _fail) + monkeypatch.setattr(mi, "get_spearman_start_params", _fail) + states = get_individual_states_from_params( + model_spec=model2, + data=model2_data, + params=model2_params, + chs_options=CHSEstimationOptions(start_params_strategy="amn"), + )["unanchored_states"]["states"] + assert len(states) == len(model2_data) + + +def test_individual_states_from_params_use_the_sigma_point_scale( + model2, model2_data, model2_params +) -> None: + """The filter runs with the likelihood settings of the given options.""" + + def _states(scale): + return get_individual_states_from_params( + model_spec=model2, + data=model2_data, + params=model2_params, + chs_options=CHSEstimationOptions( + start_params_strategy="none", sigma_points_scale=scale + ), + )["unanchored_states"]["states"]["fac1"].to_numpy() + + assert not np.allclose(_states(2.0), _states(0.5)) + + +def test_get_individual_states_requires_chs_options_for_chs_result( + model2, model2_data, model2_params +) -> None: + result = CHSEstimationResult( + model_spec=model2, + params=model2_params, + success=True, + loglikelihood=0.0, + optimize_result=None, + ) + with pytest.raises(TypeError, match="chs_options"): + get_individual_states(data=model2_data, result=result) diff --git a/tests/test_visualize_factor_distributions.py b/tests/test_visualize_factor_distributions.py index e021dff2..d181372a 100644 --- a/tests/test_visualize_factor_distributions.py +++ b/tests/test_visualize_factor_distributions.py @@ -31,7 +31,7 @@ def _load_model2_filtered() -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: params = params.loc[max_inputs["params_template"].index] states = get_individual_states_from_params( - model_spec=MODEL2, data=data, params=params + model_spec=MODEL2, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS )["anchored_states"]["states"] return data, params, states diff --git a/tests/test_visualize_transition_equations.py b/tests/test_visualize_transition_equations.py index b0dfe37f..1880734a 100644 --- a/tests/test_visualize_transition_equations.py +++ b/tests/test_visualize_transition_equations.py @@ -35,7 +35,7 @@ def test_visualize_transition_equations_runs() -> None: params["value"] = params["value"].fillna(0) states = get_individual_states_from_params( - model_spec=model, data=data, params=params + model_spec=model, data=data, params=params, chs_options=MODEL2_CHS_OPTIONS )["anchored_states"]["states"] subplots = get_transition_plots( @@ -143,7 +143,10 @@ def test_get_transition_plots_runs_for_correction_model() -> None: assert not template["value"].isna().any() states = get_individual_states_from_params( - model_spec=model, data=panel, params=template + model_spec=model, + data=panel, + params=template, + chs_options=CHSEstimationOptions(), )["anchored_states"]["states"] plots = get_transition_plots( From c162965939be171f8a60e4b5b4fae50942e4ebbe Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Wed, 14 Oct 2026 23:26:35 +0200 Subject: [PATCH 11/18] Move split_dataset onto CHSEstimationOptions so estimate_chs can use it estimate_chs never forwarded the gradient split, so every estimation computed the full-sample gradient at once. Models with many possibly-missing observed factors (augmented sigma points) need several hundred GB for it. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/maximization_inputs.py | 4 +--- src/skillmodels/chs/options.py | 4 ++++ tests/test_likelihood_regression.py | 4 +++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index 5e34b227..a89b402c 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -55,7 +55,6 @@ def get_maximization_inputs( # noqa: C901, PLR0915 model_spec: ModelSpec, data: pd.DataFrame, - split_dataset: int = 1, *, chs_options: CHSEstimationOptions, fixed_params: pd.DataFrame | None = None, @@ -65,8 +64,6 @@ def get_maximization_inputs( # noqa: C901, PLR0915 Args: model_spec: The model specification. See: :ref:`model_specs` data: Dataset in long format. - split_dataset: Controls into how many slices to split the dataset - during the gradient computation. chs_options: CHS-specific tuning parameters. Its `start_params_strategy` decides how the `params_template` values are seeded; pass `CHSEstimationOptions(start_params_strategy="none")` when only the @@ -183,6 +180,7 @@ def loglike_and_gradient( n_obs = processed_data["measurements"].shape[1] _grad = jnp.zeros_like(params_vec) start = 0 + split_dataset = chs_options.split_dataset stop = int(n_obs / split_dataset) step = int(n_obs / split_dataset) for i in range(split_dataset): diff --git a/src/skillmodels/chs/options.py b/src/skillmodels/chs/options.py index 0b3bbb24..f8058df1 100644 --- a/src/skillmodels/chs/options.py +++ b/src/skillmodels/chs/options.py @@ -20,6 +20,10 @@ class CHSEstimationOptions: """Distance for bounds. Zeroed out if `robust_bounds` is False.""" sigma_points_scale: float = 2 """Scaling factor for sigma points in unscented transform.""" + split_dataset: int = 1 + """Number of slices of the individuals over which the gradient is accumulated. + + More slices lower the peak memory of the gradient; the result is the same.""" clipping_lower_bound: float = -1e30 """Lower bound for soft clipping.""" clipping_upper_bound: float | None = None diff --git a/tests/test_likelihood_regression.py b/tests/test_likelihood_regression.py index 791d3c02..5a2da600 100644 --- a/tests/test_likelihood_regression.py +++ b/tests/test_likelihood_regression.py @@ -103,7 +103,9 @@ def test_splitting_does_not_change_gradient(model2, model2_data) -> None: model2, model2_data, chs_options=MODEL2_CHS_OPTIONS ) inputs_split = get_maximization_inputs( - model2, model2_data, 13, chs_options=MODEL2_CHS_OPTIONS + model2, + model2_data, + chs_options=replace(MODEL2_CHS_OPTIONS, split_dataset=13), ) params = inputs["params_template"] From 69bae8687d901f67a083783d1934d56ce1a8f827 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Thu, 15 Oct 2026 07:59:47 +0200 Subject: [PATCH 12/18] Missing observed factors: endogenous guard, root precision, shared columns Repairs from the round-0 review of integrating out missing observed factors: - Endogenous CHS models reject positive-variance missing inputs in periods with an outgoing transition: both half-steps would need the input's joint posterior. Zero-variance fills and complete inputs stay supported. - Promote fill variances to the covariance dtype before the square root, so float32-stored moments meet float64 precision on both predict paths. - Select each physical data column once when moment columns are shared across factors or alias another selected column. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/kalman_filters.py | 8 +- src/skillmodels/chs/maximization_inputs.py | 16 ++++ src/skillmodels/common/model_spec.py | 7 +- src/skillmodels/common/process_data.py | 29 ++++--- tests/test_missing_observed_factors.py | 93 ++++++++++++++++++++-- 5 files changed, 133 insertions(+), 20 deletions(-) diff --git a/src/skillmodels/chs/kalman_filters.py b/src/skillmodels/chs/kalman_filters.py index 3ba5d91e..ad47fd6a 100644 --- a/src/skillmodels/chs/kalman_filters.py +++ b/src/skillmodels/chs/kalman_filters.py @@ -406,7 +406,10 @@ def linear_kalman_predict( # Row j is sqrt(v_j) * B[:, j], so the rows add B diag(v) B' to the cov. b_missing = f_mat[:, n_latent + jnp.array(missing_observed_indices)] b_missing = b_missing / s_out[:, jnp.newaxis] # (n_latent, n_missing) - rows = jnp.sqrt(observed_factor_vars)[..., jnp.newaxis] * b_missing.T + # Promote before the root: loader moments are float32, while the + # covariance factor can be float64. A later cast cannot undo root rounding. + missing_sds = jnp.sqrt(observed_factor_vars.astype(chol_g.dtype)) + rows = missing_sds[..., jnp.newaxis] * b_missing.T blocks.append( jnp.broadcast_to(rows[:, jnp.newaxis], (n_obs, n_mix, *b_missing.T.shape)) ) @@ -560,7 +563,8 @@ def _augment_by_missing_observed_factors( [states, jnp.broadcast_to(means[:, None], (n_obs, n_mixtures, n_missing))], axis=-1, ) - sds = jnp.sqrt(observed_factor_vars) # (n_obs, n_missing) + # Preserve the covariance-factor precision during the root, not just scatter. + sds = jnp.sqrt(observed_factor_vars.astype(upper_chols.dtype)) aug_chols = jnp.zeros((n_obs, n_mixtures, n_aug, n_aug), dtype=upper_chols.dtype) aug_chols = aug_chols.at[..., :n_latent, :n_latent].set(upper_chols) diag = jnp.arange(n_latent, n_aug) diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index a89b402c..e51cc2ff 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -128,6 +128,22 @@ def get_maximization_inputs( # noqa: C901, PLR0915 purpose="estimation", ) + # A physical observed input is repeated in both endogenous half-steps. + # Marginalizing it separately in each half drops its posterior/cross-covariance. + # Until it is carried jointly, reject uncertain inputs in prediction periods. + # The final two augmented rows have no outgoing transition. + if processed_model.endogenous_factors_info.has_endogenous_factors and np.any( + np.asarray(processed_data["observed_factor_vars"][:-2]) > 0 + ): + msg = ( + "CHS endogenous half-steps require a " + "persistent joint missing-observed-factor state for " + "positive-variance missing inputs. " + "This combination is not yet supported; zero-variance fills and " + "complete observed factors remain supported." + ) + raise NotImplementedError(msg) + # Possibly-missing observed factors augment the sigma-point input. sigma_scaling_factor, sigma_weights = calculate_sigma_scaling_factor_and_weights( n_states=processed_model.dimensions.n_latent_factors diff --git a/src/skillmodels/common/model_spec.py b/src/skillmodels/common/model_spec.py index b90cd789..993db665 100644 --- a/src/skillmodels/common/model_spec.py +++ b/src/skillmodels/common/model_spec.py @@ -223,7 +223,12 @@ class ModelSpec: An observed factor not listed here must never be missing. Only the CHS likelihood (and the filtered states built on it) integrates missing values out; simulation and transition plots still require complete observed - factors.""" + factors. In CHS models with endogenous factors, positive-variance missing + inputs in a physical period with an outgoing transition are not yet + supported: both half-steps would need to carry the same missing input's + joint posterior. Such inputs raise `NotImplementedError`; zero-variance + fills, complete inputs, and terminal-only missing inputs remain supported. + """ def __init__( self, diff --git a/src/skillmodels/common/process_data.py b/src/skillmodels/common/process_data.py index f1c4a075..73e955da 100644 --- a/src/skillmodels/common/process_data.py +++ b/src/skillmodels/common/process_data.py @@ -143,18 +143,23 @@ def _get_period_data_for_endogenous_factors( col for cols in labels.missing_observed_factor_moments.values() for col in cols ] - out = df.query(f"period == {period}")[ - [ - "id", - *meas, - *controls, - *observed, - *moment_cols, - "period", - "__old_id__", - "__old_period__", - ] - ] + # Several factor mappings may share a physical moment column, or reuse a + # column already selected above. Preserve its first position, not duplicates. + columns = list( + dict.fromkeys( + [ + "id", + *meas, + *controls, + *observed, + *moment_cols, + "period", + "__old_id__", + "__old_period__", + ] + ) + ) + out = df.query(f"period == {period}")[columns] out["aug_period"] = aug_period return out diff --git a/tests/test_missing_observed_factors.py b/tests/test_missing_observed_factors.py index 8f627a5d..2751823a 100644 --- a/tests/test_missing_observed_factors.py +++ b/tests/test_missing_observed_factors.py @@ -220,6 +220,82 @@ def test_loglike_without_missings_matches_baseline(transition) -> None: assert_allclose(loglike, BASELINE_LOGLIKE[transition], rtol=1e-12) +def _endogenous_model(moments: dict | None) -> ModelSpec: + """Endogenous investment `inv` and state `s` share the observed inputs.""" + return ModelSpec( + factors={ + "s": FactorSpec( + measurements=(("y",), ("y",)), transition_function="linear" + ), + "inv": FactorSpec( + measurements=(("i",), ()), + is_endogenous=True, + transition_function="linear", + ), + }, + observed_factors=("x", "z"), + missing_observed_factor_moments=moments, + ) + + +def _endogenous_data(variance: float | None = None) -> pd.DataFrame: + idx = pd.MultiIndex.from_product([[0, 1], [0, 1]], names=["id", "period"]) + data = pd.DataFrame( + {"y": [-1.0, 0.0, 1.0, 2.0], "i": np.nan, "x": 0.0, "z": 0.0}, index=idx + ) + data[["m", "v"]] = 0.0 + if variance is not None: + data.loc[(0, 0), ["x", "v"]] = [np.nan, variance] + return data + + +def _endogenous_loglike(model: ModelSpec, data: pd.DataFrame) -> float: + inputs = get_maximization_inputs(model, data, chs_options=OPTIONS) + params = inputs["params_template"].copy() + params["value"] = np.where( + params.index.get_level_values("category").isin( + ["meas_sds", "shock_sds", "loadings"] + ), + 0.5, + 0.1, + ) + for idx in params.query("category == 'initial_cholcovs'").index: + row, col = idx[3].split("-") + params.loc[idx, "value"] = float(row == col) + return float(inputs["loglike"](params)) + + +@pytest.mark.parametrize( + "moments", + [ + {"x": ("m", "v"), "z": ("m", "v")}, + {"x": ("z", "v"), "z": ("m", "v")}, + ], + ids=["shared_pair", "moment_is_observed_factor"], +) +def test_moment_columns_shared_or_aliased_leave_loglike_unchanged(moments) -> None: + """Moment columns may be shared across factors or alias a selected column.""" + data = _endogenous_data() + plain = _endogenous_loglike(_endogenous_model(None), data) + declared = _endogenous_loglike(_endogenous_model(moments), data) + assert_allclose(declared, plain, rtol=0, atol=1e-12) + + +def test_endogenous_model_rejects_uncertain_missing_input() -> None: + """Both half-steps would need the input's joint posterior, which is not carried.""" + model = _endogenous_model({"x": ("m", "v")}) + with pytest.raises(NotImplementedError, match="joint missing-observed-factor"): + get_maximization_inputs(model, _endogenous_data(1.0), chs_options=OPTIONS) + + +def test_endogenous_model_plugs_in_zero_variance_fill() -> None: + """A zero-variance fill acts like observing the fill mean.""" + model = _endogenous_model({"x": ("m", "v")}) + filled = _endogenous_loglike(model, _endogenous_data(0.0)) + plugged_in = _endogenous_loglike(model, _endogenous_data()) + assert_allclose(filled, plugged_in, rtol=0, atol=1e-12) + + def test_moments_where_factor_is_observed_are_ignored() -> None: """Fill moments only matter in cells where the observed factor is missing.""" data = _data() @@ -297,7 +373,8 @@ def _predict_inputs(): chols[i, k] = np.linalg.cholesky(covs[i, k]).T # Means are already substituted where the factor is missing. observed = rng.normal(size=(n_obs, 2)) - obs_vars = np.array([[0.7, 0.0], [0.4, 1.3], [0.0, 0.0]]) + # float32-representable variances with irrational roots. + obs_vars = np.array([[0.75, 0.0], [0.375, 1.25], [0.0, 0.0]]) f_obs = np.array([[0.9, -0.6], [0.5, 0.8]]) f_lat = np.array([[0.7, 0.2], [-0.3, 0.6]]) const = np.array([0.1, -0.2]) @@ -348,7 +425,10 @@ def _linear_trans_coeffs(inp): } -def test_linear_predict_integrates_out_missing_observed_factors() -> None: +@pytest.mark.parametrize("variance_dtype", [jnp.float64, jnp.float32]) +def test_linear_predict_integrates_out_missing_observed_factors( + variance_dtype, +) -> None: inp = _predict_inputs() exp_means, exp_covs = _analytic_prediction(inp) means, chols = linear_kalman_predict( @@ -362,7 +442,7 @@ def test_linear_predict_integrates_out_missing_observed_factors() -> None: jnp.array(inp["scaling"]), jnp.array(inp["constants"]), jnp.array(inp["observed"]), - observed_factor_vars=jnp.array(inp["obs_vars"]), + observed_factor_vars=jnp.array(inp["obs_vars"], dtype=variance_dtype), latent_factors=LATENT, constant_factor_indices=frozenset(), n_all_factors=4, @@ -377,7 +457,10 @@ def _linear_transition(params, states): return jnp.column_stack([states @ params[f][:-1] + params[f][-1] for f in LATENT]) -def test_unscented_predict_integrates_out_missing_observed_factors() -> None: +@pytest.mark.parametrize("variance_dtype", [jnp.float64, jnp.float32]) +def test_unscented_predict_integrates_out_missing_observed_factors( + variance_dtype, +) -> None: inp = _predict_inputs() exp_means, exp_covs = _analytic_prediction(inp) scaling_factor, weights = calculate_sigma_scaling_factor_and_weights(4, 2) @@ -394,7 +477,7 @@ def test_unscented_predict_integrates_out_missing_observed_factors() -> None: jnp.array(inp["scaling"]), jnp.array(inp["constants"]), jnp.array(inp["observed"]), - observed_factor_vars=jnp.array(inp["obs_vars"]), + observed_factor_vars=jnp.array(inp["obs_vars"], dtype=variance_dtype), missing_observed_indices=(0, 1), ) covs = np.einsum("nkja,nkjb->nkab", chols, chols) From 83074daba9acb89fa97b357a10ede939fade3cf7 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Thu, 15 Oct 2026 12:31:02 +0200 Subject: [PATCH 13/18] Transition plots: compute ranges of factors the caller did not supply Callers pass state ranges of the latent factors, but observed factors are plotted as inputs too, which raised a KeyError. Missing ranges now come from the data. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- .../common/visualize_transition_equations.py | 17 +++++++++++------ tests/test_visualize_transition_equations.py | 14 ++++++++++++++ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/src/skillmodels/common/visualize_transition_equations.py b/src/skillmodels/common/visualize_transition_equations.py index da1d4dbd..6095d002 100644 --- a/src/skillmodels/common/visualize_transition_equations.py +++ b/src/skillmodels/common/visualize_transition_equations.py @@ -546,14 +546,19 @@ def _get_state_ranges( all_factors: tuple[str, ...], quantile_cutoff: float | None = None, ) -> dict[str, pd.DataFrame]: - """Create state ranges if none is given.""" - if state_ranges is None: - state_ranges = create_state_ranges( + """Complete the given state ranges with those of the other factors in the data.""" + given = state_ranges or {} + missing = [factor for factor in all_factors if factor not in given] + if not missing: + return given + return { + **given, + **create_state_ranges( filtered_states=states_data, - factors=list(all_factors), + factors=missing, quantile_cutoff=quantile_cutoff, - ) - return state_ranges + ), + } def _get_parsed_params( diff --git a/tests/test_visualize_transition_equations.py b/tests/test_visualize_transition_equations.py index 1880734a..cb922f7f 100644 --- a/tests/test_visualize_transition_equations.py +++ b/tests/test_visualize_transition_equations.py @@ -11,6 +11,7 @@ from skillmodels.common.config import TEST_DATA_DIR from skillmodels.common.individual_states import get_individual_states_from_params from skillmodels.common.visualize_transition_equations import ( + _get_state_ranges, combine_transition_plots, get_transition_plots, ) @@ -159,3 +160,16 @@ def test_get_transition_plots_runs_for_correction_model() -> None: ) # A target factor's plot exists -> its cf-corrected transition evaluated. assert ("inv", "fac1") in plots + + +def test_get_state_ranges_adds_factors_missing_from_given_ranges() -> None: + # Callers pass ranges of the latent factors; observed factors are inputs too. + states = pd.DataFrame( + {"aug_period": [0, 0, 1, 1], "fac1": [0.0, 1.0, 2.0, 3.0], "obs": [5, 7, 6, 9]} + ) + given = {"fac1": pd.DataFrame({"minimum": [-1.0], "maximum": [1.0]})} + ranges = _get_state_ranges( + state_ranges=given, states_data=states, all_factors=("fac1", "obs") + ) + assert ranges["fac1"] is given["fac1"] + assert ranges["obs"]["maximum"].tolist() == [7, 9] From 471ab0c01bc2449c08daa8eed392f62cb1ca7e3b Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 16 Oct 2026 07:26:11 +0200 Subject: [PATCH 14/18] Measurement correlations: one column per measurement loading on several factors Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/common/correlation_heatmap.py | 2 + tests/test_correlation_heatmap.py | 57 +++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/src/skillmodels/common/correlation_heatmap.py b/src/skillmodels/common/correlation_heatmap.py index 9225e4b0..595127b8 100644 --- a/src/skillmodels/common/correlation_heatmap.py +++ b/src/skillmodels/common/correlation_heatmap.py @@ -411,6 +411,8 @@ def _get_measurement_data_for_single_period( )["variable"].to_list() for fac in observed_factors: measurements.append(fac) + # A measurement that loads on several factors is one column. + measurements = list(dict.fromkeys(measurements)) return data.query(f"{update_info_by_period.index.names[0]}=={period}")[measurements] diff --git a/tests/test_correlation_heatmap.py b/tests/test_correlation_heatmap.py index 8251b5f4..100e0461 100644 --- a/tests/test_correlation_heatmap.py +++ b/tests/test_correlation_heatmap.py @@ -464,3 +464,60 @@ def test_get_scores_corr_single_period( periods=0, ) assert isinstance(result, pd.DataFrame) + + +def test_get_measurement_data_with_single_period_lists_a_shared_measurement_once() -> ( + None +): + """A measurement that loads on two factors appears as one column.""" + update_info = pd.DataFrame( + { + "f1": [True, True, False], + "f2": [False, True, True], + "purpose": "measurement", + }, + index=pd.MultiIndex.from_product( + [[0], ["y0", "y1", "y2"]], names=["stage", "variable"] + ), + ) + data = pd.DataFrame({"stage": [0, 0], "y0": [1, 2], "y1": [3, 4], "y2": [5, 6]}) + result = _get_measurement_data_for_single_period( + data, update_info, 0, latent_factors=["f1", "f2"], observed_factors=[] + ) + assert list(result.columns) == ["y0", "y1", "y2"] + + +def test_get_measurement_data_with_multiple_periods_lists_shared_measurement_once() -> ( + None +): + """A measurement that loads on two factors appears once per period.""" + update_info = pd.DataFrame( + { + "f1": [True, True, False] * 2, + "f2": [False, True, True] * 2, + "purpose": "measurement", + }, + index=pd.MultiIndex.from_product( + [[0, 1], ["y0", "y1", "y2"]], names=["stage", "variable"] + ), + ) + data = pd.DataFrame( + { + "stage": [0, 0, 1, 1], + "y0": [1, 2, 3, 4], + "y1": [3, 4, 5, 6], + "y2": [5, 6, 7, 8], + }, + index=[0, 1, 0, 1], + ) + result = _get_measurement_data_for_multiple_periods( + data, update_info, [0, 1], latent_factors=["f1", "f2"], observed_factors=[] + ) + assert list(result.columns) == [ + "y0, 0", + "y1, 0", + "y2, 0", + "y0, 1", + "y1, 1", + "y2, 1", + ] From 7eec29773be82eb34cf4a9a71c9d9b46da16a930 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 16 Oct 2026 10:28:28 +0200 Subject: [PATCH 15/18] Factor-distribution contours: put the first factor on the x axis `_calculate_kde_for_3d` returns `z[i, j]` at `(x[i], y[j])` (np.mgrid), while `go.Contour` reads rows as y. With equal grid sizes the transposed density plotted without error, so every pairwise contour showed the two factors' roles swapped against the axis labels and ranges. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- .../common/visualize_factor_distributions.py | 3 ++- tests/test_visualize_factor_distributions.py | 21 +++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/skillmodels/common/visualize_factor_distributions.py b/src/skillmodels/common/visualize_factor_distributions.py index 8f77f9e5..bef4becb 100644 --- a/src/skillmodels/common/visualize_factor_distributions.py +++ b/src/skillmodels/common/visualize_factor_distributions.py @@ -365,10 +365,11 @@ def bivariate_density_contours( factors=pair, n_points=n_points, ) + # `z[i, j]` is the density at `(x[i], y[j])`; plotly reads rows as y. contour = go.Contour( x=x[:, 0], y=y[0, :], - z=z, + z=z.T, line={"color": getattr(px.colors.qualitative, lines_colorscale)[i]}, ) fig.add_trace(contour) diff --git a/tests/test_visualize_factor_distributions.py b/tests/test_visualize_factor_distributions.py index d181372a..54356067 100644 --- a/tests/test_visualize_factor_distributions.py +++ b/tests/test_visualize_factor_distributions.py @@ -2,7 +2,9 @@ from pathlib import Path +import numpy as np import pandas as pd +from scipy.stats import gaussian_kde from skillmodels.chs.maximization_inputs import get_maximization_inputs from skillmodels.common.config import TEST_DATA_DIR @@ -115,3 +117,22 @@ def test_visualize_factor_distributions_with_period_indexed_states() -> None: contour_plots=contours, surface_plots=None, ) + + +def test_bivariate_density_contours_put_first_factor_on_x_axis() -> None: + """Contour value at (x[col], y[row]) is the KDE of (first, second) factor there.""" + data, _params, states = _load_model2_filtered() + contours = bivariate_density_contours( + data=data, + model_spec=MODEL2, + period=1, + filtered_states=states, + n_points=7, + ) + (fac_x, fac_y), fig = next(iter(contours.items())) + trace = fig.data[0] + period_states = states.query("period == 1") + kernel = gaussian_kde(np.vstack([period_states[fac_x], period_states[fac_y]])) + xx, yy = np.meshgrid(trace.x, trace.y) + expected = kernel(np.vstack([xx.ravel(), yy.ravel()])).reshape(xx.shape) + np.testing.assert_allclose(np.asarray(trace.z), expected, rtol=1e-10) From 5d26103c2009d2886334bdb483c28ccd5fd6f1a5 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 16 Oct 2026 21:21:21 +0200 Subject: [PATCH 16/18] Transition plots, average over draws: draw from complete rows A missing observed factor in one of the 50 draws made the mean NaN at every grid point, so the plots were empty. Draws now come from persons with all other factors observed, with a fixed seed. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- .../common/visualize_transition_equations.py | 12 +++++++-- tests/test_visualize_transition_equations.py | 26 +++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/src/skillmodels/common/visualize_transition_equations.py b/src/skillmodels/common/visualize_transition_equations.py index 6095d002..6abd3b20 100644 --- a/src/skillmodels/common/visualize_transition_equations.py +++ b/src/skillmodels/common/visualize_transition_equations.py @@ -193,7 +193,8 @@ def get_transition_plots( aggregation_method: How to handle factors not being varied in a plot. - "quantiles": Fix other factors at specified quantiles (default). - "median": Fix other factors at their median (0.5 quantile). - - "average": Average over random draws of other factors. + - "average": Average over random draws of other factors, from the + persons with all other factors observed. n_points: Number of grid points per input. Default 50. n_draws: Number of randomly drawn values of the factors that are averaged out. Only relevant if aggregation_method is "average". Default 50. @@ -765,7 +766,14 @@ def _prepare_data_for_one_plot_average_2d( period_data = states_data.query(f"aug_period == {aug_period}") sampled_factors = [factor for factor in all_factors if factor != input_factor] - draws = period_data[sampled_factors].sample(n=n_draws) + # Persons with a missing factor (e.g. an observed input) would make every mean + # NaN; draw from the complete rows, with a fixed seed for reproducible figures. + complete = period_data[sampled_factors].dropna() + draws = complete.sample( + n=n_draws, + replace=len(complete) < n_draws, + random_state=np.random.default_rng(0), + ) input_min = state_ranges[input_factor].loc[aug_period]["minimum"] input_max = state_ranges[input_factor].loc[aug_period]["maximum"] diff --git a/tests/test_visualize_transition_equations.py b/tests/test_visualize_transition_equations.py index cb922f7f..13a4969b 100644 --- a/tests/test_visualize_transition_equations.py +++ b/tests/test_visualize_transition_equations.py @@ -12,6 +12,7 @@ from skillmodels.common.individual_states import get_individual_states_from_params from skillmodels.common.visualize_transition_equations import ( _get_state_ranges, + _prepare_data_for_one_plot_average_2d, combine_transition_plots, get_transition_plots, ) @@ -173,3 +174,28 @@ def test_get_state_ranges_adds_factors_missing_from_given_ranges() -> None: ) assert ranges["fac1"] is given["fac1"] assert ranges["obs"]["maximum"].tolist() == [7, 9] + + +def test_average_over_draws_skips_persons_with_a_missing_factor() -> None: + # Observed factors can be missing; a draw with a NaN would make every mean NaN. + states = pd.DataFrame( + { + "aug_period": [0] * 4, + "fac1": [0.0, 1.0, 2.0, 3.0], + "obs": [1.0, np.nan, 1.0, 1.0], + } + ) + ranges = {"fac1": pd.DataFrame({"minimum": [0.0], "maximum": [1.0]})} + out = _prepare_data_for_one_plot_average_2d( + states_data=states, + state_ranges=ranges, + aug_period=0, + input_factor="fac1", + output_factor="fac1", + n_points=3, + n_draws=10, + transition_function=lambda _params, x: x.sum(axis=1), + transition_params={}, + all_factors=("fac1", "obs"), + ) + assert out["output_fac1"].tolist() == [1.0, 1.5, 2.0] From 5acd0a3f392c74e24964efc3974a295847b4b9fc Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sun, 18 Oct 2026 08:00:41 +0200 Subject: [PATCH 17/18] Format the estimator prerequisites table Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- .../estimator_prerequisites.md | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/docs/reference_guides/estimator_prerequisites.md b/docs/reference_guides/estimator_prerequisites.md index a119eab7..9e22066b 100644 --- a/docs/reference_guides/estimator_prerequisites.md +++ b/docs/reference_guides/estimator_prerequisites.md @@ -13,20 +13,20 @@ three estimators. ## Prerequisites matrix -| Dimension | CHS | AF | AMN | -| --------------------- | ------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Estimator | Joint maximum likelihood with square-root Gaussian-component filtering. | Sequential period-by-period likelihood with Halton integration. | Three stages: Gaussian-mixture EM → minimum distance → simulate-and-regress. | -| Latent distribution | Finite mixture of Gaussian initial states via `ModelSpec.n_mixtures`; Gaussian component filtering thereafter. | Initial finite mixture plus sequentially carried distributions; Halton nodes approximate the integrals. | Stage-1 Gaussian mixture over the augmented measurement vector; `n_mixtures` sets the component count. | -| Measurement families | **Gaussian only** (standard Kalman update): `get_maximization_inputs` / `estimate_chs` raise `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | **Initial period only**: probit/Tobit measurement families are honoured at the period-0 measurement system; transition periods ($t \geq 1$) fall back to an all-Gaussian measurement kernel. | **Gaussian only**: `estimate_amn` raises `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | -| Missing data | Gaussian measurement updates skip individually missing measurements. | Measurement masks skip missing rows in the per-step likelihood contributions. | `mixture_em_method="complete_case"` (default) or `"missing_data"` (marginalises over missing entries under MAR); never-observed columns require `allow_never_observed_measurements=True`. | -| Endogenous investment | Reconstructed endogenous factors plus a full `CorrectionSpec` control-function basis. | Reconstructed endogenous investment with independent shocks (source/destination calendar adapter); rejects `CorrectionSpec` / nonzero `kappa`. | Endogenous investment in Stage 3; correction support is a linearised control-function term, narrower than CHS. | -| Corrections | Full `CorrectionSpec` / `kappa` polynomial basis in the processed transition DAG. | Not implemented — the validator raises if a `CorrectionSpec` is attached. | Linear `cf` term only; a higher-order `kappa_terms` request raises. | -| Custom transitions | Built-ins and `@register_params` callables. | Built-ins and `@register_params` callables. | Built-ins and `@register_params` callables via the Stage-3 generic NLS path, but `log_ces` / `log_ces_with_constant` reject fixed parameters there. | -| `fixed_params` | Honoured through the shared parameter index. | Honoured (with public `constraints=` limited to `select_by_loc` equality groups). | Honoured only for the categories each stage estimates (Stage-2 loadings/intercepts/SDs, Stage-3 `transition`); other categories raise. `start_params` and `constraints` raise. | -| Normalization | Shared `Normalizations` / `fixed_params` / equality constraints. The checker is a precheck, not an identification proof. | Same public `ModelSpec`, with the AF-specific source/destination calendar and `af_state_role` metadata. | Stage-2 imposes its own structural moment restrictions and a mean-zero mixture convention. | -| Anchoring | Supported through `get_maximization_inputs` / the CHS path. | Not part of the AF likelihood; downstream visualization only. | Not wired through the AMN stages (the result reports unanchored scales). | -| Cost / scaling | Potentially expensive joint ML; the JAX square-root filter helps numerical stability. | Sequential but quadrature-heavy; cost grows with node count and state dimension. | Fast when the stages are well behaved; the cluster bootstrap is expensive because it re-estimates every stage. | -| Use it for | Likelihood benchmark, anchoring, correction-heavy models, Gaussian measurement systems. | Sequential AF-style models, initial-period limited measurements, period-by-period diagnostics. | Mixture-heavy Gaussian-measurement models, fast start values for CHS, Stage-1/2 structural diagnostics. | +| Dimension | CHS | AF | AMN | +| --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Estimator | Joint maximum likelihood with square-root Gaussian-component filtering. | Sequential period-by-period likelihood with Halton integration. | Three stages: Gaussian-mixture EM → minimum distance → simulate-and-regress. | +| Latent distribution | Finite mixture of Gaussian initial states via `ModelSpec.n_mixtures`; Gaussian component filtering thereafter. | Initial finite mixture plus sequentially carried distributions; Halton nodes approximate the integrals. | Stage-1 Gaussian mixture over the augmented measurement vector; `n_mixtures` sets the component count. | +| Measurement families | **Gaussian only** (standard Kalman update): `get_maximization_inputs` / `estimate_chs` raise `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | **Initial period only**: probit/Tobit measurement families are honoured at the period-0 measurement system; transition periods ($t \geq 1$) fall back to an all-Gaussian measurement kernel. | **Gaussian only**: `estimate_amn` raises `NotImplementedError` if `ModelSpec.measurement_models` declares any probit/Tobit measurement. | +| Missing data | Gaussian measurement updates skip individually missing measurements. | Measurement masks skip missing rows in the per-step likelihood contributions. | `mixture_em_method="complete_case"` (default) or `"missing_data"` (marginalises over missing entries under MAR); never-observed columns require `allow_never_observed_measurements=True`. | +| Endogenous investment | Reconstructed endogenous factors plus a full `CorrectionSpec` control-function basis. | Reconstructed endogenous investment with independent shocks (source/destination calendar adapter); rejects `CorrectionSpec` / nonzero `kappa`. | Endogenous investment in Stage 3; correction support is a linearised control-function term, narrower than CHS. | +| Corrections | Full `CorrectionSpec` / `kappa` polynomial basis in the processed transition DAG. | Not implemented — the validator raises if a `CorrectionSpec` is attached. | Linear `cf` term only; a higher-order `kappa_terms` request raises. | +| Custom transitions | Built-ins and `@register_params` callables. | Built-ins and `@register_params` callables. | Built-ins and `@register_params` callables via the Stage-3 generic NLS path, but `log_ces` / `log_ces_with_constant` reject fixed parameters there. | +| `fixed_params` | Honoured through the shared parameter index. | Honoured (with public `constraints=` limited to `select_by_loc` equality groups). | Honoured only for the categories each stage estimates (Stage-2 loadings/intercepts/SDs, Stage-3 `transition`); other categories raise. `start_params` and `constraints` raise. | +| Normalization | Shared `Normalizations` / `fixed_params` / equality constraints. The checker is a precheck, not an identification proof. | Same public `ModelSpec`, with the AF-specific source/destination calendar and `af_state_role` metadata. | Stage-2 imposes its own structural moment restrictions and a mean-zero mixture convention. | +| Anchoring | Supported through `get_maximization_inputs` / the CHS path. | Not part of the AF likelihood; downstream visualization only. | Not wired through the AMN stages (the result reports unanchored scales). | +| Cost / scaling | Potentially expensive joint ML; the JAX square-root filter helps numerical stability. | Sequential but quadrature-heavy; cost grows with node count and state dimension. | Fast when the stages are well behaved; the cluster bootstrap is expensive because it re-estimates every stage. | +| Use it for | Likelihood benchmark, anchoring, correction-heavy models, Gaussian measurement systems. | Sequential AF-style models, initial-period limited measurements, period-by-period diagnostics. | Mixture-heavy Gaussian-measurement models, fast start values for CHS, Stage-1/2 structural diagnostics. | ## Reading the matrix From a53ee7a2c08f9721923637adc014f8b082d356ce Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sun, 18 Oct 2026 08:58:10 +0200 Subject: [PATCH 18/18] ModelSpec.from_dict: reject unrecognized top-level keys `from_dict` silently ignored any key it did not read, so a whole `estimation_options` block (`n_mixtures`, `robust_bounds`, `bounds_distance`) was dropped and the model was built with the defaults. Check the top-level keys against the ones `from_dict` reads, including `measurement_models` and `missing_observed_factor_moments`, and raise `ModelSpecInitializationError` on anything else. Re-derived from e27a16a (branch modelspec-reject-unknown-keys, cut from an old main whose key set was narrower). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/common/model_spec.py | 28 ++++++++++++++++++++++++++++ tests/test_model_spec.py | 23 +++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/src/skillmodels/common/model_spec.py b/src/skillmodels/common/model_spec.py index 993db665..0c886cb2 100644 --- a/src/skillmodels/common/model_spec.py +++ b/src/skillmodels/common/model_spec.py @@ -22,6 +22,20 @@ Normalizations, ensure_containers_are_immutable, ) +from skillmodels.exceptions import ModelSpecInitializationError + +_ALLOWED_MODEL_SPEC_KEYS = frozenset( + { + "factors", + "observed_factors", + "controls", + "stagemap", + "anchoring", + "n_mixtures", + "measurement_models", + "missing_observed_factor_moments", + } +) def _parse_measurement_models( @@ -302,7 +316,21 @@ def from_dict(cls, d: dict[str, Any]) -> Self: Returns: A ModelSpec instance. + Raises: + ModelSpecInitializationError: If `d` contains an unrecognized top-level + key. Silently dropping unknown keys hides typos and stale config + blocks (e.g. an `estimation_options` mapping whose settings never + take effect). + """ + unknown = set(d) - _ALLOWED_MODEL_SPEC_KEYS + if unknown: + msg = ( + f"Unrecognized ModelSpec keys: {sorted(unknown)}. " + f"Allowed keys: {sorted(_ALLOWED_MODEL_SPEC_KEYS)}." + ) + raise ModelSpecInitializationError(msg) + factors = {} auto_instruments: list[str] = [] for name, spec in d["factors"].items(): diff --git a/tests/test_model_spec.py b/tests/test_model_spec.py index 97d34aa1..f0a258da 100644 --- a/tests/test_model_spec.py +++ b/tests/test_model_spec.py @@ -9,6 +9,7 @@ ModelSpec, Normalizations, ) +from skillmodels.exceptions import ModelSpecInitializationError def _minimal_dict(): @@ -81,6 +82,28 @@ def test_from_dict_with_stagemap() -> None: assert spec.stagemap == (0,) +def test_from_dict_rejects_unknown_keys() -> None: + """An unrecognized top-level key raises instead of being silently dropped.""" + d = _minimal_dict() + d["estimation_options"] = {"n_mixtures": 2, "bounds_distance": 0.001} + with pytest.raises(ModelSpecInitializationError, match="estimation_options"): + ModelSpec.from_dict(d) + + +def test_from_dict_accepts_every_key_it_reads() -> None: + d = _minimal_dict() + d |= { + "observed_factors": ["x"], + "controls": ["c"], + "stagemap": [0], + "n_mixtures": 1, + "measurement_models": {}, + "missing_observed_factor_moments": {"x": ("x_mean", "x_var")}, + } + spec = ModelSpec.from_dict(d) + assert spec.observed_factors == ("x",) + + def test_correction_spec_defaults() -> None: cf = CorrectionSpec(instruments=("z1",)) assert cf.instruments == ("z1",)