Plot a sequence of colored lines
Plot a sequence of colored lines where each line is colored according to some other property. A colorbar connects the color with the numerical values. There are a few tricks to getting this to behave as you would like, and also to look nice.
The parameters of interest are \(\mu\) and \(\sigma\).
Turn on "%matplotlib inline" for plotting inside a notebook.
In [5]:
import numpy as np
#%matplotlib inline
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.cm as cm
# Create data
# Let's plot this function
def normal_func(xs, mu, sigma):
"""normal distribution"""
return (1.0/(2 * sigma**2 * np.pi)) * np.exp(-(xs - mu)**2 / (2 *sigma**2))
# for a variety of parameters
mus = [1,2,3,4,5]
#mus = np.linspace(1.0, 6.0, 5) # for a noninteger list, may choose continuous colorbar below
sigmas = np.linspace(0.6, 2.0, 7)
# over this range of x values
xs = np.linspace(0, 10, 200)
# Compute the function for each parameter pair.
ys = np.zeros((len(mus), len(sigmas), len(xs)))
for mind, mu in enumerate(mus):
for sind, sigma in enumerate(sigmas):
ys[mind, sind, :] = normal_func(xs, mu, sigma)
Here we choose to color by \(\mu\). Incidentally, we control the linewidth with \(\sigma\) for fun.
In [6]:
# Which parameter should we represent with color?
colorparams = mus
# Choose a colormap
colormap = cm.viridis
#colormap = cm.jet
normalize = mcolors.Normalize(vmin=np.min(colorparams), vmax=np.max(colorparams))
fig = plt.figure(figsize=(12, 8))
ax = fig.gca()
for mind, mu in enumerate(mus):
for sind, sigma in enumerate(sigmas):
color = colormap(normalize(mu))
ax.plot(xs, ys[mind, sind, :], color=color, lw=3*sigma)
# Colorbar setup
s_map = cm.ScalarMappable(norm=normalize, cmap=colormap)
s_map.set_array(colorparams)
# If color parameters is a linspace, we can set boundaries in this way
halfdist = (colorparams[1] - colorparams[0])/2.0
boundaries = np.linspace(colorparams[0] - halfdist, colorparams[-1] + halfdist, len(colorparams) + 1)
# Use this to emphasize the discrete color values
cbar = fig.colorbar(s_map, spacing='proportional', ticks=colorparams, boundaries=boundaries, format='%2.2g') # format='%2i' for integer
# Use this to show a continuous colorbar
#cbar = fig.colorbar(s_map, spacing='proportional', ticks=colorparams, format='%2i')
cbarlabel = r'$\mu$'
cbar.set_label(cbarlabel, fontsize=20)
ax.set_xlabel(r'$\gamma$-wave intensity', fontsize=20)
ax.set_ylabel('Monkey derangement', fontsize=20)
ylims = ax.get_ylim()
ax.set_ylim(top=1.1*ylims[1])
filename = 'colored_lines.pdf'
#fig.savefig(filename, bbox_inches='tight')
#plt.show()