Tatik
Tatik

Reputation: 1227

How to create a simple custom heatmap?

I want to create a Heatmap from pandas DataFrame df. The df has two columns:

name  test
aa    False
bb    False
cc    True
dd    False

The heatmap should include name values in X axis, and test values in Y axis. So, basically there should be 2 rows - False and True in Y axis. When test value is False, the False square is marked in red. When test value is True, the True square is marked in Green.

This is my current code:

import pandas as pd
import seaborn as sns

Index= df["name"].values
Cols = [True, False]
r = pd.DataFrame(df, index=Index, columns=Cols)
sns.heatmap(r, annot=True)

However this code does not work properly. How can I solve my task?

Upvotes: 2

Views: 2631

Answers (1)

Chris Adams
Chris Adams

Reputation: 18647

You need to reshape your DataFrame to plot it the way you want.

For example:

import matplotlib.pyplot as plt

df_new = df.set_index('name')['test'].astype(str).str.get_dummies().T

print(df_new)
name   aa  bb  cc  dd
False   1   1   0   1
True    0   0   1   0

Then you can use seaborn.heatmap:

plt.figure(figsize=(8, 4))
sns.heatmap(df_new, cmap='RdYlGn')

enter image description here

Upvotes: 2

Related Questions