时间戳

首页 > 解决方案 > 如何向量化这个函数

performance - 如何向量化这个函数

问题描述

以下代码有效，但我想通过矢量化创建 Z。如何做到这一点？

import numpy as np
from numpy import sqrt
from math import fsum    
points = np.array([[0,0],\
                   [5,-1],\
                   [4,6],\
                   [1,3]])
d = lambda x: fsum([sqrt((x[0]-z[0])**2 + (x[1]-z[1])**2) for z in points])
x = np.linspace(min(points[:,0]),max(points[:,0]),100)
y = np.linspace(min(points[:,1]),max(points[:,1]),100)
X, Y = np.meshgrid(x,y)
Z = np.zeros(np.shape(X))
for (i,j),_ in np.ndenumerate(Z):
    Z[i,j] = d([X[i,j],Y[i,j]])
#Z=d([X,Y])  #this fails

标签： performancevectorizationmesh

解决方案

我们可以利用broadcasting直接使用1D版本，从而提高内存效率，并给自己一个矢量化的单线，就像这样 -

Z = np.sqrt((x[:,None] - points[:,0])**2 + (y[:,None,None] - points[:,1])**2).sum(2)

发布样本数据的时间 -

In [80]: %%timeit
    ...: X, Y = np.meshgrid(x,y)
    ...: Z = np.zeros(np.shape(X))
    ...: for (i,j),_ in np.ndenumerate(Z):
    ...:     Z[i,j] = d([X[i,j],Y[i,j]])
10 loops, best of 3: 101 ms per loop

In [81]: %timeit ((x[:,None] - points[:,0])**2 + (y[:,None,None] - points[:,1])**2).sum(2)
1000 loops, best of 3: 246 µs per loop

400x那里加速！

推荐阅读