Producing a heatmap from a pandas dataframe with rows of the form (x,y,z), where z is intended to be the heat value

Question:

Let’s say I have a dataframe with this structure

Example dataframe content

and I intend to transform it into something like this (done laboriously and quite manually):

Example of heatmap

Is there a simple call to (say) a seaborn or plotly function that would do this? Something like

heatmap(df, x='Dose', y='Distance', z='Passrate')

or perhaps a simple way of restructuring the dataframe to facilitate using sns.heatmap or plotly’s imshow, or similar? It seems strange to me that I cannot find a straightforward way of putting data formatted in this way into a high-level plotting function.

Asked By: Gavin Kirby

||

Answers:

Use df.pivot_table to get your data in the correct shape first.

Setup: create some random data

import pandas as pd
import numpy as np
import seaborn as sns

p_rate = np.arange(0,100)/np.arange(0,100).sum()

data = {'Dose': np.repeat(np.arange(0,3.5,0.5), 10),
        'Distance': np.tile(np.arange(0,3.5,0.5), 10),
        'Passrate': np.random.choice(np.arange(0,100), size=70,
                                     p=p_rate)}

df = pd.DataFrame(data)

Code: pivot and apply sns.heatmap

df_pivot = df.pivot_table(index='Distance', 
                          columns='Dose', 
                          values='Passrate', 
                          aggfunc='mean').sort_index(ascending=False)

sns.heatmap(df_pivot, annot=True, cmap='coolwarm')

Result:

example heatmap

Answered By: ouroboros1
Categories: questions Tags: , , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.