Pass TbTData to Hole-in-one (#557)

Author: JoschD

omc3/harpy/handler.py

@@ -56,14 +56,15 @@
 
 
 def run_per_bunch(
-    tbt_data: TbtData, harpy_input: DotDict, file: Path
+    tbt_data: TbtData, harpy_input: DotDict, output_filename: str
 ) -> dict[str, tfs.TfsDataFrame]:
     """
     Cleans data, analyses frequencies and searches for resonances.
 
     Args:
         tbt_data: single bunch `TbtData`.
         harpy_input: Analysis settings taken from the commandline.
+        output_filename: Name of the output files (placed in `harpy_input.outputdir`).
 
     Returns:
         Dictionary with a `TfsDataFrame` per plane.
@@ -73,7 +74,7 @@ def run_per_bunch(
         model = tfs.read(harpy_input.model, index=COL_NAME).loc[:, COL_S]
 
     bpm_datas, usvs, lins, bad_bpms = {}, {}, {}, {}
-    output_file_path = harpy_input.outputdir / file.name
+    output_file_path = harpy_input.outputdir / output_filename
 
     for plane in PLANES:
         bpm_data = _get_cut_tbt_matrix(tbt_data, harpy_input.turns, plane)

omc3/hole_in_one.py

@@ -55,13 +55,14 @@
 from omc3.utils.contexts import timeit
 
 if TYPE_CHECKING:
-    from collections.abc import Generator
+    from collections.abc import Generator, Iterable
 
     from generic_parser import DotDict
 
 LOGGER = logging_tools.get_logger(__name__)
 
 DEFAULT_CONFIG_FILENAME = "analysis_{time:s}.ini"
+DATATYPE_TBT: str = "tbt_data"
 
 
 def hole_in_one_params() -> EntryPointParameters:
@@ -89,7 +90,8 @@ def hole_in_one_entrypoint(opt: DotDict, rest: list[str]) -> None:
         Action: ``store_true``
 
     Harpy Kwargs:
-      - **files**: TbT files to analyse
+      - **files**: TbT files to analyse.
+        Can also be the TbtData objects directly, if 'tbt_data' is chosen as datatype.
 
         Flags: **--files**
         Required: ``True``
@@ -405,58 +407,68 @@ def _run_harpy(harpy_options: DotDict) -> list[Path]:
     iotools.create_dirs(harpy_options.outputdir)
     with timeit(lambda spanned: LOGGER.info(f"Total time for Harpy: {spanned}")):
         lins = []
-        tbt_datas = [
-            (tbt.read_tbt(file, datatype=harpy_options.tbt_datatype), file)
-            for file in harpy_options.files
-        ]
+        tbt_datas = _parse_tbt_data(harpy_options.files, harpy_options.tbt_datatype)
         for tbt_data, file in tbt_datas:
             lins.extend(
                 [
-                    handler.run_per_bunch(bunch_data, harpy_options, bunch_file)
-                    for bunch_data, bunch_file in _add_suffix_and_iter_bunches(
+                    handler.run_per_bunch(bunch_data, harpy_options, name_for_bunch)
+                    for bunch_data, name_for_bunch in _add_suffix_and_iter_bunches(
                         tbt_data, harpy_options, file
                     )
                 ]
             )
     return lins
 
 
+def _parse_tbt_data(files: Iterable[Path | str | tbt.TbtData], tbt_datatype: str
+    ) -> list[tuple[tbt.TbtData, str]]:
+    """Parse the turn-by-turn data reading given files or TbtData objects."""
+    if tbt_datatype == DATATYPE_TBT:
+        try:
+            return [(file, Path(file.meta["file"]).name) for file in files]
+        except KeyError as e:
+            raise KeyError(
+                "To determine output naming for hole-in-one, "
+                "the given TbT objects must contain a 'file' entry in their meta-data."
+            ) from e
+
+    return [(tbt.read_tbt(file, datatype=tbt_datatype), Path(file).name) for file in files]
+
+
 def _add_suffix_and_iter_bunches(
-    tbt_data: tbt.TbtData, options: DotDict, file: Path
-) -> Generator[tuple[tbt.TbtData, Path], None, None]:
-    """Add suffix to output files and iterate over bunches."""
-    dir_name: Path = file.parent
-    file_name: str = file.name
+    tbt_data: tbt.TbtData, options: DotDict, file_name: str
+) -> Generator[tuple[tbt.TbtData, str], None, None]:
+    """Add the additional suffix (if given by user) to output files and
+    split the TbT data into bunches to analyse them individually."""
     suffix: str = options.suffix or ""
 
     # Single bunch
     if tbt_data.nbunches == 1:
-        if suffix:
-            file = dir_name / f"{file_name}{suffix}"
-        yield tbt_data, file
+        file_name_out = f"{file_name}{suffix}"
+        yield tbt_data, file_name_out
         return
 
     # Multibunch
     if options.bunch_ids is not None:
         unknown_bunches = set(options.bunch_ids) - set(tbt_data.bunch_ids)
         if unknown_bunches:
-            LOGGER.warning(f"Bunch IDs {unknown_bunches} not present in multi-bunch file {file}.")
+            LOGGER.warning(f"Bunch IDs {unknown_bunches} not present in multi-bunch file {file_name}.")
 
     for index in range(tbt_data.nbunches):
         bunch_id = tbt_data.bunch_ids[index]
         if options.bunch_ids is not None and bunch_id not in options.bunch_ids:
             continue
 
         bunch_id_str = f"_bunchID{bunch_id}"
-        file = dir_name / f"{file_name}{bunch_id_str}{suffix}"
+        file_name_out = f"{file_name}{bunch_id_str}{suffix}"
         yield (
             tbt.TbtData(
                 matrices=[tbt_data.matrices[index]],
                 nturns=tbt_data.nturns,
                 bunch_ids=[bunch_id],
                 meta=tbt_data.meta,
             ),
-            file,
+            file_name_out,
         )
 
 
@@ -499,18 +511,34 @@ def _harpy_entrypoint(params: list[str]) -> tuple[DotDict, list[str]]:
         raise AttributeError(
             "The magnet order for resonance lines calculation should be between 2 and 8 (inclusive)."
         )
-    options.files = [Path(file) for file in options.files]
     options.outputdir = Path(options.outputdir)
     return options, rest
 
 
 def harpy_params() -> EntryPointParameters:
     """Create the entry point parameters for harpy."""
+    # fmt: off
     params = EntryPointParameters()
-    params.add_parameter(name="files", required=True, nargs="+", help="TbT files to analyse")
-    params.add_parameter(name="outputdir", required=True, help="Output directory.")
-    params.add_parameter(name="suffix", type=str, help="User-defined suffix for output filenames.")
-    params.add_parameter(name="model", help="Model for BPM locations")
+    params.add_parameter(
+        name="files",
+        required=True,
+        nargs="+",
+        help="TbT files to analyse. Can also be the TbtData objects directly, if 'tbt_data' is chosen as datatype."
+    )
+    params.add_parameter(
+        name="outputdir",
+        required=True,
+        help="Output directory."
+    )
+    params.add_parameter(
+        name="suffix",
+        type=str,
+        help="User-defined suffix for output filenames."
+    )
+    params.add_parameter(
+        name="model",
+        help="Model for BPM locations"
+    )
     params.add_parameter(
         name="unit",
         type=str,
@@ -541,7 +569,7 @@ def harpy_params() -> EntryPointParameters:
     params.add_parameter(
         name="tbt_datatype",
         default=HARPY_DEFAULTS["tbt_datatype"],
-        choices=list(tbt.io.TBT_MODULES.keys()),
+        choices=list(tbt.io.TBT_MODULES.keys()) + [DATATYPE_TBT],
         help="Choose the datatype from which to import. ",
     )
 
@@ -584,7 +612,11 @@ def harpy_params() -> EntryPointParameters:
         "and renormalisation in iterative SVD cleaning of dominant BPMs."
         " This is also equal to maximal number of BPMs removed per SVD mode.",
     )
-    params.add_parameter(name="bad_bpms", nargs="*", help="Bad BPMs to clean.")
+    params.add_parameter(
+        name="bad_bpms",
+        nargs="*",
+        help="Bad BPMs to clean."
+    )
     params.add_parameter(
         name="wrong_polarity_bpms",
         nargs="*",
@@ -691,6 +723,7 @@ def harpy_params() -> EntryPointParameters:
         default=HARPY_DEFAULTS["resonances"],
         help="Maximum magnet order of resonance lines to calculate.",
     )
+    # fmt: on
     return params
 
 
@@ -708,11 +741,23 @@ def _optics_entrypoint(params: list[str]) -> tuple[DotDict, list[str]]:
 
 def optics_params() -> EntryPointParameters:
     """Create the entry point parameters for optics."""
+    # fmt: off
     params = EntryPointParameters()
-    params.add_parameter(name="files", required=True, nargs="+", help="Files for analysis")
-    params.add_parameter(name="outputdir", required=True, help="Output directory")
     params.add_parameter(
-        name="calibrationdir", type=str, help="Path to calibration files directory."
+        name="files",
+        required=True,
+        nargs="+",
+        help="Files for analysis"
+    )
+    params.add_parameter(
+        name="outputdir",
+        required=True,
+        help="Output directory"
+    )
+    params.add_parameter(
+        name="calibrationdir",
+        type=str,
+        help="Path to calibration files directory."
     )
     params.add_parameter(
         name="coupling_method",
@@ -761,7 +806,9 @@ def optics_params() -> EntryPointParameters:
         help="Use 3 BPM method in beta from phase",
     )
     params.add_parameter(
-        name="only_coupling", action="store_true", help="Calculate only coupling. "
+        name="only_coupling",
+        action="store_true",
+        help="Calculate only coupling. "
     )
     params.add_parameter(
         name="compensation",
@@ -797,6 +844,7 @@ def optics_params() -> EntryPointParameters:
         help="Filter files to analyse by this value (in analysis for tune, phase, rdt and crdt). "
         "Use `None` for no filtering",
     )
+    # fmt: on
     return params
 
 

pyproject.toml

@@ -136,9 +136,6 @@ markers = [
 
 addopts = [
     "--import-mode=importlib",
-    "--cov-report=term-missing",
-    "--cov-config=pyproject.toml",
-    "--cov=omc3",
 ]
 
 # Helpful for pytest-debugging (leave commented out on commit):

tests/unit/test_harpy.py

@@ -17,7 +17,7 @@
 def test_input_suffix_and_single_bunch(suffix):
     """ Tests the function :func:`omc3.hole_in_one._add_suffix_and_loop_over_bunches`
     by checking that the suffix is attached to single-bunch files."""
-    file_name = Path("input_file.sdds")
+    input_name = "input_file.sdds"
     options = DotDict(
         suffix=suffix,
         bunch_ids=None,
@@ -28,10 +28,10 @@ def test_input_suffix_and_single_bunch(suffix):
         bunch_ids=[0],
     )
     n_data = 0
-    for data, file in _add_suffix_and_iter_bunches(tbt_data, options, file_name):
+    for data, file_name in _add_suffix_and_iter_bunches(tbt_data, options, input_name):
         suffix_str = suffix or ""
-        assert file.name.endswith(f"{file_name}{suffix_str}")
-        assert "bunchID" not in str(file)
+        assert file_name == f"{input_name}{suffix_str}"
+        assert "bunchID" not in input_name
         assert data is tbt_data
         n_data += 1
 
@@ -45,7 +45,7 @@ def test_input_suffix_and_multibunch(suffix, bunches):
     """ Tests the function :func:`omc3.hole_in_one._add_suffix_and_loop_over_bunches`
     by checking that the suffixes are attached to multi-bunch files and they are
     split up into single-bunch files correctly."""
-    file_name = Path("input_file.sdds")
+    input_name = "input_file.sdds"
     options = DotDict(
         suffix=suffix,
         bunch_ids=None if bunches is None else list(bunches),
@@ -58,10 +58,10 @@ def test_input_suffix_and_multibunch(suffix, bunches):
     n_data = 0
     bunch_ids = bunches or tbt_data.bunch_ids
     matrices =  [tbt_data.matrices[tbt_data.bunch_ids.index(id_)] for id_ in bunch_ids]
-    for (data, file), bunch_id, matrix in zip(_add_suffix_and_iter_bunches(tbt_data, options, file_name), bunch_ids, matrices):
+    for (data, filename_with_suffix), bunch_id, matrix in zip(_add_suffix_and_iter_bunches(tbt_data, options, input_name), bunch_ids, matrices):
         bunch_str = f"_bunchID{bunch_id}"
         suffix_str = suffix or ""
-        assert file.name.endswith(f"{file_name}{bunch_str}{suffix_str}")
+        assert filename_with_suffix == f"{input_name}{bunch_str}{suffix_str}"
 
         assert len(data.matrices) == 1
         assert data.matrices[0] == matrix

tests/unit/test_hole_in_one.py

@@ -25,6 +25,9 @@
 from typing import TYPE_CHECKING
 
 import pytest
+from turn_by_turn import TbtData
+from tests.accuracy.test_harpy import _get_model_dataframe
+from tests.unit.test_harpy import create_tbt_data
 
 from omc3.harpy.constants import FILE_AMPS_EXT, FILE_FREQS_EXT, FILE_LIN_EXT
 from omc3.hole_in_one import (
@@ -63,7 +66,6 @@
     "-50Hz": ["Beam1@BunchTurn@2024_03_08@18_24_02_100_250turns.sdds", "Beam1@BunchTurn@2024_03_08@18_25_23_729_250turns.sdds", "Beam1@BunchTurn@2024_03_08@18_26_41_811_250turns.sdds"],
 }
 
-
 @pytest.mark.extended
 @pytest.mark.parametrize("which_files", ("SINGLE", "0Hz", "all"))
 @pytest.mark.parametrize("clean", (True, False), ids=ids_str("clean={}"))
@@ -173,6 +175,46 @@ def test_hole_in_one(tmp_path, clean, which_files, caplog):
     )
 
 
+@pytest.mark.basic
+def test_harpy_tbtdata_ok(tmp_path):
+    """ Tests the harpy entrypoint with `tbt_datatype == 'tbt_data'`."""
+    # Mock some TbT data
+    model = _get_model_dataframe()
+    tbt_data = create_tbt_data(model=model, bunch_ids=[0])
+    tbt_data.meta = {"file": "test.sdds"}
+
+    hole_in_one_entrypoint(
+        harpy=True,
+        files=[tbt_data],
+        tbt_datatype="tbt_data",
+        unit='m',
+        autotunes="transverse",
+        clean=False,
+        outputdir=tmp_path,
+    )
+
+
+@pytest.mark.basic
+def test_harpy_tbtdata_no_name(tmp_path):
+    """ Tests the harpy entrypoint by checking that meta-field `file` is required
+    when using `tbt_datatype == 'tbt_data'`."""
+    # Mock some TbT data
+    model = _get_model_dataframe()
+    tbt_data = create_tbt_data(model=model, bunch_ids=[0])
+    tbt_data.meta = {} # in case someone adds a meta in the create_tbt_data helper
+
+    with pytest.raises(KeyError) as e:
+        hole_in_one_entrypoint(
+            harpy=True,
+            files=[tbt_data],
+            tbt_datatype="tbt_data",
+            unit='m',
+            autotunes="transverse",
+            clean=False,
+            outputdir=tmp_path,
+        )
+    assert "must contain a 'file' entry" in str(e.value)
+
 # Helper -----------------------------------------------------------------------
 
 def _check_all_harpy_files(outputdir: Path, sdds_file: Path):