def _bout_type_from_array(self, data):
"""Get the bout_type from the array 'data'
If 'data' is a BoutArray, it knows its bout_type, otherwise we
have to guess.
Parameters
----------
data : :py:obj:`~boututils.boutarray.BoutArray` or ndarray
An array with between 0 and 4 dimensions
Returns
-------
str
Either the actual bout_type or our best guess
See Also
--------
- `DataFile.bout_type`
TODO
----
- Make standalone function
"""
try:
# If data is a BoutArray, it should have a type attribute that we can use
return data.attributes["bout_type"]
except AttributeError:
# Otherwise data is a numpy.ndarray and we have to guess the bout_type
pass
try:
ndim = len(data.shape)
except AttributeError:
ndim = 0
if ndim == 4:
return "Field3D_t"
elif ndim == 3:
# not ideal, 3d field might be time-evolving 2d field,
# 'Field2D_t', but can't think of a good way to distinguish
alwayswarn(
"Warning: assuming bout_type of 3d array is Field3D. If it "
"should be a time-evolving Field2D, this may cause errors in "
"dimension sizes.")
return "Field3D"
elif ndim == 2:
return "Field2D"
elif ndim == 1:
return "scalar_t"
elif ndim == 0:
return "scalar"
else:
raise ValueError("Unrecognized variable bout_type, ndims=" +
str(ndim))
def _bout_type_from_array(self, data):
"""Get the bout_type from the array 'data'
If 'data' is a BoutArray, it knows its bout_type, otherwise we
have to guess.
Parameters
----------
data : :py:obj:`~boututils.boutarray.BoutArray` or ndarray
An array with between 0 and 4 dimensions
Returns
-------
str
Either the actual bout_type or our best guess
See Also
--------
- `DataFile.bout_type`
TODO
----
- Make standalone function
"""
try:
# If data is a BoutArray, it should have a type attribute that we can use
return data.attributes["bout_type"]
except AttributeError:
# Otherwise data is a numpy.ndarray and we have to guess the bout_type
pass
try:
ndim = len(data.shape)
except AttributeError:
ndim = 0
if ndim == 4:
return 'Field3D_t'
elif ndim == 3:
# not ideal, 3d field might be time-evolving 2d field,
# 'Field2D_t', but can't think of a good way to distinguish
alwayswarn("Warning: assuming bout_type of 3d array is Field3D. If it "
"should be a time-evolving Field2D, this may cause errors in "
"dimension sizes.")
return 'Field3D'
elif ndim == 2:
return 'Field2D'
elif ndim == 1:
return 'scalar_t'
elif ndim == 0:
return 'scalar'
else:
raise ValueError("Unrecognized variable bout_type, ndims=" + str(ndim))
def showdata(vars,
titles=[],
legendlabels=[],
surf=[],
polar=[],
tslice=0,
t_array=None,
movie=0,
fps=28,
dpi=200,
intv=1,
Ncolors=25,
x=[],
y=[],
global_colors=False,
symmetric_colors=False,
hold_aspect=False,
cmap=None,
clear_between_frames=None,
return_animation=False,
window_title=""):
"""A Function to animate time dependent data from BOUT++
To animate multiple variables on different axes:
>>> showdata([var1, var2, var3])
To animate more than one line on a single axes:
>>> showdata([[var1, var2, var3]])
The default graph types are:
2D (time + 1 spatial dimension) arrays = animated line plot
3D (time + 2 spatial dimensions) arrays = animated contour plot.
To use surface or polar plots:
>>> showdata(var, surf=1)
>>> showdata(var, polar=1)
Can plot different graph types on different axes. Default graph types will
be used depending on the dimensions of the input arrays. To specify
polar/surface plots on different axes:
>>> showdata([var1, var2], surf=[1, 0], polar=[0, 1])
Movies require FFmpeg (for .mp4) and/or ImageMagick (for .gif) to be
installed. The 'movie' option can be set to 1 (which will produce an mp4
called 'animation.mp4'), to a name with no extension (which will produce an
mp4 called '<name>.mp4')
The `tslice` variable is used to control the time value that is printed on
each frame of the animation. If the input data matches the time values
found within BOUT++'s dmp data files, then these time values will be used.
Otherwise, an integer counter is used.
The `cmap` variable (if specified) will set the colormap used in the plot
cmap must be a matplotlib colormap instance, or the name of a registered
matplotlib colormap
During animation click once to stop in the current frame. Click again to
continue.
Parameters
----------
vars : array_like or list of array_like
Variable or list of variables to plot
titles : str or list of str, optional
Title or list of titles for each axis
legendlabels : str or list of str, optional
Legend or list of legends for each variable
surf : list of int
Which axes to plot as a surface plot
polar : list of int
Which axes to plot as a polar plot
tslice : list of int
Use these time values from a dump file (see above)
t_array : array
Pass in t_array using this argument to use the simulation time in plot
titles. Otherwise, just use the t-index.
movie : int
If 1, save the animation to file
fps : int
Number of frames per second to use when saving animation
dpi : int
Dots per inch to use when saving animation
intv : int
???
Ncolors : int
Number of levels in contour plots
x, y : array_like, list of array_like
X, Y coordinates
global_colors : bool
If "vars" is a list the colorlevels are determined from the
maximum of the maxima and and the minimum of the minima in all
fields in vars
symmetric_colors : bool
Colour levels are symmetric
hold_aspect : bool
Use equal aspect ratio in plots
cmap : colormap
A matplotlib colormap instance to use
clear_between_frames : bool, optional
- Default (None) - all plots except line plots will clear between frames
- True - all plots will clear between frames
- False - no plots will clear between frames
return_animation : bool
Return the matplotlib animation instance
window_title : str
Give a title for the animation window
TODO
----
- Replace empty lists in signature with None
- Use bools in sensible places
- Put massive list of arguments in kwargs
- Speed up animations ????
- Look at theta in polar plots - periodic?!?
- Log axes, colorbars
- Figureplot
"""
plt.ioff()
# Check to see whether vars is a list or not.
if isinstance(vars, list):
Nvar = len(vars)
else:
vars = [vars]
Nvar = len(vars)
if Nvar < 1:
raise ValueError("No data supplied")
# Check to see whether each variable is a list - used for line plots only
Nlines = []
for i in range(0, Nvar):
if isinstance(vars[i], list):
Nlines.append(len(vars[i]))
else:
Nlines.append(1)
vars[i] = [vars[i]]
# Sort out titles
if len(titles) == 0:
for i in range(0, Nvar):
titles.append(('Var' + str(i + 1)))
elif len(titles) != Nvar:
raise ValueError(
'The length of the titles input list must match the length of the vars list.'
)
# Sort out legend labels
if len(legendlabels) == 0:
for i in range(0, Nvar):
legendlabels.append([])
for j in range(0, Nlines[i]):
legendlabels[i].append(chr(97 + j))
elif (isinstance(legendlabels[0], list) != 1):
if Nvar != 1:
check = 0
for i in range(0, Nvar):
if len(legendlabels) != Nlines[i]:
check = check + 1
if check == 0:
alwayswarn(
"The legendlabels list does not contain a sublist for each variable, but its length matches the number of lines on each plot. Will apply labels to each plot"
)
legendlabelsdummy = []
for i in range(0, Nvar):
legendlabelsdummy.append([])
for j in range(0, Nlines[i]):
legendlabelsdummy[i].append(legendlabels[j])
legendlabels = legendlabelsdummy
else:
alwayswarn(
"The legendlabels list does not contain a sublist for each variable, and it's length does not match the number of lines on each plot. Will default apply labels to each plot"
)
legendlabels = []
for i in range(0, Nvar):
legendlabels.append([])
for j in range(0, Nlines[i]):
legendlabels[i].append(chr(97 + j))
else:
if (Nlines[0] == len(legendlabels)):
legendlabels = [legendlabels]
elif len(legendlabels) != Nvar:
alwayswarn(
"The length of the legendlabels list does not match the length of the vars list, will continue with default values"
)
legendlabels = []
for i in range(0, Nvar):
legendlabels.append([])
for j in range(0, Nlines[i]):
legendlabels[i].append(chr(97 + j))
else:
for i in range(0, Nvar):
if isinstance(legendlabels[i], list):
if len(legendlabels[i]) != Nlines[i]:
alwayswarn(
'The length of the legendlabel (sub)list for each plot does not match the number of datasets for each plot. Will continue with default values'
)
legendlabels[i] = []
for j in range(0, Nlines[i]):
legendlabels[i].append(chr(97 + j))
else:
legendlabels[i] = [legendlabels[i]]
if len(legendlabels[i]) != Nlines[i]:
alwayswarn(
'The length of the legendlabel (sub)list for each plot does not match the number of datasets for each plot. Will continue with default values'
)
legendlabels[i] = []
for j in range(0, Nlines[i]):
legendlabels[i].append(chr(97 + j))
# Sort out surf list
if isinstance(surf, list):
if (len(surf) == Nvar):
for i in range(0, Nvar):
if surf[i] >= 1:
surf[i] = 1
else:
surf[i] = 0
elif (len(surf) == 1):
if surf[0] >= 1:
surf[0] = 1
else:
surf[0] = 0
if (Nvar > 1):
for i in range(1, Nvar):
surf.append(surf[0])
elif (len(surf) == 0):
for i in range(0, Nvar):
surf.append(0)
else:
alwayswarn(
'Length of surf list does not match number of variables. Will default to no polar plots'
)
for i in range(0, Nvar):
surf.append(0)
else:
surf = [surf]
if surf[0] >= 1:
surf[0] = 1
else:
surf[0] = 0
if (Nvar > 1):
for i in range(1, Nvar):
surf.append(surf[0])
# Sort out polar list
if isinstance(polar, list):
if (len(polar) == Nvar):
for i in range(0, Nvar):
if polar[i] >= 1:
polar[i] = 1
else:
polar[i] = 0
elif (len(polar) == 1):
if polar[0] >= 1:
polar[0] = 1
else:
polar[0] = 0
if (Nvar > 1):
for i in range(1, Nvar):
polar.append(polar[0])
elif (len(polar) == 0):
for i in range(0, Nvar):
polar.append(0)
else:
alwayswarn(
'Length of polar list does not match number of variables. Will default to no polar plots'
)
for i in range(0, Nvar):
polar.append(0)
else:
polar = [polar]
if polar[0] >= 1:
polar[0] = 1
else:
polar[0] = 0
if (Nvar > 1):
for i in range(1, Nvar):
polar.append(polar[0])
# Determine shapes of arrays
dims = []
Ndims = []
lineplot = []
contour = []
for i in range(0, Nvar):
dims.append([])
Ndims.append([])
for j in range(0, Nlines[i]):
dims[i].append(array((vars[i][j].shape)))
Ndims[i].append(dims[i][j].shape[0])
# Perform check to make sure that data is either 2D or 3D
if (Ndims[i][j] < 2):
raise ValueError(
'data must be either 2 or 3 dimensional. Exiting')
if (Ndims[i][j] > 3):
raise ValueError(
'data must be either 2 or 3 dimensional. Exiting')
if ((Ndims[i][j] == 2) & (polar[i] != 0)):
alwayswarn(
'Data must be 3 dimensional (time, r, theta) for polar plots. Will plot lineplot instead'
)
if ((Ndims[i][j] == 2) & (surf[i] != 0)):
alwayswarn(
'Data must be 3 dimensional (time, x, y) for surface plots. Will plot lineplot instead'
)
if ((Ndims[i][j] == 3) & (Nlines[i] != 1)):
raise ValueError(
'cannot have multiple sets of 3D (time + 2 spatial dimensions) on each subplot'
)
if ((Ndims[i][j] != Ndims[i][0])):
raise ValueError(
'Error, Number of dimensions must be the same for all variables on each plot.'
)
if (Ndims[i][0] == 2): # Set polar and surf list entries to 0
polar[i] = 0
surf[i] = 0
lineplot.append(1)
contour.append(0)
else:
if ((polar[i] == 1) & (surf[i] == 1)):
alwayswarn(
'Cannot do polar and surface plots at the same time. Default to contour plot'
)
contour.append(1)
lineplot.append(0)
polar[i] = 0
surf[i] = 0
elif (polar[i] == 1) | (surf[i] == 1):
contour.append(0)
lineplot.append(0)
else:
contour.append(1)
lineplot.append(0)
# Obtain size of data arrays
Nt = []
Nx = []
Ny = []
for i in range(0, Nvar):
Nt.append([])
Nx.append([])
Ny.append([])
for j in range(0, Nlines[i]):
Nt[i].append(vars[i][j].shape[0])
Nx[i].append(vars[i][j].shape[1])
if (Nt[i][j] != Nt[0][0]):
raise ValueError(
'time dimensions must be the same for all variables.')
#if (Nx[i][j] != Nx[i][0]):
# raise ValueError('Dimensions must be the same for all variables on each plot.')
if (Ndims[i][j] == 3):
Ny[i].append(vars[i][j].shape[2])
#if (Ny[i][j] != Ny[i][0]):
# raise ValueError('Dimensions must be the same for all variables.')
# Obtain number of frames
Nframes = int(Nt[0][0] / intv)
# Generate grids for plotting
# Try to use provided grids where possible
# If x and/or y are not lists, apply to all variables
if not isinstance(x, (list, tuple)):
x = [x] * Nvar # Make list of x with length Nvar
if not isinstance(y, (list, tuple)):
y = [y] * Nvar # Make list of x with length Nvar
xnew = []
ynew = []
for i in range(0, Nvar):
xnew.append([])
try:
xnew[i].append(x[i])
if not (x[i].shape == (Nx[i][0], )
or x[i].shape == (Nx[i][0], Ny[i][0])
or x[i].shape == (Nt[i][0], Nx[i][0], Ny[i], [0])):
raise ValueError(
"For variable number " + str(i) + ", " + titles[i] +
", the shape of x is not compatible with the shape of the variable. Shape of x should be (Nx), (Nx,Ny) or (Nt,Nx,Ny)."
)
except:
for j in range(0, Nlines[i]):
xnew[i].append(linspace(0, Nx[i][j] - 1, Nx[i][j]))
#x.append(linspace(0,Nx[i][0]-1, Nx[i][0]))
if (Ndims[i][0] == 3):
try:
ynew.append(y[i])
if not (y[i].shape == (Ny[i][0], )
or y[i].shape == (Nx[i][0], Ny[i][0])
or y[i].shape == (Nt[i][0], Nx[i][0], Ny[i], [0])):
raise ValueError(
"For variable number " + str(i) + ", " + titles[i] +
", the shape of y is not compatible with the shape of the variable. Shape of y should be (Ny), (Nx,Ny) or (Nt,Nx,Ny)."
)
except:
ynew.append(linspace(0, Ny[i][0] - 1, Ny[i][0]))
else:
ynew.append(0)
x = xnew
y = ynew
# Determine range of data. Used to ensure constant colour map and
# to set y scale of line plot.
fmax = []
fmin = []
xmax = []
dummymax = []
dummymin = []
clevels = []
for i in range(0, Nvar):
dummymax.append([])
dummymin.append([])
for j in range(0, Nlines[i]):
dummymax[i].append(max(vars[i][j]))
dummymin[i].append(min(vars[i][j]))
fmax.append(max(dummymax[i]))
fmin.append(min(dummymin[i]))
if (symmetric_colors):
absmax = max(abs(array(fmax[i], fmin[i])))
fmax[i] = absmax
fmin[i] = -absmax
for j in range(0, Nlines[i]):
dummymax[i][j] = max(x[i][j])
xmax.append(max(dummymax[i]))
if not (global_colors):
if isclose(fmin[i], fmax[i]):
# add/subtract very small constant in case fmin=fmax=0
thiscontourmin = fmin[i] - 3.e-15 * abs(fmin[i]) - 1.e-36
thiscontourmax = fmax[i] + 3.e-15 * abs(fmax[i]) + 1.e-36
alwayswarn(
"Contour levels too close, adding padding to colorbar range"
)
clevels.append(
linspace(thiscontourmin, thiscontourmax, Ncolors))
else:
clevels.append(linspace(fmin[i], fmax[i], Ncolors))
if (global_colors):
fmaxglobal = max(fmax)
fminglobal = min(fmin)
if isclose(fminglobal, fmaxglobal):
fminglobal = fminglobal - 3.e-15 * abs(fminglobal) - 1.e-36
fmaxglobal = fmaxglobal + 3.e-15 * abs(fmaxglobal) + 1.e-36
for i in range(0, Nvar):
clevels.append(linspace(fminglobal, fmaxglobal, Ncolors))
# Create figures for animation plotting
if (Nvar < 2):
row = 1
col = 1
h = 6.0
w = 8.0
elif (Nvar < 3):
row = 1
col = 2
h = 6.0
w = 12.0
elif (Nvar < 5):
row = 2
col = 2
h = 8.0
w = 12.0
elif (Nvar < 7):
row = 2
col = 3
h = 8.0
w = 14.0
elif (Nvar < 10):
row = 3
col = 3
h = 12.0
w = 14.0
else:
raise ValueError('too many variables...')
fig = plt.figure(window_title, figsize=(w, h))
title = fig.suptitle(r' ', fontsize=14)
# Initiate all list variables required for plotting here
ax = []
lines = []
plots = []
cbars = []
xstride = []
ystride = []
r = []
theta = []
# Initiate figure frame
for i in range(0, Nvar):
lines.append([])
if (lineplot[i] == 1):
ax.append(fig.add_subplot(row, col, i + 1))
ax[i].set_xlim((0, xmax[i]))
ax[i].set_ylim((fmin[i], fmax[i]))
for j in range(0, Nlines[i]):
lines[i].append(ax[i].plot([], [],
lw=2,
label=legendlabels[i][j])[0])
#Need the [0] to 'unpack' the line object from tuple. Alternatively:
#lines[i], = lines[i]
ax[i].set_xlabel(r'x')
ax[i].set_ylabel(titles[i])
if (Nlines[i] != 1):
legendneeded = 1
for k in range(0, i):
if (Nlines[i] == Nlines[k]):
legendneeded = 0
if (legendneeded == 1):
plt.axes(ax[i])
plt.legend(loc=0)
# Pad out unused list variables with zeros
plots.append(0)
cbars.append(0)
xstride.append(0)
ystride.append(0)
r.append(0)
theta.append(0)
elif (contour[i] == 1):
ax.append(fig.add_subplot(row, col, i + 1))
#ax[i].set_xlim((0,Nx[i][0]-1))
#ax[i].set_ylim((0,Ny[i][0]-1))
ax[i].set_xlim(min(x[i]), max(x[i]))
ax[i].set_ylim(min(y[i]), max(y[i]))
ax[i].set_xlabel(r'x')
ax[i].set_ylabel(r'y')
ax[i].set_title(titles[i])
if hold_aspect:
ax[i].set_aspect('equal')
thisx = x[i][0]
if len(thisx.shape) == 3:
thisx = thisx[0]
thisy = y[i]
if len(thisy.shape) == 3:
thisy = thisy[0]
plots.append(ax[i].contourf(thisx.T,
thisy.T,
vars[i][0][0, :, :].T,
Ncolors,
cmap=cmap,
lw=0,
levels=clevels[i]))
plt.axes(ax[i])
cbars.append(fig.colorbar(plots[i], format='%1.1e'))
# Pad out unused list variables with zeros
lines[i].append(0)
xstride.append(0)
ystride.append(0)
r.append(0)
theta.append(0)
elif (surf[i] == 1):
if (len(x[i][0].shape) == 1 and len(y[i].shape) == 1):
# plot_wireframe() requires 2d arrays for x and y coordinates
x[i][0], y[i] = meshgrid(x[i][0], y[i])
thisx = x[i][0]
if len(thisx.shape) == 3:
thisx = thisx[0]
thisy = y[i]
if len(thisy.shape) == 3:
thisy = thisy[0]
if (Nx[i][0] <= 20):
xstride.append(1)
else:
xstride.append(int(floor(Nx[i][0] / 20)))
if (Ny[i][0] <= 20):
ystride.append(1)
else:
ystride.append(int(floor(Ny[i][0] / 20)))
ax.append(fig.add_subplot(row, col, i + 1, projection='3d'))
plots.append(ax[i].plot_wireframe(thisx,
thisy,
vars[i][0][0, :, :].T,
rstride=ystride[i],
cstride=xstride[i]))
title = fig.suptitle(r'', fontsize=14)
ax[i].set_xlabel(r'x')
ax[i].set_ylabel(r'y')
ax[i].set_zlabel(titles[i])
# Pad out unused list variables with zeros
lines[i].append(0)
cbars.append(0)
r.append(0)
theta.append(0)
elif (polar[i] == 1):
r.append(linspace(1, Nx[i][0], Nx[i][0]))
theta.append(linspace(0, 2 * pi, Ny[i][0]))
r[i], theta[i] = meshgrid(r[i], theta[i])
ax.append(fig.add_subplot(row, col, i + 1, projection='polar'))
plots.append(ax[i].contourf(theta[i],
r[i],
vars[i][0][0, :, :].T,
cmap=cmap,
levels=clevels[i]))
plt.axes(ax[i])
cbars.append(fig.colorbar(plots[i], format='%1.1e'))
ax[i].set_rmax(Nx[i][0] - 1)
ax[i].set_title(titles[i])
# Pad out unused list variables with zeros
lines[i].append(0)
xstride.append(0)
ystride.append(0)
def onClick(event):
global pause
pause ^= True
def control():
global j, pause
if j == Nframes - 1: j = -1
if not pause:
j = j + 1
return j
# Animation function
def animate(i):
j = control()
index = j * intv
for j in range(0, Nvar):
#Default to clearing axis between frames on all plots except line plots
if (clear_between_frames is None
and lineplot[j] != 1) or clear_between_frames is True:
ax[j].cla(
) #Clear axis between frames so that masked arrays can be plotted
if (lineplot[j] == 1):
for k in range(0, Nlines[j]):
lines[j][k].set_data(x[j][k], vars[j][k][index, :])
elif (contour[j] == 1):
thisx = x[j][0]
if len(thisx.shape) == 3:
thisx = thisx[index]
thisy = y[j]
if len(thisy.shape) == 3:
thisy = thisy[index]
plots[j] = ax[j].contourf(x[j][0].T,
y[j].T,
vars[j][0][index, :, :].T,
Ncolors,
cmap=cmap,
lw=0,
levels=clevels[j])
ax[j].set_xlabel(r'x')
ax[j].set_ylabel(r'y')
ax[j].set_title(titles[j])
elif (surf[j] == 1):
thisx = x[j][0]
if len(thisx.shape) == 3:
thisx = thisx[index]
thisy = y[j][0]
if len(thisy.shape) == 3:
thisy = thisy[index]
ax[j] = fig.add_subplot(row, col, j + 1, projection='3d')
plots[j] = ax[j].plot_wireframe(thisx,
thisy,
vars[j][0][index, :, :].T,
rstride=ystride[j],
cstride=xstride[j])
ax[j].set_zlim(fmin[j], fmax[j])
ax[j].set_xlabel(r'x')
ax[j].set_ylabel(r'y')
ax[j].set_title(titles[j])
elif (polar[j] == 1):
plots[j] = ax[j].contourf(theta[j],
r[j],
vars[j][0][index, :, :].T,
cmap=cmap,
levels=clevels[j])
ax[j].set_rmax(Nx[j][0] - 1)
ax[j].set_title(titles[j])
if t_array is not None:
title.set_text('t = %1.2e' % t_array[index])
else:
title.set_text('t = %i' % index)
return plots
def init():
global j, pause
j = -2
pause = False
return animate(0)
# Call Animation function
fig.canvas.mpl_connect('button_press_event', onClick)
anim = animation.FuncAnimation(fig,
animate,
init_func=init,
frames=Nframes)
#If movie is not passed as a string assign the default filename
if (movie == 1):
movie = 'animation.mp4'
# Save movie with given or default name
if ((isinstance(movie, str) == 1)):
movietype = movie.split('.')[-1]
if movietype == 'mp4':
try:
anim.save(movie,
writer=FFwriter,
fps=fps,
dpi=dpi,
extra_args=['-vcodec', 'libx264'])
except Exception:
#Try specifying writer by string if ffmpeg not found
try:
anim.save(movie,
writer='ffmpeg',
fps=fps,
dpi=dpi,
extra_args=['-vcodec', 'libx264'])
except Exception:
print('Save failed: Check ffmpeg path')
raise
elif movietype == 'gif':
anim.save(movie + '.gif', writer='imagemagick', fps=fps, dpi=dpi)
else:
raise ValueError(
"Unrecognized file type for movie. Supported types are .mp4 and .gif"
)
# Show animation if not saved or returned, otherwise close the plot
if (movie == 0 and return_animation == 0):
plt.show()
else:
plt.close()
# Return animation object
if (return_animation == 1):
return (anim)
def __init__(self,
path=".",
prefix="BOUT.dmp",
suffix=None,
caching=False,
DataFileCaching=True,
**kwargs):
"""
Initialise BoutOutputs object
"""
self._path = path
# normalize prefix by removing trailing '.' if present
self._prefix = prefix.rstrip('.')
if suffix == None:
temp_file_list = glob.glob(
os.path.join(self._path, self._prefix + "*"))
latest_file = max(temp_file_list, key=os.path.getctime)
self._suffix = latest_file.split(".")[-1]
else:
# normalize suffix by removing leading '.' if present
self._suffix = suffix.lstrip('.')
self._caching = caching
self._DataFileCaching = DataFileCaching
self._kwargs = kwargs
# Label for this data
self.label = path
self._file_list = glob.glob(
os.path.join(path, self._prefix + "*" + self._suffix))
if not suffix == None:
latest_file = max(self._file_list, key=os.path.getctime)
# if suffix==None we already found latest_file
# Check that the path contains some data
if len(self._file_list) == 0:
raise ValueError("ERROR: No data files found")
# Available variables
self.varNames = []
self.dimensions = {}
self.evolvingVariableNames = []
with DataFile(latest_file) as f:
npes = f.read("NXPE") * f.read("NYPE")
if len(self._file_list) != npes:
alwayswarn("Too many data files, reading most recent ones")
if npes == 1:
# single output file
# do like this to catch, e.g. either 'BOUT.dmp.nc' or 'BOUT.dmp.0.nc'
self._file_list = [latest_file]
else:
self._file_list = [
os.path.join(
path,
self._prefix + "." + str(i) + "." + self._suffix)
for i in range(npes)
]
# Get variable names
self.varNames = f.keys()
for name in f.keys():
dimensions = f.dimensions(name)
self.dimensions[name] = dimensions
if name != "t_array" and "t" in dimensions:
self.evolvingVariableNames.append(name)
# Private variables
if self._caching:
from collections import OrderedDict
self._datacache = OrderedDict()
if self._caching is not True:
# Track the size of _datacache and limit it to a maximum of _caching
try:
# Check that _caching is a number of some sort
float(self._caching)
except ValueError:
raise ValueError(
"BoutOutputs: Invalid value for caching argument. Caching should be either a number (giving the maximum size of the cache in GB), True for unlimited size or False for no caching."
)
self._datacachesize = 0
self._datacachemaxsize = self._caching * 1.e9
self._DataFileCache = None
def __init__(self,
filename="BOUT.inp",
name="root",
gridfilename=None,
nx=None,
ny=None,
nz=None):
BoutOptions.__init__(self, name)
# Open the file
with open(filename, "r") as f:
# Go through each line in the file
section = self # Start with root section
for linenr, line in enumerate(f.readlines()):
# First remove comments, either # or ;
startpos = line.find("#")
if startpos != -1:
line = line[:startpos]
startpos = line.find(";")
if startpos != -1:
line = line[:startpos]
# Check section headers
startpos = line.find("[")
endpos = line.find("]")
if startpos != -1:
# A section heading
if endpos == -1:
raise SyntaxError("Missing ']' on line %d" %
(linenr, ))
line = line[(startpos + 1):endpos].strip()
section = self
while True:
scorepos = line.find(":")
if scorepos == -1:
break
sectionname = line[0:scorepos]
line = line[(scorepos + 1):]
section = section.getSection(sectionname)
section = section.getSection(line)
else:
# A key=value pair
eqpos = line.find("=")
if eqpos == -1:
# No '=', so just set to true
section[line.strip()] = True
else:
value = line[(eqpos + 1):].strip()
try:
# Try to convert to an integer
value = int(value)
except ValueError:
try:
# Try to convert to float
value = float(value)
except ValueError:
# Leave as a string
pass
section[line[:eqpos].strip()] = value
try:
# define arrays of x, y, z to be used for substitutions
gridfile = None
nzfromfile = None
if gridfilename:
if nx is not None or ny is not None:
raise ValueError("nx or ny given as inputs even though "
"gridfilename was given explicitly, "
"don't know which parameters to choose")
with DataFile(gridfilename) as gridfile:
self.nx = float(gridfile["nx"])
self.ny = float(gridfile["ny"])
try:
nzfromfile = gridfile["MZ"]
except KeyError:
pass
elif nx or ny:
if nx is None:
raise ValueError(
"nx not specified. If either nx or ny are given, then both must be."
)
if ny is None:
raise ValueError(
"ny not specified. If either nx or ny are given, then both must be."
)
self.nx = nx
self.ny = ny
else:
try:
self.nx = self["mesh"].evaluate_scalar("nx")
self.ny = self["mesh"].evaluate_scalar("ny")
except KeyError:
try:
# get nx, ny, nz from output files
from boutdata.collect import findFiles
file_list = findFiles(path=os.path.dirname(),
prefix="BOUT.dmp")
with DataFile(file_list[0]) as f:
self.nx = f["nx"]
self.ny = f["ny"]
nzfromfile = f["MZ"]
except (IOError, KeyError):
try:
gridfilename = self["mesh"]["file"]
except KeyError:
gridfilename = self["grid"]
with DataFile(gridfilename) as gridfile:
self.nx = float(gridfile["nx"])
self.ny = float(gridfile["ny"])
try:
nzfromfile = float(gridfile["MZ"])
except KeyError:
pass
if nz is not None:
self.nz = nz
else:
try:
self.nz = self["mesh"].evaluate_scalar("nz")
except KeyError:
try:
self.nz = self.evaluate_scalar("mz")
except KeyError:
if nzfromfile is not None:
self.nz = nzfromfile
mxg = self._keys.get("MXG", 2)
myg = self._keys.get("MYG", 2)
# make self.x, self.y, self.z three dimensional now so
# that expressions broadcast together properly.
self.x = numpy.linspace((0.5 - mxg) / (self.nx - 2 * mxg),
1. - (0.5 - mxg) / (self.nx - 2 * mxg),
self.nx)[:, numpy.newaxis, numpy.newaxis]
self.y = 2. * numpy.pi * numpy.linspace(
(0.5 - myg) / self.ny, 1. - (0.5 - myg) / self.ny,
self.ny + 2 * myg)[numpy.newaxis, :, numpy.newaxis]
self.z = 2. * numpy.pi * numpy.linspace(
0.5 / self.nz, 1. - 0.5 / self.nz, self.nz)[numpy.newaxis,
numpy.newaxis, :]
except Exception as e:
alwayswarn("While building x, y, z coordinate arrays, an "
"exception occured: " + str(e) +
"\nEvaluating non-scalar options not available")
def __init__(self, path=".", prefix="BOUT.dmp", suffix=None, caching=False,
DataFileCaching=True, **kwargs):
"""
Initialise BoutOutputs object
"""
self._path = path
# normalize prefix by removing trailing '.' if present
self._prefix = prefix.rstrip('.')
if suffix == None:
temp_file_list = glob.glob(
os.path.join(self._path, self._prefix + "*"))
latest_file = max(temp_file_list, key=os.path.getctime)
self._suffix = latest_file.split(".")[-1]
else:
# normalize suffix by removing leading '.' if present
self._suffix = suffix.lstrip('.')
self._caching = caching
self._DataFileCaching = DataFileCaching
self._kwargs = kwargs
# Label for this data
self.label = path
self._file_list = glob.glob(os.path.join(
path, self._prefix + "*" + self._suffix))
if not suffix == None:
latest_file = max(self._file_list, key=os.path.getctime)
# if suffix==None we already found latest_file
# Check that the path contains some data
if len(self._file_list) == 0:
raise ValueError("ERROR: No data files found")
# Available variables
self.varNames = []
self.dimensions = {}
self.evolvingVariableNames = []
with DataFile(latest_file) as f:
npes = f.read("NXPE")*f.read("NYPE")
if len(self._file_list) != npes:
alwayswarn("Too many data files, reading most recent ones")
if npes == 1:
# single output file
# do like this to catch, e.g. either 'BOUT.dmp.nc' or 'BOUT.dmp.0.nc'
self._file_list = [latest_file]
else:
self._file_list = [os.path.join(
path, self._prefix + "." + str(i) + "." + self._suffix) for i in range(npes)]
# Get variable names
self.varNames = f.keys()
for name in f.keys():
dimensions = f.dimensions(name)
self.dimensions[name] = dimensions
if name != "t_array" and "t" in dimensions:
self.evolvingVariableNames.append(name)
# Private variables
if self._caching:
from collections import OrderedDict
self._datacache = OrderedDict()
if self._caching is not True:
# Track the size of _datacache and limit it to a maximum of _caching
try:
# Check that _caching is a number of some sort
float(self._caching)
except ValueError:
raise ValueError(
"BoutOutputs: Invalid value for caching argument. Caching should be either a number (giving the maximum size of the cache in GB), True for unlimited size or False for no caching.")
self._datacachesize = 0
self._datacachemaxsize = self._caching*1.e9
self._DataFileCache = None
def __init__(self, filename="BOUT.inp", name="root", gridfilename=None, nx=None, ny=None, nz=None):
BoutOptions.__init__(self, name)
# Open the file
with open(filename, "r") as f:
# Go through each line in the file
section = self # Start with root section
for linenr, line in enumerate(f.readlines()):
# First remove comments, either # or ;
startpos = line.find("#")
if startpos != -1:
line = line[:startpos]
startpos = line.find(";")
if startpos != -1:
line = line[:startpos]
# Check section headers
startpos = line.find("[")
endpos = line.find("]")
if startpos != -1:
# A section heading
if endpos == -1:
raise SyntaxError("Missing ']' on line %d" % (linenr,))
line = line[(startpos+1):endpos].strip()
section = self
while True:
scorepos = line.find(":")
if scorepos == -1:
break
sectionname = line[0:scorepos]
line = line[(scorepos+1):]
section = section.getSection(sectionname)
section = section.getSection(line)
else:
# A key=value pair
eqpos = line.find("=")
if eqpos == -1:
# No '=', so just set to true
section[line.strip()] = True
else:
value = line[(eqpos+1):].strip()
try:
# Try to convert to an integer
value = int(value)
except ValueError:
try:
# Try to convert to float
value = float(value)
except ValueError:
# Leave as a string
pass
section[line[:eqpos].strip()] = value
try:
# define arrays of x, y, z to be used for substitutions
gridfile = None
nzfromfile = None
if gridfilename:
if nx is not None or ny is not None:
raise ValueError("nx or ny given as inputs even though "
"gridfilename was given explicitly, "
"don't know which parameters to choose")
with DataFile(gridfilename) as gridfile:
self.nx = float(gridfile["nx"])
self.ny = float(gridfile["ny"])
try:
nzfromfile = gridfile["MZ"]
except KeyError:
pass
elif nx or ny:
if nx is None:
raise ValueError("nx not specified. If either nx or ny are given, then both must be.")
if ny is None:
raise ValueError("ny not specified. If either nx or ny are given, then both must be.")
self.nx = nx
self.ny = ny
else:
try:
self.nx = self["mesh"].evaluate_scalar("nx")
self.ny = self["mesh"].evaluate_scalar("ny")
except KeyError:
try:
# get nx, ny, nz from output files
from boutdata.collect import findFiles
file_list = findFiles(path=os.path.dirname(), prefix="BOUT.dmp")
with DataFile(file_list[0]) as f:
self.nx = f["nx"]
self.ny = f["ny"]
nzfromfile = f["MZ"]
except (IOError, KeyError):
try:
gridfilename = self["mesh"]["file"]
except KeyError:
gridfilename = self["grid"]
with DataFile(gridfilename) as gridfile:
self.nx = float(gridfile["nx"])
self.ny = float(gridfile["ny"])
try:
nzfromfile = float(gridfile["MZ"])
except KeyError:
pass
if nz is not None:
self.nz = nz
else:
try:
self.nz = self["mesh"].evaluate_scalar("nz")
except KeyError:
try:
self.nz = self.evaluate_scalar("mz")
except KeyError:
if nzfromfile is not None:
self.nz = nzfromfile
mxg = self._keys.get("MXG", 2)
myg = self._keys.get("MYG", 2)
# make self.x, self.y, self.z three dimensional now so
# that expressions broadcast together properly.
self.x = numpy.linspace((0.5 - mxg)/(self.nx - 2*mxg),
1. - (0.5 - mxg)/(self.nx - 2*mxg),
self.nx)[:, numpy.newaxis, numpy.newaxis]
self.y = 2.*numpy.pi*numpy.linspace((0.5 - myg)/self.ny,
1.-(0.5 - myg)/self.ny,
self.ny + 2*myg)[numpy.newaxis, :, numpy.newaxis]
self.z = 2.*numpy.pi*numpy.linspace(0.5/self.nz,
1.-0.5/self.nz,
self.nz)[numpy.newaxis, numpy.newaxis, :]
except Exception as e:
alwayswarn("While building x, y, z coordinate arrays, an "
"exception occured: " + str(e) +
"\nEvaluating non-scalar options not available")
def showdata(vars, titles=[], legendlabels=[], surf=[], polar=[], tslice=0, t_array=None,
movie=0, fps=28, dpi=200, intv=1, Ncolors=25, x=[], y=[],
global_colors=False, symmetric_colors=False, hold_aspect=False,
cmap=None, clear_between_frames=None, return_animation=False, window_title=""):
"""A Function to animate time dependent data from BOUT++
To animate multiple variables on different axes:
>>> showdata([var1, var2, var3])
To animate more than one line on a single axes:
>>> showdata([[var1, var2, var3]])
The default graph types are:
2D (time + 1 spatial dimension) arrays = animated line plot
3D (time + 2 spatial dimensions) arrays = animated contour plot.
To use surface or polar plots:
>>> showdata(var, surf=1)
>>> showdata(var, polar=1)
Can plot different graph types on different axes. Default graph types will
be used depending on the dimensions of the input arrays. To specify
polar/surface plots on different axes:
>>> showdata([var1, var2], surf=[1, 0], polar=[0, 1])
Movies require FFmpeg (for .mp4) and/or ImageMagick (for .gif) to be
installed. The 'movie' option can be set to 1 (which will produce an mp4
called 'animation.mp4'), to a name with no extension (which will produce an
mp4 called '<name>.mp4')
The `tslice` variable is used to control the time value that is printed on
each frame of the animation. If the input data matches the time values
found within BOUT++'s dmp data files, then these time values will be used.
Otherwise, an integer counter is used.
The `cmap` variable (if specified) will set the colormap used in the plot
cmap must be a matplotlib colormap instance, or the name of a registered
matplotlib colormap
During animation click once to stop in the current frame. Click again to
continue.
Parameters
----------
vars : array_like or list of array_like
Variable or list of variables to plot
titles : str or list of str, optional
Title or list of titles for each axis
legendlabels : str or list of str, optional
Legend or list of legends for each variable
surf : list of int
Which axes to plot as a surface plot
polar : list of int
Which axes to plot as a polar plot
tslice : list of int
Use these time values from a dump file (see above)
t_array : array
Pass in t_array using this argument to use the simulation time in plot
titles. Otherwise, just use the t-index.
movie : int
If 1, save the animation to file
fps : int
Number of frames per second to use when saving animation
dpi : int
Dots per inch to use when saving animation
intv : int
???
Ncolors : int
Number of levels in contour plots
x, y : array_like, list of array_like
X, Y coordinates
global_colors : bool
If "vars" is a list the colorlevels are determined from the
maximum of the maxima and and the minimum of the minima in all
fields in vars
symmetric_colors : bool
Colour levels are symmetric
hold_aspect : bool
Use equal aspect ratio in plots
cmap : colormap
A matplotlib colormap instance to use
clear_between_frames : bool, optional
- Default (None) - all plots except line plots will clear between frames
- True - all plots will clear between frames
- False - no plots will clear between frames
return_animation : bool
Return the matplotlib animation instance
window_title : str
Give a title for the animation window
TODO
----
- Replace empty lists in signature with None
- Use bools in sensible places
- Put massive list of arguments in kwargs
- Speed up animations ????
- Look at theta in polar plots - periodic?!?
- Log axes, colorbars
- Figureplot
"""
plt.ioff()
# Check to see whether vars is a list or not.
if isinstance(vars, list):
Nvar = len(vars)
else:
vars = [vars]
Nvar = len(vars)
if Nvar < 1:
raise ValueError("No data supplied")
# Check to see whether each variable is a list - used for line plots only
Nlines = []
for i in range(0, Nvar):
if isinstance(vars[i], list):
Nlines.append(len(vars[i]))
else:
Nlines.append(1)
vars[i] = [vars[i]]
# Sort out titles
if len(titles) == 0:
for i in range(0,Nvar):
titles.append(('Var' + str(i+1)))
elif len(titles) != Nvar:
raise ValueError('The length of the titles input list must match the length of the vars list.')
# Sort out legend labels
if len(legendlabels) == 0:
for i in range(0,Nvar):
legendlabels.append([])
for j in range(0,Nlines[i]):
legendlabels[i].append(chr(97+j))
elif (isinstance(legendlabels[0], list) != 1):
if Nvar != 1:
check = 0
for i in range(0,Nvar):
if len(legendlabels) != Nlines[i]:
check = check+1
if check == 0:
alwayswarn("The legendlabels list does not contain a sublist for each variable, but its length matches the number of lines on each plot. Will apply labels to each plot")
legendlabelsdummy = []
for i in range(0, Nvar):
legendlabelsdummy.append([])
for j in range(0,Nlines[i]):
legendlabelsdummy[i].append(legendlabels[j])
legendlabels = legendlabelsdummy
else:
alwayswarn("The legendlabels list does not contain a sublist for each variable, and it's length does not match the number of lines on each plot. Will default apply labels to each plot")
legendlabels = []
for i in range(0,Nvar):
legendlabels.append([])
for j in range(0,Nlines[i]):
legendlabels[i].append(chr(97+j))
else:
if (Nlines[0] == len(legendlabels)):
legendlabels = [legendlabels]
elif len(legendlabels) != Nvar:
alwayswarn("The length of the legendlabels list does not match the length of the vars list, will continue with default values")
legendlabels = []
for i in range(0,Nvar):
legendlabels.append([])
for j in range(0,Nlines[i]):
legendlabels[i].append(chr(97+j))
else:
for i in range(0,Nvar):
if isinstance(legendlabels[i], list):
if len(legendlabels[i]) != Nlines[i]:
alwayswarn('The length of the legendlabel (sub)list for each plot does not match the number of datasets for each plot. Will continue with default values')
legendlabels[i] = []
for j in range(0,Nlines[i]):
legendlabels[i].append(chr(97+j))
else:
legendlabels[i] = [legendlabels[i]]
if len(legendlabels[i]) != Nlines[i]:
alwayswarn('The length of the legendlabel (sub)list for each plot does not match the number of datasets for each plot. Will continue with default values')
legendlabels[i] = []
for j in range(0,Nlines[i]):
legendlabels[i].append(chr(97+j))
# Sort out surf list
if isinstance(surf, list):
if (len(surf) == Nvar):
for i in range(0, Nvar):
if surf[i] >= 1:
surf[i] = 1
else:
surf[i] = 0
elif (len(surf) == 1):
if surf[0] >= 1:
surf[0] = 1
else:
surf[0] = 0
if (Nvar > 1):
for i in range(1,Nvar):
surf.append(surf[0])
elif (len(surf) == 0):
for i in range(0,Nvar):
surf.append(0)
else:
alwayswarn('Length of surf list does not match number of variables. Will default to no polar plots')
for i in range(0,Nvar):
surf.append(0)
else:
surf = [surf]
if surf[0] >= 1:
surf[0] = 1
else:
surf[0] = 0
if (Nvar > 1):
for i in range(1,Nvar):
surf.append(surf[0])
# Sort out polar list
if isinstance(polar, list):
if (len(polar) == Nvar):
for i in range(0, Nvar):
if polar[i] >= 1:
polar[i] = 1
else:
polar[i] = 0
elif (len(polar) == 1):
if polar[0] >= 1:
polar[0] = 1
else:
polar[0] = 0
if (Nvar > 1):
for i in range(1,Nvar):
polar.append(polar[0])
elif (len(polar) == 0):
for i in range(0,Nvar):
polar.append(0)
else:
alwayswarn('Length of polar list does not match number of variables. Will default to no polar plots')
for i in range(0,Nvar):
polar.append(0)
else:
polar = [polar]
if polar[0] >= 1:
polar[0] = 1
else:
polar[0] = 0
if (Nvar > 1):
for i in range(1,Nvar):
polar.append(polar[0])
# Determine shapes of arrays
dims = []
Ndims = []
lineplot = []
contour = []
for i in range(0,Nvar):
dims.append([])
Ndims.append([])
for j in range(0, Nlines[i]):
dims[i].append(array((vars[i][j].shape)))
Ndims[i].append(dims[i][j].shape[0])
# Perform check to make sure that data is either 2D or 3D
if (Ndims[i][j] < 2):
raise ValueError('data must be either 2 or 3 dimensional. Exiting')
if (Ndims[i][j] > 3):
raise ValueError('data must be either 2 or 3 dimensional. Exiting')
if ((Ndims[i][j] == 2) & (polar[i] != 0)):
alwayswarn('Data must be 3 dimensional (time, r, theta) for polar plots. Will plot lineplot instead')
if ((Ndims[i][j] == 2) & (surf[i] != 0)):
alwayswarn('Data must be 3 dimensional (time, x, y) for surface plots. Will plot lineplot instead')
if ((Ndims[i][j] == 3) & (Nlines[i] != 1)):
raise ValueError('cannot have multiple sets of 3D (time + 2 spatial dimensions) on each subplot')
if ((Ndims[i][j] != Ndims[i][0])):
raise ValueError('Error, Number of dimensions must be the same for all variables on each plot.')
if (Ndims[i][0] == 2): # Set polar and surf list entries to 0
polar[i] = 0
surf[i] = 0
lineplot.append(1)
contour.append(0)
else:
if ((polar[i] == 1) & (surf[i] == 1)):
alwayswarn('Cannot do polar and surface plots at the same time. Default to contour plot')
contour.append(1)
lineplot.append(0)
polar[i] = 0
surf[i] = 0
elif (polar[i] == 1) | (surf[i] == 1):
contour.append(0)
lineplot.append(0)
else:
contour.append(1)
lineplot.append(0)
# Obtain size of data arrays
Nt = []
Nx = []
Ny = []
for i in range(0, Nvar):
Nt.append([])
Nx.append([])
Ny.append([])
for j in range(0, Nlines[i]):
Nt[i].append(vars[i][j].shape[0])
Nx[i].append(vars[i][j].shape[1])
if (Nt[i][j] != Nt[0][0]):
raise ValueError('time dimensions must be the same for all variables.')
#if (Nx[i][j] != Nx[i][0]):
# raise ValueError('Dimensions must be the same for all variables on each plot.')
if (Ndims[i][j] == 3):
Ny[i].append(vars[i][j].shape[2])
#if (Ny[i][j] != Ny[i][0]):
# raise ValueError('Dimensions must be the same for all variables.')
# Obtain number of frames
Nframes = int(Nt[0][0]/intv)
# Generate grids for plotting
# Try to use provided grids where possible
# If x and/or y are not lists, apply to all variables
if not isinstance(x, (list,tuple)):
x = [x]*Nvar # Make list of x with length Nvar
if not isinstance(y, (list,tuple)):
y = [y]*Nvar # Make list of x with length Nvar
xnew = []
ynew = []
for i in range(0,Nvar):
xnew.append([])
try:
xnew[i].append(x[i])
if not (x[i].shape==(Nx[i][0],) or x[i].shape==(Nx[i][0],Ny[i][0]) or x[i].shape==(Nt[i][0],Nx[i][0],Ny[i],[0])):
raise ValueError("For variable number "+str(i)+", "+titles[i]+", the shape of x is not compatible with the shape of the variable. Shape of x should be (Nx), (Nx,Ny) or (Nt,Nx,Ny).")
except:
for j in range(0, Nlines[i]):
xnew[i].append(linspace(0,Nx[i][j]-1, Nx[i][j]))
#x.append(linspace(0,Nx[i][0]-1, Nx[i][0]))
if (Ndims[i][0] == 3):
try:
ynew.append(y[i])
if not (y[i].shape==(Ny[i][0],) or y[i].shape==(Nx[i][0],Ny[i][0]) or y[i].shape==(Nt[i][0],Nx[i][0],Ny[i],[0])):
raise ValueError("For variable number "+str(i)+", "+titles[i]+", the shape of y is not compatible with the shape of the variable. Shape of y should be (Ny), (Nx,Ny) or (Nt,Nx,Ny).")
except:
ynew.append(linspace(0, Ny[i][0]-1, Ny[i][0]))
else:
ynew.append(0)
x = xnew
y = ynew
# Determine range of data. Used to ensure constant colour map and
# to set y scale of line plot.
fmax = []
fmin = []
xmax = []
dummymax = []
dummymin = []
clevels = []
for i in range(0,Nvar):
dummymax.append([])
dummymin.append([])
for j in range(0,Nlines[i]):
dummymax[i].append(max(vars[i][j]))
dummymin[i].append(min(vars[i][j]))
fmax.append(max(dummymax[i]))
fmin.append(min(dummymin[i]))
if(symmetric_colors):
absmax =max(abs(array(fmax[i], fmin[i])))
fmax[i] = absmax
fmin[i] = -absmax
for j in range(0,Nlines[i]):
dummymax[i][j] = max(x[i][j])
xmax.append(max(dummymax[i]))
if not (global_colors):
if isclose(fmin[i], fmax[i]):
# add/subtract very small constant in case fmin=fmax=0
thiscontourmin = fmin[i] - 3.e-15*abs(fmin[i]) - 1.e-36
thiscontourmax = fmax[i] + 3.e-15*abs(fmax[i]) + 1.e-36
alwayswarn("Contour levels too close, adding padding to colorbar range")
clevels.append(linspace(thiscontourmin, thiscontourmax, Ncolors))
else:
clevels.append(linspace(fmin[i], fmax[i], Ncolors))
if(global_colors):
fmaxglobal = max(fmax)
fminglobal = min(fmin)
if isclose(fminglobal, fmaxglobal):
fminglobal = fminglobal - 3.e-15*abs(fminglobal) - 1.e-36
fmaxglobal = fmaxglobal + 3.e-15*abs(fmaxglobal) + 1.e-36
for i in range(0,Nvar):
clevels.append(linspace(fminglobal, fmaxglobal, Ncolors))
# Create figures for animation plotting
if (Nvar < 2):
row = 1
col = 1
h = 6.0
w = 8.0
elif (Nvar <3):
row = 1
col = 2
h = 6.0
w = 12.0
elif (Nvar < 5):
row = 2
col = 2
h = 8.0
w = 12.0
elif (Nvar < 7):
row = 2
col = 3
h = 8.0
w = 14.0
elif (Nvar < 10) :
row = 3
col = 3
h = 12.0
w = 14.0
else:
raise ValueError('too many variables...')
fig = plt.figure(window_title, figsize=(w,h))
title = fig.suptitle(r' ', fontsize=14 )
# Initiate all list variables required for plotting here
ax = []
lines = []
plots = []
cbars = []
xstride = []
ystride = []
r = []
theta = []
# Initiate figure frame
for i in range(0,Nvar):
lines.append([])
if (lineplot[i] == 1):
ax.append(fig.add_subplot(row,col,i+1))
ax[i].set_xlim((0,xmax[i]))
ax[i].set_ylim((fmin[i], fmax[i]))
for j in range(0,Nlines[i]):
lines[i].append(ax[i].plot([],[],lw=2, label = legendlabels[i][j])[0])
#Need the [0] to 'unpack' the line object from tuple. Alternatively:
#lines[i], = lines[i]
ax[i].set_xlabel(r'x')
ax[i].set_ylabel(titles[i])
if (Nlines[i] != 1):
legendneeded = 1
for k in range(0,i):
if (Nlines[i] == Nlines[k]):
legendneeded = 0
if (legendneeded == 1):
plt.axes(ax[i])
plt.legend(loc = 0)
# Pad out unused list variables with zeros
plots.append(0)
cbars.append(0)
xstride.append(0)
ystride.append(0)
r.append(0)
theta.append(0)
elif (contour[i] == 1):
ax.append(fig.add_subplot(row,col,i+1))
#ax[i].set_xlim((0,Nx[i][0]-1))
#ax[i].set_ylim((0,Ny[i][0]-1))
ax[i].set_xlim(min(x[i]),max(x[i]))
ax[i].set_ylim(min(y[i]),max(y[i]))
ax[i].set_xlabel(r'x')
ax[i].set_ylabel(r'y')
ax[i].set_title(titles[i])
if hold_aspect:
ax[i].set_aspect('equal')
thisx = x[i][0]
if len(thisx.shape) == 3:
thisx = thisx[0]
thisy = y[i]
if len(thisy.shape) == 3:
thisy = thisy[0]
plots.append(ax[i].contourf(thisx.T,thisy.T,vars[i][0][0,:,:].T, Ncolors, cmap=cmap, lw=0, levels=clevels[i] ))
plt.axes(ax[i])
cbars.append(fig.colorbar(plots[i], format='%1.1e'))
# Pad out unused list variables with zeros
lines[i].append(0)
xstride.append(0)
ystride.append(0)
r.append(0)
theta.append(0)
elif (surf[i] == 1):
if (len(x[i][0].shape)==1 and len(y[i].shape)==1):
# plot_wireframe() requires 2d arrays for x and y coordinates
x[i][0],y[i] = meshgrid(x[i][0],y[i])
thisx = x[i][0]
if len(thisx.shape) == 3:
thisx = thisx[0]
thisy = y[i]
if len(thisy.shape) == 3:
thisy = thisy[0]
if (Nx[i][0]<= 20):
xstride.append(1)
else:
xstride.append(int(floor(Nx[i][0]/20)))
if (Ny[i][0]<=20):
ystride.append(1)
else:
ystride.append(int(floor(Ny[i][0]/20)))
ax.append(fig.add_subplot(row,col,i+1, projection='3d'))
plots.append(ax[i].plot_wireframe(thisx, thisy, vars[i][0][0,:,:].T, rstride=ystride[i], cstride=xstride[i]))
title = fig.suptitle(r'', fontsize=14 )
ax[i].set_xlabel(r'x')
ax[i].set_ylabel(r'y')
ax[i].set_zlabel(titles[i])
# Pad out unused list variables with zeros
lines[i].append(0)
cbars.append(0)
r.append(0)
theta.append(0)
elif (polar[i] == 1):
r.append(linspace(1,Nx[i][0], Nx[i][0]))
theta.append(linspace(0,2*pi, Ny[i][0]))
r[i],theta[i] = meshgrid(r[i], theta[i])
ax.append(fig.add_subplot(row,col,i+1, projection='polar'))
plots.append(ax[i].contourf(theta[i], r[i], vars[i][0][0,:,:].T, cmap=cmap, levels=clevels[i]))
plt.axes(ax[i])
cbars.append(fig.colorbar(plots[i], format='%1.1e'))
ax[i].set_rmax(Nx[i][0]-1)
ax[i].set_title(titles[i])
# Pad out unused list variables with zeros
lines[i].append(0)
xstride.append(0)
ystride.append(0)
def onClick(event):
global pause
pause ^= True
def control():
global j, pause
if j == Nframes-1 : j = -1
if not pause:
j=j+1
return j
# Animation function
def animate(i):
j=control()
index = j*intv
for j in range(0,Nvar):
#Default to clearing axis between frames on all plots except line plots
if (clear_between_frames is None and lineplot[j] != 1 ) or clear_between_frames is True:
ax[j].cla() #Clear axis between frames so that masked arrays can be plotted
if (lineplot[j] == 1):
for k in range(0,Nlines[j]):
lines[j][k].set_data(x[j][k], vars[j][k][index,:])
elif (contour[j] == 1):
thisx = x[j][0]
if len(thisx.shape) == 3:
thisx = thisx[index]
thisy = y[j]
if len(thisy.shape) == 3:
thisy = thisy[index]
plots[j] = ax[j].contourf(x[j][0].T,y[j].T,vars[j][0][index,:,:].T, Ncolors, cmap=cmap, lw=0, levels=clevels[j])
ax[j].set_xlabel(r'x')
ax[j].set_ylabel(r'y')
ax[j].set_title(titles[j])
elif (surf[j] == 1):
thisx = x[j][0]
if len(thisx.shape) == 3:
thisx = thisx[index]
thisy = y[j][0]
if len(thisy.shape) == 3:
thisy = thisy[index]
ax[j] = fig.add_subplot(row,col,j+1, projection='3d')
plots[j] = ax[j].plot_wireframe(thisx, thisy, vars[j][0][index,:,:].T, rstride=ystride[j], cstride=xstride[j])
ax[j].set_zlim(fmin[j],fmax[j])
ax[j].set_xlabel(r'x')
ax[j].set_ylabel(r'y')
ax[j].set_title(titles[j])
elif (polar[j] == 1):
plots[j] = ax[j].contourf(theta[j], r[j], vars[j][0][index,:,:].T,cmap=cmap, levels=clevels[j])
ax[j].set_rmax(Nx[j][0]-1)
ax[j].set_title(titles[j])
if t_array is not None:
title.set_text('t = %1.2e' % t_array[index])
else:
title.set_text('t = %i' % index)
return plots
def init():
global j, pause
j=-2
pause = False
return animate(0)
# Call Animation function
fig.canvas.mpl_connect('button_press_event', onClick)
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=Nframes)
#If movie is not passed as a string assign the default filename
if (movie==1):
movie='animation.mp4'
# Save movie with given or default name
if ((isinstance(movie,str)==1)):
movietype = movie.split('.')[-1]
if movietype == 'mp4':
try:
anim.save(movie,writer = FFwriter, fps=fps, dpi=dpi, extra_args=['-vcodec', 'libx264'])
except Exception:
#Try specifying writer by string if ffmpeg not found
try:
anim.save(movie,writer = 'ffmpeg', fps=fps, dpi=dpi, extra_args=['-vcodec', 'libx264'])
except Exception:
print('Save failed: Check ffmpeg path')
raise
elif movietype == 'gif':
anim.save(movie+'.gif',writer = 'imagemagick', fps=fps, dpi=dpi)
else:
raise ValueError("Unrecognized file type for movie. Supported types are .mp4 and .gif")
# Show animation if not saved or returned, otherwise close the plot
if (movie==0 and return_animation == 0):
plt.show()
else:
plt.close()
# Return animation object
if(return_animation == 1):
return(anim)