在Python中从函数返回一个值并将其绘制在热图上

Question

我使用以下代码为给定函数生成了热图。这里是计算logistic映射变化的全局lyapunov指数-我已经添加了p的参数。这很好，我实际上不关心全局lyapunov指数的确切值是多少，而是关心它是正的还是负的。你知道吗

以下是我一直使用的代码：

# Logistic Function
def p_logistic(A, x, p):
    return (A/4) * (((p+1)**(p+1))/(p**p)) * (x**p) * (1-x)

# Lyapunov Exponent
def p_lyap(A, x, p):
    return np.log(abs((A/4) * (((p+1)**(p+1))/(p**p)) * (p*(x**(p-1)) - (p+1)*(x**p))))

n = 500
A = np.linspace(2, 4, n)
p = np.linspace(0.5, 5, n)

def F(A, p):
    A, p = np.meshgrid(A, p)
    lyapunov = 0
    x = 0.9
    N = 100

    for i in range(0,N):
        lyapunov = lyapunov + p_lyap(A, x, p)
        x = p_logistic(A, x, p)

    global_lyapunov = lyapunov/N
    return global_lyapunov

z = F(A, p)
plt.figure(figsize=(8,8))
xlabels = ['{:3.1f}'.format(x) for x in A]
ylabels = ['{:3.1f}'.format(y) for y in p]
ax = sns.heatmap(z, xticklabels = A, yticklabels = p)

我曾尝试使用if语句将全局lyapunov指数返回为1（如果它大于0）和-1（如果它小于0），但它返回了错误

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

基本上，我只希望在我的热图中有3种颜色，一种是全局lyapunov指数为负，等于0和正时的颜色。有没有可能？你知道吗

Answer 1

numpy.sign是获取所需数据的简单方法。另外，似乎您打算使用xlabel和ylabel：

sns.heatmap(np.sign(z), xticklabels=xlabel, yticklabels=ylabel)

Answer 2

假设我对你的理解正确，我会这样做：

zd = np.zeros_like(z, dtype=int)
zd[z > 0] = 1
zd[z < 0] = -1

然后绘制zd而不是z

另外请注意，最好直接从matplotlib使用imshow，这样可以获得更好的轴标签，例如：

plt.figure(figsize=(8,8))
plt.imshow(zd, extent=(2, 4, 5, 0.5), aspect='auto')
plt.colorbar()

imshow不会像heatmap那样对值进行任何重缩放，这可能更适合您的用例

此外，似乎还有很多未定义的值，因此可能需要对它们进行特殊处理，例如：

zd[~np.isfinite(z)] = 0

给出最终的绘图：

但我不确定我是否把斧子弄对了