How to retrieve a colorbar from an already existing matplotlib scatter figure

Question

I am trying to create a figure for a paper compounding previous matplotlib figures that I have saved in pickle format.

For doing so I have used the get_offsets() command, which allows me to retrieve the points of my scatter plots and get_facecolor() to get the color of the points.

Code:

### Loading my figure (visible in the enclosed picture)
with open('Figures paper/Individual spectra figures/Spectra_1000_GALL_chunk_0.pickle', 'rb') as fl:
    ax1 = pickle.load(fl)

### Getting the plot features
xd = axp.collections[0].get_offsets()[:,0]
yd = axp.collections[0].get_offsets()[:,1]
cd = axp.collections[0].get_facecolor()

ax.scatter(xd,yd, s=ar,c=cd ,alpha=0.8)

And then I can plot the points in the sub-panel that I want of my new compound figure. However, the colorbar of my initial plot does not appear, and I don't know how to retrieve it to include it in the new plot.

I have tried creating a new colorbar with the data extracted from the scatterplot, but it's not possible because .collections[0].get_facecolor() gets only the value of the colors, but not the original variable that originated them.

Any idea about how to do it?

Answer 1

Yes and maybe. Here's code demonstrating retrieval and reuse of a pickled colorbar -- if you pickled the whole figure originally this should apply.

BUT I suspect you would like not just the colorbar but the colormap from data to color. And I don't think that's explicit anywhere in the colorbar. "How to extract a colormap from an unpickled colorbar" might be your next question.

(If I didn't have any other way, I'd recreate the colormap by making a LinearSegmentedColormap using an image editor to get the colors at each tick value.)

"""
============
Unpickle a colorbar
============

To  answer https://stackoverflow.com/questions/78088551/how-to-retrieve-a-colorbar-from-an-already-existing-matplotlib-scatter-figure

Started with https://matplotlib.org/stable/_downloads/bfeb934bd7e87a776c31ca73c7218c9e/scatter.py
"""
import matplotlib.pyplot as plt
import numpy as np
import pickle

# Fixing random state for reproducibility
np.random.seed(19680802)

N = 50
x = np.random.rand(N)
y = np.random.rand(N)
colors = np.random.rand(N)
sizes = (30 * np.random.rand(N))**2  # 0 to 15 point radii
fig, ax = plt.subplots()
scatter = ax.scatter(x, y, s=sizes, c=colors, alpha=0.5)
fig.colorbar(scatter, ax=ax)
fig.suptitle("Prepickle")
fig.savefig("unpickle_colorbar_before_pickling.png")
#plt.show() # pklsc.show() at end DOES NOT WORK if this line is uncommented?

with open("unpickle_colorbar.pkl", "wb") as f:
    pickle.dump(fig, f)

with open("unpickle_colorbar.pkl", "rb") as f:
    pklsc = pickle.load(f)

plt.figure(pklsc) 
pklsc.suptitle("Unpickled!", size=14)
pklsc.savefig("unpickle_colorbar_without_editing.png")
#If we stop here, we get the original fig, with colorbar 

# now want to reuse the scatterplot and still have the colorbar 
pklsc.suptitle("Edited the unpickled", size=14)
#maybe we only want half the original points?

oldsc = pklsc.axes[0].collections[0]

# extract locations and sizes
xd = oldsc.get_offsets()[:,0]
yd = oldsc.get_offsets()[:,1]
sd = oldsc.get_sizes()
cd = oldsc.get_facecolor()

#hide what we had, but only in the scatterplot
pklsc.axes[0].clear()

#replot
M = int(N/2)
pklsc.axes[0].scatter(xd[:M], yd[:M], s=sd[:M], c=cd[:M], alpha = 0.5)
pklsc.savefig("unpickle_colorbar_after_editing.png")

#Maybe I want a really complicated new figure!
bigfig = plt.figure()
gs = gridspec.GridSpec(1,4)
oldscatter = pklsc.axes[0]
oldcbar = pklsc.axes[1]
for i, n in enumerate((oldscatter, oldcbar)):
    n.remove()
    n.figure = bigfig
    bigfig.add_axes(n)
    n.set_subplotspec(gs[0,i])


ax_new = bigfig.add_subplot(1, 3, 2)
ax_new.plot([0,2,4],[5,2,3])
ax_new.set_subplotspec(gs[0,2])

ax_more = bigfig.add_subplot(1, 3, 3)
ax_more.scatter([.3, .5, .9], [.9, .2, 1.1])
ax_more.set_subplotspec(gs[0,3])

bigfig.suptitle("New figure with more axes")
bigfig.savefig("unpickle_colorbar_more_panels.png")

And the four states of the plot: