docstring and fixing tests

Author: JoschD

tests/test_doros.py

@@ -1,5 +1,8 @@
+from __future__ import annotations
+
 from datetime import datetime
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import h5py
 import numpy as np
@@ -11,6 +14,9 @@
 from turn_by_turn.doros import DEFAULT_BUNCH_ID, DataKeys, read_tbt, write_tbt
 from turn_by_turn.structures import TbtData, TransverseData
 
+if TYPE_CHECKING:
+    from turn_by_turn.constants import MetaDict
+
 INPUTS_DIR = Path(__file__).parent / "inputs"
 
 
@@ -100,6 +106,9 @@ def _tbt_data() -> TbtData:
     """TbT data for testing. Adding random noise, so that the data is different per BPM."""
     nturns = 2000
     bpms = ["TBPM1", "TBPM2", "TBPM3", "TBPM4"]
+    meta: MetaDict = {
+        "date": datetime.now(),
+    }
 
     return TbtData(
         matrices=[
@@ -126,7 +135,7 @@ def _tbt_data() -> TbtData:
                 ),
             )
         ],
-        date=datetime.now(),
-        bunch_ids=[DEFAULT_BUNCH_ID],
         nturns=nturns,
+        bunch_ids=[DEFAULT_BUNCH_ID],
+        meta=meta,
     )

tests/test_madng.py

@@ -32,7 +32,7 @@ def test_write_ng(_ng_file: Path, tmp_path: Path, example_fake_tbt: TbtData):
 
     new_tbt = read_tbt(from_tbt, datatype="madng")
     compare_tbt(written_tbt, new_tbt, no_binary=True)
-    assert written_tbt.date == new_tbt.date
+    assert written_tbt.meta["date"] == new_tbt.meta["date"]
 
 
 def test_error_ng(_error_file: Path):

tests/test_ptc_trackone.py

@@ -24,8 +24,8 @@ def test_read_ptc_raises_on_invalid_file(_invalid_ptc_file):
 
 def test_read_ptc_defaults_date(_ptc_file_no_date):
     new = ptc.read_tbt(_ptc_file_no_date)
-    assert new.date.day == datetime.today().day
-    assert new.date.tzname() == "UTC"
+    assert new.meta["date"].day == datetime.today().day
+    assert new.meta["date"].tzname() == "UTC"
 
 
 def test_read_ptc_sci(_ptc_file_sci):
@@ -85,7 +85,7 @@ def _original_trackone(track: bool = False) -> TbtData:
             Y=pd.DataFrame(index=names, data=[[0.0011, 0.00077614, -0.00022749, -0.00103188]]),
         ),
     ]
-    return TbtData(matrix, None, [0, 1] if track else [1, 2], 4)
+    return TbtData(matrix, nturns=4, bunch_ids=[0, 1] if track else [1, 2])
 
 
 def _original_simulation_data() -> TbtData:
@@ -112,9 +112,9 @@ def _original_simulation_data() -> TbtData:
             E=pd.DataFrame(index=names, data=[[500.00088, 500.00088, 500.00088, 500.00088]]),
         ),
     ]
-    return TbtData(
-        matrices, date=None, bunch_ids=[0, 1], nturns=4
-    )  # [0, 1] for bunch_ids because it's from tracking
+
+    # [0, 1] for bunch_ids because it's from tracking
+    return TbtData(matrices, bunch_ids=[0, 1], nturns=4)
 
 
 # ----- Fixtures ----- #

tests/test_utils.py

@@ -122,7 +122,6 @@ def test_compare_average_tbtdata():
             )
             for i in range(npart)
         ],
-        date=datetime.now(),
         bunch_ids=range(npart),
         nturns=10,
     )
@@ -142,7 +141,6 @@ def test_compare_average_tbtdata():
                 ),
             )
         ],
-        date=datetime.now(),
         bunch_ids=[1],
         nturns=10,
     )

turn_by_turn/ascii.py

@@ -90,16 +90,16 @@ def _write_header(tbt_data: TbtData, bunch_id: int, output_file: TextIO) -> None
     """
     Write the appropriate headers for a ``TbtData`` object's given bunch_id in the ASCII format.
     """
+    # fmt: off
     output_file.write(f"#{ASCII_ID} v1\n")
-    output_file.write(
-        f"#Created: {datetime.now().strftime('%Y-%m-%d at %H:%M:%S')} By: Python turn_by_turn Package\n"
-    )
+    output_file.write(f"#Created: {datetime.now().strftime(ACQ_DATE_FORMAT)} By: Python turn_by_turn Package\n")
     output_file.write(f"#Number of turns: {tbt_data.nturns}\n")
-    output_file.write(
-        f"#Number of horizontal monitors: {tbt_data.matrices[bunch_id].X.index.size}\n"
-    )
+    output_file.write(f"#Number of horizontal monitors: {tbt_data.matrices[bunch_id].X.index.size}\n")
     output_file.write(f"#Number of vertical monitors: {tbt_data.matrices[bunch_id].Y.index.size}\n")
-    output_file.write(f"#Acquisition date: {tbt_data.date.strftime('%Y-%m-%d at %H:%M:%S')}\n")
+    # fmt: on
+
+    if date := tbt_data.meta.get("date"):
+        output_file.write(f"#Acquisition date: {date.strftime(ACQ_DATE_FORMAT)}\n")
 
 
 def _write_tbt_data(tbt_data: TbtData, bunch_id: int, output_file: TextIO) -> None:

turn_by_turn/iota.py

@@ -11,7 +11,7 @@
 import logging
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal
+from typing import Literal
 
 import h5py
 import numpy as np

turn_by_turn/lhc.py

@@ -18,7 +18,7 @@
 
 from turn_by_turn.ascii import is_ascii_file
 from turn_by_turn.ascii import read_tbt as read_ascii
-from turn_by_turn.constants import PLANE_TO_NUM, PLANES
+from turn_by_turn.constants import PLANE_TO_NUM, PLANES, MetaDict
 from turn_by_turn.structures import TbtData, TransverseData
 from turn_by_turn.utils import matrices_to_array
 
@@ -68,14 +68,21 @@ def read_tbt(file_path: str | Path) -> TbtData:
     bpm_names = sdds_file.values[BPM_NAMES]
     nbpms = len(bpm_names)
     data = {k: sdds_file.values[POSITIONS[k]].reshape((nbpms, nbunches, nturns)) for k in PLANES}
+
     matrices = [
         TransverseData(
             X=pd.DataFrame(index=bpm_names, data=data["X"][:, idx, :], dtype=float),
             Y=pd.DataFrame(index=bpm_names, data=data["Y"][:, idx, :], dtype=float),
         )
         for idx in range(nbunches)
     ]
-    return TbtData(matrices, date, bunch_ids, nturns)
+
+    meta: MetaDict = {
+        "file": file_path,
+        "source_datatype": "lhc",
+        "date": date,
+    }
+    return TbtData(matrices, nturns=nturns, bunch_ids=bunch_ids, meta=meta)
 
 
 def write_tbt(output_path: str | Path, tbt_data: TbtData) -> None:
