def test_rdt_calculation(df, q, order):
df = _pd_to_tfs(df, q)
rdt_names = get_all_rdts(order)
to = TwissOptics(df, quick_init=False)
with suppress_warnings(RuntimeWarning):
rdts = to.get_rdts(rdt_names)
assert all([c in rdts for c in ["S"] + rdt_names])
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 getdiff(meas_path=None, beta_file_name="getbeta"):
""" Calculates the differences between measurement, corrected and uncorrected model.
After running madx and creating the model with (twiss_cor) and without (twiss_no)
corrections, run this functions to get tfs-files with the differences between measurements
and models.
Args:
meas_path (str): Path to the measurement folder.
Needs to contain twiss_cor.dat and twiss_no.dat.
"""
if meas_path is None:
meas_path = sys.argv[1]
LOG.debug("Started 'getdiff' for measurment dir '{:s}'".format(meas_path))
if not isdir(meas_path):
raise IOError("No valid measurement directory:" + meas_path)
corrected_model_path = join(meas_path, TWISS_CORRECTED)
uncorrected_model_path = join(meas_path, TWISS_NOT_CORRECTED)
meas = OpticsMeasurement(meas_path)
twiss_cor = read_tfs(corrected_model_path).set_index('NAME')
twiss_no = read_tfs(uncorrected_model_path).set_index('NAME')
coup_cor = TwissOptics(twiss_cor,
quick_init=True).get_coupling(method='cmatrix')
coup_no = TwissOptics(twiss_no,
quick_init=True).get_coupling(method='cmatrix')
model = pd.merge(twiss_cor,
twiss_no,
how='inner',
left_index=True,
right_index=True,
suffixes=('_c', '_n'))
coupling_model = pd.merge(coup_cor,
coup_no,
how='inner',
left_index=True,
right_index=True,
suffixes=('_c', '_n'))
#coupling_model['NAME'] = coupling_model.index.values
for plane in ['x', 'y']:
_write_betabeat_diff_file(meas_path, meas, model, plane,
beta_file_name)
_write_phase_diff_file(meas_path, meas, model, plane)
_write_disp_diff_file(meas_path, meas, model, plane)
_write_closed_orbit_diff_file(meas_path, meas, model, plane)
_write_coupling_diff_file(meas_path, meas, coupling_model)
_write_norm_disp_diff_file(meas_path, meas, model)
_write_chromatic_coupling_files(meas_path, corrected_model_path)
_write_betastar_diff_file(meas_path, meas, twiss_cor, twiss_no)
LOG.debug("Finished 'getdiff'.")
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 _write_chromatic_coupling_files(meas_path, cor_path):
LOG.debug("Calculating chromatic coupling diff.")
# TODO: Add Cf1010
try:
twiss_plus = read_tfs(join(split(cor_path)[0], TWISS_CORRECTED_PLUS),
index='NAME')
twiss_min = read_tfs(join(split(cor_path)[0], TWISS_CORRECTED_MINUS),
index='NAME')
except IOError:
LOG.debug("Chromatic coupling measurements not found. Skipped.")
else:
deltap = np.abs(twiss_plus.DELTAP - twiss_min.DELTAP)
plus = TwissOptics(twiss_plus,
quick_init=True).get_coupling(method='cmatrix')
minus = TwissOptics(twiss_min,
quick_init=True).get_coupling(method='cmatrix')
model = pd.merge(plus,
minus,
how='inner',
left_index=True,
right_index=True,
suffixes=('_p', '_m'))
model['NAME'] = model.index.values
if exists(join(meas_path, "chromcoupling_free.out")):
meas = read_tfs(join(meas_path, "chromcoupling_free.out"))
else:
meas = read_tfs(join(meas_path, "chromcoupling.out"))
tw = pd.merge(meas, model, how='inner', on='NAME')
cf1001 = (tw.loc[:, 'F1001_p'] - tw.loc[:, 'F1001_m']) / deltap
tw['Cf1001r_model'] = np.real(cf1001)
tw['Cf1001i_model'] = np.imag(cf1001)
tw['Cf1001r_prediction'] = tw.loc[:,
'Cf1001r'] - tw.loc[:,
'Cf1001r_model']
tw['Cf1001i_prediction'] = tw.loc[:,
'Cf1001i'] - tw.loc[:,
'Cf1001i_model']
write_tfs(
join(meas_path, get_diff_filename('chromatic_coupling')),
tw.loc[:, [
'NAME', 'S', 'Cf1001r', 'Cf1001rERR', 'Cf1001i', 'Cf1001iERR',
'Cf1001r_model', 'Cf1001i_model', 'Cf1001r_prediction',
'Cf1001i_prediction'
]])
def test_empty_dataframe():
with pytest.raises(IndexError):
TwissOptics(pd.DataFrame())
def test_chromatic_beating(df, q):
df = _pd_to_tfs(df, q)
to = TwissOptics(df, quick_init=False)
with suppress_warnings(RuntimeWarning):
chrom_beat = to.get_chromatic_beating()
assert all([c in chrom_beat for c in ["S", "DBEATX", "DBEATY"]])
def test_linear_chromaticity(df, q):
df = _pd_to_tfs(df, q)
to = TwissOptics(df, quick_init=False)
with suppress_warnings(RuntimeWarning):
lchrom = to.get_linear_chromaticity()
assert len(lchrom) == 2
def test_linear_dispersion(df, q):
df = _pd_to_tfs(df, q)
to = TwissOptics(df, quick_init=False)
with suppress_warnings(RuntimeWarning):
disp = to.get_linear_dispersion()
assert all([c in disp for c in ["S", "DX", "DY"]])
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()
def test_import_from_file():
twiss_file = os.path.join(CURRENT_DIR, "..", "inputs", "models", "flat_beam1", "twiss.dat")
TwissOptics(twiss_file)
def test_no_index_given():
df = pd.DataFrame(np.zeros([3,3]),
columns=["S", "BETX", "BETY"])
with pytest.raises(IndexError):
TwissOptics(df)
