Cosmology For Developers

Cosmologies in Functions

It is often useful to assume a default cosmology so that the exact cosmology does not have to be specified every time a function or method is called. In this case, it is possible to specify a “default” cosmology.

You can set the default cosmology to a predefined value by using the “default_cosmology” option in the [cosmology.core] section of the configuration file (see Configuration System (astropy.config)). Alternatively, you can use the set() function of default_cosmology() to set a cosmology for the current Python session. If you have not set a default cosmology using one of the methods described above, then the cosmology module will default to using the default_cosmology._value parameters.

It is strongly recommended that you use the default cosmology through the default_cosmology science state object. An override option can then be provided using something like the following:

from astropy.cosmology import default_cosmology

def myfunc(..., cosmo=None):
    if cosmo is None:
        cosmo = default_cosmology.get()

    ... function code here ...

This ensures that all code consistently uses the default cosmology unless explicitly overridden.

Note

In general it is better to use an explicit cosmology (for example WMAP9.H(0) instead of cosmology.default_cosmology.get().H(0)). Use of the default cosmology should generally be reserved for code that will be included in the astropy core or an affiliated package.

Custom Cosmologies

In astropy.cosmology cosmologies are classes, so custom cosmologies may be implemented by subclassing Cosmology (or more likely FLRW) and adding details specific to that cosmology. Here we will review some of those details and tips and tricks to building a performant cosmology class.

from astropy.cosmology import FLRW

class CustomCosmology(FLRW):
    ...  # [details discussed below]

Parameters

An astropy Cosmology is characterized by 1) its class, which encodes the physics, and 2) its free parameter(s), which specify a cosmological realization. When defining the former, all parameters must be declared using Parameter and should have values assigned at instantiation.

A Parameter is a descriptor. When accessed from a class it transparently stores information, like the units and accepted equivalencies, that might be opaquely contained in the constructor signature or more deeply in the code. On a cosmology instance, the descriptor will return the parameter value.

There are a number of best practices. For a reference, this is excerpted from the definition of FLRW.

class FLRW(Cosmology):

    H0 = Parameter(doc="Hubble constant as an `~astropy.units.Quantity` at z=0",
                   unit="km/(s Mpc)", fvalidate="scalar")
    Om0 = Parameter(doc="Omega matter; matter density/critical density at z=0",
                    fvalidate="non-negative")
    Ode0 = Parameter(doc="Omega dark energy; dark energy density/critical density at z=0.",
                     fvalidate="float")
    Tcmb0 = Parameter(doc="Temperature of the CMB as `~astropy.units.Quantity` at z=0.",
              unit="Kelvin", fmt="0.4g", fvalidate="scalar")
    Neff = Parameter(doc="Number of effective neutrino species.", fvalidate="non-negative")
    m_nu = Parameter(doc="Mass of neutrino species.",
             unit="eV", equivalencies=u.mass_energy(), fmt="")
    Ob0 = Parameter(doc="Omega baryon; baryonic matter density/critical density at z=0.")

    def __init__(self, H0, Om0, Ode0, Tcmb0=0.0*u.K, Neff=3.04, m_nu=0.0*u.eV,
                 Ob0=None, *, name=None, meta=None):
        self.H0 = H0
        ...  # for each Parameter in turn

    @Ob0.validator
    def Ob0(self, param, value):
        """Validate baryon density to None or positive float > matter density."""
        if value is None:
            return value
        value = _validate_non_negative(self, param, value)
        if value > self.Om0:
            raise ValueError("baryonic density can not be larger than total matter density.")
        return value

First note that all the parameters are also arguments in __init__. This is not strictly necessary, but is good practice. If the parameter has units (and related equivalencies) these must be specified on the Parameter, as seen in H0 and m_nu.

The next important thing to note is how the parameter value is set, in __init__. Parameter allows for a value to be set once (before auto-locking), so self.H0 = H0 will use this setter and put the value on “._H0”. The advantage of this method over direct assignment to the private attribute is the use of validators. Parameter allows for custom value validators, using the method-decorator validator, that can check a value’s validity and modify the value, e.g to assign units. If no custom validator is specified the default is to check if the Parameter has defined units and if so, return the value as a Quantity with those units, using all enabled and the parameter’s unit equivalencies.

The last thing to note is pretty formatting for the Cosmology. Each Parameter defaults to the format specification “.3g”, but this may be overridden, like Tcmb0 does.

If a new cosmology modifies an existing Parameter, then the Parameter.clone() method is useful to deep-copy the parameter and change any constructor argument. For example, see FlatFLRWMixin in astropy.cosmology.flrw (also shown below).

