Handle TbtData objects as input files

omc3/harpy/handler.py

@@ -39,7 +39,11 @@ def run_per_bunch(tbt_data, harpy_input):
     """
     model = None if harpy_input.model is None else tfs.read(harpy_input.model, index=COL_NAME).loc[:, 'S']
     bpm_datas, usvs, lins, bad_bpms = {}, {}, {}, {}
-    output_file_path = _get_output_path_without_suffix(harpy_input.outputdir, harpy_input.files)
+
+    # Handle the case where we've got directly an object and its attributed name
+    files = harpy_input.tbt_name if harpy_input.tbt_datatype == 'tbt_data' else harpy_input.files
+    output_file_path = _get_output_path_without_suffix(harpy_input.outputdir, files)
+
     for plane in PLANES:
         bpm_data = _get_cut_tbt_matrix(tbt_data, harpy_input.turns, plane)
         bpm_data = _scale_to_meters(bpm_data, harpy_input.unit)

omc3/hole_in_one.py

@@ -117,6 +117,12 @@ def hole_in_one_entrypoint(opt, rest):
 
         Flags: **--tbt_datatype**
         Default: ``LHC``
+
+      - **tbt_name** *(list)*: Names of the turn by turn measurements. Required when `tbt_datatype`
+        is `tbt_data`.
+
+        Flags: **--tbt_name**
+        Default: None
 
       *--Cleaning--*
 
@@ -387,7 +393,13 @@ def _run_harpy(harpy_options):
     with timeit(lambda spanned: LOGGER.info(f"Total time for Harpy: {spanned}")):
         lins = []
         all_options = _replicate_harpy_options_per_file(harpy_options)
-        tbt_datas = [(tbt.read_tbt(option.files, datatype=option.tbt_datatype), option) for option in all_options]
+
+        # Read the TbT data
+        if harpy_options.tbt_datatype == "tbt_data":
+          tbt_datas = [(option.files, option) for option in all_options]
+        else:
+          tbt_datas = [(tbt.read_tbt(option.files, datatype=option.tbt_datatype), option) for option in all_options]
+
         for tbt_data, option in tbt_datas:
             lins.extend([handler.run_per_bunch(bunch_data, bunch_options)
                          for bunch_data, bunch_options in _add_suffix_and_iter_bunches(tbt_data, option)])
@@ -396,9 +408,12 @@ def _run_harpy(harpy_options):
 
 def _replicate_harpy_options_per_file(options):
     list_of_options = []
-    for input_file in options.files:
+    for i, input_file in enumerate(options.files):
         new_options = deepcopy(options)
         new_options.files = input_file
+
+        if options.tbt_name:
+          new_options.tbt_name = options.tbt_name[i]
         list_of_options.append(new_options)
     return list_of_options
 
@@ -408,13 +423,22 @@ def _add_suffix_and_iter_bunches(tbt_data: tbt.TbtData, options: DotDict
     # hint: options.files is now a single file because of _replicate_harpy_options_per_file
     # it is also only used here to define the output name, as the tbt-data is already loaded.
 
-    dir_name = dirname(options.files)
-    file_name = basename(options.files)
+    # When given directly a TbtData object, there won't be a filename! Use the argument `tbt_name`
+    # instead
+    if 'tbt_datatype' in options and options.tbt_datatype == 'tbt_data':
+      dir_name = ''
+      file_name = options.tbt_name
+    else:
+      dir_name = dirname(options.files)
+      file_name = basename(options.files)
     suffix = options.suffix or ""
 
     # Single bunch ---
     if tbt_data.nbunches == 1:
         if suffix:
+          if 'tbt_datatype' in options and options.tbt_datatype == 'tbt_data':
+            options.tbt_name = f"{file_name}{suffix}"
+          else:
             options.files = join(dir_name, f"{file_name}{suffix}")
         yield tbt_data, options
         return
@@ -475,6 +499,10 @@ def _harpy_entrypoint(params):
         options.window = "rectangle"
     if not 2 <= options.resonances <= 8:
         raise AttributeError("The magnet order for resonance lines calculation should be between 2 and 8 (inclusive).")
+    if options.tbt_datatype == "tbt_data" and options.tbt_name is None:
+        raise AttributeError("tbt_name must be specified when using TbtData objects.")
+    if options.tbt_datatype == "tbt_data" and len(options.tbt_name) != len(options.files):
+        raise AttributeError("Mismatch in length between tbt_name and files")
 
     return options, rest
 
@@ -497,8 +525,9 @@ def harpy_params():
                          choices=('lin', 'spectra', 'full_spectra', 'bpm_summary'),
                          help="Choose the type of output.")
     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()) + ["tbt_data"],
                          help="Choose the datatype from which to import. ")
+    params.add_parameter(name="tbt_name", nargs='+', help="Names of the turn by turn data objects.")
 
     # Cleaning parameters
     params.add_parameter(name="clean", action="store_true",

tests/unit/test_hole_in_one.py

@@ -48,6 +48,7 @@
     TOTAL_PHASE_NAME,
 )
 from tests.conftest import INPUTS, ids_str
+from turn_by_turn import TbtData
 
 MODEL_DIR = INPUTS / "models" / "2022_inj_b1_acd"
 SDDS_DIR = INPUTS / "lhcb1_tbt_inj_on_off_mom"
@@ -60,6 +61,75 @@
     "-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.basic
+def test_harpy_tbtdata_ok(tmp_path):
+    """ Tests the harpy entrypoint by checking that the argument `tbt_name` is required
+    when using `tbt_datatype == 'tbt_data'`."""
+
+    from tests.unit.test_harpy import create_tbt_data
+    from tests.accuracy.test_harpy import _get_model_dataframe
+
+    # Mock some TbT data
+    model = _get_model_dataframe()
+    tbt_data = create_tbt_data(model=model, bunch_ids=[0])
+
+    hole_in_one_entrypoint(
+        harpy=True,
+        files=[tbt_data],
+        tbt_name=['tbt_object'],
+        tbt_datatype="tbt_data",
+        unit='m',
+        autotunes="transverse",
+        clean=False,
+        outputdir=tmp_path,
+    )
+
+@pytest.mark.basic
+def test_harpy_tbtdata_wrong_length_name(tmp_path):
+    """ Tests the harpy entrypoint by checking that the argument `tbt_name` is required
+    when using `tbt_datatype == 'tbt_data'`."""
+
+    from tests.unit.test_harpy import create_tbt_data
+    from tests.accuracy.test_harpy import _get_model_dataframe
+
+    # Mock some TbT data
+    model = _get_model_dataframe()
+    tbt_data = create_tbt_data(model=model, bunch_ids=[0])
+
+    with pytest.raises(AttributeError):
+        hole_in_one_entrypoint(
+            harpy=True,
+            files=[tbt_data],
+            tbt_name=['tbt_object', 'wrong'],
+            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 the argument `tbt_name` is required
+    when using `tbt_datatype == 'tbt_data'`."""
+
+    from tests.unit.test_harpy import create_tbt_data
+    from tests.accuracy.test_harpy import _get_model_dataframe
+
+    # Mock some TbT data
+    model = _get_model_dataframe()
+    tbt_data = create_tbt_data(model=model, bunch_ids=[0])
+
+    with pytest.raises(AttributeError):
+        hole_in_one_entrypoint(
+            harpy=True,
+            files=[tbt_data],
+            tbt_datatype="tbt_data",
+            unit='m',
+            autotunes="transverse",
+            clean=False,
+            outputdir=tmp_path,
+        )
 
 @pytest.mark.extended
 @pytest.mark.parametrize("which_files", ("SINGLE", "0Hz", "all"))