def get_plotters(real_data_file, imag_data_file):
plotter_r = twplot.plotter(real_data_file, buffered=False, units='mks')
try:
carrier = float(imag_data_file)
plotter_i = 0.0
except ValueError:
carrier = 0.0
plotter_i = twplot.plotter(imag_data_file, buffered=False, units='mks')
return plotter_r, plotter_i, carrier
def gen_viz(primitive_file,fig_dict):
'''Figure out whether a still image or movie was requested and call appropriate routine'''
file_to_plot = fig_dict['data']
slicing_spec = fig_dict['slicing_spec']
slices = fig_dict['slices']
dyn_range = fig_dict['dyn_range']
val_range = fig_dict['val_range']
my_color_map = fig_dict['color']
plotter = twplot.plotter(file_to_plot,units=fig_dict['units'])
c = slicing_spec.find('t')
if (c==2 or c==3) and slices[c-2]==':':
gen_movie(primitive_file+'.gif',plotter,slicing_spec,slices,dyn_range,val_range,my_color_map)
return primitive_file+'.gif'
else:
gen_plot(primitive_file+'.png',plotter,slicing_spec,tuple(map(int,slices)),dyn_range,val_range,my_color_map)
return primitive_file+'.png'
def gen_viz(primitive_file,
file_to_plot,
slicing_spec,
slices,
dyn_range=0.0,
my_color_map='viridis'):
'''Figure out whether a still image or movie was requested and call appropriate routine'''
plotter = twplot.plotter(file_to_plot)
c = slicing_spec.find('t')
if (c == 2 or c == 3) and slices[c - 2] == -1:
gen_movie(primitive_file + '.gif', plotter, slicing_spec, slices,
dyn_range, my_color_map)
return primitive_file + '.gif'
else:
gen_plot(primitive_file + '.png', plotter, slicing_spec, slices,
dyn_range, my_color_map)
return primitive_file + '.png'
aval = arg.split('=')[1]
if akey == 'type':
plot_type = int(aval)
if akey == 'dr':
dyn_range = np.double(aval)
if akey == 'ask':
ask = aval
if akey == 'units':
units = aval
if plot_type != 1 and plot_type != 2:
raise ValueError('Plot type must be 1 or 2')
if ask not in ['yes', 'no']:
raise ValueError('ask key must have value of yes or no')
plotter = twplot.plotter(file_to_plot, units=units, buffered=True)
plotter.display_info()
# Set up animation slices
slices = []
axes = twplot.get_axis_info(slicing_spec)
for i, s in enumerate(primitive_slices):
ax = axes[4 - num_slice_axes + i]
num = plotter.dims4()[ax]
if s == 'x':
slices.append(np.arange(0, num))
else:
slices.append(np.array([np.int(s)]))
# Check existing image files and clean
if akey not in ['type', 'dr', 'ask']:
raise KeyError('Invalide key ' + akey + ' in arguments')
aval = arg.split('=')[1]
if akey == 'type':
plot_type = int(aval)
if akey == 'dr':
dyn_range = np.double(aval)
if akey == 'ask':
ask = aval
if plot_type != 1 and plot_type != 2:
raise ValueError('Plot type must be 1 or 2')
if ask not in ['yes', 'no']:
raise ValueError('ask key must have value of yes or no')
plotter = twplot.plotter(file_to_plot, buffered=True)
plotter.display_info()
# Set up animation slices
slices = []
axes = twplot.get_axis_info(slicing_spec)
for i, s in enumerate(primitive_slices):
ax = axes[4 - num_slice_axes + i]
num = plotter.dims4()[ax]
if s == 'x':
slices.append(np.arange(0, num))
else:
slices.append(np.array([np.int(s)]))
# Check existing image files and clean
A,
bounds_error=False,
fill_value=0.0)
pts = x.reshape(-1, 2)
return func(pts).reshape(U)
print('Get boosted frame data...')
# Prescribe the unboosted nodes
t = np.linspace(0, 1400, Nt)
xi = [x for x in xslc]
yi = [y for y in yslc]
z = np.linspace(-20, 0, Nz)
unboosted_data = np.zeros((Nt, len(xi), len(yi), Nz))
p_ex = twplot.plotter(filename + '_' + tcomp + '.npy')
p_ex.display_info()
if zcomp != None:
p_by = twplot.plotter(filename + '_' + zcomp + '.npy')
for ix, x0 in enumerate(xi):
for iy, y0 in enumerate(yi):
print('Unboost slice', x0, y0, '...')
Ex, plot_dict = p_ex.falsecolor2d('ztxy', (x0, y0), dyn_range=0.0)
if zcomp != None:
By, plot_dict = p_by.falsecolor2d('ztxy', (x0, y0), dyn_range=0.0)
else:
By = None
# N.b. the plotter gives the data back with axes swapped for compatibility with pyplot.imshow routine
tp = plot_dict['ypts']
zp = plot_dict['xpts']
def main():
if len(sys.argv) < 2 or '--help' in sys.argv or '-h' in sys.argv:
print_usage()
exit(0)
if '--version' in sys.argv or '-v' in sys.argv:
print_version()
exit(0)
# Matplotlib setup
mpl.rcParams.update({'text.usetex': False, 'font.size': 10})
proportional = False
if proportional:
my_aspect = 'equal'
else:
my_aspect = 'auto'
# Process command line arguments and setup plotter object
slicing_spec = []
primitive_slices = []
for spec in sys.argv[1].split('/'):
slicing_spec.append(spec.split('=')[0])
primitive_slices.append(spec.split('=')[1].split(','))
file_to_plot = sys.argv[2].split(',')
N = len(file_to_plot)
panels = []
rows = 1
cols = N
dyn_range = []
color = []
roi = []
val_range = []
mult = []
units = []
labels = []
texlabels = []
ask = 'yes'
keylist = [
'panels', 'layout', 'dr', 'color', 'roi', 'range', 'mult', 'units',
'labels', 'texlabels'
]
for keyval in sys.argv[3:]:
key = keyval.split('=')[0]
arg = keyval.split('=')[1]
if key not in keylist:
raise ValueError('Invalid key in optional arguments.')
if key == 'panels':
panels = arg.split(',')
if key == 'layout':
layout = arg.split('x')
rows = int(layout[0])
cols = int(layout[1])
if key == 'dr':
dyn_range = arg.split(',')
if key == 'color':
color = arg.split(',')
if key == 'roi':
roi = arg.split('/')
if key == 'range':
val_range = arg.split('/')
if key == 'mult':
mult = arg.split('/')
if key == 'units':
units = arg.split(',')
if key == 'labels':
labels = arg.split('/')
if key == 'texlabels':
texlabels = arg.split('/')
for i in range(N - len(slicing_spec)):
slicing_spec.append(slicing_spec[-1])
primitive_slices.append(primitive_slices[-1])
for i in range(N - len(panels)):
panels.append('')
for i in range(N - len(dyn_range)):
dyn_range.append('0')
for i in range(N - len(val_range)):
val_range.append('')
for i in range(N - len(color)):
color.append('viridis')
for i in range(N - len(roi)):
roi.append('')
for i in range(N - len(mult)):
mult.append('1.0,1.0,1.0')
for i in range(N - len(labels)):
labels.append(',,')
for i in range(N - len(texlabels)):
texlabels.append(',,')
for i in range(N - len(units)):
units.append('plasma')
plotter = []
for i, f in enumerate(file_to_plot):
needs_buffer = slicing_spec[i][0] == 't' or (
slicing_spec[i][1] == 't' and len(primitive_slices[i]) == 2)
plotter.append(twplot.plotter(f, buffered=needs_buffer,
units=units[i]))
plotter[-1].display_info()
# Set up animation slices
slice_tuples = []
movie = []
for i in range(N):
axes = twplot.get_axis_info(slicing_spec[i])
dims = plotter[i].dims4()
st, m = ParseSlices(dims, axes, primitive_slices[i])
slice_tuples.append(st)
movie.append(m)
# Check existing files and clean if there is a movie
if movie[0]:
img_files = glob.glob('frame*.png')
if len(img_files) > 0 and ask == 'yes':
ans = ''
while ans != 'y' and ans != 'n':
ans = input('Found some frame*.png files, OK to clean (y/n) ?')
if ans == 'n':
print(
'STOPPED. Please run script in a directory where there are no important files of the form frame*.png.'
)
exit(1)
for img_file in img_files:
os.remove(img_file)
mov_files = glob.glob('mov.gif')
if len(mov_files) > 0 and ask == 'yes':
ans = ''
while ans != 'y' and ans != 'n':
ans = input('Found mov.gif, OK to overwrite (y/n) ?')
if ans == 'n':
print(
'STOPPED. Please do something with existing mov.gif and try again.'
)
exit(1)
# Determine the global color scale bounds
global_min = 1e50 * np.ones(N)
global_max = -1e50 * np.ones(N)
for i, p in enumerate(plotter):
if val_range[i] == '':
for slice_now in slice_tuples[i]:
dmult = float(mult[i].split(',')[2])
if len(primitive_slices[i]) == 2:
test_array, plot_dict = p.falsecolor2d(
slicing_spec[i], slice_now, float(dyn_range[i]))
local_min = plot_dict['vmin'] * dmult
local_max = plot_dict['vmax'] * dmult
if len(primitive_slices[i]) == 3:
abcissa, test_array, plot_dict = p.lineout(
slicing_spec[i], slice_now, float(dyn_range[i]))
local_min = np.min(test_array) * dmult
local_max = np.max(test_array) * dmult
if local_min > local_max:
local_min, local_max = local_max, local_min
if local_min < global_min[i]:
global_min[i] = local_min
if local_max > global_max[i]:
global_max[i] = local_max
else:
global_min[i] = float(val_range[i].split(',')[0])
global_max[i] = float(val_range[i].split(',')[1])
# Generate the movie or show a single frame.
# The number of frames is dictated by the first panel.
for file_idx in range(len(slice_tuples[0])):
sz = (cols * 5, rows * 4)
plt.figure(file_idx, figsize=sz, dpi=100)
for i, p in enumerate(plotter):
slice_now = slice_tuples[i][file_idx]
exm = np.zeros(4)
exm[:2] = float(mult[i].split(',')[0])
exm[2:] = float(mult[i].split(',')[1])
dmult = float(mult[i].split(',')[2])
plt.subplot(rows, cols, i + 1)
if len(primitive_slices[i]) == 2:
data_slice, plot_dict = p.falsecolor2d(slicing_spec[i],
slice_now,
float(dyn_range[i]))
plt.imshow(data_slice * dmult,
origin='lower',
aspect=my_aspect,
extent=np.array(plot_dict['extent']) * exm,
vmin=global_min[i],
vmax=global_max[i],
cmap=color[i])
b = plt.colorbar()
l = labels[i].split(',')
tl = texlabels[i].split(',')
if l[0] == '' and tl[0] == '':
plt.xlabel(plot_dict['xlabel'], size=12)
if not l[0] == '':
plt.xlabel(format_label(l[0]), size=12)
if not tl[0] == '':
plt.xlabel(format_label(r'$' + tl[0] + r'$'), size=12)
if l[1] == '' and tl[1] == '':
plt.ylabel(plot_dict['ylabel'], size=12)
if not l[1] == '':
plt.ylabel(format_label(l[1]), size=12)
if not tl[1] == '':
plt.ylabel(format_label(r'$' + tl[1] + r'$'), size=12)
if l[2] == '' and tl[2] == '':
b.set_label(plot_dict['blabel'], size=12)
if not l[2] == '':
b.set_label(format_label(l[2]), size=12)
if not tl[2] == '':
b.set_label(format_label(r'$' + tl[2] + r'$'), size=12)
if roi[i] == '':
roi_i = np.array(plot_dict['extent']) * exm
else:
roi_i = []
for s in roi[i].split(','):
roi_i.append(float(s))
plt.xlim(roi_i[0], roi_i[1])
plt.ylim(roi_i[2], roi_i[3])
if not panels[i] == '':
plt.text(roi_i[0], roi_i[3] + 0.03 * (roi_i[3] - roi_i[2]),
'(' + panels[i] + ')')
if len(primitive_slices[i]) == 3:
abcissa, ordinate, plot_dict = p.lineout(
slicing_spec[i], slice_now, float(dyn_range[i]))
plt.plot(abcissa * exm[0], ordinate * exm[2])
l = labels[i].split(',')
tl = texlabels[i].split(',')
if l[0] == '' and tl[0] == '':
plt.xlabel(plot_dict['xlabel'], size=12)
if not l[0] == '':
plt.xlabel(format_label(l[0]), size=12)
if not tl[0] == '':
plt.xlabel(format_label(r'$' + tl[0] + r'$'), size=12)
if l[1] == '' and tl[1] == '':
plt.ylabel(plot_dict['ylabel'], size=12)
if not l[1] == '':
plt.ylabel(format_label(l[1]), size=12)
if not tl[1] == '':
plt.ylabel(format_label(r'$' + tl[1] + r'$'), size=12)
if roi[i] == '':
roi_i = [
abcissa[0] * exm[0], abcissa[-1] * exm[0],
global_min[i] * exm[2], global_max[i] * exm[2]
]
else:
s = roi[i].split(',')
roi_i = [
float(s[0]),
float(s[1]),
float(s[2]),
float(s[3])
]
plt.xlim(roi_i[0], roi_i[1])
plt.ylim(roi_i[2], roi_i[3])
if not panels[i] == '':
plt.text(roi_i[0], roi_i[3] + 0.03 * (roi_i[3] - roi_i[2]),
'(' + panels[i] + ')')
plt.tight_layout()
if movie[0]:
img_file = 'frame{:03d}.png'.format(file_idx)
print('saving', img_file, '...')
plt.savefig(img_file)
plt.close()
if movie[0]:
print('Consolidating into movie file...')
images = []
frameRateHz = 5
img_files = sorted(glob.glob('frame*.png'))
for f in img_files:
images.append(PIL.Image.open(f))
images[0].save('mov.gif',
save_all=True,
append_images=images[1:],
duration=int(1000 / frameRateHz),
loop=0)
cleanup('frame*.png')
print('Done.')
else:
plt.show()
import numpy as np
import matplotlib.pyplot as plt
import twutils.plot as twplot
# Simple example of how to use the plotter
file_to_plot = '/Users/gordon/Run/air.dvdat'
# Make it buffered if you want to have time available as a plotting axis
# For huge files make it unbuffered
plotter = twplot.plotter(file_to_plot, buffered=False)
plotter.display_info()
# 2D Figure Example
plt.figure(1, figsize=(7, 5), dpi=75)
# slicing_spec : first two characters are axes to plot, second two are the slice axes
# slice_to_plot : select the slices to plot (order is whatever is in slicing_spec)
data_slice, plot_dict = plotter.falsecolor2d(slicing_spec='xyzt',
slice_to_plot=(64, -1),
dyn_range=2.0)
plt.imshow(data_slice,
origin='lower',
extent=plot_dict['extent'],
vmin=plot_dict['vmin'],
vmax=plot_dict['vmax'],
cmap='jet')
plt.colorbar(label=file_to_plot.split('/')[-1])
plt.xlabel(plot_dict['xlabel'], fontsize=18)
plt.ylabel(plot_dict['ylabel'], fontsize=18)
# Lineout Example
for i in range(N - len(color)):
color.append('viridis')
for i in range(N - len(roi)):
roi.append('')
for i in range(N - len(mult)):
mult.append('1.0,1.0,1.0')
for i in range(N - len(labels)):
labels.append('')
for i in range(N - len(texlabels)):
texlabels.append('')
plotter = []
for i, f in enumerate(file_to_plot):
needs_buffer = slicing_spec[i][0] == 't' or (slicing_spec[i][1] == 't' and
len(primitive_slices[i]) == 2)
plotter.append(twplot.plotter(f, buffered=needs_buffer))
plotter[0].display_info()
# Set up animation slices
slice_tuples = []
movie = []
for i in range(N):
axes = twplot.get_axis_info(slicing_spec[i])
dims = plotter[i].dims4()
st, m = ParseSlices(dims, axes, primitive_slices[i])
slice_tuples.append(st)
movie.append(m)
# Check existing files and clean if there is a movie
panels = val.split(',')
if key == 'layout':
layout = val
if key == 'dr':
dyn_range = []
dyn_range.append(float(val.split(',')[0]))
dyn_range.append(float(val.split(',')[1]))
if key == 'color':
color = val.split(',')
if key == 'roi':
for s in val.split('/')[0].split(','):
roi[0].append(int(s))
for s in val.split('/')[1].split(','):
roi[1].append(int(s))
plotter_r = twplot.plotter(real_data_file, buffered=False)
try:
carrier = float(imag_data_file)
plotter_i = 0.0
except ValueError:
carrier = 0.0
plotter_i = twplot.plotter(imag_data_file, buffered=False)
plotter_r.display_info()
# Set up animation slices
axes = twplot.get_axis_info(slicing_spec)
dims = plotter_r.dims4()
slice_tuples, movie = ParseSlices(dims, axes, primitive_slices)
# Check existing image files and clean
