python - 如何在 python 中从三个列表创建热图？

Question

我有 3 个数据数组:initial_R、initial_v 和 E_total。

   import numpy as np
   import matplotlib
   import matplotlib as mpl
   import matplotlib.pyplot as plt
   import seaborn as sns
   import pandas as pd
   
   initial_R=[]
   initial_v=[]
   E_total=[]

然后我有一段代码来填充这三个数组。我不会将其发布在这里，因为它很长，但它确实有效。

initial_R 是 x 坐标，initial_v 是 y 坐标，E_total 是这些坐标处的强度值。我想将其绘制为一个简单的热图，每个方 block 都着色为两种颜色之一，具体取决于 E_total 是正数还是负数。

    df = pd.read_clipboard()
    df.drop_duplicates(['initial_R','initial_v'], inplace=True)
    pivot = df.pivot(index='initial_R', columns='initial_v', values='E_total')
    ax = sns.heatmap(pivot,annot=True)
    plt.show()

    print (pivot)

我以前见过上面的代码，但它不适用于我的数据，说实话我不确定它是如何工作的。

Answer 1

您可以使用numpy.sign来获取值的符号:

import pandas as pd
import seaborn as sns

df = pd.DataFrame({'initial_R': [1, 2, 3, 1, 3, 1, 2],
                   'initial_v': [0, 0, 0, 1, 1, 2, 2],
                   'E_total': [-1, 2, -4, -3, 0, -1, 5]})

df.drop_duplicates(['initial_R','initial_v'], inplace=True)

pivot = df.pivot(index='initial_R', columns='initial_v', values='E_total')

sns.heatmap(np.sign(pivot), annot=True)

如果您想组合 0 和正值:

sns.heatmap(pivot.ge(0).astype(int).mask(pivot.isna()), annot=True)

输出1:

输出2:

如果您想要自定义标签:

pivot2 = pivot.ge(0).astype(int).mask(pivot.isna())

sns.heatmap(pivot2, annot=pivot2.replace({0: 'negative', 1: 'positive'}), fmt='s')

输出: