python - 如何使用不同大小的向量对操作进行向量化

Question

我有不同大小的向量，想进行元素方面的操作。如何在 Python 中优化以下 for 循环？ (例如使用 np.vectorize())

import numpy as np

n = 1000000

vec1 = np.random.rand(n)
vec2 = np.random.rand(3*n)
vec3 = np.random.rand(3*n)

for i in range(len(vec1)):
    if vec1[i] < 0.5:
        vec2[3*i : 3*(i+1)] = vec1[i]*vec3[3*i : 3*(i+1)]
    else:
        vec2[3*i : 3*(i+1)] = [0,0,0]

非常感谢您的帮助。

Answer 1

我们可以利用 broadcasting -

v = vec3.reshape(-1,3)*vec1[:,None]
m = vec1<0.5
vec2_out = (v*m[:,None]).ravel()

另一种表达方式是——

mask = vec1<0.5
vec2_out = (vec3.reshape(-1,3)*(vec1*mask)[:,None]).ravel()

并使用多核 numexpr module -

import numexpr as ne

d = {'V3r':vec3.reshape(-1,3),'vec12D':vec1[:,None]}
out = ne.evaluate('V3r*vec12D*(vec12D<0.5)',d).ravel()

时间 -

In [84]: n = 1000000
    ...: np.random.seed(0)
    ...: vec1 = np.random.rand(n)
    ...: vec2 = np.random.rand(3*n)
    ...: vec3 = np.random.rand(3*n)

In [86]: %%timeit
    ...: v = vec3.reshape(-1,3)*vec1[:,None]
    ...: m = vec1<0.5
    ...: vec2_out = (v*m[:,None]).ravel()
10 loops, best of 3: 23.2 ms per loop

In [87]: %%timeit
    ...: mask = vec1<0.5
    ...: vec2_out = (vec3.reshape(-1,3)*(vec1*mask)[:,None]).ravel()
100 loops, best of 3: 13.1 ms per loop

In [88]: %%timeit
    ...: d = {'V3r':vec3.reshape(-1,3),'vec12D':vec1[:,None]}
    ...: out = ne.evaluate('V3r*vec12D*(vec12D<0.5)',d).ravel()
100 loops, best of 3: 4.11 ms per loop

---

对于一般情况，else 部分可能不是零，它将是 -

mask = vec1<0.5
IF_vals = vec3.reshape(-1,3)*vec1[:,None]
ELSE_vals = np.array([1,1,1])
out = np.where(mask[:,None],IF_vals,ELSE_vals).ravel()