Improve Mandelung constant code

Question

I'm learning and improving my Python skills.

I did a program in Python about Mandelung constant. But, I'm having kind of a problem. The Mangelung constant is calculated using this sum:

$$ V_{total} = \sum_{\substack{i,j,k = -L\\ i, j, k \neq 0}}^L V(i,j,k)= \frac e{4\pi\epsilon_0} M $$

Or

$$M =\sum_{\substack{i,j,k = -L\\ i, j, k \neq 0}}^L \frac{(-1)^{i + j + k}}{\sqrt {i^2 + j^2 + k^2}} $$

When I run it, it takes a long time, when I put a huge number. So, I need to improve my code, to run it faster. Can someone explain me a way to do it? (importing others libs, using other stuff)

The code that I did:

import time

start_time = time.time()
L = int(input("Put the number of L:")) # size of the lattice
L = L+1 # this is for the vector (0,0,0)
# n = 0 # number of atoms
M = 0 # Madelung constant
for i in range(-L,L+1):
    for j in range(-L,L+1):
        for k in range(-L,L+1):
            # n += 1 #counter for number of atoms
            if i == j == k == 0: # doesn't count the origin (0,0,0)
                continue
            r = (i**2 + j**2 + k**2)**(-0.5)
            if (i + j + k) % 2 == 1: # odd number
                r *= -1
            M += r
print ("Mandelung Constant is::", M)
print("It takes %s seconds" % (time.time() - start_time))

When I put a $L = 300$, it takes more than 7 minutes. This is why I'm trying to improve it.

Answer 1

As mentioned by @Richard, loops in Python are slow. Two solutions come to my mind:

• Use NumPy and vectorize the operations. This will speed up your calculations at the cost of storing your intermediate arrays in memory.

• Wrap your calculation in a function and use Numba's jit decorator to (magically) accelerate your calculation. This implies refactoring the code sometimes.

---

Using NumPy capabilities you can rewrite to have something like the following:

import numpy as np
import time


start_time = time.time()
L = 300
i = np.array(range(-L, L + 1), dtype=np.float)
I, J, K = np.meshgrid(i, i, i)

M = (-1)**(I + J + K)/np.sqrt(I**2 + J**2 + K**2)
M[(I == 0)*(J == 0)*(K == 0)] = 0
M = np.sum(M)
print ("Madelung Constant is::", M)
print("It takes %s seconds" % (time.time() - start_time))

And this gives a result

Madelung Constant is:: -1.7456432959005515
It takes 17.750624418258667 seconds

Compared with

Madelung Constant is:: -1.745643295911936
It takes 310.843811750412 seconds

from your code.

You could also use Numba.

import time
from numba import jit


@jit
def sum_madelung(L):
    M = 0 # Madelung constant
    for i in range(-L , L + 1):
        for j in range(-L, L + 1):
            for k in range(-L, L + 1):
                if i == j == k == 0: # doesn't count the origin (0,0,0)
                    continue
                r = (i**2 + j**2 + k**2)**(-0.5)
                if (i + j + k) % 2 == 1: # odd number
                    r *= -1
                M += r
    return M

start_time = time.time()
M = sum_madelung(300)
print ("Madelung Constant is::", M)
print("It takes %s seconds" % (time.time() - start_time))

with the result

Madelung Constant is:: -1.7456432959126562
It takes 2.8461902141571045 seconds