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.
- emutrust-0.1.0/.gitignore +11 -0
- emutrust-0.1.0/CHANGELOG.md +12 -0
- emutrust-0.1.0/CITATION.cff +16 -0
- emutrust-0.1.0/LICENSE +21 -0
- emutrust-0.1.0/PKG-INFO +301 -0
- emutrust-0.1.0/README.md +262 -0
- emutrust-0.1.0/pyproject.toml +67 -0
- emutrust-0.1.0/src/emutrust/__init__.py +23 -0
- emutrust-0.1.0/src/emutrust/emulator.py +311 -0
- emutrust-0.1.0/src/emutrust/likelihood.py +183 -0
- emutrust-0.1.0/src/emutrust/runner.py +151 -0
- emutrust-0.1.0/src/emutrust/space.py +122 -0
- emutrust-0.1.0/src/emutrust/toy.py +29 -0
- emutrust-0.1.0/src/emutrust/validation.py +245 -0
- emutrust-0.1.0/tests/conftest.py +26 -0
- emutrust-0.1.0/tests/test_emulator.py +94 -0
- emutrust-0.1.0/tests/test_likelihood.py +102 -0
- emutrust-0.1.0/tests/test_runner.py +117 -0
- emutrust-0.1.0/tests/test_space.py +73 -0
- emutrust-0.1.0/tests/test_validation.py +91 -0
|
@@ -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.
|
emutrust-0.1.0/PKG-INFO
ADDED
|
@@ -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
|
+
[](https://github.com/mqurban/emutrust/actions/workflows/tests.yml)
|
|
45
|
+
[](https://pypi.org/project/emutrust/)
|
|
46
|
+
[](https://pypi.org/project/emutrust/)
|
|
47
|
+
[](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
|
+

|
|
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
|
+

|
|
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
|
+

|
|
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
|
+

|
|
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).
|
emutrust-0.1.0/README.md
ADDED
|
@@ -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
|
+
[](https://github.com/mqurban/emutrust/actions/workflows/tests.yml)
|
|
6
|
+
[](https://pypi.org/project/emutrust/)
|
|
7
|
+
[](https://pypi.org/project/emutrust/)
|
|
8
|
+
[](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
|
+

|
|
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
|
+

|
|
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
|
+

|
|
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
|
+

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