Using matplotlib.animate to animate a contour plot in python - python

I have a 3D array of data (2 spatial dimensions and 1 time dimension) and I'm trying to produce an animated contour plot using matplotlib.animate. I'm using this link as a basis:
http://jakevdp.github.io/blog/2012/08/18/matplotlib-animation-tutorial/
And here's my attempt:
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
from numpy import array, zeros, linspace, meshgrid
from boutdata import collect
# First collect data from files
n = collect("n") # This is a routine to collect data
Nx = n.shape[1]
Nz = n.shape[2]
Ny = n.shape[3]
Nt = n.shape[0]
fig = plt.figure()
ax = plt.axes(xlim=(0, 200), ylim=(0, 100))
cont, = ax.contourf([], [], [], 500)
# initialisation function
def init():
cont.set_data([],[],[])
return cont,
# animation function
def animate(i):
x = linspace(0, 200, Nx)
y = linspace(0, 100, Ny)
x,y = meshgrid(x,y)
z = n[i,:,0,:].T
cont.set_data(x,y,z)
return cont,
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=200, interval=20, blit=True)
plt.show()
But when I do this, I get the following error:
Traceback (most recent call last):
File "showdata.py", line 16, in <module>
cont, = ax.contourf([], [], [], 500)
File "/usr/lib/pymodules/python2.7/matplotlib/axes.py", line 7387, in contourf
return mcontour.QuadContourSet(self, *args, **kwargs)
File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1112, in __init__
ContourSet.__init__(self, ax, *args, **kwargs)
File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 703, in __init__
self._process_args(*args, **kwargs)
File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1125, in _process_args
x, y, z = self._contour_args(args, kwargs)
File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1172, in _contour_args
x,y,z = self._check_xyz(args[:3], kwargs)
File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1204, in _check_xyz
raise TypeError("Input z must be a 2D array.")
TypeError: Input z must be a 2D array.
So I've tried replacing all the [] by [[],[]] but this then produces:
Traceback (most recent call last):
File "showdata.py", line 16, in <module>
cont, = ax.contourf([[],[]], [[],[]], [[],[]],500)
File "/usr/lib/pymodules/python2.7/matplotlib/axes.py", line 7387, in contourf
return mcontour.QuadContourSet(self, *args, **kwargs)
File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1112, in __init__
ContourSet.__init__(self, ax, *args, **kwargs)
File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 703, in __init__
self._process_args(*args, **kwargs)
File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1125, in _process_args
x, y, z = self._contour_args(args, kwargs)
File "/usr/lib/pymodules/python2.7/matplotlib/contour.py", line 1177, in _contour_args
self.zmax = ma.maximum(z)
File "/usr/lib/python2.7/dist-packages/numpy/ma/core.py", line 5806, in __call__
return self.reduce(a)
File "/usr/lib/python2.7/dist-packages/numpy/ma/core.py", line 5824, in reduce
t = self.ufunc.reduce(target, **kargs)
ValueError: zero-size array to maximum.reduce without identity
Thanks in advance!

Felix Schneider is correct about the animation becoming very slow. His solution of setting ax.collections = [] removes all old (and superseded) "artist"s. A more surgical approach is to only remove the artists involved in the drawing the contours:
for c in cont.collections:
c.remove()
which is useful in more complicated cases, in lieu of reconstructing the entire figure for each frame. This also works in Rehman Ali's example; instead of clearing the entire figure with clf() the value returned by contourf() is saved and used in the next iteration. Here is an example code similar to Luke's from Jun 7 '13, demonstrating removing the contours only:
import pylab as plt
import numpy
import matplotlib.animation as animation
#plt.rcParams['animation.ffmpeg_path'] = r"C:\some_path\ffmpeg.exe" # if necessary
# Generate data for plotting
Lx = Ly = 3
Nx = Ny = 11
Nt = 20
x = numpy.linspace(0, Lx, Nx)
y = numpy.linspace(0, Ly, Ny)
x,y = numpy.meshgrid(x,y)
z0 = numpy.exp(-(x-Lx/2)**2-(y-Ly/2)**2) # 2 dimensional Gaussian
def some_data(i): # function returns a 2D data array
return z0 * (i/Nt)
fig = plt.figure()
ax = plt.axes(xlim=(0, Lx), ylim=(0, Ly), xlabel='x', ylabel='y')
cvals = numpy.linspace(0,1,Nt+1) # set contour values
cont = plt.contourf(x, y, some_data(0), cvals) # first image on screen
plt.colorbar()
# animation function
def animate(i):
global cont
z = some_data(i)
for c in cont.collections:
c.remove() # removes only the contours, leaves the rest intact
cont = plt.contourf(x, y, z, cvals)
plt.title('t = %i: %.2f' % (i,z[5,5]))
return cont
anim = animation.FuncAnimation(fig, animate, frames=Nt, repeat=False)
anim.save('animation.mp4', writer=animation.FFMpegWriter())

