Merge pull request #16 from pylhc/sps_remove_trailing_planes

remove trailing planes from sps data

Author: JoschD

tests/test_sps.py

@@ -8,11 +8,12 @@
 from turn_by_turn import sps, TbtData, TransverseData
 
 
-def test_read_write_real_data(_sps_file, tmp_path):
-    input_sdds = sps.read_tbt(_sps_file)
+@pytest.mark.parametrize("remove_planes", [True, False])
+def test_read_write_real_data(_sps_file, tmp_path, remove_planes):
+    input_sdds = sps.read_tbt(_sps_file, remove_trailing_bpm_plane=remove_planes)
     tmp_sdds = tmp_path / "sps.sdds"
-    sps.write_tbt(tmp_sdds, input_sdds)
-    read_sdds = sps.read_tbt(tmp_sdds)
+    sps.write_tbt(tmp_sdds, input_sdds, add_trailing_bpm_plane=remove_planes)
+    read_sdds = sps.read_tbt(tmp_sdds, remove_trailing_bpm_plane=remove_planes)
     compare_tbt(input_sdds, read_sdds, no_binary=True)
 
 
@@ -25,21 +26,44 @@ def test_write_read(tmp_path):
         matrices=[
             TransverseData(
                 X=pd.DataFrame(
-                    index=[f"BPM{i}.H" for i in range(nbpms_x)],
+                    index=[f"BPMH{i}.H" for i in range(nbpms_x)],
                     data=create_data(np.linspace(-np.pi, np.pi, nturns, endpoint=False), nbpms_x, np.sin)
                 ),
                 Y=pd.DataFrame(
-                    index=[f"BPM{i}.V" for i in range(nbpms_y)],
+                    index=[f"BPMV{i}.V" for i in range(nbpms_y)],
                     data=create_data(np.linspace(-np.pi, np.pi, nturns, endpoint=False), nbpms_y, np.cos)
                 ),
             )
         ],
     )
     tmp_sdds = tmp_path / "sps_fake_data.sdds"
-    sps.write_tbt(tmp_sdds, original)
-    read_sdds = sps.read_tbt(tmp_sdds)
+    # Normal read/write test
+    sps.write_tbt(tmp_sdds, original, add_trailing_bpm_plane=False) 
+    read_sdds = sps.read_tbt(tmp_sdds, remove_trailing_bpm_plane=False)
     compare_tbt(original, read_sdds, no_binary=True)
 
+    # Test no name changes when writing and planes already present
+    sps.write_tbt(tmp_sdds, original, add_trailing_bpm_plane=True) 
+    read_sdds = sps.read_tbt(tmp_sdds, remove_trailing_bpm_plane=False)
+    compare_tbt(original, read_sdds, no_binary=True)
+    
+    # Test plane removal on reading
+    read_sdds = sps.read_tbt(tmp_sdds, remove_trailing_bpm_plane=True)
+    assert not any(read_sdds.matrices[0].X.index.str.endswith(".H"))
+    assert not any(read_sdds.matrices[0].Y.index.str.endswith(".V"))
+    
+    # Test planes stay off when writing
+    sps.write_tbt(tmp_sdds, read_sdds, add_trailing_bpm_plane=False)
+    read_sdds = sps.read_tbt(tmp_sdds, remove_trailing_bpm_plane=False)
+    assert not any(read_sdds.matrices[0].X.index.str.endswith(".H"))
+    assert not any(read_sdds.matrices[0].Y.index.str.endswith(".V"))
+
+    # Test adding planes again
+    sps.write_tbt(tmp_sdds, read_sdds, add_trailing_bpm_plane=True) 
+    read_sdds = sps.read_tbt(tmp_sdds, remove_trailing_bpm_plane=False)
+    compare_tbt(original, read_sdds, no_binary=True)
+
+
 
 @pytest.fixture()
 def _sps_file() -> Path:

turn_by_turn/__init__.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.5.0"
+__version__ = "0.6.0"
 __author__ = "pylhc"
 __author_email__ = "pylhc@github.com"
 __license__ = "MIT"

turn_by_turn/ascii.py

