def _compensate_by_model(input_files, meas_input, df, plane):
df = pd.merge(df, pd.DataFrame(meas_input.accelerator.model_driven.loc[:, [f"MU{plane}"]]),
how='inner', left_index=True, right_index=True, suffixes=("", "comp"))
phase_compensation = df_diff(df, f"MU{plane}", f"MU{plane}comp")
df[input_files.get_columns(df, f"MU{plane}")] = ang_sum(
input_files.get_data(df, f"MU{plane}"), phase_compensation[:, np.newaxis])
return df
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_ang():
a, b = _arand(), _arand()
diff = toolbox.df_ang_diff(*_df(a, b))
sum_angles = toolbox.ang_sum(a, b)
for res in (diff, sum_angles):
assert len(res) == len(a)
assert all(-0.5 <= res)
assert all(res <= 0.5)