Search code examples
pandasmatplotlibplot

Getting Colorbar instance of scatter plot in pandas/matplotlib


How do I get the internally created colorbar instance of a plot created by pandas.DataFrame.plot?

Here is an example for generating a colored scatter plot:

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import itertools as it

# [ (0,0), (0,1), ..., (9,9) ]
xy_positions = list( it.product( range(10), range(10) ) )

df = pd.DataFrame( xy_positions, columns=['x','y'] )

# draw 100 floats
df['score'] = np.random.random( 100 )

ax = df.plot( kind='scatter',
              x='x',
              y='y',
              c='score',
              s=500)
ax.set_xlim( [-0.5,9.5] )
ax.set_ylim( [-0.5,9.5] )

plt.show()

which gives me a figure like this: enter image description here

How do I get the colorbar instance in order to manipulate it, for instance for changing the label or setting the ticks?


Solution

  • pandas does not return the axis for the colorbar, therefore we have to locate it:

    1st, let's get the figure instance: i.e., use plt.gcf()

    In [61]:
    
    import matplotlib.pyplot as plt
    import pandas as pd
    import numpy as np
    import itertools as it
    
    # [ (0,0), (0,1), ..., (9,9) ]
    xy_positions = list( it.product( range(10), range(10) ) )
    
    df = pd.DataFrame( xy_positions, columns=['x','y'] )
    
    # draw 100 floats
    df['score'] = np.random.random( 100 )
    
    ax = df.plot( kind='scatter',
                  x='x',
                  y='y',
                  c='score',
                  s=500)
    ax.set_xlim( [-0.5,9.5] )
    ax.set_ylim( [-0.5,9.5] )
    
    f = plt.gcf()
    

    2, how many axes does this figure have?

    In [62]:
    
    f.get_axes()
    Out[62]:
    [<matplotlib.axes._subplots.AxesSubplot at 0x120a4d450>,
     <matplotlib.axes._subplots.AxesSubplot at 0x120ad0050>]
    

    3, The first axes (that is, the first one created), contains the plot

    In [63]:
    
    ax
    Out[63]:
    <matplotlib.axes._subplots.AxesSubplot at 0x120a4d450>
    

    4, Therefore, the second axis is the colorbar axes

    In [64]:
    
    cax = f.get_axes()[1]
    #and we can modify it, i.e.:
    cax.set_ylabel('test')