@@ -7,12 +7,13 @@
 containing the columns:
 - Plane (0 for horizontal, 1 for vertical)
 - Observation point (i.e. BPM name)
-- BPM index,
+- BPM index/longitunial location
 - Value Turn 1, Turn 2, etc.
 """
 import logging
 from datetime import datetime
 from pathlib import Path
+import re
 from typing import TextIO, Union, Tuple, List, Optional
 
 import numpy as np
@@ -98,23 +99,27 @@ def _write_tbt_data(tbt_data: TbtData, bunch_id: int, output_file: TextIO) -> No
 
 # ----- Reader ----- #
 
-def read_tbt(
-        file_path: Union[str, Path]
-) -> Tuple[List[TransverseData], Optional[datetime]]:
+def read_tbt(file_path: Union[str, Path], bunch_id: int = 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.
 
     Args:
         file_path (Union[str, Path]): path to the turn-by-turn measurement file.
+        bunch_id (int, optional): the bunch id associated with this file. 
+                                  Defaults to `None`, but is then attempted to parsed
+                                  from the filename. If not found, `0` is used.
 
     Returns:
-        Turn-by-turn data matrices and
+        Turn-by-turn data
     """
     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
+    
+    if bunch_id is None:
+        bunch_id = _parse_bunch_id(file_path)
 
     for line in data_lines:
         line = line.strip()
@@ -144,17 +149,18 @@ def read_tbt(
             Y=pd.DataFrame(index=bpm_names["Y"], data=np.array(bpm_data["Y"])),
         )
     ]
-    return matrices, date
+    return TbtData(matrices=matrices, date=date, bunch_ids=[bunch_id], nturns=matrices[0].X.shape[1])
 
 
 # ----- Helpers ----- #
 
+
 def _parse_samples(line: str) -> Tuple[str, str, np.ndarray]:
     """Parse a line into its different elements."""
     parts = line.split()
     plane_num = parts[0]
     bpm_name = parts[1]
-    # bunch_id = part[2]  # not used, comment for clarification
+    # bpm_location = part[2]  # not used, comment for clarification
     bpm_samples = np.array([float(part) for part in parts[3:]])
     return plane_num, bpm_name, bpm_samples
 
@@ -172,6 +178,17 @@ def _parse_date(line: str) -> datetime:
         return datetime.today().replace(tzinfo=tz.tzutc())
 
 
+def _parse_bunch_id(file_path: Path) -> int:
+    """Parse the bunch_id from the filename."""
+    bunch_id_match = re.match(r".*_(?P<bunch_id>\d+)(.sdds)?$", file_path.name)
+    if bunch_id_match:
+        try:
+            return int(bunch_id_match.group("bunch_id"))
+        except ValueError:
+            pass
+    return 0
+
+
 # For backwards compatibility <0.4.2:
 write_ascii = write_tbt
 read_ascii = read_tbt

turn_by_turn/lhc.py

@@ -51,8 +51,7 @@ def read_tbt(file_path: Union[str, Path]) -> TbtData:
     LOGGER.debug(f"Reading LHC file at path: '{file_path.absolute()}'")
 
     if is_ascii_file(file_path):
-        matrices, date = read_ascii(file_path)
-        return TbtData(matrices, date, [0], matrices[0].X.shape[1])
+        return read_ascii(file_path)
 
     sdds_file = sdds.read(file_path)
     nbunches = sdds_file.values[N_BUNCHES]

turn_by_turn/sps.py

@@ -7,6 +7,7 @@
 import logging
 from datetime import datetime
 from pathlib import Path
+import re
 from typing import Union
 
 import numpy as np
@@ -26,13 +27,17 @@
 BPM_PLANES: str = "MonPlanes"
 
 
