Skip to content

SBU-COSMOLIKE/axionHMcode

 
 

Repository files navigation

axionHMcode

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

Theory

The model computes the non-linear power spectrum by using the fully expanded power spectrum

codesketch

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.

How Does the Code Work?

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 $h/\mathrm{cMpc}$ and $\left(\mathrm{cMpc}/h\right)^3$ respectively. The code also produces a plot of the ratio between the MDM and ΛCDM linear and non-linear power spectra.

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.

HMCode-2020 Parameters

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:

  1. 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.
  2. 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.
  3. 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 with full_2h = False.
  4. A bug in the cold density profile when using the halo bloating term $\eta$ was corrected.
  5. The relations for the critical density threshold, $\delta_c$, and the virial overdensity, $\Delta_{\mathrm{vir}}$, are now calculated as in HMCode-2020 (see Mead et al. (2020), Eq. A1 and A2). This ensures that axionHMcode and HMCode-2020 agree in the case of a ΛCDM cosmology.
  6. The minimum concentration is now $B=5.196$ as found in Mead et al. (2020).
  7. Alex Laguë implemented Numba in this updated version to increase the speed of the code.
  8. Alex Laguë and Keir Rogers also included the optional parameters alpha_1, alpha_2, gamma_1, gamma_2 defined in Dentler er al. (2021), Eq. (36). To use them, just include them in your dictionary of cosmological parameters (the cosmo_dic file) before running the params and power spectra calculation 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 $1 < z < 8$, $0.01 < f_{\mathrm{ax}} = \Omega_{\mathrm{ax}} / \Omega_{\mathrm{m}} < 0.3$ and for axion masses that are close to the pivot mass $m_a=10^{-24.5}$ eV. In addition, in the dome version, we recommend to use alpha = True and concentration_param = True.

Contact data

If you find any bugs or have any questions about the code, please send me a message via GitHub or open an issue.

Fork note (SBU-COSMOLIKE): VM-SPEEDUP

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_formation inverts 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_param and NFW_profile memoize exact repeated evaluations inside the central-density root solver.
  • cosmology/overdensities.py: func_D_z_unnorm_int evaluates 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.py also exposes set_accuracy_boost(x): a single multiplier on the table node counts (default 1.0 = the counts validated below). The AxiECAMB boost theory forwards its yaml accuracy_boost option 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.

Appendix (SBU-COSMOLIKE): mathematics and implementation of the VM-SPEEDUP changes

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, $n_k \simeq 212$, 100-point halo-mass grid) with cProfile; the profiling script is axionhmcode_boost/dev_scripts/profile_axionhmcode.py in the AxiECAMB repository.

A.1 Where the time went in upstream

One dome-version evaluation at $z=0$ cost ~3.5 s, split as:

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, $z=0$): 109,334 growth-factor integrations, 99,055 full $\sigma(M)$ integrations, 6,451 brentq calls, 2.7 million Python function calls in total. Fortran HMcode-2020 evaluates the same class of physics in milliseconds because it tabulates $\sigma(M)$ and the growth once per redshift and gets the formation redshift by inverse interpolation; the fork applies exactly that strategy to the axion stage, without changing any formula.

A.2 The growth-factor table (cosmology/fast_tables.py: growth_tables, D_norm_fast, z_of_D_norm)

The linear growth factor used throughout the code is (upstream overdensities.func_D_z_unnorm, a standard matter-era integral representation)

$$D(z) = \frac{5,\Omega_m}{2}, E(z) \int_z^{100} \frac{1+x}{E(x)^3}, dx, \qquad E(z) = \frac{H(z)}{H_0},$$

which upstream evaluates with a 2000-point trapezoid rule per call, and normalizes per call as $D_{\rm norm}(z) = D(z)/D(0)$ — recomputing the constant $D(0)$ every time.

