[python] Scatter plot and Color mapping in Python

I have a range of points x and y stored in numpy arrays. Those represent x(t) and y(t) where t=0...T-1

I am plotting a scatter plot using

import matplotlib.pyplot as plt

plt.scatter(x,y)
plt.show()

I would like to have a colormap representing the time (therefore coloring the points depending on the index in the numpy arrays)

What is the easiest way to do so?

This question is related to python matplotlib

The answer is


To add to wflynny's answer above, you can find the available colormaps here

Example:

import matplotlib.cm as cm
plt.scatter(x, y, c=t, cmap=cm.jet)

or alternatively,

plt.scatter(x, y, c=t, cmap='jet')

Subplot Colorbar

For subplots with scatter, you can trick a colorbar onto your axes by building the "mappable" with the help of a secondary figure and then adding it to your original plot.

As a continuation of the above example:

import numpy as np
import matplotlib.pyplot as plt

x = np.arange(10)
y = x
t = x
fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.scatter(x, y, c=t, cmap='viridis')
ax2.scatter(x, y, c=t, cmap='viridis_r')


# Build your secondary mirror axes:
fig2, (ax3, ax4) = plt.subplots(1, 2)

# Build maps that parallel the color-coded data
# NOTE 1: imshow requires a 2-D array as input
# NOTE 2: You must use the same cmap tag as above for it match
map1 = ax3.imshow(np.stack([t, t]),cmap='viridis')
map2 = ax4.imshow(np.stack([t, t]),cmap='viridis_r')

# Add your maps onto your original figure/axes
fig.colorbar(map1, ax=ax1)
fig.colorbar(map2, ax=ax2)
plt.show()

Scatter subplots with COLORBAR

Note that you will also output a secondary figure that you can ignore.