This is what I got to work:
# Generate grid for plotting
x = linspace(0, Lx, Nx)
y = linspace(0, Ly, Ny)
x,y = meshgrid(x,y)
fig = plt.figure()
ax = plt.axes(xlim=(0, Lx), ylim=(0, Ly))
plt.xlabel(r'x')
plt.ylabel(r'y')
# animation function
def animate(i):
z = var[i,:,0,:].T
cont = plt.contourf(x, y, z, 25)
if (tslice == 0):
plt.title(r't = %1.2e' % t[i] )
else:
plt.title(r't = %i' % i)
return cont
anim = animation.FuncAnimation(fig, animate, frames=Nt)
anim.save('animation.mp4')
I found that removing the blit=0 argument in the FuncAnimation call also helped...

This is the line:
cont, = ax.contourf([], [], [], 500)
change to:
x = linspace(0, 200, Nx)
y = linspace(0, 100, Ny)
x, y = meshgrid(x, y)
z = n[i,:,0,:].T
cont, = ax.contourf(x, y, z, 500)
You need to intilize with sized arrays.

Here is another way of doing the same thing if matplotlib.animation don't work for you. If you want to continuously update the colorbar and everything else in the figure, use plt.ion() at the very beginning to enable interactive plotting and use a combo of plt.draw() and plt.clf() to continuously update the plot.
import matplotlib.pyplot as plt
import numpy as np
plt.ion(); plt.figure(1);
for k in range(10):
plt.clf(); plt.subplot(121);
plt.contourf(np.random.randn(10,10)); plt.colorbar();
plt.subplot(122,polar=True)
plt.contourf(np.random.randn(10,10)); plt.colorbar();
plt.draw();
Note that this works with figures containing different subplots and various types of plots (i.e. polar or cartesian)