class FlatFLRWMixin(FlatCosmologyMixin):
    ...

    Ode0 = FLRW.Ode0.clone(derived=True)  # now a derived param.

Mixins

Mixins are used in cosmology to reuse code across multiple classes in different inheritance lines. We use the term loosely as mixins are meant to be strictly orthogonal, but may not be, particularly in __init__.

Currently the only mixin is FlatCosmologyMixin and its FLRW-specific subclass FlatFLRWMixin. “Flat” cosmologies should use this mixin. FlatFLRWMixin must precede the base class in the multiple-inheritance so that this mixin’s __init__ proceeds the base class’.

Speeding up Integrals in Custom Cosmologies

The supplied cosmology classes use a few tricks to speed up distance and time integrals. It is not necessary for anyone subclassing FLRW to use these tricks – but if they do, such calculations may be a lot faster.

The first, more basic, idea is that, in many cases, it’s a big deal to provide explicit formulae for inv_efunc() rather than simply setting up de_energy_scale – assuming there is a nice expression. As noted above, almost all of the provided classes do this, and that template can pretty much be followed directly with the appropriate formula changes.

The second, and more advanced, option is to also explicitly provide a scalar only version of inv_efunc(). This results in a fairly large speedup (>10x in most cases) in the distance and age integrals, even if only done in python, because testing whether the inputs are iterable or pure scalars turns out to be rather expensive. To take advantage of this, the key thing is to explicitly set the instance variables self._inv_efunc_scalar and self._inv_efunc_scalar_args in the constructor for the subclass, where the latter are all the arguments except z to _inv_efunc_scalar. The provided classes do use this optimization, and in fact go even further and provide optimizations for no radiation, and for radiation with massless neutrinos coded in cython. Consult the FLRW subclasses for details, and scalar_inv_efuncs for the details.

However, the important point is that it is not necessary to do this.

Astropy Interoperability: I/O and your Cosmology Package

If you are developing a package and want to be able to interoperate with astropy.cosmology.Cosmology, you’re in the right place! Here we will discuss and provide examples for enabling Astropy to read and write your file formats, but also convert your cosmology objects to and from Astropy’s Cosmology.

The following presumes knowledge of how Astropy structures I/O functions. For a quick tutorial see Read, Write, and Convert Cosmology Objects.

Now that we know how to build and register functions into read(), write(), from_format(), to_format(), we can do this in your package.

Consider a package – since this is mine, it’s cleverly named mypackage – with the following file structure: a module for cosmology codes and a module for defining related input/output functions. In the cosmology module are defined cosmology classes and a file format – myformat – and everything should interoperate with astropy. The tests are done with pytest and are integrated within the code structure.

mypackage/
    __init__.py
    cosmology/
        __init__.py
        ...
    io/
        __init__.py
        astropy_convert.py
        astropy_io.py
        ...
        tests/
            __init__.py
            test_astropy_convert.py
            test_astropy_io.py
            ...

For a fully implemented example mypackage, see https://github.com/astropy/astropy/tree/main/astropy/cosmology/tests/mypackage

Converting Objects Between Packages

We want to enable conversion between cosmology objects from mypckage to/from Cosmology. All the Astropy interface code is defined in mypackage/io/astropy_convert.py. The following is a rough outline of the necessary functions and how to register them with astropy’s unified I/O to be automatically available to astropy.cosmology.Cosmology.from_format() and astropy.cosmology.Cosmology.to_format().

# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see Astropy LICENSE.rst

"""
Register conversion methods for cosmology objects with Astropy Cosmology.

With this registered, we can start with a Cosmology from
``mypackage`` and convert it to an astropy Cosmology instance.

    >>> from mypackage.cosmology import myplanck
    >>> from astropy.cosmology import Cosmology
    >>> cosmo = Cosmology.from_format(myplanck, format="mypackage")
    >>> cosmo

We can also do the reverse: start with an astropy Cosmology and convert it
to a ``mypackage`` object.

    >>> from astropy.cosmology import Planck18
    >>> myplanck = Planck18.to_format("mypackage")
    >>> myplanck

"""

# THIRD PARTY
import astropy.cosmology.units as cu
import astropy.units as u
from astropy.cosmology import FLRW, Cosmology, FlatLambdaCDM
from astropy.cosmology.connect import convert_registry

# LOCAL
from mypackage.cosmology import MyCosmology

__doctest_skip__ = ['*']


