Quickly creating heatmaps with pandas

panda’s style functionality is often a good enough replacement for a more sophisticated plotting library when all we want is a simple heatmap of two-dimensional data.

Let’s say we have this fake time series data, were for each year and month, multiple measurements are recorded.

import pandas as pd
import numpy as np

n = 1_000
df = pd.DataFrame(
    {
        "year": np.random.randint(2000, 2020, n),
        "month": np.random.randint(1, 13, n),
        "measurement": np.random.randn(n),
    }
)

We want to create a heatmap, showing the averages of each year-month-combination.

Let’s aggregate measurements from the same months and years and then create a pivot table.

df_pivot = (
    df.groupby(["year", "month"])
    .sum()
    .reset_index()
    .pivot(index="year", columns="month", values="measurement")
    .reset_index()
)

Now, we can use .style.background_gradient to color the output. Additionally, let’s use .format to only show one digit after the decimal point.

df_pivot.style.format(
    precision=1,
).background_gradient(
    cmap="RdBu",
    vmin=-1,
    vmax=1,
    subset=[i for i in range(1, 13)],
)

month	year	1	2	3	4	5	6	7	8	9	10	11	12
0	2000	-3.1	-1.4	-2.0	-0.8	-0.9	0.1	3.6	-2.8	1.4	-0.3	0.8	1.0
1	2001	-2.2	1.5	0.1	0.2	1.0	1.5	4.0	-0.4	3.5	0.7	2.6	nan
2	2002	0.7	-0.9	-1.4	-2.3	3.0	2.5	0.8	-0.0	3.5	0.7	-1.7	-2.1
3	2003	3.7	-1.4	-0.0	0.3	1.2	0.6	3.2	-2.7	-1.1	-3.1	-1.4	4.4
4	2004	-1.9	0.0	0.2	1.9	0.0	0.3	3.8	3.2	0.1	-1.6	2.5	0.2
5	2005	1.5	0.1	-0.8	-2.6	-1.4	0.9	-0.2	0.4	1.3	-1.8	-1.7	-3.1
6	2006	1.6	-2.8	0.3	1.0	5.1	3.2	2.7	2.6	1.9	1.3	-0.3	-2.1
7	2007	-1.4	3.0	-0.6	-1.2	1.4	1.3	1.3	0.7	-1.7	-0.0	0.8	-0.6
8	2008	4.2	1.7	-2.8	0.8	-1.1	-1.2	0.1	-0.5	-2.3	nan	-0.6	1.3
9	2009	2.2	-0.3	-3.2	-4.9	-3.3	4.9	4.7	0.3	nan	1.4	-0.5	0.4
10	2010	1.6	0.4	-2.9	2.2	-2.5	0.4	2.4	2.0	3.9	nan	-2.3	1.3
11	2011	3.7	-2.1	2.2	1.0	-1.3	-2.0	-0.7	0.2	2.4	-1.9	-3.4	-3.8
12	2012	-1.2	4.0	1.2	-3.7	0.8	-1.6	4.9	0.0	2.6	-1.1	-1.5	1.9
13	2013	-0.4	-1.7	1.1	1.0	-2.4	-2.3	-0.9	-0.5	0.5	-0.4	-2.6	-2.2
14	2014	-0.9	0.2	-2.8	1.1	1.6	2.1	1.9	-1.9	1.9	-4.6	0.9	0.3
15	2015	2.6	-0.7	-0.3	0.4	1.0	-1.3	-2.4	1.0	1.1	0.3	3.3	0.2
16	2016	3.1	-3.3	-3.8	1.2	-1.8	-0.4	-0.5	1.1	1.9	1.0	-0.2	1.9
17	2017	-0.8	-2.6	3.8	-0.4	-1.5	3.2	-2.4	-1.3	0.1	1.0	-0.3	-1.7
18	2018	nan	-0.1	-1.3	2.7	-0.1	-1.1	-1.3	1.1	1.5	-1.3	1.5	1.9
19	2019	-0.2	3.9	-0.9	-2.3	-1.6	-2.4	1.5	-0.2	0.5	-0.3	-0.3	-2.6

That’s already good enough in a lot of cases :)