I used Lukes approach (from Jun 7 '13 at 8:08 ), but added
ax.collections = []
right before
cont = plt.contourf(x, y, z, 25).
Otherwise I experienced that creating the animation will become very slow for large frame numbers.

I have been looking at this a while ago. I my situation I had a few subplots with contours which I wanted to animate. I did not want to use the plt.clf() solution as Rehman ali suggest as I used some special setup of my axis (with pi symbols etc) which would be cleaned as well, so I preferred the 'remove()' approach suggest be Felix. The thing is that only using 'remove' does not clean up memory and will clog your computer eventually, so you need to explicitly delete of the contours by setting it to an empty list as well.
In order to have a generic remove routine which is able to take away contours as well as text, I wrote the routine 'clean_up_artists' which you should use on every time step on all the axis.
This routine cleans up the artists which are passed in a list called 'artist_list' in a given axis 'axis'. This means that for animating multiple subplots, we need to store the lists of artists for each axis which we need to clean every time step.
Below the full code to animate a number of subplots of random data. It is pretty self-explanatory, so hopefully it becomes clear what happens. Anyhow, I just thought to post it, as it combines several ideas I found on stack overflow which I just to come up with this working example.
Anybody with suggestions to improve the code, please shoot-)
import matplotlib.pyplot as plt
from matplotlib import cm
import matplotlib.animation as animation
import string
import numpy as np
def clean_up_artists(axis, artist_list):
"""
try to remove the artists stored in the artist list belonging to the 'axis'.
:param axis: clean artists belonging to these axis
:param artist_list: list of artist to remove
:return: nothing
"""
for artist in artist_list:
try:
# fist attempt: try to remove collection of contours for instance
while artist.collections:
for col in artist.collections:
artist.collections.remove(col)
try:
axis.collections.remove(col)
except ValueError:
pass
artist.collections = []
axis.collections = []
except AttributeError:
pass
# second attempt, try to remove the text
try:
artist.remove()
except (AttributeError, ValueError):
pass
def update_plot(frame_index, data_list, fig, axis, n_cols, n_rows, number_of_contour_levels, v_min, v_max,
changed_artists):
"""
Update the the contour plots of the time step 'frame_index'
:param frame_index: integer required by animation running from 0 to n_frames -1. For initialisation of the plot,
call 'update_plot' with frame_index = -1
:param data_list: list with the 3D data (time x 2D data) per subplot
:param fig: reference to the figure
:param axis: reference to the list of axis with the axes per subplot
:param n_cols: number of subplot in horizontal direction
:param n_rows: number of subplot in vertical direction
:param number_of_contour_levels: number of contour levels
:param v_min: minimum global data value. If None, take the smallest data value in the 2d data set
:param v_max: maximum global data value. If None, take the largest value in the 2d data set
:param changed_artists: list of lists of artists which need to be updated between the time steps
:return: the changed_artists list
"""
nr_subplot = 0 # keep the index of the current subplot (nr_subplot = 0,1, n_cols x n_rows -1)
# loop over the subplots
for j_col in range(n_cols):
for i_row in range(n_rows):
# set a short reference to the current axis
ax = axis[i_row][j_col]
# for the first setup call, add and empty list which can hold the artists belonging to the current axis
if frame_index < 0:
# initialise the changed artist list
changed_artists.append(list())
else:
# for the next calls of update_plot, remove all artists in the list stored in changed_artists[nr_subplot]
clean_up_artists(ax, changed_artists[nr_subplot])
# get a reference to 2d data of the current time and subplot
data_2d = data_list[nr_subplot][frame_index]
# manually set the levels for better contour range control
if v_min is None:
data_min = np.nanmin(data_2d)
else:
data_min = v_min
if v_max is None:
data_max = np.nanmax(data_2d)
else:
data_max = v_max
# set the contour levels belonging to this subplot
levels = np.linspace(data_min, data_max, number_of_contour_levels + 1, endpoint=True)
# create the contour plot
cs = ax.contourf(data_2d, levels=levels, cmap=cm.rainbow, zorder=0)
cs.cmap.set_under("k")
cs.cmap.set_over("k")
cs.set_clim(v_min, v_max)
# store the contours artists to the list of artists belonging to the current axis
changed_artists[nr_subplot].append(cs)
# set some grid lines on top of the contours
ax.xaxis.grid(True, zorder=0, color="black", linewidth=0.5, linestyle='--')
ax.yaxis.grid(True, zorder=0, color="black", linewidth=0.5, linestyle='--')
# set the x and y label on the bottom row and left column respectively
if i_row == n_rows - 1:
ax.set_xlabel(r"Index i ")
if j_col == 0:
ax.set_ylabel(r"Index j")
# set the changing time counter in the top left subplot
if i_row == 0 and j_col == 1:
# set a label to show the current time
time_text = ax.text(0.6, 1.15, "{}".format("Time index : {:4d}".format(frame_index)),
transform=ax.transAxes, fontdict=dict(color="black", size=14))
# store the artist of this label in the changed artist list
changed_artists[nr_subplot].append(time_text)
# for the initialisation call only, set of a contour bar
if frame_index < 0:
# the first time we add this (make sure to pass -1 for the frame_index
cbar = fig.colorbar(cs, ax=ax)
cbar.ax.set_ylabel("Random number {}".format(nr_subplot))
ax.text(0.0, 1.02, "{}) {}".format(string.ascii_lowercase[nr_subplot],
"Random noise {}/{}".format(i_row, j_col)),
transform=ax.transAxes, fontdict=dict(color="blue", size=12))
nr_subplot += 1
return changed_artists
def main():
n_pixels_x = 50
n_pixels_y = 30
number_of_time_steps = 100
number_of_contour_levels = 10
delay_of_frames = 1000
n_rows = 3 # number of subplot rows
n_cols = 2 # number of subplot columns
min_data_value = 0.0
max_data_value = 1.0
# list containing the random plot per sub plot. Insert you own data here
data_list = list()
for j_col in range(n_cols):
for i_row in range(n_rows):
data_list.append(np.random.random_sample((number_of_time_steps, n_pixels_x, n_pixels_y)))
# set up the figure with the axis
fig, axis = plt.subplots(nrows=n_rows, ncols=n_cols, sharex=True, sharey=True, figsize=(12,8))
fig.subplots_adjust(wspace=0.05, left=0.08, right=0.98)
# a list used to store the reference to the axis of each subplot with a list of artists which belong to this subplot
# this list will be returned and will be updated every time plot which new artists
changed_artists = list()
# create first image by calling update_plot with frame_index = -1
changed_artists = update_plot(-1, data_list, fig, axis, n_cols, n_rows, number_of_contour_levels,
min_data_value, max_data_value, changed_artists)
# call the animation function. The fargs argument equals the parameter list of update_plot, except the
# 'frame_index' parameter.
ani = animation.FuncAnimation(fig, update_plot, frames=number_of_time_steps,
fargs=(data_list, fig, axis, n_cols, n_rows, number_of_contour_levels, min_data_value,
max_data_value, changed_artists),
interval=delay_of_frames, blit=False, repeat=True)
plt.show()
if __name__ == "__main__":
main()

