emutrust 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,11 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ build/
5
+ dist/
6
+ .venv/
7
+ .pytest_cache/
8
+ .ruff_cache/
9
+ .ipynb_checkpoints/
10
+ *.npz
11
+ !examples/camb/data/*.npz
@@ -0,0 +1,12 @@
1
+ # Changelog
2
+
3
+ ## 0.1.0 (2026-10-15)
4
+
5
+ First release.
6
+
7
+ - `ParameterSpace` and `latin_hypercube` for sampling training points.
8
+ - `run_model` with atomic checkpointing and output validation; `load_runs` to reload saved runs.
9
+ - `Emulator`: deep-ensemble neural network emulator (PyTorch) with automatic input and output standardisation, automatic log-space outputs, early stopping, and save/load.
10
+ - `validate` and `ValidationReport`: accuracy, worst points, error across each parameter's range, uncertainty calibration, and the measured error covariance.
11
+ - `GaussianLikelihood`: adds emulator error (validation-measured or ensemble spread) to the data noise, rejects parameters outside the training range, and evaluates batches for vectorised `emcee` sampling.
12
+ - Coverage tests on a toy model and on CAMB matter power spectra.
@@ -0,0 +1,16 @@
1
+ cff-version: 1.2.0
2
+ message: "If you use emutrust in your work, please cite it as below."
3
+ title: "emutrust: neural network emulators with calibrated uncertainty for Bayesian inference"
4
+ type: software
5
+ authors:
6
+ - family-names: Qurban
7
+ given-names: Muhammad
8
+ version: 0.1.0
9
+ date-released: 2026-10-15
10
+ license: MIT
11
+ repository-code: "https://github.com/mqurban/emutrust"
12
+ keywords:
13
+ - emulator
14
+ - uncertainty quantification
15
+ - Bayesian inference
16
+ - cosmology
emutrust-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Muhammad Qurban
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,301 @@
1
+ Metadata-Version: 2.5
2
+ Name: emutrust
3
+ Version: 0.1.0
4
+ Summary: Neural network emulators for slow scientific models, with calibrated uncertainty you can trust in Bayesian inference.
5
+ Project-URL: Repository, https://github.com/mqurban/emutrust
6
+ Project-URL: Documentation, https://github.com/mqurban/emutrust/blob/main/docs/guide.md
7
+ Project-URL: Issues, https://github.com/mqurban/emutrust/issues
8
+ Project-URL: Changelog, https://github.com/mqurban/emutrust/blob/main/CHANGELOG.md
9
+ Author-email: Muhammad Qurban <mianqurban525@gmail.com>
10
+ License-Expression: MIT
11
+ License-File: LICENSE
12
+ Keywords: bayesian inference,cosmology,emulator,mcmc,neural networks,surrogate model,uncertainty quantification
13
+ Classifier: Development Status :: 3 - Alpha
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: License :: OSI Approved :: MIT License
16
+ Classifier: Operating System :: OS Independent
17
+ Classifier: Programming Language :: Python :: 3
18
+ Classifier: Programming Language :: Python :: 3.9
19
+ Classifier: Programming Language :: Python :: 3.10
20
+ Classifier: Programming Language :: Python :: 3.11
21
+ Classifier: Programming Language :: Python :: 3.12
22
+ Classifier: Programming Language :: Python :: 3.13
23
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
24
+ Classifier: Topic :: Scientific/Engineering :: Astronomy
25
+ Classifier: Topic :: Scientific/Engineering :: Physics
26
+ Requires-Python: >=3.9
27
+ Requires-Dist: numpy>=1.22
28
+ Requires-Dist: torch>=2.0
29
+ Provides-Extra: dev
30
+ Requires-Dist: emcee>=3.1; extra == 'dev'
31
+ Requires-Dist: matplotlib>=3.5; extra == 'dev'
32
+ Requires-Dist: pytest>=7; extra == 'dev'
33
+ Requires-Dist: ruff<0.17,>=0.16; extra == 'dev'
34
+ Provides-Extra: mcmc
35
+ Requires-Dist: emcee>=3.1; extra == 'mcmc'
36
+ Provides-Extra: plot
37
+ Requires-Dist: matplotlib>=3.5; extra == 'plot'
38
+ Description-Content-Type: text/markdown
39
+
40
+ # emutrust
41
+
42
+ **Neural network emulators for slow scientific models, with calibrated uncertainty you can trust in Bayesian inference.**
43
+
44
+ [![tests](https://github.com/mqurban/emutrust/actions/workflows/tests.yml/badge.svg)](https://github.com/mqurban/emutrust/actions/workflows/tests.yml)
45
+ [![PyPI](https://img.shields.io/pypi/v/emutrust)](https://pypi.org/project/emutrust/)
46
+ [![Python](https://img.shields.io/pypi/pyversions/emutrust)](https://pypi.org/project/emutrust/)
47
+ [![License: MIT](https://img.shields.io/badge/license-MIT-blue)](https://github.com/mqurban/emutrust/blob/main/LICENSE)
48
+
49
+ > **Status:** alpha (`0.1.0`). All components are working, tested on Python 3.9 to 3.13, and demonstrated on CAMB. The API may still change before 1.0.
50
+
51
+ **Documentation:** [user guide](https://github.com/mqurban/emutrust/blob/main/docs/guide.md) · [API reference](https://github.com/mqurban/emutrust/blob/main/docs/api.md) · [quickstart example](https://github.com/mqurban/emutrust/blob/main/examples/quickstart.py) · [changelog](https://github.com/mqurban/emutrust/blob/main/CHANGELOG.md)
52
+
53
+ ## The problem
54
+
55
+ Many scientific models are slow: a single evaluation of a Boltzmann code such as CAMB takes seconds, and more complex models can take minutes. Bayesian parameter inference with MCMC needs tens of thousands of evaluations, so running the model directly inside the sampler can take days.
56
+
57
+ An **emulator** is a fast stand-in: run the slow model a few hundred times, train a neural network on the results, and use the network inside the sampler instead. This can cut inference from days to seconds.
58
+
59
+ But an emulator is only approximately right. If the inference treats it as exact, its errors silently shrink the error bars and the result becomes **overconfident**. `emutrust` is built around closing that gap:
60
+
61
+ 1. **Smart sampling** of training points across the parameter space, with checkpointing so slow runs are never lost.
62
+ 2. **Training** of a neural network emulator that also estimates **its own uncertainty**.
63
+ 3. **A validation report** showing accuracy on held-out runs, the worst cases, and where in parameter space the emulator is weak.
64
+ 4. **An honest likelihood** that adds the emulator's measured error to the data's noise, ready for samplers such as `emcee`.
65
+
66
+ It works with **any** Python function: a cosmology code, a simulation wrapper, or any other expensive model.
67
+
68
+ ## Installation
69
+
70
+ ```bash
71
+ pip install emutrust # core: NumPy and PyTorch
72
+ pip install "emutrust[mcmc,plot]" # also emcee and matplotlib, used in the examples
73
+ ```
74
+
75
+ `emutrust` needs NumPy and PyTorch. On a machine without a GPU, installing the CPU-only build of PyTorch first (`pip install torch --index-url https://download.pytorch.org/whl/cpu`) avoids a large download.
76
+
77
+ ## Usage
78
+
79
+ A complete, runnable version of the steps below is in [`examples/quickstart.py`](https://github.com/mqurban/emutrust/blob/main/examples/quickstart.py) (about 15 seconds on a laptop).
80
+
81
+ ### 1. Define the parameter space and sample training points
82
+
83
+ ```python
84
+ from emutrust import ParameterSpace
85
+
86
+ space = ParameterSpace(
87
+ {
88
+ "H0": (60.0, 75.0),
89
+ "omch2": (0.10, 0.14),
90
+ "ns": (0.92, 1.00),
91
+ }
92
+ )
93
+
94
+ X = space.sample(300, seed=42) # Latin hypercube, shape (300, 3)
95
+ ```
96
+
97
+ Points are drawn with **Latin hypercube sampling**: each parameter's range is split into 300 equal slices and every slice is used exactly once. This covers the full range of every parameter evenly, which plain random sampling does not, and avoids the exponential cost of a grid (10 values per parameter in 5 dimensions is already 100,000 runs).
98
+
99
+ ### 2. Run the slow model, with checkpointing
100
+
101
+ ```python
102
+ from emutrust import run_model
103
+
104
+
105
+ def my_model(params):
106
+ # params is a dict, e.g. {"H0": 67.3, "omch2": 0.12, "ns": 0.965}
107
+ return slow_calculation(params) # returns an array of outputs
108
+
109
+
110
+ training = run_model(my_model, X, space, checkpoint="runs.npz")
111
+ training.X, training.Y # inputs (300, 3) and outputs (300, n_outputs)
112
+ training.seconds_per_run # measured cost of one model evaluation
113
+ ```
114
+
115
+ Progress is saved after every run. If the session crashes after 200 runs, calling `run_model` again with the same points resumes at run 201. The checkpoint is written atomically, so an interruption during a save cannot corrupt it, and it refuses to resume if the points or parameters have changed.
116
+
117
+ `run_model` also checks every output: a run that returns NaN or infinity, or a different number of outputs from earlier runs, stops with a clear error instead of silently polluting the training set.
118
+
119
+ ### 3. Train an emulator that knows its own uncertainty
120
+
121
+ ```python
122
+ from emutrust import Emulator
123
+
124
+ emu = Emulator.fit(training, n_members=5, seed=0)
125
+
126
+ mean, std = emu.predict({"H0": 67.3, "omch2": 0.12, "ns": 0.965})
127
+ emu.save("emulator.pt") # reload later with Emulator.load
128
+ ```
129
+
130
+ `Emulator.fit` trains a **deep ensemble**: several neural networks on the same training set, each starting from different random weights and holding out a different random part of the data for early stopping. `predict` returns the ensemble average and the spread between members, both in the model's original units.
131
+
132
+ - **Where the training data pin the answer down, the members agree** and the spread is small.
133
+ - **Where data are thin or the function is hard to learn, they disagree**, and the spread grows. This spread is the emulator's estimate of its own error.
134
+
135
+ Inputs and outputs are standardised before training. Outputs that are all positive and span more than two orders of magnitude, such as power spectra, are learned in `log10` space automatically.
136
+
137
+ On the toy spectrum below, an ensemble of 5 networks trained on 300 runs gives, on 200 unseen test points ([`examples/toy_benchmark.py`](https://github.com/mqurban/emutrust/blob/main/examples/toy_benchmark.py)):
138
+
139
+ | Metric | Result | Ideal |
140
+ |---|---:|---:|
141
+ | Mean fractional error | 0.32% | |
142
+ | Median fractional error | 0.19% | |
143
+ | Worst fractional error | 12.8% | |
144
+ | Truth within 1 std of prediction | 67% | 68% |
145
+ | Truth within 2 std of prediction | 92% | 95% |
146
+
147
+ The last two rows test whether the uncertainty is *honest*: if it is well calibrated, the true value should fall within one predicted standard deviation 68% of the time, and within two 95% of the time. The worst cases are near the edges of the parameter space, where training points are sparsest.
148
+
149
+ **A known limitation.** Outside the parameter ranges used for training, the members spread apart, which correctly signals unfamiliar ground, but they can all extrapolate in the same wrong direction, so the spread no longer measures the true error ([`examples/ensemble_uncertainty.py`](https://github.com/mqurban/emutrust/blob/main/examples/ensemble_uncertainty.py)):
150
+
151
+ ![Ensemble members agree inside the training range and spread apart outside it](https://raw.githubusercontent.com/mqurban/emutrust/main/docs/images/ensemble_uncertainty.png)
152
+
153
+ An emulator should therefore only be used inside the ranges it was trained on; the likelihood below enforces this.
154
+
155
+ ### 4. Validate it on runs it has never seen
156
+
157
+ ```python
158
+ from emutrust import validate
159
+
160
+ validation = run_model(my_model, space.sample(300, seed=1), space)
161
+ report = validate(emu, validation)
162
+ print(report)
163
+ report.plot("validation.png")
164
+ ```
165
+
166
+ The report covers accuracy (mean, median and worst error), the worst parameter points, how the error varies across each parameter's range, and whether the uncertainty is honest. For an emulator trained on only 80 runs of the toy spectrum ([`examples/coverage_test.py`](https://github.com/mqurban/emutrust/blob/main/examples/coverage_test.py)), part of it reads:
167
+
168
+ ```
169
+ Uncertainty calibration
170
+ truth within 1 std: 65.1% (ideal 68.3%)
171
+ truth within 2 std: 87.6% (ideal 95.4%)
172
+ suggested std_scale: 1.08 (well calibrated)
173
+
174
+ Median worst-output error across each parameter's range (low -> high)
175
+ amplitude: 3.65% 4.45% 4.06% 4.96% 4.68%
176
+ slope: 4.16% 5.61% 4.68% 3.76% 3.34%
177
+ cutoff: 14.7% 4.8% 4.68% 2.48% 3%
178
+ ```
179
+
180
+ The last line points straight at the weak region: low `cutoff`, where the toy spectrum falls off most steeply. That is where more training runs would help most.
181
+
182
+ If the uncertainty is too small or too large, `report.suggested_std_scale` gives the correction, applied with `emu.std_scale *= report.suggested_std_scale`. Ideally use separate validation runs for this and for the final assessment.
183
+
184
+ ### 5. Infer parameters with an honest likelihood
185
+
186
+ ```python
187
+ import emcee
188
+ from emutrust import GaussianLikelihood
189
+
190
+ like = GaussianLikelihood(emu, data, noise_std, emulator_error=report)
191
+ sampler = emcee.EnsembleSampler(32, like.ndim, like, vectorize=True)
192
+ sampler.run_mcmc(start, 2000)
193
+ ```
194
+
195
+ The likelihood compares the data with the emulator's prediction using a covariance `C = C_noise + C_emulator`:
196
+
197
+ - **`emulator_error=report`** (recommended): the emulator's error pattern measured on the validation runs, including **correlations between outputs**.
198
+ - **`emulator_error="ensemble"`**: the ensemble's spread at each point, with each output's error treated as independent.
199
+ - **`emulator_error=None`**: the emulator is treated as exact. This is the usual approach, kept for comparison.
200
+
201
+ Parameters outside the range the emulator was trained on are rejected, so the sampler never relies on extrapolation.
202
+
203
+ **Does it work? A coverage test.** [`examples/coverage_test.py`](https://github.com/mqurban/emutrust/blob/main/examples/coverage_test.py) simulates 30 observations with true parameters drawn from the prior (50 data points each, 1% noise), infers the parameters through a deliberately imperfect emulator (80 training runs), and counts how often the reported 68% interval contains the truth. Honest error bars contain it about 68% of the time:
204
+
205
+ ![Coverage of the 68% interval for three likelihoods](https://raw.githubusercontent.com/mqurban/emutrust/main/docs/images/coverage.png)
206
+
207
+ | Likelihood | Truth inside the 68% interval |
208
+ |---|---:|
209
+ | Emulator treated as exact | 20% |
210
+ | + ensemble spread | 38% |
211
+ | + validation-measured error | 74% |
212
+
213
+ Treating the emulator as exact gives error bars far too small: the truth falls outside them most of the time. The ensemble spread helps, but not enough. The reason is that emulator errors are **systematic**: an emulator that is 1% too high at one scale is usually too high at the neighbouring scales too, so the error does not average away over 50 data points the way independent noise does. A 5-member ensemble can only describe a handful of such error patterns. The covariance measured on 300 validation runs captures them, and the result is close to honest (74% against the ideal 68%), slightly on the cautious side because a single error estimate is applied everywhere, including regions where the emulator is better than average.
214
+
215
+ With 30 observations of 3 parameters, each coverage figure is uncertain by about ±5 percentage points.
216
+
217
+ The true parameters are drawn from the whole prior, here the range the emulator was trained on. That is what makes 68% the right target: a Bayesian credible interval is only guaranteed to contain the truth 68% of the time *on average over truths drawn from the prior*. Drawing them only from the middle of the ranges inflates coverage whenever posteriors are wide enough to feel the edges.
218
+
219
+ For validation-measured error, the likelihood describes the measured error by its main correlated patterns (99.9% of the error variance by default, set with `error_variance_kept`) and treats the small remainder as independent per-output error. This is mathematically almost identical and much faster: on the CAMB example below, one evaluation takes about 0.3–0.5 ms instead of several milliseconds.
220
+
221
+ ## Demonstration on CAMB
222
+
223
+ [`examples/camb/`](https://github.com/mqurban/emutrust/tree/main/examples/camb) applies `emutrust` to a real cosmology problem: the linear matter power spectrum P(k) from the Boltzmann code [CAMB](https://camb.info/), at 200 scales between k = 10⁻⁴ and 1 h/Mpc, over five ΛCDM parameters (H0, ombh2, omch2, ns, As). The setup matches the earlier [cosmo-emulator](https://github.com/mqurban/cosmo-emulator) project, which treated its emulator as exact.
224
+
225
+ - [`generate.py`](https://github.com/mqurban/emutrust/blob/main/examples/camb/generate.py) runs CAMB for 300 training and 300 validation points (CAMB 2.0.5). The runs are included in the repository, so CAMB itself is only needed to simulate new observations.
226
+ - [`demo.py`](https://github.com/mqurban/emutrust/blob/main/examples/camb/demo.py) trains a 5-member emulator, validates it, times it against CAMB and runs a coverage test.
227
+
228
+ **Accuracy and speed.**
229
+
230
+ | | Result |
231
+ |---|---:|
232
+ | Mean error on 300 unseen spectra | 0.05% |
233
+ | Median error | 0.03% |
234
+ | Worst error (near the edges of the parameter ranges) | 1.5% |
235
+ | Truth within 1 / 2 predicted std | 71% / 92% |
236
+ | One CAMB call | 0.58 s |
237
+ | One likelihood evaluation, including emulator error | 0.3–0.5 ms |
238
+ | One MCMC run (192,000 evaluations) | about 1 minute, against about 31 hours with CAMB in the loop |
239
+
240
+ The speed-up is roughly 1,500–2,000x; the exact figure varies with machine load.
241
+
242
+ ![Validation of the CAMB emulator](https://raw.githubusercontent.com/mqurban/emutrust/main/docs/images/camb_validation.png)
243
+
244
+ **Are the error bars honest?** For 30 simulated observations at each of two noise levels, with true parameters drawn from the prior, MCMC (32 walkers, 6,000 steps, at least 47 autocorrelation times after burn-in) gives:
245
+
246
+ ![Coverage for CAMB at two noise levels](https://raw.githubusercontent.com/mqurban/emutrust/main/docs/images/camb_coverage.png)
247
+
248
+ | Noise per data point | Emulator treated as exact | + validation-measured error |
249
+ |---|---:|---:|
250
+ | 2% | 75% | 75% |
251
+ | 0.2% | **55%** | **71%** |
252
+
253
+ (ideal: 68%; each figure is uncertain by about ±4 percentage points)
254
+
255
+ - **At 2% noise, the emulator's error does not matter.** The emulator is accurate to about 0.05%, forty times smaller than the noise, and both likelihoods agree. Treating this emulator as exact was safe for data of this quality.
256
+ - **At 0.2% noise, it does.** The emulator's systematic errors become comparable to the information in the data, and treating it as exact makes the error bars too small: the truth falls inside the 68% interval only 55% of the time. Including the validation-measured error restores honest error bars (71%).
257
+
258
+ As surveys reach sub-percent precision, an emulator that was good enough yesterday stops being good enough, and its error has to be part of the inference.
259
+
260
+ ### Toy model for testing
261
+
262
+ `emutrust.toy.toy_spectrum` is a cheap function shaped like a matter power spectrum (spanning several orders of magnitude and varying smoothly with three parameters). It is used to test every component on a problem where the right answer is known exactly, before trusting it on a real, slow model.
263
+
264
+ ## Roadmap
265
+
266
+ - [x] Parameter spaces and Latin hypercube sampling
267
+ - [x] Checkpointed model runner with output validation
268
+ - [x] Neural network emulator (PyTorch) with input/output standardisation and log-space outputs
269
+ - [x] Emulator uncertainty from a deep ensemble
270
+ - [x] Validation report: held-out accuracy, worst cases, uncertainty calibration
271
+ - [x] Likelihood that includes emulator error (ensemble or validation-measured), ready for `emcee`
272
+ - [x] Coverage test showing the effect on inference
273
+ - [x] Demonstration on CAMB: accuracy, speed, and coverage with and without emulator error
274
+ - [x] User guide, API reference and PyPI release
275
+ - [ ] Adaptive sampling: add training runs where the validation report shows the largest errors
276
+ - [ ] Error estimates that vary across the parameter space, instead of one average
277
+
278
+ ## Related work
279
+
280
+ - **[CosmoPower](https://github.com/alessiospuriomancini/cosmopower)** (Spurio Mancini et al. 2022): neural network emulators specialised for cosmological power spectra.
281
+ - **[Emukit](https://github.com/emukit/emukit)**: general emulation and uncertainty quantification based on Gaussian processes.
282
+
283
+ `emutrust` aims to be a general, lightweight tool for *any* slow model, focused on making emulator uncertainty an explicit part of the inference.
284
+
285
+ ## Development
286
+
287
+ ```bash
288
+ git clone https://github.com/mqurban/emutrust
289
+ cd emutrust
290
+ pip install -e ".[dev]"
291
+ pytest # run the tests
292
+ ruff check . # lint
293
+ ```
294
+
295
+ ## License
296
+
297
+ MIT. See [LICENSE](https://github.com/mqurban/emutrust/blob/main/LICENSE).
298
+
299
+ ## Citing
300
+
301
+ If `emutrust` is useful in your work, please cite it using the metadata in [`CITATION.cff`](https://github.com/mqurban/emutrust/blob/main/CITATION.cff).
@@ -0,0 +1,262 @@
1
+ # emutrust
2
+
3
+ **Neural network emulators for slow scientific models, with calibrated uncertainty you can trust in Bayesian inference.**
4
+
5
+ [![tests](https://github.com/mqurban/emutrust/actions/workflows/tests.yml/badge.svg)](https://github.com/mqurban/emutrust/actions/workflows/tests.yml)
6
+ [![PyPI](https://img.shields.io/pypi/v/emutrust)](https://pypi.org/project/emutrust/)
7
+ [![Python](https://img.shields.io/pypi/pyversions/emutrust)](https://pypi.org/project/emutrust/)
8
+ [![License: MIT](https://img.shields.io/badge/license-MIT-blue)](https://github.com/mqurban/emutrust/blob/main/LICENSE)
9
+
10
+ > **Status:** alpha (`0.1.0`). All components are working, tested on Python 3.9 to 3.13, and demonstrated on CAMB. The API may still change before 1.0.
11
+
12
+ **Documentation:** [user guide](https://github.com/mqurban/emutrust/blob/main/docs/guide.md) · [API reference](https://github.com/mqurban/emutrust/blob/main/docs/api.md) · [quickstart example](https://github.com/mqurban/emutrust/blob/main/examples/quickstart.py) · [changelog](https://github.com/mqurban/emutrust/blob/main/CHANGELOG.md)
13
+
14
+ ## The problem
15
+
16
+ Many scientific models are slow: a single evaluation of a Boltzmann code such as CAMB takes seconds, and more complex models can take minutes. Bayesian parameter inference with MCMC needs tens of thousands of evaluations, so running the model directly inside the sampler can take days.
17
+
18
+ An **emulator** is a fast stand-in: run the slow model a few hundred times, train a neural network on the results, and use the network inside the sampler instead. This can cut inference from days to seconds.
19
+
20
+ But an emulator is only approximately right. If the inference treats it as exact, its errors silently shrink the error bars and the result becomes **overconfident**. `emutrust` is built around closing that gap:
21
+
22
+ 1. **Smart sampling** of training points across the parameter space, with checkpointing so slow runs are never lost.
23
+ 2. **Training** of a neural network emulator that also estimates **its own uncertainty**.
24
+ 3. **A validation report** showing accuracy on held-out runs, the worst cases, and where in parameter space the emulator is weak.
25
+ 4. **An honest likelihood** that adds the emulator's measured error to the data's noise, ready for samplers such as `emcee`.
26
+
27
+ It works with **any** Python function: a cosmology code, a simulation wrapper, or any other expensive model.
28
+
29
+ ## Installation
30
+
31
+ ```bash
32
+ pip install emutrust # core: NumPy and PyTorch
33
+ pip install "emutrust[mcmc,plot]" # also emcee and matplotlib, used in the examples
34
+ ```
35
+
36
+ `emutrust` needs NumPy and PyTorch. On a machine without a GPU, installing the CPU-only build of PyTorch first (`pip install torch --index-url https://download.pytorch.org/whl/cpu`) avoids a large download.
37
+
38
+ ## Usage
39
+
40
+ A complete, runnable version of the steps below is in [`examples/quickstart.py`](https://github.com/mqurban/emutrust/blob/main/examples/quickstart.py) (about 15 seconds on a laptop).
41
+
42
+ ### 1. Define the parameter space and sample training points
43
+
44
+ ```python
45
+ from emutrust import ParameterSpace
46
+
47
+ space = ParameterSpace(
48
+ {
49
+ "H0": (60.0, 75.0),
50
+ "omch2": (0.10, 0.14),
51
+ "ns": (0.92, 1.00),
52
+ }
53
+ )
54
+
55
+ X = space.sample(300, seed=42) # Latin hypercube, shape (300, 3)
56
+ ```
57
+
58
+ Points are drawn with **Latin hypercube sampling**: each parameter's range is split into 300 equal slices and every slice is used exactly once. This covers the full range of every parameter evenly, which plain random sampling does not, and avoids the exponential cost of a grid (10 values per parameter in 5 dimensions is already 100,000 runs).
59
+
60
+ ### 2. Run the slow model, with checkpointing
61
+
62
+ ```python
63
+ from emutrust import run_model
64
+
65
+
66
+ def my_model(params):
67
+ # params is a dict, e.g. {"H0": 67.3, "omch2": 0.12, "ns": 0.965}
68
+ return slow_calculation(params) # returns an array of outputs
69
+
70
+
71
+ training = run_model(my_model, X, space, checkpoint="runs.npz")
72
+ training.X, training.Y # inputs (300, 3) and outputs (300, n_outputs)
73
+ training.seconds_per_run # measured cost of one model evaluation
74
+ ```
75
+
76
+ Progress is saved after every run. If the session crashes after 200 runs, calling `run_model` again with the same points resumes at run 201. The checkpoint is written atomically, so an interruption during a save cannot corrupt it, and it refuses to resume if the points or parameters have changed.
77
+
78
+ `run_model` also checks every output: a run that returns NaN or infinity, or a different number of outputs from earlier runs, stops with a clear error instead of silently polluting the training set.
79
+
80
+ ### 3. Train an emulator that knows its own uncertainty
81
+
82
+ ```python
83
+ from emutrust import Emulator
84
+
85
+ emu = Emulator.fit(training, n_members=5, seed=0)
86
+
87
+ mean, std = emu.predict({"H0": 67.3, "omch2": 0.12, "ns": 0.965})
88
+ emu.save("emulator.pt") # reload later with Emulator.load
89
+ ```
90
+
91
+ `Emulator.fit` trains a **deep ensemble**: several neural networks on the same training set, each starting from different random weights and holding out a different random part of the data for early stopping. `predict` returns the ensemble average and the spread between members, both in the model's original units.
92
+
93
+ - **Where the training data pin the answer down, the members agree** and the spread is small.
94
+ - **Where data are thin or the function is hard to learn, they disagree**, and the spread grows. This spread is the emulator's estimate of its own error.
95
+
96
+ Inputs and outputs are standardised before training. Outputs that are all positive and span more than two orders of magnitude, such as power spectra, are learned in `log10` space automatically.
97
+
98
+ On the toy spectrum below, an ensemble of 5 networks trained on 300 runs gives, on 200 unseen test points ([`examples/toy_benchmark.py`](https://github.com/mqurban/emutrust/blob/main/examples/toy_benchmark.py)):
99
+
100
+ | Metric | Result | Ideal |
101
+ |---|---:|---:|
102
+ | Mean fractional error | 0.32% | |
103
+ | Median fractional error | 0.19% | |
104
+ | Worst fractional error | 12.8% | |
105
+ | Truth within 1 std of prediction | 67% | 68% |
106
+ | Truth within 2 std of prediction | 92% | 95% |
107
+
108
+ The last two rows test whether the uncertainty is *honest*: if it is well calibrated, the true value should fall within one predicted standard deviation 68% of the time, and within two 95% of the time. The worst cases are near the edges of the parameter space, where training points are sparsest.
109
+
110
+ **A known limitation.** Outside the parameter ranges used for training, the members spread apart, which correctly signals unfamiliar ground, but they can all extrapolate in the same wrong direction, so the spread no longer measures the true error ([`examples/ensemble_uncertainty.py`](https://github.com/mqurban/emutrust/blob/main/examples/ensemble_uncertainty.py)):
111
+
112
+ ![Ensemble members agree inside the training range and spread apart outside it](https://raw.githubusercontent.com/mqurban/emutrust/main/docs/images/ensemble_uncertainty.png)
113
+
114
+ An emulator should therefore only be used inside the ranges it was trained on; the likelihood below enforces this.
115
+
116
+ ### 4. Validate it on runs it has never seen
117
+
118
+ ```python
119
+ from emutrust import validate
120
+
121
+ validation = run_model(my_model, space.sample(300, seed=1), space)
122
+ report = validate(emu, validation)
123
+ print(report)
124
+ report.plot("validation.png")
125
+ ```
126
+
127
+ The report covers accuracy (mean, median and worst error), the worst parameter points, how the error varies across each parameter's range, and whether the uncertainty is honest. For an emulator trained on only 80 runs of the toy spectrum ([`examples/coverage_test.py`](https://github.com/mqurban/emutrust/blob/main/examples/coverage_test.py)), part of it reads:
128
+
129
+ ```
130
+ Uncertainty calibration
131
+ truth within 1 std: 65.1% (ideal 68.3%)
132
+ truth within 2 std: 87.6% (ideal 95.4%)
133
+ suggested std_scale: 1.08 (well calibrated)
134
+
135
+ Median worst-output error across each parameter's range (low -> high)
136
+ amplitude: 3.65% 4.45% 4.06% 4.96% 4.68%
137
+ slope: 4.16% 5.61% 4.68% 3.76% 3.34%
138
+ cutoff: 14.7% 4.8% 4.68% 2.48% 3%
139
+ ```
140
+
141
+ The last line points straight at the weak region: low `cutoff`, where the toy spectrum falls off most steeply. That is where more training runs would help most.
142
+
143
+ If the uncertainty is too small or too large, `report.suggested_std_scale` gives the correction, applied with `emu.std_scale *= report.suggested_std_scale`. Ideally use separate validation runs for this and for the final assessment.
144
+
145
+ ### 5. Infer parameters with an honest likelihood
146
+
147
+ ```python
148
+ import emcee
149
+ from emutrust import GaussianLikelihood
150
+
151
+ like = GaussianLikelihood(emu, data, noise_std, emulator_error=report)
152
+ sampler = emcee.EnsembleSampler(32, like.ndim, like, vectorize=True)
153
+ sampler.run_mcmc(start, 2000)
154
+ ```
155
+
156
+ The likelihood compares the data with the emulator's prediction using a covariance `C = C_noise + C_emulator`:
157
+
158
+ - **`emulator_error=report`** (recommended): the emulator's error pattern measured on the validation runs, including **correlations between outputs**.
159
+ - **`emulator_error="ensemble"`**: the ensemble's spread at each point, with each output's error treated as independent.
160
+ - **`emulator_error=None`**: the emulator is treated as exact. This is the usual approach, kept for comparison.
161
+
162
+ Parameters outside the range the emulator was trained on are rejected, so the sampler never relies on extrapolation.
163
+
164
+ **Does it work? A coverage test.** [`examples/coverage_test.py`](https://github.com/mqurban/emutrust/blob/main/examples/coverage_test.py) simulates 30 observations with true parameters drawn from the prior (50 data points each, 1% noise), infers the parameters through a deliberately imperfect emulator (80 training runs), and counts how often the reported 68% interval contains the truth. Honest error bars contain it about 68% of the time:
165
+
166
+ ![Coverage of the 68% interval for three likelihoods](https://raw.githubusercontent.com/mqurban/emutrust/main/docs/images/coverage.png)
167
+
168
+ | Likelihood | Truth inside the 68% interval |
169
+ |---|---:|
170
+ | Emulator treated as exact | 20% |
171
+ | + ensemble spread | 38% |
172
+ | + validation-measured error | 74% |
173
+
174
+ Treating the emulator as exact gives error bars far too small: the truth falls outside them most of the time. The ensemble spread helps, but not enough. The reason is that emulator errors are **systematic**: an emulator that is 1% too high at one scale is usually too high at the neighbouring scales too, so the error does not average away over 50 data points the way independent noise does. A 5-member ensemble can only describe a handful of such error patterns. The covariance measured on 300 validation runs captures them, and the result is close to honest (74% against the ideal 68%), slightly on the cautious side because a single error estimate is applied everywhere, including regions where the emulator is better than average.
175
+
176
+ With 30 observations of 3 parameters, each coverage figure is uncertain by about ±5 percentage points.
177
+
178
+ The true parameters are drawn from the whole prior, here the range the emulator was trained on. That is what makes 68% the right target: a Bayesian credible interval is only guaranteed to contain the truth 68% of the time *on average over truths drawn from the prior*. Drawing them only from the middle of the ranges inflates coverage whenever posteriors are wide enough to feel the edges.
179
+
180
+ For validation-measured error, the likelihood describes the measured error by its main correlated patterns (99.9% of the error variance by default, set with `error_variance_kept`) and treats the small remainder as independent per-output error. This is mathematically almost identical and much faster: on the CAMB example below, one evaluation takes about 0.3–0.5 ms instead of several milliseconds.
181
+
182
+ ## Demonstration on CAMB
183
+
184
+ [`examples/camb/`](https://github.com/mqurban/emutrust/tree/main/examples/camb) applies `emutrust` to a real cosmology problem: the linear matter power spectrum P(k) from the Boltzmann code [CAMB](https://camb.info/), at 200 scales between k = 10⁻⁴ and 1 h/Mpc, over five ΛCDM parameters (H0, ombh2, omch2, ns, As). The setup matches the earlier [cosmo-emulator](https://github.com/mqurban/cosmo-emulator) project, which treated its emulator as exact.
185
+
186
+ - [`generate.py`](https://github.com/mqurban/emutrust/blob/main/examples/camb/generate.py) runs CAMB for 300 training and 300 validation points (CAMB 2.0.5). The runs are included in the repository, so CAMB itself is only needed to simulate new observations.
187
+ - [`demo.py`](https://github.com/mqurban/emutrust/blob/main/examples/camb/demo.py) trains a 5-member emulator, validates it, times it against CAMB and runs a coverage test.
188
+
189
+ **Accuracy and speed.**
190
+
191
+ | | Result |
192
+ |---|---:|
193
+ | Mean error on 300 unseen spectra | 0.05% |
194
+ | Median error | 0.03% |
195
+ | Worst error (near the edges of the parameter ranges) | 1.5% |
196
+ | Truth within 1 / 2 predicted std | 71% / 92% |
197
+ | One CAMB call | 0.58 s |
198
+ | One likelihood evaluation, including emulator error | 0.3–0.5 ms |
199
+ | One MCMC run (192,000 evaluations) | about 1 minute, against about 31 hours with CAMB in the loop |
200
+
201
+ The speed-up is roughly 1,500–2,000x; the exact figure varies with machine load.
202
+
203
+ ![Validation of the CAMB emulator](https://raw.githubusercontent.com/mqurban/emutrust/main/docs/images/camb_validation.png)
204
+
205
+ **Are the error bars honest?** For 30 simulated observations at each of two noise levels, with true parameters drawn from the prior, MCMC (32 walkers, 6,000 steps, at least 47 autocorrelation times after burn-in) gives:
206
+
207
+ ![Coverage for CAMB at two noise levels](https://raw.githubusercontent.com/mqurban/emutrust/main/docs/images/camb_coverage.png)
208
+
209
+ | Noise per data point | Emulator treated as exact | + validation-measured error |
210
+ |---|---:|---:|
211
+ | 2% | 75% | 75% |
212
+ | 0.2% | **55%** | **71%** |
213
+
214
+ (ideal: 68%; each figure is uncertain by about ±4 percentage points)
215
+
216
+ - **At 2% noise, the emulator's error does not matter.** The emulator is accurate to about 0.05%, forty times smaller than the noise, and both likelihoods agree. Treating this emulator as exact was safe for data of this quality.
217
+ - **At 0.2% noise, it does.** The emulator's systematic errors become comparable to the information in the data, and treating it as exact makes the error bars too small: the truth falls inside the 68% interval only 55% of the time. Including the validation-measured error restores honest error bars (71%).
218
+
219
+ As surveys reach sub-percent precision, an emulator that was good enough yesterday stops being good enough, and its error has to be part of the inference.
220
+
221
+ ### Toy model for testing
222
+
223
+ `emutrust.toy.toy_spectrum` is a cheap function shaped like a matter power spectrum (spanning several orders of magnitude and varying smoothly with three parameters). It is used to test every component on a problem where the right answer is known exactly, before trusting it on a real, slow model.
224
+
225
+ ## Roadmap
226
+
227
+ - [x] Parameter spaces and Latin hypercube sampling
228
+ - [x] Checkpointed model runner with output validation
229
+ - [x] Neural network emulator (PyTorch) with input/output standardisation and log-space outputs
230
+ - [x] Emulator uncertainty from a deep ensemble
231
+ - [x] Validation report: held-out accuracy, worst cases, uncertainty calibration
232
+ - [x] Likelihood that includes emulator error (ensemble or validation-measured), ready for `emcee`
233
+ - [x] Coverage test showing the effect on inference
234
+ - [x] Demonstration on CAMB: accuracy, speed, and coverage with and without emulator error
235
+ - [x] User guide, API reference and PyPI release
236
+ - [ ] Adaptive sampling: add training runs where the validation report shows the largest errors
237
+ - [ ] Error estimates that vary across the parameter space, instead of one average
238
+
239
+ ## Related work
240
+
241
+ - **[CosmoPower](https://github.com/alessiospuriomancini/cosmopower)** (Spurio Mancini et al. 2022): neural network emulators specialised for cosmological power spectra.
242
+ - **[Emukit](https://github.com/emukit/emukit)**: general emulation and uncertainty quantification based on Gaussian processes.
243
+
244
+ `emutrust` aims to be a general, lightweight tool for *any* slow model, focused on making emulator uncertainty an explicit part of the inference.
245
+
246
+ ## Development
247
+
248
+ ```bash
249
+ git clone https://github.com/mqurban/emutrust
250
+ cd emutrust
251
+ pip install -e ".[dev]"
252
+ pytest # run the tests
253
+ ruff check . # lint
254
+ ```
255
+
256
+ ## License
257
+
258
+ MIT. See [LICENSE](https://github.com/mqurban/emutrust/blob/main/LICENSE).
259
+
260
+ ## Citing
261
+
262
+ If `emutrust` is useful in your work, please cite it using the metadata in [`CITATION.cff`](https://github.com/mqurban/emutrust/blob/main/CITATION.cff).