def calculate_chromatic_coupling(couplings, dpps, input_files, measure_input):
# TODO how to treat the model values?
columns = [
f"{pref}{col}{part}" for pref in ("", ERR)
for col in ("F1001", "F1010") for part in ("RE", "IM")
]
joined = couplings[0].loc[:, columns]
for i, coup in enumerate(couplings[1:]):
joined = pd.merge(joined,
coup.loc[:, columns],
how="inner",
left_index=True,
right_index=True,
suffixes=('', '__' + str(i + 1)))
for column in columns:
joined.rename(columns={column: column + '__0'}, inplace=True)
for col in ("F1001", "F1010"):
for part in ("RE", "IM"):
fit = np.polyfit(np.repeat(dpps, 2),
np.repeat(input_files.get_data(
joined, f"{col}{part}").T,
2,
axis=0),
1,
cov=True)
joined[f"D{col}{part}"] = fit[0][-2, :].T
joined[f"{ERR}D{col}{part}"] = np.sqrt(fit[1][-2, -2, :].T)
joined[f"D{col}"] = np.sqrt(
np.square(joined.loc[:, f"D{col}RE"].to_numpy()) +
np.square(joined.loc[:, f"D{col}IM"].to_numpy()))
joined[f"{ERR}D{col}"] = np.sqrt(
np.square(joined.loc[:, f"D{col}RE"].to_numpy() *
df_ratio(joined, f"{ERR}D{col}RE", f"D{col}")) +
np.square(joined.loc[:, f"D{col}IM"].to_numpy() *
df_ratio(joined, f"{ERR}D{col}IM", f"D{col}")))
output_df = pd.merge(measure_input.accelerator.model.loc[:, ["S"]],
joined.loc[:, [
f"{pref}{col}{part}" for pref in ("", ERR)
for col in ("F1001", "F1010")
for part in ("", "RE", "IM")
]],
how="inner",
left_index=True,
right_index=True)
return output_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 _get_delta_columns(beta_df, plane):
beta_df[f"{DELTA}BET{plane}"] = df_rel_diff(beta_df, f"BET{plane}",
f"BET{plane}{MDL}")
beta_df[f"{ERR}{DELTA}BET{plane}"] = df_ratio(beta_df, f"{ERR}BET{plane}",
f"BET{plane}{MDL}")
beta_df[f"{DELTA}ALF{plane}"] = df_diff(beta_df, f"ALF{plane}",
f"ALF{plane}{MDL}")
beta_df[
f"{ERR}{DELTA}ALF{plane}"] = beta_df.loc[:,
f"{ERR}ALF{plane}"].to_numpy(
)
return beta_df
def phase_to_amp_ratio(measure_input, beta_phase, beta_amp, plane):
ratio = pd.merge(beta_phase.loc[:, [f"BET{plane}"]],
beta_amp.loc[:, [f"BET{plane}"]],
how='inner',
left_index=True,
right_index=True,
suffixes=("ph", "amp"))
ph_over_amp = df_ratio(ratio, f"BET{plane}ph", f"BET{plane}amp")
mask = (np.array(0.1 < np.abs(ph_over_amp))
& np.array(np.abs(ph_over_amp) < 10.0) & np.array(
measure_input.accelerator.get_element_types_mask(
ratio.index, ["arc_bpm"])))
x_ratio = np.mean(ph_over_amp[mask])
return x_ratio
def _compensate_by_model(input_files, meas_input, df, plane):
df = pd.merge(
df,
pd.DataFrame(
meas_input.accelerator.model_driven.loc[:, [f"BET{plane}"]].rename(
columns={f"BET{plane}": f"BET{plane}comp"})),
how='inner',
left_index=True,
right_index=True)
amp_compensation = np.sqrt(
df_ratio(df, f"BET{plane}{MDL}", f"BET{plane}comp"))
df[input_files.get_columns(
df, f"AMP{plane}")] = (input_files.get_data(df, f"AMP{plane}") *
amp_compensation[:, np.newaxis])
return 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 beta_from_amplitude(meas_input, input_files, plane, tunes):
df = pd.DataFrame(
meas_input.accelerator.model).loc[:,
["S", f"MU{plane}", f"BET{plane}"]]
df.rename(columns={
f"MU{plane}": f"MU{plane}{MDL}",
f"BET{plane}": f"BET{plane}{MDL}"
},
inplace=True)
dpp_value = meas_input.dpp if "dpp" in meas_input.keys() else 0
df = pd.merge(df,
input_files.joined_frame(plane,
[f"AMP{plane}", f"MU{plane}"],
dpp_value=dpp_value),
how='inner',
left_index=True,
right_index=True)
df['COUNT'] = len(input_files.get_columns(df, f"AMP{plane}"))
if meas_input.compensation == "model":
df = _compensate_by_model(input_files, meas_input, df, plane)
if meas_input.compensation == "equation":
df = _compensate_by_equation(input_files, meas_input, df, plane, tunes)
amps_squared = np.square(input_files.get_data(df, f"AMP{plane}"))
mask = meas_input.accelerator.get_element_types_mask(df.index, ["arc_bpm"])
actions = amps_squared / df.loc[:, f"BET{plane}{MDL}"].values[:,
np.newaxis]
betas = amps_squared / np.mean(actions[mask], axis=0, keepdims=True)
df[f"BET{plane}"] = np.mean(betas, axis=1)
df[f"{ERR}BET{plane}"] = np.std(betas, axis=1)
df[f"{DELTA}BET{plane}"] = df_rel_diff(df, f"BET{plane}",
f"BET{plane}{MDL}")
df[f"{ERR}{DELTA}BET{plane}"] = df_ratio(df, f"{ERR}BET{plane}",
f"BET{plane}{MDL}")
return df.loc[:, [
'S', 'COUNT', f"BET{plane}", f"{ERR}BET{plane}", f"BET{plane}{MDL}",
f"MU{plane}{MDL}", f"{DELTA}BET{plane}", f"{ERR}{DELTA}BET{plane}"
]]
def append_model(df: pd.DataFrame,
df_model: pd.DataFrame,
parameter: str,
planes: str = '',
beat: bool = False) -> pd.DataFrame:
""" Add the values to the measurement. """
LOG.debug(f"Appending model to fake measurement for {parameter}.")
df[S] = df_model[S]
for plane in planes:
df[f'{PHASE_ADV}{plane}{MDL}'] = df_model[f'{PHASE_ADV}{plane}']
df[f"{parameter}{MDL}"] = df_model[f'{parameter}']
if beat:
df[f"{DELTA}{parameter}"] = df_rel_diff(df, parameter,
f"{parameter}{MDL}")
df[f"{ERR}{DELTA}{parameter}"] = df_ratio(df, f"{ERR}{parameter}",
f"{parameter}{MDL}")
else:
df[f"{DELTA}{parameter}"] = df_diff(df, f'{parameter}',
f'{parameter}{MDL}')
df[f"{ERR}{DELTA}{parameter}"] = df[f'{ERR}{parameter}']
return df
def test_df_prod_ratio():
a, b = _arand(), _arand()
prod = toolbox.df_prod(*_df(a, b))
ratio = toolbox.df_ratio(*_df(prod, b))
assert _numerically_equal(a, ratio)
def test_df_ratio_zero(random, zeros):
assert not sum(toolbox.df_ratio(*_df(zeros, random)))
with pytest.warns(RuntimeWarning):
toolbox.df_ratio(*_df(random, zeros))
def test_df_ratio_one(random, ones):
assert all(random == toolbox.df_ratio(*_df(random, ones)))