Kosmulator: A Python framework for cosmological inference with MCMC

We present Kosmulator, a modular and vectorised Python framework designed to accelerate the statistical testing of cosmological models. As the theoretical landscape expands beyond standard $Λ$CDM, implementing new expansion histories into traditional Einstein–Boltzmann solvers becomes a significant computational bottleneck. Kosmulator addresses this by leveraging array-native execution and efficient ensemble slice sampling (via Zeus) to perform rapid Bayesian inference. We validate the framework against the industry-standard Cobaya code using a combination of Type Ia Supernovae, Cosmic Chronometers, and Baryon Acoustic Oscillation (BAO) data. Our results demonstrate that Kosmulator reproduces Cobaya’s posterior constraints to within $\leq0.3σ$ statistical agreement on $H_{0}$ and $Ω_{m}$ and $<0.6%$ precision on $χ^{2}$, while achieving a $\sim 4.5\times$ reduction in wall-clock time on a single CPU core compared to a standard MPI-parallelised baseline. Furthermore, we showcase the framework’s utility by constraining the implicit power-law $f(Q)$ “$f_1$CDM” model and demonstrating its automated model selection capabilities (AIC/BIC). Kosmulator is introduced as a “scientific sieve” for rapid hypothesis testing, allowing researchers to efficiently filter theoretical candidates before deploying high-precision resources.

💡 Research Summary

The paper introduces Kosmulator, a pure‑Python, modular, and vectorised framework designed to speed up Bayesian inference for cosmological models, especially those that lie beyond the standard ΛCDM paradigm. Traditional pipelines rely on Einstein–Boltzmann solvers such as CLASS or CAMB, which require C/Fortran extensions for non‑standard physics and thus create a substantial bottleneck when testing new theories. Kosmulator circumvents this by allowing users to define any background expansion history as a simple Python function that returns the dimensionless Hubble rate E(z). Even implicitly defined Friedmann equations, which lack analytic solutions, can be handled by embedding a scalar root‑finder (e.g., fsolve) or, for better performance, a fully vectorised Newton–Raphson iteration that operates on whole redshift arrays at once.

The framework ships with two interchangeable MCMC back‑ends: the affine‑invariant ensemble sampler EMCEE and the ensemble slice sampler Zeus. EMCEE excels for relatively smooth, high‑dimensional posteriors, while Zeus is particularly efficient for highly non‑linear or multimodal distributions often encountered in modified‑gravity models. Users can let Kosmulator automatically select the most appropriate sampler or force a specific choice.

Kosmulator’s data handling is extensive. It provides built‑in likelihood modules for Type Ia supernova compilations (JLA, Pantheon, Pantheon+, Union3, DES‑Y5), Baryon Acoustic Oscillation measurements (including DESI DR1/DR2), Cosmic Chronometers / observational Hubble data, and Planck 2018 CMB likelihoods (TT, TE, EE, lensing) via a thin wrapper around CLASS and the CLIK interface. Big‑Bang Nucleosynthesis constraints (AlterBBN) and phenomenological growth observables f(z) and fσ₈(z) are also supported. The growth module adopts the parametrisation f(z)=Ωₘ(z)^γ, treating the growth index γ as a free parameter or fixing it to the General‑Relativity value, thereby enabling rapid “geometric consistency” checks without full perturbation calculations.

Performance validation is carried out by reproducing a baseline analysis performed with the industry‑standard Cobaya code. Using the same combination of supernovae, cosmic chronometers, and BAO data, Kosmulator recovers posterior means for H₀ and Ωₘ within ≤0.3 σ of Cobaya and matches χ² to better than 0.6 %. Crucially, on a single CPU core Kosmulator achieves a ~4.5× reduction in wall‑clock time compared with a standard MPI‑parallelised Cobaya run. The speed gain stems from (i) vectorised evaluation of background quantities and observables, (ii) elimination of Python‑level loops, and (iii) forced single‑threaded BLAS/OpenMP execution to avoid hidden oversubscription.

A key feature is automated model comparison. By designating a reference model (typically ΛCDM), Kosmulator computes for each candidate the best‑fit likelihood, reduced χ², Akaike Information Criterion (AIC), and Bayesian Information Criterion (BIC). ΔAIC and ΔBIC are reported relative to the reference, giving immediate insight into whether added parameters are justified. While the current release limits itself to these two criteria, the authors outline plans to incorporate DIC, WAIC, LOO‑CV, Bayesian evidence, and other small‑sample corrections in future versions.

The framework’s utility is demonstrated on an implicit power‑law f(Q) model, dubbed “f₁CDM”, defined by the equation
E(z)² = Ωₘ(1+z)³ + (1‑Ωₘ) E(z)^{2n}.
For n ≠ 0 the equation has no closed‑form solution; Kosmulator solves it with a vectorised Newton–Raphson routine and simultaneously fits the growth index γ. The resulting posterior is statistically indistinguishable from the Cobaya analysis, and the AIC/BIC scores indicate no significant improvement over ΛCDM, illustrating how the framework can quickly flag non‑promising theories before committing expensive resources.

In summary, Kosmulator offers: (1) a fully Pythonic, modular architecture that lowers the barrier to implementing novel Friedmann equations; (2) dual, high‑performance MCMC engines; (3) a comprehensive suite of late‑time and early‑universe likelihoods, including seamless CLASS integration for CMB; (4) automatic, publication‑ready statistical outputs and model‑selection diagnostics; and (5) demonstrable speed gains without sacrificing accuracy. By acting as a “scientific sieve,” it enables researchers to rapidly prototype, test, and discard theoretical candidates, reserving full‑scale Boltzmann‑solver pipelines for only the most promising models.