The fork evaluates the same upstream function once at 4001 linearly spaced nodes $z \in [0, 100]$ (the interval upstream's brentq searched), stores $D_{\rm norm}$ at the nodes, and interpolates linearly (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 $\Delta z = 0.025$ grid is negligible against the halo-model calibration accuracy (measured: $\le 1.6\times10^{-6}$ relative on the validation points). $D_{\rm norm}(z)$ is strictly decreasing, so it is also invertible by the same interpolation run backwards (z_of_D_norm), which is what eliminates the inner brentq (section A.5). The table is cached per $(\Omega_{m,0}, \Omega_{w,0})$ — it depends on nothing else — so its cost is paid once per cosmology, not per redshift.

A.3 The integrated growth (fast_tables.G_integral_fast, replacing scipy.integrate.dblquad)

HMcode-2020 (arXiv:2009.01858, eq. A5) needs the integrated growth

$$G(z) = \frac{5,\Omega_m}{2} \int_z^{z_{\max}} dx, \frac{E(x)}{1+x} \int_x^{z_{\max}} dy, \frac{1+y}{E(y)^3}, \qquad z_{\max} = 10^4,$$

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 $G(z)$ — not just one value — is obtained from a single dense grid by two cumulative passes. Define the inner integral $I(x) = \int_x^{z_{\max}} (1+y),E(y)^{-3}, dy$. On nodes $y_0 &lt; y_1 &lt; \dots &lt; y_N$, the reversed cumulative trapezoid sum

$$I(y_i) = \sum_{j \ge i} \tfrac12,[g(y_j) + g(y_{j+1})],(y_{j+1} - y_j), \qquad g(y) = \frac{1+y}{E(y)^3},$$

gives $I$ at every node in one vectorized pass; a second identical pass over $E(x)(1+x)^{-1} I(x)$ gives $G$ at every node. The fork uses 20,000 geometrically spaced nodes (dense near $z=0$ where the integrand varies fastest, reaching $z=0$ exactly by an offset), builds the table once per cosmology, and interpolates. Measured agreement with the upstream dblquad value: $\le 1.6\times10^{-6}$ relative at $z \in {0, 1, 5, 9.8}$. The upstream code remains in overdensities.func_D_z_unnorm_int after the return statement as the reference implementation.

A.4 The mass-variance table (fast_tables.sigma_M_fast)

The mass variance is (upstream variance.func_sigma_M/func_sigma_r)

$$\sigma^2(M) = \frac{1}{2\pi^2} \int d\ln k; P(k), W^2!\big(k R(M)\big), k^3, \qquad R(M) = \Big(\frac{3M}{4\pi\bar\rho}\Big)^{1/3},$$

with $W(x) = 3(\sin x - x\cos x)/x^3$ the spherical top-hat window; upstream integrates over the full k grid per call.

The fork evaluates that same upstream function at 321 nodes uniform in $\log_{10} M \in [3, 19.5]$ and interpolates $\ln \sigma$ linearly in $\log_{10} M$. The log-log choice is what makes 321 nodes sufficient: $\sigma(M)$ is close to a power law in $M$, so its log-log graph is close to a straight line and piecewise-linear interpolation is nearly exact. Requests outside the tabulated range fall back to the exact upstream evaluation, so the table can never extrapolate. The cache lives inside 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.

A.5 Formation redshift by monotone inversion (halo_model/cold_density_profile.py: func_z_formation)

HMcode-2020 defines the halo formation redshift (arXiv:2009.01858 eq. 21; Dome et al. arXiv:2409.11469 eq. 50) as the solution $z_f$ of

$$\frac{D(z_f)}{D(z)},\sigma(f M, z) = \delta_c(z), \qquad f = 0.01 .$$

Equivalently, with $t \equiv D_{\rm norm}(z),\delta_c(z)/\sigma(fM)$: $D_{\rm norm}(z_f) = t$. Upstream solved this with brentq on $z_f \in [z, 100]$, paying the full growth and variance quadratures at every iteration (section A.1); the fork computes $t$ once (with $\sigma$ from the A.4 table) and reads $z_f = D_{\rm norm}^{-1}(t)$ from the A.2 table.

The subtle part is not the inversion but the clamp. Upstream only ran brentq when its objective changed sign across $[z, 100]$ and returned $z$ otherwise; that branch implements HMcode's minimum-concentration rule (a formal $z_f &lt; z$ would give $c &lt; B$, which the model truncates at $c = B$, $B = 5.196$). The fork reproduces the branch through the same endpoint sign products, evaluated exactly on the tabulated values — scalar path: positive product returns $z$; array path: the inversion is used only where the product is negative. The branch decision is a comparison, never an interpolation, so no halo can land on the wrong side of the clamp because of table smoothing.

A.6 Exact memoization (func_conc_param, NFW_profile)

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 $c(M) = B,(1+z_f)/(1+z)$ do not depend on it at all. Upstream therefore recomputed identical quantities ~30 times per mass, 100 masses per redshift.

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 the axion_dic M_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 the eta_given switch.

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 $\rho_c$ such that the integrated axion profile matches the target axion halo mass; its objective is

$$M_{\rm ax}(\rho_c) = 4\pi \int_0^{r_{\rm vir}} \rho_{\rm ax}(r;\rho_c), r^2, dr,$$

which upstream evaluated by sampling $\rho_{\rm ax}$ on a 2000-point geomspace grid and calling 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, $\int f \approx \sum_i w_i f(r_i)$ with weights depending only on the grid. The fork precomputes $w$ once per grid (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 $(h_0, h_1)$ the unique parabola through the three points integrates to

$$\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 $\ln r$ was implemented and measured first: it is spectrally accurate on each smooth branch of the profile and reproduced upstream to $6\times10^{-6}$ in some cases — but non-monotonically in the node count, with excursions to $|\Delta B/B| \sim 10^{-2}$ at $k \gtrsim 8,h/{\rm cMpc}$. The mechanism: 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 $\rho_c$, and the boost. That is a behavior change relative to upstream, not a quadrature error, so the fork keeps upstream's exact 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 $|\Delta B/B| \sim 7\times10^{-2}$ at high $k$ (measured). That sensitivity is a property of the released, calibrated model — the Dome et al. calibration ran with this grid — so it is treated as part of the model definition, like the fitted constants, and not as a numerical convergence parameter of this fork.

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 $|{\rm guess} - \rho_c| &gt; 100$ rejection heuristic, which must stay bit-identical to upstream.

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

$$\tilde u(k, M) = \frac{4\pi}{M_{\rm ax}} \int_0^{r_{\rm vir}} \rho_{\rm ax}(r), r^2, \frac{\sin kr}{kr}, dr,$$

the spherical Bessel $j_0$ transform of the composed profile. Upstream materialized the kernel $\sin(kr)/(kr)$ as a dense $(n_k \times 2000)$ outer-product array per halo mass and reduced it with Simpson — a direct $O(n_k n_r)$ Hankel transform whose ~$2\times10^6$ transcendental kernel evaluations per mass dominated the stage (~0.38 s per redshift, ~45% of the round-2 cost).

FFTLog (Hamilton 2000) exploits that on a log-uniform grid — which the pinned geomspace grid already is — the biased samples $F(r_n) r_n^{-q}$, $F = \rho r^3$, decompose by FFT into power laws $r^{q+i\eta_m}$, each of which passes through $j_0$ analytically via the Mellin kernel

$$\int_0^\infty x^{s-1} j_0(x), dx = 2^{s-2}\sqrt{\pi};\frac{\Gamma(s/2)}{\Gamma((3-s)/2)}, \qquad 0 < {\rm Re},s < 2,$$

so one forward and one inverse FFT deliver $I(k)$ on the reciprocal log-uniform k grid — every k at once. This is the algorithm of cosmolike's 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):

  1. 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}$, via scipy.special.sici (the same special function upstream uses for the cold NFW profile).
  2. 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.
  3. 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 $|\Delta I|/I(k\to0) \le 2\times10^{-5}$ over $k \in [10^{-4}, 15],h$/cMpc, insensitive to the bias $q$, the window, and the padded length. Cost: 0.11 ms per mass evaluation plus 0.36 ms per-grid setup (the loggamma evaluations of the Mellin kernel, cached in 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.

A.9 The accuracy knob (fast_tables.set_accuracy_boost)

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, $z=0$): max $|\Delta B/B| = 0.7!-!1.1\times10^{-3}$ between the two, dominated by the halo-mass integration grid (also doubled by the option), with the tables of this fork contributing $\lesssim 2\times10^{-6}$ (measured unchanged after the A.8 round). The radial 2000-point grid of A.7/A.8 is deliberately excluded from the scaling (see A.7).

A.10 Validation protocol and results

The gate for every change: max $|\Delta B/B| \le 10^{-4}$ on the final nonlinear boost $B(k,z) = P_{\rm NL}/P_{\rm L}$ against unmodified upstream (commit a85ba26), same transfer inputs, across dome/basic $\times$ $z \in {0, 2}$ $\times$ three cosmologies (a Gaughan-like axion cosmology, the upstream input-file cosmology, and the LCDM limit). The driver (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 $|\Delta B/B| = 1.6\times10^{-5}$ over all cases (deviations carried by the table interpolations — they were unchanged by the memoization round, proving A.6 exact, and unchanged again by the A.8 round); growth/G primitives $\le 1.6\times10^{-6}$; the AxiECAMB boost pytest suite passes against the fork. Measured warm single-redshift times (same machine as A.1):

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.

A.11 What is not changed in strict mode

  • In strict mode (the default) the soliton central-density solver uses scipy.optimize.root with upstream's initial guess and its rejection heuristic (solutions with $|{\rm guess} - \rho_c| &gt; 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 explicit aggressive_optimization flag (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.simpson in func_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).

A.12 The default solver and the legacy_root_finder flag (round 4)

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 $M_{\rm ax}(\rho_c) = M_a(M_c)$ is solved by geometric bracket expansion around the upstream guess plus 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 $&lt; 10^{-6}$), 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):

  1. $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_dic deletes zero-density halos from the M_int integration grid, and interior holes in that grid corrupt every downstream mass integral (measured before the gap-target acceptance was adopted: frac_cluster diverging to $-1.3\times10^5$ and NaN boosts). Aggressive mode makes the closest-achievable acceptance explicit, deterministic, and counted.
  2. 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 $&gt;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 $|\Delta B/B|$ vs upstream $\le 7.2\times10^{-5}$ across the validation cases — table-level agreement, now deterministic and diagnosable. Flag on, basic: up to $2.3\times10^{-2}$ at $k \gtrsim 10,h$/cMpc, attributed to finding 2. Posterior-level: all 115 accepted points of a real-likelihood micro-MCMC (Planck lite + lowl + DESI BAO + DES-Y5 SN + ACT DR6 lensing; dome version) re-evaluated with the default solver against the legacy-solver chain values: max $|\Delta\chi^2| = 2.8\times10^{-2}$, rms $5\times10^{-3}$ — far below statistical relevance, so chains under the two solvers are indistinguishable. Timing: aggressive dome is currently ~0.40-0.44 s per redshift, slightly slower than strict (bracket probing costs a few extra objective evaluations); the mode's value is robustness and diagnosability, not speed.

About

A fully expanded halo model to calculate the non-linear power spectrum in a mixed dark matter cosmology with ultra-light axions

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages

  • Jupyter Notebook 59.5%
  • Python 40.5%