diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..a254a7f --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,97 @@ +name: tests + +# cpm has to give the same results with and without numba, so the test suite +# runs in both setups: with numba installed, and without it. +on: + pull_request: + push: + branches: [main] + workflow_dispatch: + +concurrency: + group: tests-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + test: + name: ${{ matrix.os }}, Python ${{ matrix.python }}, ${{ matrix.numba && 'numba' || 'no numba' }} + runs-on: ${{ matrix.os }} + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest] + python: ["3.11", "3.12", "3.13", "3.14"] + numba: [true, false] + include: + - { os: windows-latest, python: "3.14", numba: true } + - { os: windows-latest, python: "3.14", numba: false } + - { os: macos-latest, python: "3.14", numba: true } + - { os: macos-latest, python: "3.14", numba: false } + steps: + - uses: actions/checkout@v7 + + - uses: actions/setup-python@v7 + with: + python-version: ${{ matrix.python }} + cache: pip + + - name: Install cpm + run: pip install -e "${{ matrix.numba && '.[numba]' || '.' }}" pytest + + - name: Check that numba is used + if: matrix.numba + run: python -c "from cpm.core import _jit; assert _jit.JIT_ENABLED" + + - name: Run the tests + run: python -m pytest -q + + - name: Run the kernel and application tests with NUMBA_DISABLE_JIT=1 + if: matrix.numba + env: + NUMBA_DISABLE_JIT: "1" + run: python -m pytest -q test/models/test_kernels.py test/applications/test_backends.py + + lowest: + name: lowest supported versions, Python 3.11 + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - uses: actions/checkout@v7 + + - uses: actions/setup-python@v7 + with: + python-version: "3.11" + + - name: Install cpm with the lowest versions pyproject.toml allows + run: | + pip install uv + uv pip install --system --resolution lowest-direct -e ".[numba]" "pytest>=8" + + - name: Run the tests + run: python -m pytest -q + + prerelease: + # the fast priors in cpm.generators._fast_priors follow scipy's internals, + # so pre-releases show early when a release would change them + name: pre-releases of NumPy, SciPy and pandas + runs-on: ubuntu-latest + timeout-minutes: 30 + continue-on-error: true + steps: + - uses: actions/checkout@v7 + + - uses: actions/setup-python@v7 + with: + python-version: "3.14" + + - name: Install cpm and the pre-releases + run: | + pip install -e . pytest + pip install --pre --upgrade numpy scipy pandas + + - name: Run the tests + run: python -m pytest -q diff --git a/CHANGELOG.md b/CHANGELOG.md index 45679af..dd27499 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,50 +5,75 @@ 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.25.7.dev0] - 2026-08-13 +## [0.26.0.dev0] - 2026-09-28 ### Added -- Added test units for `cpm.hierarchical.EmpiricalBayes` -- Added test units for `cpm.hierarchical.VariationalBayes` -- Added `number_of_starts` and `initial_guess_supplied` attributes to `cpm.optimisation.Fmin`, `cpm.optimisation.FminBound`, `cpm.optimisation.Minimize` and `cpm.optimisation.Bads`, so that the number of starts and the origin of the initial guesses are recoverable from a constructed or fitted optimiser. Previously `number_of_starts` was consumed in `__init__` and discarded, recoverable only as `initial_guess.shape[0]`, and there was no record of whether the guesses were supplied by the user or drawn from the parameter bounds - which meant a fit started from a fixed point was indistinguishable from one started from random restarts (#83) -- Added test units for the new optimiser provenance attributes, covering all four optimisers -- Added a `cpm.optimisation.Bads` smoke test that runs a real fit and asserts it reaches `scipy.optimize.minimize` through pybads and gpyreg. `Bads.optimise()` was previously never executed by the test suite, so breakage in that chain - such as a SciPy or NumPy release changing an API it depends on - went undetected -- Added `log_likelihood` and `log_prior` to the fitted output of all optimisers when run with `prior=True`, so that `fun` (the negative log posterior) can be split into its two terms ([#28](https://github.com/DevComPsy/cpm/issues/28)) -- Added a `metrics` argument to all optimisers, which evaluates user-supplied goodness-of-fit metrics (such as `PenalisedLikelihoods.BIC`) at the optimum and adds them to `export()` ([#28](https://github.com/DevComPsy/cpm/issues/28)) -- Added `cpm.applications.reinforcement_learning.HybridMBMF`, the 6-parameter hybrid model-based / model-free model for the two-step task (Kool et al., 2016; Smid et al., 2022) -- Added `cpm.models.learning.SARSATrace`, a SARSA learning rule with an eligibility trace for two-stage tasks +- Added test units for `cpm.hierarchical.EmpiricalBayes` and `cpm.hierarchical.VariationalBayes` +- Added `number_of_starts` and `initial_guess_supplied` attributes to all optimisers, recording how many starts a fit used and whether the initial guesses were supplied ([#83](https://github.com/DevComPsy/cpm/issues/83)) +- Added test units for the new optimiser attributes, and a smoke test that runs a real `cpm.optimisation.Bads` fit +- Added `log_likelihood` and `log_prior` to the output of all optimisers fitted with `prior=True` ([#28](https://github.com/DevComPsy/cpm/issues/28)) +- Added a `metrics` argument to all optimisers, for goodness-of-fit metrics such as `PenalisedLikelihoods.BIC` evaluated at the optimum ([#28](https://github.com/DevComPsy/cpm/issues/28)) +- Added `cpm.applications.reinforcement_learning.HybridMBMF`, the hybrid model-based / model-free model of the two-step task (Kool et al., 2016; Smid et al., 2022) +- Added `cpm.models.learning.SARSATrace`, a SARSA learning rule with an eligibility trace - Added test units for `HybridMBMF` and `SARSATrace` -- Added a two-step task example notebook replicating Smid et al. (2022) -- Added `cpm.datasets.load_two_step_data`, which loads the adult two-step task data of Smid et al. (2022) that ship with the package +- Added a two-step task example replicating Smid et al. (2022), with its data in `cpm.datasets.load_two_step_data` +- Added `cpm.generators.SessionWrapper`, a `Wrapper` for models that compute all trials of a participant at once +- Added `cpm.models.kernels`, the formulas of the `cpm.models` classes as functions that numba can compile, used by the classes and the built-in applications +- Added numba as an optional dependency (`pip install "cpm-toolbox[numba]"`) that compiles the built-in applications; without it, they give the same results as plain Python +- Added the numba install option to the installation guide, a how-to guide on numba, and troubleshooting entries +- Added a how-to guide on speeding up your own model with `SessionWrapper`, with a test that runs its example +- Added test units for `SessionWrapper`, `cpm.models.kernels`, the fast priors, and the built-in applications with and without numba +- Added a benchmark suite (`benchmarks/run.py` for one evaluation, `benchmarks/hierarchical.py` for the hierarchical tutorials) and `scripts/local_tests.py`, which runs the tests with and without numba +- Added a GitHub Actions workflow that runs the tests with and without numba, on Python 3.11 to 3.14 ### Changed -- Rebuilt the documentation with Sphinx and the PyData theme, replacing MkDocs. The site now separates the API reference from the tutorials (a numbered learning path from fitting a first model to hierarchical estimation), a gallery of research examples, and how-to guides; adds a quickstart and core concept pages; documents every public object; and links API pages to the tutorials that use them. All tutorials and examples were rewritten or revised and re-run on the current version. Pages of the old site redirect to their new locations -- Converted all docstrings from Markdown to reStructuredText (NumPy style), fixed incorrect examples in the docstrings of `RLRW`, `EmpiricalBayes`, `ProspectUtility`, `bin_ratings`, `count_trials` and `convert_to_RLRW`, added an example to `VariationalBayes`, and corrected the column list of `load_bandit_data`. Docstrings with LaTeX are now raw strings, which fixes a `SyntaxWarning` and a corrupted equation in the docstring of `RLRW` -- Restricted package discovery to `cpm`, so that files under `docs/` are no longer packaged, and added `docs` and `notebooks` optional dependencies -- Significantly improved the performance of `cpm.hierarchical.EmpiricalBayes` and `cpm.hierarchical.VariationalBayes` for long or multi-chain EM runs, by buffering results instead of concatenating on every iteration and vectorising the Hessian inversions +- Rebuilt the documentation with Sphinx and the PyData theme, with a separate API reference, tutorials, examples and how-to guides; pages of the old site redirect +- Converted all docstrings to reStructuredText (NumPy style), and fixed incorrect examples and a `SyntaxWarning` from LaTeX in docstrings +- Restricted package discovery to `cpm`, and added `docs` and `notebooks` optional dependencies +- Sped up long and multi-chain runs of `cpm.hierarchical.EmpiricalBayes` and `cpm.hierarchical.VariationalBayes` +- Made every model cheaper to evaluate: copies of a `cpm.generators.Value` share their prior, and the built-in priors are evaluated without scipy overhead, with identical results +- `cpm.generators.Value.update_prior()` now replaces the prior instead of changing it in place; changing `value.prior.kwds` directly now affects every copy of the `Value` +- `cpm.generators.Wrapper.reset()` now keeps the current priors and bounds of the parameters, instead of restoring the ones the model was created with +- Made the built-in applications (`RLRW`, `HybridMBMF`, `PTSM`, `PTSM1992`, `PTSM2025`) compute all trials at once: 90-220 times faster with numba, 11-23 times without, with the same interface and results +- Made per-trial `cpm.generators.Wrapper` models about twice as fast, by reading trials without `DataFrame.iloc` +- Made the classes in `cpm.models` faster, up to 17 times for large inputs, with identical results +- Made `LogLikelihood.bernoulli` and `LogLikelihood.continuous` 2.5-3 times faster, with identical results +- Made `cpm.applications.signal_detection.EstimatorMetaD` about 10 times faster, with identical results +- `cpm.models.decision.Softmax` no longer overflows: it returns the correct policy where it returned NaN and warned +- `cpm.models.activation.ProspectUtility.weights` and `.utilities` are float arrays instead of object arrays, and the `simulation` records of `HybridMBMF` hold NumPy scalars instead of Python numbers +- `PTSM2025` always uses its power utility: replacing `parameters.utility_curvature` on a model no longer changes it +- Required Python 3.11 or later (previously `>3.11.0`, which excluded 3.11.0), removed the PyPy classifier, and tagged the wheel for Python 3 only ### Fixed -- Fixed `cpm.optimisation.Bads` emitting a `DeprecationWarning` on every GP fit, which would become an error on Python 3.16, by passing `gp_fixed_mean` to pybads as a NumPy boolean ([#88](https://github.com/DevComPsy/cpm/issues/88)) -- Fixed `cpm.optimisation.Bads` failing with `ValueError: setting an array element with a sequence` under NumPy 2 by requiring `gpyreg>=1.2.1` -- Fixed `cpm.optimisation.FminBound` raising `TypeError` on SciPy 1.18.0 and later, which removed the `disp` and `iprint` options of L-BFGS-B; they are now passed only where SciPy still accepts them -- Fixed `cpm.hierarchical.VariationalBayes.ttest` raising a `NameError` when `null` was passed as a `pandas.DataFrame`, due to referencing an undefined variable from the wrong branch -- Fixed `cpm.hierarchical.VariationalBayes.lmes` recording the same, fully-grown list of log model evidence values for every iteration of a chain instead of a snapshot of that iteration's value, due to appending a reference to a still-mutating list -- Removed a dead, always-zero `mean_errorbar` column from `cpm.hierarchical.VariationalBayes.hyperparameters` that was left behind by a column-naming mismatch (values were actually being written to a separate `mean_se` column) -- Fixed `cpm.generators.Value` with `prior="uniform"` spanning `[lower, lower + upper]` instead of `[lower, upper]`, due to passing `upper` rather than `upper - lower` as the scale, which biased every log prior and let `Parameters.sample()` draw values above the upper bound -- Fixed `cpm.generators.Parameters.sample()` raising `AttributeError` when a parameter was `None`, due to a missing `None` guard that `free()`, `bounds()` and `PDF()` already had -- Fixed callable arguments to `cpm.generators.Parameters` leaking onto every other `Parameters` instance in the process, due to being set as static methods on the class rather than stored on the instance -- Fixed `cpm.models.learning.SeparableRule.error` never being populated, which left it at zeros and made `noisy_learning_rule()` a no-op because the noise scales with the error; it now holds the prediction error for each outcome-stimulus pair. Also fixed `DeltaRule` and `SeparableRule` sizing `error` before promoting 1D weights to 2D -- Fixed `cpm.generators.Wrapper.reset()` misassigning an array of parameter values whenever a non-free attribute (such as an initial state) was declared before a free parameter, due to mapping the array onto all keys instead of `Parameters.free()` -- Fixed `cpm.generators.Simulator` not raising its intended `TypeError` for an ungrouped `pandas.DataFrame`, due to type-checking a boolean flag instead of the data; the mistake previously surfaced later as a misleading group-count `ValueError` -- Fixed `cpm.generators.Simulator` raising `AttributeError` when parameters were passed as a `pandas.DataFrame` with grouped data, due to looking up parameter rows by group key instead of position; `cpm.core.data.unpack_participants` also returned every remaining row instead of a single participant's row -- Fixed `cpm.hierarchical.EmpiricalBayes.diagnostics()` and `cpm.hierarchical.VariationalBayes.diagnostics()` raising `IndexError` for models with three or more free parameters, due to hardcoded trace limits for two parameters; `convergence_diagnostics_plots` now takes the limits from the model's parameter bounds via a new `bounds` argument -- Fixed a clean install of `cpm-toolbox` raising `ModuleNotFoundError` on `diagnostics()`, due to `matplotlib` being imported by `cpm.core.diagnostics` but missing from the declared dependencies -- Fixed `cpm.utils.data` (and `cpm.utils.metad`) not being reachable after `import cpm`, due to `cpm/utils/__init__.py` not importing its submodules -- Fixed `cpm.generators.Simulator` raising `ValueError` for a single parameter set given as a `Parameters` object, `dict` or `pandas.Series`, due to checking the length of the input rather than of the cast parameters; a `Parameters` object is now used as-is for every participant, as documented, instead of drawing a new prior sample for each -- Fixed `cpm.core.diagnostics.gelman_rubin` and `cpm.core.diagnostics.psrf` failing on every call, due to the removed `pandas.DataFrame.append` and a nonexistent `unique_` attribute; `gelman_rubin` also computed the statistic incorrectly (misaligned chains, swapped within- and between-chain variances) and now reports it for both the mean and sd of each parameter, comparing chains over their common iterations +- Fixed `cpm.hierarchical.EmpiricalBayes` and `cpm.hierarchical.VariationalBayes` discarding the estimated population priors after the first evaluation of each fit; hierarchical results change +- Fixed the random starting priors of later chains of `EmpiricalBayes` and `VariationalBayes` being length-one arrays +- Fixed `numpy.random.seed` not reproducing the later chains of `EmpiricalBayes` and `VariationalBayes`, whose starting priors now also respect the lower bound and work with infinite bounds +- Fixed `cpm.optimisation.Bads` emitting a `DeprecationWarning` on every GP fit ([#88](https://github.com/DevComPsy/cpm/issues/88)) +- Fixed `cpm.optimisation.Bads` failing under NumPy 2, by requiring `gpyreg>=1.2.1` +- Fixed installs with pandas, matplotlib or pybads releases that do not work with NumPy 2, by requiring `pandas>=2.2.2`, `matplotlib>=3.8.4` and `pybads>=1.0.5` +- Fixed `cpm.optimisation.FminBound` raising `TypeError` on SciPy 1.18.0 and later +- Fixed `cpm.hierarchical.VariationalBayes.ttest` raising `NameError` when `null` is a `pandas.DataFrame` +- Fixed `cpm.hierarchical.VariationalBayes.lmes` repeating the final list of model evidences for every iteration +- Removed a dead, always-zero `mean_errorbar` column from `cpm.hierarchical.VariationalBayes.hyperparameters` +- Fixed `cpm.generators.Value` with `prior="uniform"` spanning `[lower, lower + upper]` instead of `[lower, upper]` +- Fixed `cpm.generators.Parameters.sample()` raising `AttributeError` for a parameter that is `None` +- Fixed callables passed to `cpm.generators.Parameters` leaking onto every other `Parameters` instance +- Fixed `cpm.models.learning.SeparableRule.error` never being filled, which made `noisy_learning_rule()` a no-op, and the `error` of `DeltaRule` and `SeparableRule` having the wrong shape for 1D weights +- Fixed `cpm.generators.Wrapper.reset()` misassigning an array of parameter values when a state is declared before a free parameter +- Fixed `numpy.asarray(value, dtype=...)` raising `TypeError` for a `cpm.generators.Value` +- Fixed `cpm.models.activation.ProspectUtility` with `weighting="prelec"` failing for options with several outcomes +- Fixed `cpm.applications.decision_making.PTSM2025` stopping a fit with `ValueError` when its exponentials overflowed +- Fixed `cpm.models.learning.HumbleTeacher` raising `IndexError` for 1D weights +- Fixed `cpm.generators.Simulator` not raising its intended `TypeError` for an ungrouped `pandas.DataFrame` +- Fixed `cpm.generators.Simulator` raising `AttributeError` for parameters in a `pandas.DataFrame` with grouped data +- Fixed `cpm.generators.Simulator` raising `ValueError` for a single parameter set given as `Parameters`, `dict` or `pandas.Series` +- Fixed `diagnostics()` of `EmpiricalBayes` and `VariationalBayes` failing for three or more free parameters; `convergence_diagnostics_plots` has a new `bounds` argument +- Fixed a clean install raising `ModuleNotFoundError` in `diagnostics()`, by adding `matplotlib` to the dependencies +- Fixed `cpm.utils.data` and `cpm.utils.metad` not being reachable after `import cpm` +- Fixed `cpm.core.diagnostics.gelman_rubin` and `cpm.core.diagnostics.psrf` failing on every call, and `gelman_rubin` computing the statistic incorrectly ## [0.25.6] - 2026-04-15 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 50363cb..d8b1c6a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -25,6 +25,14 @@ pip install -e . pip install -e ".[docs]" ``` +5. (Optional) Install numba, which compiles the built-in models (`cpm.applications`) and the kernels in `cpm.models.kernels`: + +```bash +pip install -e ".[numba]" +``` + +cpm has to work, with the same results, with and without numba, so changes to the models are tested both ways (see [Testing](#testing)). + ## Development workflow (with direct access to repository) 1. Create a branch from `main`. @@ -75,6 +83,16 @@ pytest If you changed specific functionality, also run targeted tests first (for example, under `test/models/` or `test/optimisation/`). +If you changed the models (`cpm.models`, `cpm.applications`) or anything they use, also run the tests with and without numba, and the benchmark suite: + +```bash +python scripts/local_tests.py +``` + +It runs the test suite with numba (if it is installed), with `CPM_DISABLE_JIT=1` (as if numba were not installed), and with `NUMBA_DISABLE_JIT=1`, and then `benchmarks/run.py`. + +numba caches compiled code in the `__pycache__` folders under `cpm/`, and only recompiles a function when its own file changes. After changing `cpm/models/kernels.py` (yourself, or by pulling), the compiled loops in `cpm/applications/_sessions.py` still use the old kernels until you delete the cached `*.nbi` and `*.nbc` files; `scripts/local_tests.py` deletes them before it runs. + ## Documentation The documentation is built with [Sphinx](https://www.sphinx-doc.org/) and the [PyData theme](https://pydata-sphinx-theme.readthedocs.io/). diff --git a/README.md b/README.md index 89a346a..e1f9fe4 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,12 @@ To install the package, run the following command: pip install cpm-toolbox ``` +To fit the built-in models many times faster, also install the optional [numba](https://numba.pydata.org/) dependency, which compiles them; nothing else changes: + +```bash +pip install "cpm-toolbox[numba]" +``` + Once the package is installed, you can import it in your Python code: ```python diff --git a/benchmarks/cases.py b/benchmarks/cases.py new file mode 100644 index 0000000..cd15444 --- /dev/null +++ b/benchmarks/cases.py @@ -0,0 +1,132 @@ +""" +The benchmark cases: every built-in application of cpm, on one participant. + +Each case builds a model for a given backend and returns it together with the +observed data and the parameter values at which the objective is evaluated. + +- "python" and "numba": the application, which computes all trials at once, as + plain Python or compiled with numba. Installing numba makes "numba" the default. +- "trial": the application's per-trial `model` function in a plain per-trial + `Wrapper`, which is how users extend an application trial by trial. + +With cpm versions before the applications computed all trials at once, only +"trial" is available, and it is the application itself. +""" + +import warnings + +import numpy as np +import pandas as pd + +import cpm +from cpm.generators import Wrapper +from cpm.datasets import load_bandit_data, load_risky_choices + + +def bandit(): + data = load_bandit_data() + data = data[data.ppt == 1].reset_index(drop=True) + data["observed"] = data["response"] + return data + + +def two_step(n=200, seed=0): + rng = np.random.default_rng(seed) + action = rng.integers(0, 2, n) + common = rng.random(n) < 1.0 # deterministic transitions in the novel task + return pd.DataFrame( + { + "s1": rng.integers(0, 2, n), + "stimuli_first": rng.integers(0, 2, n), + "action": action, + "s2": np.where(common, 1 - action, action), + "reward": rng.integers(0, 10, n) / 9, + "reward_0": rng.integers(0, 10, n) / 9, + "reward_1": rng.integers(0, 10, n) / 9, + "observed": action, + } + ) + + +def risky(): + data = load_risky_choices() + data = data[data.ppt == 1].reset_index(drop=True) + data["observed"] = data["choice"].astype(int) + return data + + +PT_SETTINGS = { + "alpha": [0.8, 1e-2, 5.0], + "lambda_loss": [1.6, 1e-2, 5.0], + "gamma": [0.7, 1e-2, 5.0], + "temperature": [3.0, 1e-2, 15.0], + "beta": [0.9, 0.0, 5.0], + "delta": [0.6, 1e-2, 5.0], + "eta": [0.1, -0.49, 0.49], + "phi_gain": [0.2, -10.0, 10.0], + "phi_loss": [-0.3, -10.0, 10.0], +} + + +def _application(module, name): + return getattr(getattr(cpm.applications, module), name) + + +def _on_backend(model, backend): + """The model on `backend`, or None if this version of cpm cannot run it there.""" + session = getattr(model, "_session_model", None) + if backend == "trial": + if session is None: + return model # an older cpm: the application is a per-trial Wrapper + return Wrapper(model=model.model, data=model.data, parameters=model.parameters) + if session is None: + return None + session.backend = backend + return model + + +def build(case, backend="trial"): + """ + The model, observed data and parameter values of a benchmark case. + + Returns None if the installed cpm has no model for this backend. + """ + warnings.simplefilter("ignore") + if case == "RLRW": + cls = _application("reinforcement_learning", "RLRW") + data = bandit() + kwargs = dict(dimensions=4, parameters_settings=[[0.3, 0, 1], [4, 0, 10]]) + elif case == "HybridMBMF": + cls = _application("reinforcement_learning", "HybridMBMF") + data = two_step() + kwargs = dict( + parameters_settings=[ + [2, 0, 5], [0.4, 0, 1], [0.6, 0, 1], [0.5, 0, 1], [0.3, -5, 5], [-0.2, -5, 5] + ] + ) + elif case in ("PTSM", "PTSM1992", "PTSM2025"): + cls = _application("decision_making", case) + data = risky() + kwargs = dict(parameters_settings=PT_SETTINGS) + else: + raise ValueError(case) + if backend == "numba" and not _jit_enabled(): + return None + model = _on_backend(cls(data=data, **kwargs), backend) + if model is None: + return None + observed = data["observed"].to_numpy() + x = np.array([model.parameters[k].value for k in model.parameters.free()], dtype=float) + return model, observed, x + + +def _jit_enabled(): + try: + from cpm.core import _jit + except ImportError: + return False + return _jit.JIT_ENABLED + + +CASES = ["RLRW", "HybridMBMF", "PTSM", "PTSM1992", "PTSM2025"] +BACKENDS = ["trial", "python", "numba"] diff --git a/benchmarks/hierarchical.py b/benchmarks/hierarchical.py new file mode 100644 index 0000000..4f5f0ea --- /dev/null +++ b/benchmarks/hierarchical.py @@ -0,0 +1,186 @@ +""" +How long does a hierarchical analysis take? The two hierarchical tutorials, timed. + + python benchmarks/hierarchical.py ANALYSIS MODEL OUT + +ANALYSIS is "eb" (docs/tutorials/hierarchical-empirical-bayes.ipynb: separate +maximum-likelihood fits, then empirical Bayes) or "vb" +(docs/tutorials/hierarchical-variational-bayes.ipynb: variational Bayes). Both +simulate 50 participants of the bandit task from a known group distribution +and fit them with `FminBound` (2 starts per participant), exactly as the +tutorials do, with the same random seed. + +MODEL is "tutorial" (the per-trial model the tutorials build from `Softmax` +and `SeparableRule`, in a `Wrapper`) or "rlrw" (the same model as the built-in +`cpm.applications.reinforcement_learning.RLRW`, which computes all trials at +once and is compiled with numba if it is installed). The data are simulated +with the tutorial model in both cases, so both fit the same data. + +Writes the timings and the results (estimates and hyperparameters) to OUT, a +pickle file. +""" + +import pickle +import platform +import sys +import time +import warnings +from functools import partial + +import numpy as np +import pandas as pd + +import cpm +from cpm.datasets import load_bandit_data +from cpm.generators import Parameters, Simulator, Value, Wrapper +from cpm.models.decision import Softmax +from cpm.models.learning import SeparableRule +from cpm.optimisation import FminBound, minimise + +ANALYSIS, MODEL, OUT = sys.argv[1:4] +assert ANALYSIS in ("eb", "vb") and MODEL in ("tutorial", "rlrw") + +np.random.seed(2026) +_ = np.seterr(all="ignore") +warnings.simplefilter("ignore") + + +## --------------------------------------------------------------------------- +## the tutorials' model and data (cell 3 of both notebooks) + + +def model(parameters, trial, generate=False): + values = np.asarray(parameters.values).copy() + stimuli = np.array([trial.arm_left, trial.arm_right]).astype(int) + rewards = np.array([trial.reward_left, trial.reward_right]) + + choice_rule = Softmax(activations=values[stimuli - 1], temperature=parameters.temperature) + choice_rule.compute() + choice = choice_rule.choice() if generate else int(trial.response) + + chosen = np.zeros(4) + chosen[stimuli[choice] - 1] = 1 + update = SeparableRule(weights=values, feedback=[rewards[choice]], input=chosen, alpha=parameters.alpha) + update.compute() + values += update.weights.flatten() + + return { + "policy": choice_rule.policies, + "response": choice, + "values": values, + "dependent": np.array([choice_rule.policies[1]]), + } + + +def make_parameters(alpha_mean, alpha_sd, temperature_mean, temperature_sd): + return Parameters( + alpha=Value(value=0.5, lower=1e-10, upper=1, prior="truncated_normal", args={"mean": alpha_mean, "sd": alpha_sd}), + temperature=Value(value=1, lower=0, upper=10, prior="truncated_normal", args={"mean": temperature_mean, "sd": temperature_sd}), + values=np.array([0.25, 0.25, 0.25, 0.25]), + ) + + +timings = {} +start = time.perf_counter() +data = load_bandit_data() +data = data[data.ppt <= 50].reset_index(drop=True) +data["observed"] = data["response"] +population = make_parameters(alpha_mean=0.6, alpha_sd=0.2, temperature_mean=2, temperature_sd=1) +true_parameters = pd.DataFrame(population.sample(size=data.ppt.nunique())) +simulator = Simulator( + wrapper=Wrapper(model=partial(model, generate=True), parameters=population, data=data[data.ppt == 1]), + parameters=true_parameters, + data=data.groupby("ppt"), +) +simulator.run() +data["response"] = simulator.export()["response"].to_numpy() +data["observed"] = data["response"] +timings["simulate"] = time.perf_counter() - start + + +def make_wrapper(): + """The model with the broad starting priors of the tutorials.""" + if MODEL == "tutorial": + parameters = make_parameters(alpha_mean=0.5, alpha_sd=0.5, temperature_mean=5, temperature_sd=5) + return Wrapper(model=model, parameters=parameters, data=data[data.ppt == 1]) + from cpm.applications.reinforcement_learning import RLRW + + wrapper = RLRW(data=data[data.ppt == 1], dimensions=4, parameters_settings=[[0.5, 1e-10, 1], [1, 0, 10]]) + wrapper.parameters.update_prior(alpha={"mean": 0.5, "sd": 0.5}, temperature={"mean": 5, "sd": 5}) + return wrapper + + +def optimiser(prior): + return FminBound( + model=make_wrapper(), + data=data.groupby("ppt"), + minimisation=minimise.LogLikelihood.bernoulli, + prior=prior, + number_of_starts=2, + ppt_identifier="ppt", + parallel=False, + display=False, + approx_grad=True, + ) + + +results = {"true_parameters": true_parameters} +if ANALYSIS == "eb": + from cpm.hierarchical import EmpiricalBayes + + start = time.perf_counter() + separate = optimiser(prior=False) + separate.optimise() + timings["separate fits"] = time.perf_counter() - start + results["maximum_likelihood"] = pd.DataFrame(separate.parameters) + + start = time.perf_counter() + eb = EmpiricalBayes(optimiser=optimiser(prior=True), iteration=6, chain=2, tolerance=1e-3, quiet=True) + eb.optimise() + timings["empirical Bayes"] = time.perf_counter() - start + results["hyperparameters"] = eb.hyperparameters + results["fit"] = eb.fit +else: + from cpm.hierarchical import VariationalBayes + + start = time.perf_counter() + number_of_parameters = 2 + vb = VariationalBayes( + optimiser=optimiser(prior=True), + iteration=6, + chain=2, + convergence="lme", + tolerance_lme=1e-2, + hyperpriors={ + "a0": np.zeros(number_of_parameters), + "b": 1, + "v": 0.5, + "s": np.repeat(0.01, number_of_parameters), + }, + quiet=True, + ) + vb.optimise() + timings["variational Bayes"] = time.perf_counter() - start + results["hyperparameters"] = vb.hyperparameters + results["fit"] = vb.fit + +try: + from cpm.core import _jit + + numba_used = MODEL == "rlrw" and bool(_jit.JIT_ENABLED) +except ImportError: + numba_used = False +record = { + "analysis": ANALYSIS, + "model": MODEL, + "numba": numba_used, + "cpm": cpm.__version__, + "python": platform.python_version(), + "numpy": np.__version__, + "timings": timings, + "results": results, +} +with open(OUT, "wb") as f: + pickle.dump(record, f) +print(f"{ANALYSIS} {MODEL} (cpm {cpm.__version__}, numba {'on' if numba_used else 'off'}): " + + ", ".join(f"{k} {v:.1f} s" for k, v in timings.items())) diff --git a/benchmarks/results/hierarchical-report.csv b/benchmarks/results/hierarchical-report.csv new file mode 100644 index 0000000..12a612d --- /dev/null +++ b/benchmarks/results/hierarchical-report.csv @@ -0,0 +1,9 @@ +analysis,variant,separate_fits_s,hierarchical_fit_s,iterations,seconds_per_iteration,fitting_total_s +empirical Bayes,"main, tutorial model",37.02,458.9981,10,45.8998,496.0181 +empirical Bayes,"branch, tutorial model",18.5401,224.9755,12,18.748,243.5156 +empirical Bayes,"branch, RLRW without numba",3.0379,45.0587,12,3.7549,48.0966 +empirical Bayes,"branch, RLRW with numba",1.0321,7.5243,12,0.627,8.5564 +variational Bayes,"main, tutorial model",,180.9569,4,45.2392,180.9569 +variational Bayes,"branch, tutorial model",,227.5081,12,18.959,227.5081 +variational Bayes,"branch, RLRW without numba",,43.6024,12,3.6335,43.6024 +variational Bayes,"branch, RLRW with numba",,8.4628,12,0.7052,8.4628 diff --git a/benchmarks/results/report-branch.csv b/benchmarks/results/report-branch.csv new file mode 100644 index 0000000..779f338 --- /dev/null +++ b/benchmarks/results/report-branch.csv @@ -0,0 +1,16 @@ +case,backend,trials,parameters,objective_ms,reset_ms,prior_ms,run_ms,loss_ms,log_posterior,speed_up +RLRW,trial,71,2,3.8262499729171395,0.004699919372797012,0.004800036549568176,3.712500096298754,0.019399914890527725,-26.034401210112605,1.0 +RLRW,python,71,2,0.7825000211596489,0.004800036549568176,0.004799920134246349,0.7339000003412366,0.01910002902150154,-26.034401210112605,4.889776190991939 +RLRW,numba,71,2,0.08020002860575914,0.004999921657145023,0.004800036549568176,0.03490003291517496,0.018799910321831703,-26.034401210112605,47.708835513337675 +HybridMBMF,trial,200,6,11.40304998261854,0.011499971151351929,0.008800067007541656,11.376600014045835,0.021999934688210487,-151.7819350468299,1.0 +HybridMBMF,python,200,6,2.066750021185726,0.011999974958598614,0.009300070814788342,1.9907000241801143,0.022000051103532314,-151.7819350468299,5.517382298647051 +HybridMBMF,numba,200,6,0.11930009350180626,0.012100092135369778,0.00870006624609232,0.0564999645575881,0.022099935449659824,-151.7819350468299,95.58290901461777 +PTSM,trial,40,4,2.216799999587238,0.0067999353632330894,0.0069999368861317635,2.084299921989441,0.019300030544400215,-36.42346848683204,1.0 +PTSM,python,40,4,0.30149996746331453,0.007600057870149612,0.00700005330145359,0.24449999909847975,0.01850002445280552,-36.42346848683204,7.352571273020022 +PTSM,numba,40,4,0.07349997758865356,0.0074999406933784485,0.0069999368861317635,0.02119992859661579,0.0171000137925148,-36.42346848683204,30.160553408514954 +PTSM1992,trial,40,6,2.256499952636659,0.009400071576237679,0.009200070053339005,2.155599999241531,0.018600025214254856,-39.43873213610623,1.0 +PTSM1992,python,40,6,0.3492499818094075,0.0097999582067132,0.009200070053339005,0.2853999612852931,0.018799910321831703,-39.43873213610623,6.460988031970966 +PTSM1992,numba,40,6,0.08310005068778992,0.009199953638017178,0.008500064723193645,0.025799963623285294,0.0198000343516469,-39.43873213610623,27.154014154749632 +PTSM2025,trial,40,5,1.9840000313706696,0.007400056347250938,0.008100061677396297,1.8681000219658017,0.019799917936325073,-33.01065168146936,1.0 +PTSM2025,python,40,5,0.30955002875998616,0.007999944500625134,0.00809994526207447,0.25210005696862936,0.019400031305849552,-33.01065168146936,6.409303333998366 +PTSM2025,numba,40,5,0.0665000407025218,0.00800006091594696,0.00809994526207447,0.014499993994832039,0.019600032828748226,-33.01065168146936,29.834568677120114 diff --git a/benchmarks/results/report-main.csv b/benchmarks/results/report-main.csv new file mode 100644 index 0000000..e9314e8 --- /dev/null +++ b/benchmarks/results/report-main.csv @@ -0,0 +1,6 @@ +case,backend,trials,parameters,objective_ms,reset_ms,prior_ms,run_ms,loss_ms,log_posterior,speed_up +RLRW,trial,71,2,8.587050018832088,0.41924999095499516,0.08154998067766428,7.541499915532768,0.04809990059584379,-26.034401210112605,1.0 +HybridMBMF,trial,200,6,25.720649980939925,0.8374500321224332,0.23040000814944506,23.526999983005226,0.054900068789720535,-151.7819350468299,1.0 +PTSM,trial,40,4,6.807900033891201,0.6519000162370503,0.16279995907098055,5.2281999960541725,0.04780001472681761,-36.42346848683204,1.0 +PTSM1992,trial,40,6,7.718699984252453,0.8669000235386193,0.2381000085733831,6.072300020605326,0.0479499576613307,-39.43873213610623,1.0 +PTSM2025,trial,40,5,76.83905004523695,0.7273000665009022,0.202000024728477,76.62820001132786,0.04780001472681761,-33.01065168146936,1.0 diff --git a/benchmarks/run.py b/benchmarks/run.py new file mode 100644 index 0000000..411b492 --- /dev/null +++ b/benchmarks/run.py @@ -0,0 +1,116 @@ +""" +Where does the time of one evaluation of the objective go, for every built-in +application of cpm and every backend? + + python benchmarks/run.py [--repeats 200] [--label NAME] + +Prints a table and writes benchmarks/results/