Matplotlib 2 subtramas, 1 barra de colores

He pasado demasiado tiempo investigando cómo obtener dos subtramas para compartir el mismo eje y con una sola barra de colores compartida entre los dos en Matplotlib.

Lo que estaba sucediendo era que cuando llamaba a la colorbar()función en uno subplot1o en otro subplot2, escalaba automáticamente el gráfico de manera que la barra de color más el gráfico encajaría dentro del cuadro delimitador 'subplot', haciendo que los dos gráficos de lado a lado fueran dos muy diferentes Tamaños.

Para evitar esto, intenté crear una tercera subtrama que luego pirateé para no representar ninguna trama con solo una barra de colores presente. El único problema es que ahora las alturas y los anchos de las dos parcelas son desiguales, y no puedo entender cómo hacer que se vea bien.

Aquí está mi código:

from __future__ import division
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import patches
from matplotlib.ticker import NullFormatter

# SIS Functions
TE = 1 # Einstein radius
g1 = lambda x,y: (TE/2) * (y**2-x**2)/((x**2+y**2)**(3/2)) 
g2 = lambda x,y: -1*TE*x*y / ((x**2+y**2)**(3/2))
kappa = lambda x,y: TE / (2*np.sqrt(x**2+y**2))

coords = np.linspace(-2,2,400)
X,Y = np.meshgrid(coords,coords)
g1out = g1(X,Y)
g2out = g2(X,Y)
kappaout = kappa(X,Y)
for i in range(len(coords)):
    for j in range(len(coords)):
        if np.sqrt(coords[i]**2+coords[j]**2) <= TE:
            g1out[i][j]=0
            g2out[i][j]=0

fig = plt.figure()
fig.subplots_adjust(wspace=0,hspace=0)

# subplot number 1
ax1 = fig.add_subplot(1,2,1,aspect='equal',xlim=[-2,2],ylim=[-2,2])
plt.title(r"$\gamma_{1}$",fontsize="18")
plt.xlabel(r"x ($\theta_{E}$)",fontsize="15")
plt.ylabel(r"y ($\theta_{E}$)",rotation='horizontal',fontsize="15")
plt.xticks([-2.0,-1.5,-1.0,-0.5,0,0.5,1.0,1.5])
plt.xticks([-2.0,-1.5,-1.0,-0.5,0,0.5,1.0,1.5])
plt.imshow(g1out,extent=(-2,2,-2,2))
plt.axhline(y=0,linewidth=2,color='k',linestyle="--")
plt.axvline(x=0,linewidth=2,color='k',linestyle="--")
e1 = patches.Ellipse((0,0),2,2,color='white')
ax1.add_patch(e1)

# subplot number 2
ax2 = fig.add_subplot(1,2,2,sharey=ax1,xlim=[-2,2],ylim=[-2,2])
plt.title(r"$\gamma_{2}$",fontsize="18")
plt.xlabel(r"x ($\theta_{E}$)",fontsize="15")
ax2.yaxis.set_major_formatter( NullFormatter() )
plt.axhline(y=0,linewidth=2,color='k',linestyle="--")
plt.axvline(x=0,linewidth=2,color='k',linestyle="--")
plt.imshow(g2out,extent=(-2,2,-2,2))
e2 = patches.Ellipse((0,0),2,2,color='white')
ax2.add_patch(e2)

# subplot for colorbar
ax3 = fig.add_subplot(1,1,1)
ax3.axis('off')
cbar = plt.colorbar(ax=ax2)

plt.show()

Simplemente coloque la barra de colores en su propio eje y úsela subplots_adjustpara dejar espacio.

Como un ejemplo rápido:

import numpy as np
import matplotlib.pyplot as plt

fig, axes = plt.subplots(nrows=2, ncols=2)
for ax in axes.flat:
    im = ax.imshow(np.random.random((10,10)), vmin=0, vmax=1)

fig.subplots_adjust(right=0.8)
cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7])
fig.colorbar(im, cax=cbar_ax)

plt.show()

Tenga en cuenta que el rango de color se establecerá en la última imagen trazada (que dio lugar a im) incluso si el rango de valores se establece con vminy vmax. Si otro gráfico tiene, por ejemplo, un valor máximo más alto, los puntos con valores más altos que el máximo de imse mostrarán en color uniforme.

