Skip to content

ColmTalbot/wcosmo

Repository files navigation

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 or analytic expressions using hypergeometric functions.
  • Support for jax and cupy 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 the cupy backend. Units are supported with the jax backend using unxt.

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 [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.

Basic usage

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>

Changing backend

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.

GWPopulation

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