Issue with Matplotlib scatterplot and Color maps

Question

I am working on a project that involves applying colormaps to scatterplots generated in matplotlib. My code works as expected, unless the scatterplot being generated has exactly four points. This is illustrated in the following code:

import numpy as np
import matplotlib.pyplot as plt

cmap = plt.get_cmap('rainbow_r')

z = np.arange(20)
plt.close()
plt.figure(figsize=[8,6])

for i in range(1,11):
    x = np.arange(i)
    y = np.zeros(i) + i
    plt.scatter(x, y, c=cmap(i / 10), edgecolor='k', label=i, s=200)

plt.legend()
plt.show()

This code generates the following plot:

Each row should consist of points of the same color, but that is not the case for the row with four points.

I assume it has to do with the fact that the color selected from the colormap is returned as a tuple of 4 floats, as illustrated below:

print(cmap(0.4))
>>  (0.69999999999999996, 0.95105651629515364, 0.58778525229247314, 1.0)

Assuming that this is the source of the issue, I have no idea how to fix it.

Answer 1

I played around a little with your code and basically you were on to the problem; the similar dimensions of the cmap values with the data caused them to be interpreted differently for the "4" case.

This seems to work correctly:

plt.scatter(x, y, c=[cmap(i / 10)], edgecolor='k', label=i, s=200)

(List constructor around return value from cmap)