CHS: per-observation scores by AD, for BHHH and ML inference

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

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,15 +157,23 @@ 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
-    # 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),
     }
 

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,
     }

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,68 @@ 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))
+
+
+@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
 ):

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(

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")