Repository: European-XFEL/EXtra-data

Issue:
Provide a means to get train timestamps
The new DAQ version adds a timestamp field in the INDEX section, with a timestamp matching each train ID. We should have a way to access these.

Comments:

Patch:
diff --git a/docs/reading_files.rst b/docs/reading_files.rst
index faabbfa..f8b6b4a 100644
--- a/docs/reading_files.rst
+++ b/docs/reading_files.rst
@@ -67,6 +67,8 @@ to refer to all data associated with that 0.1 second window.
 
    .. automethod:: get_data_counts
 
+   .. automethod:: train_timestamps
+
    .. automethod:: info
 
 .. _data-by-source-and-key:
diff --git a/extra_data/reader.py b/extra_data/reader.py
index 916382f..e6d885b 100644
--- a/extra_data/reader.py
+++ b/extra_data/reader.py
@@ -1084,6 +1084,46 @@ class DataCollection:
         print('\tControls')
         [print('\t-', d) for d in sorted(ctrl)] or print('\t-')
 
+    def train_timestamps(self, labelled=False):
+        """Get approximate timestamps for each train
+
+        Timestamps are stored and returned in UTC (not local time).
+        Older files (before format version 1.0) do not have timestamp data,
+        and the returned data in those cases will have the special value NaT
+        (Not a Time).
+
+        If *labelled* is True, they are returned in a pandas series, labelled
+        with train IDs. If False (default), they are returned in a NumPy array
+        of the same length as data.train_ids.
+        """
+        arr = np.zeros(len(self.train_ids), dtype=np.uint64)
+        id_to_ix = {tid: i for (i, tid) in enumerate(self.train_ids)}
+        missing_tids = np.array(self.train_ids)
+        for fa in self.files:
+            tids, file_ixs, _ = np.intersect1d(
+                fa.train_ids, missing_tids, return_indices=True
+            )
+            if not self.inc_suspect_trains:
+                valid = fa.validity_flag[file_ixs]
+                tids, file_ixs = tids[valid], file_ixs[valid]
+            if tids.size == 0 or 'INDEX/timestamp' not in fa.file:
+                continue
+            file_tss = fa.file['INDEX/timestamp'][:]
+            for tid, ts in zip(tids, file_tss[file_ixs]):
+                arr[id_to_ix[tid]] = ts
+
+            missing_tids = np.setdiff1d(missing_tids, tids)
+            if missing_tids.size == 0:  # We've got a timestamp for every train
+                break
+
+        arr = arr.astype('datetime64[ns]')
+        epoch = np.uint64(0).astype('datetime64[ns]')
+        arr[arr == epoch] = 'NaT'  # Not a Time
+        if labelled:
+            import pandas as pd
+            return pd.Series(arr, index=self.train_ids)
+        return arr
+
     def write(self, filename):
         """Write the selected data to a new HDF5 file
 


Test patch:
diff --git a/extra_data/tests/mockdata/base.py b/extra_data/tests/mockdata/base.py
index 0b10de9..ad725cc 100644
--- a/extra_data/tests/mockdata/base.py
+++ b/extra_data/tests/mockdata/base.py
@@ -1,4 +1,4 @@
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 import os.path as osp
 import re
 
@@ -191,7 +191,7 @@ def write_base_index(f, N, first=10000, chunksize=16, format_version='0.5'):
         # timestamps
         ds = f.create_dataset('INDEX/timestamp', (Npad,), 'u8', maxshape=(None,))
         # timestamps are stored as a single uint64 with nanoseconds resolution
-        ts = datetime.utcnow().timestamp() * 10**9
+        ts = datetime.now(tz=timezone.utc).timestamp() * 10**9
         ds[:N] = [ts + i * 10**8 for i in range(N)]
 
     # trainIds
diff --git a/extra_data/tests/test_reader_mockdata.py b/extra_data/tests/test_reader_mockdata.py
index 0863028..2e8544e 100644
--- a/extra_data/tests/test_reader_mockdata.py
+++ b/extra_data/tests/test_reader_mockdata.py
@@ -1,3 +1,4 @@
+from datetime import datetime, timedelta, timezone
 from itertools import islice
 
 import h5py
@@ -588,6 +589,37 @@ def test_select_trains(mock_fxe_raw_run):
         run.select_trains(by_index[[480]])
 
 
+def test_train_timestamps(mock_scs_run):
+    run = RunDirectory(mock_scs_run)
+
+    tss = run.train_timestamps(labelled=False)
+    assert isinstance(tss, np.ndarray)
+    assert tss.shape == (len(run.train_ids),)
+    assert tss.dtype == np.dtype('datetime64[ns]')
+    assert np.all(np.diff(tss).astype(np.uint64) > 0)
+
+    # Convert numpy datetime64[ns] to Python datetime (dropping some precision)
+    dt0 = tss[0].astype('datetime64[ms]').item().replace(tzinfo=timezone.utc)
+    now = datetime.now(timezone.utc)
+    assert dt0 > (now - timedelta(days=1))  # assuming tests take < 1 day to run
+    assert dt0 < now
+
+    tss_ser = run.train_timestamps(labelled=True)
+    assert isinstance(tss_ser, pd.Series)
+    np.testing.assert_array_equal(tss_ser.values, tss)
+    np.testing.assert_array_equal(tss_ser.index, run.train_ids)
+
+
+def test_train_timestamps_nat(mock_fxe_control_data):
+    f = H5File(mock_fxe_control_data)
+    tss = f.train_timestamps()
+    assert tss.shape == (len(f.train_ids),)
+    if f.files[0].format_version == '0.5':
+        assert np.all(np.isnat(tss))
+    else:
+        assert not np.any(np.isnat(tss))
+
+
 def test_union(mock_fxe_raw_run):
     run = RunDirectory(mock_fxe_raw_run)