使用numpy更新节点值而不进行for循环

时间:2019-04-01 16:09:36

标签: python performance numpy vectorization

我正在尝试根据元素值更新mesh上的节点值。

在数组faces中,我定义了一个元素的节点ID(假设我只有两个元素):

faces = np.array([[0, 1, 2], [1, 3, 2]])

force_el数组包含的力平均作用于元素的每个节点:

force_el = np.array([[0.7, 1.1], [1.2, 0.3]])

现在我想更新节点力force_node

force_node = np.zeros((4, force_el.shape[1]))
for face, fel in zip(faces, force_el):
    force_node[face.ravel(), :] += fel

结果是:

>>> force_node
array([[0.7, 1.1],
       [1.9, 1.4],
       [1.9, 1.4],
       [1.2, 0.3]])

由于此更新必须执行多次(大约100k-1m次),因此我正在尝试对其进行优化,但是我看不到一个好的解决方案。

2 个答案:

答案 0 :(得分:1)

您应该尽可能利用numpy广播。

使用np.add.at

np.add.at(force_node, faces, force_el[:,None])

答案 1 :(得分:1)

您可以使用一些matrix-multiplication force -

out_nrows = 4 # number of nodes
mask = np.zeros((len(faces),out_nrows),dtype=bool)
np.put_along_axis(mask,faces,True,axis=1)
force_node_out = mask.T.dot(force_el)

force_el中的列数较少时,我们也可以使用np.bincount获得更好的性能-

out_nrows = 4 # number of nodes
out = np.zeros((out_nrows, force_el.shape[1]))
n = faces.shape[1]
l = force_el.shape[1]
for i in range(n):
    for j in range(l):
        out[:,j] += np.bincount(faces[:,i],force_el[:,j],minlength=out_nrows)

时间-

In [35]: # Setup data (from OP's comments)
    ...: np.random.seed(0)
    ...: faces=np.array([np.random.choice(1800,3,replace=0) for i in range(3500)])
    ...: force_el = np.random.rand(len(faces),3)

In [36]: %%timeit # Original loopy soln
    ...: out_nrows = 1800
    ...: force_node = np.zeros((out_nrows, force_el.shape[1]))
    ...: for face, fel in zip(faces, force_el):
    ...:     force_node[face.ravel(), :] += fel
100 loops, best of 3: 16.1 ms per loop

In [37]: %%timeit # @RafaelC's soln with np.add.at
    ...: force_node = np.zeros((1800, force_el.shape[1]))
    ...: np.add.at(force_node, faces, force_el[:,None])
100 loops, best of 3: 2.45 ms per loop

In [38]: %%timeit # Posted in this post that uses matrix-multiplication
    ...: out_nrows = 1800
    ...: mask = np.zeros((len(faces),out_nrows),dtype=bool)
    ...: np.put_along_axis(mask,faces,True,axis=1)
    ...: force_node_out = mask.T.dot(force_el)
10 loops, best of 3: 38.4 ms per loop

In [39]: %%timeit # Posted in this post that uses bincount
    ...: out_nrows = 1800
    ...: out = np.zeros((out_nrows, force_el.shape[1]))
    ...: n = faces.shape[1]
    ...: l = force_el.shape[1]
    ...: for i in range(n):
    ...:     for j in range(l):
    ...:         out[:,j]+=np.bincount(faces[:,i],force_el[:,j],minlength=out_nrows)
10000 loops, best of 3: 149 µs per loop