How do I speed up a Python nested loop?

Question

I'm trying to calculate the gravity effect of a buried object by calculating the effect on each side of the body and then summing up the contributions to get one measurement at one station, and repeating for a number of stations. The code is as follows (the body is a square and the code calculates clockwise around it. That's why it goes from -x back to -x coordinates).

grav = []
x = si.arange(-30.0, 30.0, 0.5)

# -9.79742526     9.78716693    22.32153704    27.07382349  2138.27146193
xcorn = (-9.79742526,  9.78716693,  9.78716693, -9.79742526, -9.79742526)
zcorn = (22.32153704, 22.32153704, 27.07382349, 27.07382349, 22.32153704)
gamma = (6.672*(10**-11)) # 'N m^2 / Kg^2'
rho = 2138.27146193 # 'Kg / m^3'
grav = []
iter_time = []
def procedure():
    for i in si.arange(len(x)): # Cycles position
        t0 = time.clock()
        sum_lines = 0.0

        for n in si.arange(len(xcorn)-1): # Cycles corners
            x1 = xcorn[n]   - x[i]
            x2 = xcorn[n+1] - x[i]
            z1 = zcorn[n]   -0.0  # Just depth to corner since all observations are on the surface.
            z2 = zcorn[n+1] -0.0
            r1 = ((z1**2) + (x1**2))**0.5
            r2 = ((z2**2) + (x2**2))**0.5
            O1 = si.arctan2(z1, x1)
            O2 = si.arctan2(z2, x2)
            denom = z2-z1
            if denom == 0.0:
                denom = 1.0e-6

            alpha = (x2-x1)/denom

            beta = ((x1*z2)-(x2*z1))/denom
            factor = (beta/(1.0 + (alpha**2)))
            term1 = si.log(r2/r1) # Log base 10
            term2 = alpha*(O2-O1)
            sum_lines = sum_lines + (factor*(term1-term2))
        sum_lines = sum_lines*2*gamma*rho
        grav.append(sum_lines)
        t1 = time.clock()
        dt = t1-t0
        iter_time.append(dt)

How can I speed this loop up?

Suggesting codereview.stackexchange.com A better forum for this type questions. — mjv
– mjv, Commented Oct 12, 2011 at 17:32
@aix: si.arange() might suggest that he uses numpy arrays already. — jfs
– jfs, Commented Oct 12, 2011 at 17:44
@J.F.Sebastian: Good catch. I've removed my comment to not cause any further confusion. — NPE
– NPE, Commented Oct 12, 2011 at 19:00

Zack Bloom · Accepted Answer · 2011-10-12 18:16:48Z

1

Your xcorn and zcorn values repeat, so consider caching the result of some of the computations.

Take a look at the timeit and profile modules to get more information about what is taking the most computational time.

answered Oct 12, 2011 at 18:16

Zack Bloom

8,4372 gold badges23 silver badges27 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Peter Mortensen · Accepted Answer · 2023-11-10 00:53:23Z

1

It is very inefficient to access individual elements of a NumPy array in a Python loop. For example, this Python loop:

for i in xrange(0, len(a), 2):
    a[i] = i

would be much slower than:

a[::2] = np.arange(0, len(a), 2)

You could use a better algorithm (less time complexity) or use vector operations on NumPy arrays as in the example above. But the quicker way might be just to compile the code using Cython:

#cython: boundscheck=False, wraparound=False
#procedure_module.pyx
import numpy as np
cimport numpy as np

ctypedef np.float64_t dtype_t

def procedure(np.ndarray[dtype_t,ndim=1] x,
              np.ndarray[dtype_t,ndim=1] xcorn):
    cdef:
        Py_ssize_t i, j
        dtype_t x1, x2, z1, z2, r1, r2, O1, O2
        np.ndarray[dtype_t,ndim=1] grav = np.empty_like(x)
    for i in range(x.shape[0]):
        for j in range(xcorn.shape[0]-1):
            x1 = xcorn[j]-x[i]
            x2 = xcorn[j+1]-x[i]
            ...
        grav[i] = ...
    return grav

It is not necessary to define all types, but if you need a significant speedup compared to Python, you should define at least types of arrays and loop indexes.

You could use cProfile (Cython supports it) instead of manual calls to time.clock().

To call procedure():

#!/usr/bin/env python
import pyximport; pyximport.install() # pip install cython
import numpy as np

from procedure_module import procedure

x = np.arange(-30.0, 30.0, 0.5)
xcorn = np.array((-9.79742526, 9.78716693, 9.78716693, -9.79742526, -9.79742526))
grav = procedure(x, xcorn)

edited Nov 10, 2023 at 0:53

Peter Mortensen

31.4k22 gold badges110 silver badges134 bronze badges

answered Oct 12, 2011 at 18:17

jfs

417k210 gold badges1k silver badges1.7k bronze badges

1 Comment

user991926 Over a year ago

Thanks for the help, il take a look at cython and see if that helps, im not worried about the time.clock() because ive only got that in there to see how fast each iteration is taking while i try and speed it up, its not a necessary part of the code.

Collectives™ on Stack Overflow

How do I speed up a Python nested loop?

2 Answers 2

Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related