Register in “table” for Cosmology <-> QTable

Signed-off-by: Nathaniel Starkman (@nstarman) <nstarkman@protonmail.com>

astropy/cosmology/io/__init__.py

@@ -6,4 +6,4 @@
 """
 
 # Import the interchange to register them into the io registry.
-from . import mapping  # noqa: F403
+from . import mapping, table  # noqa: F403

astropy/cosmology/io/table.py

@@ -0,0 +1,246 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+
+import copy
+
+import numpy as np
+
+from astropy.table import Table, QTable
+from astropy.cosmology.connect import convert_registry
+from astropy.cosmology.core import Cosmology
+
+from .mapping import from_mapping, to_mapping
+
+
+def from_table(table, index=None, *, move_to_meta=False, cosmology=None):
+    """Instantiate a `~astropy.cosmology.Cosmology` from a |QTable|.
+
+    Parameters
+    ----------
+    table : `~astropy.table.QTable`
+    index : int, str, or None, optional
+        Needed to select the row in tables with multiple rows. ``index`` can be
+        an integer for the row number or, if the table is indexed by a column,
+        the value of that column. If the table is not indexed and ``index``
+        is a string, the "name" column is used as the indexing column.
+
+    move_to_meta : bool (optional, keyword-only)
+        Whether to move keyword arguments that are not in the Cosmology class'
+        signature to the Cosmology's metadata. This will only be applied if the
+        Cosmology does NOT have a keyword-only argument (e.g. ``**kwargs``).
+        Arguments moved to the metadata will be merged with existing metadata,
+        preferring specified metadata in the case of a merge conflict
+        (e.g. for ``Cosmology(meta={'key':10}, key=42)``, the ``Cosmology.meta``
+        will be ``{'key': 10}``).
+
+    cosmology : str, `~astropy.cosmology.Cosmology` class, or None (optional, keyword-only)
+        The cosmology class (or string name thereof) to use when constructing
+        the cosmology instance. The class also provides default parameter values,
+        filling in any non-mandatory arguments missing in 'table'.
+
+    Returns
+    -------
+    `~astropy.cosmology.Cosmology` subclass instance
+
+    Examples
+    --------
+    To see loading a `~astropy.cosmology.Cosmology` from a Table with
+    ``from_table``, we will first make a |QTable| using
+    :func:`~astropy.cosmology.Cosmology.to_format`.
+
+        >>> from astropy.cosmology import Cosmology, Planck18
+        >>> ct = Planck18.to_format("astropy.table")
+        >>> ct
+        <QTable length=1>
+          name        H0        Om0    Tcmb0    Neff    m_nu [3]    Ob0
+                 km / (Mpc s)            K                 eV
+          str8     float64    float64 float64 float64   float64   float64
+        -------- ------------ ------- ------- ------- ----------- -------
+        Planck18        67.66 0.30966  2.7255   3.046 0.0 .. 0.06 0.04897
+
+    Now this table can be used to load a new cosmological instance identical
+    to the ``Planck18`` cosmology from which it was generated.
+
+        >>> cosmo = Cosmology.from_format(ct, format="astropy.table")
+        >>> cosmo
+        FlatLambdaCDM(name="Planck18", H0=67.7 km / (Mpc s), Om0=0.31,
+                      Tcmb0=2.725 K, Neff=3.05, m_nu=[0. 0. 0.06] eV, Ob0=0.049)
+
+    Specific cosmology classes can be used to parse the data. The class'
+    default parameter values are used to fill in any information missing in the
+    data.
+
+        >>> from astropy.cosmology import FlatLambdaCDM
+        >>> del ct["Tcmb0"]  # show FlatLambdaCDM provides default
+        >>> FlatLambdaCDM.from_format(ct)
+        FlatLambdaCDM(name="Planck18", H0=67.7 km / (Mpc s), Om0=0.31,
+                      Tcmb0=0 K, Neff=3.05, m_nu=None, Ob0=0.049)
+
+    For tables with multiple rows of cosmological parameters, the ``index``
+    argument is needed to select the correct row. The index can be an integer
+    for the row number or, if the table is indexed by a column, the value of
+    that column. If the table is not indexed and ``index`` is a string, the
+    "name" column is used as the indexing column.
+
+    Here is an example where ``index`` is needed and can be either an integer
+    (for the row number) or the name of one of the cosmologies, e.g. 'Planck15'.
+
+        >>> from astropy.cosmology import Planck13, Planck15, Planck18
+        >>> from astropy.table import vstack
+        >>> cts = vstack([c.to_format("astropy.table")
+        ...               for c in (Planck13, Planck15, Planck18)],
+        ...              metadata_conflicts='silent')
+        >>> cts
+        <QTable length=3>
+          name        H0        Om0    Tcmb0    Neff    m_nu [3]    Ob0
+                 km / (Mpc s)            K                 eV
+          str8     float64    float64 float64 float64   float64   float64
+        -------- ------------ ------- ------- ------- ----------- --------
+        Planck13        67.77 0.30712  2.7255   3.046 0.0 .. 0.06 0.048252
+        Planck15        67.74  0.3075  2.7255   3.046 0.0 .. 0.06   0.0486
+        Planck18        67.66 0.30966  2.7255   3.046 0.0 .. 0.06  0.04897
+
+        >>> cosmo = Cosmology.from_format(cts, index=1, format="astropy.table")
+        >>> cosmo == Planck15
+        True
+
+    For further examples, see :doc:`astropy:cosmology/io`.
+    """
+    # ------------------
+    # Get row from table
+
+    # string index uses the indexed column on the table to find the row index.
+    if isinstance(index, str):
+        if not table.indices:  # no indexing column, find by string match
+            index = np.where(table['name'] == index)[0][0]
+        else:  # has indexing column
+            index = table.loc_indices[index]  # need to convert to row index (int)
+
+    # no index is needed for a 1-row table. For a multi-row table...
+    if index is None:
+        if len(table) != 1:  # multi-row table and no index
+            raise ValueError("need to select a specific row (e.g. index=1) when "
+                             "constructing a Cosmology from a multi-row table.")
+        else:  # single-row table
+            index = 0
+    row = table[index]  # index is now the row index (int)
+
+    # ------------------
+    # parse row to cosmo
+
+    # special values
+    name = row['name'] if 'name' in row.columns else None  # get name from column
+    meta = copy.deepcopy(row.meta)
+
+    # turn row into mapping, filling cosmo if not in a column
+    mapping = dict(row)
+    mapping["name"] = name
+    mapping.setdefault("cosmology", meta.pop("cosmology", None))
+    mapping["meta"] = meta
+
+    # build cosmology from map
+    return from_mapping(mapping, move_to_meta=move_to_meta, cosmology=cosmology)
+
+
+def to_table(cosmology, *args, cls=QTable, cosmology_in_meta=True):
+    """Serialize the cosmology into a `~astropy.table.QTable`.
+
+    Parameters
+    ----------
+    cosmology : `~astropy.cosmology.Cosmology` subclass instance
+    *args
+        Not used. Needed for compatibility with
+        `~astropy.io.registry.UnifiedReadWriteMethod`
+    cls: type (optional, keyword-only)
+        Astropy :class:`~astropy.table.Table` class or subclass type to return.
+        Default is :class:`~astropy.table.QTable`.
+    cosmology_in_meta : bool
+        Whether to put the cosmology class in the Table metadata (if `True`,
+        default) or as the first column (if `False`).
+
+    Returns
+    -------
+    `~astropy.table.QTable`
+        With columns for the cosmology parameters, and metadata and
+        cosmology class name in the Table's ``meta`` attribute
+
+    Raises
+    ------
+    TypeError
+        If kwarg (optional) 'cls' is not a subclass of `astropy.table.Table`
+
+    Examples
+    --------
+    A Cosmology as a `~astropy.table.QTable` will have the cosmology's name and
+    parameters as columns.
+
+        >>> from astropy.cosmology import Planck18
+        >>> ct = Planck18.to_format("astropy.table")
+        >>> ct
+        <QTable length=1>
+          name        H0        Om0    Tcmb0    Neff    m_nu [3]    Ob0
+                 km / (Mpc s)            K                 eV
+          str8     float64    float64 float64 float64   float64   float64
+        -------- ------------ ------- ------- ------- ----------- -------
+        Planck18        67.66 0.30966  2.7255   3.046 0.0 .. 0.06 0.04897
+
+    The cosmological class and other metadata, e.g. a paper reference, are in
+    the Table's metadata.
+
+        >>> ct.meta
+        OrderedDict([..., ('cosmology', 'FlatLambdaCDM')])
+
+    To move the cosmology class from the metadata to a Table row, set the
+    ``cosmology_in_meta`` argument to `False`:
+
+        >>> Planck18.to_format("astropy.table", cosmology_in_meta=False)
+        <QTable length=1>
+          cosmology     name        H0        Om0    Tcmb0    Neff    m_nu [3]    Ob0
+                               km / (Mpc s)            K                 eV
+            str13       str8     float64    float64 float64 float64   float64   float64
+        ------------- -------- ------------ ------- ------- ------- ----------- -------
+        FlatLambdaCDM Planck18        67.66 0.30966  2.7255   3.046 0.0 .. 0.06 0.04897
+
+    Astropy recommends `~astropy.table.QTable` for tables with
+    `~astropy.units.Quantity` columns. However the returned type may be
+    overridden using the ``cls`` argument:
+
+        >>> from astropy.table import Table
+        >>> Planck18.to_format("astropy.table", cls=Table)
+        <Table length=1>
+        ...
+    """
+    if not issubclass(cls, Table):
+        raise TypeError(f"'cls' must be a (sub)class of Table, not {type(cls)}")
+
+    # start by getting a map representation. This requires minimal repackaging.
+    p = to_mapping(cosmology)
+    p["cosmology"] = p["cosmology"].__qualname__
+    # create metadata from mapping
+    meta = p.pop("meta")
+    if cosmology_in_meta:  # move class to Table meta
+        meta["cosmology"] = p.pop("cosmology")
+    # package parameters into lists for Table parsing
+    params = {k: [v] for k, v in p.items()}
+
+    return cls(params, meta=meta)
+
+
+def table_identify(origin, format, *args, **kwargs):
+    """Identify if object uses the Table format.
+
+    Returns
+    -------
+    bool
+    """
+    itis = False
+    if origin == "read":
+        itis = isinstance(args[1], Table) and (format in (None, "astropy.table"))
+    return itis
+
+
+# ===================================================================
+# Register
+
+convert_registry.register_reader("astropy.table", Cosmology, from_table)
+convert_registry.register_writer("astropy.table", Cosmology, to_table)
+convert_registry.register_identifier("astropy.table", Cosmology, table_identify)