matplotlib 等高线图中的条件函数变化

Question

我想使用 2 变量函数 Z(X,Y) 生成等高线图。但是，我想强加一个条件，当 X 比某个值 less/greater 时改变函数。这样的更改将允许我仅使用单个 plt.contour 行（即我不想创建两个单独定义的函数，这导致使用两个单独的绘图命令行） .
我继续运行进入（我认为是）truth/logic 错误。我的猜测是 numpy meshgrid 的某些方面不符合函数的条件“开关”。下面附上的是显示概念的短代码，以及完整的 Traceback 错误。如果有什么不清楚的地方，或者我提供的内容不足以解释我的问题，欢迎在下方评论。

import numpy as np
import matplotlib.pyplot as plt


X = np.linspace(0,50,100)
Y = np.linspace(0,50,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()```

Answer 1

一切正常，但您必须更改 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,50,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()

Answer 2

numpy 数组中的真实性测试是元素明智的，

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

输出: [真真假假]

您可以使用 .all() 或 .any() 取决于您是否需要 all 或任何的元素满足条件

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

输出：假

Answer 3

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

一个条件作为第一个参数
如果该条件作为第二个参数为真，可供选择的值
如果该条件作为第三个参数为假，从中选择的值

这样，你的z()函数就变成了：

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

生成的情节：

matplotlib 等高线图中的条件函数变化

Conditional function change in matplotlib's contour plot

python

numpy

matplotlib

conditional-statements

contourf