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