Retrieve XY data from seaborn heatmap figure

Question

I ve searched the forum but could find anything.

My code is as follow:

def my_function(df):
    plt.figure()
    heatmap=sns.heatmap(df,cmap='coolwarm',)
    plt.title('title')
    plt.show()
    return heatmap

I woud like to retrieve the data from heatmap. I have seen how to do with matplotlib but i couldnt figure out how to do with sns/seaborn heatmap

Edit: the heatmap variable type is <class 'matplotlib.axes._subplots.AxesSubplot'>

Edit2: I know i can retrieve the data in dataframe but i want to unit test my function, that's why i try to retrieve the data in seaborn heatmap

Note that after `plt.show()` the image gets deleted, so, `return heatmap` (which in reality is an `ax` (a subplot), so you could name it `ax_heatmap` or so) doesn't make sense. If you just want the data, you should `return df` or `return df.to_numpy()`. — JohanC, Mar 28 '22 at 12:56
I know, but im actually trying to unit test the function with pytest. so i want to make sure the heatmap has the data i have put in it — , Mar 28 '22 at 14:05
You might check how seaborn does unit tests: https://github.com/mwaskom/seaborn/blob/master/seaborn/tests/test_matrix.py — JohanC, Mar 28 '22 at 14:24
thank you...they use p = mat._HeatMapper(self.df_norm, **self.default_kws) npt.assert_array_equal(p.plot_data, self.x_norm) but I dont know what is HeatMapper. I dont find plot_data in my seaborn object — , Mar 28 '22 at 14:35
Maybe `ax = sns.heatmap(....)` and then `ax.collections[0].get_array()`? As seaborn's test code shows, there is a huge amount of information inside a heatmap, and it is unclear what exactly you want to test. — JohanC, Mar 28 '22 at 14:47

Firefighting Physicist · Answer 1 · 2022-03-28T12:49:33.587

What exactly do you want? The XY data you are displaying is already available in your DataFrame. When you read the documentation you can see that the plot_data variable is just the values of the DataFrame that you put in.

If you are asking to get the color palette for the given values, take a look at https://stackoverflow.com/a/42126859/18482459

sns.heatmap only returns the axes, which allows you to manipulate the heatmap like so:

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

def my_function(df):
    plt.figure()
    heatmap=sns.heatmap(df,cmap='coolwarm',)
    plt.title('title')
    return heatmap

df = pd.DataFrame(np.random.rand(100).reshape(10,10))
heatmap = my_function(df)
heatmap.set_title("my_function call with df")
plt.show()

As it’s currently written, your answer is unclear. Please [edit] to add additional details that will help others understand how this addresses the question asked. You can find more information on how to write good answers [in the help center](/help/how-to-answer). — Community, Mar 28 '22 at 12:25

Retrieve XY data from seaborn heatmap figure

1 Answers1