Установить равный аспект на графике с цветовой шкалой

Question

Мне нужно сгенерировать график с одинаковым аспектом по оси и цветовой шкале справа. Я пытался установить aspect='auto', aspect=1, а также aspect='equal' без хороших результатов. Ниже приведены примеры и MWE.

С помощью aspect='auto' цветные полосы имеют правильную высоту, но графики искажены:

С помощью aspect=1 или же aspect='equal' графики являются квадратными (одинаковый аспект на обеих осях), но цветные полосы искажены:

На обоих графиках цветные полосы по какой-то причине расположены слишком далеко вправо. Как я могу получить квадратный график с цветными полосами соответствующей высоты?

---

MWE

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

def col_plot(params):

    gs, i, data = params

    xarr, yarr, zarr = zip(*data)[0], zip(*data)[1], zip(*data)[2]

    xmin, xmax = min(xarr), max(xarr)
    ymin, ymax = min(yarr), max(yarr)

    #plt.subplot(gs[i], aspect='auto')
    plt.subplot(gs[i], aspect=1)
    #plt.subplot(gs[i], aspect='equal')

    plt.xlim(xmin, xmax)
    plt.ylim(xmin, xmax)
    plt.xlabel('$x axis$', fontsize=20)
    plt.ylabel('$y axis$', fontsize=20)
    # Scatter plot.
    cm = plt.cm.get_cmap('RdYlBu_r')
    SC = plt.scatter(xarr, yarr, marker='o', c=zarr, s=60, lw=0.25, cmap=cm,
        zorder=3)
    # Colorbar.
    ax0 = plt.subplot(gs[i + 1])
    cbar = plt.colorbar(SC, cax=ax0)
    cbar.set_label('$col bar$', fontsize=21, labelpad=-2)

# Generate data.
data0 = np.random.uniform(0., 1., size=(50, 3))
data1 = np.random.uniform(0., 1., size=(50, 3))

# Create the top-level container
fig = plt.figure(figsize=(14, 25))
gs = gridspec.GridSpec(4, 4, width_ratios=[1, 0.05, 1, 0.05])

# Generate plots.
par_lst = [[gs, 0, data0], [gs, 2, data1]]
for pl_params in par_lst:
    col_plot(pl_params)

# Output png file.
fig.tight_layout()
plt.savefig('colorbar_aspect.png', dpi=300)

Answer 1

Вы можете использовать AxesDivider для этого. Я немного изменил ваш код, чтобы использовать AxesDivider.

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

def col_plot(data):

    xarr, yarr, zarr = zip(*data)[0], zip(*data)[1], zip(*data)[2]
    xarr = [2*x for x in xarr]

    xmin, xmax = min(xarr), max(xarr)
    ymin, ymax = min(yarr), max(yarr)

    fig = plt.figure()

    ax0 = fig.add_subplot(111, aspect='equal')

    plt.xlim(xmin, xmax)
    plt.ylim(ymin, ymax)
    plt.xlabel('$x axis$', fontsize=20)
    plt.ylabel('$y axis$', fontsize=20)
    # Scatter plot.
    cm = plt.cm.get_cmap('RdYlBu_r')
    SC = ax0.scatter(xarr, yarr, marker='o', c=zarr, s=60, lw=0.25, cmap=cm,
        zorder=3)

    the_divider = make_axes_locatable(ax0)
    color_axis = the_divider.append_axes("right", size="5%", pad=0.1)

    # Colorbar.
    cbar = plt.colorbar(SC, cax=color_axis)
    cbar.set_label('$col bar$', fontsize=21, labelpad=-2)

# Generate data.
data0 = np.random.uniform(0., 1., size=(20, 3))

col_plot(data0)

И вот результат (я изменил ваши данные, чтобы они охватывали диапазон [0, 2] в направлении x для демонстрационных целей):

Answer 2

В блоге Джозефа Лонга есть следующее хорошее решение.

1) Определите colorbar функционировать как:

from mpl_toolkits.axes_grid1 import make_axes_locatable

def colorbar(mappable):
    ax = mappable.axes
    fig = ax.figure
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)
    return fig.colorbar(mappable, cax=cax)

2) Позвонить colorbar(thing) когда вы хотите сделать цветную полосу. В твоем случае:

SC = ax0.scatter(xarr, yarr, marker='o', c=zarr, s=60, lw=0.25, cmap=cm,
    zorder=3)

colorbar(SC)

3) И вы получите: