added meta, some classes still missing

Author: JoschD

tests/test_iota.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from pathlib import Path
 
 import h5py
 import numpy as np
@@ -11,9 +12,9 @@
 from turn_by_turn.structures import TbtData, TransverseData
 
 
-def test_tbt_read_hdf5(_hdf5_file):
+def test_tbt_read_hdf5(_hdf5_file_v1):
     origin = _hdf5_file_content()
-    new = iota.read_tbt(_hdf5_file, hdf5_version=1)
+    new = iota.read_tbt(_hdf5_file_v1, version=1)
     compare_tbt(origin, new, no_binary=False)
 
 
@@ -23,9 +24,9 @@ def test_tbt_read_hdf5_v2(_hdf5_file_v2):
     compare_tbt(origin, new, no_binary=False)
 
 
-def test_tbt_raises_on_wrong_hdf5_version(_hdf5_file):
+def test_tbt_raises_on_wrong_hdf5_version(_hdf5_file_v1):
     with pytest.raises(HDF5VersionError):
-        iota.read_tbt(_hdf5_file, hdf5_version=2)
+        iota.read_tbt(_hdf5_file_v1, version=2)
 
 
 def _hdf5_file_content() -> TbtData:
@@ -45,15 +46,16 @@ def _hdf5_file_content() -> TbtData:
                 ),
             )
         ],
-        date=datetime.now(),
         bunch_ids=[1],
         nturns=2000,
     )
 
 
 @pytest.fixture()
-def _hdf5_file(tmp_path) -> h5py.File:
+def _hdf5_file_v1(tmp_path) -> Path:
     """IOTA File standard."""
+    content: TransverseData = _hdf5_file_content().matrices[0]
+
     with h5py.File(tmp_path / "test_file.hdf5", "w") as hd5_file:
         hd5_file.create_dataset(
             "N:IBE2RH",
@@ -80,11 +82,11 @@ def _hdf5_file(tmp_path) -> h5py.File:
             "N:IBA1CS",
             data=create_data(np.linspace(0, 20, 2000, endpoint=False), 1, np.exp).flatten(),
         )
-    yield tmp_path / "test_file.hdf5"
+    return tmp_path / "test_file.hdf5"
 
 
 @pytest.fixture()
-def _hdf5_file_v2(tmp_path) -> h5py.File:
+def _hdf5_file_v2(tmp_path) -> Path:
     """IOTA File standard."""
     with h5py.File(tmp_path / "test_file_v2.hdf5", "w") as hd5_file:
         hd5_file.create_group("A1C")
@@ -114,4 +116,4 @@ def _hdf5_file_v2(tmp_path) -> h5py.File:
             "Intensity",
             data=create_data(np.linspace(0, 20, 2000, endpoint=False), 1, np.exp).flatten(),
         )
-    yield tmp_path / "test_file_v2.hdf5"
+    return tmp_path / "test_file_v2.hdf5"

turn_by_turn/ascii.py

@@ -23,7 +23,13 @@
 import pandas as pd
 from dateutil import tz
 
-from turn_by_turn.constants import FORMAT_STRING, NUM_TO_PLANE, PLANE_TO_NUM, PLANES
+from turn_by_turn.constants import (
+    FORMAT_STRING,
+    NUM_TO_PLANE,
+    PLANE_TO_NUM,
+    PLANES,
+    MetaDict,
+)
 from turn_by_turn.structures import TbtData, TransverseData
 
 LOGGER = logging.getLogger(__name__)
@@ -108,7 +114,7 @@ def _write_tbt_data(tbt_data: TbtData, bunch_id: int, output_file: TextIO) -> No
 # ----- Reader ----- #
 
 
