perf(hierarchical): faster options for EmpiricalBayes and VariationalBayes

CHANGELOG.md

@@ -5,6 +5,21 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.26.0.dev3] - 2026-10-14
+
+### Added
+
+- Added `benchmarks/hierarchical_profile.py`, which shows where the time of a `VariationalBayes` fit goes and counts the objective evaluations per participant and iteration; against cbm (RLRW, 100 participants, 120 trials), cpm makes about as many evaluations but each one costs more, at 17.3 against 5.0 ms per participant and iteration
+- Added a `hessian` argument to `cpm.optimisation.FminBound`, `Fmin`, `Minimize` and `Bads`; `hessian="finite_differences"` computes the Hessian at the optimum with bounded finite differences (`cpm.core.optimisers.finite_difference_hessian`), cutting Hessian evaluations from 22 to 7 per participant and iteration in the speed comparison. The default, `"numdifftools"`, keeps the previous Hessian
+- Added a `start` argument to `cpm.hierarchical.EmpiricalBayes` and `cpm.hierarchical.VariationalBayes`; `start="prior_mean"` starts the first optimisation of every participant from the current prior's location, clipped into the bounds, as cbm does, so fits stop sooner. The default, `"random"`, keeps the previous starts
+- Added a `positive_definite` column to the `fit` of `cpm.hierarchical.EmpiricalBayes` and `cpm.hierarchical.VariationalBayes`, recording whether each participant's Hessian was positive definite, and a `not_positive_definite` column to their `hyperparameters` with the count; an estimate on a bound can have a non-positive-definite Hessian and is kept, as the bounds require
+- Added the results of the speed comparison with cbm (100 participants, 120 trials) in `benchmarks/results/hierarchical-comparison-*.csv`: combining `hessian="finite_differences"`, `start="prior_mean"` and `number_of_starts=1` brings `EmpiricalBayes`/`VariationalBayes` from 20.8/20.1 to 6.8/7.0 ms per participant and iteration, against cbm's 4.4 ms, with RMSE against 0.26.0.dev1 staying within bootstrap 95% intervals throughout
+
+### Changed
+
+- `cpm.optimisation.FminBound` with `approx_grad=True` now computes the finite-difference gradient itself, together with the objective, exactly as SciPy does, with identical results; with `LogLikelihood.bernoulli`, this makes the hierarchical fits of the speed comparison 19% faster
+- The defaults of the new arguments, `hessian="numdifftools"` and `start="random"`, give the same results as before for the same seed; the faster options have to be chosen
+
 ## [0.26.0.dev2] - 2026-10-06
 
 ### Added

benchmarks/hierarchical_profile.py

