Correlation Heatmaps with Significance in Python

While Pandas and Seaborn offer very quick ways to calculate correlations and show them in a heatmap. Whether those
correlations are statistically significant or not is omitted from those plots. Over the years I collected bits and
pieces of code, like this, that turn out to be quite useful. Though them being scattered across a few dozen projects
isn鈥檛 very convenient when I actually need them. So I’ll start to add some documentation and put them here with the tag
Code Nugget, so they can easily be found by myself and others.

Normally you can use corr_df = df.corr() to get a correlation matrix for numerical columns in a Pandas data frame.
These in turn can be shown in a heatmap using sns.clustermap(corr_df, cmap="vlag", vmin=-1, vmax=1), leveraging
SeaBorn clustermap. Easy, though the significance of those correlations isn’t reported. To get those you can rely on built-in functions and a bit more effort is required.

from sklearn.datasets import load_iris
from scipy.stats import spearmanr
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from statsmodels.stats.multitest import multipletests

iris_obj = load_iris()
iris_df = pd.DataFrame(iris_obj.data, columns=iris_obj.feature_names)


def get_correlations(df):
    df = df.dropna()._get_numeric_data()
    dfcols = pd.DataFrame(columns=df.columns)
    pvalues = dfcols.transpose().join(dfcols, how="outer")
    correlations = dfcols.transpose().join(dfcols, how="outer")
    for ix, r in enumerate(df.columns):
        for jx, c in enumerate(df.columns):
            sp = spearmanr(df[r], df[c])
            correlations[c][r] = sp[0]
            pvalues[c][r] = sp[1] if ix > jx else np.nan  # Only store values below the diagonal
    return correlations.astype("float"), pvalues.astype("float")


correlations, uncorrected_p_values = get_correlations(iris_df)

# Correct p-values for multiple testing and check significance (True if the corrected p-value < 0.05)
shape = uncorrected_p_values.values.shape
significant_matrix = multipletests(uncorrected_p_values.values.flatten())[0].reshape(
    shape
)

# Here we start plotting
g = sns.clustermap(correlations, cmap="vlag", vmin=-1, vmax=1)

# Here labels on the y-axis are rotated
for tick in g.ax_heatmap.get_yticklabels():
    tick.set_rotation(0)

# Here we add asterisks onto cells with signficant correlations
for i, ix in enumerate(g.dendrogram_row.reordered_ind):
    for j, jx in enumerate(g.dendrogram_row.reordered_ind):
        if i != j:
            text = g.ax_heatmap.text(
                j + 0.5,
                i + 0.5,
                "*" if significant_matrix[ix, jx] or significant_matrix[jx, ix] else "",
                ha="center",
                va="center",
                color="black",
            )
            text.set_fontsize(20)

# Save a high-res copy of the image to disk
plt.tight_layout()
plt.savefig("clustermap.png", dpi=200)

In this example we’ll load the iris dataset and convert it to a Pandas data frame, next a new function get_correlations
is defined that will return two new dataframes, one with the correlations (here spearman rank is used, see below) and
another one with the p-values for those correlations. Note we don’t store p-values for
combinations we don’t want to test (values on the diagonal) or don’t need to test (correlations are symmetrical, only
values below the diagonal are stored). Including these would make corrections for multiple testing unnecessarily harsh.

While we have p-values for all those values, shown below, these are not corrected for multiple testing. The function
multipletests from the statsmodels package can correct these for us and report which ones are significant
(default cutoff <0.05), but the function requires a flat list of values. So we convert the matrix to a one-dimensional
array, apply the function and transform it back to the original shape using reshape.

Finally, the correlations need to be drawn and the clustermap function is great here. Though we need a few
extra lines of code to put an asterisk in the cells which are significant. While it is exactly rocket science how
this is done, it did require a fair bit of digging in the clustermap code to find exactly how to hook this in.
There are a ton of tweaks that could still be done here, but this will depend on your personal style and preference. The
hard part is done! Check out the result below!