def from_mypackage(mycosmo):
    """Load `~astropy.cosmology.Cosmology` from ``mypackage`` object.

    Parameters
    ----------
    mycosmo : `~mypackage.cosmology.MyCosmology`

    Returns
    -------
    `~astropy.cosmology.Cosmology`
    """
    m = dict(mycosmo)
    m["name"] = mycosmo.name

    # ----------------
    # remap Parameters
    m["H0"] = m.pop("hubble_parameter") * (u.km / u.s / u.Mpc)
    m["Om0"] = m.pop("initial_matter_density")
    m["Tcmb0"] = m.pop("initial_temperature") * u.K
    # m["Neff"] = m.pop("Neff")  # skip b/c unchanged
    m["m_nu"] = m.pop("neutrino_masses") * u.eV
    m["Ob0"] = m.pop("initial_baryon_density")

    # ----------------
    # remap metadata
    m["t0"] = m.pop("current_age") * u.Gyr

    # optional
    if "reionization_redshift" in m:
        m["z_reion"] = m.pop("reionization_redshift")

    # ...  # keep building `m`

    # ----------------
    # Detect which type of Astropy cosmology to build.
    # TODO! CUSTOMIZE FOR DETECTION
    # Here we just force FlatLambdaCDM, but if your package allows for
    # non-flat cosmologies...
    m["cosmology"] = FlatLambdaCDM

    # build cosmology
    return Cosmology.from_format(m, format="mapping", move_to_meta=True)


def to_mypackage(cosmology, *args):
    """Return the cosmology as a ``mycosmo``.

    Parameters
    ----------
    cosmology : `~astropy.cosmology.Cosmology`

    Returns
    -------
    `~mypackage.cosmology.MyCosmology`
    """
    if not isinstance(cosmology, FLRW):
        raise TypeError("format 'mypackage' only supports FLRW cosmologies.")

    # ----------------
    # Cosmology provides a nice method "mapping", so all that needs to
    # be done here is initialize from the dictionary
    m = cosmology.to_format("mapping")

    # Detect which type of MyCosmology to build.
    # Here we have forced FlatLambdaCDM, but if your package allows for
    # non-flat cosmologies...
    m.pop("cosmology")

    # MyCosmology doesn't support metadata. If your cosmology class does...
    meta = m.pop("meta")
    m = {**meta, **m}  # merge, preferring current values

    # ----------------
    # remap values
    # MyCosmology doesn't support units, so take values.
    m["hubble_parameter"] = m.pop("H0").to_value(u.km/u.s/u.Mpc)
    m["initial_matter_density"] = m.pop("Om0")
    m["initial_temperature"] = m.pop("Tcmb0").to_value(u.K)
    # m["Neff"] = m.pop("Neff")  # skip b/c unchanged
    m["neutrino_masses"] = m.pop("m_nu").to_value(u.eV)
    m["initial_baryon_density"] = m.pop("Ob0")
    m["current_age"] = m.pop("t0", cosmology.age(0 * cu.redshift)).to_value(u.Gyr)

    # optional
    if "z_reion" in m:
        m["reionization_redshift"] = (m.pop("z_reion") << cu.redshift).value

    # ...  # keep remapping

    return MyCosmology(**m)


def mypackage_identify(origin, format, *args, **kwargs):
    """Identify if object uses format "mypackage"."""
    itis = False
    if origin == "read":
        itis = isinstance(args[1], MyCosmology) and (format in (None, "mypackage"))
    return itis


# -------------------------------------------------------------------
# Register to/from_format & identify methods with Astropy Unified I/O

convert_registry.register_reader("mypackage", Cosmology, from_mypackage, force=True)
convert_registry.register_writer("mypackage", Cosmology, to_mypackage, force=True)
convert_registry.register_identifier("mypackage", Cosmology, mypackage_identify, force=True)

Reading and Writing

Everything Astropy read/write related is defined in mypackage/io/astropy_io.py. The following is a rough outline of the read, write, and identify functions and how to register them with astropy’s unified io to be automatically available to astropy.cosmology.Cosmology.read() and astropy.cosmology.Cosmology.write().

# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see Astropy LICENSE.rst

"""
Register Read/Write methods for "myformat" (JSON) with Astropy Cosmology.

With this format registered, we can start with a Cosmology from
``mypackage``, write it to a file, and read it with Astropy to create an
astropy Cosmology instance.

    >>> from mypackage.cosmology import myplanck
    >>> from mypackage.io import file_writer
    >>> file_writer('<file name>', myplanck)

    >>> from astropy.cosmology import Cosmology
    >>> cosmo = Cosmology.read('<file name>', format="myformat")
    >>> cosmo

We can also do the reverse: start with an astropy Cosmology, save it and
read it with ``mypackage``.

    >>> from astropy.cosmology import Planck18
    >>> Planck18.write('<file name>', format="myformat")

    >>> from mypackage.io import file_reader
    >>> cosmo2 = file_reader('<file name>')
    >>> cosmo2

"""

