Adding Colorbars to the CEBRA embeddings

[BrianNGitahi]

Hi, this might be a trivial question but I can't find an example that shows how you can add a colorbar to the figure of the embedding produced by CEBRA. Has anyone been able to do this? Thanks

[stes]

Hi @BrianNGitahi , thanks a lot for using CEBRA!

To your question, adding a colorbar is possible. Here is an example:

# create example data
time = np.linspace(0, 20, 1000)
x = np.sin(time)
y = np.cos(time)
z = time
embedding = np.stack([x, y, z], axis = 1)

# the actual plotting code
plt.figure(figsize = (3,3))
ax = cebra.plot_embedding(embedding, embedding_labels = time, alpha = 1, title = "Colorbar demo")
scatter = ax.collections[-1] # this gets the scatter plot
plt.colorbar(scatter, ax = ax, pad=0.2, shrink=0.5) # this adds the colorbar. Also check out the cax= arg for more customization!
plt.show()

We will consider adding this to the API in the future, but hopefully for now this works for you! (please mark my comment as an answer if it does)