python - 在 python 中构建对象列表以进行矢量化 : Can a list of structures(objects) be vectorized, 或者需要显式数组

Question

分子模拟中的能量计算本质上充满了“for”循环。传统上，每个原子/分子的坐标存储在数组中。数组的矢量化相当简单，但结构很适合编码。将分子视为单独的对象，每个对象都有自己的坐标和其他属性，就簿记而言非常方便且清晰得多。

我使用的是 Python 3.6 版本

我的问题是，当我使用对象数组时，我无法弄清楚如何矢量化计算......似乎无法避免 for 循环。我是否有必要使用数组才能利用 numpy 并对我的代码进行矢量化？

这是一个利用数组(代码第 121 行)的 Python 示例，并显示了快速(numpy)和慢速(“正常”)Python 能量计算。

https://github.com/Allen-Tildesley/examples/blob/master/python_examples/mc_lj_module.py

使用 numpy 加速方法计算速度要快得多，因为它是矢量化的。

如果我不使用数组，而是使用对象数组(每个对象都有自己的坐标)，我将如何矢量化能量计算？这似乎需要使用较慢的 for 循环。

下面是一个简单的示例代码，其中包含工作缓慢的 for 循环版本，以及尝试的矢量化，但不起作用:

import numpy as np
import time

class Mol:  
    num = 0    
    def __init__(self, r):
        Mol.num += 1
        self.r       = np.empty((3),dtype=np.float_)
        self.r[0]     = r[0]
        self.r[1]     = r[1] 
        self.r[2]     = r[2]
    """ Alot more useful things go in here in practice"""

################################################
#                                              #
#               Main Program                   #
#                                              #
################################################
L = 5.0            # Length of simulation box (arbitrary)
r_cut_box_sq = L/2 # arbitrary cutoff - required
mol_list=[]
nmol = 1000    # number of molecules
part = 1    # arbitrary molecule to interact with rest of molecules

""" make 1000 molecules (1 atom per molecule), give random coordinates """
for i in range(nmol):
    r = np.random.rand(3) * L
    mol_list.append( Mol( r ) )

energy = 0.0

start = time.time()
################################################
#                                              #
#   Slow but functioning loop                  #
#                                              #
################################################
for i in range(nmol):
    if i == part:
        continue

    rij = mol_list[part].r - mol_list[i].r
    rij = rij - np.rint(rij/L)*L                # apply periodic boundary conditions
    rij_sq = np.sum(rij**2)  # Squared separations

    in_range = rij_sq < r_cut_box_sq                
    sr2      = np.where ( in_range, 1.0 / rij_sq, 0.0 )
    sr6  = sr2 ** 3
    sr12 = sr6 ** 2
    energy  += sr12 - sr6                    

end = time.time()
print('slow: ', end-start)
print('energy: ', energy)

start = time.time()
################################################
#                                              #
#   Failed vectorization attempt               #
#                                              #
################################################


    """ The next line is my problem, how do I vectorize this so I can avoid the for loop all together?
Leads to error AttributeError: 'list' object has no attribute 'r' """

""" I also must add in that part cannot interact with itself in mol_list"""
rij = mol_list[part].r - mol_list[:].r
rij = rij - np.rint(rij/L)*L                # apply periodic boundary conditions
rij_sq = np.sum(rij**2) 

in_range = rij_sq < r_cut_box_sq
sr2      = np.where ( in_range, 1.0 / rij_sq, 0.0 )
sr6  = sr2 ** 3
sr12 = sr6 ** 2
energy  = sr12 - sr6                    

energy = sum(energy)
end = time.time()
print('faster??: ', end-start)
print('energy: ', energy)

最后

如果在能量计算中，有必要循环每个分子中的每个原子，其中每个原子现在每个分子超过 1 个原子，并且并非所有分子都具有相同数量的原子，因此任何可能的解决方案都会受到影响吗？用于分子与分子相互作用的双 for 循环，而不是当前使用的简单的配对相互作用。

Answer 1

利用 itertools 库可能是前进的方向。假设您将一对分子的能量计算包装在一个函数中:

def calc_pairwise_energy((mol_a,mol_b)):
    # function takes a 2 item tuple of molecules
    # energy calculating code here
    return pairwise_energy

然后你可以使用 itertools.combinations 来获取所有分子对和 python 内置的列表推导式(下面最后一行 [ ] 内的代码):

from itertools import combinations
pairs = combinations(mol_list,2)
energy = sum( [calc_pairwise_energy(pair) for pair in pairs] )

我回到这个答案是因为我意识到我没有正确回答你的问题。根据我已经发布的内容，成对能量计算函数如下所示(我对您的代码进行了一些优化):

def calc_pairwise_energy(molecules):
    rij = molecules[0].r - molecules[1].r
    rij = rij - np.rint(rij/L)*L
    rij_sq = np.sum(rij**2)  # Squared separations
    if rij_sq < r_cut_box_sq:
        return (rij_sq ** -6) - (rij_sq ** - 3)
    else:
        return 0.0

在一次调用中执行所有成对计算的矢量化实现可能如下所示:

def calc_all_energies(molecules):
    energy = 0
    for i in range(len(molecules)-1):
        mol_a = molecules[i]
        other_mols = molecules[i+1:]
        coords = np.array([mol.r for mol in other_mols])
        rijs = coords - mol_a.r
        # np.apply_along_axis replaced as per @hpaulj's  comment (see below)
        #rijs = np.apply_along_axis(lambda x: x - np.rint(x/L)*L,0,rijs)
        rijs = rijs - np.rint(rijs/L)*L
        rijs_sq = np.sum(rijs**2,axis=1)
        rijs_in_range= rijs_sq[rijs_sq < r_cut_box_sq]
        energy += sum(rijs_in_range ** -6 - rijs_in_range ** -3)
    return energy

这要快得多，但这里仍然有很多需要优化的地方。