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")