# STDLIB
import json
import os

# THIRD PARTY
import astropy.units as u
from astropy.cosmology import Cosmology
from astropy.cosmology.connect import readwrite_registry

# LOCAL
from .core import file_reader, file_writer

__doctest_skip__ = ['*']


def read_myformat(filename, **kwargs):
    """Read files in format 'myformat'.

    Parameters
    ----------
    filename : str
    **kwargs
        Keyword arguments into `astropy.cosmology.Cosmology.from_format`
        with ``format="mypackage"``.

    Returns
    -------
    `~mypackage.cosmology.MyCosmology` instance
    """
    mycosmo = file_reader(filename)  # ← read file  ↓ build Cosmology
    return Cosmology.from_format(mycosmo, format="mypackage", **kwargs)


def write_myformat(cosmology, file, *, overwrite=False):
    """Write files in format 'myformat'.

    Parameters
    ----------
    cosmology : `~astropy.cosmology.Cosmology` instance
    file : str, bytes, or `~os.PathLike`
    overwrite : bool (optional, keyword-only)
        Whether to overwrite an existing file. Default is False.
    """
    cosmo = cosmology.to_format("mypackage")  # ← convert Cosmology ↓ write file
    file_writer(file, cosmo, overwrite=overwrite)


def myformat_identify(origin, filepath, fileobj, *args, **kwargs):
    """Identify if object uses ``myformat`` (JSON)."""
    return filepath is not None and filepath.endswith(".myformat")


# -------------------------------------------------------------------
# Register read/write/identify methods with Astropy Unified I/O

readwrite_registry.register_reader("myformat", Cosmology, read_myformat, force=True)
readwrite_registry.register_writer("myformat", Cosmology, write_myformat, force=True)
readwrite_registry.register_identifier("myformat", Cosmology, myformat_identify, force=True)

If Astropy is an optional dependency

The astropy_io and astropy_convert modules are written assuming Astropy is installed. If in mypackage it is an optional dependency then it is important to detect if Astropy is installed (and the correct version) before importing astropy_io and astropy_convert. We do this in mypackage/io/__init__.py:

# -*- coding: utf-8 -*-

"""
Readers, Writers, and I/O Miscellany.
"""

__all__ = ["file_reader", "file_writer"]  # from `mypackage`

# e.g. a file reader and writer for ``myformat``
# this will be used in ``astropy_io.py``
from .core import file_reader, file_writer

# Register read and write methods into Astropy:
# determine if it is 1) installed and 2) the correct version (v5.0+)
try:
    import astropy
    from astropy.utils.introspection import minversion
except ImportError:
    ASTROPY_GE_5 = False
else:
    ASTROPY_GE_5 = minversion(astropy, "5.0")

if ASTROPY_GE_5:
    # Astropy is installed and v5.0+ so we import the following modules
    # to register "myformat" with Cosmology read/write and "mypackage"
    # with Cosmology to/from_format.
    from . import astropy_convert, astropy_io

Astropy Interoperability Tests

Lastly, it’s important to test that everything works. In this example package all such tests are contained in mypackage/io/tests/test_astropy_io.py. These tests require Astropy and will be skipped if it is not installed (and not the correct version), so at least one test in the test matrix should include astropy >= 5.0.

# -*- coding: utf-8 -*-

"""
Test that the conversion interface with Astropy works as expected.
"""

# STDLIB
import os

# THIRD PARTY
import pytest

# LOCAL
from mypackage.cosmology import MyCosmology, myplanck

# skip all tests in module if import is missing
astropy = pytest.importorskip("astropy", minversion="4.3")  # isort: skip
# can now import freely from astropy
from astropy import cosmology
from astropy.cosmology import Cosmology
from astropy.io import registry as io_registry
from astropy.utils.compat.optional_deps import HAS_SCIPY

# cosmology instances to test reading and writing
astropy_cosmos = [
    getattr(cosmology.realizations, n) for n in cosmology.parameters.available
]

mypackage_cosmos = [myplanck]  # list of ``mypackage`` cosmology realizations


