Mubin17 · May 31, 2023 · May 31, 2023 · May 31, 2023 · May 31, 2023 · May 31, 2023
@@ -5,4 +5,4 @@
 """
 
 # Import to register with the I/O machinery
-from . import cosmology, ecsv, html, latex, mapping, model, row, table, yaml
+from . import cosmology, ecsv, html, latex, mapping, model, mrt, row, table, yaml
@@ -0,0 +1,147 @@
+# THIRD PARTY
+import numpy as np
+
+import astropy.cosmology.units as cu
+import astropy.units as u
+from astropy.cosmology.connect import readwrite_registry
+from astropy.cosmology.core import Cosmology
+from astropy.table import QTable
+from astropy.table.serialize import represent_mixins_as_columns
+
+from .table import from_table, to_table
+
+
+def read_mrt(filename, index=None, *, move_to_meta=False, cosmology=None, **kwargs):
+    """Read a |Cosmology| from a MRT file.
+
+    Parameters
+    ----------
+    filename : path-like or file-like
+        The file from which to read the MRT data.
+    index : int or 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 ``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 or |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'.
+    **kwargs : Any
+        Additional keyword arguments passed to :attr:`astropy.table.QTable.read`.
+        The 'format' is set to 'mrt', regardless of input.
+
+    Returns
+    -------
+    |Cosmology| subclass instance
+
+    Raises
+    ------
+    ValueError
+        If the keyword argument 'format' is given and is not "ascii.mrt".
+    """
+    # Check that the format is 'mrt' (or not specified)
+    format = kwargs.pop("format", "ascii.mrt")
+    if format != "ascii.mrt":
+        raise ValueError(f"format must be 'ascii.mrt', not {format}")
+
+    # Reading is handled by `QTable`.
+    table = QTable.read(filename, format="ascii.mrt", **kwargs)
+
+    # Create a single column named 'm_nu' by combining the values from the 'm_nu_i' columns
+    m_nu_data = []
+    i, more_m_nu = 0, True
+    while more_m_nu:
+        cn = f"m_nu[{i}]"  # column name
+        if cn not in table.colnames:
+            more_m_nu = False
+            continue
+
+        m_nu_data.append(table[cn][0])
+        table.remove_column(cn)
+
+        i += 1  # increment the m_nu index
+
+    col = (
+        [m_nu_data]
+        if not isinstance(m_nu_data[0], u.Quantity)
+        else [u.Quantity(m_nu_data)]
+    )
+    table.add_column(col, name="m_nu", index=-2)
+
+    # Build the cosmology from table, using the private backend.
+    return from_table(
+        table, index=index, move_to_meta=move_to_meta, cosmology=cosmology
+    )
+
+
+def write_mrt(cosmology, file, *, overwrite=False, cls=QTable, **kwargs):
+    """Serialize the |Cosmology| into an MRT file.
+
+    Parameters
+    ----------
+    cosmology : |Cosmology| subclass instance
+        The |Cosmology| object to serialize into an MRT file.
+    file : path-like or file-like
+        The file path or file-like object where the MRT file will be written.
+    overwrite : bool, optional keyword-only
+        Whether to overwrite the file if it already exists.
+    cls : |Table| class, optional keyword-only
+        The |Table| class to use when writing the MRT file. Default is |QTable|.
+    **kwargs : Any
+        Additional keyword arguments passed to ``cls.write``.
+
+    Raises
+    ------
+    TypeError
+        If kwarg (optional) 'cls' is not a subclass of `astropy.table.Table`
+    ValueError
+        If the keyword argument 'format' is given and is not "ascii.mrt".
+    """
+    # Check that the format is 'mrt' (or not specified)
+    format = kwargs.pop("format", "ascii.mrt")
+    if format != "ascii.mrt":
+        raise ValueError(f"format must be 'ascii.mrt', not {format}")
+
+    # Set cosmology_in_meta as false for now since there is no metadata being kept
+    table = represent_mixins_as_columns(
+        to_table(cosmology, cls=cls, cosmology_in_meta=False)
+    )
+
+    # CDS can't serialize redshift units, so remove them  # TODO: fix this
+    for k, col in table.columns.items():
+        if col.unit is cu.redshift:
+            table[k] <<= u.dimensionless_unscaled
+
+    # Replace the m_nu column with individual columns
+    if "m_nu" in table.colnames:
+        m_nu = table["m_nu"]
+
+        index = table.colnames.index("m_nu")
+        table.remove_column("m_nu")
+        m_nu = np.atleast_1d(m_nu)  # Ensure m_nu is an array
+        if m_nu.shape == (1,):
+            # If m_nu is a scalar, add a single column
+            table.add_column(m_nu[0], name="m_nu[0]", index=index)
+        else:
+            # If m_nu is an array, add multiple columns with names 'm_nu_[i]'
+            cols, names = tuple(
+                zip(*((m, f"m_nu[{i}]") for i, m in enumerate(m_nu[0])))
+            )
+            table.add_columns(cols, names=names, indexes=(index,) * len(m_nu[0]))
+
+    table.write(file, overwrite=overwrite, format="ascii.mrt", **kwargs)
+
+
+# ===================================================================
+# Register
+
+readwrite_registry.register_reader("ascii.mrt", Cosmology, read_mrt)
+readwrite_registry.register_writer("ascii.mrt", Cosmology, write_mrt)
@@ -12,7 +12,7 @@ def test_expected_readwrite_io():
     """Test that ONLY the expected I/O is registered."""
 
     got = {k for k, _ in readwrite_registry._readers.keys()}
-    expected = {"ascii.ecsv", "ascii.html"}
+    expected = {"ascii.ecsv", "ascii.html", "ascii.mrt"}
 
     assert got == expected
 

@@ -0,0 +1,140 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+
+# THIRD PARTY
+import pytest
+
+# LOCAL
+from astropy.cosmology.io.mrt import read_mrt, write_mrt
+from astropy.table import QTable, Table
+
+from .base import ReadWriteDirectTestBase, ReadWriteTestMixinBase
+
+
+class ReadWriteMRTTestMixin(ReadWriteTestMixinBase):
+    """
+    Tests for a Cosmology[Read/Write] with ``format="mrt"``.
+    This class will not be directly called by :mod:`pytest` since its name does
+    not begin with ``Test``. To activate the contained tests this class must
+    be inherited in a subclass. Subclasses must dfine a :func:`pytest.fixture`
+    ``cosmo`` that returns/yields an instance of a |Cosmology|.
+    See ``TestCosmology`` for an example.
+    """
+
+    def test_to_mrt_bad_index(self, read, write, tmp_path):
+        """Test if argument ``index`` is incorrect"""
+        fp = tmp_path / "test_to_mrt_bad_index.mrt"
+
+        write(fp, format="ascii.mrt")
+
+        # single-row table and has a non-0/None index
+        with pytest.raises(IndexError, match="index 2 out of range"):
+            read(fp, index=2, format="ascii.mrt")
+
+        # string index where doesn't match
+        with pytest.raises(KeyError, match="No matches found for key"):
+            read(fp, index="row 0", format="ascii.mrt")
+
+    # -----------------------
+
+    def test_to_mrt_failed_cls(self, write, tmp_path):
+        """Test failed table type."""
+        fp = tmp_path / "test_to_mrt_failed_cls.mrt"
+
+        with pytest.raises(TypeError, match="'cls' must be"):
+            write(fp, format="ascii.mrt", cls=list)
+
+    # -----------------------
+
+    @pytest.mark.parametrize("tbl_cls", [QTable, Table])
+    def test_to_mrt_cls(self, write, tbl_cls, tmp_path):
+        fp = tmp_path / "test_to_mrt_cls.mrt"
+        write(fp, format="ascii.mrt", cls=tbl_cls)
+
+    # -----------------------
+
+    def test_readwrite_mrt_instance(self, cosmo_cls, cosmo, read, write, tmp_path):
+        """Test cosmology -> ascii.mrt -> cosmology."""
+        fp = tmp_path / "test_readwrite_mrt_instance.mrt"
+
+        # ------------
+        # To Table
+
+        write(fp, format="ascii.mrt")
+
+        # some checks on the saved file
+        tbl = Table.read(fp, format="ascii.mrt")
+        assert tbl["name"] == cosmo.name
+
+        # ------------
+        # From Table
+
+        tbl["mismatching"] = "will error"
+        tbl.write(fp, format="ascii.mrt", overwrite=True)
+
+        # tests are different if the last argument is a **kwarg
+        if tuple(cosmo._init_signature.parameters.values())[-1].kind == 4:
+            got = read(fp, format="ascii.mrt")
+
+            assert got.__class__ is cosmo_cls
+            assert got.name == cosmo.name
+            assert "mismatching" not in got.meta
+
+            return  # don't continue testing
+
+        # read with mismatching parameters errors
+        with pytest.raises(TypeError, match="there are unused parameters"):
+            read(fp, format="ascii.mrt")
+
+        # unless mismatched are moved to meta
+        got = read(fp, format="ascii.mrt", move_to_meta=True)
+        assert got == cosmo
+        assert got.meta["mismatching"] == "will error"
+
+        # it won't error if everything matches up
+        tbl.remove_column("mismatching")
+        tbl.write(fp, format="ascii.mrt", overwrite=True)
+        got = read(fp, format="ascii.mrt")
+        assert got == cosmo
+
+    # -----------------------
+
+    def test_readwrite_mrt_subclass_partial_info(
+        self, cosmo_cls, cosmo, read, write, tmp_path
+    ):
+        """
+        Test writing from an instance and reading from that class.
+        This works with missing information.
+        """
+        fp = tmp_path / "test_read_mrt_subclass_partial_info.mrt"
+
+        # test write
+        write(fp, format="ascii.mrt")
+
+        # partial information
+        tbl = Table.read(fp, format="ascii.mrt")
+        del tbl["Tcmb0"]
+        tbl.write(fp, format="ascii.mrt", overwrite=True)
+
+        # read with the same class that wrote fills in the missing info with
+        # the default value
+        got = cosmo_cls.read(fp, format="ascii.mrt")
+        got2 = read(fp, format="ascii.mrt", cosmology=cosmo_cls)
+
+        assert got == got2  # internal consistency
+
+        # not equal, because Tcmb0 is changed, which also changes m_nu
+        assert got != cosmo
+        assert got.Tcmb0 == cosmo_cls._init_signature.parameters["Tcmb0"].default
+        assert got.clone(name=cosmo.name, Tcmb0=cosmo.Tcmb0, m_nu=cosmo.m_nu) == cosmo
+
+
+class TestReadWriteMRT(ReadWriteDirectTestBase, ReadWriteMRTTestMixin):
+    """
+    Directly test ``read/write_mrt``.
+    These are not public API and are discouraged from use, in favor of
+    ``Cosmology.read/write(..., format="mrt")``, but should be
+    tested regardless b/c they are used internally.
+    """
+
+    def setup_class(self):
+        self.functions = {"read": read_mrt, "write": write_mrt}
@@ -16,6 +16,7 @@
     test_latex,
     test_mapping,
     test_model,
+    test_mrt,
     test_row,
     test_table,
     test_yaml,
@@ -34,6 +35,7 @@
     ("ascii.ecsv", True, True),
     ("ascii.html", False, HAS_BS4),
     ("ascii.latex", False, True),
+    ("ascii.mrt", False, True),
     ("json", True, True),
     ("latex", False, True),
 }
@@ -59,6 +61,7 @@ class ReadWriteTestMixin(
     test_html.ReadWriteHTMLTestMixin,
     test_json.ReadWriteJSONTestMixin,
     test_latex.WriteLATEXTestMixin,
+    test_mrt.ReadWriteMRTTestMixin,
 ):
     """
     Tests for a CosmologyRead/Write on a |Cosmology|.

diff --git a/docs/changes/cosmology/14883.feature.rst b/docs/changes/cosmology/14883.feature.rst
@@ -0,0 +1,2 @@
+Registered the ``read_mrt()`` and ``write_mrt()`` methods with the Cosmology class,
+enabling the import and export of cosmological data to and from MRT files.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		Registered the ``read_mrt()`` and ``write_mrt()`` methods with the Cosmology class,
		enabling the import and export of cosmological data to and from MRT files.