# jax-cosmo
[](https://gitter.im/DifferentiableUniverseInitiative/jax_cosmo?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge) [](https://jax-cosmo.readthedocs.io/en/latest/?badge=latest) []() [](https://github.com/psf/black) [](https://pypi.org/project/jax-cosmo/) [](CODE_OF_CONDUCT.md) [](https://github.com/google/jax-cosmo/blob/master/LICENSE) <!-- ALL-CONTRIBUTORS-BADGE:START - Do not remove or modify this section -->
[](#contributors-)
<!-- ALL-CONTRIBUTORS-BADGE:END -->
<h3 align="center">Finally a differentiable cosmology library, and it's in JAX!</h3>
Have a look at the [GitHub issues](https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues) to see what is needed or if you have any thoughts on the design, and don't hesitate to join the [Gitter room](https://gitter.im/DifferentiableUniverseInitiative/jax_cosmo) for discussions.
## TL;DR
This is what `jax-cosmo` aims to do:
```python
...
def likelihood(cosmo):
# Compute mean and covariance of angular Cls, for specific probes
mu, cov = jax_cosmo.angular_cl.gaussian_cl_covariance_and_mean(cosmo, ell, probes)
# Return likelihood value
return jax_cosmo.likelihood.gaussian_log_likelihood(data, mu, cov)
# Compute derivatives of the likelihood with respect to cosmological parameters
g = jax.grad(likelihood)(cosmo)
# Compute Fisher matrix of cosmological parameters
F = - jax.hessian(likelihood)(cosmo)
```
This is how you can compute **gradients and hessians of any functions in `jax-cosmo`**,
all of this without any finite differences.
Check out a full example here: [](https://colab.research.google.com/github/DifferentiableUniverseInitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb)
Have a look at the [design document](design.md) to learn more about the structure of the code.
## What is JAX?
[JAX](https://github.com/google/jax) = NumPy + autodiff + GPU
JAX is a framework for automatic differentiation (like TensorFlow or PyTorch) but following the NumPy API, and using the GPU/TPU enable XLA backend.
What does that mean?
- You write plain Python/NumPy code, no need to learn a different language
- It runs on GPU, you don't need to do anything particular
- You can take derivatives of any quantity with respect to any parameters by
automatic differentiation.
Checkout the [JAX](https://github.com/google/jax) project page to learn more!
## Install
`jax-cosmo` is pure Python, so installing is a breeze:
```bash
$ pip install jax-cosmo
```
## Philosophy
Here are some of the design guidelines:
- Implementation of equations should be human readable, and documentation should always live next to the implementation.
- Should always be trivially installable: external dependencies should be kept
to a minimum, especially the ones that require compilation or with restrictive licenses.
- Keep API and implementation simple and intuitive, minimize user and developer
surprise.
- “Debugging is twice as hard as writing the code in the first place. Therefore, if you write the code as cleverly as possible, you are, by definition, not smart enough to debug it.” -Brian Kernighan, quote stolen from
[here](https://flax.readthedocs.io/en/latest/philosophy.html).
## Contributing
`jax-cosmo` aims to be a community effort, contributions are most welcome and
can come in several forms
- Bug reports
- API design suggestions
- (Pull) requests for more features
- Examples and notebooks of cool things that can be done with the code
You can chime-in on any aspects of the design by proposing a PR to the [design document](design.md). The issue page is a good place to start, but don't hesitate to come chat in the Gitter room.
Please take a look at the [Contributing Document](CONTRIBUTING.md) for more information.
This project follows the [All Contributors](https://allcontributors.org/) guidelines aiming at recognizing and valorizing
contributions at any levels.
## Contributors ✨
Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):
<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
<!-- prettier-ignore-start -->
<!-- markdownlint-disable -->
<table>
<tr>
<td align="center"><a href="http://flanusse.net"><img src="https://avatars0.githubusercontent.com/u/861591?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Francois Lanusse</b></sub></a><br /><a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=EiffL" title="Code">💻</a></td>
<td align="center"><a href="http://www.cosmostat.org/people/santiago-casas"><img src="https://avatars0.githubusercontent.com/u/6987716?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Santiago Casas</b></sub></a><br /><a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues?q=author%3Asantiagocasas" title="Bug reports">🐛</a> <a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=santiagocasas" title="Code">💻</a></td>
<td align="center"><a href="https://github.com/austinpeel"><img src="https://avatars0.githubusercontent.com/u/17024310?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Austin Peel</b></sub></a><br /><a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=austinpeel" title="Code">💻</a></td>
<td align="center"><a href="https://minaskaramanis.com"><img src="https://avatars2.githubusercontent.com/u/23280751?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Minas Karamanis</b></sub></a><br /><a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=minaskar" title="Code">💻</a></td>
<td align="center"><a href="https://faculty.sites.uci.edu/dkirkby/"><img src="https://avatars1.githubusercontent.com/u/185007?v=4?s=100" width="100px;" alt=""/><br /><sub><b>David Kirkby</b></sub></a><br /><a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=dkirkby" title="Code">💻</a> <a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues?q=author%3Adkirkby" title="Bug reports">🐛</a></td>
<td align="center"><a href="https://aboucaud.github.io"><img src="https://avatars0.githubusercontent.com/u/3065310?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Alexandre Boucaud</b></sub></a><br /><a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=aboucaud" title="Code">💻</a></td>
<td align="center"><a href="https://www.cosmostat.org/people/denise-lanzieri"><img src="https://avatars.githubusercontent.com/u/72620117?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Denise Lanzieri</b></sub></a><br /><a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=dlanzieri" title="Code">💻</a></td>
</tr>
<tr>
<td align="center"><a href="https://github.com/jecampagne"><img src="https://avatars.githubusercontent.com/u/20539759?v=4?s=100" width="100px;" alt=""/><br /><sub><b>jecampagne</b></sub></a><br /><a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues?q=author%3Ajecampagne" title="Bug reports">🐛</a></td>
<td align="center"><a href="https://github.com/eelregit"><img src="https://avatars.githubusercontent.com/u/7311098?v=4?s=100" width="100px;" alt=""/><br /><sub><b>Yin Li</b></sub></a><br /><a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=eelregit" title="Code">💻</a> <a href="https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues?q=author%3Aeelregit" title="Bug reports">🐛</a></td>
</tr>
</table>
<!-- markdownlint-restore -->
<!-- prettier-ignore-end -->
<!-- ALL-CONTRIBUTORS-LIST:END -->
This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
Raw data
{
"_id": null,
"home_page": "https://github.com/DifferentiableUniverseInitiative/jax_cosmo",
"name": "jax-cosmo",
"maintainer": "",
"docs_url": null,
"requires_python": "",
"maintainer_email": "",
"keywords": "",
"author": "jax-cosmo developers",
"author_email": "",
"download_url": "https://files.pythonhosted.org/packages/b1/ff/3d1c8d348ce2eb07cc990142cbd4dd2e72070dfd30e2dbe3660f1b392cd2/jax-cosmo-0.1.0.tar.gz",
"platform": null,
"description": "# jax-cosmo\n\n[](https://gitter.im/DifferentiableUniverseInitiative/jax_cosmo?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge) [](https://jax-cosmo.readthedocs.io/en/latest/?badge=latest) []() [](https://github.com/psf/black) [](https://pypi.org/project/jax-cosmo/) [](CODE_OF_CONDUCT.md) [](https://github.com/google/jax-cosmo/blob/master/LICENSE) <!-- ALL-CONTRIBUTORS-BADGE:START - Do not remove or modify this section -->\n[](#contributors-)\n\n<!-- ALL-CONTRIBUTORS-BADGE:END -->\n\n<h3 align=\"center\">Finally a differentiable cosmology library, and it's in JAX!</h3>\n\nHave a look at the [GitHub issues](https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues) to see what is needed or if you have any thoughts on the design, and don't hesitate to join the [Gitter room](https://gitter.im/DifferentiableUniverseInitiative/jax_cosmo) for discussions.\n\n## TL;DR\n\nThis is what `jax-cosmo` aims to do:\n\n```python\n...\ndef likelihood(cosmo):\n # Compute mean and covariance of angular Cls, for specific probes\n mu, cov = jax_cosmo.angular_cl.gaussian_cl_covariance_and_mean(cosmo, ell, probes)\n # Return likelihood value\n return jax_cosmo.likelihood.gaussian_log_likelihood(data, mu, cov)\n\n# Compute derivatives of the likelihood with respect to cosmological parameters\ng = jax.grad(likelihood)(cosmo)\n\n# Compute Fisher matrix of cosmological parameters\nF = - jax.hessian(likelihood)(cosmo)\n```\n\nThis is how you can compute **gradients and hessians of any functions in `jax-cosmo`**,\nall of this without any finite differences.\n\nCheck out a full example here: [](https://colab.research.google.com/github/DifferentiableUniverseInitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb)\n\nHave a look at the [design document](design.md) to learn more about the structure of the code.\n\n## What is JAX?\n\n[JAX](https://github.com/google/jax) = NumPy + autodiff + GPU\n\nJAX is a framework for automatic differentiation (like TensorFlow or PyTorch) but following the NumPy API, and using the GPU/TPU enable XLA backend.\n\nWhat does that mean?\n\n- You write plain Python/NumPy code, no need to learn a different language\n- It runs on GPU, you don't need to do anything particular\n- You can take derivatives of any quantity with respect to any parameters by\n automatic differentiation.\n\nCheckout the [JAX](https://github.com/google/jax) project page to learn more!\n\n## Install\n\n`jax-cosmo` is pure Python, so installing is a breeze:\n\n```bash\n$ pip install jax-cosmo\n```\n\n## Philosophy\n\nHere are some of the design guidelines:\n\n- Implementation of equations should be human readable, and documentation should always live next to the implementation.\n- Should always be trivially installable: external dependencies should be kept\n to a minimum, especially the ones that require compilation or with restrictive licenses.\n- Keep API and implementation simple and intuitive, minimize user and developer\n surprise.\n- \u201cDebugging is twice as hard as writing the code in the first place. Therefore, if you write the code as cleverly as possible, you are, by definition, not smart enough to debug it.\u201d -Brian Kernighan, quote stolen from\n [here](https://flax.readthedocs.io/en/latest/philosophy.html).\n\n## Contributing\n\n`jax-cosmo` aims to be a community effort, contributions are most welcome and\ncan come in several forms\n\n- Bug reports\n- API design suggestions\n- (Pull) requests for more features\n- Examples and notebooks of cool things that can be done with the code\n\nYou can chime-in on any aspects of the design by proposing a PR to the [design document](design.md). The issue page is a good place to start, but don't hesitate to come chat in the Gitter room.\n\nPlease take a look at the [Contributing Document](CONTRIBUTING.md) for more information.\n\nThis project follows the [All Contributors](https://allcontributors.org/) guidelines aiming at recognizing and valorizing\ncontributions at any levels.\n\n## Contributors \u2728\n\nThanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):\n\n<!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->\n<!-- prettier-ignore-start -->\n<!-- markdownlint-disable -->\n<table>\n <tr>\n <td align=\"center\"><a href=\"http://flanusse.net\"><img src=\"https://avatars0.githubusercontent.com/u/861591?v=4?s=100\" width=\"100px;\" alt=\"\"/><br /><sub><b>Francois Lanusse</b></sub></a><br /><a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=EiffL\" title=\"Code\">\ud83d\udcbb</a></td>\n <td align=\"center\"><a href=\"http://www.cosmostat.org/people/santiago-casas\"><img src=\"https://avatars0.githubusercontent.com/u/6987716?v=4?s=100\" width=\"100px;\" alt=\"\"/><br /><sub><b>Santiago Casas</b></sub></a><br /><a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues?q=author%3Asantiagocasas\" title=\"Bug reports\">\ud83d\udc1b</a> <a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=santiagocasas\" title=\"Code\">\ud83d\udcbb</a></td>\n <td align=\"center\"><a href=\"https://github.com/austinpeel\"><img src=\"https://avatars0.githubusercontent.com/u/17024310?v=4?s=100\" width=\"100px;\" alt=\"\"/><br /><sub><b>Austin Peel</b></sub></a><br /><a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=austinpeel\" title=\"Code\">\ud83d\udcbb</a></td>\n <td align=\"center\"><a href=\"https://minaskaramanis.com\"><img src=\"https://avatars2.githubusercontent.com/u/23280751?v=4?s=100\" width=\"100px;\" alt=\"\"/><br /><sub><b>Minas Karamanis</b></sub></a><br /><a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=minaskar\" title=\"Code\">\ud83d\udcbb</a></td>\n <td align=\"center\"><a href=\"https://faculty.sites.uci.edu/dkirkby/\"><img src=\"https://avatars1.githubusercontent.com/u/185007?v=4?s=100\" width=\"100px;\" alt=\"\"/><br /><sub><b>David Kirkby</b></sub></a><br /><a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=dkirkby\" title=\"Code\">\ud83d\udcbb</a> <a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues?q=author%3Adkirkby\" title=\"Bug reports\">\ud83d\udc1b</a></td>\n <td align=\"center\"><a href=\"https://aboucaud.github.io\"><img src=\"https://avatars0.githubusercontent.com/u/3065310?v=4?s=100\" width=\"100px;\" alt=\"\"/><br /><sub><b>Alexandre Boucaud</b></sub></a><br /><a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=aboucaud\" title=\"Code\">\ud83d\udcbb</a></td>\n <td align=\"center\"><a href=\"https://www.cosmostat.org/people/denise-lanzieri\"><img src=\"https://avatars.githubusercontent.com/u/72620117?v=4?s=100\" width=\"100px;\" alt=\"\"/><br /><sub><b>Denise Lanzieri</b></sub></a><br /><a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=dlanzieri\" title=\"Code\">\ud83d\udcbb</a></td>\n </tr>\n <tr>\n <td align=\"center\"><a href=\"https://github.com/jecampagne\"><img src=\"https://avatars.githubusercontent.com/u/20539759?v=4?s=100\" width=\"100px;\" alt=\"\"/><br /><sub><b>jecampagne</b></sub></a><br /><a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues?q=author%3Ajecampagne\" title=\"Bug reports\">\ud83d\udc1b</a></td>\n <td align=\"center\"><a href=\"https://github.com/eelregit\"><img src=\"https://avatars.githubusercontent.com/u/7311098?v=4?s=100\" width=\"100px;\" alt=\"\"/><br /><sub><b>Yin Li</b></sub></a><br /><a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/commits?author=eelregit\" title=\"Code\">\ud83d\udcbb</a> <a href=\"https://github.com/DifferentiableUniverseInitiative/jax_cosmo/issues?q=author%3Aeelregit\" title=\"Bug reports\">\ud83d\udc1b</a></td>\n </tr>\n</table>\n\n<!-- markdownlint-restore -->\n<!-- prettier-ignore-end -->\n\n<!-- ALL-CONTRIBUTORS-LIST:END -->\n\nThis project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!\n",
"bugtrack_url": null,
"license": "",
"summary": "Differentiable Python Cosmology Library",
"version": "0.1.0",
"split_keywords": [],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "b1ff3d1c8d348ce2eb07cc990142cbd4dd2e72070dfd30e2dbe3660f1b392cd2",
"md5": "fd3bc8164c69fb53cee3c8a9991672b4",
"sha256": "a44003ff38381347059327b88e3fd820499150195acbf341b463db52bb461372"
},
"downloads": -1,
"filename": "jax-cosmo-0.1.0.tar.gz",
"has_sig": false,
"md5_digest": "fd3bc8164c69fb53cee3c8a9991672b4",
"packagetype": "sdist",
"python_version": "source",
"requires_python": null,
"size": 287446,
"upload_time": "2023-02-05T14:04:06",
"upload_time_iso_8601": "2023-02-05T14:04:06.937604Z",
"url": "https://files.pythonhosted.org/packages/b1/ff/3d1c8d348ce2eb07cc990142cbd4dd2e72070dfd30e2dbe3660f1b392cd2/jax-cosmo-0.1.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2023-02-05 14:04:06",
"github": true,
"gitlab": false,
"bitbucket": false,
"github_user": "DifferentiableUniverseInitiative",
"github_project": "jax_cosmo",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "jax-cosmo"
}
JSON