@@ -101,7 +108,7 @@ def write_tbt(output_path: str | Path, tbt_data: TbtData) -> None:
         sdds.classes.Array(POSITIONS["Y"], "float"),
     ]
     values = [
-        tbt_data.date.timestamp() * 1e9,
+        tbt_data.meta.get("date", datetime.now(tz=tz.tzutc())).timestamp() * 1e9,
         tbt_data.nbunches,
         tbt_data.nturns,
         tbt_data.bunch_ids,

turn_by_turn/madng.py

@@ -23,12 +23,15 @@
 
 import pandas as pd
 
+from turn_by_turn.structures import TbtData, TransverseData
+
 if TYPE_CHECKING:
-    from pathlib import Path  # Only used for type hinting
+    from pathlib import Path
 
     import tfs
 
-from turn_by_turn.structures import TbtData, TransverseData
+    from turn_by_turn.constants import MetaDict
+
 
 LOGGER = logging.getLogger(__name__)
 
@@ -72,7 +75,10 @@ def read_tbt(file_path: str | Path) -> TbtData:
 
     LOGGER.debug("Starting to read TBT data from dataframe")
     df = tfs.read(file_path)
-    return convert_to_tbt(df)
+    tbt_data = convert_to_tbt(df)
+    tbt_data.meta["file"] = file_path
+    tbt_data.meta["source_datatype"] = "madng"
+    return tbt_data
 
 
 def convert_to_tbt(df: pd.DataFrame | tfs.TfsDataFrame) -> TbtData:
@@ -104,6 +110,8 @@ def convert_to_tbt(df: pd.DataFrame | tfs.TfsDataFrame) -> TbtData:
         LOGGER.debug("The 'tfs' package is not installed. Assuming a pandas DataFrame.")
         is_tfs_df = False
 
+    meta: MetaDict = {}
+
     if is_tfs_df:
         date_str = df.headers.get(DATE)
         time_str = df.headers.get(TIME)
@@ -112,11 +120,10 @@ def convert_to_tbt(df: pd.DataFrame | tfs.TfsDataFrame) -> TbtData:
         time_str = df.attrs.get(TIME)
 
     # Combine the date and time into a datetime object
-    date = None
     if date_str and time_str:
-        date = datetime.strptime(f"{date_str} {time_str}", "%d/%m/%y %H:%M:%S")
+        meta["date"] = datetime.strptime(f"{date_str} {time_str}", "%d/%m/%y %H:%M:%S")
     elif date_str:
-        date = datetime.strptime(date_str, "%d/%m/%y")
+        meta["date"] = datetime.strptime(date_str, "%d/%m/%y")
 
     nturns = int(df.iloc[-1].loc[TURN])
     npart = int(df.iloc[-1].loc[PARTICLE_ID])
@@ -161,7 +168,7 @@ def convert_to_tbt(df: pd.DataFrame | tfs.TfsDataFrame) -> TbtData:
         matrices.append(TransverseData(**tracking_data_dict))
 
     LOGGER.debug("Finished reading TBT data")
-    return TbtData(matrices=matrices, bunch_ids=list(particle_ids), nturns=nturns, date=date)
+    return TbtData(matrices=matrices, bunch_ids=list(particle_ids), nturns=nturns, meta=meta)
 
 
 def write_tbt(output_path: str | Path, tbt_data: TbtData) -> None:
@@ -236,8 +243,10 @@ def write_tbt(output_path: str | Path, tbt_data: TbtData) -> None:
     headers = {
         HNAME: "TbtData",
         ORIGIN: "Python",
-        DATE: tbt_data.date.strftime("%d/%m/%y"),
-        TIME: tbt_data.date.strftime("%H:%M:%S"),
         REFCOL: NAME,
     }
+    if date := tbt_data.meta.get("date"):
+        headers[DATE] = date.strftime("%d/%m/%y")
+        headers[TIME] = date.strftime("%H:%M:%S")
+
     tfs.write(output_path, merged_df, headers_dict=headers, save_index=NAME)

turn_by_turn/ptc.py

@@ -13,6 +13,7 @@
 import copy
 import logging
 from collections import namedtuple
+from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -21,7 +22,7 @@
 import pandas as pd
 from dateutil import tz
 
-from turn_by_turn.constants import PLANES
+from turn_by_turn.constants import PLANES, MetaDict
 from turn_by_turn.errors import PTCFormatError
 from turn_by_turn.structures import TbtData, TransverseData
 
@@ -66,21 +67,27 @@ def read_tbt(file_path: str | Path) -> TbtData:
     lines = lines[header_length:]
 
     # parameters
-    bpms, particles, column_indices, n_turns, n_particles = _read_from_first_turn(lines)
+    params = _read_from_first_turn(lines)
 
     # read into dict first for speed then convert to DFs
     matrices = [
-        {p: {bpm: np.zeros(n_turns) for bpm in bpms} for p in PLANES} for _ in range(n_particles)
+        {p: {bpm: np.zeros(params.n_turns) for bpm in params.bpms} for p in PLANES} for _ in range(params.n_particles)
     ]
-    matrices = _read_data(lines, matrices, column_indices)
-    for bunch in range(n_particles):
+    matrices = _read_data(lines, matrices, params.column_indices)
+    for bunch in range(params.n_particles):
         matrices[bunch] = TransverseData(
             X=pd.DataFrame(matrices[bunch]["X"]).transpose(),
             Y=pd.DataFrame(matrices[bunch]["Y"]).transpose(),
         )
 
     LOGGER.debug(f"Read Tbt matrices from: '{file_path.absolute()}'")
-    return TbtData(matrices=matrices, date=date, bunch_ids=particles, nturns=n_turns)
+
+    meta: MetaDict = {
+        "date": date,
+        "file": file_path,
+        "source_datatype": "ptc",
+    }
+    return TbtData(matrices=matrices, nturns=params.n_turns, bunch_ids=params.particles, meta=meta)
 
 
 def _read_header(lines: Sequence[str]) -> tuple[datetime, int]:
@@ -105,19 +112,23 @@ def _read_header(lines: Sequence[str]) -> tuple[datetime, int]:
     return datetime.strptime(f"{date_str[DATE]} {date_str[TIME]}", TIME_FORMAT), idx_line
 
 
-def _read_from_first_turn(
-    lines: Sequence[str],
-) -> tuple[list[str], list[int], dict[Any, Any], int, int]:
+@dataclass(slots=True)
+class TbTParams:
+    """ Parameters read from the first turn of the file. """
+    bpms: list[str] = field(default_factory=list)
+    particles: list[int] = field(default_factory=list)
+    column_indices: dict[Any, Any] | None = None
+    n_turns: int = 0
+    n_particles: int = 0
+
+
+def _read_from_first_turn(lines: Sequence[str]) -> TbTParams:
     """
     Reads the BPMs, particles, column indices and number of turns and particles from the matrices of
     the first turn.
     """
     LOGGER.debug("Reading first turn to define boundary parameters.")
-    bpms = []
-    particles = []
-    column_indices = None
-    n_turns = 0
-    n_particles = 0
+    data = TbTParams()
     first_segment = True
 
     for line in lines:
@@ -126,37 +137,38 @@ def _read_from_first_turn(
             continue
 
         if parts[0] == NAMES:  # read column names
-            if column_indices is not None:
+            if data.column_indices is not None:
                 raise KeyError(f"{NAMES} are defined twice in tbt file!")
-            column_indices = _parse_column_names_to_indices(parts[1:])
+            data.column_indices = _parse_column_names_to_indices(parts[1:])
             continue
 
         if parts[0] == SEGMENTS:  # read segments, append to bunch_id
             segment = Segment(*parts[1:])
             if segment.name == SEGMENT_MARKER[0]:  # start of first segment
-                n_turns = int(segment.turns) - 1
-                n_particles = int(segment.particles)
+                data.n_turns = int(segment.turns) - 1
+                data.n_particles = int(segment.particles)
 
             elif segment.name == SEGMENT_MARKER[1]:  # end of first segment
                 break
 
             else:
                 first_segment = False
-                bpms.append(segment.name)
+                data.bpms.append(segment.name)
 
         elif first_segment:
-            if column_indices is None:
+            if data.column_indices is None:
                 LOGGER.error("Columns not defined in Tbt file")
                 raise PTCFormatError
 
-            new_data = _parse_data(column_indices, parts)
+            new_data = _parse_data(data.column_indices, parts)
             particle = int(float(new_data[COLPARTICLE]))
-            particles.append(particle)
+            data.particles.append(particle)
 
-    if len(particles) == 0:
-        LOGGER.error("No matrices found in TbT file")
-        raise PTCFormatError
-    return bpms, particles, column_indices, n_turns, n_particles
+    if len(data.particles) == 0:
+        msg = "No particles found in TbT file"
+        LOGGER.error(msg)
+        raise PTCFormatError(msg)
+    return data
 
 
 def _read_data(