How to modify the colormap of a scatter plot after it was drawn

Question

You can assign a colormap to a scatter plot by using scatter(X, Y, c=..., cmap="rainbow") for example, but is there a way to modify the colormap later?

I know we can change the transparency and color via set_color() and set_alpha() but I did not find the equivalent for colormaps. A workaround would be to save the data of the scatter, erase it, then scatter() again with new parameters but that is going to be expensive. That is an issue since the goal of this question is to be able to change the c parameter in real time so that old points "cool down" to darker colors.

Thanks for your time.

---

Edit to add more information: I would like to be able to switch back and forth between a scatter made with a single color and a scatter with a variable colormap and normalization. Between scatter(X, Y, color="#FF0000") and scatter(X, Y, c=age, cmap="rainbow") for example, with a variable age = [i for i,x in enumerate(X)].

Answer 1

Instead of changing the colormap, you can actually change the values. If you set an under color to the colormap, you can change the values to anything below the norm and have it appear in that color.

The following would set the values to one less than the minimal value when n is pressed. They are set back to their original value upon pressing m.

import matplotlib.pyplot as plt
import numpy as np; np.random.seed(42)

a = np.random.randint(10,30, size=20)
x = np.random.rand(20)
y = np.random.rand(20)

cmap = plt.get_cmap("rainbow")
cmap.set_under("black")

fig, ax = plt.subplots()
ax.set_title("Press n and m keys to change colors")
sc = ax.scatter(x,y, c=a, cmap=cmap)

def change(evt):
    if evt.key in list("nm"):
        if evt.key == "n":
            sc.set_array(np.ones_like(a)*a.min()-1)
        if evt.key == "m":
            sc.set_array(a)
        fig.canvas.draw_idle()


fig.canvas.mpl_connect("key_press_event", change)

plt.show()

---

If you really need to start with a scatter that doesn't have its array defined, you can do that as shown in the below.

import matplotlib.pyplot as plt
import numpy as np; np.random.seed(42)

a = np.random.randint(10,30, size=20)
x = np.random.rand(20)
y = np.random.rand(20)

cmap = plt.get_cmap("rainbow")
cmap.set_under("black")
norm = plt.Normalize(a.min(), a.max())

fig, ax = plt.subplots()
ax.set_title("Press n and m keys to change colors")
sc = ax.scatter(x,y, color="black")

def change(evt):
    if evt.key in list("nm"):
        sc.set_cmap(cmap)
        sc.set_norm(norm)
        if evt.key == "n":
            sc.set_array(np.ones_like(a)*a.min()-1)
        if evt.key == "m":
            sc.set_array(a)
        fig.canvas.draw_idle()
        

fig.canvas.mpl_connect("key_press_event", change)

plt.show()