Merge pull request #4717 from cphyc/find-by-key

[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