Merge pull request #20 from pylhc/doros_oscillations_positions

DOROS: oscillations and positions

Author: JoschD

tests/inputs/doros_shrink_data.py

@@ -6,10 +6,10 @@
 import subprocess
 import h5py
 
-from turn_by_turn.doros import DEFAULT_OSCILLATION_DATA, N_ORBIT_SAMPLES, N_OSCILLATION_SAMPLES, OSCILLATIONS, POSITIONS
+from turn_by_turn.doros import DataKeys
 
 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"
+file_out = "./tests/inputs/test_doros.h5"
 
 # Values need to coincide with the values in the test_doros.py
 N_BPMS= 3
@@ -18,26 +18,25 @@
 # 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)
+
+data_keys_list = [DataKeys.get_data_keys(datatype) for datatype in DataKeys.types()]
+
 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()]
+
+    bpms = [name for name in hdf_file["/"].keys() if data_keys_list[0].n_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 data_keys in data_keys_list:
+            del hdf_file[bpm][data_keys.n_samples]
+            hdf_file[bpm].create_dataset(data_keys.n_samples, data=[NTURNS])
+            data_x = hdf_file[bpm][data_keys.data["X"]][:NTURNS]
+            data_y = hdf_file[bpm][data_keys.data["Y"]][:NTURNS]
+
+            del hdf_file[bpm][data_keys.data["X"]]
+            del hdf_file[bpm][data_keys.data["Y"]]
+            hdf_file[bpm].create_dataset(data_keys.data["X"], data=data_x)
+            hdf_file[bpm].create_dataset(data_keys.data["Y"], data=data_y)
+
 
     for bpm in bpms[N_BPMS:]:
         del hdf_file[bpm]

tests/inputs/test_doros.h5

@@ -2,81 +2,96 @@
 from datetime import datetime
 from pathlib import Path
 
+import h5py
 import numpy as np
 import pandas as pd
 import pytest
-import h5py
 
+import turn_by_turn as tbt
+from tests.test_lhc_and_general import compare_tbt, create_data
+from turn_by_turn.doros import DEFAULT_BUNCH_ID, DataKeys, read_tbt, write_tbt
 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"
 
-@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)
+@pytest.mark.parametrize("datatype", DataKeys.types())
+def test_read_write_real_data(tmp_path, datatype):
+    tbt_data = read_tbt(INPUTS_DIR / "test_doros.h5", bunch_id=10, data_type=datatype)
 
-    assert tbt.nbunches == 1
-    assert len(tbt.matrices) == 1
-    assert tbt.nturns == 50000
-    assert tbt.matrices[0].X.shape == (3, tbt.nturns)
-    assert tbt.matrices[0].Y.shape == (3, tbt.nturns)
-    assert len(set(tbt.matrices[0].X.index)) == 3
-    assert np.all(tbt.matrices[0].X.index == tbt.matrices[0].Y.index)
+    assert tbt_data.nbunches == 1
+    assert len(tbt_data.matrices) == 1
+    assert tbt_data.nturns == 50000
+    assert tbt_data.matrices[0].X.shape == (3, tbt_data.nturns)
+    assert tbt_data.matrices[0].Y.shape == (3, tbt_data.nturns)
+    assert len(set(tbt_data.matrices[0].X.index)) == 3
+    assert np.all(tbt_data.matrices[0].X.index == tbt_data.matrices[0].Y.index)
 
     file_path = tmp_path / "test_file.h5"
-    write_tbt(tbt, file_path)
-    new = read_tbt(file_path, bunch_id=10)
-    compare_tbt(tbt, new, no_binary=False)
+    write_tbt(file_path, tbt_data, data_type=datatype)
+    new = read_tbt(file_path, bunch_id=10, data_type=datatype)
+    compare_tbt(tbt_data, new, no_binary=False)
 
 
-def test_write_read(tmp_path):
-    tbt = _tbt_data()
+@pytest.mark.parametrize("datatype", DataKeys.types())
+def test_write_read(tmp_path, datatype):
+    tbt_data = _tbt_data()
     file_path = tmp_path / "test_file.h5"
