如何提高python3中这个矩阵生成循环的运行时间？

Question

我有模拟大量粒子之间相互作用的代码。使用分析，我发现导致最慢的函数是一个循环，它遍历我的所有粒子并计算出每个粒子之间的碰撞时间。这会生成一个对称矩阵，然后我会从中取出最小值。

def find_next_collision(self, print_matrix = False):
    """
    Sets up a matrix of collision times
    Returns the indices of the balls in self.list_of_balls that are due to
    collide next and the time to the next collision
    """
    self.coll_time_matrix = np.zeros((np.size(self.list_of_balls), np.size(self.list_of_balls)))

    for i in range(np.size(self.list_of_balls)):
        for j in range(i+1):
            if (j==i):
                self.coll_time_matrix[i][j] = np.inf
            else:
                self.coll_time_matrix[i][j] = self.list_of_balls[i].time_to_collision(self.list_of_balls[j])
    matrix = self.coll_time_matrix + self.coll_time_matrix.T
    self.coll_time_matrix = matrix
    ind = np.unravel_index(np.argmin(self.coll_time_matrix, axis = None), self.coll_time_matrix.shape)
    dt = self.coll_time_matrix[ind]
    if (print_matrix):
        print(self.coll_time_matrix)
    return dt, ind

这段代码是一个类中的一个方法，它定义了所有粒子的位置。这些粒子中的每一个都是保存在self.list_of_balls（这是一个列表）中的一个对象。正如你所看到的，我已经只迭代了这个矩阵的一半，但它仍然是一个相当慢的函数。我试过使用 numba，但这是一段相当大的代码，我不想在速度慢的时候用 numba 优化每个函数。

有人对编写此函数的更有效方法有任何想法吗？

提前谢谢你！

Answer 1

像Raubsauger mentioned in their answer，评估ifs 很慢

for j in range(i+1): 
    if (j==i):

您只需执行for j in range(i) 即可摆脱此if。这样j就从0变成i-1

您还应该尽可能避免循环。您可以通过以矢量化方式表达您的问题，并使用 leverage SIMD operations 的 numpy 或 scipy 函数来加快计算速度。这是一个简化的示例，假设time_to_collision 只是将欧几里得距离除以速度。如果您将球的坐标和速度存储在一个 numpy 数组中，而不是将球对象存储在列表中，您可以这样做：

from scipy.spatial.distance import pdist
rel_distances = pdist(ball_coordinates)
rel_speeds = pdist(ball_speeds)
time = rel_distances / rel_speeds

pdist documentation

当然，如果您的 time_to_collision 函数更复杂，这将无法逐字执行，但它应该为您指明正确的方向。

Answer 2

第一个问题：你有多少粒子？

如果你有很多粒子：一个改进是

    for i in range(np.size(self.list_of_balls)):
        for j in range(i):
            self.coll_time_matrix[i][j] = self.list_of_balls[i].time_to_collision(self.list_of_balls[j])
        self.coll_time_matrix[i][i] = np.inf

经常执行ifs 会减慢一切。在内部循环中避免它们

第二个问题：有必要每次都计算吗？计算时间点并只刷新那些涉及碰撞的行和列不是更快吗？

编辑：

这里的想法是最初计算剩余时间或（更好的解决方案）碰撞的时间戳，因为你已经和订单一样。但是，为了放弃计算结果，您只需要在需要时更新值。这样你只需要计算 2*n 而不是 n^2/2 值。

草图：

# init step, done once at the beginning, might need an own function
matrix ... # calculate matrix like before; I asume that you use timestamps instead of time left
min_times = np.zeros(np.size(self.list_of_balls))
for i in range(np.size(self.list_of_balls)):
    min_times[i] = min(self.coll_time_matrix[i])
order_coll = np.argsort(min_times)

ind = order_coll[0]
dt = self.coll_time_matrix[ind]
return dt, ind


# function step: if a collision happened, order_coll[0] and order_coll[1] hit each other

for balls in order_coll[0:2]:
    for i in range(np.size(self.list_of_balls)):
        self.coll_time_matrix[balls][i] = self.list_of_balls[balls].time_to_collision(self.list_of_balls[i])
        self.coll_time_matrix[i][balls] = self.coll_time_matrix[balls][i]
    self.coll_time_matrix[balls][balls] = np.inf
for i in range(np.size(self.list_of_balls)):
    min_times[i] = min(self.coll_time_matrix[i])
order_coll = np.argsort(min_times)

ind = order_coll[0]
dt = self.coll_time_matrix[ind]
return dt, ind

如果您计算矩阵中剩余的时间，则必须从矩阵中减去经过的时间。此外，您还需要以某种方式存储矩阵和（可选）min_times 和 order_coll。