Search code examples
matplotlibcolorbar

How to plot a colorbar in a given axis in Matplotlib?

I want to plot there subplots in a figure using GridSpec in matplotlib.

In the first rows (varying number of rows), I want to place a grid with circles.

In the second last row, I want to place a legend.

In the last row, I want to place a colorbar.

But it seems that my colorbar is taking more space than I have given.

Do you have any idea to solve it?

Here is the code to regenerate it:

from matplotlib import  pyplot as plt
import matplotlib.cm as matplotlib_cm
import numpy as np
import matplotlib as mpl
import os

def plot_legend_in_given_axis(ax, fontsize):
    diameter_labels = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
    row_labels = ['circle']
    ax.grid(which="major", color="white", linestyle='-', linewidth=3)
    ax.set_aspect(1.0)

    for row_index, row_label in enumerate(row_labels):
        for diameter_index, diameter_label in enumerate(diameter_labels):
            circle=plt.Circle((diameter_index + 0.5, row_index + 0.5), radius=(diameter_label/(2*1.09)), color='gray', fill=True)
            ax.add_artist(circle)

    ax.set_xlim([0, len(diameter_labels)])
    ax.set_xticklabels([])

    ax.tick_params(axis='x', which='minor', length=0, labelsize=12)
    ax.set_xticks(np.arange(0, len(diameter_labels), 1))
    ax.xaxis.set_ticks_position('bottom')
    ax.set_xlabel('Very long\nx axis label', fontsize=fontsize, labelpad=5)

    ax.set_ylim([0, len(row_labels)])
    ax.set_yticklabels([])
    ax.tick_params(axis='y', which='minor', length=0, labelsize=12)
    ax.set_yticks(np.arange(0, len(row_labels), 1))

    ax.grid(which='major', color='black', linestyle='-', linewidth=1)


def main(width, height, xaxis_labels, yaxis_labels):
    width_multiply = 1.5
    height_multiply = 1.5

    figure_width = int(width_multiply * width)
    figure_height = int(height_multiply * height)

    fig = plt.figure(figsize=(figure_width, figure_height))

    grid = plt.GridSpec(width, height, hspace=0, wspace=0)  
    ax = fig.add_subplot(grid[0:-2, :])
    ax.set_aspect(1.0)
    legend_ax = fig.add_subplot(grid[-2, :])
    plot_legend_in_given_axis(legend_ax, 20)

    color_bar_ax = fig.add_subplot(grid[-1, :])
    cmap = matplotlib_cm.get_cmap('YlOrRd')
    v_min = 2
    v_max = 20
    norm = plt.Normalize(v_min, v_max)

    bounds = np.arange(v_min, v_max + 1, 2)
    cb = mpl.colorbar.ColorbarBase(color_bar_ax, cmap=cmap, norm=norm, ticks=bounds, spacing='proportional', orientation='horizontal')
    cb.ax.tick_params(labelsize=20)
    cb.set_label("-log10 (q-value)", horizontalalignment='center', rotation=0, fontsize=20)

    plt.xlim([0, width])
    ax.set_xticks(np.arange(0, width + 1, 1))

    plt.ylim([1, height])
    ax.set_yticks(np.arange(0, height + 1, 1))

    # Plot the circles with color in the grid
    for yaxis_index, yaxis_label in enumerate(yaxis_labels):
        for xaxis_index, xaxis_label in enumerate(xaxis_labels):
            circle = plt.Circle((xaxis_index + 0.5, yaxis_index + 0.5), 0.4, color=cmap(norm(10)), fill=True)
            ax.add_artist(circle)

    ax.grid()
    for edge, spine in ax.spines.items():
        spine.set_visible(True)
        spine.set_color('black')

    figures_path = os.path.join('/Users', 'burcakotlu', 'Desktop')
    figFile = os.path.join(figures_path, 'ColorBar_Using_GridSpec.png')
    fig.savefig(figFile, dpi=100, bbox_inches="tight")

    plt.cla()
    plt.close(fig)

width = 5
height = 5
xaxis_labels = ["x%d" %i for i in range(width)]
yaxis_labels = ["y%d" %i for i in range(height)]
main(width, height, xaxis_labels, yaxis_labels)

enter image description here

Solution

  • This issue comes from the fact that you are using plt.xlim([0, width]) and plt.ylim([1, height]) to set the range of the axes of your first subplot (with axes ax) immediately after defining the axes color_bar_ax of your colorbar. If instead you use ax.set_xlim([0, width]) and ax.set_ylim([1, height]) then you are specifying that you are now referring to your fist subplot (ax) and everything works well.

    The output I get from making those changes can be seen below:

    enter image description here

    Optional modifications:

    On top of that, you can change the way you distribute you plots on your GridSpec such that there is no overlap between plots and labels. Below, I changed the location of the first subplot on the grid from grid[0:-2, :]) to grid[0:-3, :] and the second subplot went from grid[-2,:] to grid[-3:-1, :]. The output then becomes:

    enter image description here