Removing the blit=0 or blit = True argument in the FuncAnimation call also helped
is important!!!

Related

Draw a circle on the plot that follows the mouse [duplicate]

I tried to write a simple script which updates a scatter plot for every timestep t. I wanted to do it as simple as possible. But all it does is to open a window where I can see nothing. The window just freezes. It is maybe just an small error, but I can not find it.
The the data.dat has the format
x y
Timestep 1 1 2
3 1
Timestep 2 6 3
2 1
(the file contains just the numbers)
import numpy as np
import matplotlib.pyplot as plt
import time
# Load particle positioins
with open('//home//user//data.dat', 'r') as fp:
particles = []
for line in fp:
line = line.split()
if line:
line = [float(i) for i in line]
particles.append(line)
T = 100
numbParticles = 2
x, y = np.array([]), np.array([])
plt.ion()
plt.figure()
plt.scatter(x,y)
for t in range(T):
plt.clf()
for k in range(numbP):
x = np.append(x, particles[numbParticles*t+k][0])
y = np.append(y, particles[numbParticles*t+k][1])
plt.scatter(x,y)
plt.draw()
time.sleep(1)
x, y = np.array([]), np.array([])
The simplest, cleanest way to make an animation is to use the matplotlib.animation module.
Since a scatter plot returns a matplotlib.collections.PathCollection, the way to update it is to call its set_offsets method. You can pass it an array of shape (N, 2) or a list of N 2-tuples -- each 2-tuple being an (x,y) coordinate.
For example,
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
T = 100
numbParticles = 2
particles = np.random.random((T,numbParticles)).tolist()
x, y = np.array([]), np.array([])
def init():
pathcol.set_offsets([[], []])
return [pathcol]
def update(i, pathcol, particles):
pathcol.set_offsets(particles[i])
return [pathcol]
fig = plt.figure()
xs, ys = zip(*particles)
xmin, xmax = min(xs), max(xs)
ymin, ymax = min(ys), max(ys)
ax = plt.axes(xlim=(xmin, xmax), ylim=(ymin, ymax))
pathcol = plt.scatter([], [], s=100)
anim = animation.FuncAnimation(
fig, update, init_func=init, fargs=(pathcol, particles), interval=1000, frames=T,
blit=True, repeat=True)
plt.show()
I finally found a solution. You can do it simply by using this script. I tried to keep it simple:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
# Helps me to get the data from the file I want to plot
N = 0
# Load particle positioins
with open('//home//user//data.dat', 'r') as fp:
particles = []
for line in fp:
line = line.split()
particles.append(line)
# Create new Figure and an Axes which fills it.
fig = plt.figure(figsize=(7, 7))
ax = fig.add_axes([0, 0, 1, 1], frameon=True)
border = 100
ax.set_xlim(-border, border), ax.set_xticks([])
ax.set_ylim(-border, border), ax.set_yticks([])
# particle data
p = 18 # number of particles
myPa = np.zeros(p, dtype=[('position', float, 2)])
# Construct the scatter which we will update during animation
scat = ax.scatter(myPa['position'][:, 0], myPa['position'][:, 1])
def update(frame_number):
# New positions
myPa['position'][:] = particles[N*p:N*p+p]
# Update the scatter collection, with the new colors, sizes and positions.
scat.set_offsets(myPa['position'])
increment()
def increment():
global N
N = N+1
# Construct the animation, using the update function as the animation director.
animation = FuncAnimation(fig, update, interval=20)
plt.show()

