Adding colorbar to scatterplot after loop

Question

I am trying to add a continuous colorbar to a seaborn scatterplot (similar to the answers here and in here). For my purposes, I am building the scatterplot with a loop, and then trying to add the continuous colorbar, but I dont know what object to include as argument of fig.colorbar(). How would you do this?

import pandas as pd
import seaborn as sb
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(42)

df=pd.DataFrame(np.random.rand(2, 100), index=['S1','S2']).T
tars=np.random.choice([0,0.3,0.5,0.8,1], 100)
df=pd.concat([df,pd.Series(tars, name='group')],1)
colors = matplotlib.cm.viridis(np.linspace(0,1,len(pd.unique(tars))))

fig = plt.figure(figsize = (10,8), dpi=300)
ax = fig.add_subplot(1,1,1) 
targets=pd.unique(tars)
for target, color in zip(targets,colors):
    ...
    g=ax.scatter(
        df.loc[df.group==target, 'S1'], 
        df.loc[df.group==target, 'S2'],
        color = [color]
    )
fig.colorbar(g)
plt.show()

If I add ax.legend(targets) instead of fig.colorbar(g), the legend displays correctly but is categorical.

df=pd.DataFrame(np.random.rand(2, 100), index=['S1','S2']).T
tars=np.random.choice([0,0.3,0.5,0.8,1], 100)
df=pd.concat([df,pd.Series(tars, name='group')],1)

cmap=matplotlib.cm.gnuplot2
colors = cmap(np.linspace(0,1,len(pd.unique(tars))))

fig = plt.figure(figsize = (10,8), dpi=300)
ax = fig.add_subplot(1,1,1) 
targets=pd.unique(tars)
for target, color in zip(targets,colors):
    ...
    g=ax.scatter(
        df.loc[df.group==target, 'S1'], 
        df.loc[df.group==target, 'S2'],
        color = [color]
    )
ax.legend(targets)
plt.show()

user_1356000 · Answer

I am not entirely sure to understand what you are trying to achieve with your for-loop.

Is this the output that you are looking for?

fig = plt.figure()
ax = fig.add_subplot(1,1,1) 
g = ax.scatter(df['S1'],df['S2'],c=df['group'],cmap='viridis')
cbar = fig.colorbar(g)
plt.show()

user_1379826 · Answer

Thanks to this answer I could edit my code to show the continuous colorbar.

df=pd.DataFrame(np.random.rand(2, 100), index=['S1','S2']).T
tars=np.random.choice([0,0.3,0.5,0.8,1], 100)
df=pd.concat([df,pd.Series(tars, name='group')],1)

cmap=matplotlib.cm.viridis
colors = cmap(np.linspace(0,1,len(pd.unique(tars))))

fig = plt.figure(figsize = (10,8), dpi=300)
ax = fig.add_subplot(1,1,1) 
targets=pd.unique(tars)
for target, color in zip(targets,colors):
    g=ax.scatter(
        df.loc[df.group==target, 'S1'], 
        df.loc[df.group==target, 'S2'],
        color=[color]
    )

norm = plt.Normalize(np.min(tars), np.max(tars))
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
ax.figure.colorbar(sm)
plt.show()

Adding colorbar to scatterplot after loop

2 Answers