我希望我的 Python 函数接受 numpy ndarrays

Question

我正在编写一个脚本来绘制一些热力学属性 Z = f(pr, Tr) 的 3D 表示。 pr 和 Tr 用[numpy.]arange() 创建，然后用[numpy.]meshgrid() 映射如下：

    Tr = arange(1.0, 2.6, 0.10)
    pr = arange(0.5, 9.0, 0.25)
    Xpr, YTr = meshgrid(pr, Tr)

然后将 Xpr 和 YTr 传递给计算上述属性的函数：

    Z = function(Xpr, YTr)

（“function”只是一个通用名称，稍后会替换为实际的函数名称）。

Z 中存储的值最终被绘制出来：

    fig = plt.figure(1, figsize=(7, 6))
    ax = fig.add_subplot(projection='3d')
    surf = ax.plot_surface(Xpr, YTr, Z, cmap=cm.jet, linewidth=0, antialiased=False)

当“功能”非常简单时，一切正常：

def zshell(pr_, Tr_):
    A = -0.101 - 0.36*Tr_ + 1.3868*sqrt(Tr_ - 0.919)
    B = 0.021 + 0.04275/(Tr_ - 0.65)
    E = 0.6222 - 0.224*Tr_
    F = 0.0657/(Tr_ - 0.85) - 0.037
    G = 0.32*exp(-19.53*(Tr_ - 1.0))
    D = 0.122*exp(-11.3*(Tr_ - 1.0))
    C = pr_*(E + F*pr_ + G*pr_**4)

    z = A + B*pr_ + (1.0 - A)*exp(-C) - D*(pr_/10.0)**4

    return z

但是当函数是这样的时候它会失败：

def zvdw(pr_, Tr_):
    A = 0.421875*pr_/Tr_**2     # 0.421875 = 27.0/64.0
    B = 0.125*pr_/Tr_           # 0.125 = 1.0/8.0

    z = 9.5e-01
    erro = 1.0

    while erro >= 1.0e-06:
        c2 = -(B + 1.0)
        c1 = A
        c0 = -A*B
        f = z**3 + c2*z**2 + c1*z + c0
        df = 3.0e0*z**2 + 2.0e0*c2*z + c1
        zf = z - f/df
        erro = abs((zf - z)/z)
        z = zf

    return z

我强烈怀疑失败是由 function zvdw(pr_, Tr_) 内部的迭代方法引起的（zvdw 之前已经过测试，并且在将浮点参数传递给它时效果很好）。这是我得到的错误信息：

Traceback (most recent call last):

  File "/home/fausto/.../TestMesh.py", line 81, in <module>
    Z = zvdw(Xpr, YTr)

  File "/home/fausto/.../TestMesh.py", line 63, in zvdw
    while erro >= 1.0e-06:

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

然而，错误消息似乎（在我看来）与while 语句没有直接关系。

有什么想法吗？

Answer 1

Tl;博士

为此有一个numpy 功能。你只需替换

def zvdw(pr_, Tr_):

通过

@np.vectorize
def zvdw(pr_, Tr_):

而且它有效。

速度更快

不幸的是，生成的图片看起来很难看，因为您的网格太稀疏了。我用TR 和pr 替换了你的步长。不幸的是，我们在这里遇到了np.vectorize 的限制。来自 numpy 文档

提供矢量化功能主要是为了方便，而不是为了 表现。该实现本质上是一个 for 循环。

即它非常缓慢。已经在 0.0010 和 0.0025 时花了 17.5 秒。因此，每小 10 倍是不现实的，因为这将花费大约 100 倍的时间。幸运的是，您的代码很简单，我可以使用 @numba.vectorize，这在我的机器上快了约 23 倍。

请注意，这只适用于某些 python 代码。 Numba 将 python 代码编译为 llvm 代码，因此它可以快速运行。但它无法为任意 python 代码做到这一点。见https://numba.pydata.org/numba-doc/dev/user/vectorize.html

这是你的照片

代码似乎不适用于np.vectorize/numba.vectorize

这很奇怪。这是在我的机器上工作的代码的文字复制粘贴：

import numpy as np
import matplotlib.pyplot as plt

Tr = np.arange(1.0, 2.6, 0.10)
pr = np.arange(0.5, 9.0, 0.25)
Xpr, YTr = np.meshgrid(pr, Tr)

def zshell(pr_, Tr_):
    A = -0.101 - 0.36*Tr_ + 1.3868*sqrt(Tr_ - 0.919)
    B = 0.021 + 0.04275/(Tr_ - 0.65)
    E = 0.6222 - 0.224*Tr_
    F = 0.0657/(Tr_ - 0.85) - 0.037
    G = 0.32*exp(-19.53*(Tr_ - 1.0))
    D = 0.122*exp(-11.3*(Tr_ - 1.0))
    C = pr_*(E + F*pr_ + G*pr_**4)

    z = A + B*pr_ + (1.0 - A)*exp(-C) - D*(pr_/10.0)**4

    return z

@np.vectorize
def zvdw(pr_, Tr_):
    A = 0.421875*pr_/Tr_**2     # 0.421875 = 27.0/64.0
    B = 0.125*pr_/Tr_           # 0.125 = 1.0/8.0

    z = 9.5e-01
    erro = 1.0

    while erro >= 1.0e-06:
        c2 = -(B + 1.0)
        c1 = A
        c0 = -A*B
        f = z**3 + c2*z**2 + c1*z + c0
        df = 3.0e0*z**2 + 2.0e0*c2*z + c1
        zf = z - f/df
        erro = abs((zf - z)/z)
        z = zf

    return z

Z = zvdw(Xpr, YTr)
fig = plt.figure(1, figsize=(7, 6))
ax = fig.add_subplot(projection='3d')
surf = ax.plot_surface(Xpr, YTr, Z, cmap=plt.cm.jet, linewidth=0, antialiased=False)