def _maybe_add_coupling_to_model(model, keys):
if any([key for key in keys if key.startswith("F1")]):
tw_opt = TwissOptics(model)
couple = tw_opt.get_coupling(method="cmatrix")
model["F1001R"] = couple["F1001"].apply(np.real).astype(np.float64)
model["F1001I"] = couple["F1001"].apply(np.imag).astype(np.float64)
model["F1010R"] = couple["F1010"].apply(np.real).astype(np.float64)
model["F1010I"] = couple["F1010"].apply(np.imag).astype(np.float64)
return model
def _add_coupling(dict_of_tfs):
""" Adds coupling to the tfs. QUICK FIX VIA LOOP!"""
with timeit(lambda t: LOG.debug(" Time adding coupling: {:f}s".format(t))):
for var in dict_of_tfs:
twopt = TwissOptics(dict_of_tfs[var])
cpl = twopt.get_coupling("cmatrix")
dict_of_tfs[var]["1001"] = cpl["F1001"]
dict_of_tfs[var]["1010"] = cpl["F1010"]
return dict_of_tfs
def _add_coupling(dict_of_tfs):
""" Adds coupling to the tfs. QUICK FIX VIA LOOP!"""
with timeit(
lambda t: LOG.debug(" Time adding coupling: {:f}s".format(t))):
for var in dict_of_tfs:
twopt = TwissOptics(dict_of_tfs[var], keep_all_elem=True)
cpl = twopt.get_coupling("cmatrix")
dict_of_tfs[var]["1001"] = cpl["F1001"]
dict_of_tfs[var]["1010"] = cpl["F1010"]
return dict_of_tfs
def test_calculate_cmatrix(df):
df.loc[:, "R11"] = np.sin(df["R"])
df.loc[:, "R22"] = df["R11"]
df.loc[:, "R21"] = np.cos(df["R"])
df.loc[:, "R12"] = -df["R21"]
df = _pd_to_tfs(df, (0, 0))
to = TwissOptics(df.copy())
with suppress_warnings(RuntimeWarning):
cmat = to.get_cmatrix()
couple = to.get_coupling(method="cmatrix")
gamma = to.get_gamma()
assert all([c in cmat for c in ['S', 'C11', 'C12', 'C21', 'C22']])
assert all([c in couple for c in ['S', 'F1001', 'F1010']])
assert all([c in gamma for c in ['S', 'GAMMA']])
assert len(cmat.index.values) == len(df.index.values)
def test_missing_columns():
df = pd.DataFrame(index=["BPM" + str(i) for i in range(10)],
columns=["S", "BETX", "BETY"])
to = TwissOptics(df)
with pytest.raises(KeyError):
to.get_coupling()