-    write_tbt(tbt, file_path)
-    new = read_tbt(file_path)
-    compare_tbt(tbt, new, no_binary=False)
+    write_tbt(file_path, tbt_data, data_type=datatype)
+    new = read_tbt(file_path, data_type=datatype)
+    compare_tbt(tbt_data, new, no_binary=False)
+
+
+@pytest.mark.parametrize("datatype", ["doros_oscillations", "doros_positions"])
+def test_write_read_via_io_module(tmp_path, datatype):
+    tbt_data = _tbt_data()
+    file_path = tmp_path / "test_file.h5"
+    tbt.write(file_path, tbt_data, datatype=datatype)
+    new = tbt.read(file_path, datatype=datatype)
+    compare_tbt(tbt_data, new, no_binary=False)
 
 
 def test_read_raises_different_bpm_lengths(tmp_path):
-    tbt = _tbt_data()
+    tbt_data = _tbt_data()
     file_path = tmp_path / "test_file.h5"
-    write_tbt(tbt, file_path)
+    data_type = DataKeys.OSCILLATIONS
+    write_tbt(file_path, tbt_data, data_type=data_type)
+    keys = DataKeys.get_data_keys(data_type)
 
-    bpm = tbt.matrices[0].X.index[0]
+    bpm = tbt_data.matrices[0].X.index[0]
 
     # modify the BPM lengths in the file
     with h5py.File(file_path, "r+") as h5f:
         delta = 10
-        del h5f[bpm][N_ORBIT_SAMPLES]
-        h5f[bpm][N_ORBIT_SAMPLES] = [tbt.matrices[0].X.shape[1] - delta]
-        for key in POSITIONS.values():
+        del h5f[bpm][keys.n_samples]
+        h5f[bpm][keys.n_samples] = [tbt_data.matrices[0].X.shape[1] - delta]
+        for key in keys.data.values():
             data = h5f[bpm][key][:-delta]
             del h5f[bpm][key]
             h5f[bpm][key] = data
 
     with pytest.raises(ValueError) as e:
-        read_tbt(file_path)
+        read_tbt(file_path, data_type=DataKeys.OSCILLATIONS)
     assert "Not all BPMs have the same number of turns!" in str(e)
 
 
 def test_read_raises_on_different_bpm_lengths_in_data(tmp_path):
-    tbt = _tbt_data()
+    tbt_data = _tbt_data()
     file_path = tmp_path / "test_file.h5"
-    write_tbt(tbt, file_path)
+    data_type = DataKeys.OSCILLATIONS
+    keys = DataKeys.get_data_keys(data_type)
+
+    write_tbt(file_path, tbt_data, data_type=data_type)
 
-    bpms = [tbt.matrices[0].X.index[i] for i in (0, 2)]
+    bpms = [tbt_data.matrices[0].X.index[i] for i in (0, 2)]
 
     # modify the BPM lengths in the file
     with h5py.File(file_path, "r+") as h5f:
         for bpm in bpms:
-            del h5f[bpm][N_ORBIT_SAMPLES]
-            h5f[bpm][N_ORBIT_SAMPLES] = [tbt.matrices[0].X.shape[1] + 10]
+            del h5f[bpm][keys.n_samples]
+            h5f[bpm][keys.n_samples] = [tbt_data.matrices[0].X.shape[1] + 10]
 
     with pytest.raises(ValueError) as e:
-        read_tbt(file_path)
+        read_tbt(file_path, data_type=data_type)
     assert "Found BPMs with different data lengths" in str(e)
     assert all(bpm in str(e) for bpm in bpms)
 

tests/inputs/test_doros_2024-09-29.h5

@@ -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.1"
+__version__ = "0.7.2"
 __author__ = "pylhc"
 __author_email__ = "[email protected]"
 __license__ = "MIT"