@pytest.mark.skipif(not HAS_SCIPY, reason="test requires scipy.")
class TestAstropyCosmologyConvert:
    """Test conversion to/from Astropy."""

    @pytest.mark.parametrize("expected", astropy_cosmos)
    def test_roundtrip_from_astropy(self, expected):
        # convert to ``mypackage``
        mycosmo = expected.to_format("mypackage")

        # Read back
        got = Cosmology.from_format(mycosmo, format="mypackage")

        # test round-tripped as expected
        assert got == expected  # tests immutable parameters, e.g. H0

        # NOTE: if your package's cosmology supports metadata
        # assert got.meta == expected.meta  # (metadata not tested above)

    @pytest.mark.parametrize("expected", mypackage_cosmos)
    def test_roundtrip_from_mypackage(self, expected):
        # convert to Astropy
        acosmo = Cosmology.from_format(expected, format="mypackage")

        # convert back to ``mypackage```
        got = acosmo.to_format("mypackage")

        # test round-tripped as expected
        assert isinstance(got, MyCosmology)
        assert got == expected  # assuming ``MyCosmology`` has an __eq__ method

        ...  # more equality tests

    @pytest.mark.parametrize("acosmo", astropy_cosmos)
    def test_package_equality(self, acosmo):
        """
        The most important test: are the ``mypackage`` and ``astropy``
        cosmologies equivalent!? They should be if read from the same source
        file.
        """
        # convert to ``mypacakge``
        mycosmo = acosmo.to_format("mypackage")

        # ------------
        # test that the mypackage and astropy cosmologies are equivalent

        assert mycosmo.name == acosmo.name
        assert mycosmo.hubble_parameter == acosmo.H0.value

        ...  # continue with the tests

        # e.g. test some methods
        # assert mycosmo.age_in_Gyr(1100) == acosmo.age(1100).to_value(u.Gyr)
# -*- coding: utf-8 -*-

"""
Test that the interface with Astropy works as expected.
"""

# STDLIB
import os

# THIRD PARTY
import pytest

# LOCAL
from mypackage.cosmology import myplanck
from mypackage.io import file_reader

# skip all tests in module if import is missing
astropy = pytest.importorskip("astropy", minversion="4.3")  # isort: skip
# can now import freely from astropy
import astropy.units as u
from astropy import cosmology
from astropy.cosmology import Cosmology
from astropy.io import registry as io_registry
from astropy.utils.compat.optional_deps import HAS_SCIPY

# the formats being registered with Astropy
readwrite_formats = ["myformat"]
# cosmology instances to test reading and writing
astropy_cosmos = cosmology.parameters.available

# list of ``mypackage`` cosmology realizations
mypackage_cosmos = [myplanck]


@pytest.mark.skipif(not HAS_SCIPY, reason="test requires scipy.")
class TestAstropyCosmologyIO:
    """Test read/write interoperability with Astropy."""

    @pytest.mark.parametrize("format", readwrite_formats)
    @pytest.mark.parametrize("instance", astropy_cosmos)
    def test_roundtrip_from_astropy(self, tmp_path, instance, format):
        cosmo = getattr(cosmology.realizations, instance)
        fname = tmp_path / f"{instance}.{format}"

        # write to file
        cosmo.write(str(fname), format=format)

        # also test kwarg "overwrite"
        assert os.path.exists(str(fname))  # file exists
        with pytest.raises(IOError):
            cosmo.write(str(fname), format=format, overwrite=False)

        assert os.path.exists(str(fname))  # overwrite file existing file
        cosmo.write(str(fname), format=format, overwrite=True)

        # Read back
        got = Cosmology.read(fname, format=format)

        # test round-tripped as expected
        assert got == cosmo  # tests immutable parameters, e.g. H0

        # NOTE: if your package's cosmology supports metadata
        # assert got.meta == expected.meta  # (metadata not tested above)

    @pytest.mark.parametrize("format", readwrite_formats)
    @pytest.mark.parametrize("instance", astropy_cosmos)
    def test_package_equality(self, tmp_path, instance, format):
        """
        The most important test: are the ``mypackage`` and ``astropy``
        cosmologies equivalent!? They should be if read from the same source
        file.
        """
        original = getattr(cosmology.realizations, instance)
        fname = tmp_path / f"{instance}.{format}"

        # write with Astropy
        original.write(str(fname), format=format)

        # Read back with ``myformat``
        cosmo = file_reader(str(fname))  # read to instance from mypackage

        # and a read comparison from Astropy
        cosmo2 = cosmology.Cosmology.read(str(fname), format=format)

        # ------------
        # test that the mypackage and astropy cosmologies are equivalent
        assert original.H0.value == cosmo.hubble_parameter
        assert cosmo2.H0.value == cosmo.hubble_parameter

        ...  # continue with the tests

        # e.g. test some methods
        # assert original.age(1100).to_value(u.Gyr) == cosmo.age_in_Gyr(1100)
        # assert cosmo2.age(1100) == cosmo.age(1100)