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()
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
img_files = glob.glob('tempfile*.png')
if len(img_files) > 0 and ask == 'yes':
ans = ''
while ans != 'y' and ans != 'n':
ans = input('Found some tempfile*.png files, OK to clean (y/n) ?')
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 and default args
mpl.rcParams.update({'text.usetex': False, 'font.size': 10})
proportional = False
if proportional:
my_aspect = 'equal'
else:
my_aspect = 'auto'
# Required arguments
slicing_spec = []
primitive_slices = []
real_data_file = []
imag_data_file = []
for s in sys.argv[1].split('/'):
slicing_spec += [s.split('=')[0]]
primitive_slices += [(s.split('=')[1]).split(',')]
for f in sys.argv[2].split(','):
real_data_file += [f]
for f in sys.argv[3].split(','):
imag_data_file += [f]
N = len(real_data_file)
if len(imag_data_file) != N:
raise ValueError('Number of files does not match.')
# Optional arguments
color = []
dyn_range = []
roi = []
val_range = []
ask = 'yes'
panels = []
type = []
cols = 1
rows = N
for keyval in sys.argv[4:]:
key = keyval.split('=')[0]
val = keyval.split('=')[1]
if key == 'panels':
panels = val.split(',')
if key == 'layout':
layout = val.split('x')
rows = int(layout[0])
cols = int(layout[1])
if key == 'dr':
for dr in val.split(','):
dyn_range += [float(dr)]
if key == 'color':
color = val.split(',')
if key == 'roi':
roi = val.split('/')
if key == 'range':
val_range = val.split('/')
if key == 'type':
type = val.split(',')
# Fill in defaults
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.0)
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(val_range)):
val_range.append('')
for i in range(N - len(type)):
type.append('wigner')
# Check existing image files and clean
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)
# Determine the global color scale bounds for all plots
# If a movie we have to do all the transforms first
# We don't save the data, just do redundant transforms later
num_frames = []
global_min = [1e50] * N
global_max = [-1e50] * N
print('First pass...')
print('-------------')
for i in range(N):
plotter_r, plotter_i, carrier = get_plotters(real_data_file[i],
imag_data_file[i])
axes = twplot.get_axis_info(slicing_spec[i])
dims = plotter_r.dims4()
slice_tuples, movie = ParseSlices(dims, axes, primitive_slices[i])
num_frames += [len(slice_tuples)]
if val_range[i] == '':
for slice_now in slice_tuples:
dat, ext = get_plot_data(plotter_r, plotter_i, carrier,
type[i], slicing_spec[i], slice_now,
dyn_range[i])
local_min = np.min(dat)
local_max = np.max(dat)
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])
# Make a movie or display single frame
print('Generate images...')
print('------------------')
for file_idx in range(num_frames[0]):
sz = (cols * 5, rows * 4)
plt.figure(file_idx, figsize=sz, dpi=100)
for i in range(N):
plt.subplot(rows, cols, i + 1)
plotter_r, plotter_i, carrier = get_plotters(
real_data_file[i], imag_data_file[i])
axes = twplot.get_axis_info(slicing_spec[i])
dims = plotter_r.dims4()
slice_tuples, movie = ParseSlices(dims, axes, primitive_slices[i])
slice_now = slice_tuples[file_idx]
dat, ext = get_plot_data(plotter_r, plotter_i, carrier, type[i],
slicing_spec[i], slice_now, dyn_range[i])
if type[i] == 'pulse':
ext_scaled = list(1e15 * np.array(ext[:2]) / C.c) + list(
1e6 * np.array(ext[2:]))
plt.imshow(dat,
origin='lower',
aspect=my_aspect,
extent=ext_scaled,
vmin=global_min[i],
vmax=global_max[i],
cmap=color[i])
b = plt.colorbar()
b.set_label(r'${\cal E}$ (TV/cm)', size=12)
plt.xlabel(r'$z/c - t$ (fs)', size=12)
plt.ylabel(r'$\rho$ ($\mu$m)', size=12)
if type[i] == 'wigner':
ext_scaled = list(1e15 * np.array(ext[:2]) / C.c) + list(
1e-15 * np.array(ext[2:]))
plt.imshow(dat.swapaxes(0, 1),
origin='lower',
aspect=my_aspect,
extent=ext_scaled,
vmin=global_min[i],
vmax=global_max[i],
cmap=color[i])
b = plt.colorbar()
b.set_label(r'${\cal N}$ (MJ/cm$^2$)', size=12)
plt.xlabel(r'$z/c - t$ (fs)', size=12)
plt.ylabel(r'$\delta\omega$ (fs$^{-1}$)', size=12)
if type[i] == 'spectrum':
ext *= 1e-15
ext_scaled = [ext[0], ext[-1], global_min[i], global_max[i]]
plt.plot(ext, dat)
plt.xlabel(r'$\delta\omega$ (fs$^{-1}$)', size=12)
plt.ylabel(r'$|E(\omega)|^2$', size=12)
if roi[i] == '':
roi_i = ext_scaled
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] + ')')
plt.tight_layout()
if movie:
img_file = 'frame{:03d}.png'.format(file_idx)
print('saving', img_file, '...')
plt.savefig(img_file)
plt.close()
if movie:
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()