1# -*- coding: utf-8 -*-
2# Licensed under a 3-clause BSD style license - see Astropy LICENSE.rst
3
4"""
5Register conversion methods for cosmology objects with Astropy Cosmology.
6
7With this registered, we can start with a Cosmology from
8``mypackage`` and convert it to an astropy Cosmology instance.
9
10    >>> from mypackage.cosmology import myplanck
11    >>> from astropy.cosmology import Cosmology
12    >>> cosmo = Cosmology.from_format(myplanck, format="mypackage")
13    >>> cosmo
14
15We can also do the reverse: start with an astropy Cosmology and convert it
16to a ``mypackage`` object.
17
18    >>> from astropy.cosmology import Planck18
19    >>> myplanck = Planck18.to_format("mypackage")
20    >>> myplanck
21
22"""
23
24# THIRD PARTY
25import astropy.cosmology.units as cu
26import astropy.units as u
27from astropy.cosmology import FLRW, Cosmology, FlatLambdaCDM
28from astropy.cosmology.connect import convert_registry
29
30# LOCAL
31from mypackage.cosmology import MyCosmology
32
33__doctest_skip__ = ['*']
34
35
36def from_mypackage(mycosmo):
37    """Load `~astropy.cosmology.Cosmology` from ``mypackage`` object.
38
39    Parameters
40    ----------
41    mycosmo : `~mypackage.cosmology.MyCosmology`
42
43    Returns
44    -------
45    `~astropy.cosmology.Cosmology`
46    """
47    m = dict(mycosmo)
48    m["name"] = mycosmo.name
49
50    # ----------------
51    # remap Parameters
52    m["H0"] = m.pop("hubble_parameter") * (u.km / u.s / u.Mpc)
53    m["Om0"] = m.pop("initial_matter_density")
54    m["Tcmb0"] = m.pop("initial_temperature") * u.K
55    # m["Neff"] = m.pop("Neff")  # skip b/c unchanged
56    m["m_nu"] = m.pop("neutrino_masses") * u.eV
57    m["Ob0"] = m.pop("initial_baryon_density")
58
59    # ----------------
60    # remap metadata
61    m["t0"] = m.pop("current_age") * u.Gyr
62
63    # optional
64    if "reionization_redshift" in m:
65        m["z_reion"] = m.pop("reionization_redshift")
66
67    # ...  # keep building `m`
68
69    # ----------------
70    # Detect which type of Astropy cosmology to build.
71    # TODO! CUSTOMIZE FOR DETECTION
72    # Here we just force FlatLambdaCDM, but if your package allows for
73    # non-flat cosmologies...
74    m["cosmology"] = FlatLambdaCDM
75
76    # build cosmology
77    return Cosmology.from_format(m, format="mapping", move_to_meta=True)
78
79
80def to_mypackage(cosmology, *args):
81    """Return the cosmology as a ``mycosmo``.
82
83    Parameters
84    ----------
85    cosmology : `~astropy.cosmology.Cosmology`
86
87    Returns
88    -------
89    `~mypackage.cosmology.MyCosmology`
90    """
91    if not isinstance(cosmology, FLRW):
92        raise TypeError("format 'mypackage' only supports FLRW cosmologies.")
93
94    # ----------------
95    # Cosmology provides a nice method "mapping", so all that needs to
96    # be done here is initialize from the dictionary
97    m = cosmology.to_format("mapping")
98
99    # Detect which type of MyCosmology to build.
100    # Here we have forced FlatLambdaCDM, but if your package allows for
101    # non-flat cosmologies...
102    m.pop("cosmology")
103
104    # MyCosmology doesn't support metadata. If your cosmology class does...
105    meta = m.pop("meta")
106    m = {**meta, **m}  # merge, preferring current values
107
108    # ----------------
109    # remap values
110    # MyCosmology doesn't support units, so take values.
111    m["hubble_parameter"] = m.pop("H0").to_value(u.km/u.s/u.Mpc)
112    m["initial_matter_density"] = m.pop("Om0")
113    m["initial_temperature"] = m.pop("Tcmb0").to_value(u.K)
114    # m["Neff"] = m.pop("Neff")  # skip b/c unchanged
115    m["neutrino_masses"] = m.pop("m_nu").to_value(u.eV)
116    m["initial_baryon_density"] = m.pop("Ob0")
117    m["current_age"] = m.pop("t0", cosmology.age(0 * cu.redshift)).to_value(u.Gyr)
118
119    # optional
120    if "z_reion" in m:
121        m["reionization_redshift"] = (m.pop("z_reion") << cu.redshift).value
122
123    # ...  # keep remapping
124
125    return MyCosmology(**m)
126
127
128def mypackage_identify(origin, format, *args, **kwargs):
129    """Identify if object uses format "mypackage"."""
130    itis = False
131    if origin == "read":
132        itis = isinstance(args[1], MyCosmology) and (format in (None, "mypackage"))
133    return itis
134
135
136# -------------------------------------------------------------------
137# Register to/from_format & identify methods with Astropy Unified I/O
138
139convert_registry.register_reader("mypackage", Cosmology, from_mypackage, force=True)
140convert_registry.register_writer("mypackage", Cosmology, to_mypackage, force=True)
141convert_registry.register_identifier("mypackage", Cosmology, mypackage_identify, force=True)
142