diff --git a/omc3/harpy/handler.py b/omc3/harpy/handler.py index 038b8f038..932272322 100644 --- a/omc3/harpy/handler.py +++ b/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) diff --git a/omc3/hole_in_one.py b/omc3/hole_in_one.py index 74d80aebd..e1858028f 100644 --- a/omc3/hole_in_one.py +++ b/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", diff --git a/tests/unit/test_hole_in_one.py b/tests/unit/test_hole_in_one.py index 800e0a819..6b649a84a 100644 --- a/tests/unit/test_hole_in_one.py +++ b/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"))