Puede simplificar el código de Joe Kington usando el axparámetro de figure.colorbar()con una lista de ejes. De la documentación :

hacha

Ninguno | objeto (s) principal (es) desde el cual se robará espacio para una nueva barra de colores. Si se proporciona una lista de ejes, todos se redimensionarán para dejar espacio para los ejes de la barra de colores.

import numpy as np
import matplotlib.pyplot as plt

fig, axes = plt.subplots(nrows=2, ncols=2)
for ax in axes.flat:
    im = ax.imshow(np.random.random((10,10)), vmin=0, vmax=1)

fig.colorbar(im, ax=axes.ravel().tolist())

plt.show()

Esta solución no requiere ajustes manuales de las ubicaciones de los ejes o el tamaño de la barra de colores, funciona con diseños de varias filas y una sola fila, y funciona con tight_layout(). Está adaptado a partir de un ejemplo galería , con el ImageGridde de matplotlib AxesGrid Caja de herramientas .

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import ImageGrid

# Set up figure and image grid
fig = plt.figure(figsize=(9.75, 3))

grid = ImageGrid(fig, 111,          # as in plt.subplot(111)
                 nrows_ncols=(1,3),
                 axes_pad=0.15,
                 share_all=True,
                 cbar_location="right",
                 cbar_mode="single",
                 cbar_size="7%",
                 cbar_pad=0.15,
                 )

# Add data to image grid
for ax in grid:
    im = ax.imshow(np.random.random((10,10)), vmin=0, vmax=1)

# Colorbar
ax.cax.colorbar(im)
ax.cax.toggle_label(True)

#plt.tight_layout()    # Works, but may still require rect paramater to keep colorbar labels visible
plt.show()

Usar make_axeses aún más fácil y ofrece un mejor resultado. También ofrece posibilidades para personalizar el posicionamiento de la barra de colores. También tenga en cuenta la opción de subplotscompartir los ejes x e y.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl

fig, axes = plt.subplots(nrows=2, ncols=2, sharex=True, sharey=True)
for ax in axes.flat:
    im = ax.imshow(np.random.random((10,10)), vmin=0, vmax=1)

cax,kw = mpl.colorbar.make_axes([ax for ax in axes.flat])
plt.colorbar(im, cax=cax, **kw)

plt.show()

Como principiante que tropezó con este hilo, me gustaría agregar una adaptación de python-for-dummies de la muy clara respuesta de abevieiramota (porque estoy en el nivel en el que tuve que buscar 'ravel' para averiguar qué su código estaba haciendo):

import numpy as np
import matplotlib.pyplot as plt

fig, ((ax1,ax2,ax3),(ax4,ax5,ax6)) = plt.subplots(2,3)

axlist = [ax1,ax2,ax3,ax4,ax5,ax6]

first = ax1.imshow(np.random.random((10,10)), vmin=0, vmax=1)
third = ax3.imshow(np.random.random((12,12)), vmin=0, vmax=1)

fig.colorbar(first, ax=axlist)

plt.show()

Mucho menos pitónico, mucho más fácil para novatos como yo para ver lo que realmente está sucediendo aquí.

Como se señaló en otras respuestas, la idea suele ser definir un eje para que resida la barra de colores. Hay varias formas de hacerlo; uno que no ha sido mencionado todavía habría que especificar directamente los ejes de barra de colores en la creación trama secundaria con plt.subplots(). La ventaja es que no es necesario establecer manualmente la posición de los ejes y, en todos los casos con aspecto automático, la barra de colores tendrá exactamente la misma altura que las subtramas. Incluso en muchos casos donde se usan imágenes, el resultado será satisfactorio como se muestra a continuación.

Cuando se usa plt.subplots(), el uso del gridspec_kwargumento permite hacer que los ejes de la barra de colores sean mucho más pequeños que los otros ejes.

fig, (ax, ax2, cax) = plt.subplots(ncols=3,figsize=(5.5,3), 
                  gridspec_kw={"width_ratios":[1,1, 0.05]})

Ejemplo:

import matplotlib.pyplot as plt
import numpy as np; np.random.seed(1)

