A package for cosmology calculations with arbitrary numpy-like APIs
Project description
An efficient implementation of astropy
-like cosmology compatible
with numpy
-like backends, e.g., jax
and cupy
.
There are two main features leading to superior efficiency to astropy
:
Integrals of \(E(z)\) and related functions are performed analytically with Pade approximations.
Support for
jax
andcupy
backends allow hardware acceleration, just-in-time compilation, and automatic differentiation.We don’t use explicit units, which can be a source of inefficiency.
The primary limitations are:
Only flat cosmologies are supported with two components with constant equations of state, e.g.,
FlatwCDM
.Approximations to the various integrals generally agree with
astropy
at the 0.1% level.The units we use are documented, but users need to be careful.
Installation and contribution
wcosmo
can be installed via conda-forge
, pypi
or from
source.
$ mamba install -c conda-forge wcosmo
$ pip install wcosmo
$ pip install git+https://github.com/ColmTalbot/wcosmo.git
for development you should follow a standard fork-and-pull workflow.
First create a new fork at
github.com/UserName/wcosmo
.Clone your fork
$ git clone git@github.com:UserName/wcosmo.git
or use a GitHub codespace.
Install the local version with
$ python -m pip install .
Make any desired edits and push to your fork.
Open a pull request into
git@github.com:ColmTalbot/wcosmo.git
.
Basic usage
To import an astropy-like cosmology (without units)
from wcosmo import FlatwCDM
cosmology = FlatwCDM(H0=70, Om0=0.3, w0=-1)
GWPopulation
The primary intention for this package is for use with GWPopulation
.
This code is automatically used in GWPopulation
when using either
gwpopulation.experimental.cosmo_models.CosmoModel
and/or
PowerLawRedshift
Changing backend
The backend can be switched automatically using, e.g.,
import gwpopulation
gwpopulation.backend.set_backend("jax")
Manual backend setting can be done as follows:
import jax.numpy as jnp
from jax.scipy.linalg.toeplitz import toeplitz
from wcosmo import wcosmo, utils
wcosmo.xp = jnp
utils.xp = jnp
utils.toeplitz = toeplitz
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.