时间戳

首页 > 解决方案 > 如何用不同大小的向量向量化操作

python - 如何用不同大小的向量向量化操作

问题描述

我有不同大小的向量，并且想要进行元素操作。如何在 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]

非常感谢你的帮助。

标签： pythonnumpyvectorization

解决方案

我们可以利用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()

推荐阅读