使用列表索引时如何加速“for”循环？（Python）

Question

我尝试使用 Numpy 函数或向量而不是 for 循环来加速此代码：

sommes = []

for j in range(vertices.shape[0]):

    terme = new_vertices[j] - new_vertices[vertex_neighbors[j]]

    somme_j = np.sum(terme)

    sommes.append(somme_j)

E_int = np.sum(sommes)

（它是迭代算法的一部分，并且有很多“顶点”，所以我认为 for 循环花费的时间太长。）

例如，要计算 j = 0 时的“terme”，我有：

In: new_vertices[0]

Out: array([ 10.2533888 , -42.32279717,  68.27230793])

In: vertex_neighbors[0]

Out: [1280, 2, 1511, 511, 1727, 1887, 759, 509, 1023]

In: new_verties[vertex_neighbors[0]]

Out: array([[ 10.47121043, -42.00123956,  68.218715  ],

            [ 10.2533888 , -43.26905874,  62.59473849],

            [ 10.69773735, -41.26464083,  68.09594854],

            [ 10.37030712, -42.16729601,  68.24639107],

            [ 10.12158146, -42.46624547,  68.29621598],

            [  9.81850836, -42.71158695,  68.33710623],

            [  9.97615447, -42.59625943,  68.31788497],

            [ 10.37030712, -43.11676015,  62.54960623],

            [ 10.55512696, -41.82622703,  68.18954624]])

In: new_vertices[0] - new_vertices[vertex_neighbors[0]]

Out: array([[-0.21782162, -0.32155761,  0.05359293],

             [ 0.        ,  0.94626157,  5.67756944],

             [-0.44434855, -1.05815634,  0.17635939],

             [-0.11691832, -0.15550116,  0.02591686],

             [ 0.13180734,  0.1434483 , -0.02390805],

             [ 0.43488044,  0.38878979, -0.0647983 ],

             [ 0.27723434,  0.27346227, -0.04557704],

             [-0.11691832,  0.79396298,  5.7227017 ],

             [-0.30173816, -0.49657014,  0.08276169]])

问题是 new_vertices[vertex_neighbors[j]] 并不总是具有相同的大小。例如，当 j = 7 时：

In: new_vertices[7]

Out: array([ 10.74106112, -63.88592276, -70.15593947])

In: vertex_neighbors[7]

Out: [1546, 655, 306, 1879, 920, 925]

没有for循环可以吗？我的想法已经用完了，所以任何帮助将不胜感激！

Answer 1

是的，这是可能的。这个想法是用来np.repeat创建一个向量，其中的项目重复可变次数。这是代码：# The two following lines can be done only once if the indices are constant between iterations (precomputation)counts = np.array([len(e) for e in vertex_neighbors])flatten_indices = np.concatenate(vertex_neighbors)E_int = np.sum(np.repeat(new_vertices, counts, axis=0) - new_vertices[flatten_indices])这是一个基准：import numpy as npfrom time import *n = 32768vertices = np.random.rand(n, 3)indices = []count = np.random.randint(1, 10, size=n)for i in range(n):    indices.append(np.random.randint(0, n, size=count[i]))def initial_version(vertices, vertex_neighbors):    sommes = []    for j in range(vertices.shape[0]):        terme = vertices[j] - vertices[vertex_neighbors[j]]        somme_j = np.sum(terme)        sommes.append(somme_j)    return np.sum(sommes)def optimized_version(vertices, vertex_neighbors):    # The two following lines can be precomputed    counts = np.array([len(e) for e in indices])    flatten_indices = np.concatenate(indices)    return np.sum(np.repeat(vertices, counts, axis=0) - vertices[flatten_indices])def more_optimized_version(vertices, vertex_neighbors, counts, flatten_indices):    return np.sum(np.repeat(vertices, counts, axis=0) - vertices[flatten_indices])timesteps = 20a = time()for t in range(timesteps):    res = initial_version(vertices, indices)b = time()print("V1: time:", b - a)print("V1: result", res)a = time()for t in range(timesteps):    res = optimized_version(vertices, indices)b = time()print("V2: time:", b - a)print("V2: result", res)a = time()counts = np.array([len(e) for e in indices])flatten_indices = np.concatenate(indices)for t in range(timesteps):    res = more_optimized_version(vertices, indices, counts, flatten_indices)b = time()print("V3: time:", b - a)print("V3: result", res)这是我机器上的基准测试结果：V1: time: 3.656714916229248V1: result -395.8416223057596V2: time: 0.19800186157226562V2: result -395.8416223057595V3: time: 0.07983255386352539V3: result -395.8416223057595正如您所看到的，这一优化版本比参考实现快 18 倍，而预先计算索引的版本比参考实现快 46 倍。请注意，优化版本应该需要更多 RAM（特别是如果每个顶点的邻居数量很大）。