在执行 numpy 广播操作时应用元素规则

Question

我想将两个数字 numpy 对象 t 和 speed 相乘，但事先不知道每个对象是标量还是数组。问题是 0 是 t（或 t 的元素）的合法值，而 inf 是 speed（或 [= 的元素）的合法值12=]).我正在使用的模型对这种情况有一个规则：只要速度是无限的并且 t 是零，乘积就被定义为 0.

我的问题是在应用规则的同时处理两个操作数的可能标量可能不是。到目前为止，我能想到的最好的就是这个丑陋的级联：

if t.size == 1:
    if t.flat[0]:
        t = t * speed
else:
    nonzero = t != 0
    if speed.size == 1:
        t[nonzero] *= speed
    else:
        t[nonzero] *= speed[nonzero]

我忍不住认为必须有一种我所缺少的更高效、更 Numpythonic 的方式。有吗？

Answer 1

0 * np.inf 抛出一个 RuntimeWarning 并评估为 np.nan 只做乘法，然后替换 np.nan 怎么样？

import numpy as np

def multiply(a, b):
    temp = np.array(a*b)
    temp[np.isnan(temp)] = 0
    return temp

测试：

np.random.seed(123)
t = np.random.rand(10)
t[0] = 0
speed = np.random.rand(10)
speed[0] = np.inf
speed[6] = np.inf

输入：

multiply(t, speed)

输出：

array([ 0.        ,  0.09201602,  0.02440781,  0.49796297,  0.06170694,
    0.57372377,         inf,  0.03658583,  0.53141369,  0.1746193 ])

Answer 2

几天过去了，hpaulj和SubhaneilLahiri的评论里的答案都没有人写出来，还是自己写吧。我将此作为公认答案的原因是，即使在超出问题中给出的特定模型的情况下，它也足以解决问题的标题。

# prepare zeros of the correct shape:
out = np.zeros(np.broadcast(t, speed).shape), dtype=float)

# write the product values into the prepared array, but ONLY
# according to the mask t!=0 
np.multiply(t, speed, where=t!=0, out=out)

在问题的特定模型中，pre-prepared 默认输出值 0 对于测试 t!=0 失败的所有地方已经正确：0 * speed 是 0 无论如何，只要 speed 是有限的，当 speed 是无限的时候，它也是 0（根据模型的规则）。但这种方法适用于更一般的情况：原则上，可以使用额外的 numpy ufunc 调用来使用任意不同规则的结果填充输出数组的 masked-out 部分。

在执行 numpy 广播操作时应用元素规则

applying an elementwise rule while performing a numpy broadcasting operation

python

numpy

array-broadcasting