-def read_tbt(file_path: str | Path, bunch_id: int = None) -> TbtData:
+def read_tbt(file_path: str | Path, bunch_id: int | None = None) -> TbtData:
     """
     Reads turn-by-turn data from an ASCII turn-by-turn format file, and return the date as well as
     parsed matrices for construction of a ``TbtData`` object.
@@ -125,7 +131,10 @@ def read_tbt(file_path: str | Path, bunch_id: int = None) -> TbtData:
     data_lines = Path(file_path).read_text().splitlines()
     bpm_names = {"X": [], "Y": []}
     bpm_data = {"X": [], "Y": []}
-    date = None  # will switch to TbtData.date's default if not found in file
+    meta: MetaDict = {
+        "file": file_path,
+        "source_datatype": "ascii",
+    }
 
     if bunch_id is None:
         bunch_id = _parse_bunch_id(file_path)
@@ -135,7 +144,7 @@ def read_tbt(file_path: str | Path, bunch_id: int = None) -> TbtData:
 
         if ACQ_DATE_PREFIX in line:
             LOGGER.debug("Acquiring date from file")
-            date = _parse_date(line)
+            meta["date"] = _parse_date(line)
             continue
 
         if line == "" or line.startswith(ASCII_COMMENT):  # empty or comment line
@@ -159,7 +168,7 @@ def read_tbt(file_path: str | Path, bunch_id: int = None) -> TbtData:
         )
     ]
     return TbtData(
-        matrices=matrices, date=date, bunch_ids=[bunch_id], nturns=matrices[0].X.shape[1]
+        matrices=matrices, meta=meta, bunch_ids=[bunch_id], nturns=matrices[0].X.shape[1]
     )
 
 

turn_by_turn/constants.py

@@ -7,10 +7,34 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, TypedDict
+
+if TYPE_CHECKING:
+    from datetime import datetime
+    from pathlib import Path
+
 PLANES: tuple[str, str] = ("X", "Y")
 NUM_TO_PLANE: dict[str, str] = {"0": "X", "1": "Y"}
 PLANE_TO_NUM: dict[str, int] = {"X": 0, "Y": 1}
 
+# ----- Common Meta Keys ----- #
+
+class MetaDict(TypedDict, total=False):
+    """ Metadata dictionary, to type-hint known entries.
+    None of the entries are required (``total=False``).
+
+    Attributes:
+        date (datetime): Date of the measurement/creation of the data
+        file (Path | str): Path to the file the data was loaded from (if available).
+        machine (str): Name of the machine the data was measured/simulated on.
+        source_datatype (str): The datatype this data was loaded from.
+        comment (str): Any comment on the measurement.
+    """
+    date: datetime
+    file: Path | str
+    machine: str
+    source_datatype: str
+    comment: str
 
 # ----- Miscellaneous ----- #
 

turn_by_turn/doros.py

@@ -38,6 +38,7 @@
 import logging
 from datetime import datetime
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import h5py
 import pandas as pd
@@ -46,6 +47,9 @@
 from turn_by_turn.structures import TbtData, TransverseData
 from turn_by_turn.utils import all_elements_equal
 
+if TYPE_CHECKING:
+    from turn_by_turn.constants import MetaDict
+
 LOGGER = logging.getLogger(__name__)
 
 DEFAULT_BUNCH_ID: int = 0  # bunch ID not saved in the DOROS file
@@ -132,6 +136,10 @@ def read_tbt(
     file_path = Path(file_path)
     LOGGER.debug(f"Reading DOROS {data_type} data at path: '{file_path.absolute()}'")
     data_keys = DataKeys.get_data_keys(data_type)
+    meta: MetaDict = {
+        "file": file_path,
+        "source_datatype": data_type,
+    }
 
     with h5py.File(file_path, "r") as hdf_file:
         # use "/" to keep track of bpm order, see https://github.com/h5py/h5py/issues/1471
@@ -141,7 +149,7 @@ def read_tbt(
         _check_data_lengths(hdf_file, data_keys, bpm_names)
 
         time_stamps = [hdf_file[bpm][ACQ_STAMP][0] for bpm in bpm_names]
-        date = datetime.fromtimestamp(min(time_stamps) / 1e6, tz=tz.tzutc())
+        meta["date"] = datetime.fromtimestamp(min(time_stamps) / 1e6, tz=tz.tzutc())
 
         nturns = hdf_file[bpm_names[0]][data_keys.n_samples][0]  # equal lengths checked before
         matrices = [
@@ -150,7 +158,7 @@ def read_tbt(
                 Y=_create_dataframe(hdf_file, data_keys, bpm_names, plane="Y"),
             )
         ]
-    return TbtData(matrices, date, [bunch_id], nturns)
+    return TbtData(matrices, bunch_ids=[bunch_id], nturns=nturns, meta=meta)
 
 
 def write_tbt(
@@ -172,13 +180,15 @@ def write_tbt(
     data_keys = DataKeys.get_data_keys(data_type)
     other_keys = DataKeys.get_other_data_keys(data_type)
 
+    timestamp = tbt_data.meta.get("date", datetime.now(tz=tz.tzutc())).timestamp() * 1e6
+
     data = tbt_data.matrices[0]
     with h5py.File(file_path, "w", track_order=True) as hdf_file:
         hdf_file.create_group(METADATA)
         for bpm in data.X.index:
             hdf_file.create_group(bpm)
-            hdf_file[bpm].create_dataset(ACQ_STAMP, data=[tbt_data.date.timestamp() * 1e6])
-            hdf_file[bpm].create_dataset(BST_TIMESTAMP, data=[tbt_data.date.timestamp() * 1e6])
+            hdf_file[bpm].create_dataset(ACQ_STAMP, data=[timestamp])
+            hdf_file[bpm].create_dataset(BST_TIMESTAMP, data=[timestamp])
 
             hdf_file[bpm].create_dataset(data_keys.n_samples, data=[tbt_data.nturns])
             hdf_file[bpm].create_dataset(data_keys.data["X"], data=data.X.loc[bpm, :].values)

turn_by_turn/esrf.py

@@ -38,7 +38,10 @@ def read_tbt(file_path: str | Path) -> TbtData:
     file_path = Path(file_path)
     LOGGER.debug(f"Reading ESRF file at path: '{file_path.absolute()}'")
     names, matrix = load_esrf_mat_file(file_path)
-    return numpy_to_tbt(names, matrix)
+    tbt_data = numpy_to_tbt(names, matrix)
+    tbt_data.meta["file"] = file_path
+    tbt_data.meta["source_datatype"] = "esrf"
+    return tbt_data
 
 
 def load_esrf_mat_file(infile: str | Path) -> tuple[np.ndarray, np.ndarray]: