在python中构造对象列表以进行矢量化:可以对结构(对象)列表进行矢量化,还是需要显式数组

时间:2018-08-03 20:45:19

标签: python arrays numpy vectorization

分子模拟中的能量计算固有地充满了“ 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个,并且并非所有分子都具有相同的原子数,因此具有用于分子-分子相互作用而不是当前使用的简单对-对相互作用的双环。

2 个答案:

答案 0 :(得分:2)

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

这要快得多,但是这里还有很多要优化的地方。

答案 1 :(得分:1)

如果您想以坐标为输入来计算能量,我假设您正在寻找成对距离。为此,您应该查看SciPy库。具体来说,我来看看scipy.spatial.distance.pdist。可以在here中找到该文档。