{"id":24741425,"url":"https://github.com/dpiras/cosmopower-jax","last_synced_at":"2025-09-14T20:57:57.399Z","repository":{"id":154066544,"uuid":"631302429","full_name":"dpiras/cosmopower-jax","owner":"dpiras","description":"Differentiable cosmological emulators: the JAX version of CosmoPower","archived":false,"fork":false,"pushed_at":"2025-06-23T22:59:45.000Z","size":60913,"stargazers_count":40,"open_issues_count":0,"forks_count":3,"subscribers_count":5,"default_branch":"main","last_synced_at":"2025-07-14T23:42:01.696Z","etag":null,"topics":["bayesian-inference","cosmology","deep-learning","emulators","jax","machine-learning"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"gpl-3.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/dpiras.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2023-04-22T15:35:10.000Z","updated_at":"2025-07-02T14:15:00.000Z","dependencies_parsed_at":null,"dependency_job_id":"942a6f29-51a8-4ec0-964a-23c6a7009c63","html_url":"https://github.com/dpiras/cosmopower-jax","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/dpiras/cosmopower-jax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/dpiras%2Fcosmopower-jax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/dpiras%2Fcosmopower-jax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/dpiras%2Fcosmopower-jax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/dpiras%2Fcosmopower-jax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/dpiras","download_url":"https://codeload.github.com/dpiras/cosmopower-jax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/dpiras%2Fcosmopower-jax/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":275167468,"owners_count":25416990,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-09-14T02:00:10.474Z","response_time":75,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["bayesian-inference","cosmology","deep-learning","emulators","jax","machine-learning"],"created_at":"2025-01-27T23:51:59.520Z","updated_at":"2025-09-14T20:57:57.388Z","avatar_url":"https://github.com/dpiras.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"# CosmoPower-JAX\n\n\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"https://user-images.githubusercontent.com/25639122/235351711-39be2b50-dbcb-4964-adbf-f38ffc74ef5f.jpeg\" width=\"300\" height=\"202.5\"\n alt=\"CPJ_logo\"/\u003e\n\u003c/p\u003e\n\u003cdiv align=\"center\"\u003e\n  \n![](https://img.shields.io/badge/Python-181717?style=plastic\u0026logo=python)\n![](https://img.shields.io/badge/Author-Davide%20Piras%20-181717?style=plastic)\n![](https://img.shields.io/badge/Installation-pip%20install%20cosmopower--jax-181717?style=plastic)\n[![arXiv](https://img.shields.io/badge/arXiv-2305.06347-b31b1b.svg)](https://arxiv.org/abs/2305.06347)\n\n\n\u003c/div\u003e\n\n     \n`CosmoPower-JAX` in an extension of the [CosmoPower](https://github.com/alessiospuriomancini/cosmopower) framework to emulate cosmological power spectra in a differentiable way. With `CosmoPower-JAX` you 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 applications [in our paper](https://arxiv.org/abs/2305.06347). You can also have a look at [our poster](https://github.com/dpiras/dpiras.github.io/blob/master/assets/images/poster_CPJ.pdf) presented at the [ML-IAP/CCA-2023](https://indico.iap.fr/event/1/overview) conference, which includes a [video](https://github.com/user-attachments/assets/fd90feab-e3bd-415f-91d5-2ae4f424b449) on `CosmoPower-JAX`.\n\nOf course, with `CosmoPower-JAX` you 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.\n\n## Installation\n\nTo install `CosmoPower-JAX`, you can simply use `pip`:\n\n    pip install cosmopower-jax\n\nWe recommend doing it in a fresh `conda` environment, to avoid clashes (e.g. `conda create -n cpj python=3.9 \u0026\u0026 conda activate cpj`). \n\nAlternatively, you can:\n\n    git clone https://github.com/dpiras/cosmopower-jax.git\n    cd cosmopower-jax\n    pip install . \n\nThe latter will also give you access to a Jupyter notebook with some examples.\n\n## Usage \u0026 example\n\nAfter the installation, getting a cosmological power spectrum prediction is as simple as (e.g. for the CMB temperature power spectrum):\n\n    import numpy as np\n    from cosmopower_jax.cosmopower_jax import CosmoPowerJAX as CPJ\n    # omega_b, omega_cdm, h, tau, n_s, ln10^10A_s\n    cosmo_params = np.array([0.025, 0.11, 0.68, 0.1, 0.97, 3.1])\n    emulator = CPJ(probe='cmb_tt')\n    emulator_predictions = emulator.predict(cosmo_params)\n\nSimilarly, we can also compute derivatives like:\n\n    emulator_derivatives = emulator.derivative(cosmo_params)\n\nRather than passing an array, as in the original `CosmoPower` syntax you can also pass a dictionary:\n\n    cosmo_params = {'omega_b': np.array([0.025]),\n                    'omega_cdm': np.array([0.11]),\n                    'h': np.array([0.68]),\n                    'tau_reio': np.array([0.1]),\n                    'n_s': np.array([0.97]),\n                    'ln10^{10}A_s': np.array([3.1]),\n                    }\n    emulator = CPJ(probe='cmb_tt')\n    emulator_predictions = emulator.predict(cosmo_params)\n\nWe also support reusing original `CosmoPower` models, which you can now use in JAX without retraining. In that case, you should: \n\n```\n   git clone https://github.com/dpiras/cosmopower-jax.git\n   cd cosmopower-jax\n```\n\nand move your model(s) `.pkl` files into the folder `cosmopower_jax/trained_models`. At this point:\n\n- if you can call your models from the `cosmopower-jax` folder you are in, you should be good to go;\n- otherwise, run first `pip install .`, and then you should be able to call your custom models from anywhere.\n \nTo finally call a custom model, you can run:\n\n```\nfrom cosmopower_jax.cosmopower_jax import CosmoPowerJAX as CPJ\nemulator_custom = CPJ(probe='custom_log', filename='\u003ccustom_filename\u003e.pkl')\n```\n\nwhere `\u003ccustom_filename\u003e.pkl` is the filename (only, no path) with your custom model, and `custom_log` indicates that your model was trained on log-spectra, so all predictions will be returned elevated to the power of 10. Alternatively, you can pass `custom_pca`, and you will automatically get the predictions for a model trained with `PCAplusNN`. In this case the parameter dictionary should of course contain the parameter keys corresponding to your trained model. We also allow the full `filepath` of the trained model to be indicated: in this case, do not specify `filename` and only indicate the full `filepath` including the suffix.\n\nWe provide a full walkthrough and all instructions in the accompanying [Jupyter notebook](https://github.com/dpiras/cosmopower-jax/blob/main/notebooks/emulators_example.ipynb), and we describe `CosmoPower-JAX` in 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 issue](https://github.com/dpiras/cosmopower-jax/issues) or contact [Davide Piras](mailto:davide.piras@unige.ch).\n\n### Note if you are using `TensorFlow\u003e=2.14`\nIf you are reusing a model trained with `CosmoPower` and have a `TensorFlow` version higher or equal to 2.14, you might get an error when trying to load the model, even in `CosmoPower-JAX`. This is [a known issue](https://github.com/alessiospuriomancini/cosmopower/issues/22). In this case, you should run the `convert_tf214.py` script available in this repository to transform your `.pkl` file into a different format (based on `NumPy`) that will then be read by `CosmoPower-JAX`. You only have to do the conversion once for each `.pkl` file you have, make sure you `pip install .` after the conversion, and everything else should remain unchanged.\n\n\n## Contributing and contacts\n\nFeel free to [fork](https://github.com/dpiras/cosmopower-jax/fork) this repository to work on it; otherwise, please [raise an issue](https://github.com/dpiras/cosmopower-jax/issues) or contact [Davide Piras](mailto:davide.piras@unige.ch).\n\n## Citation\nIf you use `CosmoPower-JAX` in your work, please cite both papers as follows:\n\n    @article{Piras23,\n             author = {{Piras}, Davide and {Spurio Mancini}, Alessio},\n             title = \"{CosmoPower-JAX: high-dimensional Bayesian inference \n             with differentiable cosmological emulators}\",\n             journal = {The Open Journal of Astrophysics},\n             keywords = {Astrophysics - Cosmology and Nongalactic Astrophysics, \n             Astrophysics - Instrumentation and Methods for Astrophysics, \n             Computer Science - Machine Learning},\n             year = 2023,\n             month = jul,\n             volume = {6},\n             eid = {20},\n             pages = {20},\n             doi = {10.21105/astro.2305.06347},\n             archivePrefix = {arXiv},\n             eprint = {2305.06347},\n             primaryClass = {astro-ph.CO}\n             }\n    \n    @article{SpurioMancini2022,\n             title={CosmoPower: emulating cosmological power spectra for \n             accelerated Bayesian inference from next-generation surveys},\n             volume={511},\n             ISSN={1365-2966},\n             url={http://dx.doi.org/10.1093/mnras/stac064},\n             DOI={10.1093/mnras/stac064},\n             number={2},\n             journal={Monthly Notices of the Royal Astronomical Society},\n             publisher={Oxford University Press (OUP)},\n             author={Spurio Mancini, Alessio and Piras, Davide and \n             Alsing, Justin and Joachimi, Benjamin and Hobson, Michael P},\n             year={2022},\n             month={Jan},\n             pages={1771–1788}\n             }\n             \n\n## License\n\n`CosmoPower-JAX` is released under the GPL-3 license - see [LICENSE](https://github.com/dpiras/cosmopower-jax/blob/main/LICENSE)-, subject to \nthe non-commercial use condition - see [LICENSE_EXT](https://github.com/dpiras/cosmopower-jax/blob/main/LICENSE_EXT).\n\n     CosmoPower-JAX     \n     Copyright (C) 2023 Davide Piras \u0026 contributors\n\n     This program is released under the GPL-3 license (see LICENSE), \n     subject to a non-commercial use condition (see LICENSE_EXT).\n\n     This program is distributed in the hope that it will be useful,\n     but WITHOUT ANY WARRANTY; without even the implied warranty of\n     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdpiras%2Fcosmopower-jax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fdpiras%2Fcosmopower-jax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fdpiras%2Fcosmopower-jax/lists"}