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()