Commit 1f012d13 authored by Mathias Diez's avatar Mathias Diez
Browse files

New method plot_3d_data_colormap.

parent 0a6b6595
......@@ -517,6 +517,114 @@ def plot_2d_data(data, labels=None, x=None, y=None, constants={}):
ax.set_title(', '.join((i[0] + '=' + str(i[1]) for i in constants)))
return fig
def plot_3d_data_colormap(data, x, y, z, xvals=None, yvals=None, vmin=None,
vmax=None, cmap=None, constants={}, ax=None):
"""Create a colormap figure from array data.
Create a `matplotlib.pyplot.figure` containing a colormap a scalar
function of to variables. The input structure is similar to the output
structure of split_data.
Parameters
----------
data : list of numpy arrays or `SplitData`
Data, containing the data to plot. Length of data must equal length
of labels.
x : string or None
Name of the variable used for x-axis. If x is None, the first
unused variable from the data array is used.
y : string or None
Name of the variable used for y-axis. If y is None, the first
unused variable from the data array is used. (so the second one, if
the first one is used for x)
z : string or None
Name of the variable which is mapped to a color scale. If z is None,
the first unused variable from the data array is used. (so the second
/third one, if the first/second one is used for x and or y).
xvals : list of floats [x0, x1, ...] or None
List of (evenly spaced) x-values for which z is ploted. If xvals is
given, the list of x_i in data['x'] need not be complete. If xvals is
None, unique(data['x']) is used which must be complete and evenly
spaced.
yvals : list of floats [y0, y1, ...] or None
Same as xvals, only for the y-values.
vmin : float or None
Maximum value (upper cutoff) of the range of z which is mapped to
the color scale.
vmax : float or None
Minimum value (lower cutoff) of the range of z which is mapped to
the color scale.
cmap : Colormap or None
The matplotlib.colors.Colormap instance, eg. cm.jet, used for the
color scale of the map. If None the default one is used.
constants : list of tuples (name, value) or `SplitData`
List with values of all the other parameters (used to create the
figure title)
ax : `matplotlib.axes.Axes` instance or `None`
If `ax` is not `None`, no new figure is created, but the plot is done
within the existing Axes `ax`.
Returns
-------
fig : matplotlib figure if `ax` is not set, else None
Figure with the result, only returned if no `axis` is set.
The figure is fully modifiable, and `matplotlib.pyplot.show()` will
show the figure.
"""
if not _mpl_enabled:
raise RuntimeError('matplotlib is not installed.')
if isinstance(data, SplitData):
data = data.arrays[0]
if isinstance(constants, SplitData):
constants = constants.constants
names = list(data[0].dtype.names)
if x is None:
x = names[0]
del names[names.index(x)]
if y is None:
y = names[0]
del names[names.index(y)]
if z is None:
z = names[0]
del names[names.index(z)]
if xvals is None:
xvals = np.sort(np.unique(data[x]))
if yvals is None:
yvals = np.sort(np.unique(data[y]))
Z = np.zeros((len(yvals), len(xvals)))
for an in np.nditer(data):
xn, yn, zn = an[x], an[y], an[z]
i = np.where(xvals==xn)[0][0]
j = np.where(yvals==yn)[0][0]
Z[j][i] = zn
return_fig = False
if ax is None:
return_fig = True
fig = matplotlib.figure.Figure()
ax = fig.add_subplot(111)
extent = [xvals[0], xvals[-1], yvals[0], yvals[-1]]
extend = 'neither'
if vmax is not None:
extend = 'max'
if vmin is not None:
if vmin != 0:
extend = 'both' if extend == 'max' else 'min'
im = ax.imshow(Z, vmin=vmin, vmax=vmax, extent=extent, aspect='auto',
origin='lower', cmap=cmap, interpolation='none')
ax.get_figure().colorbar(im, ax=ax, extend=extend)
ax.set_xlabel(x)
ax.set_ylabel(y)
ax.set_title(', '.join((i[0] + '=' + str(i[1]) for i in constants)))
ax.tick_params(color='w', length=7, width=1.5)
if return_fig:
return fig
def output_fig(fig, output_mode='auto', fname=None, savefile_opts=None):
"""Output a matplotlib figure using a given output mode.
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment