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

Answer 1

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()