如何使用 if 条件在 1D 和 2D numpy 数组之间进行矢量化计算

Question

我使用一维和二维 numpy 数组进行计算。它有两个级别的 if 条件。我能够使用 np.where 来避免一个 if 语句，并进一步使用慢速列表理解来遍历每一行。

理想情况下，我想将整个计算过程向量化。可能吗？

这是我的代码：

import numpy as np

r_base = np.linspace(0, 4, 5)
np.random.seed(0)
r_mat = np.array([r_base * np.random.uniform(0.9, 1.1, 5),
                  r_base * np.random.uniform(0.9, 1.1, 5),
                  r_base * np.random.uniform(0.9, 1.1, 5)])

a_array = np.linspace(1, 3, 3)

def func_vetorized_level1(r_row, a):
    if r_row.mean() > 2:
        result = np.where((r_row >= a), r_row - a, np.nan)
    else:
        result = np.where((r_row >= a), r_row + a, 0)
    return result
# try to broadcast this func to every row of r_mat using list comprehension
res_mat = np.array([func_vetorized_level1(this_r_row, this_a) 
                    for this_r_row, this_a in zip(r_mat, a_array)])

结果是

res_mat =
array([[       nan, 0.04303787, 1.04110535, 2.02692991, 2.93892384],
       [       nan,        nan, 0.1567092 , 1.27819766, 1.90675322],
       [0.        , 0.        , 0.        , 6.25535798, 6.65682885]])

Answer 1

您的代码比您想象的更可向量化。除了对其进行vectoring，还可以更恰当地利用已有的函数

要生成整数范围，np.arange 比 np.linspace 效果更好：

r_base = np.arange(5.)
a_array = np.arange(1., 4.)

一次乘法即可生成随机数：

np.random.seed(0)
r_mat = r_base * np.random.uniform(0.9, 1.1, (3, 5))

我认为最简单的做法是制作一个输出数组并根据不同的条件填充它：

out = np.empty_like(r_mat)

将 a_array 设为与 r_mat 中的行数相匹配的列会很有帮助：

a = a_array[:, None]

接下来你要做的就是为条件制作面具。第一个是 r_row.mean() > 2 的 row-wise 掩码。第二个是element-wiser_row >= a条件：

row_mask = (r_mat.mean(axis=1) > 2)[:, None]
elem_mask = r_mat >= a

row_mask 上的索引 [:, None] 使其成为用于广播目的的列向量。现在您可以使用直接掩码和 where 关键字评估选择到适当的 ufuncs:

np.subtract(r_mat, a, out=out, where=row_mask & elem_mask)
np.add(r_mat, a, out=out, where=~row_mask & elem_mask)
out[row_mask & ~elem_mask] = np.nan
out[~row_mask & ~elem_mask] = 0

如何使用 if 条件在 1D 和 2D numpy 数组之间进行矢量化计算

How to have vectorize calculation between a 1D and 2D numpy array with if conditions

python

numpy

vectorization

conditional-statements