seaborn heatmap from pandas dataframe with NaNs

Question

Hi I really want to create a heatmap but am struggling:

# correlations between undergrad studies and occupation
data_uni = n.groupby(['Q5','Q6'])['Q6'].count().to_frame(name = 'count').reset_index()
# some participants did not answer the question in the survey
data_uni.fillna('Unknown', inplace=True)

data_uni.pivot(index='Q5', columns='Q6', values='count')
plt.figure(1, figsize=(14,10))
sns.heatmap(data_uni, cmap="YlGnBu")

The error message I get is "TypeError: ufunc 'isnan' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''".

Is this the right way to create a heatmap? If yes, what am I doing wrong, and if not, what would be the right way? Thank you for your help!

Answer 1

As per issue GH375, you can specify a mask, where data will not be shown for those cells whose mask values are True.

sns.heatmap(data_uni, cmap="YlGnBu", mask=data_uni.isnull())