updated iota tests

Author: JoschD

tests/test_iota.py

@@ -1,4 +1,3 @@
-from datetime import datetime
 from pathlib import Path
 
 import h5py
@@ -12,36 +11,39 @@
 from turn_by_turn.structures import TbtData, TransverseData
 
 
-def test_tbt_read_hdf5(_hdf5_file_v1):
-    origin = _hdf5_file_content()
+def test_tbt_read_hdf5(_hdf5_file_v1, _hdf5_file_content):
     new = iota.read_tbt(_hdf5_file_v1, version=1)
-    compare_tbt(origin, new, no_binary=False)
+    compare_tbt(_hdf5_file_content, new, no_binary=False)
 
 
-def test_tbt_read_hdf5_v2(_hdf5_file_v2):
-    origin = _hdf5_file_content()
+def test_tbt_read_hdf5_v2(_hdf5_file_v2, _hdf5_file_content):
     new = iota.read_tbt(_hdf5_file_v2)
-    compare_tbt(origin, new, no_binary=False)
+    compare_tbt(_hdf5_file_content, new, no_binary=False)
 
 
-def test_tbt_raises_on_wrong_hdf5_version(_hdf5_file_v1):
+def test_tbt_raises_on_wrong_hdf5_version(_hdf5_file_v1, _hdf5_file_v2):
     with pytest.raises(HDF5VersionError):
         iota.read_tbt(_hdf5_file_v1, version=2)
 
+    with pytest.raises(HDF5VersionError):
+        iota.read_tbt(_hdf5_file_v2, version=1)
+
 
+
+@pytest.fixture(scope="module")
 def _hdf5_file_content() -> TbtData:
     """TbT data as had been written out to hdf5 files (see below)."""
     return TbtData(
         matrices=[
             TransverseData(
                 X=pd.DataFrame(
                     index=["IBPMA1C", "IBPME2R"],
-                    data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 2, np.sin),
+                    data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 2, np.sin, noise=0.02),
                     dtype=float,
                 ),
                 Y=pd.DataFrame(
                     index=["IBPMA1C", "IBPME2R"],
-                    data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 2, np.cos),
+                    data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 2, np.cos, noise=0.015),
                     dtype=float,
                 ),
             )
@@ -52,18 +54,18 @@ def _hdf5_file_content() -> TbtData:
 
 
 @pytest.fixture()
-def _hdf5_file_v1(tmp_path) -> Path:
-    """IOTA File standard."""
-    content: TransverseData = _hdf5_file_content().matrices[0]
+def _hdf5_file_v1(tmp_path, _hdf5_file_content) -> Path:
+    """IOTA File v1 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",
-            data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 1, np.sin).flatten(),
+            data=content.X.loc["IBPME2R"].to_numpy(),
         )
         hd5_file.create_dataset(
             "N:IBE2RV",
-            data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 1, np.cos).flatten(),
+            data=content.Y.loc["IBPME2R"].to_numpy(),
         )
         hd5_file.create_dataset(
             "N:IBE2RS",
@@ -72,11 +74,11 @@ def _hdf5_file_v1(tmp_path) -> Path:
 
         hd5_file.create_dataset(
             "N:IBA1CH",
-            data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 1, np.sin).flatten(),
+            data=content.X.loc["IBPMA1C"].to_numpy(),
         )
         hd5_file.create_dataset(
             "N:IBA1CV",
-            data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 1, np.cos).flatten(),
+            data=content.Y.loc["IBPMA1C"].to_numpy(),
         )
         hd5_file.create_dataset(
             "N:IBA1CS",
@@ -86,17 +88,19 @@ def _hdf5_file_v1(tmp_path) -> Path:
 
 
 @pytest.fixture()
-def _hdf5_file_v2(tmp_path) -> Path:
-    """IOTA File standard."""
+def _hdf5_file_v2(tmp_path, _hdf5_file_content) -> Path:
+    """IOTA File v2 standard."""
+    content: TransverseData = _hdf5_file_content.matrices[0]
+
     with h5py.File(tmp_path / "test_file_v2.hdf5", "w") as hd5_file:
         hd5_file.create_group("A1C")
         hd5_file["A1C"].create_dataset(
             "Horizontal",
-            data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 1, np.sin).flatten(),
+            data=content.X.loc["IBPMA1C"].to_numpy(),
         )
         hd5_file["A1C"].create_dataset(
             "Vertical",
-            data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 1, np.cos).flatten(),
+            data=content.Y.loc["IBPMA1C"].to_numpy(),
         )
         hd5_file["A1C"].create_dataset(
             "Intensity",
@@ -106,11 +110,11 @@ def _hdf5_file_v2(tmp_path) -> Path:
         hd5_file.create_group("E2R")
         hd5_file["E2R"].create_dataset(
             "Horizontal",
-            data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 1, np.sin).flatten(),
+            data=content.X.loc["IBPME2R"].to_numpy(),
         )
         hd5_file["E2R"].create_dataset(
             "Vertical",
-            data=create_data(np.linspace(-np.pi, np.pi, 2000, endpoint=False), 1, np.cos).flatten(),
+            data=content.Y.loc["IBPME2R"].to_numpy(),
         )
         hd5_file["E2R"].create_dataset(
             "Intensity",

tests/test_lhc_and_general.py

@@ -78,8 +78,8 @@ def compare_tbt(
                 assert np.all(origin_mat == new_mat)
 
 
-def create_data(phases, nbpm, function, noise: float = 0) -> np.ndarray:
-    rng = np.random.default_rng()
+def create_data(phases, nbpm, function, noise: float = 0, seed: int = None) -> np.ndarray:
+    rng = np.random.default_rng(seed=seed)
     return np.ones((nbpm, len(phases))) * function(phases) + noise * rng.standard_normal(
         size=(nbpm, len(phases))
     )