如何绘制 Python 中给定函数的热图

Question

我有一个名为 func(mu, gamma) 的函数。对于 mu 和 gamma 的每个组合，函数将 return 一个值，我们称它为 return_value。

现在我已经为 mu 和 gamma 设置了范围：

mu = np.linspace(0,1,100)
gamma = np.linspace(0,1,100)

现在我们有1e4个组合，每个组合对应一个return_value。我想绘制 return_value.

的热图

我尝试在 Python 中使用 pcolor。但是，从文档中的示例来看：

import matplotlib.pyplot as plt
import numpy as np

# make these smaller to increase the resolution
dx, dy = 0.15, 0.05

# generate 2 2d grids for the x & y bounds
y, x = np.mgrid[slice(-3, 3 + dy, dy),
                slice(-3, 3 + dx, dx)]
z = (1 - x / 2. + x ** 5 + y ** 3) * np.exp(-x ** 2 - y ** 2)
# x and y are bounds, so z should be the value *inside* those bounds.
# Therefore, remove the last value from the z array.
z = z[:-1, :-1]
z_min, z_max = -np.abs(z).max(), np.abs(z).max()

因为我的脚本中定义的函数 fun 不能将数组作为输入，所以它不起作用，如果我按照示例进行操作，我会收到此消息：

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

下面是我的 func 代码：

def fun(mu, gamma2):
    Isolation_Ratio = []
    kappa1 = gamma2
    kappa2 = gamma2                         
    gamma1 = gamma2
    g0 = gamma2 + kappa2 + gamma1 + kappa1
    gammag = kappa1/2. + gamma1/2.
    gamma = gamma2/2. + kappa2/2.
    for ii in range(len(rangedeltaw)):

        deltaw = rangedeltaw[ii]

        Forward_delta = forward_delta(mu, deltaw)
        Backward_delta = backward_delta(mu, deltaw)

        forward_root1, forward_root2, forward_root3 = forward_root(mu, deltaw)
        test_D, backward_root1, backward_root2, backward_root3 = backward_root(mu, deltaw)

        Root1.append(backward_root1)
        Root2.append(backward_root2)
        Root3.append(backward_root3)

        root1.append(forward_root1)
        root2.append(forward_root2)
        root3.append(forward_root3)
        if Forward_delta >= 0 and Backward_delta >= 0:
            a2sq = [max([forward_root1.real, forward_root2.real, forward_root3.real])]
            b1sq = [max([backward_root1.real, backward_root2.real, backward_root3.real])]
            A2sq.append(max([forward_root1.real, forward_root2.real, forward_root3.real]))
            B1sq.append(max([backward_root1.real, backward_root2.real, backward_root3.real]))
            for ii in range(len(a2sq)):
                for jj in range(len(b1sq)):

                    Isolation_Ratio.append(kappa2*a2sq[ii]/(kappa1*b1sq[jj]))
        elif Forward_delta >= 0 and Backward_delta < 0:
            a2sq = [max([forward_root1.real, forward_root2.real, forward_root3.real])]
            b1sq = [backward_root1.real]
            A2sq.append(max([forward_root1.real, forward_root2.real, forward_root3.real]))
            B1sq.append(backward_root1.real)
            for ii in range(len(a2sq)):
                for jj in range(len(b1sq)):
                    Isolation_Ratio.append(kappa2*a2sq[ii]/(kappa1*b1sq[jj]))

        elif Forward_delta < 0 and Backward_delta >= 0:
            a2sq = [forward_root1.real]
            b1sq = [max([backward_root1.real, backward_root2.real, backward_root3.real])]
            A2sq.append(forward_root1.real)
            B1sq.append(max([backward_root1.real, backward_root2.real, backward_root3.real]))
            for ii in range(len(a2sq)):
                for jj in range(len(b1sq)):
                    Isolation_Ratio.append(kappa2*a2sq[ii]/(kappa1*b1sq[jj]))

        else:   
            A2sq.append(forward_root1.real)
            B1sq.append(backward_root1.real)
            Isolation_Ratio.append(kappa2*forward_root1.real/(kappa1*backward_root1.real))   



    x = Isolation_RangeDeltaw
    y = Isolation_Ratio


    return max(y)

Answer 1

那么，首先，如何获取热图。 fun() 仍然不是自包含的（forward_delta() 等未定义），所以我无法执行它，并且您没有指定错误发生在哪一行，但我可以猜测违规者是

if Forward_delta >= 0 and Backward_delta >= 0:

意味着 forward_delta() 等函数与数组一起使用。从表面上看，完全向量化函数是可能的，但这是一项不平凡的任务，它本身就是一个问题（如果你问过它，请确保 self contained example)。一个更简单但效率较低的解决方案是按值填充热图值：

import matplotlib.pyplot as plt
import numpy

def fun(mu, gamma):
    # your function

mu = numpy.linspace(0,1,100)
gamma = numpy.linspace(0,1,100)

# filling the heatmap, value by value
fun_map = numpy.empty((mu.size, gamma.size))
for i in range(mu.size):
    for j in range(gamma.size):
        fun_map[i,j] = fun(mu[i], gamma[j])

现在你有了数组，问题的第二部分是如何绘制它。 pcolor() is used to visualize discreet arrays; imshow() 更适合您的目的：

fig = plt.figure()
s = fig.add_subplot(1, 1, 1, xlabel='$\gamma$', ylabel='$\mu$')
im = s.imshow(
    fun_map,
    extent=(gamma[0], gamma[-1], mu[0], mu[-1]),
    origin='lower')
fig.colorbar(im)
fig.savefig('t.png')

请注意，在数组中 X 维度在最后（对应于 gamma），但 imshow 将 X 维度放在第一位。

简单函数的结果

def fun(mu, gamma):
    return numpy.sin(mu) + numpy.cos(gamma)

看起来像

如何绘制 Python 中给定函数的热图

How to plot the heat map for a given function in Python

python

matplotlib

heatmap