Scatterplot with regression fit and auto-positioned labels in Matplotlib

A custom scatterplot with an overlayed regression fit and auto-positioned labels to explore the relationship between the Corruption Perceptions Index and Human Development Index made with Python and Matplotlib. This post guides you through a beautiful and very informative scatterplot that includes a variety of custom colors, markers, and layout adjustments. The library adjustText is used to automatically adjust the position of labels in the plots and scikit-learn is used to obtain a regression fit.

About

This page showcases the work of Claus O. Wilke in his R package practicalgg that contains step-by-step examples demonstrating how to get the most out of ggplot2. You can find the original R code for this example here.

Thanks to him for accepting sharing his work here! Thanks also to Tomás Capretto who translated this work from R to Python! 🙏🙏

As a teaser, here is the plot we’re gonna try building:

scatterplot

Load libraries

In addition to Matplotlib, today's chart also uses adjustText to automatically adjust the position of labels in the plot, and scikit-learn to obtain a linear regression fit.

Note: the adjustText library can be installed with pip install adjustText

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from adjustText import adjust_text
from matplotlib.lines import Line2D # for the legend
from sklearn.linear_model import LinearRegression

Load and prepare data

Today's chart uses the corruption dataset in the practicalgg package. This data contains information about Corruption Perceptions Index (CPI) and Human Development Index (HDI) for 176 countries, from 2012 to 2015.

The original source are the Corruption Perceptions Index 2016 released by Transparency International and the Human Development Index made available in the Human Development Reports by the United Nations Development Programme. These datasets were merged and made available by Claus O. Wilke as the corruption dataset in his practicalgg package. Thanks to Claus for all the work and making this possible!

corruption = pd.read_csv("https://raw.githubusercontent.com/holtzy/The-Python-Graph-Gallery/master/static/data/corruption.csv")

The next step is to keep only observations for the 2015 year and drop any row that contains a missing value.

corrupt = corruption.query("year == 2015").dropna()
corrupt.head()

	country	region	year	cpi	iso3c	hdi
0	Denmark	Europe and Central Asia	2015	91.0	DNK	0.925
1	New Zealand	Asia Pacific	2015	91.0	NZL	0.915
2	Finland	Europe and Central Asia	2015	90.0	FIN	0.895
3	Sweden	Europe and Central Asia	2015	89.0	SWE	0.913
4	Switzerland	Europe and Central Asia	2015	86.0	CHE	0.939

Basic Scatterplot

Let's get started by creating a very basic scatterplot to show the relationship between the Corruption Perceptions Index and the Human Development Index.

CPI = corrupt["cpi"].values
HDI = corrupt["hdi"].values

fig, ax = plt.subplots(figsize=(12, 8));
ax.scatter(CPI, HDI);

It's also possible to use ax.scatter("cpi", "hdi", data=corrupt)

Albeit being simple, the plot above is enough to see there's a positive association between the two indexes under study.

But this is just the start, let's make it better!

Customize marker colors

The next step is to use a different color for each region so it is possible to see the association between the region, the Corrupt Perceptions Index, and the Human Development Index.

As its name says, the following function can be used to adjust the lightness of a color. The amount argument controls this adjustment. Values below 1 mean the color darkens, while values above 1 indicate the color is lightened.

def adjust_lightness(color, amount=0.5):
    import matplotlib.colors as mc
    import colorsys
    try:
        c = mc.cnames[color]
    except:
        c = color
    c = colorsys.rgb_to_hls(*mc.to_rgb(c))
    return colorsys.hls_to_rgb(c[0], c[1] * amount, c[2])

Now let's select the colors for fill and the edges:

# Okabe Ito colors
REGION_COLS = ["#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2"]

# Category values for the colors
CATEGORY_CODES = pd.Categorical(corrupt["region"]).codes

# Select colors for each region according to its category.
COLORS = np.array(REGION_COLS)[CATEGORY_CODES]

# Compute colors for the edges: simply darker versions of the original colors
EDGECOLORS = [adjust_lightness(color, 0.6) for color in COLORS]

And make a new plot:

fig, ax = plt.subplots(figsize=(12, 8));
ax.scatter(
    CPI, HDI, color=COLORS, edgecolors=EDGECOLORS,
    s=80, alpha=0.5, zorder=10
);
# zorder = 10 is used to make sure markers are on top of the regression line added later