matplotlib轮廓图中的条件函数变化

编程问答 2022-06-13

问题描述

我想使用2变量函数Z（X，Y）生成轮廓图。但是，我想施加一个条件，当X小于/大于某个特定值时，该条件会改变函数。这样的更改将允许我仅使用单条plt.contour行（即我不想想要创建两个单独定义的函数，从而导致使用两个单独的绘图命令行）。
我继续遇到（我相信是）真理/逻辑错误。我的猜测是numpy meshgrid的某些方面不符合该功能的条件“开关”。下面随附的是显示概念的简短代码，以及完整的Traceback错误。如果有任何不清楚的地方，或者我提供的内容不足以解释我的问题，请随时在下面发表评论。

import numpy as np
import matplotlib.pyplot as plt


X = np.linspace(0,50,100)
Y = np.linspace(0,100)
X,Y = np.meshgrid(X,Y)


def z(x,y):

    if x < 20:
        return np.sin(x) + np.cos(y)
    
    else:
        return np.tan(x * y)


Z = z(X,Y)


plt.contourf(X,Y,Z)
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar()

ValueError                                Traceback (most recent call last)
<ipython-input-29-7e200be093e6> in <module>
     16 
     17 
---> 18 Z = z(X,Y)
     19 
     20 plt.figure(figsize=(8,6))

<ipython-input-29-7e200be093e6> in z(x,y)
      9 
     10 def z(x,y):
---> 11     if x < 20:
     12         return np.sin(x) + np.cos(y)
     13 

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

解决方法

您可以简单地使用numpy.where()。
您必须通过：

  • 条件作为第一个参数
  • 该条件为真时从中选择的值作为第二个参数
  • 该条件为假时从中选择的值作为第三个参数

这样，您的z()函数将变为：

def z(x,y):
    return np.where(x < 20,np.sin(x) + np.cos(y),np.tan(x*y))

生成的情节：

enter image description here

,

一切都很好，但是您必须更改np.any(x<20)上的比较。这意味着，如果x的任何元素大于20。您也可以使用np.all，只要数组的每个元素都满足条件就可以这样做

import numpy as np
import matplotlib.pyplot as plt


X = np.linspace(0,50,100)
Y = np.linspace(0,100)
X,Y = np.meshgrid(X,Y)


def z(x,y):

    if np.any(x < 20):
        return np.sin(x) + np.cos(y)
    
    else:
        return np.tan(x * y)


Z = z(X,Y)


plt.contourf(X,Y,Z)
plt.xlabel('x')
plt.ylabel('y')
plt.colorbar()
,

在numpy数组中进行真实测试是明智的，

import numpy as np
X =  np.array([1,2,3,4])
print(X<3)

输出：[True True False False]

您可以使用 .all（）或 .any（），具体取决于您要全部还是任何满足条件的元素

print((X<3).all())

输出：错误

conditional-statementscontourfmatplotlibnumpypython