fig, (ax, ax2, cax) = plt.subplots(ncols=3,figsize=(5.5,3), 
                  gridspec_kw={"width_ratios":[1,1, 0.05]})
fig.subplots_adjust(wspace=0.3)
im  = ax.imshow(np.random.rand(11,8), vmin=0, vmax=1)
im2 = ax2.imshow(np.random.rand(11,8), vmin=0, vmax=1)
ax.set_ylabel("y label")

fig.colorbar(im, cax=cax)

plt.show()

Esto funciona bien, si el aspecto de los gráficos se escala automáticamente o las imágenes se reducen debido a su aspecto en la dirección del ancho (como en el anterior). Sin embargo, si las imágenes son más anchas que altas, el resultado sería el siguiente, que podría no ser deseado.

Una solución para fijar la altura de la barra de colores a la altura de la subtrama sería utilizar mpl_toolkits.axes_grid1.inset_locator.InsetPositionpara establecer los ejes de la barra de colores en relación con los ejes de la subtrama de la imagen.

import matplotlib.pyplot as plt
import numpy as np; np.random.seed(1)
from mpl_toolkits.axes_grid1.inset_locator import InsetPosition

fig, (ax, ax2, cax) = plt.subplots(ncols=3,figsize=(7,3), 
                  gridspec_kw={"width_ratios":[1,1, 0.05]})
fig.subplots_adjust(wspace=0.3)
im  = ax.imshow(np.random.rand(11,16), vmin=0, vmax=1)
im2 = ax2.imshow(np.random.rand(11,16), vmin=0, vmax=1)
ax.set_ylabel("y label")

ip = InsetPosition(ax2, [1.05,0,0.05,1]) 
cax.set_axes_locator(ip)

fig.colorbar(im, cax=cax, ax=[ax,ax2])

plt.show()

La solución de usar una lista de ejes por abevieiramota funciona muy bien hasta que use solo una fila de imágenes, como se señala en los comentarios. Usar una relación de aspecto razonable para obtener figsizeayuda, pero aún está lejos de ser perfecto. Por ejemplo:

import numpy as np
import matplotlib.pyplot as plt

fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(9.75, 3))
for ax in axes.flat:
    im = ax.imshow(np.random.random((10,10)), vmin=0, vmax=1)

fig.colorbar(im, ax=axes.ravel().tolist())

plt.show()

La función de barra de colores proporciona el shrinkparámetro que es un factor de escala para el tamaño de los ejes de la barra de colores. Requiere algún tipo de prueba y error manual. Por ejemplo:

fig.colorbar(im, ax=axes.ravel().tolist(), shrink=0.75)

Para agregar a la excelente respuesta de @ abevieiramota, puede obtener el euqivalente de tight_layout con restrictined_layout. Todavía obtendrá grandes espacios horizontales si lo usa en imshowlugar de pcolormeshdebido a la relación de aspecto 1: 1 impuesta por imshow.

import numpy as np
import matplotlib.pyplot as plt

fig, axes = plt.subplots(nrows=2, ncols=2, constrained_layout=True)
for ax in axes.flat:
    im = ax.pcolormesh(np.random.random((10,10)), vmin=0, vmax=1)

fig.colorbar(im, ax=axes.flat)
plt.show()

Noté que casi todas las soluciones publicadas involucraban ax.imshow(im, ...)y no normalizaban los colores mostrados en la barra de colores para las múltiples subfiguras. El immapeable se toma de la última instancia, pero ¿qué pasa si los valores de los múltiples im-s son diferentes? (Supongo que estos mapas se tratan de la misma manera que se tratan los conjuntos de contorno y los conjuntos de superficie). Tengo un ejemplo que usa un diagrama de superficie en 3D a continuación que crea dos barras de colores para una subtrama de 2x2 (una barra de colores por una fila ) Aunque la pregunta pide explícitamente un arreglo diferente, creo que el ejemplo ayuda a aclarar algunas cosas. plt.subplots(...)Aún no he encontrado una manera de hacerlo usando los ejes 3D, desafortunadamente.

Si tan solo pudiera posicionar las barras de color de una mejor manera ... (Probablemente haya una forma mucho mejor de hacer esto, pero al menos no debería ser demasiado difícil de seguir).

