时间戳

首页 > 解决方案 > 在 python 中构建对象列表以进行矢量化：可以对结构（对象）列表进行矢量化，还是需要显式数组

python - 在 python 中构建对象列表以进行矢量化：可以对结构（对象）列表进行矢量化，还是需要显式数组

问题描述

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

我正在使用 Python 3.6 版

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

这是一个使用数组的 python 示例（代码的第 121 行），并显示了快速（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 个，并且并非所有分子都具有相同数量的原子，因此对于循环用于分子 - 分子相互作用，而不是目前使用的简单的配对配对相互作用。

标签： pythonarraysnumpyvectorization

解决方案

利用 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

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

推荐阅读