Custom Cosmology I/O formats#
Custom Cosmology To/From Formats#
Custom representation formats may also be registered into the Astropy Cosmology
I/O framework for use by these methods. For details of the framework see
I/O Registry (astropy.io.registry). Note Cosmology to/from_format uses a custom registry,
available at Cosmology.<to/from>_format.registry.
As an example, the following is an implementation of an Row converter. We can
and should use inbuilt parsers, like QTable, but to show a more complete
example we limit ourselves to only the “mapping” parser.
We start by defining the function to parse a Row into a Cosmology. This
function should take 1 positional argument, the row object, and 2 keyword
arguments, for how to handle extra metadata and which Cosmology class to use.
Details about metadata treatment are in
Cosmology.from_format.help("mapping").
>>> import copy
>>> from astropy.cosmology import Cosmology
>>> def from_table_row(row, *, move_to_meta=False, cosmology=None):
... # get name from column
... name = row['name'] if 'name' in row.columns else None
... meta = copy.deepcopy(row.meta)
... # turn row into mapping (dict of the arguments)
... mapping = dict(row)
... mapping['name'] = name
... mapping.setdefault("cosmology", meta.pop("cosmology", None))
... mapping["meta"] = meta
... # build cosmology from map
... return Cosmology.from_format(mapping, move_to_meta=move_to_meta,
... cosmology=cosmology)
convert_registry.register_reader("astropy.row", Cosmology, from_table_row)
The next step is a function to perform the reverse operation: parse a
Cosmology into a Row. This function requires only the cosmology object and
a *args to absorb unneeded information passed by
astropy.io.registry.UnifiedReadWrite (which implements
to_format()).
>>> from astropy.table import QTable
>>> def to_table_row(cosmology, *args):
... p = cosmology.to_format("mapping", cosmology_as_str=True)
... meta = p.pop("meta")
... # package parameters into lists for Table parsing
... params = {k: [v] for k, v in p.items()}
... return QTable(params, meta=meta)[0] # return row
convert_registry.register_writer("astropy.row", Cosmology, to_table_row)
Last we write a function to help with format auto-identification and then
register everything into astropy.io.registry.
>>> from astropy.cosmology import Cosmology
>>> from astropy.table import Row
>>> def row_identify(origin, format, *args, **kwargs):
... """Identify if object uses the Table format."""
... if origin == "read":
... return isinstance(args[1], Row) and (format in (None, "astropy.row"))
... return False
convert_registry.register_identifier("astropy.row", Cosmology, row_identify)
Now the registered functions can be used in from_format() and
to_format().
>>> from astropy.cosmology import Planck18
>>> row = Planck18.to_format("astropy.row")
>>> row
<Row index=0>
cosmology name H0 Om0 Tcmb0 Neff m_nu Ob0
km / (Mpc s) K eV
str13 str8 float64 float64 float64 float64 float64[3] float64
------------- -------- ------------ ------- ------- ------- ----------- -------
FlatLambdaCDM Planck18 67.66 0.30966 2.7255 3.046 0.0 .. 0.06 0.04897
>>> cosmo = Cosmology.from_format(row)
>>> cosmo == Planck18 # test it round-trips
True
Custom Cosmology Readers/Writers#
Custom read / write formats may be registered into the Astropy
Cosmology I/O framework. For details of the framework see I/O Registry (astropy.io.registry).
Note Cosmology read/write uses a custom registry, available at
Cosmology.<read/write>.registry.
As an example, in the following we will fully work out a Cosmology <-> JSON
(de)serializer. Note that we can use other registered parsers – here “mapping”
– to make the implementation much simpler.
We start by defining the function to parse JSON into a Cosmology. This
function should take 1 positional argument, the file object or file path. We
will also pass kwargs through to from_format(), which handles
metadata and which Cosmology class to use. Details of which are in
Cosmology.from_format.help("mapping").
>>> import json, os
>>> import astropy.units as u
>>> from astropy.cosmology import Cosmology
>>> from astropy.cosmology.connect import readwrite_registry
>>> def read_json(filename, **kwargs):
... # read file, from path-like or file-like
... if isinstance(filename, (str, bytes, os.PathLike)):
... with open(filename, "r") as file:
... data = file.read()
... else: # file-like : this also handles errors in dumping
... data = filename.read()
... mapping = json.loads(data) # parse json mappable to dict
... # deserialize Quantity
... for k, v in mapping.items():
... if isinstance(v, dict) and "value" in v and "unit" in v:
... mapping[k] = u.Quantity(v["value"], v["unit"])
... for k, v in mapping.get("meta", {}).items(): # also the metadata
... if isinstance(v, dict) and "value" in v and "unit" in v:
... mapping["meta"][k] = u.Quantity(v["value"], v["unit"])
... return Cosmology.from_format(mapping, **kwargs)
>>> readwrite_registry.register_reader("json", Cosmology, read_json)
The next step is a function to write a Cosmology to JSON. This function
requires the cosmology object and a file object/path. We also require the
boolean flag “overwrite” to set behavior for existing files. Note that
Quantity is not natively compatible with JSON. In both the write and
read methods we have to create custom parsers.
>>> def write_json(cosmology, file, *, overwrite=False, **kwargs):
... data = cosmology.to_format("mapping", cosmology_as_str=True) # start by turning into dict
... # serialize Quantity
... for k, v in data.items():
... if isinstance(v, u.Quantity):
... data[k] = {"value": v.value.tolist(), "unit": str(v.unit)}
... for k, v in data.get("meta", {}).items(): # also serialize the metadata
... if isinstance(v, u.Quantity):
... data["meta"][k] = {"value": v.value.tolist(), "unit": str(v.unit)}
...
... if isinstance(file, (str, bytes, os.PathLike)):
... # check that file exists and whether to overwrite.
... if os.path.exists(file) and not overwrite:
... raise IOError(f"{file} exists. Set 'overwrite' to write over.")
... with open(file, "w") as write_file:
... json.dump(data, write_file)
... else:
... json.dump(data, file)
>>> readwrite_registry.register_writer("json", Cosmology, write_json)
Last we write a function to help with format auto-identification and then
register everything into astropy.io.registry.
>>> def json_identify(origin, filepath, fileobj, *args, **kwargs):
... """Identify if object uses the JSON format."""
... return filepath is not None and filepath.endswith(".json")
>>> readwrite_registry.register_identifier("json", Cosmology, json_identify)
Now the registered functions can be used in read() and
write().
>>> import tempfile
>>> from astropy.cosmology import Planck18
>>>
>>> file = tempfile.NamedTemporaryFile()
>>> Planck18.write(file.name, format="json", overwrite=True)
>>> with open(file.name) as f: f.readlines()
['{"cosmology": "FlatLambdaCDM", "name": "Planck18",
"H0": {"value": 67.66, "unit": "km / (Mpc s)"}, "Om0": 0.30966,
...
>>>
>>> cosmo = Cosmology.read(file.name, format="json")
>>> file.close()
>>> cosmo == Planck18 # test it round-trips
True
:hide:
>>> from astropy.io.registry import IORegistryError
>>> readwrite_registry.unregister_reader("json", Cosmology)
>>> readwrite_registry.unregister_writer("json", Cosmology)
>>> readwrite_registry.unregister_identifier("json", Cosmology)
>>> try:
... readwrite_registry.get_reader("json", Cosmology)
... except IORegistryError:
... pass