def test_backup_and_restore(tmp_path): """Test that the backup and restore functionality works.""" columns = [COL_TUNE] plane_columns = [f"{col}{p}" for col in columns for p in PLANES] linfiles = _copy_and_modify_linfiles(tmp_path) unfiltered = {p: tfs.read(f) for p, f in linfiles.items()} _assert_nlinfiles(tmp_path, 1) clean_columns(files=linfiles.values(), columns=plane_columns, backup=True) _assert_nlinfiles(tmp_path, 2) clean_columns(files=linfiles.values(), columns=plane_columns, backup=True) _assert_nlinfiles(tmp_path, 3) filtered = {p: tfs.read(f) for p, f in linfiles.items()} restore_files(files=linfiles.values()) _assert_nlinfiles(tmp_path, 2) restore_files(files=linfiles.values()) _assert_nlinfiles(tmp_path, 1) restored = {p: tfs.read(f) for p, f in linfiles.items()} for plane in PLANES: assert_frame_equal(unfiltered[plane], restored[plane]) with pytest.raises(AssertionError): assert_frame_equal(unfiltered[plane], filtered[plane]) with pytest.raises(IOError): restore_files(files=linfiles.values())
def _create_fake_measurement(tmp_path, model_path, twiss_path, error_val, optics_params, seed): model_df = tfs.read(model_path / "twiss.dat", index=NAME) model_df = add_coupling_to_model(model_df) twiss_df = tfs.read(twiss_path, index=NAME) twiss_df = add_coupling_to_model(twiss_df) # create fake measurement data fake_measurement( model=model_df, twiss=twiss_df, randomize=[VALUES, ERRORS], relative_errors=[error_val], seed=seed, outputdir=tmp_path, ) # load the fake data into a dict _, meas_dict = get_measurement_data( optics_params, meas_dir=tmp_path, beta_file_name='beta_phase_', ) # map to VALUE, ERROR and WEIGHT, similar to filter_measurement # but without the filtering for col, meas in meas_dict.items(): if col[:-1] in (F1010, F1001): col = {c[0]: c for c in (REAL, IMAG, PHASE, AMPLITUDE)}[col[-1]] if col != TUNE: meas[VALUE] = meas.loc[:, col].to_numpy() meas[ERROR] = meas.loc[:, f"{ERR}{col}"].to_numpy() meas[WEIGHT] = 1. return twiss_df, model_df, meas_dict
def test_switch_signs_for_beam4_madx_data(): input_dir = INPUT / "twiss_optics" file_twiss = "twiss.lhc.b{:d}.unsliced.tfs" file_errors = "errors.lhc.b{:d}.unsliced.tfs" df_twiss_b2 = tfs.read(input_dir / file_twiss.format(2), index=NAME) df_errors_b2 = tfs.read(input_dir / file_errors.format(2), index=NAME) df_twiss_b4 = tfs.read(input_dir / file_twiss.format(4), index=NAME) df_errors_b4 = tfs.read(input_dir / file_errors.format(4), index=NAME) # Reverse index to compare with beam 2 df_twiss_b4 = df_twiss_b4.iloc[::-1, :] df_errors_b4 = df_errors_b4.iloc[::-1, :] df_twiss_b4switched, df_errors_b4switched = switch_signs_for_beam4( df_twiss_b4, df_errors_b4) twiss_cols, err_cols = get_twiss_and_error_columns(6) # as the values are not exact and not all signs perfect: check if more signs are equal than before... # Not the most impressive test. Other ideas welcome. assert ((np.sign(df_twiss_b4switched[twiss_cols]) == np.sign( df_twiss_b2[twiss_cols])).sum().sum() > (np.sign(df_twiss_b4[twiss_cols]) == np.sign( df_twiss_b2[twiss_cols])).sum().sum() + 0.9 * len(df_twiss_b2)) assert ((np.sign(df_errors_b4switched[err_cols]) == np.sign( df_errors_b2[err_cols])).sum().sum() > (np.sign(df_errors_b4[err_cols]) == np.sign( df_errors_b2[err_cols])).sum().sum() + len(df_twiss_b2))
def test_harpy_3d(_test_file, _model_file): model = _get_model_dataframe() tfs.write(_model_file, model, save_index="NAME") _write_tbt_file(model, os.path.dirname(_test_file)) hole_in_one_entrypoint(harpy=True, clean=True, autotunes="all", outputdir=os.path.dirname(_test_file), files=[_test_file], model=_model_file, to_write=["lin"], turn_bits=18, unit="m") lin = dict(X=tfs.read(f"{_test_file}.linx"), Y=tfs.read(f"{_test_file}.liny")) model = tfs.read(_model_file) _assert_spectra(lin, model) assert _rms(_diff(lin["X"].loc[:, "TUNEZ"].to_numpy(), TUNEZ)) < LIMITS["F2"] assert _rms( _rel_diff( lin["X"].loc[:, f"AMPZ"].to_numpy() * lin["X"].loc[:, f"AMPX"].to_numpy() * 2, AMPZ * BASEAMP)) < LIMITS["A2"] assert _rms(_angle_diff(lin["X"].loc[:, f"MUZ"].to_numpy(), MUZ)) < LIMITS["P2"]
def test_simple_diff(model_error_path): """Asserts that the diff_twiss_parameters functions perform the correct calculations by applying the respective inverse calculations on model, model_errors and delta.""" model = tfs.read(MODEL_INJ_BEAM1, index=NAME) model = add_coupling_to_model(model) model_errors = tfs.read(model_error_path, index=NAME) model_errors = add_coupling_to_model(model_errors) diff = diff_twiss_parameters(model_errors, model, OPTICS_PARAMS_CHOICES) for param in OPTICS_PARAMS_CHOICES: delta = f"{DELTA}{param}" if param[:-1] == BETA: check = model[param] * (1 + diff[delta]) - model_errors[param] elif param[:-1] == NORM_DISP: beta, disp = f"{BETA}{param[-1]}", f"{DISP}{param[-1]}" check = model[disp]/np.sqrt(model[beta]) + diff[delta] - model_errors[disp]/np.sqrt(model_errors[beta]) elif param[:-1] in (PHASE, PHASE_ADV): param = f"{PHASE_ADV}{param[-1]}" check = ang_diff(ang_sum(np.diff(model[param]), diff[delta][1:]), np.diff(model_errors[param])) elif param == TUNE: check = [model[f"{param}{i}"] + diff[f"{delta}{i}"] - model_errors[f"{param}{i}"] for i in "12"] else: check = model[param] + diff[delta] - model_errors[param] assert all(np.abs(check) < EPS)
def test_freekick_harpy(_test_file, _model_file): model = _get_model_dataframe() tfs.write(_model_file, model, save_index="NAME") _write_tbt_file(model, os.path.dirname(_test_file)) hole_in_one_entrypoint(harpy=True, clean=True, autotunes="transverse", is_free_kick=True, outputdir=os.path.dirname(_test_file), files=[_test_file], model=_model_file, to_write=["lin"], unit='m', turn_bits=18) lin = dict(X=tfs.read(f"{_test_file}.linx"), Y=tfs.read(f"{_test_file}.liny")) model = tfs.read(_model_file) for plane in PLANES: # main and secondary frequencies assert _rms( _diff(lin[plane].loc[:, f"TUNE{plane}"].to_numpy(), model.loc[:, f"TUNE{plane}"].to_numpy())) < LIMITS["F1"] # main and secondary amplitudes # TODO remove factor 2 - only for backwards compatibility with Drive assert _rms( _rel_diff(lin[plane].loc[:, f"AMP{plane}"].to_numpy() * 2, model.loc[:, f"AMP{plane}"].to_numpy())) < LIMITS["A1"] # main and secondary phases assert _rms( _angle_diff(lin[plane].loc[:, f"MU{plane}"].to_numpy(), model.loc[:, f"MU{plane}"].to_numpy())) < LIMITS["P1"]
def test_harpy(_test_file, _model_file): [ clean, keep_exact_zeros, singval, peak_to_peak, window, max_peak, svd_dominance_limit, num_svd_iterations, tolerance, tune_clean_limit, turn_bits, output_bits ] = HARPY_INPUT[0] model = _get_model_dataframe() tfs.write(_model_file, model, save_index="NAME") _write_tbt_file(model, os.path.dirname(_test_file)) hole_in_one_entrypoint(harpy=True, clean=clean, keep_exact_zeros=keep_exact_zeros, singval=singval, peak_to_peak=peak_to_peak, window=window, max_peak=max_peak, svd_dominance_limit=svd_dominance_limit, num_svd_iterations=num_svd_iterations, tolerance=tolerance, tune_clean_limit=tune_clean_limit, turn_bits=turn_bits, output_bits=output_bits, autotunes="transverse", outputdir=os.path.dirname(_test_file), files=[_test_file], model=_model_file, to_write=["lin"], unit="m") lin = dict(X=tfs.read(f"{_test_file}.linx"), Y=tfs.read(f"{_test_file}.liny")) model = tfs.read(_model_file) _assert_spectra(lin, model)
def get_combined_model_and_tunes(model_dir): free = tfs.read(join(model_dir, 'twiss.dat')) driven = tfs.read(join(model_dir, 'twiss_ac.dat')) nattune = {"X": np.remainder(free.headers['Q1'], 1), "Y": np.remainder(free.headers['Q2'], 1)} tune = {"X": np.remainder(driven.headers['Q1'], 1), "Y": np.remainder(driven.headers['Q2'], 1)} model = pd.merge(free, driven, how='inner', on='NAME', suffixes=MOTION.values()) model['S'] = model.loc[:, 'S_f'] return model, tune, nattune
def test_polar_plot_interpolated(tmp_path): df_angles = tfs.read(DA_RESULTS_DIR / 'da_per_angle.tfs', index=ANGLE) df_da = tfs.read(DA_RESULTS_DIR / 'da.tfs') fig = plot_polar(df_angles=df_angles, df_da=df_da, interpolated=True, fill=False) assert len(fig.axes) == 1 assert len(fig.axes[0].lines) == 63 # 60 Seeds, MEAN, MIN, MAX
def init_from_model_dir(self, model_dir): LOGGER.debug("Creating accelerator instance from model dir") self.model_dir = model_dir # Elements ##################################### elements_path = join(model_dir, TWISS_ELEMENTS_DAT) if not isfile(elements_path): raise AcceleratorDefinitionError("Elements twiss not found") self.elements = tfs.read(elements_path, index="NAME") LOGGER.debug(f" model path = {join(model_dir, TWISS_DAT)}") try: self.model = tfs.read(join(model_dir, TWISS_DAT), index="NAME") except IOError: bpm_index = [ idx for idx in self.elements.index.to_numpy() if idx.startswith(self.BPM_INITIAL) ] self.model = self.elements.loc[bpm_index, :] self.nat_tunes = [ float(self.model.headers["Q1"]), float(self.model.headers["Q2"]) ] # Excitations ##################################### driven_filenames = dict(acd=join(model_dir, TWISS_AC_DAT), adt=join(model_dir, TWISS_ADT_DAT)) if isfile(driven_filenames["acd"]) and isfile(driven_filenames["adt"]): raise AcceleratorDefinitionError( "ADT as well as ACD models provided. Choose only one.") for key in driven_filenames.keys(): if isfile(driven_filenames[key]): self._model_driven = tfs.read(driven_filenames[key], index="NAME") self.excitation = DRIVEN_EXCITATIONS[key] if not self.excitation == AccExcitationMode.FREE: self.drv_tunes = [ self.model_driven.headers["Q1"], self.model_driven.headers["Q2"] ] # Best Knowledge ##################################### best_knowledge_path = join(model_dir, TWISS_BEST_KNOWLEDGE_DAT) if isfile(best_knowledge_path): self.model_best_knowledge = tfs.read(best_knowledge_path, index="NAME") # Optics File ######################################### opticsfilepath = join(self.model_dir, MODIFIERS_MADX) if isfile(opticsfilepath): self.modifiers = opticsfilepath # Error Def ##################################### errordefspath = join(self.model_dir, ERROR_DEFFS_TXT) if isfile(errordefspath): self.error_defs_file = errordefspath
def test_crdt_amp(order, _create_input): omc3_input = _create_input (optics_opt, path_to_lin) = omc3_input[order] ptc_crdt = tfs.read(path_to_lin / 'ptc_crdt.tfs', index="NAME") for crdt_dict in crdt.CRDTS: if order == crdt_dict["order"]: hio_crdt = tfs.read(optics_opt["outputdir"] / "crdt" / order / f'{crdt_dict["term"]}.tfs', index="NAME") assert _max_dev(hio_crdt["AMP"].to_numpy(), ptc_crdt[f"{crdt_dict['term']}_ABS"].to_numpy(), NOISELEVEL_AMP[order]) < ACCURACY_LIMIT[order]
def test_headers_merging_right(self, _tfs_file_x_pathlib, _tfs_file_y_pathlib, how): headers_left = tfs.read(_tfs_file_x_pathlib).headers headers_right = tfs.read(_tfs_file_y_pathlib).headers result = merge_headers(headers_left, headers_right, how=how) assert isinstance(result, OrderedDict) assert len(result) >= len(headers_left) and len(result) >= len( headers_right) # no key disappeared for key in result: # check that we prioritized headers_right's contents if key in headers_left and key in headers_right: assert result[key] == headers_right[key]
def test_correct_appending(self, _tfs_file_x_pathlib, _tfs_file_y_pathlib, how_headers): dframe_x = tfs.read(_tfs_file_x_pathlib) dframe_y = tfs.read(_tfs_file_y_pathlib) result = dframe_x.append(dframe_y, how_headers=how_headers) assert isinstance(result, TfsDataFrame) assert isinstance(result.headers, OrderedDict) assert_dict_equal( result.headers, merge_headers(dframe_x.headers, dframe_y.headers, how=how_headers)) assert_frame_equal( result, pd.DataFrame(dframe_x).append(pd.DataFrame(dframe_y)))
def test_correct_concatenating(self, _tfs_file_x_pathlib, _tfs_file_y_pathlib, how_headers, axis, join): dframe_x = tfs.read(_tfs_file_x_pathlib) dframe_y = tfs.read(_tfs_file_y_pathlib) objs = [dframe_x] * 4 + [dframe_y] * 4 result = concat(objs, how_headers=how_headers, axis=axis, join=join) merger = partial(merge_headers, how=how_headers) all_headers = (tfsdframe.headers for tfsdframe in objs) assert isinstance(result, TfsDataFrame) assert isinstance(result.headers, OrderedDict) assert_dict_equal(result.headers, reduce(merger, all_headers)) assert_frame_equal(result, pd.concat(objs, axis=axis, join=join))
def test_remove_some_not_found(): update_nattune( files=[str(_get_input_file())], interval=[0.2631, 0.265], # some here rename_suffix=RENAME_SUFFIX, not_found_action='remove') newx = tfs.read( str(_get_input_dir() / f'spec_test.sdds{RENAME_SUFFIX}.linx')) assert len(newx.index) == 2 # specific to this test-set newy = tfs.read( str(_get_input_dir() / f'spec_test.sdds{RENAME_SUFFIX}.liny')) assert len(newy.index) == 3 # specific to this test-set
def _run_analysis(output_dir: Union[str, Path], beam: int, sdds_input: str): """Run hole_in_one on provided data, return the loaded result coupling files for f1001 and f1010.""" hole_in_one_entrypoint( optics=True, accel="lhc", year="2018", beam=beam, model_dir=COUPLING_INPUTS / f"model_b{beam}", files=[f"{COUPLING_INPUTS}/{sdds_input}.sdds"], compensation="none", outputdir=output_dir, only_coupling=True, ) f1001 = tfs.read(output_dir / f"{F1001.lower()}.tfs", index=NAME) f1010 = tfs.read(output_dir / f"{F1010.lower()}.tfs", index=NAME) return f1001, f1010
def prepare_run(cls, lhc_instance, output_path): if lhc_instance.fullresponse: cls._prepare_fullresponse(lhc_instance, output_path) macros_path = join(output_path, MACROS_DIR) iotools.create_dirs(macros_path) lib_path = join(os.path.dirname(__file__), os.pardir, os.pardir, "lib") shutil.copy(join(lib_path, GENERAL_MACROS), join(macros_path, GENERAL_MACROS)) shutil.copy(join(lib_path, LHC_MACROS), join(macros_path, LHC_MACROS)) if lhc_instance.energy is not None: core = f"{int(lhc_instance.energy*1000):04d}" file_path = lhc_instance.get_lhc_error_dir() shutil.copy(join(file_path, f"{core}GeV.tfs"), join(output_path, ERROR_DEFFS_TXT)) shutil.copy( join(file_path, "b2_errors_settings", f"beam{lhc_instance.beam}_{core}GeV.madx"), join(output_path, B2_SETTINGS_MADX)) b2_table = tfs.read(join(lhc_instance.get_lhc_error_dir(), f"b2_errors_beam{lhc_instance.beam}.tfs"), index="NAME") gen_df = pd.DataFrame(data=np.zeros( (b2_table.index.size, len(_b2_columns()))), index=b2_table.index, columns=_b2_columns()) gen_df["K1L"] = b2_table.loc[:, f"K1L_{core}"].to_numpy() tfs.write(join(output_path, B2_ERRORS_TFS), gen_df, headers_dict={ "NAME": "EFIELD", "TYPE": "EFIELD" }, save_index="NAME")
def __init__(self, files_to_analyse, optics_opt): super(InputFiles, self).__init__(zip(PLANES, ([], []))) read_files = isinstance(files_to_analyse[0], str) for file_in in files_to_analyse: for plane in PLANES: df_to_load = ( tfs.read(f"{file_in}.lin{plane.lower()}").set_index( "NAME", drop=False) if read_files else file_in[plane]) df_to_load.index.name = None self[plane].append( self._repair_backwards_compatible_frame(df_to_load, plane)) if len(self['X']) + len(self['Y']) == 0: raise IOError("No valid input files") dpp_values = dpp.calculate_dpoverp(self, optics_opt) LOGGER.info(f"DPPS: {dpp_values}") amp_dpp_values = dpp.calculate_amp_dpoverp(self, optics_opt) LOGGER.info(f"DPP_AMPS: {amp_dpp_values}") for plane in PLANES: if optics_opt.isolation_forest: self[plane] = iforest.clean_with_isolation_forest( self[plane], optics_opt, plane) self[plane] = dpp.append_dpp(self[plane], dpp.arrange_dpps(dpp_values)) self[plane] = dpp.append_amp_dpp(self[plane], amp_dpp_values)
def test_kmod_simulation_ip1b1(tmp_path, _kmod_inputs_path): analyse_kmod( betastar_and_waist=[0.25, 0.0], working_directory=str(_kmod_inputs_path), beam="B1", simulation=True, no_sig_digits=True, no_plots=False, interaction_point="ip1", cminus=0.0, misalignment=0.0, errorK=0.0, errorL=0.0, tune_uncertainty=0.0e-5, outputdir=str(tmp_path), ) results = tfs.read(tmp_path / "ip1B1" / f"{RESULTS_FILE_NAME}{EXT}") beta_twiss = {"X": 0.25, "Y": 0.25} for plane in PLANES: beta_sim = beta_twiss[plane] beta_meas = results[f"{BETA}{STAR}{plane}"].loc[0] assert (np.abs(beta_meas - beta_sim)) / beta_sim < LIMITS["Accuracy"] beta_err_meas = results[f"{ERR}{BETA}{STAR}{plane}"].loc[0] assert (np.abs(beta_err_meas)) < LIMITS["Num Precision"]
def test_kmod_simulation_ip4b1(tmp_path, _kmod_inputs_path): analyse_kmod( betastar_and_waist=[200.0, -100.0], working_directory=str(_kmod_inputs_path), beam="B1", simulation=True, no_sig_digits=True, no_plots=False, circuits=["RQ6.R4B1", "RQ7.R4B1"], cminus=0.0, misalignment=0.0, errorK=0.0, errorL=0.0, tune_uncertainty=0.5e-5, outputdir=str(tmp_path), ) results = tfs.read(tmp_path / "MQY.6R4.B1-MQM.7R4.B1" / f"{INSTRUMENTS_FILE_NAME}{EXT}", index="NAME") original = { "BPMCS.7R4.B1": (3.64208332528655e01, 9.46041254954643e01), "BPM.7R4.B1": (3.61317067929723e01, 9.48945562104017e01), "BQSH.7R4.B1": (5.07121388372368e02, 9.07140610660815e01), "BPLH.7R4.B1": (4.79632975072045e02, 8.65331699893341e01), } for inst in results.index: beta_x, beta_y = original[inst] betas = dict(X=beta_x, Y=beta_y) for plane in PLANES: beta_meas = results[f"{BETA}{plane}"].loc[inst] assert (np.abs(beta_meas - betas[plane]) ) / betas[plane] < LIMITS["Meas Accuracy"] beta_err_meas = results[f"{ERR}{BETA}{plane}"].loc[inst] assert (beta_err_meas / beta_meas) < LIMITS["Meas Precision"]
def test_kmod_meas_ip4b1(_workdir_path): analyse_kmod(betastar_and_waist=[200.0, -100.0], working_directory=_workdir_path, beam='B1', simulation=False, no_sig_digits=True, no_plots=False, circuits=['RQ6.R4B1', 'RQ7.R4B1'], cminus=0.0, misalignment=0.0, errorK=0.0, errorL=0.0, tune_uncertainty=0.5E-5) results = tfs.read(join(_workdir_path, "MQY.6R4.B1-MQM.7R4.B1", f"{INSTRUMENTS_FILE_NAME}{EXT}"), index='NAME') original = { 'BPMCS.7R4.B1': (17.5074335336, 157.760070696), 'BPM.7R4.B1': (17.6430538896, 157.972911909), 'BQSH.7R4.B1': (455.457631868, 124.586686684), 'BPLH.7R4.B1': (423.68951095, 123.578577484) } for inst in results.index: beta_x, beta_y = original[inst] betas = dict(X=beta_x, Y=beta_y) for plane in PLANES: beta_meas = results[f"{BETA}{plane}"].loc[inst] assert (np.abs(beta_meas-betas[plane]))/betas[plane] < LIMITS['Meas Accuracy'] beta_err_meas = results[f"{ERR}{BETA}{plane}"].loc[inst] assert (beta_err_meas/beta_meas) < LIMITS['Meas Precision']
def test_coupling_bump_sextupole_rdts(): input_dir = INPUT / "coupling_bump" df_twiss = tfs.read(input_dir / "twiss.lhc.b1.coupling_bump.tfs", index=NAME) df_ptc_rdt = tfs.read(input_dir / "ptc_rdt.lhc.b1.coupling_bump.tfs", index=NAME) df_twiss = prepare_twiss_dataframe(df_twiss=df_twiss) rdt_names = ["F1002", "F3000"] df_rdt = calculate_rdts(df_twiss, rdt_names) for rdt in rdt_names: rdt_ptc = df_ptc_rdt[f"{rdt}{REAL}"] + 1j * df_ptc_rdt[f"{rdt}{IMAG}"] assert arrays_are_close_almost_everywhere(df_rdt[rdt], rdt_ptc, rtol=1e-2, percentile=0.9)
def test_kmod_meas_ip4b2(_workdir_path): analyse_kmod(betastar_and_waist=[200.0, -100.0], working_directory=_workdir_path, beam='B2', simulation=False, no_sig_digits=True, no_plots=False, circuits=['RQ7.L4B2', 'RQ6.L4B2'], cminus=0.0, misalignment=0.0, errorK=0.0, errorL=0.0, tune_uncertainty=0.5E-5) results = tfs.read(join(_workdir_path, "MQM.7L4.B2-MQY.6L4.B2", f"{INSTRUMENTS_FILE_NAME}{EXT}"), index='NAME') original = { 'BPMYA.6L4.B2': (456.789268726, 149.073169556), 'BGVCA.B7L4.B2': (119.359634764, 152.116072289), 'BPLH.B7L4.B2': (434.440558008, 148.460642194), 'BPLH.A7L4.B2': (441.781928671, 148.654814221) } for inst in results.index: beta_x, beta_y = original[inst] betas = dict(X=beta_x, Y=beta_y) for plane in PLANES: beta_meas = results[f"{BETA}{plane}"].loc[inst] assert (np.abs(beta_meas - betas[plane])) / betas[plane] < LIMITS['Meas Accuracy'] beta_err_meas = results[f"{ERR}{BETA}{plane}"].loc[inst] assert (beta_err_meas / beta_meas) < LIMITS['Meas Precision']
def convert_old_closed_orbit( inputdir: Union[Path, str], outputdir: Union[Path, str], plane: str, old_file_name: str = "CO", new_file_name: str = ORBIT_NAME, ) -> None: """ Looks in the provided directory for expected closed orbit file from ``BetaBeat.src`` for a given plane, converts it to the output format used by ``omc3`` and write them to the new location. The file naming should be getCO(x,y).out, with the following expected columns: NAME, S, COUNT, X, STDX, XMDL, MUXMDL. Args: inputdir (Union[Path, str]): Location of the directory with BetaBeat.src output files. outputdir (Union[Path, str]): Location of the output directory for converted files. plane (str): the transverse plane for which to look for the output file. old_file_name (str): the standard naming for the old output file. new_file_name (str): the standard naming for the new converted file. """ LOGGER.info("Converting old closed orbit file") old_file_path = Path( inputdir) / f"get{old_file_name}{plane.lower()}{OLD_EXT}" if not old_file_path.is_file(): LOGGER.debug( f"Expected BetaBeat.src output at '{old_file_path.absolute()}' is not a file, skipping" ) return dframe = tfs.read(old_file_path) dframe = dframe.rename(columns={f"STD{plane}": f"{ERR}{plane}"}) dframe[f"{DELTA}{plane}"] = df_diff(dframe, f"{plane}", f"{plane}{MDL}") dframe[f"{ERR}{DELTA}{plane}"] = dframe.loc[:, f"{ERR}{plane}"].to_numpy() tfs.write(Path(outputdir) / f"{new_file_name}{plane.lower()}{EXT}", dframe)
def _read_data(file_path, x_col, y_col, err_col): tfs_data = tfs.read(str(file_path)) return DotDict( x=tfs_data[x_col], y=tfs_data[y_col], err=tfs_data[err_col] if err_col is not None else None, )
def test_harpy_without_model(_test_file, _model_file): model = _get_model_dataframe() tfs.write(_model_file, model, save_index="NAME") _write_tbt_file(model, os.path.dirname(_test_file)) hole_in_one_entrypoint(harpy=True, clean=True, autotunes="transverse", outputdir=os.path.dirname(_test_file), files=[_test_file], to_write=["lin"], turn_bits=18, unit="m") lin = dict(X=tfs.read(f"{_test_file}.linx"), Y=tfs.read(f"{_test_file}.liny")) model = tfs.read(_model_file) _assert_spectra(lin, model)
def test_kmod_meas_ip1b2(tmp_path, _kmod_inputs_path): analyse_kmod( betastar_and_waist=[0.44, 0.0], working_directory=str(_kmod_inputs_path), beam="B2", simulation=False, no_sig_digits=True, no_plots=False, interaction_point="ip1", cminus=0.0, misalignment=0.0, errorK=0.0, errorL=0.0, tune_uncertainty=2.5e-5, outputdir=str(tmp_path), ) results = tfs.read(tmp_path / "ip1B2" / f"{RESULTS_FILE_NAME}{EXT}") beta_prev = {"X": 0.387, "Y": 0.410} for plane in PLANES: beta_meas = results[f"{BETA}{STAR}{plane}"].loc[0] assert (np.abs(beta_meas - beta_prev[plane]) ) / beta_prev[plane] < LIMITS["Meas Accuracy"] beta_err_meas = results[f"{ERR}{BETA}{STAR}{plane}"].loc[0] assert (beta_err_meas / beta_meas) < LIMITS["Meas Precision"]
def get_variables(self, frm=None, to=None, classes=None): correctors_dir = LHC_DIR / "2012" / "correctors" all_corrs = _merge_jsons( correctors_dir / f"correctors_b{self.beam}" / "beta_correctors.json", correctors_dir / f"correctors_b{self.beam}" / "coupling_correctors.json", self._get_triplet_correctors_file(), ) my_classes = classes if my_classes is None: my_classes = all_corrs.keys() vars_by_class = set( _flatten_list([all_corrs[corr_cls] for corr_cls in my_classes if corr_cls in all_corrs]) ) if frm is None and to is None: return list(vars_by_class) elems_matrix = tfs.read(self._get_corrector_elems()).sort_values("S") if frm is not None and to is not None: if frm > to: elems_matrix = elems_matrix[(elems_matrix.S >= frm) | (elems_matrix.S <= to)] else: elems_matrix = elems_matrix[(elems_matrix.S >= frm) & (elems_matrix.S <= to)] elif frm is not None: elems_matrix = elems_matrix[elems_matrix.S >= frm] elif to is not None: elems_matrix = elems_matrix[elems_matrix.S <= to] vars_by_position = _remove_dups_keep_order( _flatten_list([raw_vars.split(",") for raw_vars in elems_matrix.loc[:, "VARS"]]) ) return _list_intersect_keep_order(vars_by_position, vars_by_class)
def merge_data(kmod_input_params): magnet_df = [] work_dir = kmod_input_params.working_directory ip = kmod_input_params.interaction_point beam = kmod_input_params.beam for (filepaths, magnet) in zip( return_ip_filename(work_dir, ip, beam) if ip is not None else return_circuit_filename( work_dir, kmod_input_params.circuits, beam), kmod_input_params.magnets): LOG.debug(f'Loading tunes from {filepaths[0]} and {filepaths[1]}') tune_dfs = dict(X=tfs.read(filepaths[0]), Y=tfs.read(filepaths[1])) LOG.debug(f'Loading k from {filepaths[2]}') k_df = tfs.read(filepaths[2]) LOG.debug('Binning data') magnet_df.append(bin_tunes_and_k(tune_dfs, k_df, magnet)) return magnet_df
def test_filter_tune_limit(tmp_path): """ Test filtering works on outlier created by modify linfiles function. """ columns = [COL_TUNE] plane_columns = [f"{col}{p}" for col in columns for p in PLANES] linfiles = _copy_and_modify_linfiles(tmp_path, columns=columns, index=[2], by=0.1) unfiltered = {p: tfs.read(f) for p, f in linfiles.items()} # choose limit greater than the changes made clean_columns(files=linfiles.values(), columns=plane_columns, limit=0.2) filtered = {p: tfs.read(f) for p, f in linfiles.items()} for plane in PLANES: assert_frame_equal(unfiltered[plane], filtered[plane])