From 7d3a406f10483293ceb3907fad1fb7d1cc7430dc Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Mon, 2 Nov 2026 22:20:35 +0100 Subject: [PATCH 1/3] Add per-observation scores to the CHS maximization inputs `scores(params)` returns the derivatives of the per-observation log-likelihoods (observations x parameters). Observations are independent, so the gradient of one observation's log-likelihood, vmapped over observations, gives all scores at about the cost of one gradient. The data are split as for the gradient. BHHH needs the scores at every iteration; so far they were only available by finite differences. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- AGENTS.md | 4 +- src/skillmodels/chs/maximization_inputs.py | 97 +++++++++++++++------- tests/test_likelihood_regression.py | 38 +++++++++ tests/test_tobit_interval_observations.py | 15 ++++ 4 files changed, 125 insertions(+), 29 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 6b73b8de..c54f0246 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -121,7 +121,7 @@ simulate_dataset() → Simulate states (with optional policy effects) ### `get_maximization_inputs()` Return Dict -Returns a dict with 6 keys: +Returns a dict with 7 keys: - `"loglike"`: `(params: pd.DataFrame) -> float` — jitted scalar log-likelihood - `"loglikeobs"`: `(params: pd.DataFrame) -> NDArray` — jitted per-observation @@ -130,6 +130,8 @@ Returns a dict with 6 keys: keys `value`, `contributions`, `residuals`, `residual_sds`, `filtered_states`, `state_ranges`, etc. - `"loglike_and_gradient"`: `(params: pd.DataFrame) -> tuple[float, NDArray]` +- `"scores"`: `(params: pd.DataFrame) -> NDArray` — derivatives of the per-observation + log-likelihoods (observations x parameters), from a vmapped gradient - `"constraints"`: list of optimagic constraint objects - `"params_template"`: `pd.DataFrame` with correct MultiIndex and bounds; fixed constraints pre-applied diff --git a/src/skillmodels/chs/maximization_inputs.py b/src/skillmodels/chs/maximization_inputs.py index 10c61778..645652da 100644 --- a/src/skillmodels/chs/maximization_inputs.py +++ b/src/skillmodels/chs/maximization_inputs.py @@ -2,7 +2,7 @@ import functools import warnings -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from types import MappingProxyType from typing import Any @@ -59,6 +59,16 @@ jax.config.update("jax_enable_x64", True) # noqa: FBT003 +# The per-observation data, with the observations on axis 1. +_SLICED_DATA = ( + "measurements", + "controls", + "observed_factors", + "observed_factor_vars", + "events", + "interval_bounds", +) + @beartype(conf=ESTIMATION_CONF) def get_maximization_inputs( # noqa: C901, PLR0915 @@ -102,6 +112,8 @@ def get_maximization_inputs( # noqa: C901, PLR0915 function with respect to the parameters. loglike_and_gradient: Combination of loglike and loglike_gradient that is faster than calling the two functions separately. + scores: The derivatives of the per-observation log likelihoods, an array + with one row per observation and one column per parameter. constraints: List of optimagic constraints that are implied by the model specification, extended by any user-supplied ``fixed_params``. params_template: Parameter DataFrame with correct index and @@ -201,6 +213,41 @@ def get_maximization_inputs( # noqa: C901, PLR0915 _jitted_loglikeobs = jax.jit(partialed_loglikes["llo"]) _gradient = jax.jit(jax.grad(partialed_loglikes["ll"])) + def _person_loglike(params_vec: Array, *person_data: Array | None) -> Array: + # Observations are independent, so the gradient of one observation's + # log-likelihood, vmapped over observations, gives all scores in about the + # time of one gradient. + return partialed_loglikes["llo"]( + params_vec, + **{ + key: None if value is None else value[:, None] + for key, value in zip(_SLICED_DATA, person_data, strict=True) + }, + )[0] + + _person_scores = jax.jit( + jax.vmap(jax.grad(_person_loglike), in_axes=(None, *[1] * len(_SLICED_DATA))) + ) + + def _data_slices() -> Iterator[dict[str, Array | None]]: + """Yield the per-observation data in `chs_options.split_dataset` slices.""" + data = { + "measurements": processed_data["measurements"], + "controls": processed_data["controls"], + "observed_factors": processed_data["observed_factors"], + "observed_factor_vars": processed_data["observed_factor_vars"], + "events": events, + "interval_bounds": interval_bounds, + } + n_obs = data["measurements"].shape[1] + step = int(n_obs / chs_options.split_dataset) + for i in range(chs_options.split_dataset): + stop = n_obs if i == chs_options.split_dataset - 1 else (i + 1) * step + yield { + key: None if value is None else value[:, i * step : stop] + for key, value in data.items() + } + def loglike(params: pd.DataFrame) -> float: params_vec = partialed_get_jnp_params_vec(params) return float(_jitted_loglike(params_vec)) @@ -214,34 +261,9 @@ def loglike_and_gradient( ) -> tuple[float, NDArray[np.floating]]: params_vec = partialed_get_jnp_params_vec(params) crit = float(_jitted_loglike(params_vec)) - 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): - stop = n_obs if i == split_dataset - 1 else stop - events_slice = {} if events is None else {"events": events[:, start:stop]} - if interval_bounds is not None: - events_slice["interval_bounds"] = interval_bounds[:, start:stop] - measurements_slice = processed_data["measurements"][:, start:stop] - controls_slice = processed_data["controls"][:, start:stop, :] - observed_factors_slice = processed_data["observed_factors"][ - :, start:stop, : - ] - _grad += _gradient( - params_vec, - measurements=measurements_slice, - controls=controls_slice, - observed_factors=observed_factors_slice, - observed_factor_vars=processed_data["observed_factor_vars"][ - :, start:stop, : - ], - **events_slice, - ) - start += step - stop += step + for data_slice in _data_slices(): + _grad += _gradient(params_vec, **data_slice) grad = _to_numpy(_grad) if not (np.isfinite(crit) and np.isfinite(grad).all()): # Raise rather than return, so `error_handling="continue"` can replace @@ -253,6 +275,24 @@ def loglike_and_gradient( raise FloatingPointError(msg) return crit, grad + def scores(params: pd.DataFrame) -> NDArray[np.floating]: + """Return the derivatives of the per-observation log-likelihoods. + + Rows are observations, columns the parameters of `params`. + + """ + params_vec = partialed_get_jnp_params_vec(params) + return np.concatenate( + [ + _to_numpy( + _person_scores( + params_vec, *(data_slice[key] for key in _SLICED_DATA) + ) + ) + for data_slice in _data_slices() + ] + ) + def debug_loglike(params: pd.DataFrame) -> dict[str, Any]: params_vec = partialed_get_jnp_params_vec(params) jax_output = partialed_loglikes["debug_ll"](params_vec) @@ -353,6 +393,7 @@ def debug_loglike(params: pd.DataFrame) -> dict[str, Any]: "loglikeobs": loglikeobs, "debug_loglike": debug_loglike, "loglike_and_gradient": loglike_and_gradient, + "scores": scores, "constraints": constraints, "params_template": params_template, } diff --git a/tests/test_likelihood_regression.py b/tests/test_likelihood_regression.py index 5a2da600..b55b6df8 100644 --- a/tests/test_likelihood_regression.py +++ b/tests/test_likelihood_regression.py @@ -117,6 +117,44 @@ def test_splitting_does_not_change_gradient(model2, model2_data) -> None: aaae(gradient, gradient_split) +def test_scores_are_the_derivatives_of_the_per_observation_loglikes( + model2, model2_data +) -> None: + inputs = get_maximization_inputs( + model2, model2_data, chs_options=MODEL2_CHS_OPTIONS + ) + params = inputs["params_template"] + params["value"] = 0.1 + + scores = inputs["scores"](params) + + _, gradient = inputs["loglike_and_gradient"](params) + assert scores.shape == (len(model2_data.index.unique("caseid")), len(params)) + aaae(scores.sum(axis=0), gradient) + step = 1e-5 + for position in (0, len(params) // 2, len(params) - 1): + up, down = params.copy(), params.copy() + up.iloc[position, 0] += step + down.iloc[position, 0] -= step + central = (inputs["loglikeobs"](up) - inputs["loglikeobs"](down)) / (2 * step) + np.testing.assert_allclose(scores[:, position], central, rtol=0, atol=1e-5) + + +def test_splitting_does_not_change_scores(model2, model2_data) -> None: + inputs = get_maximization_inputs( + model2, model2_data, chs_options=MODEL2_CHS_OPTIONS + ) + inputs_split = get_maximization_inputs( + model2, + model2_data, + chs_options=replace(MODEL2_CHS_OPTIONS, split_dataset=13), + ) + params = inputs["params_template"] + params["value"] = 0.1 + + aaae(inputs["scores"](params), inputs_split["scores"](params)) + + 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( diff --git a/tests/test_tobit_interval_observations.py b/tests/test_tobit_interval_observations.py index 7c251056..e6eff8fe 100644 --- a/tests/test_tobit_interval_observations.py +++ b/tests/test_tobit_interval_observations.py @@ -408,6 +408,21 @@ def test_split_dataset_gradient_equals_unsplit_gradient(interval_inputs) -> None ) +def test_split_scores_sum_to_the_gradient(interval_inputs) -> None: + data, inputs, params = interval_inputs + split = get_maximization_inputs( + model_spec=_model(), + data=data, + chs_options=CHSEstimationOptions(start_params_strategy="none", split_dataset=3), + ) + np.testing.assert_allclose( + split["scores"](params).sum(axis=0), + inputs["loglike_and_gradient"](params)[1], + rtol=1e-11, + atol=1e-11, + ) + + def test_without_interval_observations_the_model_equals_the_plain_tobit() -> None: data = _data().assign(c1_lo=np.nan, c1_hi=np.nan) options = CHSEstimationOptions(start_params_strategy="none") From 422920e35b8444b57613d454efd64ac5ae8ff7a7 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Tue, 3 Nov 2026 05:38:29 +0100 Subject: [PATCH 2/3] Give likelihood optimizers the scores as their Jacobian For an optimizer that works on likelihood contributions (BHHH), `estimate_chs` passes the per-observation scores, marked as a likelihood Jacobian, instead of the scalar log-likelihood and its gradient. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/estimate.py | 9 ++++++++- tests/test_chs_estimate.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/skillmodels/chs/estimate.py b/src/skillmodels/chs/estimate.py index e5982d58..3639e5ba 100644 --- a/src/skillmodels/chs/estimate.py +++ b/src/skillmodels/chs/estimate.py @@ -16,6 +16,8 @@ import pandas as pd from beartype import beartype from estimagic import estimate_ml +from optimagic.algorithms import ALL_ALGORITHMS +from optimagic.typing import AggregationLevel from skillmodels._beartype_conf import ESTIMATION_CONF from skillmodels.chs.maximization_inputs import get_maximization_inputs @@ -155,9 +157,14 @@ def estimate_chs( optimize_options = { "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, } + solver_type = ALL_ALGORITHMS[options.optimizer_algorithm].algo_info.solver_type + if solver_type == AggregationLevel.LIKELIHOOD: + # BHHH-type optimizers need the derivatives of the contributions. + optimize_options["jac"] = om.mark.likelihood(max_inputs["scores"]) + else: + optimize_options["fun_and_jac"] = max_inputs["loglike_and_gradient"] # Default to OPG/jacobian-based inference (`hessian=False`): the numerical # Hessian costs O(n_params**2) Kalman passes and is prohibitive on real diff --git a/tests/test_chs_estimate.py b/tests/test_chs_estimate.py index 4061eaf3..c0d0d2e4 100644 --- a/tests/test_chs_estimate.py +++ b/tests/test_chs_estimate.py @@ -7,6 +7,7 @@ import optimagic as om import pandas as pd import pytest +from optimagic.typing import AggregationLevel from skillmodels.chs.estimate import estimate_chs from skillmodels.chs.options import CHSEstimationOptions @@ -154,6 +155,35 @@ def fake_estimate_ml(**kwargs): assert captured["error_handling"] == "continue" +def test_estimate_chs_gives_a_likelihood_optimizer_the_scores( + model2_data, anchoring_start_params, monkeypatch +): + """BHHH-type optimizers get the per-observation scores as their Jacobian.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + 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", optimizer_algorithm="bhhh" + ), + start_params=anchoring_start_params, + require_initial_anchors=False, + ) + options = captured["optimize_options"] + assert "fun_and_jac" not in options + params = captured["params"] + assert options["jac"]._problem_type == AggregationLevel.LIKELIHOOD # noqa: SLF001 + scores = options["jac"](params) + assert scores.shape == (len(model2_data.index.unique("caseid")), len(params)) + + def test_estimate_chs_applies_lower_bounds_and_clips_start( model2_data, anchoring_start_params, monkeypatch ): From efefa3428f18aa44cf63ca343ba1d7da44ed28e0 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Tue, 3 Nov 2026 08:27:17 +0100 Subject: [PATCH 3/3] Use the scores as the Jacobian of the ML inference estimate_ml computed the score matrix of the OPG covariance by finite differences, about one likelihood evaluation per free parameter (an hour on a real model). The vmapped AD scores cost about one gradient. estimate_ml_options={"jacobian": None} restores finite differences. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Y3CuwvNUhYTSpriKYSwqmq --- src/skillmodels/chs/estimate.py | 5 ++++- tests/test_chs_estimate.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/skillmodels/chs/estimate.py b/src/skillmodels/chs/estimate.py index 3639e5ba..21f111c3 100644 --- a/src/skillmodels/chs/estimate.py +++ b/src/skillmodels/chs/estimate.py @@ -168,9 +168,12 @@ def estimate_chs( # Default to OPG/jacobian-based inference (`hessian=False`): the numerical # Hessian costs O(n_params**2) Kalman passes and is prohibitive on real - # models. Overridable via `options.estimate_ml_options`. + # models. The OPG uses the AD scores, which cost about one gradient; a + # numerical Jacobian costs about n_params likelihood evaluations. Both are + # overridable via `options.estimate_ml_options`. estimate_ml_kwargs = { "hessian": False, + "jacobian": max_inputs["scores"], **to_plain_dict(options.estimate_ml_options), } diff --git a/tests/test_chs_estimate.py b/tests/test_chs_estimate.py index c0d0d2e4..ff393807 100644 --- a/tests/test_chs_estimate.py +++ b/tests/test_chs_estimate.py @@ -184,6 +184,39 @@ def fake_estimate_ml(**kwargs): assert scores.shape == (len(model2_data.index.unique("caseid")), len(params)) +@pytest.mark.parametrize( + ("estimate_ml_options", "uses_scores"), + [({}, True), ({"jacobian": None}, False)], +) +def test_estimate_chs_gives_inference_the_scores( + model2_data, anchoring_start_params, monkeypatch, estimate_ml_options, uses_scores +): + """Inference uses the scores unless `estimate_ml_options` asks otherwise.""" + captured = {} + + def fake_estimate_ml(**kwargs): + captured.update(kwargs) + 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", estimate_ml_options=estimate_ml_options + ), + start_params=anchoring_start_params, + require_initial_anchors=False, + ) + if uses_scores: + params = captured["params"] + scores = captured["jacobian"](params) + assert scores.shape == (len(model2_data.index.unique("caseid")), len(params)) + else: + assert captured["jacobian"] is None + + def test_estimate_chs_applies_lower_bounds_and_clips_start( model2_data, anchoring_start_params, monkeypatch ):