-def read_tbt(file_path: Union[str, Path]) -> TbtData:
+def read_tbt(file_path: Union[str, Path], remove_trailing_bpm_plane: bool = True) -> TbtData:
     """
     Reads turn-by-turn data from the ``SPS``'s **SDDS** format file.
     Will first determine if it is in ASCII format to figure out which reading method to use.
 
     Args:
         file_path (Union[str, Path]): path to the turn-by-turn measurement file.
+        remove_trailing_bpm_plane (bool, optional): if ``True``, will remove the trailing
+            BPM plane ('.H', '.V') from the BPM-names. 
+            This makes the measurement data compatible with the madx-models.
+            Defaults to ``True``.
 
     Returns:
         A ``TbTData`` object with the loaded data.
@@ -41,8 +46,7 @@ def read_tbt(file_path: Union[str, Path]) -> TbtData:
     LOGGER.debug(f"Reading SPS file at path: '{file_path.absolute()}'")
 
     if is_ascii_file(file_path):
-        matrices, date = read_ascii(file_path)
-        return TbtData(matrices, date, [0], matrices[0].X.shape[1])
+        return read_ascii(file_path)
 
     sdds_file = sdds.read(file_path)
 
@@ -53,30 +57,39 @@ def read_tbt(file_path: Union[str, Path]) -> TbtData:
     bpm_names = np.array(sdds_file.values[BPM_NAMES])
     bpm_planes = np.array(sdds_file.values[BPM_PLANES]).astype(bool)
 
-    ver_bpms = bpm_names[bpm_planes]
-    hor_bpms = bpm_names[~bpm_planes]
+    bpm_names_y = bpm_names[bpm_planes]
+    bpm_names_x = bpm_names[~bpm_planes]
 
-    tbt_data_x = [sdds_file.values[bpm] for bpm in hor_bpms]
-    tbt_data_y = [sdds_file.values[bpm] for bpm in ver_bpms]
+    tbt_data_x = [sdds_file.values[bpm] for bpm in bpm_names_x]
+    tbt_data_y = [sdds_file.values[bpm] for bpm in bpm_names_y]
+    
+    if remove_trailing_bpm_plane:
+        pattern = re.compile("\.[HV]$", flags=re.IGNORECASE)
+        bpm_names_x  = [pattern.sub("", bpm) for bpm in bpm_names_x]
+        bpm_names_y  = [pattern.sub("", bpm) for bpm in bpm_names_y]
 
     matrices = [
         TransverseData(
-            X=pd.DataFrame(index=hor_bpms, data=tbt_data_x, dtype=float),
-            Y=pd.DataFrame(index=ver_bpms, data=tbt_data_y, dtype=float),
+            X=pd.DataFrame(index=bpm_names_x, data=tbt_data_x, dtype=float),
+            Y=pd.DataFrame(index=bpm_names_y, data=tbt_data_y, dtype=float),
         )
     ]
 
     return TbtData(matrices, date, [0], nturns)
 
 
-def write_tbt(output_path: Union[str, Path], tbt_data: TbtData) -> None:
+def write_tbt(output_path: Union[str, Path], tbt_data: TbtData, add_trailing_bpm_plane: bool = True) -> None:
     """
     Write a ``TbtData`` object's data to file, in a ``SPS``'s **SDDS** format.
     The format is reduced to the necessary parameters used by the reader.
 
     Args:
         output_path (Union[str, Path]): path to a the disk location where to write the data.
         tbt_data (TbtData): the ``TbtData`` object to write to disk.
+        add_trailing_bpm_plane (bool, optional): if ``True``, will add the trailing
+            BPM plane ('.H', '.V') to the BPM-names. This assures that all BPM-names are unique,
+            and that the measurement data is compatible with the sdds files from the FESA-class.
+            Defaults to ``True``.
     """
     output_path = Path(output_path)
     LOGGER.info(f"Writing TbTdata in binary SDDS (SPS) format at '{output_path.absolute()}'")
@@ -85,6 +98,11 @@ def write_tbt(output_path: Union[str, Path], tbt_data: TbtData) -> None:
 
     # bpm names
     bpm_names_x, bpm_names_y = df_x.index.to_list(), df_y.index.to_list()
+
+    if add_trailing_bpm_plane:
+        bpm_names_x = [f"{bpm_name}.H" if not bpm_name.endswith(".H") else bpm_name for bpm_name in bpm_names_x]
+        bpm_names_y = [f"{bpm_name}.V" if not bpm_name.endswith(".V") else bpm_name for bpm_name in bpm_names_y]
+
     bpm_names = bpm_names_x + bpm_names_y
 
     # bpm planes