I want my Python functions to accept numpy ndarrays

Question

I'm writing a script to plot a 3D representation of some thermodynamic property, Z = f(pr, Tr). pr and Tr are created with [numpy.]arange() and then they are mapped with [numpy.]meshgrid() the following way:

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

Xpr and YTr are then passed to the function that calculates the aforementioned property:

    Z = function(Xpr, YTr)

("function" is just a generic name that is later replaced by the actual function name).

The values stored in Z are finally plotted:

    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)

Everything works fine when "function" is something quite straightforward like:

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

But it fails when the function is something like this:

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

I strongly suspect that the failure is caused by the iterative method inside function zvdw(pr_, Tr_) (zvdw has been previously tested and works perfectly well when float arguments are passed to it). That is the error message I get:

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()

Yet the error message doesn't seem (to me) to be directly related to the while statement.

Any ideas?

zf and z are 2 dimensional arrays inside while loop. after erro = abs((zf - z)/z) erro is also 2 dimensional. after one loop erro >= 1.0e-06 will be a 2 dimensional array with truth values, possibly both True and False. Thats were the error message comes from. What should be the condition to end the while loop? Every value equals False, most of them or is one enough? — Michael Szczesny
– Michael Szczesny, Commented Aug 29, 2021 at 22:53
@MichaelSzczesny, thanks for the enlightenment. I would say that one False should be enough, but maybe I'm being too hasty in my judgement. What should I do, supposing I'm right? — Fausto Arinos Barbuto
– Fausto Arinos Barbuto, Commented Aug 29, 2021 at 23:15
while all(erro >= 1.0e-06): where all should be the numpy.all function, not the python function. If you want to work with numpy you should look into np.all and np.any. To avoid confusion you should use the common way to import numpy with import numpy as np. — Michael Szczesny
– Michael Szczesny, Commented Aug 29, 2021 at 23:25
@MichaelSzczesny, that worked, and while (erro >= 1.0e-06).all(): also works. But first I had to redefine the initial value of errofrom erro = 1.0 to erro = array(1.0). Thanks a lot for your help! — Fausto Arinos Barbuto
– Fausto Arinos Barbuto, Commented Aug 30, 2021 at 1:23

Lukas S · Accepted Answer · 2021-08-30 14:29:03Z

Tl;dr

There is a numpy functionality for that. You just replace

def zvdw(pr_, Tr_):

by

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

And it works.

Get it faster

Unfortunately the resulting picture looked ugly since your mesh is to sparse. I replaced your step size in TR and pr. Unfortunately here we run into the limitation of np.vectorize. From the numpy documentation

The vectorize function is provided primarily for convenience, not for performance. The implementation is essentially a for loop.

I.e. it is painfully slow. Already at 0.0010 and 0.0025 it took 17.5 seconds. So it is not realistic to go a factor of 10 each smaller since that would take ~100 times longer. Fortunately your code is simple enough that I can use @numba.vectorize which on my machine was a factor of ~23 faster.

Notice that this only works for some python code. Numba compiles the python code to llvm code so it can run fast. But it is unable to do that for arbitrary python code. See https://numba.pydata.org/numba-doc/dev/user/vectorize.html

And here is your picture

Code seems not to work with `np.vectorize/numba.vectorize`

That is very odd. Here is a literal copy paste of code that works on my machine:

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)

Collectives™ on Stack Overflow

I want my Python functions to accept numpy ndarrays

1 Answer 1

Tl;dr

Get it faster

Code seems not to work with `np.vectorize/numba.vectorize`

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Tl;dr

Get it faster

Code seems not to work with np.vectorize/numba.vectorize

Comments

Your Answer

Sign up or log in

Post as a guest

Related

Code seems not to work with `np.vectorize/numba.vectorize`