From c1cd0c1d6bf1571ad26a10f8043df8c753ab3634 Mon Sep 17 00:00:00 2001 From: "Nathaniel Starkman (@nstarman)" Date: Tue, 28 Sep 2021 15:27:51 -0400 Subject: [PATCH] Add method ``__astropy_table__`` Allows for ``Table(cosmo)``. Signed-off-by: Nathaniel Starkman (@nstarman) --- astropy/cosmology/core.py | 20 ++++++++++++++++++ astropy/cosmology/io/tests/test_table.py | 7 +++--- astropy/cosmology/tests/test_core.py | 24 +++++++++++++++++++++ docs/changes/cosmology/12213.feature.rst | 1 + docs/cosmology/io.rst | 19 +++++++++++++++-- docs/table/construct_table.rst | 2 ++ docs/whatsnew/5.0.rst | 27 +++++++++++++++++++++--- 7 files changed, 92 insertions(+), 8 deletions(-) diff --git a/astropy/cosmology/core.py b/astropy/cosmology/core.py index 925ef17e8686..240ea3081263 100644 --- a/astropy/cosmology/core.py +++ b/astropy/cosmology/core.py @@ -428,6 +428,26 @@ def __repr__(self): return namelead + ", ".join(fmtps) + ")" + def __astropy_table__(self, cls, copy, **kwargs): + """Return a `~astropy.table.Table` of type ``cls``. + + Parameters + ---------- + cls : type + Astropy ``Table`` class or subclass. + copy : bool + Ignored. + **kwargs : dict, optional + Additional keyword arguments. Passed to ``self.to_format()``. + See ``Cosmology.to_format.help("astropy.table")`` for allowed kwargs. + + Returns + ------- + `astropy.table.Table` or subclass instance + Instance of type ``cls``. + """ + return self.to_format("astropy.table", cls=cls, **kwargs) + class FlatCosmologyMixin(metaclass=abc.ABCMeta): """ diff --git a/astropy/cosmology/io/tests/test_table.py b/astropy/cosmology/io/tests/test_table.py index 128bf83e88c8..af24b17f708b 100644 --- a/astropy/cosmology/io/tests/test_table.py +++ b/astropy/cosmology/io/tests/test_table.py @@ -110,8 +110,9 @@ def test_to_from_table_instance(self, cosmo, to_format, from_format): got = from_format(tbl, format="astropy.table") assert got == cosmo - # and it will also work if the cosmology is a string - tbl.meta["cosmology"] = _COSMOLOGY_CLASSES[tbl.meta["cosmology"]].__qualname__ + # and it will also work if the cosmology is a class + # Note this is not the default output of ``to_format``. + tbl.meta["cosmology"] = _COSMOLOGY_CLASSES[tbl.meta["cosmology"]] got = from_format(tbl, format="astropy.table") assert got == cosmo @@ -194,4 +195,4 @@ class TestToFromTable(ToFromFormatTestBase, ToFromTableTestMixin): """ def setup_class(self): - self.io_functions = {"to": to_table, "from": from_table} + self.functions = {"to": to_table, "from": from_table} diff --git a/astropy/cosmology/tests/test_core.py b/astropy/cosmology/tests/test_core.py index c1d471d8bad8..92dfb69ddc26 100644 --- a/astropy/cosmology/tests/test_core.py +++ b/astropy/cosmology/tests/test_core.py @@ -306,6 +306,30 @@ def test_is_equivalent(self, cosmo): # different class assert not cosmo.is_equivalent(2) + # ------------------------------------------------ + + @pytest.mark.parametrize("in_meta", [True, False]) + @pytest.mark.parametrize("table_cls", [Table, QTable]) + def test_astropy_table(self, cosmo, table_cls, in_meta): + """Test ``astropy.table.Table(cosmology)``.""" + tbl = table_cls(cosmo, cosmology_in_meta=in_meta) + + assert isinstance(tbl, table_cls) + # the name & all parameters are columns + for n in ("name", *cosmo.__parameters__): + assert n in tbl.colnames + assert all(tbl[n] == getattr(cosmo, n)) + # check if Cosmology is in metadata or a column + if in_meta: + assert tbl.meta["cosmology"] == cosmo.__class__.__qualname__ + assert "cosmology" not in tbl.colnames + else: + assert "cosmology" not in tbl.meta + assert tbl["cosmology"][0] == cosmo.__class__.__qualname__ + # the metadata is transferred + for k, v in cosmo.meta.items(): + assert np.all(tbl.meta[k] == v) + class CosmologySubclassTest(TestCosmology): """ diff --git a/docs/changes/cosmology/12213.feature.rst b/docs/changes/cosmology/12213.feature.rst index 0a3cb01f9aa5..0c3d3f56a25e 100644 --- a/docs/changes/cosmology/12213.feature.rst +++ b/docs/changes/cosmology/12213.feature.rst @@ -1,2 +1,3 @@ Register Astropy Table into Cosmology's ``to/from_format`` I/O, allowing a Cosmology instance to be parsed from or converted to a Table instance. +Also adds the ``__astropy_table__`` method allowing ``Table(cosmology)``. diff --git a/docs/cosmology/io.rst b/docs/cosmology/io.rst index d6ba84b080c0..baa9f6d2fa47 100644 --- a/docs/cosmology/io.rst +++ b/docs/cosmology/io.rst @@ -142,8 +142,8 @@ to the ``Planck18`` cosmology from which it was created. .. EXAMPLE START: Planck18 to QTable and back - Another pre-registered format is "table", for converting a |Cosmology| to - and from a |QTable|. +Another pre-registered format is "table", for converting a |Cosmology| to +and from a |QTable|. .. code-block:: @@ -156,6 +156,21 @@ to the ``Planck18`` cosmology from which it was created. -------- ------------ ------- ------- ------- ----------- ------- Planck18 67.66 0.30966 2.7255 3.046 0.0 .. 0.06 0.04897 +Cosmology supports the astropy Table-like protocol (see :ref:`Table-like Objects`) +to the same effect: + +.. code-block:: + + >>> from astropy.table import QTable + >>> ct = QTable(Planck18) + >>> ct + + 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 :class:|QTable| can be used to load a new cosmological instance identical to the ``Planck18`` cosmology from which it was created. diff --git a/docs/table/construct_table.rst b/docs/table/construct_table.rst index eb40de2e43ad..3af12e0a4150 100644 --- a/docs/table/construct_table.rst +++ b/docs/table/construct_table.rst @@ -1184,6 +1184,8 @@ units, see the :ref:`columns_with_units` section. .. EXAMPLE END +.. _Table-like Objects: + Table-like Objects ================== diff --git a/docs/whatsnew/5.0.rst b/docs/whatsnew/5.0.rst index a423c1bc6165..a2d86bbaf2bb 100644 --- a/docs/whatsnew/5.0.rst +++ b/docs/whatsnew/5.0.rst @@ -16,7 +16,7 @@ Support for reading, writing, and converting ``Cosmology`` Four new methods -- ``read``, ``write``, ``to_format``, ``from_format`` -- have been added to the ``Cosmology`` class, enabling reading from and writing to -files and converting between different python objects. +files and converting between different Python objects. The methods use Astropy's Unified I/O registry so custom formats can be registered. Details are discussed in an addition to the docs. @@ -36,7 +36,7 @@ Currently no file formats are registered, but the syntax is as follows: The transformations between ``Cosmology`` and `dict`/``QTable`` are pre-registered, e.g. enabling:: - >>> from astropy.cosmology import Planck18 + >>> from astropy.cosmology import Cosmology, Planck18 >>> cm = Planck18.to_format("mapping") >>> cm {'cosmology': , @@ -44,6 +44,9 @@ pre-registered, e.g. enabling:: 'H0': , 'Om0': 0.30966, ... + + or to a Table with: + >>> ct = Planck18.to_format("astropy.table") >>> ct @@ -53,11 +56,29 @@ pre-registered, e.g. enabling:: -------- ------------ ------- ------- ------- ----------- ------- Planck18 67.66 0.30966 2.7255 3.046 0.0 .. 0.06 0.04897 - >>> from astropy.cosmology import Cosmology + Even more succinctly: + + >>> from astropy.table import QTable + >>> QTable(Planck18) + + name H0 Om0 Tcmb0 Neff m_nu [3] Ob0 + km / (Mpc s) eV + str8 float64 float64 float64 float64 float64 float64 + -------- ------------ ------- ------- ------- ----------- ------- + Planck18 67.66 0.30966 2.7255 3.046 0.0 .. 0.06 0.04897 + + All the format conversions can be reversed: + >>> cosmo = Cosmology.from_format(cm, format="mapping") + >>> 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) + + >>> cosmo = Cosmology.from_format(ct, format="astropy.table") >>> cosmo == Planck18 True + .. _whatsnew-5.0-cosmology-units: ``Cosmology`` units module