Generic way to set values depending on indexes in ndarray?

user_2300597 0
Viewed 71

I wrote this below but it looks ugly/clumsy to me.
Is there a better way to loop through all elements while keeping track of indexes
(i.e. keeping track at which address I am) and set a value which is a function
of the indexes?

I can imagine this getting even uglier if I have more dimensions.
There should be a better way designed to do just this. 

import numpy as np

def f(i,j,k,l):
    return i + 2 * j+ 3 * k + 4 * l

A = np.zeros((3,4,7,8), dtype=int)

i1 = i2 = i3 = i4 = -1

i1 = -1
for r1 in A:
    i1 += 1

    i2 = -1
    for r2 in r1:
        i2 += 1

        i3 = -1
        for r3 in r2:

            i3 += 1

            i4 = -1
            for r4 in r3:
                i4 += 1

                v = f(i1, i2, i3, i4)

                print ("Setting value {0} at ({1}, {2}, {3}, {4})".format(v, i1, i2, i3, i4))

                A[i1, i2, i3, i4] = v

print(A)
multidimensional-array python-3.x numpy python
3 Answers

There is np.fromfunction:

np.fromfunction(f,(3,4,7,8),dtype=int)

This works as long as f is "numpy aware", i.e can handle array arguments.

user_72073920

Solving the problem with np.apply_along_axis manner

I break it down to a more easy example to explain the logic behind my idea and move on from that.

I want to use a combination of np.indices which returns a grid of indices of the given shape and np.apply_along_axis with which we can apply a function along a given axis of a multi-dimensional array.

Example:

def f(arr):
    return arr[0] + 2 * arr[1]

A=np.zeros((2,3))

Note that the input of the function f is changed since it's input is a np.array of shape (4,) and not 4-element tuple.

we will use

np.indices(A.shape)

to get the a grid of indices of matrix A. We can easily apply the function as

np.apply_along_axis(f, axis=0, arr=np.indices(A.shape))

which results in

array([[0, 2, 4],    # 0 + 2*0, 0 + 2*1, 0 +2*2
       [1, 3, 5]])   # 1 + 2*0, 1 + 2*1, 1 +2*2

Now let's move to the example given by OP. The code is the same, we just exchange f and A

def f(arr):
   return arr[0] + 2 * arr[1]+ 3 * arr[2] + 4* arr[3]

A = np.zeros((3,4,7,8), dtype=int)
np.apply_along_axis(f, axis=0, arr=np.indices(A.shape))

Since the result is of shape (3, 4, 7, 8) I won't post it here.

Timing results

np.apply_along_axis proves to be a lot slower than the manual iteration even for arrays of considerable size.

Here are some timing results:

shape            manual            np.apply_along_axis
100,100,100,10  4.18 s ± 39.7 ms     34.8 s ± 298 ms
100,100,10,10   412 ms ± 6.23 ms     3.49 s ± 13.4 ms
3,4,7,8         285 µs ± 1.79 µs     2.43 ms ± 35.6 µs
user_95343900

Your f, at least this example, works with arrays, since it uses the numpy vectorized operators + and *. 

def f(i,j,k,l):
    return i + 2 * j+ 3 * k + 4 * l

Using np.ix_ to generate indices that can broadcast against each other, I can produce an array that matches your iterative A:

In [14]: A1 = f(*np.ix_(np.arange(3),np.arange(4),np.arange(7),np.arange(8)))                                   
In [15]: np.allclose(A,A1)                                                                                      
Out[15]: True

Now if f can work with arrays, we can't do this. A classic example is one that uses math.sin instead of np.sin. Or one that is if tests.

That ix_ way is fast, since it calls f only once. I can also generate all indices as tuples with np.ndindex, and call f many times:

In [24]: A2 = np.zeros_like(A)                                                                                  
In [25]: for ijk in np.ndindex(*A2.shape): A2[ijk]=f(*ijk)                                                      
In [26]: np.allclose(A,A2)                                                                                      
Out[26]: True

The speed will be similar to your A evaluation.

In [29]: %%timeit  
    ...: A2 = np.zeros_like(A) 
    ...: for ijk in np.ndindex(*A2.shape): A2[ijk]=f(*ijk)  
    ...:  
    ...:                                                                                                        
758 µs ± 11.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [30]: timeit A1 = f(*np.ix_(np.arange(3),np.arange(4),np.arange(7),np.arange(8)))                            
54.3 µs ± 2.44 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [31]: %%timeit 
    ...: A = np.zeros((3,4,7,8), dtype=int) 
    ...:  
    ...: i1 = i2 = i3 = i4 = -1 
    ...:  
    ...: i1 = -1 
    ...: for r1 in A: 
    ...:     i1 += 1 
    ...:  
    ...:     i2 = -1 
    ...:     for r2 in r1: 
    ...:         i2 += 1 
 ....
    ...:                 A[i1, i2, i3, i4] = v 
    ...:                                                                                                        
506 µs ± 2.03 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Come to think of it, we can clean up your loops with a more typical Python range loop:

In [35]: %%timeit 
    ...: A = np.zeros((3,4,7,8), dtype=int) 
    ...: for i1 in range(A.shape[0]): 
    ...:     for i2 in range(A.shape[1]): 
    ...:         for i3 in range(A.shape[2]): 
    ...:             for i4 in range(A.shape[3]): 
    ...:                 v = f(i1, i2, i3, i4) 
    ...:                 A[i1, i2, i3, i4] = v 
    ...:                                                                                                        
343 µs ± 17.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
user_9019250
Related