How to speed up multidimensional array access in scipy.weave?

Question

Hi All:

I'm weaving my c code in python to speed up the loop:

from scipy import weave
from numpy import *

#1) create the array
a=zeros((200,300,400),int)
for i in range(200):
    for j in range(300):
        for k in range(400):    
            a[i,j,k]=i*300*400+j*400+k
#2) test on c code to access the array
code="""
for(int i=0;i<200;++i){
for(int j=0;j<300;++j){
for(int k=0;k<400;++k){
printf("%ld,",a[i*300*400+j*400+k]);    
}
printf("\\n");
}
printf("\\n\\n");
}
"""
test =weave.inline(code, ['a'])

It's working all well, but it is still costly when the array is big. Someone suggested me to use a.strides instead of the nasty "a[i*300*400+j*400+k]" I can't understand the document about .strides.

Any ideas

Thanks in advance

Answer 1

You could replace the 3 for-loops with

grid=np.ogrid[0:200,0:300,0:400]
a=grid[0]*300*400+grid[1]*400+grid[2]

The following suggests this may result in a ~68x (or better? see below) speedup:

% python -mtimeit -s"import test" "test.m1()"
100 loops, best of 3: 17.5 msec per loop
% python -mtimeit -s"import test" "test.m2()"
1000 loops, best of 3: 247 usec per loop

test.py:

import numpy as np

n1,n2,n3=20,30,40
def m1():
    a=np.zeros((n1,n2,n3),int)
    for i in range(n1):
        for j in range(n2):
            for k in range(n3):    
                a[i,j,k]=i*300*400+j*400+k
    return a

def m2():    
    grid=np.ogrid[0:n1,0:n2,0:n3]
    b=grid[0]*300*400+grid[1]*400+grid[2]
    return b 

if __name__=='__main__':
    assert(np.all(m1()==m2()))

With n1,n2,n3 = 200,300,400,

python -mtimeit -s"import test" "test.m2()"

took 182 ms on my machine, and

python -mtimeit -s"import test" "test.m1()"

has yet to finish.

Answer 2

The problem is that you are printing out 2.4 million numbers to the screen in your C code. This is of course going to take a while because the numbers have to be converted into strings and then printed to the screen. Do you really need to print them all to the screen? What is your end goal here?

For a comparison, I tried just setting another array as each of the elements in a. This process took about .05 seconds in weave. I gave up on timing the printing of all elements to the screen after 30 seconds or so.