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

Question

我有一个名为的函数func(mu, gamma)。对于每个组合mu和gamma，该函数会返回一个值，让我们把它return_value。

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

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

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

我尝试pcolor在Python中使用。但是，根据文档中的示例：

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()etc函数可用于数组。从外观上看，可能可以对功能进行完全向量化，但这是一项艰巨的任务，并且是其自身的问题（如果您提出要求，请确保制作一个独立的示例）。一个简单但有效的解决方案是按值填充热图值：

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()用于可视化离散数组；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)

看起来像