def hack_plot_fieldmaps(astra_input, sections=['cavity', 'solenoid'], fieldmaps = {}, verbose=False):
"""
Plots Cavity and Solenoid fielmaps.
return: plot figure and axis objects
TODO: quadrupoles
"""
fmaps = load_fieldmaps(astra_input, sections=sections, verbose=verbose)
assert len(sections) == 2, 'TODO: more general'
# plt.figure(figsize = (12,8))
fig, ax0 = plt.subplots(figsize = (20,4))
# Make RHS axis for the solenoid field.
ax1 = ax0.twinx()
ax = [ax0, ax1]
ylabel = {'cavity': '$E_z$ (MV/m)', 'solenoid':'$B_z$ (T)'}
color = {'cavity': 'green', 'solenoid':'blue'}
for i, section in enumerate(sections):
a = ax[i]
ixlist = find_fieldmap_ixlist(astra_input, section)
for ix in ixlist:
dat = fieldmap_data(astra_input, section=section, index=ix, fieldmaps=fmaps, verbose=verbose)
label = f'{section}_{ix}'
c = color[section]
dat2plot = dat.T
# a.plot(*dat.T, label=label, color=c)
a.plot(dat2plot[0],abs(dat2plot[1]), label=label, color=c)
a.set_ylabel(ylabel[section])
ax0.set_xlabel('$z$ (m)')
# TODO: is this the correct thing to return?
return fig, ax
def old_plot_fieldmaps(astra_input, sections=['cavity', 'solenoid'], fieldmap_dict = {}, verbose=False):
"""
Plots Cavity and Solenoid fielmaps.
TODO: quadrupoles
"""
if fieldmap_dict:
fmaps = fieldmap_dict
else:
fmaps = load_fieldmaps(astra_input, sections=sections, verbose=verbose)
assert len(sections) == 2, 'TODO: more general'
fig, ax0 = plt.subplots()
# Make RHS axis for the solenoid field.
ax1 = ax0.twinx()
ax = [ax0, ax1]
ylabel = {'cavity': '$E_z$ (MV/m)', 'solenoid':'$B_z$ (T)'}
color = {'cavity': 'green', 'solenoid':'blue'}
for i, section in enumerate(sections):
a = ax[i]
ixlist = find_fieldmap_ixlist(astra_input, section)
for ix in ixlist:
dat = fieldmap_data(astra_input, section=section, index=ix, fieldmaps=fmaps, verbose=verbose)
label = f'{section}_{ix}'
c = color[section]
a.plot(*dat.T, label=label, color=c)
a.set_ylabel(ylabel[section])
ax0.set_xlabel('$z$ (m)')
def plot_stats(astra_object, keys=['norm_emit_x', 'sigma_z'], sections=['cavity', 'solenoid'], fieldmaps = {}, verbose=False):
"""
Plots stats, with fieldmaps plotted from seections.
TODO: quadrupoles
"""
astra_input = astra_object.input
fmaps = load_fieldmaps(astra_input, sections=sections, verbose=verbose)
assert len(sections) == 2, 'TODO: more general'
nplots = len(keys) + 1
fig, axs = plt.subplots(nplots)
# Make RHS axis for the solenoid field.
xdat = astra_object.stat('mean_z')
xmin = min(xdat)
xmax = max(xdat)
for i, key in enumerate(keys):
ax = axs[i]
unit = astra_object.units(key)
ydat = astra_object.stat(key)
ndat, factor, prefix = nice_array(ydat)
label = f'{key} ({prefix}{unit})'
ax.set_ylabel(label)
ax.set_xlim(xmin, xmax)
ax.plot(xdat, ndat)
ax1 = axs[-1]
ax1rhs = ax1.twinx()
ax = [ax1, ax1rhs]
ylabel = {'cavity': '$E_z$ (MV/m)', 'solenoid':'$B_z$ (T)'}
color = {'cavity': 'green', 'solenoid':'blue'}
for i, section in enumerate(sections):
a = ax[i]
ixlist = find_fieldmap_ixlist(astra_input, section)
for ix in ixlist:
dat = fieldmap_data(astra_input, section=section, index=ix, fieldmaps=fmaps, verbose=verbose)
label = f'{section}_{ix}'
c = color[section]
a.plot(*dat.T, label=label, color=c)
a.set_ylabel(ylabel[section])
ax1.set_xlabel('$z$ (m)')
ax1.set_xlim(xmin, xmax)
def plot_stats(astra_object, keys=['norm_emit_x', 'sigma_z'], sections=['cavity', 'solenoid'], fieldmaps = {}, verbose=False):
"""
Plots stats, with fieldmaps plotted from seections.
TODO: quadrupoles
"""
astra_input = astra_object.input
fmaps = load_fieldmaps(astra_input, sections=sections, verbose=verbose)
assert len(sections) == 2, 'TODO: more general'
nplots = len(keys) + 1
fig, axs = plt.subplots(nplots)
# Make RHS axis for the solenoid field.
xdat = astra_object.stat('mean_z')
xmin = min(xdat)
xmax = max(xdat)
for i, key in enumerate(keys):
ax = axs[i]
unit = astra_object.units(key)
ydat = astra_object.stat(key)
ndat, factor, prefix = nice_array(ydat)
label = f'{key} ({prefix}{unit})'
ax.set_ylabel(label)
ax.set_xlim(xmin, xmax)
ax.plot(xdat, ndat)
add_fieldmaps_to_axes(astra_object, axs[-1], bounds=(xmin, xmax),
sections=['cavity', 'solenoid'],
include_labels=True)
def add_fieldmaps_to_axes(astra_object, axes, bounds=None,
sections=['cavity', 'solenoid'],
include_labels=True):
"""
Adds fieldmaps to an axes.
"""
astra_input = astra_object.input
verbose=astra_object.verbose
if astra_object.fieldmap:
fmaps = astra_object.fieldmap
else:
fmaps = load_fieldmaps(astra_input, sections=sections, verbose=verbose)
ax1 = axes
ax1rhs = ax1.twinx()
ax = [ax1, ax1rhs]
ylabel = {'cavity': '$E_z$ (MV/m)', 'solenoid':'$B_z$ (T)'}
color = {'cavity': 'green', 'solenoid':'blue'}
for i, section in enumerate(sections):
a = ax[i]
ixlist = find_fieldmap_ixlist(astra_input, section)
for ix in ixlist:
dat = fieldmap_data(astra_input, section=section, index=ix, fieldmaps=fmaps, verbose=verbose)
label = f'{section}_{ix}'
c = color[section]
a.plot(*dat.T, label=label, color=c)
a.set_ylabel(ylabel[section])
ax1.set_xlabel('$z$ (m)')
if bounds:
ax1.set_xlim(bounds[0], bounds[1])