def set_yaxis_label(plot, axis, subnode):
labels_map = { # mapping for labels known by plot_style.py
'Phase_PhMdl': 'phase',
'Diff_Phase_PhMdl': 'phase',
'Beta_beat': 'betabeat',
'Beta_BMdl': 'beta',
'Disp_DMdl': 'dispersion',
'diff_Disp_DMdl': 'dispersion',
'NDisp_NDMdl': 'norm_dispersion',
'diff_NDisp_NDMdl': 'norm_dispersion',
'CO': 'co',
'ChromaticAmplitude': 'chromamp',
'ChromaticCouplingReal': 'real',
'ChromaticCouplingImaginary': 'imag',
'ChromaticCouplingAmp': 'absolute',
'amp': 'absolute',
'real': 'real',
'imaginary': 'imag',
}
coupling_map = {'x': r'f_{{1001}}', 'y': r'f_{{1010}}'}
delta = False
chromcoup = False
if subnode.lower().startswith('diff') or subnode.lower() == "co":
delta = True
if subnode.lower().startswith('chromaticcoupling'):
delta = True
chromcoup = True
axis = coupling_map[axis]
if subnode.lower() in ['amp', 'real', 'imaginary']:
axis = coupling_map[axis]
ps.set_yaxis_label(
param=labels_map[subnode],
plane=axis,
ax=plot,
delta=delta,
chromcoup=chromcoup,
)
def plot_linear_dispersion(self, combined=True):
""" Plot the Linear Dispersion.
Available after calc_linear_dispersion!
Args:
combined (bool): If 'True' plots x and y into the same axes.
"""
LOG.debug("Plotting Linear Dispersion")
lin_disp = self.get_linear_dispersion().dropna()
title = 'Linear Dispersion'
pstyle.set_style(self._plot_options.style, self._plot_options.manual)
if combined:
ax_dx = lin_disp.plot(x='S')
ax_dx.set_title(title)
pstyle.set_yaxis_label('dispersion', 'x,y', ax_dx)
ax_dy = ax_dx
else:
ax_dx = lin_disp.plot(x='S', y='DX')
ax_dx.set_title(title)
pstyle.set_yaxis_label('dispersion', 'x', ax_dx)
ax_dy = lin_disp.plot(x='S', y='DY')
ax_dy.set_title(title)
pstyle.set_yaxis_label('dispersion', 'y', ax_dy)
for ax in (ax_dx, ax_dy):
self._nice_axes(ax)
ax.legend()
def plot_rdts(self, rdt_names=None, apply_fun=np.abs, combined=True):
""" Plot Resonance Driving Terms """
LOG.debug("Plotting Resonance Driving Terms")
rdts = self.get_rdts(rdt_names)
is_s = regex_in(r'\AS$', rdts.columns)
rdts = rdts.dropna()
rdts.loc[:, ~is_s] = rdts.loc[:, ~is_s].applymap(apply_fun)
pstyle.set_style(self._plot_options.style, self._plot_options.manual)
if combined:
ax = rdts.plot(x='S')
ax.set_title('Resonance Driving Terms')
pstyle.small_title(ax)
pstyle.set_name('Resonance Driving Terms', ax)
pstyle.set_yaxis_label(apply_fun.__name__, 'F_{{jklm}}', ax)
self._nice_axes(ax)
else:
for rdt in rdts.loc[:, ~is_s]:
ax = rdts.plot(x='S', y=rdt)
ax.set_title('Resonance Driving Term ' + rdt)
pstyle.small_title(ax)
pstyle.set_name('Resonance Driving Term ' + rdt, ax)
pstyle.set_yaxis_label(apply_fun.__name__, rdt, ax)
self._nice_axes(ax)
def plot_phase_advance(self, combined=True):
""" Plots the phase advances between two consecutive elements
Args:
combined (bool): If 'True' plots x and y into the same axes.
"""
raise NotImplementedError(
'Plotting Phase Advance Shift is not Implemented yet.')
#TODO: reimplement the phase-advance shift calculations (if needed??)
LOG.debug("Plotting Phase Advance")
tw = self.mad_twiss
pa = self._phase_advance
dpa = self._dphase_advance
phase_advx = np.append(pa['X'].iloc[0, -1] + tw.Q1,
pa['X'].values.diagonal(offset=-1))
dphase_advx = np.append(dpa['X'].iloc[0, -1],
dpa['X'].values.diagonal(offset=-1))
phase_advy = np.append(pa['Y'].iloc[0, -1] + tw.Q2,
pa['Y'].values.diagonal(offset=-1))
dphase_advy = np.append(dpa['Y'].iloc[0, -1],
dpa['Y'].values.diagonal(offset=-1))
phase_adv = tw[["S"]].copy()
phase_adv['MUX'] = np.cumsum(phase_advx + dphase_advx) % 1 - .5
phase_adv['MUY'] = np.cumsum(phase_advy + dphase_advy) % 1 - .5
title = 'Phase'
pstyle.set_style(self._plot_options.style, self._plot_options.manual)
if combined:
ax_dx = phase_adv.plot(x='S')
ax_dx.set_title(title)
pstyle.small_title(ax_dx)
pstyle.set_name(title, ax_dx)
pstyle.set_yaxis_label('phase', 'x,y', ax_dx, delta=False)
ax_dy = ax_dx
else:
ax_dx = phase_adv.plot(x='S', y='MUX')
ax_dx.set_title(title)
pstyle.small_title(ax_dx)
pstyle.set_name(title, ax_dx)
pstyle.set_yaxis_label('phase', 'x', ax_dx, delta=False)
ax_dy = phase_adv.plot(x='S', y='MUY')
ax_dy.set_title(title)
pstyle.small_title(ax_dy)
pstyle.set_name(title, ax_dy)
pstyle.set_yaxis_label('phase', 'y', ax_dy, delta=False)
for ax in (ax_dx, ax_dy):
self._nice_axes(ax)
ax.legend()
def plot_chromatic_beating(self, combined=True):
""" Plot the Chromatic Beating
Available after calc_chromatic_beating
Args:
combined (bool): If 'True' plots x and y into the same axes.
"""
LOG.debug("Plotting Chromatic Beating")
chrom_beat = self.get_chromatic_beating().dropna()
title = 'Chromatic Beating'
pstyle.set_style(self._plot_options.style, self._plot_options.manual)
if combined:
ax_dx = chrom_beat.plot(x='S')
ax_dx.set_title(title)
pstyle.small_title(ax_dx)
pstyle.set_name(title, ax_dx)
pstyle.set_yaxis_label('dbetabeat', 'x,y', ax_dx)
ax_dy = ax_dx
else:
ax_dx = chrom_beat.plot(x='S', y='DBEATX')
ax_dx.set_title(title)
pstyle.small_title(ax_dx)
pstyle.set_name(title, ax_dx)
pstyle.set_yaxis_label('dbetabeat', 'x', ax_dx)
ax_dy = chrom_beat.plot(x='S', y='DBEATY')
ax_dy.set_title(title)
pstyle.small_title(ax_dy)
pstyle.set_name(title, ax_dy)
pstyle.set_yaxis_label('dbetabeat', 'y', ax_dy)
for ax in (ax_dx, ax_dy):
self._nice_axes(ax)
ax.legend()
def _create_single_plot(x_cols, y_cols, e_cols, twiss_data, legends, labels,
xy, change_marker, no_legend, auto_scale):
""" Create plots per parameter """
# create layout
if xy:
plane_map = {0: "X", 1: "Y"}
gs = gridspec.GridSpec(2, 1, height_ratios=[1, 1])
else:
plane_map = None
gs = gridspec.GridSpec(1, 1, height_ratios=[1])
ir_pos = None
x_is_length = all([xc == "S" for xc in x_cols])
if x_is_length:
ir_pos = _find_ir_pos(twiss_data)
# create figure
fig = plt.figure()
data = twiss_data[
0] # kept it to be more similar with other plotting function
p_title = labels[0]
fig.canvas.set_window_title("File '{:s}'".format(p_title))
for plt_idx in range(1 + xy):
ax = fig.add_subplot(gs[plt_idx])
for idx_col, (x_col, y_col, e_col,
legend) in enumerate(zip(x_cols, y_cols, e_cols,
legends)):
LOG.debug("Plotting parameter '{:s}'".format(y_col))
# plot data
if xy:
y_name = y_col
y_full = y_col + plane_map[plt_idx]
e_full = e_col + plane_map[plt_idx]
else:
y_name = y_col[:-1]
y_full = y_col
e_full = e_col
x_val = data[x_col]
y_val = data[y_full]
try:
e_val = data[e_full]
except KeyError:
e_val = None
_, caps, bars = ax.errorbar(x_val,
y_val,
yerr=e_val,
ls=rcParams[u"lines.linestyle"],
fmt=get_marker(idx_col, change_marker),
label=legend)
# loop through bars and caps and set the alpha value
[bar.set_alpha(ERROR_ALPHA) for bar in bars]
[cap.set_alpha(ERROR_ALPHA) for cap in caps]
if x_is_length and (idx_col + 1) == len(x_cols):
try:
ps.set_xLimits(data.SEQUENCE, ax)
except (AttributeError, ps.ArgumentError):
pass
if auto_scale:
lim = _get_auto_scale(y_val, auto_scale)
ax.set_ylim(*lim)
# manage layout
if len(legends) == 1:
if legend is None:
try:
# if it's a recognized column make nice label
ps.set_yaxis_label(_map_proper_name(y_name),
y_full[-1], ax)
except (KeyError, ps.ArgumentError):
ax.set_ylabel(y_full)
else:
# use given legend/label
if xy:
legend += plane_map[plt_idx]
ax.set_ylabel(legend)
if xy and plt_idx == 0:
ax.axes.get_xaxis().set_visible(False)
if x_is_length and ir_pos:
ps.show_ir(ir_pos, ax, mode='lines')
else:
if x_is_length:
ps.set_xaxis_label(ax)
if ir_pos:
ps.show_ir(ir_pos, ax, mode='outside')
if not no_legend and plt_idx == 0:
ax.legend(loc='upper center',
bbox_to_anchor=(0.5, 1.25),
fancybox=True,
shadow=True,
ncol=3)
return fig
def _create_plots(x_cols, y_cols, e_cols, twiss_data, legends, labels, xy,
change_marker, no_legend, auto_scale):
""" Create plots per parameter """
# create layout
if xy:
plane_map = {0: "X", 1: "Y"}
gs = gridspec.GridSpec(2, 1, height_ratios=[1, 1])
else:
plane_map = None
gs = gridspec.GridSpec(1, 1, height_ratios=[1])
ir_pos = _find_ir_pos(twiss_data)
# create individual figures
figs = {}
for idx_col, (x_col, y_col,
e_col) in enumerate(zip(x_cols, y_cols, e_cols)):
LOG.debug("Plotting parameter '{:s}'".format(y_col))
# create figure
fig = plt.figure()
p_title = y_col
if xy:
p_title += " X/Y"
fig.canvas.set_window_title("Parameter '{:s}'".format(p_title))
# plot data
for plt_idx in range(1 + xy):
if xy:
y_name = y_col
y_full = y_col + plane_map[plt_idx]
e_full = e_col + plane_map[plt_idx]
else:
y_name = y_col[:-1]
y_full = y_col
e_full = e_col
ax = fig.add_subplot(gs[plt_idx])
for idx, data in enumerate(twiss_data):
x_val = data[x_col]
y_val = data[y_full]
try:
e_val = data[e_full]
except KeyError:
e_val = None
_, caps, bars = ax.errorbar(x_val,
y_val,
yerr=e_val,
ls=rcParams[u"lines.linestyle"],
fmt=get_marker(idx, change_marker),
label=legends[idx])
# loop through bars and caps and set the alpha value
[bar.set_alpha(ERROR_ALPHA) for bar in bars]
[cap.set_alpha(ERROR_ALPHA) for cap in caps]
if x_col == "S" and (idx + 1) == len(twiss_data):
try:
ps.set_xLimits(data.SEQUENCE, ax)
except (AttributeError, ps.ArgumentError):
pass
if auto_scale:
current_y_lims = _get_auto_scale(y_val, auto_scale)
if idx == 0:
y_lims = current_y_lims
else:
y_lims = [
min(y_lims[0], current_y_lims[0]),
max(y_lims[1], current_y_lims[1])
]
# manage layout
if auto_scale:
ax.set_ylim(*y_lims)
if labels[idx_col] is None:
try:
# if it's a recognized column make nice label
ps.set_yaxis_label(_map_proper_name(y_name), y_full[-1],
ax)
except (KeyError, ps.ArgumentError):
ax.set_ylabel(y_full)
else:
try:
# if it's a recognized name make nice label
ps.set_yaxis_label(_map_proper_name(labels[idx_col][:-1]),
labels[idx_col][-1], ax)
except (KeyError, ps.ArgumentError):
# use given label
ax.set_ylabel(labels[idx_col])
if xy and plt_idx == 0:
ax.axes.get_xaxis().set_visible(False)
if x_col == "S" and ir_pos:
ps.show_ir(ir_pos, ax, mode='lines')
else:
if x_col == "S":
ps.set_xaxis_label(ax)
if ir_pos:
ps.show_ir(ir_pos, ax, mode='outside')
if not no_legend and plt_idx == 0:
ax.legend(loc='upper center',
bbox_to_anchor=(0.5, 1.25),
fancybox=True,
shadow=True,
ncol=3)
figs[y_col] = fig
return figs
def _create_plots(x_cols, y_cols, e_cols, twiss_data, legends, labels, xy,
change_marker, no_legend):
""" Create plots per parameter """
_param_map = {
"BET": "beta",
"BB": "betabeat",
"D": "dispersion",
"ND": "norm_dispersion",
"MU": "phase",
"X": "co",
"Y": "co",
}
# create layout
if xy:
plane_map = {0: "X", 1: "Y"}
gs = gridspec.GridSpec(2, 1, height_ratios=[1, 1])
else:
plane_map = None
gs = gridspec.GridSpec(1, 1, height_ratios=[1])
ir_pos = _find_ir_pos(twiss_data)
# create individual figures
figs = {}
for x_col, y_col, e_col in zip(x_cols, y_cols, e_cols):
LOG.debug("Plotting parameter '{:s}'".format(y_col))
# create figure
fig = plt.figure()
p_title = y_col
if xy:
p_title += " X/Y"
fig.canvas.set_window_title("Parameter '{:s}'".format(p_title))
# plot data
for plt_idx in range(1 + xy):
if xy:
y_name = y_col
y_full = y_col + plane_map[plt_idx]
e_full = e_col + plane_map[plt_idx]
else:
y_name = y_col[:-1]
y_full = y_col
e_full = e_col
ax = fig.add_subplot(gs[plt_idx])
for idx, data in enumerate(twiss_data):
x_val = data[x_col]
y_val = data[y_full]
try:
e_val = data[e_full]
except KeyError:
e_val = None
ax.errorbar(x_val,
y_val,
yerr=e_val,
fmt=get_marker(idx, change_marker),
label=legends[idx])
if (idx + 1) == len(twiss_data):
try:
ps.set_xLimits(data.SEQUENCE, ax)
except (AttributeError, ps.ArgumentError):
pass
# manage layout
if labels[plt_idx] is None:
try:
# if it's a recognized column make nice label
ps.set_yaxis_label(_param_map[y_name], y_full[-1], ax)
except (KeyError, ps.ArgumentError):
ax.set_ylabel(y_full)
else:
# use given label
ax.set_ylabel(labels[plt_idx])
if xy and plt_idx == 0:
ax.axes.get_xaxis().set_visible(False)
if ir_pos:
ps.show_ir(ir_pos, ax, mode='lines')
else:
ps.set_xaxis_label(ax)
if ir_pos:
ps.show_ir(ir_pos, ax, mode='outside')
if not no_legend and plt_idx == 0:
ax.legend(loc='upper center',
bbox_to_anchor=(0.5, 1.25),
fancybox=True,
shadow=True,
ncol=3)
figs[y_col] = fig
return figs