如何优化分段连续、多参数、包含条件语句的 Python 函数以对 NumPy 数组进行操作？

Question

我对 Python 编程和 NumPy 比较陌生，但我有工程背景并且广泛使用了 VBA。 我目前正在使用 Python 来分析气体动力学问题，这些问题涉及数学上复杂且有时不连续的函数（冲击波和过渡流）。 我需要有关分段连续函数编程的帮助。

这是我正在做的一个方面的简化版本。 我有两个变量的分段连续函数：

f(Re, r) = fL(Re, r), Re<2000 for all r

f(Re, r) = fT(Re, r)，对于所有 r，Re>= 2000

我测试了计算 f 的各种选项，包括循环、向量化标量函数、numpy.where()、numpy.select()、numpy.piecewise() 和各种形式的布尔索引。 numpy.piecewise() 似乎提供了最优雅的解决方案，但我不知道如何使用两个参数。 希望下面的代码能够清楚地说明这一点；非常感谢任何帮助或评论。

有没有办法以元素方式实现带有两个参数的 numpy.piecewise() 函数？

我的布尔索引解决方案是否低于合理的解决方案，注意到数组可能会变得相当大？

import numpy as np

def fL(Re, r):
    return 16.0/Re

def fT(Re,r):
    return ((-1.8*np.log10((r/3.7)**1.11 + 6.9/Re))**-2.0)/4.0

Re = np.logspace(3., 7., 5)

# This works fine for a scalar value of r:
r = 0.0001
def f0(Re, r):
    return np.piecewise(Re, [Re<2000.0, Re>=2000.0], [fL, fT], r)
print('\n       Re: ', Re)
print('        r: ', r)
print('f0(Re, r): ', f0(Re, r))

# This works element-wise for arrays of Re and r (which is what I want):
r = np.full_like(Re, 0.0001)
def f1(Re, r):
    fct = np.zeros(len(Re))
    fct[Re < 2000.0] = fL(Re[Re < 2000.0], r[Re < 2000.0])
    fct[Re >= 2000.0] = f0(Re[Re >= 2000.0], r[Re >= 2000.0])
    return fct
print('\n       Re: ', Re)
print('        r: ', r)
print('f1(Re, r): ', f1(Re, r))

# but this does not work element-wise:
print('\n      Re: ', Re)
print('       r: ', r)
print('f(Re, r): ', f0(Re, r))

Answer 1

一般来说，我建议使用最简洁地描述您想要执行的操作的

numpy

函数。在这种情况下，我会选择

out = np.where(
    Re < 2000.0,
    fL(Re, r),
    fT(Re, r)
)

如果你有更多的条件，我会去

np.select

。

无论哪种情况，我都会接受我们正在计算

fL

和

fT

至少部分在它们未定义的域上，如果它“足够快”。

np.piecewise

，内部执行的操作与您已经使用布尔索引解决方案执行的操作完全相同，只是更通用一些。问题是

np.piecewise

干净地将完整的

r

传递到

fL

和

fT

，但仅将条件成立的

Re

值的子集传递。

我想不出任何方法可以让

np.piecewise

解决您的问题，所以在这种情况下，您的

f1

是您最好的选择。

话虽这么说，如果性能是您的瓶颈，那么通过例如某种（即时）编译函数可能会为您提供更好的服务。

numba