@@ -0,0 +1,261 @@
+"""
+Where does the time of a hierarchical fit go?
+
+    python benchmarks/hierarchical_profile.py [--data DATASET.csv] [--iterations 5] [--label NAME]
+
+Fits one dataset of the four-armed bandit (100 participants, 120 trials) with
+`RLRW` and `VariationalBayes` for a fixed number of iterations, one chain and
+two optimiser starts per participant, under `cProfile`. Prints the share of
+time in the optimiser, the Hessian, `evaluate_fit`, the population update and
+the parts of `objective`, and the number of objective evaluations per
+participant and iteration, separately for the optimiser, the Hessian and
+`evaluate_fit`. Writes the table to benchmarks/results/<label>.csv.
+
+`--data` is a dataset in the format `RLRW` reads (columns ppt, trial, arm_left,
+arm_right, reward_left, reward_right, response). Without it, the task and the
+responses are simulated from a known population. The settings match the speed
+comparison with cbm: bounds alpha [0.01, 1] and temperature [0, 10], a starting
+prior of mean 0.5 and sd 0.5 for alpha and mean 5 and sd 5 for temperature.
+
+The profiler adds overhead to Python-heavy code, so the shares are indicative;
+the fit is also timed once without it.
+"""
+
+import argparse
+import cProfile
+import pathlib
+import pstats
+import sys
+import time
+import warnings
+
+import numpy as np
+import pandas as pd
+from scipy.stats import truncnorm
+
+sys.path.insert(0, str(pathlib.Path(__file__).parent))
+
+import cpm  # noqa: E402
+import cpm.optimisation.fmin as fmin_module  # noqa: E402
+from cpm.applications.reinforcement_learning import RLRW  # noqa: E402
+from cpm.hierarchical import VariationalBayes  # noqa: E402
+from cpm.optimisation import FminBound, minimise  # noqa: E402
+
+RESULTS = pathlib.Path(__file__).parent / "results"
+PARTICIPANTS, TRIALS, STIMULI = 100, 120, 4
+BOUNDS = {"alpha": (1e-2, 1.0), "temperature": (0.0, 10.0)}
+STARTING_PRIOR = {"alpha": {"mean": 0.5, "sd": 0.5}, "temperature": {"mean": 5.0, "sd": 5.0}}
+POPULATION = {"alpha": {"mean": 0.5, "sd": 0.15}, "temperature": {"mean": 4.5, "sd": 1.5}}
+
+
+def simulate(seed=2026):
+    """A random four-armed bandit task with responses of `RLRW` from a known population."""
+    rng = np.random.default_rng(seed)
+    n, trials = PARTICIPANTS, TRIALS
+    arms = np.array([rng.choice(STIMULI, size=2, replace=False) for _ in range(n * trials)]).reshape(n, trials, 2)
+    reward_probability = np.array([0.2, 0.4, 0.6, 0.8])
+    rewards = (rng.random((n, trials, 2)) < reward_probability[arms]).astype(float)
+    parameters = {
+        name: truncnorm.rvs(
+            (BOUNDS[name][0] - p["mean"]) / p["sd"], (BOUNDS[name][1] - p["mean"]) / p["sd"],
+            loc=p["mean"], scale=p["sd"], size=n, random_state=rng,
+        )
+        for name, p in POPULATION.items()
+    }
+    values = np.full((n, STIMULI), 1 / STIMULI)
+    rows = np.arange(n)
+    response = np.empty((n, trials), dtype=int)
+    for t in range(trials):
+        scaled = values[rows[:, None], arms[:, t]] * parameters["temperature"][:, None]
+        scaled -= scaled.max(axis=1, keepdims=True)
+        policy = np.exp(scaled)
+        policy /= policy.sum(axis=1, keepdims=True)
+        choice = (rng.random(n) >= policy[:, 0]).astype(int)
+        stimulus = arms[rows, t, choice]
+        values[rows, stimulus] += parameters["alpha"] * (rewards[rows, t, choice] - values[rows, stimulus])
+        response[:, t] = choice
+    return pd.DataFrame(
+        {
+            "ppt": np.repeat(np.arange(1, n + 1), trials),
+            "trial": np.tile(np.arange(1, trials + 1), n),
+            "arm_left": arms[:, :, 0].ravel() + 1,
+            "arm_right": arms[:, :, 1].ravel() + 1,
+            "reward_left": rewards[:, :, 0].ravel(),
+            "reward_right": rewards[:, :, 1].ravel(),
+            "response": response.ravel(),
+            "observed": response.ravel(),
+        }
+    )
+
+
+def load(path):
+    data = pd.read_csv(path)
+    ppts = np.sort(data.ppt.unique())[:PARTICIPANTS]
+    data = data[data.ppt.isin(ppts) & (data.trial <= TRIALS)].reset_index(drop=True)
+    data["observed"] = data["response"]
+    return data
+
+
+def setup(data, iterations, start="random", **options):
+    """`VariationalBayes` with the settings of the speed comparison, forced to run `iterations` iterations; `options` go to `FminBound`."""
+    first = data[data.ppt == data.ppt.iloc[0]]
+    model = RLRW(data=first, dimensions=STIMULI, parameters_settings=[[0.5, *BOUNDS["alpha"]], [2.0, *BOUNDS["temperature"]]])
+    model.parameters.update_prior(**STARTING_PRIOR)
+    optimiser = FminBound(
+        model=model,
+        data=data.groupby("ppt"),
+        minimisation=minimise.LogLikelihood.bernoulli,
+        prior=True,
+        number_of_starts=2,
+        ppt_identifier="ppt",
+        parallel=False,
+        display=False,
+        approx_grad=True,
+        **options,
+    )
+    return VariationalBayes(
+        optimiser=optimiser,
+        iteration=iterations,
+        chain=1,
+        tolerance_lme=0.0,
+        tolerance_param=0.0,
+        hyperpriors={"a0": np.array([p["mean"] for p in STARTING_PRIOR.values()]), "b": 1, "v": 0.5, "s": np.repeat(0.01, 2)},
+        quiet=True,
+        start=start,
+    )
+
+
+class Counter:
+    """Counts the objective evaluations of `FminBound` by what asked for them."""
+
+    def __init__(self):
+        self.phase = "optimiser"
+        self.counts = {"optimiser": 0, "hessian": 0, "evaluate_fit": 0}
+        self.originals = (fmin_module.objective, fmin_module.numerical_hessian, fmin_module.finite_difference_hessian,
+                          fmin_module.evaluate_fit)
+
+    def __enter__(self):
+        objective, numerical_hessian, finite_difference_hessian, evaluate_fit = self.originals
+
+        def counted_objective(*args, **kwargs):
+            self.counts[self.phase] += 1
+            return objective(*args, **kwargs)
+
+        def counted(hessian):
+            def counted_hessian(*args, **kwargs):
+                self.phase = "hessian"
+                try:
+                    return hessian(*args, **kwargs)
+                finally:
+                    self.phase = "optimiser"
+
+            return counted_hessian
+
+        def counted_evaluate_fit(*args, **kwargs):
+            self.counts["evaluate_fit"] += 1  # one model run per call
+            return evaluate_fit(*args, **kwargs)
+
+        fmin_module.objective = counted_objective
+        fmin_module.numerical_hessian = counted(numerical_hessian)
+        fmin_module.finite_difference_hessian = counted(finite_difference_hessian)
+        fmin_module.evaluate_fit = counted_evaluate_fit
+        return self
+
+    def __exit__(self, *exc):
+        (fmin_module.objective, fmin_module.numerical_hessian, fmin_module.finite_difference_hessian,
+         fmin_module.evaluate_fit) = self.originals
+
+
+def profile(vb):
+    profiler = cProfile.Profile()
+    profiler.enable()
+    vb.optimise()
+    profiler.disable()
+    stats = pstats.Stats(profiler)
+    stats.calc_callees()
+    return stats
+
+
+def find(stats, filename, function):
+    """The (file, line, name) key of a profiled function, by the basename of its file."""
+    for key in stats.stats:
+        if key[2] == function and pathlib.Path(key[0]).name == filename:
+            return key
+    return None
+
+
+def inclusive(stats, filename, function):
+    key = find(stats, filename, function)
+    return stats.stats[key][3] if key else 0.0
+
+
+def inclusive_from(stats, caller, filename, function):
+    """Seconds spent in `function` when called from `caller`."""
+    callees = stats.all_callees.get(caller, {})
+    return sum(value[3] for key, value in callees.items() if key[2] == function and pathlib.Path(key[0]).name == filename)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--data", help="a dataset CSV; simulated if not given")
+    parser.add_argument("--iterations", type=int, default=5)
+    parser.add_argument("--label", default="hierarchical-profile")
+    parser.add_argument("--hessian", default="numdifftools", help="the `hessian` option of FminBound")
+    parser.add_argument("--start", default="random", help="the `start` option of VariationalBayes")
+    args = parser.parse_args()
+
+    warnings.simplefilter("ignore")
+    np.seterr(all="ignore")
+    data = load(args.data) if args.data else simulate()
+    participants, iterations = data.ppt.nunique(), args.iterations
+    options = {"hessian": args.hessian, "start": args.start}
+    print(f"cpm {cpm.__version__}: {participants} participants, {data.trial.max()} trials, {iterations} iterations, 2 starts, {options}")
+
+    # the compiled model and everything else that loads on the first call
+    setup(data[data.ppt.isin(data.ppt.unique()[:5])], 1, **options).optimise()
+
+    # once without the profiler, for the time
+    np.random.seed(2026)
+    start = time.perf_counter()
+    with Counter() as counter:
+        setup(data, iterations, **options).optimise()
+    seconds = time.perf_counter() - start
+    per = participants * iterations
+    counts = {k: v / per for k, v in counter.counts.items()}
+
+    np.random.seed(2026)
+    stats = profile(setup(data, iterations, **options))
+    total = inclusive(stats, "variational.py", "optimise")
+    objective_key = find(stats, "optimisers.py", "objective")
+    rows = [
+        ("fit", "seconds without profiler", seconds),
+        ("fit", "ms per participant and iteration", seconds * 1e3 / per),
+        ("evaluations per participant and iteration", "optimiser", counts["optimiser"]),
+        ("evaluations per participant and iteration", "hessian", counts["hessian"]),
+        ("evaluations per participant and iteration", "evaluate_fit", counts["evaluate_fit"]),
+        ("share of profiled time", "fmin_l_bfgs_b", inclusive(stats, "_lbfgsb_py.py", "fmin_l_bfgs_b") / total),
+        ("share of profiled time", "hessian", (inclusive(stats, "optimisers.py", "numerical_hessian")
+                                               + inclusive(stats, "optimisers.py", "finite_difference_hessian")) / total),
+        ("share of profiled time", "evaluate_fit", inclusive(stats, "optimisers.py", "evaluate_fit") / total),
+        ("share of profiled time", "update_population", inclusive(stats, "variational.py", "update_population") / total),
+        ("share of profiled time", "get_lme", inclusive(stats, "variational.py", "get_lme") / total),
+        ("share of profiled time", "objective (all callers)", inclusive(stats, "optimisers.py", "objective") / total),
+        ("share of objective", "SessionWrapper.run", inclusive_from(stats, objective_key, "session.py", "run")),
+        ("share of objective", "  of which the RLRW model function", inclusive(stats, "reinforcement_learning.py", "__call__")),
+        ("share of objective", "Wrapper.reset", inclusive_from(stats, objective_key, "session.py", "reset")),
+        ("share of objective", "deepcopy", inclusive_from(stats, objective_key, "copy.py", "deepcopy")),
+        ("share of objective", "loss", inclusive_from(stats, objective_key, "minimise.py", "bernoulli")),
+        ("share of objective", "Parameters.PDF", inclusive_from(stats, objective_key, "parameters.py", "PDF")),
+    ]
+    objective_seconds = inclusive(stats, "optimisers.py", "objective")
+    rows = [(s, i, v / objective_seconds if s == "share of objective" else v) for s, i, v in rows]
+    table = pd.DataFrame(rows, columns=["section", "item", "value"])
+    table["value"] = table["value"].round(4)
+    print(table.to_string(index=False))
+
+    RESULTS.mkdir(exist_ok=True)
+    table.to_csv(RESULTS / f"{args.label}.csv", index=False)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/results/hierarchical-comparison-individual.csv

