Tom Tom - 2 months ago 35
Python Question

How to speed up three dimensional sum for Madelung Constant?

For my Computational Physics class we have to compute the Madelung Constant for NaCl. My code to do this uses three nested for loops and therefore runs very slowly. I was wondering if there was a way to use arrays or some other method to increase the speed of computation. Thanks

from math import sqrt
l= int(input("The lattice size is :"))
M = 0.0
for i in range(-L,L+1):
for j in range(-L,L+1):
for k in range(-L,L+1):
if not (i==j==k==0):
M += ((-1)**(i+j+k+1))/sqrt(i*i +j*j +k*k)

print("The Madelung constant for NaCl with lattice size",l,"is",M)

Answer

Since you've noted in a comment that you may use numpy, I suggest doing that. You can construct a 3d grid for your integers, and compute each term simultaneously, thereby vectorizing your calculation. You only need to watch out for the singular case where every integer is 0, for instance using numpy.where:

import numpy as np

ran = np.arange(-L,L+1)
i,j,k = np.meshgrid(ran,ran,ran)
M = np.where((i!=0) | (j!=0) | (k!=0),
             (-1)**(i+j+k+1)/np.sqrt(i**2+j**2+k**2),
             0).sum()

ran is a numpy array with the same elements as in range() (if cast to a list in python 3). meshgrid then constructs three 3d arrays, which together span the 3d space where you need to perform your sum.

Note that for large domains this approach has much higher memory need. This is ususal with vectorization: you can spare CPU time at the cost of increased memory need.