在 python/NumPy 中广播/矢量化内部和外部 for 循环 [英] Broadcasting/Vectorizing inner and outer for loops in python/NumPy

问题描述

我使用 vectorization 将双 for 循环 变成了单个 for 循环.我现在想去掉最后一个loop.

I have turned a double for loop into a single for loop using vectorization. I would like to now get rid of the last loop.

我想切片一个Nx3数组坐标并计算切片部分和剩余部分之间的距离不使用for循环.

I want to slice an Nx3 array of coordinates and calculate distances between the sliced portion and the remaining portion without using a for loop.

(1) 切片总是 3x3.

(2) 切片是可变的，即 Mx3，其中 M 总是明显小于 N

(2) the slice is variable i.e., Mx3 where M is always significantly smaller than N

将切片的 1 行与其余部分的交互矢量化很简单.但是，我坚持使用 for 循环来执行(在大小为 3 的切片的情况下)3 个循环，以计算所有距离.

Vectorizing the interaction of 1 row of the slice interacting with the remainder is straightforward. However, I am stuck using a for loop to do (in the case of the slice of size 3) 3 loops, to calculate all distances.

Nx3 数组是原子坐标，切片是特定分子中的所有原子.我想计算给定分子与系统其余部分相互作用的能量.第一步是计算分子中每个原子与所有其他原子之间的距离.第二部分是在函数中使用这些距离来计算能量，这超出了本问题的范围.

The Nx3 array is atom coordinates, the slice is all atoms in a specific molecule. I want to calculate the energy of a given molecule interacting with the rest of the system. The first step is calculating the distances between each atom in the molecule, with all other atoms. The second part is to use those distances in a function to calculate energy, and that is outside the scope of this question.

这是我的最小工作示例(我有 vectorized 内循环，但是，需要(真的很想...)vectorize外循环.该循环不会总是只有 3 大小，而且 python 在 for 循环中很慢.

Here is what I have for a working minimal example (I have vectorized the inner loop, but, need to (would really like to...) vectorize the outer loop. That loop won't always be of only size 3, and python is slow at for loops.

import numpy as np 

box=10 # simulation box is size 10 for this example
r = np.random.rand(1000,3) * box  # avoids huge numbers later by scaling  coords

start=0 #fixed starting index for example (first atom)
end=2   #fixed ending index for example   (last atom)

rj=np.delete(r, np.arange(start,end), 0)
ri = r[np.arange(start,end),:]

atoms_in_molecule, coords = np.shape(ri)
energy = 0
for a in range(atoms_in_molecule):
    rij = ri[a,:] - rj                # I want to get rid of this 'a' index dependance
    rij = rij - np.rint(rij/box)*box  # periodic boundary conditions - necessary
    rij_sq = np.sum(rij**2,axis=1)

    # perform energy calculation using rij_sq
    ener = 4 * ((1/rij_sq)**12 - (1/rij_sq)**6)  # dummy LJ, do not optimize
    energy += np.sum(ener)

print(energy)

这个问题与优化我已有的矢量化无关.我玩过 pdist/cdist 和其他人.我想要的只是摆脱原子上讨厌的 for 循环.我会优化其余的.

This question is not about optimizing the vectorizing I already have. I have played around with pdist/cdist and others. All I want is to get rid of the pesky for loop over atoms. I will optimize the rest.

推荐答案

这里你可以这样做:

R = ri[:,None] - rj[None, :]
R = R - np.rint(R/box)*box
R_sq = np.sum(np.square(R), axis=2)

energy = np.sum(4 * ((1/R_sq)**12 - (1/R_sq)**6))

这篇关于在 python/NumPy 中广播/矢量化内部和外部 for 循环的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！