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 coupling_via_rdts(df: TfsDataFrame, complex_columns: bool = True, **kwargs) -> TfsDataFrame:
"""Returns the coupling terms.
.. warning::
This function changes sign of the real part of the RDTs compared to
[FranchiAnalyticFormulas2017]_ to be consistent with the RDT
calculations from [CalagaBetatronCoupling2005]_ .
Args:
df (TfsDataFrame): Twiss Dataframe.
complex_columns (bool): Output complex values in single column of type complex.
If ``False``, split complex columns into
two real-valued columns.
Keyword Args:
**kwargs: Remaining arguments from :func:`~optics_functions.rdt.rdts`
i.e. ``qx``, ``qy``, ``feeddown``, ``loop_phases``
and ``hamiltionian_terms``.
Returns:
A new ``TfsDataFrame`` with Coupling Columns.
"""
df_res = calculate_rdts(df, rdts=COUPLING_RDTS, **kwargs)
for rdt in COUPLING_RDTS:
rdt_array = df_res[rdt].to_numpy() # might return a copy!
rdt_array.real *= -1 # definition
df_res.loc[:, rdt] = rdt_array
if not complex_columns:
df_res = split_complex_columns(df_res, COUPLING_RDTS)
return df_res
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_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_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_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_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_missing_columns():
df = get_df(n=8)
with pytest.raises(KeyError) as e:
calculate_rdts(df, rdts=["F1002"])
assert "K2L" in str(e.value)
df["K2L"] = 1
with pytest.raises(KeyError) as e:
calculate_rdts(df, rdts=["F1002", "F2001"])
assert "K2SL" in str(e.value)
df["K2SL"] = 1
with pytest.raises(KeyError) as e:
calculate_rdts(df, rdts=["F2001", "F1002"], feeddown=1)
assert "K3L" in str(e.value)
df = df.drop(columns=[f"{PHASE_ADV}{X}"])
with pytest.raises(KeyError) as e:
calculate_rdts(df, rdts=["F2001", "F1002"])
assert f"{PHASE_ADV}{X}" in str(e.value)
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_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
