[ENH] Access output in timeseries by its time or redshift

doc/source/analyzing/time_series_analysis.rst

@@ -77,6 +77,29 @@ see:
  * The cookbook recipe for :ref:`cookbook-time-series-analysis`
  * :class:`~yt.data_objects.time_series.DatasetSeries`
 
+In addition, the :class:`~yt.data_objects.time_series.DatasetSeries` object allows to
+select an output based on its time or by its redshift (if defined) as follows:
+
+.. code-block:: python
+
+   import yt
+
+   ts = yt.load("*/*.index")
+   # Get output at 3 Gyr
+   ds = ts.get_by_time((3, "Gyr"))
+   # This will fail if no output is found within 100 Myr
+   ds = ts.get_by_time((3, "Gyr"), tolerance=(100, "Myr"))
+   # Get the output at the time right before and after 3 Gyr
+   ds_before = ts.get_by_time((3, "Gyr"), prefer="smaller")
+   ds_after = ts.get_by_time((3, "Gyr"), prefer="larger")
+
+   # For cosmological simulations, you can also select an output by its redshift
+   # with the same options as above
+   ds = ts.get_by_redshift(0.5)
+
+For more information, see :meth:`~yt.data_objects.time_series.DatasetSeries.get_by_time` and
+:meth:`~yt.data_objects.time_series.DatasetSeries.get_by_redshift`.
+
 .. _analyzing-an-entire-simulation:
 
 Analyzing an Entire Simulation

nose_ignores.txt

@@ -47,3 +47,4 @@
 --ignore-file=test_disks\.py
 --ignore-file=test_offaxisprojection_pytestonly\.py
 --ignore-file=test_sph_pixelization_pytestonly\.py
+--ignore-file=test_time_series\.py

tests/tests.yaml

@@ -223,6 +223,7 @@ other_tests:
      - "--ignore-file=test_gadget_pytest\\.py"
      - "--ignore-file=test_vr_orientation\\.py"
      - "--ignore-file=test_particle_trajectories_pytest\\.py"
+     - "--ignore-file=test_time_series\\.py"
      - "--exclude-test=yt.frontends.gdf.tests.test_outputs.TestGDF"
      - "--exclude-test=yt.frontends.adaptahop.tests.test_outputs"
      - "--exclude-test=yt.frontends.stream.tests.test_stream_particles.test_stream_non_cartesian_particles"

yt/data_objects/tests/test_time_series.py

@@ -1,7 +1,7 @@
 import tempfile
 from pathlib import Path
 
-from numpy.testing import assert_raises
+import pytest
 
 from yt.data_objects.static_output import Dataset
 from yt.data_objects.time_series import DatasetSeries
@@ -29,76 +29,127 @@ def test_pattern_expansion():
 def test_no_match_pattern():
     with tempfile.TemporaryDirectory() as tmpdir:
         pattern = Path(tmpdir).joinpath("fake_data_file_*")
-        assert_raises(
-            FileNotFoundError, DatasetSeries._get_filenames_from_glob_pattern, pattern
-        )
+        with pytest.raises(FileNotFoundError):
+            DatasetSeries._get_filenames_from_glob_pattern(pattern)
 
 
-def test_init_fake_dataseries():
-    file_list = [f"fake_data_file_{str(i).zfill(4)}" for i in range(10)]
+@pytest.fixture
+def FakeDataset():
+    class _FakeDataset(Dataset):
+        """A minimal loadable fake dataset subclass"""
+
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+
+        @classmethod
+        def _is_valid(cls, *args, **kwargs):
+            return True
+
+        def _parse_parameter_file(self):
+            return
+
+        def _set_code_unit_attributes(self):
+            return
+
+        def set_code_units(self):
+            i = int(Path(self.filename).name.split("_")[-1])
+            self.current_time = i
+            self.current_redshift = 1 / (i + 1)
+            return
+
+        def _hash(self):
+            return
+
+        def _setup_classes(self):
+            return
+
+    try:
+        yield _FakeDataset
+    finally:
+        output_type_registry.pop("_FakeDataset")
+
+
+@pytest.fixture
+def fake_datasets():
+    file_list = [f"fake_data_file_{i:04d}" for i in range(10)]
     with tempfile.TemporaryDirectory() as tmpdir:
         pfile_list = [Path(tmpdir) / file for file in file_list]
         sfile_list = [str(file) for file in pfile_list]
         for file in pfile_list:
             file.touch()
         pattern = Path(tmpdir) / "fake_data_file_*"
 
