Backend agnostic astropy-like cosmology#

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 and cupy 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

Examples#

Examples
- Spectral Siren Population Inference on GWTC-3
- Compare wcosmo and astropy timing.