Matplotlib: colorbar outside plot in subplot2grid

Question

I'm trying to draw two gaussians distributions in some subplots. This is the result:

And here is the code

import numpy as np
import numpy.random as rnd
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable



mu_x1 = 3
sig_x1 = 2

mu_x2 = 3
sig_x2 = 0.5

N = int(1e3)

gaussian1 = rnd.normal(mu_x1, sig_x1, N)
gaussian2 = rnd.normal(mu_x2, sig_x2, N)
fig = plt.figure()

ax1 = plt.subplot2grid(shape = (3, 3),
                       loc = (1, 0),
                       rowspan = 2,
                       colspan = 2,
                       fig = fig
                       )
hist2d = ax1.hist2d(gaussian1,
                    gaussian2,
                    bins = 100,
                    density = True,
                    cmap=plt.cm.jet
                    )
fig.colorbar(hist2d[3], ax=ax1)
ax1.set_xlabel("$x_1$")
ax1.set_ylabel("$x_2$")



ax2 = plt.subplot2grid(shape = (3, 3),
                       loc = (0, 0),
                       rowspan = 1,
                       colspan = 2,
                       fig = fig
                       )
ax2.hist(gaussian1,
         bins = 100
         )
ax2.xaxis.set_ticks_position('top')
ax2.set_xlabel("$x_1$")
ax2.xaxis.set_label_position('top')
ax2.set_ylabel("Cuentas")

ax3 = plt.subplot2grid(shape = (3, 3),
                       loc = (1, 2),
                       rowspan = 2,
                       colspan = 1,
                       fig = fig
                       )
ax3.hist(gaussian2,
         bins = 100,
         orientation=u'horizontal')
ax3.yaxis.set_ticks_position('right')
ax3.yaxis.set_label_position('right')
ax3.set_xlabel("Cuentas")
ax3.set_ylabel("$x_2$")

plt.show()

It looks good but not as I want. I want the x axis of the middle subplot to be as large as the the x axis of the subplot above, so they coincide. The problem is the colorbar. I would like to move the colorbar outside the middle subplot to do what I want. I looked some info and I arrived to:

Here the code

fig = plt.figure()

ax1 = plt.subplot2grid(shape = (4, 3),
                       loc = (1, 0),
                       rowspan = 2,
                       colspan = 2,
                       fig = fig
                       )
hist2d = ax1.hist2d(gaussian1,
                    gaussian2,
                    bins = 100,
                    density = True,
                    cmap=plt.cm.jet
                    )
#fig.colorbar(hist2d[3], ax=ax1)
ax4 = plt.subplot2grid(shape = (4, 3),
                       loc = (3, 0),
                       rowspan = 1,
                       colspan = 2,
                       fig = fig
                       )    
fig.colorbar(hist2d[3], cax=ax4, orientation="horizontal")
ax1.set_xlabel("$x_1$")
ax1.set_ylabel("$x_2$")
ax1.axhline(mu_x2 + sig_x2*n2,c="k")
ax1.axhline(mu_x2 - sig_x2*n2,c="k")
ax1.axvline(mu_x1 + sig_x1*n1,c="k")
ax1.axvline(mu_x1 - sig_x1*n1,c="k")



ax2 = plt.subplot2grid(shape = (4, 3),
                       loc = (0, 0),
                       rowspan = 1,
                       colspan = 2,
                       fig = fig
                       )
ax2.hist(gaussian1,
         bins = 100
         )
ax2.xaxis.set_ticks_position('top')
ax2.set_xlabel("$x_1$")
ax2.xaxis.set_label_position('top')
ax2.set_ylabel("Cuentas")
ax2.axvline(mu_x1 + sig_x1*n1,c="k")
ax2.axvline(mu_x1 - sig_x1*n1,c="k")
#ax2.set_xticks([])

ax3 = plt.subplot2grid(shape = (4, 3),
                       loc = (1, 2),
                       rowspan = 2,
                       colspan = 1,
                       fig = fig
                       )
ax3.hist(gaussian2,
         bins = 100,
         orientation=u'horizontal')
ax3.yaxis.set_ticks_position('right')
ax3.yaxis.set_label_position('right')
ax3.set_xlabel("Cuentas")
ax3.set_ylabel("$x_2$")
ax3.axhline(mu_x2 + sig_x2*n2,c="k")
ax3.axhline(mu_x2 - sig_x2*n2,c="k")

The problem now is that the colorbar is filling a full subplot, it is too big and it is hidding the ticks of the subplot above.

Any idea to solve it?

Answer 1

I had the same issue. Here is how I manage it using add_axes function from this answer.

import numpy as np
import numpy.random as rnd
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable



mu_x1 = 3
sig_x1 = 2

mu_x2 = 3
sig_x2 = 0.5

N = int(1e3)

gaussian1 = rnd.normal(mu_x1, sig_x1, N)
gaussian2 = rnd.normal(mu_x2, sig_x2, N)
fig = plt.figure(figsize=(10,10))

ax1 = plt.subplot2grid(shape = (3, 3),
                       loc = (1, 0),
                       rowspan = 2,
                       colspan = 2,
                       fig = fig
                       )
hist2d = ax1.hist2d(gaussian1,
                    gaussian2,
                    bins = 100,
                    density = True,
                    cmap=plt.cm.jet
                    )
#fig.colorbar(hist2d[3], ax=ax1)
# ax4 = plt.subplot2grid(shape = (4, 3),
#                        loc = (3, 0),
#                        rowspan = 1,
#                        colspan = 2,
#                        fig = fig
#                        )    
# fig.colorbar(hist2d[3], cax=ax4, orientation="horizontal")
ax1.set_xlabel("$x_1$")
ax1.set_ylabel("$x_2$")

w=0.5
h=0.03

w1=w*(ax1.get_position().x1-ax1.get_position().x0)
x1=ax1.get_position().x0+w1/2
y1=ax1.get_position().y0-5*h

cax = fig.add_axes([x1,y1,w1,h])

plt.colorbar(hist2d[3],orientation='horizontal',cax=cax,aspect=10,shrink=0.5)



ax2 = plt.subplot2grid(shape = (3, 3),
                       loc = (0, 0),
                       rowspan = 1,
                       colspan = 2,
                       fig = fig
                       )
ax2.hist(gaussian1,
         bins = 100
         )
ax2.xaxis.set_ticks_position('top')
ax2.set_xlabel("$x_1$")
ax2.xaxis.set_label_position('top')
ax2.set_ylabel("Cuentas")
#ax2.set_xticks([])

ax3 = plt.subplot2grid(shape = (3, 3),
                       loc = (1, 2),
                       rowspan = 2,
                       colspan = 1,
                       fig = fig
                       )
ax3.hist(gaussian2,
         bins = 100,
         orientation=u'horizontal')
ax3.yaxis.set_ticks_position('right')
ax3.yaxis.set_label_position('right')
ax3.set_xlabel("Cuentas")
ax3.set_ylabel("$x_2$")