Make a point move on the plot without clearing earlier plots in matplotlib [duplicate]

I tried to write a simple script which updates a scatter plot for every timestep t. I wanted to do it as simple as possible. But all it does is to open a window where I can see nothing. The window just freezes. It is maybe just an small error, but I can not find it.
The the data.dat has the format
x y
Timestep 1 1 2
3 1
Timestep 2 6 3
2 1
(the file contains just the numbers)
import numpy as np
import matplotlib.pyplot as plt
import time
# Load particle positioins
with open('//home//user//data.dat', 'r') as fp:
particles = []
for line in fp:
line = line.split()
if line:
line = [float(i) for i in line]
particles.append(line)
T = 100
numbParticles = 2
x, y = np.array([]), np.array([])
plt.ion()
plt.figure()
plt.scatter(x,y)
for t in range(T):
plt.clf()
for k in range(numbP):
x = np.append(x, particles[numbParticles*t+k][0])
y = np.append(y, particles[numbParticles*t+k][1])
plt.scatter(x,y)
plt.draw()
time.sleep(1)
x, y = np.array([]), np.array([])
The simplest, cleanest way to make an animation is to use the matplotlib.animation module.
Since a scatter plot returns a matplotlib.collections.PathCollection, the way to update it is to call its set_offsets method. You can pass it an array of shape (N, 2) or a list of N 2-tuples -- each 2-tuple being an (x,y) coordinate.
For example,
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
T = 100
numbParticles = 2
particles = np.random.random((T,numbParticles)).tolist()
x, y = np.array([]), np.array([])
def init():
pathcol.set_offsets([[], []])
return [pathcol]
def update(i, pathcol, particles):
pathcol.set_offsets(particles[i])
return [pathcol]
fig = plt.figure()
xs, ys = zip(*particles)
xmin, xmax = min(xs), max(xs)
ymin, ymax = min(ys), max(ys)
ax = plt.axes(xlim=(xmin, xmax), ylim=(ymin, ymax))
pathcol = plt.scatter([], [], s=100)
anim = animation.FuncAnimation(
fig, update, init_func=init, fargs=(pathcol, particles), interval=1000, frames=T,
blit=True, repeat=True)
plt.show()
I finally found a solution. You can do it simply by using this script. I tried to keep it simple:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
# Helps me to get the data from the file I want to plot
N = 0
# Load particle positioins
with open('//home//user//data.dat', 'r') as fp:
particles = []
for line in fp:
line = line.split()
particles.append(line)
# Create new Figure and an Axes which fills it.
fig = plt.figure(figsize=(7, 7))
ax = fig.add_axes([0, 0, 1, 1], frameon=True)
border = 100
ax.set_xlim(-border, border), ax.set_xticks([])
ax.set_ylim(-border, border), ax.set_yticks([])
# particle data
p = 18 # number of particles
myPa = np.zeros(p, dtype=[('position', float, 2)])
# Construct the scatter which we will update during animation
scat = ax.scatter(myPa['position'][:, 0], myPa['position'][:, 1])
def update(frame_number):
# New positions
myPa['position'][:] = particles[N*p:N*p+p]
# Update the scatter collection, with the new colors, sizes and positions.
scat.set_offsets(myPa['position'])
increment()
def increment():
global N
N = N+1
# Construct the animation, using the update function as the animation director.
animation = FuncAnimation(fig, update, interval=20)
plt.show()