@@ -0,0 +1,31 @@
+config,method,parameter,individual_rmse
+0.26.0.dev1 (Phase 0 commit),cpm-empirical,alpha,"0.085 [0.073, 0.097]"
+0.26.0.dev1 (Phase 0 commit),cpm-empirical,temperature,"0.763 [0.640, 0.895]"
+0.26.0.dev1 (Phase 0 commit),cpm-variational,alpha,"0.085 [0.074, 0.097]"
+0.26.0.dev1 (Phase 0 commit),cpm-variational,temperature,"0.763 [0.642, 0.895]"
+"new, defaults",cpm-empirical,alpha,"0.085 [0.073, 0.097]"
+"new, defaults",cpm-empirical,temperature,"0.763 [0.640, 0.895]"
+"new, defaults",cpm-variational,alpha,"0.085 [0.074, 0.097]"
+"new, defaults",cpm-variational,temperature,"0.763 [0.642, 0.895]"
+"new, hessian",cpm-empirical,alpha,"0.084 [0.073, 0.096]"
+"new, hessian",cpm-empirical,temperature,"0.760 [0.637, 0.892]"
+"new, hessian",cpm-variational,alpha,"0.085 [0.074, 0.096]"
+"new, hessian",cpm-variational,temperature,"0.758 [0.636, 0.889]"
+"new, start",cpm-empirical,alpha,"0.085 [0.074, 0.097]"
+"new, start",cpm-empirical,temperature,"0.760 [0.637, 0.893]"
+"new, start",cpm-variational,alpha,"0.085 [0.074, 0.097]"
+"new, start",cpm-variational,temperature,"0.762 [0.640, 0.892]"
+"new, start + hessian",cpm-empirical,alpha,"0.084 [0.073, 0.096]"
+"new, start + hessian",cpm-empirical,temperature,"0.759 [0.635, 0.889]"
+"new, start + hessian",cpm-variational,alpha,"0.085 [0.074, 0.096]"
+"new, start + hessian",cpm-variational,temperature,"0.755 [0.635, 0.884]"
+"new, start + hessian, 1 start",cpm-empirical,alpha,"0.085 [0.073, 0.096]"
+"new, start + hessian, 1 start",cpm-empirical,temperature,"0.758 [0.635, 0.890]"
+"new, start + hessian, 1 start",cpm-variational,alpha,"0.085 [0.074, 0.097]"
+"new, start + hessian, 1 start",cpm-variational,temperature,"0.756 [0.634, 0.886]"
+comparison (cpm-empirical),cpm-empirical,alpha,"0.085 [0.073, 0.097]"
+comparison (cpm-empirical),cpm-empirical,temperature,"0.763 [0.639, 0.895]"
+comparison (cpm-variational),cpm-variational,alpha,"0.085 [0.074, 0.097]"
+comparison (cpm-variational),cpm-variational,temperature,"0.764 [0.642, 0.895]"
+comparison (cbm),cbm,alpha,"0.085 [0.074, 0.097]"
+comparison (cbm),cbm,temperature,"0.765 [0.634, 0.906]"

