axionHMcode is a code to compute the non-linear matter power spectrum in a mixed dark matter cosmology with an ultra-light axion (ULA) component of the dark matter as described in Vogt et al. (2022). A very accurate halo model for ΛCDM and massive neutrino cosmologies is given in Mead et al. (2020), referred to as HMCode-2020. Since the axionHMcode model is inspired by their code and uses some of their fitting parameters, our model is named after HMcode. More recently, the code was updated and calibrated to simulations in Dome et al. (2024).
The model computes the non-linear power spectrum by using the fully expanded power spectrum
The cold part can be computed as usually with the standard halo model (see Massara et al. (2014) or Mead et al. (2020)). In contrast, the cross and axion parts have to take into account the non clustering of axions on small scales due to free-streaming. This is done by splitting the axion overdensity into a clustered and linear component. For details see Massara et al. (2014) where the same full treatment was used for massive neutrinos, but can be translated to any other warm/hot, i.e. free streaming/(partially) non-clustering, matter component.
The code expects an input file called input_file.txt which contains information about the cosmology in the style of a Python dictionary. Alternatively, one can also provide the corresponding dictionary directly. Have a look into load_cosmology.py to see what information has to be stored in this dictionary. At the least, the linear matter power spectrum for the target cosmology broken down into different components such as, axions, CDM and baryons saved in a dictionary is needed (see example notebook for the structure). If the linear matter power spectrum is not available, the built-in wrapper to axionCAMB can be used. In other words, if you have a working axionCAMB executable, the code can directly compute the linear power spectrum for the cosmology as specified by input_file.txt.
An example Python file is given in example_file.py. To run the file you have to change the input_file_path and the axionCAMB_exe_path fields (provide the complete path). If the paths are not correct the code will produce an error message. Besides the non-linear total matter power spectrum, the example file also computes the non-linear power spectrum in a ΛCDM cosmology where the axion density is transformed into CDM density. Both power spectra are saved to a file, whose name can be set by the user via the datafile_path variable. The units of the wavenumber and the power spectra are
Moreover, for consistency the example notebook makes a comparison to HMcode-2020 in the ΛCDM case. A comparison is made using the implementation in CAMB and the Python package hmcode.
The axionHMcode can also use the parameters from HMCode-2020 in Mead et al. (2020) which improves the predictions in the case of a CDM cosmology with massive neutrinos. The parameters can be switched on by setting the corresponding variables to True or False in the function for the non-linear power spectrum (make sure the parameters are consistent in the different functions). The parameters are the smoothing parameters, alpha, the halo bloating term, eta_given, the one-halo damping on large scales, one_halo_damping, and the two-halo damping on large scales, two_halo_damping.
Since these parameters are not calibrated to mixed axion - cold dark matter cosmologies, we suggest to not use them in the base version.
Code Updates including calibration from Dome et al. (2024)
Dome et al. (2024) presented an improved version of axionHMcode, which introduced new parameters calibrated to MDM simulations.
Moreover, several updates were made to speed up the code. The changes are summarised in the following:
- The axion mass - cold mass relation,
$M_a(M_c)$ , is now modelled as a broken power law (see Eq. 51 in Dome et al. (2024)) which ensures the old relation of$M_a(M_c) = (\Omega_a/\Omega_c) M_c$ is satisfied above a defined cut-off mass. This new relation is inspired by simulations and is calibrated by them in the redshift range$1 < z < 8$ and for axion fractions of$0.01 < f_{\mathrm{ax}} = \Omega_{\mathrm{ax}} / \Omega_{\mathrm{m}} < 0.3$ around the pivot axion mass$m_a=10^{-24.5}$ eV. -
Dome et al. (2024) introduced new smoothing parameters
$\alpha$ for the cold-cold power spectrum and the cross-power spectrum which depend on the axion mass and axion density. The exact form of these parameters was calibrated to the MDM simulations in the redshift range$1 < z < 8$ and for axion fractions of$0.01 < f_{\mathrm{ax}} = \Omega_{\mathrm{ax}} / \Omega_{\mathrm{m}} < 0.3$ around the pivot axion mass$m_a=10^{-24.5}$ eV. - The two-halo term is now calculated by the linear power spectrum only. The difference between this and the total two-halo term is minor and the speed up is increased. There is still the option to use the full two-halo term by setting the parameter
full_2h = True, but note that the calibrations were performed withfull_2h = False. - A bug in the cold density profile when using the halo bloating term
$\eta$ was corrected. - The relations for the critical density threshold,
$\delta_c$ , and the virial overdensity,$\Delta_{\mathrm{vir}}$ , are now calculated as inHMCode-2020(see Mead et al. (2020), Eq. A1 and A2). This ensures thataxionHMcodeandHMCode-2020agree in the case of a ΛCDM cosmology. - The minimum concentration is now
$B=5.196$ as found in Mead et al. (2020). - Alex Laguë implemented Numba in this updated version to increase the speed of the code.
- Alex Laguë and Keir Rogers also included the optional parameters
alpha_1,alpha_2,gamma_1,gamma_2defined in Dentler er al. (2021), Eq. (36). To use them, just include them in your dictionary of cosmological parameters (thecosmo_dicfile) before running theparamsandpower spectracalculation using e.g.cosmo_dic['alpha_1'] = X. They are not yet included in the input file.
Both the base version and calibrated version from Dome et al. (2024) can be used. To distinguish, set the parameter version in the input file to basic or dome, respectively.
When using the calibrated version from Dome et al. (2024), one has to take into account the redshift range, the axion fraction as well as the axion mass for which the model was calibrated. Therefore, we only recommend to use the dome version when dome version, we recommend to use alpha = True and concentration_param = True.
If you find any bugs or have any questions about the code, please send me a message via GitHub or open an issue.
This fork adds physics-preserving performance fixes to the upstream code, all
marked with greppable # VM-SPEEDUP fences (grep -rn "VM-SPEEDUP" . lists
every change site):
cosmology/fast_tables.py(new): lazily built lookup tables for the growth factor D(z), the HMcode-2020 eq. A5 integrated growth G(z), and sigma(M). Table nodes are computed with the original upstream functions.halo_model/cold_density_profile.py:func_z_formationinverts the tabulated growth instead of running a brentq whose objective re-integrated the growth (including its constant normalization) and sigma(f*M) at every iteration; exact clamp/branch semantics preserved.func_conc_paramandNFW_profilememoize exact repeated evaluations inside the central-density root solver.cosmology/overdensities.py:func_D_z_unnorm_intevaluates the same double integral from a dense cumulative-trapezoid table instead of scipy.dblquad (~0.15 s per call).halo_model/axion_density_profile.py:func_ax_halo_mass(the objective of the soliton central-density solver) integrates the composed profile on upstream's identical 2000-point geomspace grid, but with precomputed canonical composite-Simpson weights (fast_tables.geom_simpson_grid) instead of a scipy.integrate.simpson call — the integral becomes a dot product, removing ~60 us of per-call bookkeeping from the solver's ~3000 objective evaluations per redshift.halo_model/axion_density_profile.py:func_dens_profile_ax_kspace(the k-space profile needed by the one-halo and cross terms) evaluates the spherical Bessel j0 transform with FFTLog on the identical radial samples (fast_tables.fftlog_j0_grid/fftlog_j0_eval; the algorithm of cosmolike's cfftlog with the Bessel order pinned to zero), replacing the dense (n_k x 2000) sin(kr)/(kr) kernel array and its Simpson reduction — two FFTs per mass instead of ~2e6 transcendental kernel evaluations. The truncation edge is handled by a continuous constant extension plus an analytic Si/Ci tail subtraction (appendix A.8).cosmology/fast_tables.pyalso exposesset_accuracy_boost(x): a single multiplier on the table node counts (default 1.0 = the counts validated below). The AxiECAMB boost theory forwards its yamlaccuracy_boostoption here; running the same setup at 1 and 2 is the convergence check. The radial 2000-point grid is deliberately excluded from the scaling (see appendix section A.7).- default solver (appendix A.12; released behavior behind
legacy_root_finder: true): the re-engineered soliton central-density machinery — bracketed brentq with residual classification and per-evaluation diagnostics instead of unbracketed hybr with the |guess - rho_c| > 100 net, plus a continuous crossover-cell correction. Results are not bit-comparable to upstream; dome-version agreement is <= 7e-5 anyway, and the mode's diagnostics exposed silent upstream solver failures in the basic version (A.12).
Validation against unmodified upstream (commit a85ba26): nonlinear boost
grids B(k,z) agree to max |dB/B| = 1.6e-5 across dome/basic, z = 0/2, three
cosmologies including the LCDM limit; growth-table primitives agree with the
original quadratures to <= 1.6e-6. Measured single-redshift speedups: ~8x
(dome, z=0: 2.9-3.5 -> 0.34 s), ~7x (basic, z=0), ~4x (z=2), 16-24x (LCDM
path) — a dome likelihood evaluation drops from ~135 s to ~17 s
single-threaded. The remaining cost is dominated by the soliton
central-density root solver structure (optimize.root re-evaluating the
full profile per iteration), untouched because solver replacement changes
the guess-based no-solution rejection behavior (appendix A.11).
The complete mathematical and implementation documentation of every change, including the measured performance anatomy that motivated it and the validation protocol, is in the appendix below.
This appendix documents each change in full: the quantity involved, how the
unmodified upstream code evaluated it, what the fork does instead, why the
result is either bit-identical or bounded, and the measurements. Every change
site in the source carries a # VM-SPEEDUP fence with a condensed version of
the same information. All profiling numbers were measured on one halo-model
evaluation (one cosmology, one redshift, axionhmcode_boost/dev_scripts/profile_axionhmcode.py in the AxiECAMB
repository.
One dome-version evaluation at
| stage | dome | basic |
|---|---|---|
HMCode_param_dic (cold HMcode-2020 parameters) |
0.00 s | 0.00 s |
func_axion_param_dic (cut mass, soliton central density, clustered fraction) |
2.35 s | 1.22 s |
func_full_halo_model_ax (the halo-model integrals themselves) |
0.43 s | 0.33 s |
The halo-model sums are not the problem; ~85% of the runtime is the axion parameter stage, and the mechanism is structural. That stage solves implicit equations with root finders nested inside root finders:
func_cut_mass_axion_halo: brentq over log10 M
-> func_jeans_virial_ratio(M_trial)
-> func_halo_jeans_kscale(M_trial)
-> func_conc_param(M_trial) [trial masses are off any grid]
-> func_z_formation: a second brentq, over z_f in [z, 100]
-> every iteration re-integrates the growth factor D(z_f)
(including its constant normalization D(0)) and re-integrates
sigma(f*M) over the whole k grid
and the same pattern repeats inside the soliton central-density solver
(axion_density_profile.func_central_density_param). The innermost
quantities are smooth one-variable functions, but upstream re-derived them
from quadratures at every innermost iteration. Measured per redshift (dome,
brentq calls, 2.7 million Python function calls in
total. Fortran HMcode-2020 evaluates the same class of physics in
milliseconds because it tabulates
The linear growth factor used throughout the code is (upstream
overdensities.func_D_z_unnorm, a standard matter-era integral
representation)
which upstream evaluates with a 2000-point trapezoid rule per call, and
normalizes per call as
The fork evaluates the same upstream function once at 4001 linearly spaced
nodes brentq searched), stores
np.interp).
Because the nodes are computed with the unmodified upstream quadrature, the
table reproduces upstream numerics at the nodes exactly; between nodes the
interpolation error of a smooth function on a z_of_D_norm), which is what eliminates the inner brentq
(section A.5). The table is cached per
HMcode-2020 (arXiv:2009.01858, eq. A5) needs the integrated growth
which upstream computed with scipy.integrate.dblquad at ~0.15 s per call,
once per (cosmology, redshift) evaluation.
The key observation is that both integrals run from a moving lower limit to
the same fixed upper limit, so the whole function
gives dblquad
value: overdensities.func_D_z_unnorm_int after the
return statement as the reference implementation.
The mass variance is (upstream variance.func_sigma_M/func_sigma_r)
with
The fork evaluates that same upstream function at 321 nodes uniform in
cosmo_dic under a
'_vm_' key and is fingerprinted by the identity and endpoints of the power
spectrum array, so its lifetime is exactly one (cosmology, redshift)
evaluation — there is no possibility of a stale table leaking across
likelihood points, and multiprocessing fork workers inherit nothing they
should not.
HMcode-2020 defines the halo formation redshift (arXiv:2009.01858 eq. 21;
Dome et al. arXiv:2409.11469 eq. 50) as the solution
Equivalently, with brentq on
The subtle part is not the inversion but the clamp. Upstream only ran
brentq when its objective changed sign across
The soliton central-density solver
(axion_density_profile.func_central_density_param) finds, for each halo
mass, the soliton amplitude such that the integrated axion halo profile
matches the target axion halo mass. scipy.optimize.root evaluates the
objective ~30 times per halo mass, and each evaluation rebuilds the full
density profile — but the solver's unknown only scales the soliton term; the
NFW-shaped part and the concentration
The fork memoizes the two quantities exactly:
-
func_conc_param(scalar masses): keyed on the mass,$\Omega_0$ ,c_min, the identity of the power-spectrum array (id(PS)is stable because the arrays are constructed once per evaluation and never mutated), and theaxion_dicM_cut state that gates the Dentler gamma correction. -
NFW_profile(scalar mass, radial grid): additionally fingerprints the r grid by its length and both endpoints — which identifies a geomspace grid uniquely — plus theeta_givenswitch.
A memoization is exact by construction when the key captures every input the
result depends on; the remaining inputs (redshift, cosmology) are constant
within the memo's lifetime because both memos live in cosmo_dic under
'_vm_' keys, i.e. for one (cosmology, redshift) evaluation. The proof that
nothing was missed is empirical and strict: the fork validation measured
bit-identical boost deviations before and after the memoization round — the
memos changed the runtime and nothing else.
A.7 The soliton mass integral: precomputed weights on upstream's grid (func_ax_halo_mass, fast_tables.geom_simpson_grid)
The soliton central-density solver finds, per halo mass, the soliton
amplitude
which upstream evaluated by sampling scipy.integrate.simpson. With ~30 objective
evaluations per mass and 100 masses per redshift, simpson's per-call
bookkeeping (~60 us of argument validation, slice construction and masked
divisions) recurs ~3000 times per redshift and dominated the solver's cost.
Quadrature by a weight vector. Any fixed-node quadrature rule is a linear
functional of the samples, fast_tables.geom_simpson_grid, cached in cosmo_dic) and evaluates each
integral as a dot product. The weights implement the canonical composite
Simpson rule for unevenly spaced samples — for each pair of adjacent
intervals
$$\frac{h_0+h_1}{6}\left[\Big(2-\frac{h_1}{h_0}\Big) f_i
- \frac{(h_0+h_1)^2}{h_0 h_1} f_{i+1}
- \Big(2-\frac{h_0}{h_1}\Big) f_{i+2}\right],$$
with the leftover final interval closed by a trapezoid. Both pieces are textbook mathematics: no scipy call remains in the hot loop, so the result does not depend on scipy's per-version policy for even point counts (the default handling changed in scipy 1.11 and the compatibility keyword was removed in 1.14; that policy dependence is why cloning scipy's exact weights was rejected). The difference from scipy's current rule is confined to the outermost interval and is orders of magnitude below the gate.
Why the grid must be upstream's, not a better quadrature's. A
Gauss-Legendre rewrite in func_dens_profile_ax composes
the profile as soliton inside, NFW outside, with the crossover radius
detected from sign changes of the sampled difference. The composed profile
is therefore a function of the node positions; for halos where the two
branches barely cross, moving the nodes flips the detection, discretely
changing the profile, the solved np.geomspace(1e-15, r_vir, 2000) nodes —
bit-identical composition, crossover snapping, and rejection behavior — and
changes only the summation rule applied to the identical samples.
For the same reason the 2000-point radial grid is pinned and excluded from
the accuracy_boost scaling: doubling it re-snaps the crossovers and flips
marginal rejections, moving the dome boost by up to
The guess-construction integrals of the solver
(integrate.quad/integrate.simpson in func_central_density_param) are
untouched: they run once per mass (negligible cost) and feed the
A.8 The profile Fourier transform via FFTLog (func_dens_profile_ax_kspace, fast_tables.fftlog_j0_grid / fftlog_j0_eval)
The halo-model one-halo and cross terms need the normalized k-space profile
the spherical Bessel
FFTLog (Hamilton 2000) exploits that on a log-uniform grid — which the
pinned geomspace grid already is — the biased samples
so one forward and one inverse FFT deliver cfftlog (cosmo2D.c) with the Bessel order pinned to zero; the per-order
kernel loop of the non-Limber case (their phase 2) does not exist here.
The radial samples, and hence the profile composition and crossover
snapping, are bit-identical to upstream's; only the reduction changes.
Three details carry the accuracy (all established in the round-3 prototype
scan, dev_scripts):
-
Edge treatment. The profile truncates at
$r_{\rm vir}$ where$F$ is at its maximum; a bare Fourier representation integrates that jump$O(\Delta\ln r)$ differently from Simpson (~$5\times10^{-3}$). The padding is therefore filled with the constant$F(r_{\rm vir})$ — a continuous extension whose biased amplitude decays by$e^{-36q}$ long before the periodic wrap — and its exact tail is subtracted analytically:$\int_{r_{\rm vir}}^\infty F(r_{\rm vir}), j_0(kr), dr/r = F(r_{\rm vir}),[\sin a/a - {\rm Ci}(a)]$ ,$a = k r_{\rm vir}$ , viascipy.special.sici(the same special function upstream uses for the cold NFW profile). -
Padding to
$n_{\rm fft} = 4096$ : a power of two (FFT-friendly small primes — the cosmo2D.c line-4725 consideration) whose extra ~36 e-folds simultaneously push the reciprocal k grid far below$1/r_{\rm vir}$ , so a single transform covers every target k down to$10^{-4},h$ /cMpc. -
Cubic interpolation onto the target k (closed-form 4-point Lagrange
weights on the uniform
$\ln k$ grid): linear interpolation of the oscillatory tail was the dominant residual ($3\times10^{-5}$ →$5\times10^{-6}$ ).
Measured accuracy against upstream's dense-Simpson rule on the same
samples: max cosmo_dic per grid) — the stage drops from ~0.38 s to ~0.05 s per
redshift. FFT count: 2 per mass, 200 per redshift, ~$10^4$ per likelihood
evaluation; all share one padded length, so pocketfft's internal
per-length twiddle cache plays the role of cosmo2D.c's static FFTW plans
(rebuilt never, warmed once per worker process). The transforms are not
additionally threaded: the boost theory's process-level fork over the
redshift loop owns the core budget (one worker per OMP_NUM_THREADS
core), and each FFT batch is ~5 ms — the same resource discipline as
MPI x OpenMP on Cocoa.
All table node counts in A.2-A.4 are scaled by a single multiplier, set from
the AxiECAMB boost theory's yaml option accuracy_boost (modeled on CAMB's
AccuracyBoost and cosmolike's Ntable.high_def_integration). The default
1.0 reproduces the validated node counts exactly. Rerunning the same setup
at accuracy_boost: 1 and 2 and comparing boost grids measures the total
internal-discretization error in one shot; measurement (dome,
The gate for every change: max axionhmcode_boost/dev_scripts/fork_validate.py in the AxiECAMB repo)
builds one set of transfers, runs upstream and fork in separate child
processes, and compares the grids; a primitive-level check compares the A.2
and A.3 tables against the original quadratures directly. The gate is what
caught the Gauss-Legendre crossover flips of A.7 before they shipped.
Results (all rounds, including A.7 and A.8): max
| case | upstream | fork | speedup |
|---|---|---|---|
| dome, z=0 | 2.9-3.5 s | 0.34 s | ~8x |
| basic, z=0 | 1.6-2.0 s | 0.21 s | ~7x |
| dome, z=2 | 1.4-1.8 s | 0.33 s | ~4x |
| basic, z=2 | 0.8-1.0 s | 0.21 s | ~3.6x |
| LCDM limit | 0.1 s | 0.005 s | 16-24x |
Per likelihood evaluation in the AxiECAMB boost pipeline (~50 redshifts),
the dome cost drops from ~135 s (upstream) to ~17 s, before the boost
theory's z-loop fork parallelism divides the wall time (always one worker
per OMP_NUM_THREADS core; the environment's budget is never overridden).
The outsized LCDM-limit ratio deserves a note, because it looks anomalous.
Profiled, upstream's LCDM path spends ~75% of its time in one
scipy.integrate.dblquad evaluation of the HMcode-2020 eq. A5 integrated
growth — 0.15 s to produce a single number (A.3) — and most of the
remainder in the formation-redshift brentq cascade (A.5); the actual
halo-model sums cost ~0.01 s. Everything slow on this path is of the
"re-derive a smooth function from scratch" kind, so nearly all of it
tabulates away and only the honest halo-model arithmetic remains; on the
axion path, by contrast, the remaining cost is dominated by the soliton
central-density solver, which is deliberately not restructured (A.11) —
that is why the LCDM ratio exceeds the axion-case ratio. The ratio implies
nothing about Fortran HMcode-2020: that code already tabulates the growth
and obtains the formation redshift by inverse interpolation — this fork
imported the Fortran code's structure into the python port (A.2, A.5), not
a new trick that would transfer back. Note also that the fork times here
are warm-cache marginal costs: a fresh cosmology pays a one-time ~0.1 s
growth/G table build per worker process, amortized to ~2 ms per redshift
over the pipeline's ~50-node grid.
- In strict mode (the default) the soliton central-density solver uses
scipy.optimize.rootwith upstream's initial guess and its rejection heuristic (solutions with$|{\rm guess} - \rho_c| > 100$ are declared unphysical and set to zero), bit-faithfully. Replacing the solver changes which halos are assigned a soliton — behavior, not precision — so the re-engineered solver lives behind the explicitaggressive_optimizationflag (A.12), with the strict path preserved verbatim. The same reasoning pinned the radial grid in A.7/A.8. The guess-construction integrals (integrate.quad/integrate.simpsoninfunc_central_density_param) are untouched in both modes: once per mass, negligible cost, and the guess seeds both solvers. - The input k grid is never thinned or resampled in either mode: the boost
is evaluated on the transfer grid it receives (halving that grid was
measured to corrupt the boost by 11% through the
$\sigma(M)$ integrals).
The re-engineered soliton solver described here is the default (PI
decision, 2026-07-03, after the validation evidence below); the released
code's verbatim behavior is preserved behind legacy_root_finder: true —
a yaml option of the boost theory, forwarded to the fork like
accuracy_boost (fast_tables.set_legacy_root_finder), displayed
explicitly in the example yamls, recorded in every chain's updated.yaml,
and required to be constant within a chain. The legacy path is what the
A.10 gate certifies against upstream, and it is the only mode a plain
upstream checkout supports (the boost theory hard-errors if the default
solver is requested without the fork).
What the default solver changes. (i) The soliton central-density equation
brentq — guaranteed, tolerance-controlled
convergence — instead of unbracketed hybr with no success check. (ii) A
residual test classifies every outcome, counted per evaluation in
cosmo_dic['_vm_agg_diag']: solved (residual gap_target
(unreachable target, see below; the closest-achievable amplitude is
accepted), no_bracket (no soliton regime found; halo zeroed). (iii) The
objective is made smooth within branches by the interpolated crossover-cell
correction of _ax_halo_mass_aggressive (log-linear crossing radius,
two-piece trapezoid cell split).
What the investigation found (the reason the mode exists):
-
$M_{\rm ax}(\rho_c)$ is discontinuous by construction: at the amplitude where a soliton/NFW crossing is first detected, the integral jumps by the mass of the NFW envelope. Targets inside the jump are exactly unreachable — for the dome Gaughan case at$z=0$ , 33 of 100 halos. Upstream's hybr lands near the jump and its distance-from-guess test silently accepts the closest-achievable amplitude (integrated mass off by up to the jump). Rejecting such halos instead is not an option:func_axion_param_dicdeletes zero-density halos from theM_intintegration grid, and interior holes in that grid corrupt every downstream mass integral (measured before the gap-target acceptance was adopted:frac_clusterdiverging to$-1.3\times10^5$ and NaN boosts). Aggressive mode makes the closest-achievable acceptance explicit, deterministic, and counted. - In the basic version (inputfile cosmology,
$z=2$ ), upstream's hybr returns non-root amplitudes for 18 of 59 halos — the worst off by a factor$3.5\times10^3$ ; the$>100$ rejection net is absolute while$\rho_c \sim 10^{-5}$ , so it filters essentially nothing — and rejects 5 halos that possess perfectly valid roots. Net effect on the clustered fraction: 0.147 (upstream) vs 0.200 (aggressive). The flag-on differences in basic-mode boosts (below) are upstream solver noise, not aggressive-mode error.
Validation. Flag off: byte-identical to the A.10 results (gate rerun
after the round-4 code landed; pytest 4/4). Flag on, dome (the calibrated
version): max