animating a stem plot in matplotlib

I'm trying to animate a stem plot in matplotlib and I can't find the necessary documentation to help me. I have a series of data files which each look like this:
1 0.345346
2 0.124325
3 0.534585
and I want plot each file as a separate frame.
According to this and this other tutorial, I should create a function which updates the data contained in each plot object (artist? I'm not sure about the terminology)
From the second link, this is the update function
def update(frame):
global P, C, S
# Every ring is made more transparent
C[:,3] = np.maximum(0, C[:,3] - 1.0/n)
# Each ring is made larger
S += (size_max - size_min) / n
# Reset ring specific ring (relative to frame number)
i = frame % 50
P[i] = np.random.uniform(0,1,2)
S[i] = size_min
C[i,3] = 1
# Update scatter object
scat.set_edgecolors(C)
scat.set_sizes(S)
scat.set_offsets(P)
# Return the modified object
return scat,
How can I adapt this kind of update function for a stem plot? The documentation for stem is horribly brief (in fact this is a recurring issue as I'm learning matplotlib), but the example code shows that the output of stem is a tuple markerline, stemlines, baseline rather than an artist object like for plt.plot or plt.imshow.
So when I write my update function for the animation, how can I update the data inside the stem plot?
Here you go!
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np
fig, ax = plt.subplots()
x = np.linspace(0.1, 2*np.pi, 10)
markerline, stemlines, baseline = ax.stem(x, np.cos(x), '-.')
def update(i):
ax.cla()
markerline, stemlines, baseline = ax.stem(x, np.cos(x+i/10), '-.')
ax.set_ylim((-1, 1))
anim = FuncAnimation(fig, update, frames=range(10, 110, 10), interval=500)
anim.save('so.gif', dpi=80, writer='imagemagick')
I think there can be better ways of achieving this- not requiring to clear the plot each time. However, this works!
When using the keyword use_line_collection=True (default behavior since Matplotlib 3.3) one can update the three elements
markerline
stemlines
baseline
individualy. Here is the code for the sine wave example:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
fig, ax = plt.subplots()
x = np.linspace(0.1, 2*np.pi, 10)
y = np.cos(x)
bottom = 0
h_stem = ax.stem(x, y, bottom=bottom, use_line_collection=True, linefmt='-.')
def update(i):
y = np.cos(x+i/10)
# markerline
h_stem[0].set_ydata(y)
h_stem[0].set_xdata(x) # not necessary for constant x
# stemlines
h_stem[1].set_paths([np.array([[xx, bottom],
[xx, yy]]) for (xx, yy) in zip(x, y)])
# baseline
h_stem[2].set_xdata([np.min(x), np.max(x)])
h_stem[2].set_ydata([bottom, bottom]) # not necessary for constant bottom
anim = FuncAnimation(fig, update, frames=range(10, 110, 10), interval=1)
anim.save('so.gif', dpi=80, writer='imagemagick')
Depending on what values (x, y, bottom) should be updated you can omit some parts of this update or reuse the current values. I wrote a more general function, where you can pass an arbitrary combination of these values:
def update_stem(h_stem, x=None, y=None, bottom=None):
if x is None:
x = h_stem[0].get_xdata()
else:
h_stem[0].set_xdata(x)
h_stem[2].set_xdata([np.min(x), np.max(x)])
if y is None:
y = h_stem[0].get_ydata()
else:
h_stem[0].set_ydata(y)
if bottom is None:
bottom = h_stem[2].get_ydata()[0]
else:
h_stem[2].set_ydata([bottom, bottom])
h_stem[1].set_paths([np.array([[xx, bottom],
[xx, yy]]) for (xx, yy) in zip(x, y)])

python matplotlib with a line color gradient and colorbar

I've been toying around with this problem and am close to what I want but missing that extra line or two.
Basically, I'd like to plot a single line whose color changes given the value of a third array. Lurking around I have found this works well (albeit pretty slowly) and represents the problem
import numpy as np
import matplotlib.pyplot as plt
c = np.arange(1,100)
x = np.arange(1,100)
y = np.arange(1,100)
cm = plt.get_cmap('hsv')
fig = plt.figure(figsize=(5,5))
ax1 = plt.subplot(111)
no_points = len(c)
ax1.set_color_cycle([cm(1.*i/(no_points-1))
for i in range(no_points-1)])
for i in range(no_points-1):
bar = ax1.plot(x[i:i+2],y[i:i+2])
plt.show()
Which gives me this:
I'd like to be able to include a colorbar along with this plot. So far I haven't been able to crack it just yet. Potentially there will be other lines included with different x,y's but the same c, so I was thinking that a Normalize object would be the right path.
Bigger picture is that this plot is part of a 2x2 sub plot grid. I am already making space for the color bar axes object with matplotlib.colorbar.make_axes(ax4), where ax4 with the 4th subplot.
Take a look at the multicolored_line example in the Matplotlib gallery and dpsanders' colorline notebook:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.collections as mcoll
def multicolored_lines():
"""
http://nbviewer.ipython.org/github/dpsanders/matplotlib-examples/blob/master/colorline.ipynb
http://matplotlib.org/examples/pylab_examples/multicolored_line.html
"""
x = np.linspace(0, 4. * np.pi, 100)
y = np.sin(x)
fig, ax = plt.subplots()
lc = colorline(x, y, cmap='hsv')
plt.colorbar(lc)
plt.xlim(x.min(), x.max())
plt.ylim(-1.0, 1.0)
plt.show()
def colorline(
x, y, z=None, cmap='copper', norm=plt.Normalize(0.0, 1.0),
linewidth=3, alpha=1.0):
"""
http://nbviewer.ipython.org/github/dpsanders/matplotlib-examples/blob/master/colorline.ipynb
http://matplotlib.org/examples/pylab_examples/multicolored_line.html
Plot a colored line with coordinates x and y
Optionally specify colors in the array z
Optionally specify a colormap, a norm function and a line width
"""
# Default colors equally spaced on [0,1]:
if z is None:
z = np.linspace(0.0, 1.0, len(x))
# Special case if a single number:
# to check for numerical input -- this is a hack
if not hasattr(z, "__iter__"):
z = np.array([z])
z = np.asarray(z)
segments = make_segments(x, y)
lc = mcoll.LineCollection(segments, array=z, cmap=cmap, norm=norm,
linewidth=linewidth, alpha=alpha)
ax = plt.gca()
ax.add_collection(lc)
return lc
def make_segments(x, y):
"""
Create list of line segments from x and y coordinates, in the correct format
for LineCollection: an array of the form numlines x (points per line) x 2 (x
and y) array
"""
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
return segments
multicolored_lines()
Note that calling plt.plot hundreds of times tends to kill performance.
Using a LineCollection to build multi-colored line segments is much much faster.

Why is my line clipping in matplotlib?

I am trying to draw a series of lines. The lines are all the same length, and randomly switch colors for a random length (blue to orange). I am drawing the lines in blue and then overlaying orange on top. You can see from my picture there are clipped parts of the lines where it is grey. I cannot figure out why this is happening. Also related I believe is that my labels are not moving to a left alignment like they should. Any help is greatly appreciated.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import random
plt.close('all')
fig, ax = plt.subplots(figsize=(15,11))
def label(xy, text):
y = xy[1] - 2
ax.text(xy[0], y, text, ha="left", family='sans-serif', size=14)
def draw_chromosome(start, stop, y, color):
x = np.array([start, stop])
y = np.array([y, y])
line = mlines.Line2D(x , y, lw=10., color=color)
ax.add_line(line)
x = 50
y = 100
chr = 1
for i in range(22):
draw_chromosome(x, 120, y, "#1C2F4D")
j = 0
while j < 120:
print j
length = 1
if random.randint(1, 100) > 90:
length = random.randint(1, 120-j)
draw_chromosome(j, j+length, y, "#FA9B00")
j = j+length+1
label([x, y], "Chromosome%i" % chr)
y -= 3
chr += 1
plt.axis('equal')
plt.axis('off')
plt.tight_layout()
plt.show()
You're only drawing the blue background from x = 50 to x = 120.
Replace this line:
draw_chromosome(x, 120, y, "#1C2F4D")
with this:
draw_chromosome(0, 120, y, "#1C2F4D")
To draw the blue line all the way across.
Alternately, if you also want to move your labels to the left, you can just set x=0 instead of setting it to 50.
I suggest using LineCollection for this. Below is a little helper function I wrote based on the example at http://matplotlib.org/examples/pylab_examples/multicolored_line.html (it looks long, but there is a lot of comments + docstrings)
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import ListedColormap, BoundaryNorm
from matplotlib.ticker import NullLocator
from collections import OrderedDict
def binary_state_lines(ax, chrom_data, xmin=0, xmax=120,
delta_y=3,
off_color = "#1C2F4D",
on_color = "#FA9B00"):
"""
Draw a whole bunch of chromosomes
Parameters
----------
ax : Axes
The axes to draw stuff to
chrom_data : OrderedDict
The chromosome data as a dict, key on the label with a list of pairs
of where the data is 'on'. Data is plotted top-down
xmin, xmax : float, optional
The minimum and maximum limits for the x values
delta_y : float, optional
The spacing between lines
off_color, on_color : color, optional
The colors to use for the the on/off state
Returns
-------
collections : dict
dictionary of the collections added keyed on the label
"""
# base offset
y_val = 0
# make the color map and norm
cmap = ListedColormap([off_color, on_color])
norm = BoundaryNorm([0, 0.5, 1], cmap.N)
# sort out where the text should be
txt_x = (xmax + xmin) / 2
# dictionary to hold the returned artists
ret = dict()
# loop over the input data draw each collection
for label, data in chrom_data.items():
# increment the y offset
y_val += delta_y
# turn the high windows on to alternating
# high/low regions
x = np.asarray(data).ravel()
# assign the high/low state to each one
state = np.mod(1 + np.arange(len(x)), 2)
# deal with boundary conditions to be off
# at start/end
if x[0] > xmin:
x = np.r_[xmin, x]
state = np.r_[0, state]
if x[-1] < xmax:
x = np.r_[x, xmax]
state = np.r_[state, 0]
# make the matching y values
y = np.ones(len(x)) * y_val
# call helper function to create the collection
coll = draw_segments(ax, x, y, state,
cmap, norm)
ret[label] = coll
# set up the axes limits
ax.set_xlim(xmin, xmax)
ax.set_ylim(0, y_val + delta_y)
# turn off x-ticks
ax.xaxis.set_major_locator(NullLocator())
# make the y-ticks be labeled as per the input
ax.yaxis.set_ticks((1 + np.arange(len(chrom_data))) * delta_y)
ax.yaxis.set_ticklabels(list(chrom_data.keys()))
# invert so that the first data is at the top
ax.invert_yaxis()
# turn off the frame and patch
ax.set_frame_on(False)
# return the added artists
return ret
def draw_segments(ax, x, y, state, cmap, norm, lw=10):
"""
helper function to turn boundary edges into the input LineCollection
expects.
Parameters
----------
ax : Axes
The axes to draw to
x, y, state : array
The x edges, the y values and the state of each region
cmap : matplotlib.colors.Colormap
The color map to use
norm : matplotlib.ticker.Norm
The norm to use with the color map
lw : float, optional
The width of the lines
"""
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(segments, cmap=cmap, norm=norm)
lc.set_array(state)
lc.set_linewidth(lw)
ax.add_collection(lc)
return lc
An example:
synthetic_data = OrderedDict()
for j in range(21):
key = 'data {:02d}'.format(j)
synthetic_data[key] = np.cumsum(np.random.randint(1, 10, 20)).reshape(-1, 2)
fig, ax = plt.subplots(tight_layout=True)
binary_state_lines(ax, synthetic_data, xmax=120)
plt.show()
Separating the plotting logic from everything else will make your code easier to maintain and more reusable.
I also took the liberty of moving your labels from between the lines (where they can be ambiguous) to the yaxis tick labels.

Categories

Resources