def write(df, header, meas_input, plane, rdt): outputdir = join(meas_input.outputdir, "rdt", _rdt_to_order_and_type(rdt)) iotools.create_dirs(outputdir) tfs.write(join(outputdir, f"f{_rdt_to_str(rdt)}_{plane.lower()}{EXT}"), df, header, save_index="NAME")
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}"].values, model.loc[:, f"TUNE{plane}"].values)) < 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}"].values * 2, model.loc[:, f"AMP{plane}"].values)) < LIMITS["A1"] # main and secondary phases assert _rms( _angle_diff(lin[plane].loc[:, f"MU{plane}"].values, model.loc[:, f"MU{plane}"].values)) < LIMITS["P1"]
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 chromatic_beating(input_files, measure_input, tune_dict): """ Main function to compute chromatic optics beating Args: tune_dict: input_files: InputFiles object containing frequency spectra files (linx/y) measure_input: OpticsInput object containing analysis settings Returns: """ dpps = np.array([dpp_val for dpp_val in set(input_files.dpps("X"))]) if np.max(dpps) - np.min(dpps) == 0.0: return for plane in PLANES: betas = [] for dpp_val in dpps: dpp_meas_input = deepcopy(measure_input) dpp_meas_input["dpp"] = dpp_val phase_dict, out_dfs = phase.calculate(dpp_meas_input, input_files, tune_dict, plane) beta_df, _ = beta_from_phase.calculate(dpp_meas_input, tune_dict, phase_dict, OrderedDict(), plane) betas.append(beta_df) output_df = chromatic.calculate_w_and_phi(betas, dpps, input_files, measure_input, plane) tfs.write(os.path.join(measure_input.outputdir, f"{CHROM_BETA_NAME}{plane.lower()}{EXT}"), output_df, {}, save_index="NAME")
def calculate(meas_input, input_files, tune_dict, beta_phase, header_dict, plane): """ Calculates beta and fills the following `TfsFiles`: ``f"{AMP_BETA_NAME}{plane.lower()}{EXT}"`` Args: meas_input: `OpticsInput` object. input_files: `InputFiles` object contains measurement files. tune_dict: `TuneDict` contains measured tunes. beta_phase: contains beta functions from measured from phase. header_dict: dictionary of header items common for all output files. plane: marking the horizontal or vertical plane, **X** or **Y**. Returns: """ beta_amp = beta_from_amplitude(meas_input, input_files, plane, tune_dict) x_ratio = phase_to_amp_ratio(meas_input, beta_phase, beta_amp, plane) beta_amp = add_rescaled_beta_columns(beta_amp, x_ratio, plane) header_d = _get_header( header_dict, np.std(beta_amp.loc[:, f"{DELTA}BET{plane}"].values), x_ratio) tfs.write(join(meas_input.outputdir, f"{AMP_BETA_NAME}{plane.lower()}{EXT}"), beta_amp, header_d, save_index='NAME') return x_ratio
def write(dfs, headers, output, plane): LOGGER.info(f"Writing phases: {len(dfs)}") for head, df, name in zip(headers, dfs, (PHASE_NAME, TOTAL_PHASE_NAME, PHASE_NAME + "driven_", TOTAL_PHASE_NAME + "driven_")): tfs.write(join(output, f"{name}{plane.lower()}{EXT}"), df, head) LOGGER.info(f"Phase advance beating in {name}{plane.lower()}{EXT} = " f"{stats.weighted_rms(df.loc[:, f'{DELTA}PHASE{plane}'])}")
def convert_old_coupling(inputdir: Union[Path, str], outputdir: Union[Path, str], suffix: str, old_file_name: str = "couple", new_file_name: str = "f") -> None: """ Looks in the provided directory for expected coupling file from ``BetaBeat.src``, converts it to the output format used by ``omc3`` and write them to the new location. The file naming should be getcouple(x,y).out, with the following expected columns: NAME, S, COUNT, F1001W, FWSTD1, F1001R, F1001I, F1010W, FWSTD2, F1010R, F1010I, Q1001, Q1001STD, Q1010, Q1010STD, MDLF1001R, MDLF1001I, MDLF1010R, MDLF1010I. 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. suffix (str): Compensation suffix used in the provided BetaBeat.src output files. 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 coupling files") old_file_path = Path(inputdir) / f"get{old_file_name}{suffix}{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) rdt_dfs = { "1001": dframe.loc[:, [ "S", "COUNT", "F1001W", "FWSTD1", "F1001R", "F1001I", "Q1001", "Q1001STD", "MDLF1001R", "MDLF1001I" ]], "1010": dframe.loc[:, [ "S", "COUNT", "F1010W", "FWSTD2", "F1010R", "F1010I", "Q1010", "Q1010STD", "MDLF1010R", "MDLF1010I" ]], } for i, rdt in enumerate(("1001", "1010")): LOGGER.debug(f"Converting F{rdt} file") rdt_dfs[rdt] = rdt_dfs[rdt].rename( columns={ f"F{rdt}W": AMPLITUDE, f"FWSTD{i+1}": f"{ERR}{AMPLITUDE}", f"Q{rdt}": PHASE, f"Q{rdt}STD": f"{ERR}{PHASE}", f"F{rdt}R": REAL, f"F{rdt}I": IMAG, f"MDLF{rdt}R": f"{REAL}{MDL}", f"MDLF{rdt}I": f"{IMAG}{MDL}", }) tfs.write(Path(outputdir) / f"{new_file_name}{rdt}{EXT}", rdt_dfs[rdt])
def create_lsa_results_file(betastar_required, instruments_found, results_df, instrument_beta_df, output_dir): lsa_results_df = pd.DataFrame(columns=LSA_COLUMNS) if betastar_required: exporting_columns=['LABEL', f'{BETA}{STAR}X', f'{ERR}{BETA}{STAR}X', f'{BETA}{STAR}Y', f'{ERR}{BETA}{STAR}Y'] lsa_results_df=results_df[exporting_columns].rename(columns=dict(zip(exporting_columns, LSA_COLUMNS))) if instruments_found: lsa_results_df = lsa_results_df.append(instrument_beta_df, sort=False, ignore_index=True) if not lsa_results_df.empty: tfs.write(join(output_dir, f'{LSA_FILE_NAME}{EXT}'), lsa_results_df)
def convert_old_beta_from_phase( inputdir: Union[Path, str], outputdir: Union[Path, str], suffix: str, plane: str, old_file_name: str = "beta", new_file_name: str = BETA_NAME, ) -> None: """ Looks in the provided directory for expected beta from phase 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 getbeta(x,y).out, with the following expected columns: NAME, S, COUNT, BETX, SYSBETX, STATBETX, ERRBETX, CORR_ALFABETA, ALFX, SYSALFX, STATALFX, ERRALFX, BETXMDL, ALFXMDL, MUXMDL, NCOMBINATIONS. 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. suffix (str): Compensation suffix used in the provided BetaBeat.src output 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 beta from phase file") old_file_path = Path( inputdir) / f"get{old_file_name}{plane.lower()}{suffix}{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) if "CORR_ALFABETA" in dframe.columns.to_numpy(): dframe = dframe.drop(columns=[ f"STATBET{plane}", f"SYSBET{plane}", "CORR_ALFABETA", f"STATALF{plane}", f"SYSALF{plane}" ]) else: dframe[f"{ERR}BET{plane}"] = df_err_sum(dframe, f"{ERR}BET{plane}", f"STDBET{plane}") dframe[f"{ERR}ALF{plane}"] = df_err_sum(dframe, f"{ERR}ALF{plane}", f"STDALF{plane}") dframe[f"{DELTA}BET{plane}"] = df_rel_diff(dframe, f"BET{plane}", f"BET{plane}{MDL}") dframe[f"{ERR}{DELTA}BET{plane}"] = df_ratio(dframe, f"{ERR}BET{plane}", f"BET{plane}{MDL}") dframe[f"{DELTA}ALF{plane}"] = df_diff(dframe, f"ALF{plane}", f"ALF{plane}{MDL}") dframe[ f"{ERR}{DELTA}ALF{plane}"] = dframe.loc[:, f"{ERR}ALF{plane}"].to_numpy() tfs.write(Path(outputdir) / f"{new_file_name}{plane.lower()}{EXT}", dframe)
def generate(opt) -> Dict[str, tfs.TfsDataFrame]: """ Takes a twiss file and writes the parameters in optics_parameters to Output_dir in the format global_correction_entrypoint uses (same format you would get from hole_in_one). """ LOG.info("Generating fake measurements.") # prepare data np.random.seed(opt.seed) randomize = opt.randomize if opt.randomize is not None else [] df_twiss, df_model = _get_data(opt.twiss, opt.model, add_coupling=(F1001 in opt.parameters) or (F1010 in opt.parameters)) # headers headers = { f"{TUNE}1": df_twiss[f"{TUNE}1"], f"{TUNE}2": df_twiss[f"{TUNE}2"] } if isinstance(opt.twiss, PathOrStr): headers[FAKED_HEADER] = str(opt.twiss) # create defaults results = {} for parameter, error in _get_loop_parameters(opt.parameters, opt.relative_errors): create = CREATOR_MAP[parameter[:-1]] new_dfs = create(df_twiss, df_model, parameter, relative_error=error, randomize=randomize, headers=headers) results.update(new_dfs) # maybe create normalized dispersion for plane in PLANES: nd_param = f'{NORM_DISP}{plane}' if nd_param in opt.parameters: nd_df = create_normalized_dispersion( results[f'{DISPERSION_NAME}{plane.lower()}'], results[f'{BETA_NAME}{plane.lower()}'], df_model, nd_param, headers) results.update(nd_df) # output if opt.outputdir is not None: LOG.info("Writing fake measurements to files.") output_path = Path(opt.outputdir) for filename, df in results.items(): full_path = output_path / f"{filename}{EXT}" tfs.write(full_path, df, save_index=NAME) return results
def _generate_jobs(basedir, jobid_mask, **kwargs) -> tfs.TfsDataFrame: """ Generates product matrix for job-values and stores it as TfsDataFrame. """ LOG.debug("Creating Jobs") job_names, values_grid = get_jobs_and_values(jobid_mask, **kwargs) job_df = tfs.TfsDataFrame( headers={HEADER_BASEDIR: basedir}, index=job_names, columns=list(kwargs.keys()), data=values_grid, ) tfs.write(basedir / JOBSUMMARY_FILE, job_df, save_index=COLUMN_JOBID) return job_df
def prepare_run(cls, accel: Lhc) -> None: if accel.year in ["2018", "2022" ]: # these years should be handled by the fetcher symlink_dst = Path(accel.model_dir) / LHC_REPOSITORY_NAME if not symlink_dst.exists(): LOGGER.debug(f"Symlink destination: {symlink_dst}") symlink_dst.absolute().symlink_to( (ACCELERATOR_MODEL_REPOSITORY / f"{accel.year}")) cls.check_accelerator_instance(accel) LOGGER.debug("Preparing model creation structure") macros_path = accel.model_dir / MACROS_DIR iotools.create_dirs(macros_path) LOGGER.debug("Copying macros to model directory") lib_path = Path(__file__).parent.parent / "madx_macros" shutil.copy(lib_path / GENERAL_MACROS, macros_path / GENERAL_MACROS) shutil.copy(lib_path / LHC_MACROS, macros_path / LHC_MACROS) shutil.copy(lib_path / LHC_MACROS_RUN3, macros_path / LHC_MACROS_RUN3) if accel.energy is not None: LOGGER.debug( "Copying B2 error files for given energy in model directory") core = f"{int(accel.energy * 1000):04d}" error_dir_path = accel.get_lhc_error_dir() shutil.copy(error_dir_path / f"{core}GeV.tfs", accel.model_dir / ERROR_DEFFS_TXT) shutil.copy( error_dir_path / "b2_errors_settings" / f"beam{accel.beam}_{core}GeV.madx", accel.model_dir / B2_SETTINGS_MADX, ) b2_table = tfs.read(error_dir_path / f"b2_errors_beam{accel.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( accel.model_dir / B2_ERRORS_TFS, gen_df, headers_dict={ "NAME": "EFIELD", "TYPE": "EFIELD" }, save_index="NAME", )
def write_special(meas_input, phase_advances, plane_tune, plane): # TODO REFACTOR AND SIMPLIFY accel = meas_input.accelerator meas = phase_advances["MEAS"] bd = accel.beam_direction elements = accel.elements special_phase_columns = ['ELEMENT1', 'ELEMENT2', f'PHASE{plane}', f'{ERR}PHASE{plane}', f'PHASE{plane}_DEG', f'{ERR}PHASE{plane}_DEG', f'PHASE{plane}{MDL}', f'PHASE{plane}{MDL}_DEG', 'BPM1', 'BPM2', f'BPM_PHASE{plane}', f'BPM_{ERR}PHASE{plane}',] special_phase_df = pd.DataFrame(columns=special_phase_columns) for elem1, elem2 in accel.important_phase_advances(): mus1 = elements.loc[elem1, f"MU{plane}"] - elements.loc[:, f"MU{plane}"] minmu1 = abs(mus1.loc[meas.index]).idxmin() mus2 = elements.loc[:, f"MU{plane}"] - elements.loc[elem2, f"MU{plane}"] minmu2 = abs(mus2.loc[meas.index]).idxmin() bpm_phase_advance = meas.loc[minmu1, minmu2] model_value = elements.loc[elem2, f"MU{plane}"] - elements.loc[elem1, f"MU{plane}"] if (elements.loc[elem1, "S"] - elements.loc[elem2, "S"]) * bd > 0.0: bpm_phase_advance += plane_tune model_value += plane_tune bpm_err = phase_advances["ERRMEAS"].loc[minmu1, minmu2] elems_to_bpms = -mus1.loc[minmu1] - mus2.loc[minmu2] ph_result = ((bpm_phase_advance + elems_to_bpms) * bd) model_value = (model_value * bd) % 1 special_phase_df=special_phase_df.append(dict(zip(special_phase_columns,[ elem1, elem2, ph_result % 1, bpm_err, _to_deg(ph_result), bpm_err * 360, model_value, _to_deg(model_value), minmu1, minmu2, bpm_phase_advance, elems_to_bpms, ])), ignore_index=True) tfs.write(join(meas_input.outputdir, f"{SPECIAL_PHASE_NAME}{plane.lower()}{EXT}"), special_phase_df)
def main(): #folder = "/Users/anagtv/Documents/OneDrive/046 - Medical Devices/Mantenimientos ciclotrones/MRS/LOGS/2018" folder = "/Users/anagtv/Documents/OneDrive/046 - Medical Devices/Mantenimientos ciclotrones/TCP/LOGS/2021" filename_completed = [] for file in os.listdir(folder): filename_completed.append(os.path.join(folder, file)) target = target_parameters() central_region_summary = central_region() file_summary = file_information() week_list = [] filename_completed_current = [] transmission_all = [] vacuum_transmission = [] probe_current_all = [] probe_current_std_all = [] j = -1 for file in filename_completed: real_values = saving_files_summary_list_20200420.get_irradiation_information( file) data_df = saving_files_summary_list_20200420.get_data(real_values) target_max_current = np.max(data_df.Target_I.astype(float)) target_average_current = np.average(data_df.Target_I.astype(float)) if target_average_current > 10.0: j += 1 file_summary.selecting_file_summary(file) df_isochronism = saving_files_summary_list_20200420.get_isochronism( data_df) central_region_summary.selecting_central_region_data( data_df, target_max_current, np.max(df_isochronism.Foil_I)) target.selecting_data( data_df, target_max_current, central_region_summary.transmission_option_a[j]) df_source = pd.DataFrame(list( zip(file_summary.date_stamp, file_summary.file_number, central_region_summary.source_current_estable, central_region_summary.source_current, target.probe_current_estimated, target.target_current)), columns=[ "DATE", "FILE", "SOURCE_CURRENT_AVE", "SOURCE_CURRENT_CUMULATIVE", "PROBE", "TARGET" ]) df_source_sort = df_source.sort_values(by="FILE", ignore_index=True) cummulative_source = [] for i in range(len(df_source_sort.PROBE)): cummulative_source.append( df_source_sort.SOURCE_CURRENT_CUMULATIVE[0:i].sum() * 3 / 3600) df_source_sort["CUMULATIVE_SOURCE_CURRENT"] = cummulative_source print(df_source_sort) tfs.write("source_performance_values.out", df_source_sort)
def write_summary(output_path: str, acc_time: AccDatetime, bp_info: DotDict, o_info: DotDict, trims: dict): """ Write summary into file. """ info_tfs = tfs.TfsDataFrame(trims.items(), columns=[const.column_knob, const.column_value]) info_tfs.headers = OrderedDict([ (const.head_time, acc_time.utc_string()), (const.head_beamprocess, bp_info.name), (const.head_fill, bp_info.fill), (const.head_beamprocess_start, bp_info.start.utc_string()), (const.head_context_category, bp_info.contextCategory), (const.head_beamprcess_description, bp_info.description), (const.head_optics, o_info.name), (const.head_optics_start, o_info.start.utc_string()), ("Hint:", "All times given in UTC.") ]) tfs.write(output_path, info_tfs)
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_no_tunes_in_files_plot(file_path, bpms): with tempfile.TemporaryDirectory() as out_dir: for f in glob(f'{file_path}*'): copy(f, out_dir) file_path = join(out_dir, basename(file_path)) for p in PLANES: fname = f'{file_path}.lin{p.lower()}' df = tfs.read(fname) tfs.write( fname, df.drop(columns=[f'TUNE{p.upper()}', f'NATTUNE{p.upper()}'])) plot_spectrum( files=[file_path], bpms=bpms, combine_by=['files', 'bpms'], )
def convert_old_total_phase( inputdir: Union[Path, str], outputdir: Union[Path, str], suffix: str, plane: str, old_file_name: str = "phasetot", new_file_name: str = TOTAL_PHASE_NAME, ) -> None: """ Looks in the provided directory for expected total phase 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 getphasetot(x,y).out, with the following expected columns: NAME, NAME2, S, S1, COUNT, PHASEX, STDPHX, PHXMDL, 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. suffix (str): Compensation suffix used in the provided BetaBeat.src output 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 total phase file") old_file_path = Path( inputdir) / f"get{old_file_name}{plane.lower()}{suffix}{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"STDPH{plane}": f"{ERR}{PHASE}{plane}", f"PH{plane}{MDL}": f"{PHASE}{plane}{MDL}", "S1": "S2", }, ) dframe[f"{DELTA}{PHASE}{plane}"] = df_ang_diff(dframe, f"{PHASE}{plane}", f"{PHASE}{plane}{MDL}") dframe[ f"{ERR}{DELTA}{PHASE}{plane}"] = dframe.loc[:, f"{ERR}{PHASE}{plane}"].to_numpy( ) tfs.write(Path(outputdir) / f"{new_file_name}{plane.lower()}{EXT}", dframe)
def _calculate_dispersion_3d(meas_input, input_files, header_dict, plane): """Computes dispersion from 3D kicks.""" output, accelerator = meas_input.outputdir, meas_input.accelerator model = accelerator.model df_orbit = _get_merged_df(meas_input, input_files, plane, ['AMPZ', 'MUZ', f"AMP{plane}"]) # work around due to scaling to main line in lin files unscaled_amps = (df_orbit.loc[:, input_files.get_columns(df_orbit, 'AMPZ')].values * df_orbit.loc[:, input_files.get_columns(df_orbit, f"AMP{plane}")].values) mask = accelerator.get_element_types_mask(df_orbit.index, ["arc_bpm"]) global_factors = np.array([1 / df.headers["DPPAMP"] for df in input_files[plane]]) # scaling to the model, and getting the synchrotron phase in the arcs df_orbit[f"D{plane}"], df_orbit[f"{ERR}D{plane}"] = _get_signed_dispersion( input_files, df_orbit, unscaled_amps * global_factors, mask) df_orbit[f"DP{plane}"] = _calculate_dp(model, df_orbit.loc[:, [f"D{plane}", f"{ERR}D{plane}"]], plane) df_orbit = _get_delta_columns(df_orbit, plane) output_df = df_orbit.loc[:, _get_output_columns(plane, df_orbit)] tfs.write(join(output, f"{DISPERSION_NAME}{plane.lower()}{EXT}"), output_df, header_dict, save_index='NAME') return output_df
def test_no_tunes_in_files_plot(tmp_path, file_path, bpms): from glob import glob for f in glob(f"{file_path}*"): print(f) for f in INPUT_DIR.glob(f"{file_path.name}*"): copy(f, tmp_path) file_path = tmp_path / file_path.name for p in PLANES: fname = file_path.with_suffix(f"{file_path.suffix}.lin{p.lower()}") df = tfs.read(fname) tfs.write(fname, df.drop(columns=[f"TUNE{p.upper()}", f"NATTUNE{p.upper()}"])) plot_spectrum( files=[file_path], bpms=bpms, combine_by=["files", "bpms"], )
def _calculate_normalised_dispersion_3d(meas_input, input_files, beta, header): """It computes horizontal normalised dispersion from 3 D kicks, it performs model based compensation, i.e. as in _free2 files""" output, accelerator = meas_input.outputdir, meas_input.accelerator model = accelerator.model driven_model = accelerator.model_driven if accelerator.excitation else model plane = "X" df_orbit = _get_merged_df(meas_input, input_files, plane, ['AMPZ', 'MUZ']) df_orbit[f"ND{plane}{MDL}"] = df_orbit.loc[:, f"D{plane}{MDL}"] / np.sqrt( df_orbit.loc[:, f"BET{plane}{MDL}"]) df_orbit = pd.merge(df_orbit, beta.loc[:, [f"BET{plane}", f"{ERR}BET{plane}"]], how='inner', left_index=True, right_index=True) df_orbit = pd.merge(df_orbit, driven_model.loc[:, [f"BET{plane}"]], how='inner', left_index=True, right_index=True, suffixes=('', '_driven')) mask = accelerator.get_element_types_mask(df_orbit.index, ["arc_bpm"]) compensation = np.sqrt(df_orbit.loc[:, f"BET{plane}_driven"].values / df_orbit.loc[:, f"BET{plane}{MDL}"].values) global_factors = np.sum( df_orbit.loc[mask, f"ND{plane}{MDL}"].values) / np.sum( df_orbit.loc[mask, input_files.get_columns(df_orbit, 'AMPZ')].values * compensation[mask, None], axis=0) # scaling to the model, and getting the synchrotron phase in the arcs scaled_amps = ( df_orbit.loc[:, input_files.get_columns(df_orbit, 'AMPZ')].values * global_factors) * compensation[:, None] df_orbit[f"ND{plane}"], df_orbit[ f"{ERR}ND{plane}"] = _get_signed_dispersion(input_files, df_orbit, scaled_amps, mask) df_orbit = _calculate_from_norm_disp(df_orbit, model, plane) output_df = df_orbit.loc[:, _get_output_columns(plane, df_orbit)] tfs.write(join(output, f"{NORM_DISP_NAME}{plane.lower()}{EXT}"), output_df, header, save_index='NAME') return output_df
def convert_old_beta_from_amplitude( inputdir: Union[Path, str], outputdir: Union[Path, str], suffix: str, plane: str, old_file_name: str = "ampbeta", new_file_name: str = AMP_BETA_NAME, ) -> None: """ Looks in the provided directory for expected beta from amplitude 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 getampbeta(x,y).out, with the following expected columns: NAME, S, COUNT, BETX, BETXSTD, BETXMDL, MUXMDL, BETXRES, BETXSTDRES. 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. suffix (str): Compensation suffix used in the provided BetaBeat.src output 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 beta from amplitude file") old_file_path = Path( inputdir) / f"get{old_file_name}{plane.lower()}{suffix}{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"BET{plane}STD": f"{ERR}BET{plane}", f"BET{plane}STDRES": f"{ERR}BET{plane}RES" }, ) dframe[f"{DELTA}BET{plane}"] = df_rel_diff(dframe, f"BET{plane}", f"BET{plane}{MDL}") dframe[f"{ERR}{DELTA}BET{plane}"] = df_ratio(dframe, f"{ERR}BET{plane}", f"BET{plane}{MDL}") tfs.write(Path(outputdir) / f"{new_file_name}{plane.lower()}{EXT}", dframe)
def calculate_orbit(meas_input, input_files, header, plane): """ Calculates orbit. Args: meas_input: `OpticsInput` object input_files: Stores the input files tfs. header: `OrderedDict` containing information about the analysis. plane: marking the horizontal or vertical plane, **X** or **Y**. Returns: `TfsDataFrame` corresponding to output file. """ df_orbit = _get_merged_df(meas_input, input_files, plane, ['CO', 'CORMS']) df_orbit[plane] = stats.weighted_mean(input_files.get_data(df_orbit, 'CO'), axis=1) df_orbit[f"{ERR}{plane}"] = stats.weighted_error(input_files.get_data(df_orbit, 'CO'), axis=1) df_orbit = _get_delta_columns(df_orbit, plane) output_df = df_orbit.loc[:, _get_output_columns(plane, df_orbit)] tfs.write(join(meas_input.outputdir, f"{ORBIT_NAME}{plane.lower()}{EXT}"), output_df, header, save_index='NAME') return output_df
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, "madx_macros") 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 calculate(measure_input, input_files, scale, header_dict, plane): """ Args: measure_input: Optics_input object input_files: Stores the input files tfs scale: measured beta functions header_dict: OrderedDict containing information about the analysis plane: "X" or "Y" Returns: DataFrame containing actions and their errors """ try: kick_frame = _getkick(measure_input, input_files, plane) except IndexError: # occurs if either no x or no y files exist return pd.DataFrame kick_frame = _rescale_actions(kick_frame, scale, plane) header = _get_header(header_dict, plane, scale) tfs.write(join(measure_input.outputdir, f"{KICK_NAME}{plane.lower()}{EXT}"), kick_frame, header) return kick_frame.loc[:, [f"{SQRT_ACTION}{plane}", f"{ERR}{SQRT_ACTION}{plane}"]].to_numpy()
def calculate(measure_input, input_files, scale, header_dict, plane): """ Args: measure_input: `OpticsInput` object. input_files: Stores the input files tfs. scale: measured beta functions. header_dict: `OrderedDict` containing information about the analysis. plane: marking the horizontal or vertical plane, **X** or **Y**. Returns: `TfsDataFrame` containing actions and their errors. """ try: kick_frame = _getkick(measure_input, input_files, plane) except IndexError: # occurs if either no x or no y files exist return pd.DataFrame kick_frame = _rescale_actions(kick_frame, scale, plane) header = _get_header(header_dict, plane, scale) tfs.write(join(measure_input.outputdir, f"{KICK_NAME}{plane.lower()}{EXT}"), kick_frame, header) return kick_frame.loc[:, [f"{SQRT_ACTION}{plane}", f"{ERR}{SQRT_ACTION}{plane}"]].to_numpy()
def get_rdts(madx: Madx, order: int = 4, file: Union[Path, str] = None) -> tfs.TfsDataFrame: """ INITIAL IMPLEMENTATION CREDITS GO TO JOSCHUA DILLY (@JoschD). Calculate the RDTs via PTC_TWISS. Args: madx (cpymad.madx.Madx): an instanciated cpymad Madx object. order (int): maximum derivative order (only 0, 1 or 2 implemented in PTC). Defaults to `2`. file (Union[Path, str]): path to output file. Default `None` Returns: A TfsDataframe with results. """ logger.info(f"Entering PTC to calculate RDTs up to order {order}") madx.ptc_create_universe() logger.trace("Creating PTC layout") madx.ptc_create_layout(model=3, method=4, nst=3, exact=True) # madx.ptc_create_layout(model=3, method=6, nst=1) # from Michi logger.trace("Incorporating MAD-X alignment errors") madx.ptc_align() # use madx alignment errors # madx.ptc_setswitch(fringe=True) # include fringe effects logger.debug("Executing PTC Twiss") madx.ptc_twiss(icase=6, no=order, normal=True, trackrdts=True) madx.ptc_end() logger.debug("Extracting results to TfsDataFrame") dframe = tfs.TfsDataFrame(madx.table.twissrdt.dframe()) dframe.columns = dframe.columns.str.upper() dframe.NAME = dframe.NAME.str.upper() if file: logger.debug(f"Exporting results to disk at '{Path(file).absolute()}'") tfs.write(file, dframe) return dframe
def _calculate_dispersion_2d(meas_input, input_files, header, plane): order = 2 if meas_input.second_order_dispersion else 1 dpps = input_files.dpps(plane) if np.max(dpps) - np.min(dpps) == 0.0: return # temporary solution model = meas_input.accelerator.model df_orbit = _get_merged_df(meas_input, input_files, plane, ['CO', 'CORMS']) fit = np.polyfit(dpps, input_files.get_data(df_orbit, 'CO').T, order, cov=True) # in the fit results the coefficients are sorted by power in decreasing order if order > 1: df_orbit[f"D2{plane}"] = fit[0][-3, :].T df_orbit[f"{ERR}D2{plane}"] = np.sqrt(fit[1][-3, -3, :].T) df_orbit[f"D{plane}"] = fit[0][-2, :].T df_orbit[f"{ERR}D{plane}"] = np.sqrt(fit[1][-2, -2, :].T) df_orbit[plane] = fit[0][-1, :].T df_orbit[f"{ERR}{plane}"] = np.sqrt(fit[1][-1, -1, :].T) # since we get variances from the fit, maybe we can include the variances of fitted points df_orbit[f"DP{plane}"] = _calculate_dp(model, df_orbit.loc[:, [f"D{plane}", f"{ERR}D{plane}"]], plane) df_orbit = _get_delta_columns(df_orbit, plane) output_df = df_orbit.loc[:, _get_output_columns(plane, df_orbit)] tfs.write(join(meas_input.outputdir, f"{DISPERSION_NAME}{plane.lower()}{EXT}"), output_df, header, save_index='NAME') return output_df
def merge_kmod_results(opt) -> tfs.TfsDataFrame: """ Main function to merge K-Mod output. See :mod:`omc3.scripts.merge_kmod_results`. """ if opt.outputdir: save_config(opt.outputdir, opt, script=__file__) # creates outputdir # Get the directories we need where the tfs are stored ip_dir_names: List[pathlib.Path] = get_kmod_ip_subdirs(opt.kmod_dirs) # Combine the lsa data lsa_tfs = merge_tfs(ip_dir_names, f"{LSA_RESULTS}{EXT}") # If the TFS data contains everything we need: get the imbalance if _validate_for_imbalance(lsa_tfs): imbalance_results = get_lumi_imbalance(lsa_tfs) lsa_tfs = _add_imbalance_to_header(lsa_tfs, *imbalance_results) if opt.outputdir: output_path = opt.outputdir / f"{LSA_RESULTS}{EXT}" LOG.info(f"Writing result TFS file to disk at: {str(output_path)}") tfs.write(output_path, lsa_tfs, save_index=NAME) return lsa_tfs
def analyse_kmod(opt): """ Run Kmod analysis """ LOG.info('Getting input parameter') if opt.interaction_point is None and opt.circuits is None: raise AttributeError('No IP or circuits specified, stopping analysis') if opt.interaction_point is not None and opt.circuits is not None: raise AttributeError('Both IP and circuits specified, choose only one, stopping analysis') if not 1 < len(opt.betastar_and_waist) < 5: raise AttributeError("Option betastar_and_waist has to consist of 2 to 4 floats") opt.betastar_and_waist = convert_betastar_and_waist(opt.betastar_and_waist) for error in ("cminus", "errorK", "errorL", "misalignment"): opt = check_default_error(opt, error) if opt.measurement_dir is None and opt.model_dir is None and opt.phase_weight: raise AttributeError("Cannot use phase advance without measurement or model") if opt.outputdir is None: opt.outputdir = opt.working_directory LOG.info(f"{'IP trim' if opt.interaction_point is not None else 'Individual magnets'} analysis") opt['magnets'] = MAGNETS_IP[opt.interaction_point.upper()] if opt.interaction_point is not None else [ find_magnet(opt.beam, circuit) for circuit in opt.circuits] opt['label'] = f'{opt.interaction_point}{opt.beam}' if opt.interaction_point is not None else f'{opt.magnets[0]}-{opt.magnets[1]}' opt['instruments'] = list(map(str.upper, opt.instruments.split(","))) output_dir = join(opt.outputdir, opt.label) iotools.create_dirs(output_dir) LOG.info('Get inputfiles') magnet1_df, magnet2_df = helper.get_input_data(opt) opt, magnet1_df, magnet2_df, betastar_required = define_params(opt, magnet1_df, magnet2_df) LOG.info('Run simplex') magnet1_df, magnet2_df, results_df, instrument_beta_df = analysis.analyse(magnet1_df, magnet2_df, opt, betastar_required) LOG.info('Plot tunes and fit') if opt.no_plots: helper.plot_cleaned_data([magnet1_df, magnet2_df], join(output_dir, FIT_PLOTS_NAME), interactive_plot=False) LOG.info('Write magnet dataframes and results') for magnet_df in [magnet1_df, magnet2_df]: tfs.write(join(output_dir, f"{magnet_df.headers['QUADRUPOLE']}{EXT}"), magnet_df) tfs.write(join(output_dir, f'{RESULTS_FILE_NAME}{EXT}'), results_df) if opt.instruments_found: tfs.write(join(output_dir, f'{INSTRUMENTS_FILE_NAME}{EXT}'), instrument_beta_df) create_lsa_results_file(betastar_required, opt.instruments_found, results_df, instrument_beta_df, output_dir)
def write(beta_df, header, outputdir, plane): tfs.write(os.path.join(outputdir, f"{BETA_NAME}{plane.lower()}{EXT}"), beta_df, header, save_index="NAME")