Matplotlib, vertical space between legend symbols

Viewed 360

I have an issue with customizing the legend of my plot. I did lot's of customizing but couldnt get my head around this one. I want the symbols (not the labels) to be equally spaced in the legend. As you can see in the example, the space between the circles in the legend, gets smaller as the circles get bigger. any ideas? Also, how can I also add a color bar (in addition to the size), with smaller circles being light red (for example) and bigger circle being blue (for example) here is my code so far:

import pandas as pd
import matplotlib.pyplot as plt
from vega_datasets import data as vega_data
gap = pd.read_json(vega_data.gapminder.url)

df = gap.loc[gap['year'] == 2000]

fig, ax = plt.subplots(1, 1,figsize=[14,12])
ax=ax.scatter(df['life_expect'], df['fertility'], 
            s = df['pop']/100000,alpha=0.7, edgecolor="black",cmap="viridis")

plt.xlabel("X")
plt.ylabel("Y");
kw = dict(prop="sizes", num=6, color="lightgrey", markeredgecolor='black',markeredgewidth=2)
plt.legend(*ax.legend_elements(**kw),bbox_to_anchor=(1, 0),frameon=False,
                    loc="lower left",markerscale=1,ncol=1,borderpad=2,labelspacing=4,handletextpad=2)

plt.grid()
plt.show()

enter image description here

1 Answers

It's a bit tricky, but you could measure the legend elements and reposition them to have a constant inbetween distance. Due to the pixel positioning, the plot can't be resized afterwards.

I tested the code inside PyCharm with the 'Qt5Agg' backend. And in a Jupyter notebook, both with %matplotlib inline and with %matplotlib notebook. I'm not sure whether it would work well in all environments.

Note that ax.scatter doesn't return an ax (countrary to e.g. sns.scatterplot) but a list of the created scatter dots.

import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.transforms import IdentityTransform
from vega_datasets import data as vega_data

gap = pd.read_json(vega_data.gapminder.url)
df = gap.loc[gap['year'] == 2000]

fig, ax = plt.subplots(1, 1, figsize=[14, 12])
fig.subplots_adjust(right=0.8)
scat = ax.scatter(df['life_expect'], df['fertility'],
                  s=df['pop'] / 100000, alpha=0.7, edgecolor="black", cmap="viridis")

plt.xlabel("X")
plt.ylabel("Y")
x = 1.1
y = 0.1
is_first = True
kw = dict(prop="sizes", num=6, color="lightgrey", markeredgecolor='black', markeredgewidth=2)
handles, labels = scat.legend_elements(**kw)

inverted_transData = ax.transData.inverted()
for handle, label in zip(handles[::-1], labels[::-1]):
    plt.setp(handle, clip_on=False)
    for _ in range(1 if is_first else 2):
        plt.setp(handle, transform=ax.transAxes)
        if is_first:
            xd, yd = x, y
        else:
            xd, yd = inverted_transData.transform((x, y))
        handle.set_xdata([xd])
        handle.set_ydata([yd])
        ax.add_artist(handle)
        bbox = handle.get_window_extent(fig.canvas.get_renderer())
        y += y - bbox.y0 + 15  # 15 pixels inbetween
    x = (bbox.x0 + bbox.x1) / 2
    if is_first:
        xd_text, _ = inverted_transData.transform((bbox.x1+10, y))
    ax.text(xd_text, yd, label, transform=ax.transAxes, ha='left', va='center')
    y = bbox.y1
    is_first = False
plt.show()

resulting plot

Related