def _set_ylabel(ax, default, y_label, y_plane, chromatic):
""" Tries to set a mapped y label, otherwise the default """
try:
ps.set_yaxis_label(_map_proper_name(y_label),
y_plane, ax, chromcoup=chromatic)
except (KeyError, ps.ArgumentError):
ax.set_ylabel(default)
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 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 _finalize_plot(fig, ax, param, title, plane, opt):
ax.set_title(title)
if opt['plot']['legend']:
ymove = fig.get_figwidth() / 150 # empirical...
ax.legend(loc='upper right',
frameon=True,
bbox_to_anchor=[1 + ymove, 1])
ax.set_position(ax.get_position().translated(-ymove / 5, 0))
ax.set_xlabel(ax.get_xlabel(), labelpad=4)
ax.set_ylabel(ax.get_ylabel(), labelpad=5)
ax.ticklabel_format(axis='y', style='sci', scilimits=(-2, 3))
ps.set_xaxis_label(ax)
ps.set_yaxis_label(ax, plane, param)
ps.set_xLimits('LHCB1', ax)
ps.show_ir(_get_ip_positions(opt), ax)
_save_figure(
fig,
title.replace(" ", "_").replace("$", "").replace("{",
"").replace("}", ""),
opt)
plt.close(fig)
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 = rdts.columns.str.match(r'S$', case=False)
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()