Matplotlib: Plotting numerous disconnected line segments with different colors

Rabih Kodeih picture Rabih Kodeih · Jan 25, 2014 · Viewed 61k times · Source

I have a set of data records like this:

(s1, t1), (u1, v1), color1
(s2, t2), (u2, v2), color2
.
.
.
(sN, tN), (uN, vN), colorN

In any record, the first two values are the end-points of a line segment, the third value is the color of that line segment. More specifically, (sn, tn) are the x-y coordinates of the first end-point, (un, vn) are the x-y coordinates of the second-endpoint. Also, color is an rgb with alpha value.

In general, any two line segments are disconnected (meaning that their end-points do not necessarily coincide).

How to plot this data using matplotlib with a single plot call (or as few as possible) as there could be potentially thousands of records.

Attempts

Preparing the data in one big list and calling plot against it is way too slow. For example the following code couldn't finish in a reasonable amount of time:

import numpy as np
import matplotlib.pyplot as plt

data = []
for _ in xrange(60000):
    data.append((np.random.rand(), np.random.rand()))
    data.append((np.random.rand(), np.random.rand()))
    data.append('r')

print 'now plotting...' # from now on, takes too long
plt.plot(*data)
print 'done'
#plt.show()

I was able to speed-up the plot rendering by using the None insertion trick as follows:

import numpy as np
import matplotlib.pyplot as plt
from timeit import timeit

N = 60000
_s = np.random.rand(N)
_t = np.random.rand(N)
_u = np.random.rand(N)
_v = np.random.rand(N)
x = []
y = []
for s, t, u, v in zip(_s, _t, _u, _v):
    x.append(s)
    x.append(u)
    x.append(None)
    y.append(t)
    y.append(v)
    y.append(None)
print timeit(lambda:plt.plot(x, y), number=1)

This executes in under a second on my machine. I still have to figure out how to embed the color values (RGB with alpha channel).

Answer

HYRY picture HYRY · Jan 26, 2014

use LineCollection:

import numpy as np
import pylab as pl
from matplotlib import collections  as mc

lines = [[(0, 1), (1, 1)], [(2, 3), (3, 3)], [(1, 2), (1, 3)]]
c = np.array([(1, 0, 0, 1), (0, 1, 0, 1), (0, 0, 1, 1)])

lc = mc.LineCollection(lines, colors=c, linewidths=2)
fig, ax = pl.subplots()
ax.add_collection(lc)
ax.autoscale()
ax.margins(0.1)

here is the output:

enter image description here