How to pick colors and ticks on a color bar?

Question

I am trying to make two scatter plots using matplotlib with the same color bar and scheme. Here is my code for creating the data,

import numpy as np
import matplotlib.pyplot as plt
import random

x = np.arange(0,10)
y = np.arange(0,10)
actual = np.zeros(10, )
predicted = np.zeros(10, )

for i in range(10):
    actual[i] = random.choice([1, 2, 3, 4])
    predicted[i] = random.choice([0, 1, 3])

I then create the first scatter plot,

plt.scatter(x, y, c=actual)
cbar = plt.colorbar(ticks=[1, 2, 3, 4])
cbar.ax.set_yticklabels(["b", "c", "d", "e"])

And the second scatter plot,

plt.scatter(x, y, c=predicted)
cbar = plt.colorbar(ticks=[0, 1, 2, 3])
cbar.ax.set_yticklabels(["a", "b", "c", "d"])

I want to create both the scatter plots with the same color bar, i.e. ticks=[0,1,2,3,4], labels as ["a", "b", "c", "d", "e"], and same colors corresponding to each tick/label. However, I am unable to do so. For example, doing this for the second plot leaves the colorbar unchanged as the c variable only contains 0, 1, and 3, but not 2 and 4. Any ideas?

Answer 1

Use vmin and vmax when you do the scatter, and set them to the same for both your plots. That way, you fix the upper and lower limits of the colorbars, regardless of the data you are plotting.

For example:

import numpy as np
import matplotlib.pyplot as plt
import random

x = np.arange(0,10)
y = np.arange(0,10)
actual = np.zeros(10, )
predicted = np.zeros(10, )

for i in range(10):
    actual[i] = random.choice([1, 2, 3, 4])
    predicted[i] = random.choice([0, 1, 3])

cbticks = range(5)
cblabels = ["a", "b", "c", "d", "e"]

plt.subplot(211)
plt.scatter(x, y, c=actual, vmin=0, vmax=4)
cbar = plt.colorbar(ticks=cbticks)
cbar.ax.set_yticklabels(cblabels)

plt.subplot(212)
plt.scatter(x, y, c=predicted, vmin=0, vmax=4)
cbar = plt.colorbar(ticks=cbticks)
cbar.ax.set_yticklabels(cblabels)

plt.show()