import matplotlib
from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

cmap = 'plasma'
ncontours = 5

def get_data(row, col):
    """ get X, Y, Z, and plot number of subplot
        Z > 0 for top row, Z < 0 for bottom row """
    if row == 0:
        x = np.linspace(1, 10, 10, dtype=int)
        X, Y = np.meshgrid(x, x)
        Z = np.sqrt(X**2 + Y**2)
        if col == 0:
            pnum = 1
        else:
            pnum = 2
    elif row == 1:
        x = np.linspace(1, 10, 10, dtype=int)
        X, Y = np.meshgrid(x, x)
        Z = -np.sqrt(X**2 + Y**2)
        if col == 0:
            pnum = 3
        else:
            pnum = 4
    print("\nPNUM: {}, Zmin = {}, Zmax = {}\n".format(pnum, np.min(Z), np.max(Z)))
    return X, Y, Z, pnum

fig = plt.figure()
nrows, ncols = 2, 2
zz = []
axes = []
for row in range(nrows):
    for col in range(ncols):
        X, Y, Z, pnum = get_data(row, col)
        ax = fig.add_subplot(nrows, ncols, pnum, projection='3d')
        ax.set_title('row = {}, col = {}'.format(row, col))
        fhandle = ax.plot_surface(X, Y, Z, cmap=cmap)
        zz.append(Z)
        axes.append(ax)

## get full range of Z data as flat list for top and bottom rows
zz_top = zz[0].reshape(-1).tolist() + zz[1].reshape(-1).tolist()
zz_btm = zz[2].reshape(-1).tolist() + zz[3].reshape(-1).tolist()
## get top and bottom axes
ax_top = [axes[0], axes[1]]
ax_btm = [axes[2], axes[3]]
## normalize colors to minimum and maximum values of dataset
norm_top = matplotlib.colors.Normalize(vmin=min(zz_top), vmax=max(zz_top))
norm_btm = matplotlib.colors.Normalize(vmin=min(zz_btm), vmax=max(zz_btm))
cmap = cm.get_cmap(cmap, ncontours) # number of colors on colorbar
mtop = cm.ScalarMappable(cmap=cmap, norm=norm_top)
mbtm = cm.ScalarMappable(cmap=cmap, norm=norm_btm)
for m in (mtop, mbtm):
    m.set_array([])

# ## create cax to draw colorbar in
# cax_top = fig.add_axes([0.9, 0.55, 0.05, 0.4])
# cax_btm = fig.add_axes([0.9, 0.05, 0.05, 0.4])
cbar_top = fig.colorbar(mtop, ax=ax_top, orientation='vertical', shrink=0.75, pad=0.2) #, cax=cax_top)
cbar_top.set_ticks(np.linspace(min(zz_top), max(zz_top), ncontours))
cbar_btm = fig.colorbar(mbtm, ax=ax_btm, orientation='vertical', shrink=0.75, pad=0.2) #, cax=cax_btm)
cbar_btm.set_ticks(np.linspace(min(zz_btm), max(zz_btm), ncontours))

plt.show()
plt.close(fig)
## orientation of colorbar = 'horizontal' if done by column

Este tema está bien cubierto, pero aún me gustaría proponer otro enfoque en una filosofía ligeramente diferente.

Es un poco más complejo de configurar, pero permite (en mi opinión) un poco más de flexibilidad. Por ejemplo, uno puede jugar con las proporciones respectivas de cada subtrama / barra de colores:

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.gridspec import GridSpec

# Define number of rows and columns you want in your figure
nrow = 2
ncol = 3

# Make a new figure
fig = plt.figure(constrained_layout=True)

# Design your figure properties
widths = [3,4,5,1]
gs = GridSpec(nrow, ncol + 1, figure=fig, width_ratios=widths)

# Fill your figure with desired plots
axes = []
for i in range(nrow):
    for j in range(ncol):
        axes.append(fig.add_subplot(gs[i, j]))
        im = axes[-1].pcolormesh(np.random.random((10,10)))

# Shared colorbar    
axes.append(fig.add_subplot(gs[:, ncol]))
fig.colorbar(im, cax=axes[-1])

plt.show()