-        # init from str pattern
-        ts = DatasetSeries(pattern)
-        assert ts._pre_outputs == sfile_list
+        yield file_list, pfile_list, sfile_list, pattern
+
+
+def test_init_fake_dataseries(fake_datasets):
+    file_list, pfile_list, sfile_list, pattern = fake_datasets
+
+    # init from str pattern
+    ts = DatasetSeries(pattern)
+    assert ts._pre_outputs == sfile_list
+
+    # init from Path pattern
+    ppattern = Path(pattern)
+    ts = DatasetSeries(ppattern)
+    assert ts._pre_outputs == sfile_list
+
+    # init form str list
+    ts = DatasetSeries(sfile_list)
+    assert ts._pre_outputs == sfile_list
+
+    # init form Path list
+    ts = DatasetSeries(pfile_list)
+    assert ts._pre_outputs == pfile_list
+
+    # rejected input type (str repr of a list) "[file1, file2, ...]"
+    with pytest.raises(FileNotFoundError):
+        DatasetSeries(str(file_list))
 
-        # init from Path pattern
-        ppattern = Path(pattern)
-        ts = DatasetSeries(ppattern)
-        assert ts._pre_outputs == sfile_list
+    # finally, check that ts[0] fails to actually load
+    with pytest.raises(YTUnidentifiedDataType):
+        ts[0]
 
-        # init form str list
-        ts = DatasetSeries(sfile_list)
-        assert ts._pre_outputs == sfile_list
 
-        # init form Path list
-        ts = DatasetSeries(pfile_list)
-        assert ts._pre_outputs == pfile_list
+def test_init_fake_dataseries2(FakeDataset, fake_datasets):
+    _file_list, _pfile_list, _sfile_list, pattern = fake_datasets
+    ds = DatasetSeries(pattern)[0]
+    assert isinstance(ds, FakeDataset)
 
-        # rejected input type (str repr of a list) "[file1, file2, ...]"
-        assert_raises(FileNotFoundError, DatasetSeries, str(file_list))
+    ts = DatasetSeries(pattern, my_unsupported_kwarg=None)
 
-        # finally, check that ts[0] fails to actually load
-        assert_raises(YTUnidentifiedDataType, ts.__getitem__, 0)
+    with pytest.raises(TypeError):
+        ts[0]
 
-        class FakeDataset(Dataset):
-            """A minimal loadable fake dataset subclass"""
 
-            @classmethod
-            def _is_valid(cls, *args, **kwargs):
-                return True
+def test_get_by_key(FakeDataset, fake_datasets):
+    _file_list, _pfile_list, sfile_list, pattern = fake_datasets
+    ts = DatasetSeries(pattern)
 
-            def _parse_parameter_file(self):
-                return
+    Ntot = len(sfile_list)
 
-            def _set_code_unit_attributes(self):
-                return
+    t = ts[0].quan(1, "code_time")
 
-            def set_code_units(self):
-                self.current_time = 0
-                return
+    assert sfile_list[0] == ts.get_by_time(-t).filename
+    assert sfile_list[0] == ts.get_by_time(t - t).filename
+    assert sfile_list[1] == ts.get_by_time((0.8, "code_time")).filename
+    assert sfile_list[1] == ts.get_by_time((1.2, "code_time")).filename
+    assert sfile_list[Ntot - 1] == ts.get_by_time(t * (Ntot - 1)).filename
+    assert sfile_list[Ntot - 1] == ts.get_by_time(t * Ntot).filename
 
-            def _hash(self):
-                return
+    with pytest.raises(ValueError):
+        ts.get_by_time(-2 * t, tolerance=0.1 * t)
+    with pytest.raises(ValueError):
+        ts.get_by_time(1000 * t, tolerance=0.1 * t)
 
-            def _setup_classes(self):
-                return
+    assert sfile_list[1] == ts.get_by_redshift(1 / 2.2).filename
+    assert sfile_list[1] == ts.get_by_redshift(1 / 2).filename
+    assert sfile_list[1] == ts.get_by_redshift(1 / 1.6).filename
 
-        try:
-            ds = DatasetSeries(pattern)[0]
-            assert isinstance(ds, FakeDataset)
+    with pytest.raises(ValueError):
+        ts.get_by_redshift(1000, tolerance=0.1)
 
-            ts = DatasetSeries(pattern, my_unsupported_kwarg=None)
+    zmid = (ts[0].current_redshift + ts[1].current_redshift) / 2
 
-            assert_raises(TypeError, ts.__getitem__, 0)
-            # the exact error message is supposed to be this
-            # """__init__() got an unexpected keyword argument 'my_unsupported_kwarg'"""
-            # but it's hard to check for within the framework
-        finally:
-            # tear down to avoid possible breakage in following tests
-            output_type_registry.pop("FakeDataset")
+    assert sfile_list[1] == ts.get_by_redshift(zmid, prefer="smaller").filename
+    assert sfile_list[0] == ts.get_by_redshift(zmid, prefer="larger").filename