def test_prepare_twiss_dataframe_beams():
n_index, n_kmax, n_valmax = 5, 6, 10
df_twiss, df_errors = get_twiss_and_error_df(n_index, n_kmax, n_valmax)
df1 = prepare_twiss_dataframe(df_twiss=df_twiss, df_errors=df_errors)
df2 = prepare_twiss_dataframe(df_twiss=df_twiss, df_errors=df_errors)
df4 = prepare_twiss_dataframe(df_twiss=df_twiss,
df_errors=df_errors,
invert_signs_madx=True)
assert df1.equals(df2)
assert not df1.equals(df4)
def test_prepare_twiss_dataframe():
n_index, n_kmax, n_valmax = 5, 6, 10
n_kmax_prepare = 16
df_twiss, df_errors = get_twiss_and_error_df(n_index, n_kmax, n_valmax)
df = prepare_twiss_dataframe(df_twiss=df_twiss,
df_errors=df_errors,
max_order=n_kmax_prepare)
assert all(df[S] == df_twiss[S])
k_columns = df.columns[df.columns.str.match(r"^K\d+S?L")]
assert len(k_columns) == n_kmax_prepare * 2
for col in k_columns:
if col in df_twiss.columns and col in df_errors.columns:
assert all(df[col] == (df_twiss[col] + df_errors[col]))
elif col in df_twiss.columns:
assert all(df[col] == df_twiss[col])
elif col in df_errors.columns:
assert all(df[col] == df_errors[col])
else:
assert not any(df[col])
for plane in PLANES:
assert all(df[plane] == (df_twiss[plane] +
df_errors[f"{DELTA_ORBIT}{plane}"]))
assert df.headers == df_twiss.headers
assert not df.isna().any().any()
def _coupling_bump_teapot_cta() -> float:
"""Compute and return the CTA for teapot reference method on the lhcb1 coupling bump test case."""
df_twiss = tfs.read(COUPLING_BUMP_TWISS_BEAM_1, index=NAME)
df = prepare_twiss_dataframe(df_twiss=df_twiss, max_order=7)
df_cmatrix = coupling_via_cmatrix(df)
df_twiss[F1001] = df_cmatrix[
F1001] # ignoring F1010 in this test as it is bigger than F1001
cta_df = closest_tune_approach(df_twiss,
method="teapot") # only one column
return cta_df.mean().abs()[0] # this is the cminus
def test_rdts_skew_octupole_bump():
df = get_df(n=8)
df = prepare_twiss_dataframe(df_twiss=df)
df.loc["3", "K3SL"] = 1
df.loc["5", "K3SL"] = -1
df_rdts = calculate_rdts(df, rdts=["F4000", "F3001"])
df_diff, df_jump = get_absdiff_and_jumps(df_rdts)
assert not df_rdts.isna().any().any()
assert not df_jump["F4000"].any()
assert df_jump.loc[["3", "5"], "F3001"].all()
assert not df_jump.loc[df_jump.index.difference(["3", "5"]), "F3001"].any()
def test_rdts_normal_sextupole_bump():
df = get_df(n=7)
df = prepare_twiss_dataframe(df_twiss=df)
df.loc["3", "K2L"] = 1
df.loc["5", "K2L"] = -1
df_rdts = calculate_rdts(df, rdts=["F1002", "F2001"])
df_diff, df_jump = get_absdiff_and_jumps(df_rdts)
assert not df_rdts.isna().any().any()
assert not df_jump["F2001"].any()
assert df_jump.loc[["3", "5"], "F1002"].all()
assert not df_jump.loc[df_jump.index.difference(["3", "5"]), "F1002"].any()
def test_rdts_save_memory():
np.random.seed(2047294792)
n = 10
df = get_df(n=n)
df = prepare_twiss_dataframe(df_twiss=df)
df.loc[:, "K2L"] = np.random.rand(n)
df.loc[:, "K2SL"] = np.random.rand(n)
df_save = calculate_rdts(df, rdts=["F1002", "F2001"], loop_phases=True)
df_notsave = calculate_rdts(df, rdts=["F1002", "F2001"], loop_phases=False)
assert not df_save.isna().any().any()
assert ((df_save - df_notsave).abs() < 1e-14).all().all()
def test_wrong_rdt_names():
df = get_df(n=8)
df = prepare_twiss_dataframe(df_twiss=df)
with pytest.raises(ValueError):
calculate_rdts(df, rdts=["F10003", "F1002"])
with pytest.raises(ValueError):
calculate_rdts(df, rdts=["F102", "F1002"])
with pytest.raises(ValueError):
calculate_rdts(df, rdts=["M1003"])
with pytest.raises(ValueError):
calculate_rdts(df, rdts=["F1000"])
def test_rdts_normal_dodecapole_to_octupole_feeddown():
df = get_df(n=7)
df = prepare_twiss_dataframe(df_twiss=df)
df_comp = df.copy()
df.loc["3", "K5L"] = 1
df.loc["5", "K5L"] = -1
df_comp.loc["3", ["K3L", "K3SL"]] = 1
df_comp.loc["5", ["K3L", "K3SL"]] = -1
df_rdts = calculate_rdts(df, rdts=["F1003", "F0004"], feeddown=2)
df_rdts_comp = calculate_rdts(df_comp, rdts=["F1003", "F0004"])
assert not df_rdts["F1003"].any()
assert not df_rdts["F0004"].any()
# Feed-down K5L -> (x**2-y**2)/2 K3L , xy K3SL
# Feed-down K5L -> K3L
df["X"] = 1
df["Y"] = 0
df_rdts = calculate_rdts(df, rdts=["F1003", "F0004"], feeddown=2)
assert not df_rdts["F1003"].any()
assert df_rdts["F0004"].all()
assert all(df_rdts["F0004"] == 0.5 * df_rdts_comp["F0004"])
# Feed-down K5L -> K3L
df["X"] = 0
df["Y"] = 1
df_rdts = calculate_rdts(df, rdts=["F1003", "F0004"], feeddown=2)
assert not df_rdts["F1003"].any()
assert df_rdts["F0004"].all()
assert all(df_rdts["F0004"] == -0.5 * df_rdts_comp["F0004"])
# Feed-down K5L -> K3SL
df["X"] = 1
df["Y"] = 1
df_rdts = calculate_rdts(df, rdts=["F1003", "F0004"], feeddown=2)
assert not df_rdts["F0004"].any()
assert df_rdts["F1003"].all()
assert all(df_rdts["F1003"] == df_rdts_comp["F1003"])
# Feed-down K5L -> (x**2-y**2)/2 K3L , xy K3SL
df["X"] = 1
df["Y"] = 0.5
df_rdts = calculate_rdts(df, rdts=["F1003", "F0004"], feeddown=2)
assert all(
(df_rdts["F0004"] - 0.375 * df_rdts_comp["F0004"]).abs() <= 1e-15)
assert all(df_rdts["F1003"] == 0.5 * df_rdts_comp["F1003"])
def test_coupling_rdt_bump_cmatrix_compare():
df_twiss = tfs.read(COUPLING_BUMP_TWISS_BEAM_1, index=NAME)
df = prepare_twiss_dataframe(df_twiss=df_twiss, max_order=7)
df_rdts = coupling_via_rdts(df)
df_cmatrix = coupling_via_cmatrix(df)
# from tests.unit.debug_helper import plot_rdts_vs
# plot_rdts_vs(df_rdts, "analytical", df_cmatrix, "cmatrix", df_twiss, ["F1001", "F1010"])
for rdt in COUPLING_RDTS:
assert arrays_are_close_almost_everywhere(df_rdts[rdt],
df_cmatrix[rdt],
rtol=1e-3,
atol=1e-4,
percentile=0.99)
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_prepare_twiss_dataframe_inner():
n_index, n_kmax, n_valmax = 5, 6, 10
n_kmax_prepare = 5
df_twiss, _ = get_twiss_and_error_df(n_index, n_kmax, n_valmax)
_, df_errors_1 = get_twiss_and_error_df(n_index + 3, n_kmax, n_valmax)
df = prepare_twiss_dataframe(df_twiss=df_twiss,
df_errors=df_errors_1.iloc[3:, :],
max_order=n_kmax_prepare,
join="inner")
k_columns = df.columns[df.columns.str.match(r"^K\d+S?L")]
assert len(k_columns) == n_kmax * 2
assert len(df.index) == n_index - 3
assert all(df[S] == df_twiss.loc[df.index, S])
assert df.headers == df_twiss.headers
assert not df.isna().any().any()
def test_prepare_twiss_dataframe_outer():
n_index, n_kmax, n_valmax = 8, 6, 10
n_kmax_prepare = 16
_, df_errors = get_twiss_and_error_df(n_index, n_kmax, n_valmax)
df_twiss_1, _ = get_twiss_and_error_df(n_index, n_kmax + 2, n_valmax)
df = prepare_twiss_dataframe(df_twiss=df_twiss_1.iloc[3:, :],
df_errors=df_errors.iloc[:-3, :],
max_order=n_kmax_prepare,
join="outer")
k_columns = df.columns[df.columns.str.match(r"^K\d+S?L")]
assert len(k_columns) == n_kmax_prepare * 2
assert all(df.index == df_twiss_1.index)
assert all(df[S] == df_twiss_1.loc[df.index, S])
assert df.headers == df_twiss_1.headers
assert not df.isna().any().any()
def test_closest_tune_approach(cta_method, max_relative_error_to_teapot,
result_should_be_real,
_coupling_bump_teapot_cta):
df_twiss = tfs.read(COUPLING_BUMP_TWISS_BEAM_1, index=NAME)
df = prepare_twiss_dataframe(df_twiss=df_twiss, max_order=7)
df_cmatrix = coupling_via_cmatrix(df)
df_twiss[F1001] = df_cmatrix[
F1001] # ignoring F1010 in this test as it is bigger than F1001
cta_df = closest_tune_approach(df_twiss,
method=cta_method) # only one column
cminus = cta_df.mean().abs()[0]
relative_error = _relative_error(cminus, _coupling_bump_teapot_cta)
assert relative_error <= max_relative_error_to_teapot
assert not cta_df.isna().any().any() # check no NaNs
assert all(cta_df[df_cmatrix[F1001] != 0] != 0)
if result_should_be_real:
assert all(np.isreal(cta_df))
def test_real_terms_and_hamiltonians():
np.random.seed(2047294792)
n = 12
df = get_df(n=n)
df = prepare_twiss_dataframe(df_twiss=df)
df.loc[:, "K2L"] = np.random.rand(n)
df.loc[:, "K2SL"] = np.random.rand(n)
rdts = ["F1002", "F2001"]
df_rdts = calculate_rdts(df,
rdts=rdts,
complex_columns=False,
hamiltionian_terms=True)
n_rdts = len(rdts)
assert all(np.real(df_rdts) == df_rdts)
assert len(df_rdts.columns) == 4 * n_rdts
assert df_rdts.columns.str.match(f".+{REAL}$").sum() == 2 * n_rdts
assert df_rdts.columns.str.match(f".+{IMAG}$").sum() == 2 * n_rdts
assert df_rdts.columns.str.match("F").sum() == 2 * n_rdts
assert df_rdts.columns.str.match("H").sum() == 2 * n_rdts
def test_real_output():
n = 7
np.random.seed(474987942)
df = generate_fake_data(n)
df = prepare_twiss_dataframe(df_twiss=df)
df.loc[:, "K1L"] = np.random.rand(n)
df.loc[:, "K1SL"] = np.random.rand(n)
df_cmatrix = coupling_via_cmatrix(df, complex_columns=False)
df_rdts = coupling_via_rdts(df, qx=1.31, qy=1.32, complex_columns=False)
assert all(np.real(df_cmatrix) == df_cmatrix)
assert all(np.real(df_rdts) == df_rdts)
columns = [f"{c}{r}" for c in COUPLING_RDTS for r in (REAL, IMAG)]
assert df_rdts[columns].all().all() # check no 0-values
assert df_cmatrix[columns].all().all() # check no 0-values
assert df_cmatrix.columns.str.match(f".+{REAL}$").sum() == 2
assert df_cmatrix.columns.str.match(f".+{IMAG}$").sum() == 2
assert df_rdts.columns.str.match(f".+{REAL}$").sum() == 2
assert df_rdts.columns.str.match(f".+{IMAG}$").sum() == 2
def test_rdts_normal_octupole_to_sextupole_feeddown():
df = get_df(n=15)
df = prepare_twiss_dataframe(df_twiss=df)
df_comp = df.copy()
df.loc["3", "K3L"] = 1
df.loc["5", "K3L"] = -1
df_comp.loc["3", ["K2L", "K2SL"]] = 1
df_comp.loc["5", ["K2L", "K2SL"]] = -1
df_rdts = calculate_rdts(df, rdts=["F1002", "F2001"], feeddown=1)
df_rdts_comp = calculate_rdts(df_comp, rdts=["F1002", "F2001"])
assert not df_rdts["F2001"].any()
assert not df_rdts["F1002"].any()
# Feed-down K3L -> K2L
df["X"] = 1
df["Y"] = 0
df_rdts = calculate_rdts(df, rdts=["F1002", "F2001"], feeddown=1)
assert not df_rdts["F2001"].any()
assert df_rdts["F1002"].all()
assert all(df_rdts["F1002"] == df_rdts_comp["F1002"])
# Feed-down K3L -> K2SL
df["X"] = 0
df["Y"] = 1
df_rdts = calculate_rdts(df, rdts=["F1002", "F2001"], feeddown=1)
assert not df_rdts["F1002"].any()
assert df_rdts["F2001"].all()
assert all(df_rdts["F2001"] == df_rdts_comp["F2001"])
# Feed-down K3L -> K2L, K2SL
df["X"] = 1
df["Y"] = 1
df_rdts = calculate_rdts(df, rdts=["F1002", "F2001"], feeddown=1)
assert all(df_rdts["F1002"] == df_rdts_comp["F1002"])
assert all(df_rdts["F2001"] == df_rdts_comp["F2001"])
def test_prepare_twiss_dataframe_no_error():
n_index, n_kmax, n_valmax = 5, 3, 10
df_twiss, _ = get_twiss_and_error_df(n_index, n_kmax, n_valmax)
df = prepare_twiss_dataframe(df_twiss=df_twiss, max_order=n_kmax)
assert df_twiss.equals(df)
