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 _create_fullresponse_from_dict(var_to_twiss):
""" Convert var-tfs dictionary to fullresponse dictionary """
var_to_twiss = _add_coupling(var_to_twiss)
resp = pandas.Panel.from_dict(var_to_twiss)
resp = resp.transpose(2, 1, 0)
# After transpose e.g: resp[NDX, bpm12l1.b1, kqt3]
# The magnet called "0" is no change (nominal model)
resp['NDX'] = resp.xs('DX', axis=0).div(np.sqrt(resp.xs('BETX', axis=0)), axis="index")
resp['NDY'] = resp.xs('DY', axis=0).div(np.sqrt(resp.xs('BETY', axis=0)), axis="index")
resp['BBX'] = resp.xs('BETX', axis=0).div(resp.loc['BETX', :, '0'], axis="index")
resp['BBY'] = resp.xs('BETY', axis=0).div(resp.loc['BETY', :, '0'], axis="index")
resp = resp.subtract(resp.xs('0', axis=2), axis=2)
# Remove difference of nominal model with itself (bunch of zeros)
resp.drop('0', axis=2, inplace=True)
resp = resp.div(resp.loc['incr', :, :])
with suppress_warnings(np.ComplexWarning): # raised as everything is complex-type now
df = {'MUX': resp.xs('MUX', axis=0).astype(np.float64),
'MUY': resp.xs('MUY', axis=0).astype(np.float64),
'BETX': resp.xs('BETX', axis=0).astype(np.float64),
'BETY': resp.xs('BETY', axis=0).astype(np.float64),
'BBX': resp.xs('BBX', axis=0).astype(np.float64),
'BBY': resp.xs('BBY', axis=0).astype(np.float64),
'DX': resp.xs('DX', axis=0).astype(np.float64),
'DY': resp.xs('DY', axis=0).astype(np.float64),
'NDX': resp.xs('NDX', axis=0).astype(np.float64),
'NDY': resp.xs('NDY', axis=0).astype(np.float64),
"F1001R": tfs.TfsDataFrame(resp.xs('1001', axis=0).apply(np.real).astype(np.float64)),
"F1001I": tfs.TfsDataFrame(resp.xs('1001', axis=0).apply(np.imag).astype(np.float64)),
"F1010R": tfs.TfsDataFrame(resp.xs('1010', axis=0).apply(np.real).astype(np.float64)),
"F1010I": tfs.TfsDataFrame(resp.xs('1010', axis=0).apply(np.imag).astype(np.float64)),
'Q': resp.loc[['Q1', 'Q2'], resp.major_axis[0], :].transpose().astype(np.float64),
}
return df
def _create_fullresponse_from_dict(var_to_twiss):
""" Convert var-tfs dictionary to fullresponse dictionary """
var_to_twiss = _add_coupling(var_to_twiss)
resp = pandas.Panel.from_dict(var_to_twiss)
resp = resp.transpose(2, 1, 0)
# After transpose e.g: resp[NDX, bpm12l1.b1, kqt3]
# The magnet called "0" is no change (nominal model)
resp['NDX'] = resp.xs('DX', axis=0).div(np.sqrt(resp.xs('BETX', axis=0)),
axis="index")
resp['NDY'] = resp.xs('DY', axis=0).div(np.sqrt(resp.xs('BETY', axis=0)),
axis="index")
resp['BBX'] = resp.xs('BETX', axis=0).div(resp.loc['BETX', :, '0'],
axis="index")
resp['BBY'] = resp.xs('BETY', axis=0).div(resp.loc['BETY', :, '0'],
axis="index")
resp = resp.subtract(resp.xs('0', axis=2), axis=2)
# Remove difference of nominal model with itself (bunch of zeros)
resp.drop('0', axis=2, inplace=True)
resp = resp.div(resp.loc['incr', :, :])
with suppress_warnings(
np.ComplexWarning): # raised as everything is complex-type now
df = {
'MUX':
resp.xs('MUX', axis=0).astype(np.float64),
'MUY':
resp.xs('MUY', axis=0).astype(np.float64),
'BETX':
resp.xs('BETX', axis=0).astype(np.float64),
'BETY':
resp.xs('BETY', axis=0).astype(np.float64),
'BBX':
resp.xs('BBX', axis=0).astype(np.float64),
'BBY':
resp.xs('BBY', axis=0).astype(np.float64),
'DX':
resp.xs('DX', axis=0).astype(np.float64),
'DY':
resp.xs('DY', axis=0).astype(np.float64),
'NDX':
resp.xs('NDX', axis=0).astype(np.float64),
'NDY':
resp.xs('NDY', axis=0).astype(np.float64),
"F1001R":
tfs.TfsDataFrame(
resp.xs('1001', axis=0).apply(np.real).astype(np.float64)),
"F1001I":
tfs.TfsDataFrame(
resp.xs('1001', axis=0).apply(np.imag).astype(np.float64)),
"F1010R":
tfs.TfsDataFrame(
resp.xs('1010', axis=0).apply(np.real).astype(np.float64)),
"F1010I":
tfs.TfsDataFrame(
resp.xs('1010', axis=0).apply(np.imag).astype(np.float64)),
'Q':
resp.loc[['Q1', 'Q2'],
resp.major_axis[0], :].transpose().astype(np.float64),
}
return df
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_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 plot_bbq_data(bbq_df,
interval=None,
xmin=None,
xmax=None,
ymin=None,
ymax=None,
output=None,
show=True,
two_plots=False):
""" Plot BBQ data.
Args:
bbq_df: BBQ Dataframe with moving average columns
interval: start and end time of used interval, will be marked with red bars
xmin: Lower x limit (time)
xmax: Upper x limit (time)
ymin: Lower y limit (tune)
ymax: Upper y limit (tune)
output: Path to the output file
show: Shows plot if `True`
two_plots: Plots each tune in it's own axes if `True`
Returns:
Plotted figure
"""
LOG.debug("Plotting BBQ data.")
ps.set_style("standard", {u"lines.marker": u""})
fig = plt.figure()
if two_plots:
gs = gridspec.GridSpec(2, 1, height_ratios=[1, 1])
ax = [fig.add_subplot(gs[1]), fig.add_subplot(gs[0])]
else:
gs = gridspec.GridSpec(1, 1, height_ratios=[1])
ax = fig.add_subplot(gs[0])
ax = [ax, ax]
bbq_df.index = [
datetime.datetime.fromtimestamp(time, tz=TIMEZONE)
for time in bbq_df.index
]
for idx, plane in enumerate(PLANES):
color = ps.get_mpl_color(idx)
mask = bbq_df[COL_IN_MAV(plane)]
with suppress_warnings(UserWarning): # caused by _nolegend_
bbq_df.plot(y=COL_BBQ(plane),
ax=ax[idx],
color=color,
alpha=.2,
label="_nolegend_")
bbq_df.loc[mask, :].plot(y=COL_BBQ(plane),
ax=ax[idx],
color=color,
alpha=.4,
label="$Q_{:s}$ filtered".format(
plane.lower()))
bbq_df.plot(y=COL_MAV(plane),
ax=ax[idx],
color=color,
label="$Q_{:s}$ moving av.".format(plane.lower()))
if ymin is None and two_plots:
ax[idx].set_ylim(bottom=min(bbq_df.loc[mask, COL_BBQ(plane)]))
if ymax is None and two_plots:
ax[idx].set_ylim(top=max(bbq_df.loc[mask, COL_BBQ(plane)]))
# things to add/do only once if there is only one plot
for idx in range(1 + two_plots):
if interval:
ax[idx].axvline(x=interval[0], color="red")
ax[idx].axvline(x=interval[1], color="red")
ax[idx].set_ylabel('Tune')
ax[idx].set_ylim(bottom=ymin, top=ymax)
ax[idx].yaxis.set_major_formatter(FormatStrFormatter('%.5f'))
ax[idx].set_xlim(left=xmin, right=xmax)
ax[idx].set_xlabel('Time')
ax[idx].xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))
# don't show labels on upper plot (if two plots)
if idx:
# use the visibility to allow cursor x-position to be shown
ax[idx].tick_params(labelbottom=False)
ax[idx].xaxis.get_label().set_visible(False)
plt.tight_layout()
if output:
fig.savefig(output)
ps.set_name(os.path.basename(output))
if show:
plt.draw()
return fig