时间戳

首页 > 解决方案 > 如何使用 if 条件在 1D 和 2D numpy 数组之间进行矢量化计算

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

问题描述

我有一个使用 1D 和 2D numpy 数组的计算。它有两个级别的条件if。我能够使用np.whereto 避免一个if-statement 并进一步使用慢速列表理解来遍历每一行。

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

这是我的代码：

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

标签： pythonnumpyconditional-statementsvectorization

解决方案

你的代码比你想象的更矢量化。除了对其进行矢量化之外，您还可以更适当地使用现有功能。

要生成整数范围，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. 第二个是逐元素r_row >= a条件：

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

[:, None]上的索引row_mask使其成为用于广播目的的列向量。现在，您可以使用直接掩码和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

推荐阅读