Package: cosmopower-jax
CosmoPower-JAXCosmoPower-JAXin an extension of theCosmoPowerframework to emulate cosmological power spectra in a differentiable way. WithCosmoPower-JAXyou can efficiently run Hamiltonian Monte Carlo with hundreds of parameters (for example, nuisance parameters describing systematic effects), on CPUs and GPUs, in a fraction of the time which would be required with traditional methods. We provide some examples on how to use the neural emulators below, and more applicationsin our paper.Of course, withCosmoPower-JAXyou can also obtain efficient and differentiable predictions of cosmological power spectra. We show how to achieve this in less than 5 lines of code below.InstallationTo installCosmoPower-JAX, you can simply usepip:pip install cosmopower-jaxWe recommend doing it in a freshcondaenvironment, to avoid clashes (e.g.conda create -n cpj python=3.9 && conda activate cpj).Alternatively, you can:git clone https://github.com/dpiras/cosmopower-jax.git
cd cosmopower-jax
pip install .The latter will also give you access to a Jupyter notebook with some examples.Usage & exampleAfter the installation, getting a cosmological power spectrum prediction is as simple as (e.g. for the CMB temperature power spectrum):import numpy as np
from cosmopower_jax.cosmopower_jax import CosmoPowerJAX as CPJ
# omega_b, omega_cdm, h, tau, n_s, ln10^10A_s
cosmo_params = np.array([0.025, 0.11, 0.68, 0.1, 0.97, 3.1])
emulator = CPJ(probe='cmb_tt')
emulator_predictions = emulator.predict(cosmo_params)Similarly, we can also compute derivatives like:emulator_derivatives = emulator.derivative(cosmo_params)Rather than passing an array, as in the originalCosmoPowersyntax you can also pass a dictionary:cosmo_params = {'omega_b': [0.025],
                'omega_cdm': [0.11],
                'h': [0.68],
                'tau_reio': [0.1],
                'n_s': [0.97],
                'ln10^{10}A_s': [3.1],
                }
emulator = CPJ(probe='cmb_tt')
emulator_predictions = emulator.predict(cosmo_params)We also support reusing originalCosmoPowermodels, which you can now use in JAX without retraining. In that case, you should:git clone https://github.com/dpiras/cosmopower-jax.git
   cd cosmopower-jaxand move your model(s).pklfiles into the foldercosmopower_jax/trained_models. At this point:if you can call your models from thecosmopower-jaxfolder you are in, you should be good to go;otherwise, run firstpip install ., and then you should be able to call your custom models from anywhere.To finally call a custom model, you can run:from cosmopower_jax.cosmopower_jax import CosmoPowerJAX as CPJ
emulator_custom = CPJ(probe='custom_log', filename='<custom_filename>.pkl')where<custom_filename>.pklis the filename (only, no path) with your custom model, andcustom_logindicates that your model was trained on log-spectra, so all predictions will be returned elevated to the power of 10. Alternatively, you can passcustom_pca, and you will automatically get the predictions for a model trained withPCAplusNN.We provide a full walkthrough and all instructions in the accompanyingJupyter notebook, and we describeCosmoPower-JAXin detail in the release paper. We currently do not provide the code to train a neural-network model in JAX; if you would like to re-train a JAX-based neural network on different data,raise an issueor contactDavide Piras.Contributing and contactsFeel free toforkthis repository to work on it; otherwise, pleaseraise an issueor contactDavide Piras.CitationIf you useCosmoPower-JAXin your work, please cite both papers as follows:@article{Piras23,
         title={CosmoPower-JAX: high-dimensional Bayesian inference with
         differentiable cosmological emulators},
         author = {Piras, Davide and Spurio Mancini, Alessio},
         journal = {arXiv e-prints},
         year = 2023,
         month = may,
         eid = {arXiv:2305.06347},
         pages = {arXiv:2305.06347},
         archivePrefix = {arXiv},
         eprint = {2305.06347},
         adsurl = {https://ui.adsabs.harvard.edu/abs/2023arXiv230506347P},
        }


@article{SpurioMancini2022,
         title={CosmoPower: emulating cosmological power spectra for 
         accelerated Bayesian inference from next-generation surveys},
         volume={511},
         ISSN={1365-2966},
         url={http://dx.doi.org/10.1093/mnras/stac064},
         DOI={10.1093/mnras/stac064},
         number={2},
         journal={Monthly Notices of the Royal Astronomical Society},
         publisher={Oxford University Press (OUP)},
         author={Spurio Mancini, Alessio and Piras, Davide and 
         Alsing, Justin and Joachimi, Benjamin and Hobson, Michael P},
         year={2022},
         month={Jan},
         pages={1771–1788}
         }LicenseCosmoPower-JAXis released under the GPL-3 license - seeLICENSE-, subject to
the non-commercial use condition - seeLICENSE_EXT.CosmoPower-JAX     
 Copyright (C) 2023 Davide Piras & contributors

 This program is released under the GPL-3 license (see LICENSE), 
 subject to a non-commercial use condition (see LICENSE_EXT).

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.