Colormap in Matplotlib without using the "scatter" function

Question

I have constructed a scatter plot with x and y positions. Now I have an array with a third variable, density, and I want to assign a color for each point in my scatter plot depending on its density value. I know how to do it using the "scatter" task of matplotlib, for example:

x = [1,2,3,4]
y = [5,3,7,1]
density = [1,2,3,4]

map = plt.scatter(x, y, c=density)
colorbar = plt.colorbar(map)

Now, I would like to do the same using the "plot" function instead, something like:

map = plt.plot(x,y, '.', c=t)

I am trying to do an animation of a galaxy merger, and assign each particle a color depending of the density of that region. So far the code only works with the "plot" task, so I need to implement it that way, but all the examples I've found use the former way.

Thanks in advance!

Answer 1

First off, @tcaswell is right. You're probably wanting to animate a scatter plot. Using lots of plot calls for this will result in much worse performance than changing the collection that scatter returns.

However, here's how you'd go about using multiple plot calls to do this:

import numpy as np
import matplotlib.pyplot as plt

xdata, ydata, zdata = np.random.random((3, 10))

cmap = plt.cm.gist_earth
norm = plt.Normalize(zdata.min(), zdata.max())

fig, ax = plt.subplots()
for x, y, z in zip(xdata, ydata, zdata):
    ax.plot([x], [y], marker='o', ms=20, color=cmap(norm(z)))

sm = plt.cm.ScalarMappable(norm, cmap)
sm.set_array(zdata)
fig.colorbar(sm)

plt.show()

Just for comparison, here's the exact same thing using scatter:

import numpy as np
import matplotlib.pyplot as plt

xdata, ydata, zdata = np.random.random((3, 10))


fig, ax = plt.subplots()
scat = ax.scatter(xdata, ydata, c=zdata, s=200, marker='o')
fig.colorbar(scat)

plt.show()

If you wanted to change the position of the markers in the scatter plot, you'd use scat.set_offsets(xydata), where xydata is an Nx2 array-like sequence.