How to add colorbars to scatterplots created like this?

Question

I create scatterplots with code that, in essence, goes like this

cmap = (matplotlib.color.LinearSegmentedColormap.
        from_list(\'blueWhiteRed\', [\'blu         


        

Answer 1

I think this should do the trick. I'm pretty sure I grabbed this from one of the matplotlib cookbook examples a while back, but I can't seem to find it now...

from mpl_toolkits.axes_grid1 import make_axes_locatable

cmap = (matplotlib.color.LinearSegmentedColormap.
        from_list('blueWhiteRed', ['blue', 'white', 'red']))

fig = matplotlib.figure.Figure(figsize=(4, 4), dpi=72)
ax = fig.gca()

for record in data:
    level = record.level # a float in [0.0, 1.0]
    ax.scatter(record.x, record.y,
               c=level, vmin=0, vmax=1, cmap=cmap, **otherkwargs)

# various settings of ticks, labels, etc. omitted

divider= make_axes_locatable(ax)
cax = divider.append_axes("right", size="1%", pad=0.05)
cb = plt.colorbar(cax=cax)
cb.set_label('Cbar Label Here')

canvas = matplotlib.backends.backend_agg.FigureCanvasAgg(fig)
fig.set_canvas(canvas)
canvas.print_png('/path/to/output/fig.png')

Answer 2

If you have to use a different marker for each set, you have to do a bit of extra work and force all of the clims to be the same (otherwise they default to scaling from the min/max of the c data per scatter plot).

from pylab import *
import matplotlib.lines as mlines
import itertools
fig = gcf()
ax = fig.gca()

# make some temorary arrays
X = []
Y = []
C = []
cb = None
# generate fake data
markers = ['','o','*','^','v']
cmin = 0
cmax = 1
for record,marker in itertools.izip(range(5),itertools.cycle(mlines.Line2D.filled_markers)):
    x = rand(50)
    y = rand(50)
    c = rand(1)[0] * np.ones(x.shape)
    if cb is None:
        s = ax.scatter(x,y,c=c,marker=markers[record],linewidths=0)
        s.set_clim([cmin,cmax])
        cb = fig.colorbar(s)
    else:
        s = ax.scatter(x,y,c=c,marker=markers[record],linewidths=0)
        s.set_clim([cmin,cmax])

cb.set_label('Cbar Label Here')

thelinewidths=0 sets the width of the border on the shapes, I find that for small shapes the black border can overwhelm the color of the fill.

If you only need one shape you can do this all with a single scatter plot, there is no need to make a separate one for each pass through your loop.

from pylab import *
fig = gcf()
ax = fig.gca()

# make some temorary arrays
X = []
Y = []
C = []
# generate fake data
for record in range(5):
    x = rand(50)
    y = rand(50)
    c = rand(1)[0] * np.ones(x.shape)
    print c
    X.append(x)
    Y.append(y)
    C.append(c)

X = np.hstack(X)
Y = np.hstack(Y)
C = np.hstack(C)

once you have the data all beaten down into a 1D array, make the scatter plot, and keep the returned value:

s = ax.scatter(X,Y,c=C)

You then make your color bar and pass the object returned by scatter as the first argument.

cb = plt.colorbar(s)
cb.set_label('Cbar Label Here')

You need do this so that the color bar knows which color map (both the map and the range) to use.

Answer 3

The answer to this can be to only plot a single scatter, which would then directly allow for a colobar to be created. This involves putting the markers into the PathCollection created by the scatter a posteriori, but it can be easily placed in a function. This function comes from my answer on another question, but is directly applicable here.

Taking the data from @PaulH's post this would look like

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

def mscatter(x,y,ax=None, m=None, **kw):
    import matplotlib.markers as mmarkers
    ax = ax or plt.gca()
    sc = ax.scatter(x,y,**kw)
    if (m is not None) and (len(m)==len(x)):
        paths = []
        for marker in m:
            if isinstance(marker, mmarkers.MarkerStyle):
                marker_obj = marker
            else:
                marker_obj = mmarkers.MarkerStyle(marker)
            path = marker_obj.get_path().transformed(
                        marker_obj.get_transform())
            paths.append(path)
        sc.set_paths(paths)
    return sc


markers = ['s', 'o', '^']
records = []
for n in range(37):
    records.append([np.random.normal(), np.random.normal(), np.random.normal(), 
                    markers[np.random.randint(0, high=3)]])

records = pd.DataFrame(records, columns=['x', 'y', 'z', 'marker'])

fig, ax = plt.subplots()
sc = mscatter(records.x, records.y, c=records.z, m=records.marker, ax=ax)
fig.colorbar(sc, ax=ax)

plt.show()

Answer 4

I think your best bet will be to stuff your data into a pandas dataframe, and loop through all of your markers like so:

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

markers = ['s', 'o', '^']
records = []
for n in range(37):
    records.append([np.random.normal(), np.random.normal(), np.random.normal(), 
                    markers[np.random.randint(0, high=3)]])

records = pd.DataFrame(records, columns=['x', 'y', 'z', 'marker'])

fig, ax = plt.subplots()
for m in np.unique(records.marker):
    selector = records.marker == m
    s = ax.scatter(records[selector].x, records[selector].y, c=records[selector].z,
                   marker=m, cmap=plt.cm.coolwarm, 
                   vmin=records.z.min(), vmax=records.z.max())

cbar = plt.colorbar(mappable=s, ax=ax)
cbar.set_label('My Label')