Commit 4e149f7e authored by syz's avatar syz
Browse files

New function adds colorbar to line plots. Line family - Better legend names + optional color bar

parent b1970bd3
......@@ -39,6 +39,7 @@ if sys.version_info.major == 3:
default_cmap = plt.cm.viridis
def set_tick_font_size(axes, font_size):
"""
Sets the font size of the ticks in the provided axes
......@@ -100,6 +101,32 @@ def make_scalar_mappable(vmin, vmax, cmap=None):
return sm
def cbar_for_line_plot(axis, num_steps, discrete_ticks=True, **kwargs):
"""
Adds a colorbar next to a line plot axis
Parameters
----------
axis : axis handle
Axis with multiple line objects
num_steps : uint
Number of steps in the colorbar
discrete_ticks : (optional) bool
Whether or not to have the ticks match the number of number of steps. Default = True
"""
cmap = get_cmap_object(kwargs.pop('cmap', None))
cmap = discrete_cmap(num_steps, base_cmap=cmap.name)
sm = make_scalar_mappable(0, num_steps - 1, cmap=cmap, **kwargs)
if discrete_ticks:
kwargs.update({'ticks': np.arange(num_steps)})
cbar = plt.colorbar(sm, ax=axis, orientation='vertical',
pad=0.04, use_gridspec=True, **kwargs)
return cbar
def get_cmap_object(cmap):
"""
Get the matplotlib.colors.LinearSegmentedColormap object regardless of the input
......@@ -299,8 +326,8 @@ def rainbow_plot(axis, x_vec, y_vec, num_steps=32, **kwargs):
color=cmap(255 * num_steps / num_steps), **kwargs)
def plot_line_family(axis, x_axis, line_family, line_names=None, label_prefix='Line', label_suffix='',
cmap=default_cmap, y_offset=0, **kwargs):
def plot_line_family(axis, x_vec, line_family, line_names=None, label_prefix='', label_suffix='',
y_offset=0, show_cbar=False, **kwargs):
"""
Plots a family of lines with a sequence of colors
......@@ -308,7 +335,7 @@ def plot_line_family(axis, x_axis, line_family, line_names=None, label_prefix='L
----------
axis : axis handle
Axis to plot the curve
x_axis : array-like
x_vec : array-like
Values to plot against
line_family : 2D numpy array
family of curves arranged as [curve_index, features]
......@@ -318,27 +345,39 @@ def plot_line_family(axis, x_axis, line_family, line_names=None, label_prefix='L
prefix for the legend (before the index of the curve)
label_suffix : string / unicode
suffix for the legend (after the index of the curve)
cmap : matplotlib.colors.LinearSegmentedColormap object
Colormap to be used
y_offset : (optional) number
quantity by which the lines are offset from each other vertically (useful for spectra)
show_cbar : (optional) bool
Whether or not to show a colorbar (instead of a legend)
"""
cmap = get_cmap_object(cmap)
cmap = get_cmap_object(kwargs.pop('cmap', None))
num_lines = len(line_family)
num_lines = line_family.shape[0]
default_names = False
if line_names is None:
line_names = ['{} {} {}'.format(label_prefix, line_ind, label_suffix) for line_ind in range(num_lines)]
else:
if len(line_names) != num_lines:
warn('Line names of different length compared to provided dataset')
line_names = ['{} {} {}'.format(label_prefix, line_ind, label_suffix) for line_ind in range(num_lines)]
label_prefix = 'Line '
default_names = True
elif len(line_names) != num_lines:
warn('Line names of different length compared to provided dataset')
default_names = True
if default_names:
line_names = [str(line_ind) for line_ind in range(num_lines)]
line_names = ['{} {} {}'.format(label_prefix, cur_name, label_suffix) for cur_name in line_names]
for line_ind in range(num_lines):
axis.plot(x_axis, line_family[line_ind] + line_ind * y_offset,
axis.plot(x_vec, line_family[line_ind] + line_ind * y_offset,
label=line_names[line_ind],
color=cmap(int(255 * line_ind / (num_lines - 1))), **kwargs)
if show_cbar:
# put back the cmap parameter:
kwargs.update({'cmap': cmap})
cbar_for_line_plot(axis, num_lines, **kwargs)
def plot_map(axis, data, stdevs=None, origin='lower', **kwargs):
"""
......
Markdown is supported
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