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 or analytic expressions using hypergeometric functions.
- Support for
jax
andcupy
backends allow hardware acceleration, just-in-time compilation, and automatic differentiation.
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
astropy
units are incompatible with thecupy
backend. Units are supported with thejax
backend usingunxt
.
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 [email protected]: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
[email protected]:ColmTalbot/wcosmo.git
.
To import an astropy-like cosmology
>>> from wcosmo import FlatwCDM
>>> cosmology = FlatwCDM(H0=70, Om0=0.3, w0=-1)
>>> cosmology.luminosity_distance(1)
The built-in cosmologies in astropy
are all available, e.g.,
>>> from wcosmo import Planck18
>>> Planck18.luminosity_distance(1)
<Quantity 6797.43628659 Mpc>
they can also be accessed using wcosmo.available
>>> from wcosmo import available
>>> available["Planck18"].luminosity_distance(1)
<Quantity 6797.43628659 Mpc>
Explicit usage of astropy
units can be freely enabled/disabled.
In this case, the values will have the default units for each method.
>>> from wcosmo import FlatwCDM
>>> from wcosmo.utils import disable_units, enable_units
>>> cosmology = FlatwCDM(H0=70, Om0=0.3, w0=-1)
>>> disable_units()
>>> cosmology.luminosity_distance(1)
6607.657732077576
>>> enable_units()
>>> cosmology.luminosity_distance(1)
<Quantity 6607.65773208 Mpc>
wcosmo
mostly relies on implicit backend switching. The backend is
determined automatically based on the input arguments. When an input value is
a Python
built-in type, the default backend is chosen using the
environment variable WCOSMO_ARRAY_API
. The default is numpy
.
The original intention for this package was for use with GWPopulation
.
This code is automatically used in GWPopulation
when using either
gwpopulation.experimental.cosmo_models.CosmoModel
and/or
PowerLawRedshift