emu-hmf 1.0.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.
Files changed (61) hide show
  1. emu_hmf-1.0.0/CHANGELOG.md +18 -0
  2. emu_hmf-1.0.0/CITATION.cff +72 -0
  3. emu_hmf-1.0.0/LICENSE +29 -0
  4. emu_hmf-1.0.0/MANIFEST.in +7 -0
  5. emu_hmf-1.0.0/PKG-INFO +201 -0
  6. emu_hmf-1.0.0/README.md +160 -0
  7. emu_hmf-1.0.0/docs/Makefile +22 -0
  8. emu_hmf-1.0.0/docs/_static/figures/accuracy.png +0 -0
  9. emu_hmf-1.0.0/docs/_static/figures/correction_vs_nu.png +0 -0
  10. emu_hmf-1.0.0/docs/_static/figures/cosmology_dependence.png +0 -0
  11. emu_hmf-1.0.0/docs/_static/figures/covered_domain.png +0 -0
  12. emu_hmf-1.0.0/docs/_static/figures/growth_with_redshift.png +0 -0
  13. emu_hmf-1.0.0/docs/_static/figures/mass_function.png +0 -0
  14. emu_hmf-1.0.0/docs/_static/figures/sensitivity.png +0 -0
  15. emu_hmf-1.0.0/docs/_static/figures/shape_parameters.png +0 -0
  16. emu_hmf-1.0.0/docs/api/box.rst +19 -0
  17. emu_hmf-1.0.0/docs/api/fit.rst +21 -0
  18. emu_hmf-1.0.0/docs/api/generate.rst +22 -0
  19. emu_hmf-1.0.0/docs/api/model.rst +24 -0
  20. emu_hmf-1.0.0/docs/api/target.rst +32 -0
  21. emu_hmf-1.0.0/docs/changelog.rst +5 -0
  22. emu_hmf-1.0.0/docs/citation.rst +39 -0
  23. emu_hmf-1.0.0/docs/concepts.rst +105 -0
  24. emu_hmf-1.0.0/docs/conf.py +59 -0
  25. emu_hmf-1.0.0/docs/data/sigma_illustrative.npz +0 -0
  26. emu_hmf-1.0.0/docs/halo_model.rst +86 -0
  27. emu_hmf-1.0.0/docs/index.rst +71 -0
  28. emu_hmf-1.0.0/docs/installation.rst +68 -0
  29. emu_hmf-1.0.0/docs/make_figures.py +298 -0
  30. emu_hmf-1.0.0/docs/make_sigma_table.py +107 -0
  31. emu_hmf-1.0.0/docs/massdefs.rst +97 -0
  32. emu_hmf-1.0.0/docs/quickstart.rst +106 -0
  33. emu_hmf-1.0.0/docs/reproducing.rst +111 -0
  34. emu_hmf-1.0.0/docs/requirements.txt +3 -0
  35. emu_hmf-1.0.0/docs/testing.rst +86 -0
  36. emu_hmf-1.0.0/docs/tutorial.rst +151 -0
  37. emu_hmf-1.0.0/docs/validity.rst +143 -0
  38. emu_hmf-1.0.0/emu_hmf/__init__.py +40 -0
  39. emu_hmf-1.0.0/emu_hmf/box.py +84 -0
  40. emu_hmf-1.0.0/emu_hmf/data/emu_hmf_mlp.npz +0 -0
  41. emu_hmf-1.0.0/emu_hmf/data/emu_hmf_mlp_vir.npz +0 -0
  42. emu_hmf-1.0.0/emu_hmf/fit.py +250 -0
  43. emu_hmf-1.0.0/emu_hmf/generate.py +176 -0
  44. emu_hmf-1.0.0/emu_hmf/model.py +192 -0
  45. emu_hmf-1.0.0/emu_hmf/target.py +329 -0
  46. emu_hmf-1.0.0/emu_hmf.egg-info/PKG-INFO +201 -0
  47. emu_hmf-1.0.0/emu_hmf.egg-info/SOURCES.txt +59 -0
  48. emu_hmf-1.0.0/emu_hmf.egg-info/dependency_links.txt +1 -0
  49. emu_hmf-1.0.0/emu_hmf.egg-info/requires.txt +20 -0
  50. emu_hmf-1.0.0/emu_hmf.egg-info/top_level.txt +1 -0
  51. emu_hmf-1.0.0/environment-gen.yml +26 -0
  52. emu_hmf-1.0.0/pyproject.toml +74 -0
  53. emu_hmf-1.0.0/setup.cfg +4 -0
  54. emu_hmf-1.0.0/tests/conftest.py +71 -0
  55. emu_hmf-1.0.0/tests/test_box.py +87 -0
  56. emu_hmf-1.0.0/tests/test_fit.py +263 -0
  57. emu_hmf-1.0.0/tests/test_generate.py +258 -0
  58. emu_hmf-1.0.0/tests/test_model.py +339 -0
  59. emu_hmf-1.0.0/tests/test_packaging.py +124 -0
  60. emu_hmf-1.0.0/tests/test_public_api.py +145 -0
  61. emu_hmf-1.0.0/tests/test_target.py +423 -0
@@ -0,0 +1,18 @@
1
+ # Changelog
2
+
3
+ ## 1.0.0 — 2026-08-29
4
+
5
+ First public release.
6
+
7
+ * Recalibrated Tinker08 multiplicity function at two halo definitions,
8
+ `RockstarM200m` and `RockstarMvir`, each fitted on 2000 CSST cosmologies over
9
+ twelve redshifts and 456 526 rows inside ν ∈ [0.5, 3].
10
+ Residual in ln f, on cosmologies held out entirely: 7.00 % → 0.52 % at 200m,
11
+ 10.92 % → 0.54 % at virial.
12
+ * Pure-JAX inference path: `jax.grad`, `jax.jit` and `jax.vmap` all pass through
13
+ the cosmology. Two runtime dependencies, numpy and JAX.
14
+ * `target.nu_covered(z)` records the peak heights the training set actually
15
+ spans, which narrows with redshift as growth pushes σ down.
16
+ * Both weight files are float64 and carry their own provenance — mass
17
+ definition, sample size, held-out split and residual.
18
+ * Released under the BSD 3-Clause licence.
@@ -0,0 +1,72 @@
1
+ cff-version: 1.2.0
2
+ title: "emu_hmf: a differentiable, cosmology-dependent recalibration of the Tinker08 halo multiplicity function"
3
+ message: "If you use this software, please cite it together with the works it is built on (references below)."
4
+ type: software
5
+ authors:
6
+ - family-names: Comparat
7
+ given-names: Johan
8
+ email: jcomparat@lpsc.in2p3.fr
9
+ orcid: "https://orcid.org/0000-0001-9200-1497"
10
+ repository-code: "https://github.com/JohanComparat/emu_hmf"
11
+ url: "https://emu-hmf.readthedocs.io"
12
+ license: BSD-3-Clause
13
+ version: 1.0.0
14
+ date-released: "2026-08-31"
15
+ keywords:
16
+ - cosmology
17
+ - halo mass function
18
+ - large-scale structure
19
+ - emulator
20
+ - automatic differentiation
21
+ references:
22
+ - type: article
23
+ title: "Toward a Halo Mass Function for Precision Cosmology: The Limits of Universality"
24
+ authors:
25
+ - family-names: Tinker
26
+ given-names: Jeremy
27
+ - family-names: Kravtsov
28
+ given-names: "Andrey V."
29
+ - family-names: Klypin
30
+ given-names: Anatoly
31
+ - family-names: Abazajian
32
+ given-names: Kevork
33
+ - family-names: Warren
34
+ given-names: Michael
35
+ - family-names: Yepes
36
+ given-names: Gustavo
37
+ - family-names: Gottlöber
38
+ given-names: Stefan
39
+ - family-names: Holz
40
+ given-names: "Daniel E."
41
+ journal: "The Astrophysical Journal"
42
+ volume: 688
43
+ start: 709
44
+ year: 2008
45
+ doi: 10.1086/591439
46
+ - type: article
47
+ title: >-
48
+ CSST cosmological emulator II: Generalized accurate halo mass function
49
+ emulation
50
+ authors:
51
+ - family-names: Chen
52
+ given-names: Zhao
53
+ - family-names: Yu
54
+ given-names: Yu
55
+ journal: "Science China Physics, Mechanics, and Astronomy"
56
+ volume: 68
57
+ issue: 10
58
+ start: 109513
59
+ year: 2025
60
+ doi: 10.1007/s11433-025-2764-x
61
+ url: "https://ui.adsabs.harvard.edu/abs/2025SCPMA..6809513C/abstract"
62
+ notes: >-
63
+ The halo mass function emulation this package recalibrates against.
64
+ - type: software
65
+ title: "CSSTemu: emulators for the CSST simulation suite"
66
+ authors:
67
+ - family-names: Chen
68
+ given-names: Zhao
69
+ - family-names: Yu
70
+ given-names: Yu
71
+ year: 2025
72
+ repository-code: "https://github.com/czymh/csstemu"
emu_hmf-1.0.0/LICENSE ADDED
@@ -0,0 +1,29 @@
1
+ BSD 3-Clause License
2
+
3
+ Copyright (c) 2026, Johan Comparat
4
+ All rights reserved.
5
+
6
+ Redistribution and use in source and binary forms, with or without
7
+ modification, are permitted provided that the following conditions are met:
8
+
9
+ 1. Redistributions of source code must retain the above copyright notice, this
10
+ list of conditions and the following disclaimer.
11
+
12
+ 2. Redistributions in binary form must reproduce the above copyright notice,
13
+ this list of conditions and the following disclaimer in the documentation
14
+ and/or other materials provided with the distribution.
15
+
16
+ 3. Neither the name of the copyright holder nor the names of its contributors
17
+ may be used to endorse or promote products derived from this software
18
+ without specific prior written permission.
19
+
20
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
21
+ ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
22
+ WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
23
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
24
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
25
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
26
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
27
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
28
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
@@ -0,0 +1,7 @@
1
+ include LICENSE README.md CHANGELOG.md CITATION.cff
2
+ include environment-gen.yml
3
+ recursive-include emu_hmf/data *.npz
4
+ recursive-include tests *.py
5
+ recursive-include docs *.rst *.py *.txt *.png *.npz Makefile
6
+ prune docs/_build
7
+ global-exclude __pycache__ *.py[cod]
emu_hmf-1.0.0/PKG-INFO ADDED
@@ -0,0 +1,201 @@
1
+ Metadata-Version: 2.4
2
+ Name: emu_hmf
3
+ Version: 1.0.0
4
+ Summary: A differentiable, cosmology-dependent recalibration of the Tinker08 halo multiplicity function
5
+ Author-email: Johan Comparat <jcomparat@lpsc.in2p3.fr>
6
+ License-Expression: BSD-3-Clause
7
+ Project-URL: Homepage, https://github.com/JohanComparat/emu_hmf
8
+ Project-URL: Documentation, https://emu-hmf.readthedocs.io
9
+ Project-URL: Source, https://github.com/JohanComparat/emu_hmf
10
+ Project-URL: Issues, https://github.com/JohanComparat/emu_hmf/issues
11
+ Keywords: cosmology,halo mass function,large-scale structure,emulator,jax,automatic differentiation,Tinker08
12
+ Classifier: Development Status :: 5 - Production/Stable
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: Topic :: Scientific/Engineering :: Astronomy
15
+ Classifier: Topic :: Scientific/Engineering :: Physics
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Programming Language :: Python :: 3.10
18
+ Classifier: Programming Language :: Python :: 3.11
19
+ Classifier: Programming Language :: Python :: 3.12
20
+ Classifier: Programming Language :: Python :: 3.13
21
+ Classifier: Operating System :: OS Independent
22
+ Requires-Python: >=3.10
23
+ Description-Content-Type: text/markdown
24
+ License-File: LICENSE
25
+ Requires-Dist: numpy>=1.24
26
+ Requires-Dist: jax>=0.4.20
27
+ Provides-Extra: train
28
+ Requires-Dist: optax>=0.1.7; extra == "train"
29
+ Provides-Extra: dev
30
+ Requires-Dist: pytest>=7; extra == "dev"
31
+ Requires-Dist: pytest-cov; extra == "dev"
32
+ Requires-Dist: optax>=0.1.7; extra == "dev"
33
+ Requires-Dist: build; extra == "dev"
34
+ Requires-Dist: tomli>=2.0; python_version < "3.11" and extra == "dev"
35
+ Provides-Extra: docs
36
+ Requires-Dist: sphinx; extra == "docs"
37
+ Requires-Dist: sphinx-rtd-theme; extra == "docs"
38
+ Requires-Dist: myst-parser; extra == "docs"
39
+ Requires-Dist: matplotlib; extra == "docs"
40
+ Dynamic: license-file
41
+
42
+ # emu_hmf
43
+
44
+ [![tests](https://github.com/JohanComparat/emu_hmf/actions/workflows/ci.yml/badge.svg)](https://github.com/JohanComparat/emu_hmf/actions/workflows/ci.yml)
45
+ [![docs](https://readthedocs.org/projects/emu-hmf/badge/?version=latest)](https://emu-hmf.readthedocs.io/en/latest/)
46
+ [![PyPI](https://img.shields.io/pypi/v/emu_hmf.svg)](https://pypi.org/project/emu_hmf/)
47
+ [![Python](https://img.shields.io/pypi/pyversions/emu_hmf.svg)](https://pypi.org/project/emu_hmf/)
48
+ [![License](https://img.shields.io/badge/license-BSD--3--Clause-blue.svg)](LICENSE)
49
+
50
+ A differentiable, cosmology-dependent recalibration of the Tinker et al. (2008)
51
+ halo multiplicity function, trained against the CSST emulator
52
+ ([Chen & Yu 2025](https://github.com/czymh/csstemu)) over the box that
53
+ emulator was built on.
54
+
55
+ Tinker08 is a fit to simulations — and not to the simulations anyone compares
56
+ against now. At a Planck cosmology it is offset by a few per cent at z = 0, and
57
+ the size of that offset is itself a function of cosmology and redshift, which a
58
+ fit whose only inputs are σ(M) and z cannot express. This package learns that
59
+ offset.
60
+
61
+ What it learns is *not* a mass function. It is a correction to Tinker08's four
62
+ shape parameters (A, a, b, c) as a function of the eight CSST cosmological
63
+ parameters and redshift:
64
+
65
+ ```
66
+ f(σ) = A [ (σ/b)^-a + 1 ] exp(-c/σ²), with (A, a, b, c) → (A, a, b, c) · e^g(θ, z)
67
+ ```
68
+
69
+ Keeping Tinker08 as the carrier is the whole design. At `g = 0` the answer *is*
70
+ Tinker08, exactly, so the baseline is a point in the same parameterisation
71
+ rather than a different code. The peak-height dependence stays where the physics
72
+ put it, and the network only has to express what the simulations add. And the
73
+ result is a fit with named parameters, so you can ask what the recalibration did
74
+ to the amplitude as against the tilt.
75
+
76
+ ## Install
77
+
78
+ ```bash
79
+ pip install emu_hmf
80
+ ```
81
+
82
+ Two dependencies, numpy and JAX, and 90 kB of trained weights. No Boltzmann
83
+ solver, no Gaussian-process emulator, no training stack, no conda environment —
84
+ a forecast that wants to *evaluate* a mass function should not have to install
85
+ the machinery that fitted one. `tests/test_public_api.py` asserts that split
86
+ rather than trusting it.
87
+
88
+ ### A dedicated environment
89
+
90
+ If you want one anyway — for a reproducible box rather than because the package
91
+ needs it — [`environment.yml`](environment.yml) is the minimal one:
92
+
93
+ ```bash
94
+ mamba env create -f environment.yml # or: conda env create -f environment.yml
95
+ mamba activate emu_hmf
96
+ pip install -e .
97
+ ```
98
+
99
+ It pins the **CPU** build of `jaxlib`: 64 MB against 199 MB for the CUDA one,
100
+ and left unpinned the build depends on whether the machine that solved the
101
+ environment happened to have a driver. The file says how to swap it for a GPU,
102
+ and carries commented blocks for the `[dev]`, `[train]` and `[docs]` extras.
103
+
104
+ ## Use
105
+
106
+ ```python
107
+ import numpy as np
108
+ from emu_hmf.model import HmfCorrection
109
+
110
+ corr = HmfCorrection() # 200m; HmfCorrection(WEIGHTS["vir"]) for virial
111
+
112
+ theta = np.array([0.049, 0.31, 67.36, 0.9649, 2.1, -1.0, 0.0, 0.06])
113
+ # Ω_b Ω_cb H0 n_s 10⁹A_s w w_a Σm_ν
114
+
115
+ f = corr.fsigma(sigma=0.8, theta=theta, z=0.5) # the multiplicity function
116
+ n = corr.dndlnM(m, sigma, dlnsigma_dlnm, rho_cold, theta, z=0.5) # the abundance
117
+ ```
118
+
119
+ σ(M) is passed in, not computed: this package has no power spectrum and should
120
+ not acquire one, and the σ(M) the fit was made against is the *cold* field
121
+ against ρ̄_cb. Fitting f(σ) against one variance and evaluating it with another
122
+ is the mismatch that makes a multiplicity function look wrong when the
123
+ convention around it is what moved.
124
+
125
+ Everything is JAX, so `jax.grad`, `jax.jit` and `jax.vmap` all work through the
126
+ cosmology. That is the reason this exists rather than a table of numbers.
127
+
128
+ ## Two mass definitions, two files
129
+
130
+ The correction is not the same function at two halo boundaries, so there is no
131
+ single correction with a Δ argument. Both are fitted against a *Rockstar*
132
+ spherical-overdensity mass, so the comparison isolates the boundary rather than
133
+ mixing in a change of halo finder.
134
+
135
+ | weights | halo definition | Tinker08 unchanged | recalibrated | improvement |
136
+ | --- | --- | --- | --- | --- |
137
+ | `WEIGHTS["200m"]` | SO 200 × mean, Rockstar | 7.00 % | **0.52 %** | 13.1× |
138
+ | `WEIGHTS["vir"]` | SO virial, Rockstar | 10.92 % | **0.54 %** | 19.1× |
139
+
140
+ rms in ln f, on 200 cosmologies held out *entirely* from training — not held-out
141
+ rows. Each design contributes several hundred rows and at fixed cosmology ln f
142
+ is smooth in σ, so a random row split measures interpolation between neighbouring
143
+ masses of a cosmology the network has already seen. This correction is only ever
144
+ asked for a cosmology it has not seen.
145
+
146
+ Both files carry `WEIGHTS["200m"]`'s Δ = 200m Tinker08 as the carrier, so the
147
+ virial weights absorb the *boundary change* as well as the recalibration. They
148
+ are not a per-cent correction: at z = 0 they sit some 13 % below the 200m
149
+ carrier. Reading "correction" as "small" at `vir` is a misreading.
150
+
151
+ ## Where it is defined
152
+
153
+ Outside either bound the package refuses rather than extrapolating.
154
+
155
+ **The cosmology** must be inside CSST's box, which is copied into `box.py` and
156
+ checked against the emulator's own `param_limits` by `tests/test_box.py`:
157
+
158
+ | Ω_b | Ω_cb | H₀ | n_s | 10⁹A_s | w | w_a | Σm_ν |
159
+ | --- | --- | --- | --- | --- | --- | --- | --- |
160
+ | 0.04–0.06 | 0.24–0.40 | 60–80 | 0.92–1.00 | 1.7–2.5 | −1.3–−0.7 | −0.5–0.5 | 0–0.3 |
161
+
162
+ Note Ω_cb: CSST bounds the **cold** density, with massive neutrinos excluded.
163
+
164
+ **The peak height** must be inside ν = δ_c/σ ∈ [0.5, 3], and the mass inside
165
+ 10¹²–10¹⁴ M⊙/h. Those two cuts do not commute with redshift: growth pushes σ
166
+ down, so a fixed mass is a higher peak later, and the low-ν half of the band is
167
+ simply absent above z ≈ 0.25. `target.nu_covered(z)` records what the training
168
+ set actually spans — ν ≥ 1.4 by z = 3 — because a caller who checked only the
169
+ nominal range would be extrapolating with no warning.
170
+
171
+ The correction is a few per cent at z = 0 and grows with redshift, reaching
172
+ about 12 % rms by z = 3. Quoting the low-redshift figure alone would understate
173
+ it several-fold over most of the range it is defined on.
174
+
175
+ ## Documentation
176
+
177
+ [emu-hmf.readthedocs.io](https://emu-hmf.readthedocs.io) — concepts, a tutorial
178
+ with figures, the validity domain, and how to reproduce the training set.
179
+
180
+ ## Reproducing
181
+
182
+ The 2000-cosmology training set (11.7 MB, both mass definitions) is archived
183
+ with a DOI; the fit that turns it into the shipped weights needs only `optax`:
184
+
185
+ ```bash
186
+ pip install emu_hmf[train]
187
+ python -m emu_hmf.fit --shards ./shards --out weights.npz
188
+ ```
189
+
190
+ Regenerating the shards themselves needs CLASS and the CSST emulator; see the
191
+ documentation's *Reproducing the training set* page.
192
+
193
+ ## Citation
194
+
195
+ If you use this package, please cite Tinker et al. (2008) for the functional
196
+ form, Chen & Yu (2025) for the CSST emulator this is calibrated against, and
197
+ this package for the recalibration. See `CITATION.cff`.
198
+
199
+ ## Licence
200
+
201
+ BSD 3-Clause. See [LICENSE](LICENSE).
@@ -0,0 +1,160 @@
1
+ # emu_hmf
2
+
3
+ [![tests](https://github.com/JohanComparat/emu_hmf/actions/workflows/ci.yml/badge.svg)](https://github.com/JohanComparat/emu_hmf/actions/workflows/ci.yml)
4
+ [![docs](https://readthedocs.org/projects/emu-hmf/badge/?version=latest)](https://emu-hmf.readthedocs.io/en/latest/)
5
+ [![PyPI](https://img.shields.io/pypi/v/emu_hmf.svg)](https://pypi.org/project/emu_hmf/)
6
+ [![Python](https://img.shields.io/pypi/pyversions/emu_hmf.svg)](https://pypi.org/project/emu_hmf/)
7
+ [![License](https://img.shields.io/badge/license-BSD--3--Clause-blue.svg)](LICENSE)
8
+
9
+ A differentiable, cosmology-dependent recalibration of the Tinker et al. (2008)
10
+ halo multiplicity function, trained against the CSST emulator
11
+ ([Chen & Yu 2025](https://github.com/czymh/csstemu)) over the box that
12
+ emulator was built on.
13
+
14
+ Tinker08 is a fit to simulations — and not to the simulations anyone compares
15
+ against now. At a Planck cosmology it is offset by a few per cent at z = 0, and
16
+ the size of that offset is itself a function of cosmology and redshift, which a
17
+ fit whose only inputs are σ(M) and z cannot express. This package learns that
18
+ offset.
19
+
20
+ What it learns is *not* a mass function. It is a correction to Tinker08's four
21
+ shape parameters (A, a, b, c) as a function of the eight CSST cosmological
22
+ parameters and redshift:
23
+
24
+ ```
25
+ f(σ) = A [ (σ/b)^-a + 1 ] exp(-c/σ²), with (A, a, b, c) → (A, a, b, c) · e^g(θ, z)
26
+ ```
27
+
28
+ Keeping Tinker08 as the carrier is the whole design. At `g = 0` the answer *is*
29
+ Tinker08, exactly, so the baseline is a point in the same parameterisation
30
+ rather than a different code. The peak-height dependence stays where the physics
31
+ put it, and the network only has to express what the simulations add. And the
32
+ result is a fit with named parameters, so you can ask what the recalibration did
33
+ to the amplitude as against the tilt.
34
+
35
+ ## Install
36
+
37
+ ```bash
38
+ pip install emu_hmf
39
+ ```
40
+
41
+ Two dependencies, numpy and JAX, and 90 kB of trained weights. No Boltzmann
42
+ solver, no Gaussian-process emulator, no training stack, no conda environment —
43
+ a forecast that wants to *evaluate* a mass function should not have to install
44
+ the machinery that fitted one. `tests/test_public_api.py` asserts that split
45
+ rather than trusting it.
46
+
47
+ ### A dedicated environment
48
+
49
+ If you want one anyway — for a reproducible box rather than because the package
50
+ needs it — [`environment.yml`](environment.yml) is the minimal one:
51
+
52
+ ```bash
53
+ mamba env create -f environment.yml # or: conda env create -f environment.yml
54
+ mamba activate emu_hmf
55
+ pip install -e .
56
+ ```
57
+
58
+ It pins the **CPU** build of `jaxlib`: 64 MB against 199 MB for the CUDA one,
59
+ and left unpinned the build depends on whether the machine that solved the
60
+ environment happened to have a driver. The file says how to swap it for a GPU,
61
+ and carries commented blocks for the `[dev]`, `[train]` and `[docs]` extras.
62
+
63
+ ## Use
64
+
65
+ ```python
66
+ import numpy as np
67
+ from emu_hmf.model import HmfCorrection
68
+
69
+ corr = HmfCorrection() # 200m; HmfCorrection(WEIGHTS["vir"]) for virial
70
+
71
+ theta = np.array([0.049, 0.31, 67.36, 0.9649, 2.1, -1.0, 0.0, 0.06])
72
+ # Ω_b Ω_cb H0 n_s 10⁹A_s w w_a Σm_ν
73
+
74
+ f = corr.fsigma(sigma=0.8, theta=theta, z=0.5) # the multiplicity function
75
+ n = corr.dndlnM(m, sigma, dlnsigma_dlnm, rho_cold, theta, z=0.5) # the abundance
76
+ ```
77
+
78
+ σ(M) is passed in, not computed: this package has no power spectrum and should
79
+ not acquire one, and the σ(M) the fit was made against is the *cold* field
80
+ against ρ̄_cb. Fitting f(σ) against one variance and evaluating it with another
81
+ is the mismatch that makes a multiplicity function look wrong when the
82
+ convention around it is what moved.
83
+
84
+ Everything is JAX, so `jax.grad`, `jax.jit` and `jax.vmap` all work through the
85
+ cosmology. That is the reason this exists rather than a table of numbers.
86
+
87
+ ## Two mass definitions, two files
88
+
89
+ The correction is not the same function at two halo boundaries, so there is no
90
+ single correction with a Δ argument. Both are fitted against a *Rockstar*
91
+ spherical-overdensity mass, so the comparison isolates the boundary rather than
92
+ mixing in a change of halo finder.
93
+
94
+ | weights | halo definition | Tinker08 unchanged | recalibrated | improvement |
95
+ | --- | --- | --- | --- | --- |
96
+ | `WEIGHTS["200m"]` | SO 200 × mean, Rockstar | 7.00 % | **0.52 %** | 13.1× |
97
+ | `WEIGHTS["vir"]` | SO virial, Rockstar | 10.92 % | **0.54 %** | 19.1× |
98
+
99
+ rms in ln f, on 200 cosmologies held out *entirely* from training — not held-out
100
+ rows. Each design contributes several hundred rows and at fixed cosmology ln f
101
+ is smooth in σ, so a random row split measures interpolation between neighbouring
102
+ masses of a cosmology the network has already seen. This correction is only ever
103
+ asked for a cosmology it has not seen.
104
+
105
+ Both files carry `WEIGHTS["200m"]`'s Δ = 200m Tinker08 as the carrier, so the
106
+ virial weights absorb the *boundary change* as well as the recalibration. They
107
+ are not a per-cent correction: at z = 0 they sit some 13 % below the 200m
108
+ carrier. Reading "correction" as "small" at `vir` is a misreading.
109
+
110
+ ## Where it is defined
111
+
112
+ Outside either bound the package refuses rather than extrapolating.
113
+
114
+ **The cosmology** must be inside CSST's box, which is copied into `box.py` and
115
+ checked against the emulator's own `param_limits` by `tests/test_box.py`:
116
+
117
+ | Ω_b | Ω_cb | H₀ | n_s | 10⁹A_s | w | w_a | Σm_ν |
118
+ | --- | --- | --- | --- | --- | --- | --- | --- |
119
+ | 0.04–0.06 | 0.24–0.40 | 60–80 | 0.92–1.00 | 1.7–2.5 | −1.3–−0.7 | −0.5–0.5 | 0–0.3 |
120
+
121
+ Note Ω_cb: CSST bounds the **cold** density, with massive neutrinos excluded.
122
+
123
+ **The peak height** must be inside ν = δ_c/σ ∈ [0.5, 3], and the mass inside
124
+ 10¹²–10¹⁴ M⊙/h. Those two cuts do not commute with redshift: growth pushes σ
125
+ down, so a fixed mass is a higher peak later, and the low-ν half of the band is
126
+ simply absent above z ≈ 0.25. `target.nu_covered(z)` records what the training
127
+ set actually spans — ν ≥ 1.4 by z = 3 — because a caller who checked only the
128
+ nominal range would be extrapolating with no warning.
129
+
130
+ The correction is a few per cent at z = 0 and grows with redshift, reaching
131
+ about 12 % rms by z = 3. Quoting the low-redshift figure alone would understate
132
+ it several-fold over most of the range it is defined on.
133
+
134
+ ## Documentation
135
+
136
+ [emu-hmf.readthedocs.io](https://emu-hmf.readthedocs.io) — concepts, a tutorial
137
+ with figures, the validity domain, and how to reproduce the training set.
138
+
139
+ ## Reproducing
140
+
141
+ The 2000-cosmology training set (11.7 MB, both mass definitions) is archived
142
+ with a DOI; the fit that turns it into the shipped weights needs only `optax`:
143
+
144
+ ```bash
145
+ pip install emu_hmf[train]
146
+ python -m emu_hmf.fit --shards ./shards --out weights.npz
147
+ ```
148
+
149
+ Regenerating the shards themselves needs CLASS and the CSST emulator; see the
150
+ documentation's *Reproducing the training set* page.
151
+
152
+ ## Citation
153
+
154
+ If you use this package, please cite Tinker et al. (2008) for the functional
155
+ form, Chen & Yu (2025) for the CSST emulator this is calibrated against, and
156
+ this package for the recalibration. See `CITATION.cff`.
157
+
158
+ ## Licence
159
+
160
+ BSD 3-Clause. See [LICENSE](LICENSE).
@@ -0,0 +1,22 @@
1
+ SPHINXBUILD ?= sphinx-build
2
+ SOURCEDIR = .
3
+ BUILDDIR = _build
4
+
5
+ .PHONY: help html strict figures clean
6
+
7
+ help:
8
+ @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)"
9
+
10
+ html:
11
+ @$(SPHINXBUILD) -b html "$(SOURCEDIR)" "$(BUILDDIR)/html"
12
+
13
+ strict:
14
+ @$(SPHINXBUILD) -W -b html "$(SOURCEDIR)" "$(BUILDDIR)/html"
15
+
16
+ # Regenerate the committed figures. Not part of `html`: the documentation
17
+ # builder has no matplotlib, and the PNGs are in the repository.
18
+ figures:
19
+ @python make_figures.py
20
+
21
+ clean:
22
+ rm -rf "$(BUILDDIR)"
@@ -0,0 +1,19 @@
1
+ emu_hmf.box
2
+ ===========
3
+
4
+ The hypercube the recalibration lives in, taken from the CSST emulator rather
5
+ than chosen, and copied here rather than imported so that evaluating a mass
6
+ function does not require installing a Gaussian-process emulator to find out
7
+ what the bounds are. ``tests/test_box.py`` asserts the copy against the
8
+ emulator's own ``param_limits``.
9
+
10
+ * :data:`~emu_hmf.box.PARAMS` — the column order, which is the emulator's own.
11
+ * :data:`~emu_hmf.box.BOX` — the closed bounds. Note that ``Omegam`` is the
12
+ **cold** density.
13
+ * :func:`~emu_hmf.box.sample` — a Latin hypercube, deterministic in its seed.
14
+ * :func:`~emu_hmf.box.check` — raise, naming every parameter outside the box.
15
+ * :func:`~emu_hmf.box.inside` — the same, as data rather than an exception.
16
+
17
+ .. automodule:: emu_hmf.box
18
+ :members:
19
+ :show-inheritance:
@@ -0,0 +1,21 @@
1
+ emu_hmf.fit
2
+ ===========
3
+
4
+ Fitting the four Tinker08 parameters as functions of cosmology and redshift.
5
+ Needs ``optax``, which is the ``[train]`` extra — and nothing else, so this is
6
+ the step an outsider can rerun from the archived shards.
7
+
8
+ * :class:`~emu_hmf.fit.Shards` — one flat training set, with the record of what
9
+ it was built from.
10
+ * :func:`~emu_hmf.fit.load_shards` — every shard in a directory, cut to the
11
+ trusted peak heights, refusing to mix mass definitions.
12
+ * :func:`~emu_hmf.fit.fit` — the optimisation. Splits on whole cosmologies,
13
+ pins float64, and writes its provenance into the output.
14
+
15
+ .. code-block:: bash
16
+
17
+ python -m emu_hmf.fit --shards ./shards --out weights.npz
18
+
19
+ .. automodule:: emu_hmf.fit
20
+ :members:
21
+ :show-inheritance:
@@ -0,0 +1,22 @@
1
+ emu_hmf.generate
2
+ ================
3
+
4
+ Building the training set: the emulator's abundance, converted into a
5
+ multiplicity function in this package's variance convention. Needs the
6
+ generation stack — see :doc:`../reproducing`.
7
+
8
+ * :func:`~emu_hmf.generate.solve_one` — one cosmology: one CLASS solve and one
9
+ emulator call.
10
+ * :func:`~emu_hmf.generate.shard` — one contiguous slice of the design, written
11
+ in chunks and resumable.
12
+ * :data:`~emu_hmf.generate.M_GRID`, :data:`~emu_hmf.generate.K_GRID`,
13
+ :data:`~emu_hmf.generate.CHUNK` — the grids and the write cadence.
14
+
15
+ .. code-block:: bash
16
+
17
+ python -m emu_hmf.generate --shard 0 --n-per-shard 250 --n-total 2000 \
18
+ --out shards/hmf_000.npz --massdef RockstarM200m
19
+
20
+ .. automodule:: emu_hmf.generate
21
+ :members:
22
+ :show-inheritance:
@@ -0,0 +1,24 @@
1
+ emu_hmf.model
2
+ =============
3
+
4
+ Inference: Tinker08 with cosmology-dependent parameters, in pure JAX. This is
5
+ the only module a user of the released package needs, and the only one that
6
+ imports nothing beyond numpy and JAX.
7
+
8
+ * :func:`~emu_hmf.model.tinker08` — the carrier, at
9
+ :math:`\Delta_{\rm m} = 200`, optionally corrected. With ``g=None`` it is the
10
+ published fit, unchanged.
11
+ * :class:`~emu_hmf.model.HmfCorrection` — the recalibrated multiplicity
12
+ function. :meth:`~emu_hmf.model.HmfCorrection.fsigma` gives
13
+ :math:`f(\sigma)`, :meth:`~emu_hmf.model.HmfCorrection.dndlnM` the abundance,
14
+ and :meth:`~emu_hmf.model.HmfCorrection.g` the four log-corrections
15
+ themselves.
16
+ * :data:`~emu_hmf.model.WEIGHTS` — one file per halo definition, keyed
17
+ ``"200m"`` and ``"vir"``. See :doc:`../massdefs`.
18
+ * :func:`~emu_hmf.model.load_weights` — the arrays and the provenance, cached
19
+ and read-only.
20
+ * :func:`~emu_hmf.model.normalise` — :math:`(\theta, z)` onto the unit cube.
21
+
22
+ .. automodule:: emu_hmf.model
23
+ :members:
24
+ :show-inheritance:
@@ -0,0 +1,32 @@
1
+ emu_hmf.target
2
+ ==============
3
+
4
+ What is being fitted, and the conversions that reach it. The functions that
5
+ call the CSST emulator import it lazily, so this module is importable without
6
+ it — but calling them is not.
7
+
8
+ * :data:`~emu_hmf.target.FIDUCIAL` — a Planck-like point in the middle of the
9
+ box; the cosmology every worked example and quoted number here uses.
10
+ * :data:`~emu_hmf.target.MASSDEFS`, :data:`~emu_hmf.target.DEFAULT_MASSDEF` —
11
+ the three halo definitions the emulator offers, and what each one *is*.
12
+ * :data:`~emu_hmf.target.M_TRUSTED`, :data:`~emu_hmf.target.NU_TRUSTED`,
13
+ :data:`~emu_hmf.target.DELTA_C` — the fitted domain.
14
+ * :data:`~emu_hmf.target.NU_COVERED`, :func:`~emu_hmf.target.nu_covered` — what
15
+ the training set *actually* spans at each redshift, which is narrower. See
16
+ :doc:`../validity`.
17
+ * :func:`~emu_hmf.target.to_ggah_cosmology`,
18
+ :func:`~emu_hmf.target.theta_from_cosmology` — the two directions of the
19
+ cosmology conversion, and the only places it happens.
20
+ * :func:`~emu_hmf.target.csst_dndlnM`, :func:`~emu_hmf.target.csst_tinker08`,
21
+ :func:`~emu_hmf.target.sigma_chain` — the offline half; these need the
22
+ generation stack.
23
+
24
+ .. note::
25
+
26
+ ``ggah_mod`` in the text below is the halo-model code this package was built
27
+ for, and ``emu_pk`` its linear-spectrum emulator. Neither is needed to
28
+ *use* ``emu_hmf``; see :doc:`../halo_model`.
29
+
30
+ .. automodule:: emu_hmf.target
31
+ :members:
32
+ :show-inheritance:
@@ -0,0 +1,5 @@
1
+ Changelog
2
+ =========
3
+
4
+ .. include:: ../CHANGELOG.md
5
+ :parser: myst_parser.sphinx_