benchmarks/results/hierarchical-comparison-population.csv

@@ -0,0 +1,31 @@
+config,method,parameter,rmse_mean,rmse_sd
+0.26.0.dev1 (Phase 0 commit),cpm-empirical,alpha,"0.013 [0.008, 0.018]","0.026 [0.016, 0.035]"
+0.26.0.dev1 (Phase 0 commit),cpm-empirical,temperature,"0.199 [0.129, 0.262]","0.212 [0.136, 0.281]"
+0.26.0.dev1 (Phase 0 commit),cpm-variational,alpha,"0.012 [0.008, 0.015]","0.028 [0.020, 0.035]"
+0.26.0.dev1 (Phase 0 commit),cpm-variational,temperature,"0.184 [0.122, 0.240]","0.220 [0.134, 0.296]"
+"new, defaults",cpm-empirical,alpha,"0.013 [0.008, 0.018]","0.026 [0.016, 0.035]"
+"new, defaults",cpm-empirical,temperature,"0.199 [0.129, 0.262]","0.212 [0.136, 0.281]"
+"new, defaults",cpm-variational,alpha,"0.012 [0.008, 0.015]","0.028 [0.020, 0.035]"
+"new, defaults",cpm-variational,temperature,"0.184 [0.122, 0.240]","0.220 [0.134, 0.296]"
+"new, hessian",cpm-empirical,alpha,"0.012 [0.008, 0.017]","0.021 [0.011, 0.028]"
+"new, hessian",cpm-empirical,temperature,"0.197 [0.128, 0.261]","0.191 [0.109, 0.256]"
+"new, hessian",cpm-variational,alpha,"0.012 [0.008, 0.015]","0.025 [0.017, 0.031]"
+"new, hessian",cpm-variational,temperature,"0.180 [0.121, 0.235]","0.170 [0.105, 0.225]"
+"new, start",cpm-empirical,alpha,"0.013 [0.008, 0.017]","0.027 [0.016, 0.037]"
+"new, start",cpm-empirical,temperature,"0.199 [0.130, 0.262]","0.211 [0.136, 0.281]"
+"new, start",cpm-variational,alpha,"0.012 [0.009, 0.015]","0.029 [0.021, 0.036]"
+"new, start",cpm-variational,temperature,"0.186 [0.124, 0.242]","0.211 [0.133, 0.282]"
+"new, start + hessian",cpm-empirical,alpha,"0.012 [0.008, 0.017]","0.021 [0.012, 0.028]"
+"new, start + hessian",cpm-empirical,temperature,"0.200 [0.129, 0.262]","0.191 [0.108, 0.256]"
+"new, start + hessian",cpm-variational,alpha,"0.012 [0.008, 0.015]","0.025 [0.017, 0.031]"
+"new, start + hessian",cpm-variational,temperature,"0.183 [0.123, 0.236]","0.170 [0.105, 0.225]"
+"new, start + hessian, 1 start",cpm-empirical,alpha,"0.012 [0.008, 0.016]","0.022 [0.013, 0.028]"
+"new, start + hessian, 1 start",cpm-empirical,temperature,"0.200 [0.130, 0.262]","0.180 [0.094, 0.249]"
+"new, start + hessian, 1 start",cpm-variational,alpha,"0.011 [0.009, 0.014]","0.025 [0.018, 0.031]"
+"new, start + hessian, 1 start",cpm-variational,temperature,"0.183 [0.122, 0.236]","0.171 [0.107, 0.226]"
+comparison (cpm-empirical),cpm-empirical,alpha,"0.013 [0.008, 0.018]","0.026 [0.016, 0.035]"
+comparison (cpm-empirical),cpm-empirical,temperature,"0.199 [0.127, 0.259]","0.213 [0.133, 0.279]"
+comparison (cpm-variational),cpm-variational,alpha,"0.012 [0.008, 0.015]","0.028 [0.021, 0.035]"
+comparison (cpm-variational),cpm-variational,temperature,"0.184 [0.123, 0.240]","0.205 [0.126, 0.279]"
+comparison (cbm),cbm,alpha,"0.023 [0.018, 0.029]","0.022 [0.015, 0.028]"
+comparison (cbm),cbm,temperature,"0.279 [0.183, 0.362]","0.176 [0.115, 0.234]"