python - 以 3D 形式绘制分段函数

Question

我是 python 新手，正在尝试 3d 绘制分段函数。我正在尝试在 z 轴上绘制下面的 3d 'mainformula' 函数，因为它随着 x 和 y 的变化而变化，范围从 0 到 10，常数 = 1。但我似乎不太清楚这里的绘图方法。

from sympy import *
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib.pyplot as plt
import numpy as np

def mainformula(x,y,constant):
    return Piecewise((subformula1(x,y,constant), y >= 0 and y < 3),(subformula2(x,y,constant),y>=3 and y <= 10))

def subformula1(x,y,constant):
    return x + y + constant

def subformula2(x,y,constant):       
    return x - y - constant

fig = plt.figure()
ax = fig.gca(projection='3d')
X = np.arange(0, 10, 0.25)
Y = np.arange(0, 10, 0.25)
constant = 1
X, Y = np.meshgrid(X, Y)
Z = mainformula(X,Y,constant)
surf = ax.plot_surface(X, Y, Z)
plt.show()

运行该代码时出现的错误是:“ValueError:具有多个元素的数组的真值不明确。使用 a.any() 或 a.all()”

Answer 1

您正在处理数组，因此在 bool 上下文中使用 array > 3 永远不会起作用(例如使用 and)，这将始终给您带来错误你收到。但您始终可以将条件定义为 bool 掩码，并在适当的元素上运行公式:

def mainformula(x,y,constant):
    z = np.zeros_like(x)
    # Condition 1 indexes all elements where subformula 1 is valid
    condition1 = np.logical_and(y >= 0, y < 3)
    # condition1 = (y >= 0) & (y < 3) # is another way of writing it
    z[condition1] = x[condition1] + y[condition1] + constant
    # now do it in the range where subformula 2 is valid
    condition2 = np.logical_and(y >= 3, y <= 10)
    # condition1 = (y >= 3) & (y <= 10) # is another way of writing it
    z[condition2] = x[condition2] - y[condition2] - constant

    return z

这不使用 sympy.Piecewise ，但仅在尝试绘图时工作正常。如果您想要单独的函数而不是在主公式中完成所有操作，您需要对其进行一些更改:

z[condition1] = subformula1(x[condition1], y[condition1], constant)

与条件2类似。