Merge pull request #19 from pylhc/doros_update

Update Doros for new entries

Author: JoschD

tests/inputs/doros_shrink_data.py

@@ -0,0 +1,47 @@
+""" 
+Quick script to remove some data from the DOROS files, to make them smaller for the tests.
+"""
+from pathlib import Path
+import shutil
+import subprocess
+import h5py
+
+from turn_by_turn.doros import DEFAULT_OSCILLATION_DATA, N_ORBIT_SAMPLES, N_OSCILLATION_SAMPLES, OSCILLATIONS, POSITIONS
+
+file_in = "/afs/cern.ch/work/j/jdilly/dorostest/DOROS-2024-09-29_01_37_13_522358-NO_USER.h5"
+file_out = "./tests/inputs/test_doros_2024-09-29.h5"
+
+# Values need to coincide with the values in the test_doros.py
+N_BPMS= 3
+NTURNS = 50000
+
+# Copy file to temporary location and "delete" unneccessary data (only removes references to data)
+file_temp = file_out + ".tmp"
+shutil.copyfile(file_in, file_temp)
+with h5py.File(file_temp, "r+", track_order=True) as hdf_file:
+    bpms = [name for name in hdf_file["/"].keys() if N_ORBIT_SAMPLES in hdf_file[f"/{name}"].keys()]
+
+    for bpm in bpms[:N_BPMS]:
+        del hdf_file[bpm][N_ORBIT_SAMPLES]
+        hdf_file[bpm].create_dataset(N_ORBIT_SAMPLES, data=[NTURNS])
+        data_x = hdf_file[bpm][POSITIONS["X"]][:NTURNS]
+        data_y = hdf_file[bpm][POSITIONS["Y"]][:NTURNS]
+
+        del hdf_file[bpm][POSITIONS["X"]]
+        del hdf_file[bpm][POSITIONS["Y"]]
+        hdf_file[bpm].create_dataset(POSITIONS["X"], data=data_x)
+        hdf_file[bpm].create_dataset(POSITIONS["Y"], data=data_y)
+
+        del hdf_file[bpm][N_OSCILLATION_SAMPLES]
+        del hdf_file[bpm][OSCILLATIONS["X"]]
+        del hdf_file[bpm][OSCILLATIONS["Y"]]
+        hdf_file[bpm].create_dataset(N_OSCILLATION_SAMPLES, data=0)
+        hdf_file[bpm].create_dataset(OSCILLATIONS["X"], data=[DEFAULT_OSCILLATION_DATA])
+        hdf_file[bpm].create_dataset(OSCILLATIONS["Y"], data=[DEFAULT_OSCILLATION_DATA])
+    
+    for bpm in bpms[N_BPMS:]:
+        del hdf_file[bpm]
+
+# Repack file (which actually removes the data in the file)
+subprocess.run(["h5repack", file_temp, file_out])
+Path(file_temp).unlink()

tests/inputs/test_doros_2024-09-29.h5

@@ -7,18 +7,16 @@
 import pytest
 import h5py
 
-from turn_by_turn.constants import PRINT_PRECISION
-from turn_by_turn.errors import DataTypeError
 from turn_by_turn.structures import TbtData, TransverseData
 from tests.test_lhc_and_general import create_data, compare_tbt
 
 from turn_by_turn.doros import N_ORBIT_SAMPLES, read_tbt, write_tbt, DEFAULT_BUNCH_ID, POSITIONS
 
 INPUTS_DIR = Path(__file__).parent / "inputs"
 
-
-def test_read_write_real_data(tmp_path):
-    tbt = read_tbt(INPUTS_DIR / "test_doros.h5", bunch_id=10)
+@pytest.mark.parametrize("filename", ["test_doros.h5", "test_doros_2024-09-29.h5"])
+def test_read_write_real_data(tmp_path, filename):
+    tbt = read_tbt(INPUTS_DIR / filename, bunch_id=10)
 
     assert tbt.nbunches == 1
     assert len(tbt.matrices) == 1

tests/test_doros.py

@@ -5,7 +5,7 @@
 __title__ = "turn_by_turn"
 __description__ = "Read and write turn-by-turn measurement files from different particle accelerator formats."
 __url__ = "https://github.com/pylhc/turn_by_turn"
-__version__ = "0.7.0"
+__version__ = "0.7.1"
 __author__ = "pylhc"
 __author_email__ = "pylhc@github.com"
 __license__ = "MIT"

turn_by_turn/__init__.py

@@ -50,6 +50,7 @@
 DEFAULT_BUNCH_ID: int = 0  # bunch ID not saved in the DOROS file 
 
 METADATA: str = "METADATA"
+BPM_NAME_END: str = "_DOROS"
 
 # tiemstamps
 BST_TIMESTAMP: str = "bstTimestamp"   # microseconds
@@ -85,7 +86,8 @@ def read_tbt(file_path: str|Path, bunch_id: int = DEFAULT_BUNCH_ID) -> TbtData:
     LOGGER.debug(f"Reading DOROS file at path: '{file_path.absolute()}'")
     with h5py.File(file_path, "r") as hdf_file:
         # use "/" to keep track of bpm order, see https://github.com/h5py/h5py/issues/1471
-        bpm_names = [name for name in hdf_file["/"].keys() if name != METADATA]
+        bpm_names = [name for name in hdf_file["/"].keys() if N_ORBIT_SAMPLES in hdf_file[f"/{name}"].keys()]
+        LOGGER.debug(f"Found BPMs in DOROS-type file: {bpm_names}")
 
         _check_data_lengths(hdf_file, bpm_names)