Use numpy to create matrix based on conditional indices of an array

Viewed 868

I would like to create an n by m matrix based on elements of an n + m length array.

Here a simple double for loop suffices, but I wish for an expedient solution. The matrix will be relatively small.

n = 4
m = 6
s = n + m

array = np.arange(s)  # note: arange is only for example. real array varies.
matrix = np.zeros((n,m))

for i in range(n):
    for j in range (m):
        matrix[i,j] = array[i+j]

I have found that a comprehension is faster than the double for loops

matrix3 = [[array[i+j] for i in range(m)] for j in range(n)]

Is there a faster way?

An additional bonus would be to incorporate the modulo operator. I only actually need the indices where i+j % 2 == 0. In the double for loop the modulo method seems a little bit faster, but this may not be convenient or expedient for generating this matrix via numpy.

It is fine to not do this, since matrix multiplcation will occur after and the necessary elements will be multiplied by zero anyways. My mentioning the modulo is only in the off-case that this leads to a faster solution.

for this MWE

for i in range(n):
    for j in range (m):
        if (i + j) % 2 == 0:
            matrix[i,j] = array[i+j]

note:

I ask for a numpy solution under the assumption numpy will be fastest, but any pure python (including numpy/scipy) solution is fine so long as it is faster than pure python double for loops

motivation:

I am trying to remove all dependencies on arrays from a double for loop so that I can use broadcasting rather than a double for loop. This is the last array left

5 Answers

You can create a hankel matrix:

>>> from scipy.linalg import hankel
>>> matrix = hankel(array[0:n], array[n:s])
>>> matrix
array([[0, 1, 2, 3, 4, 6],
       [1, 2, 3, 4, 6, 7],
       [2, 3, 4, 6, 7, 8],
       [3, 4, 6, 7, 8, 9]])

If you absolutely want to set elements where (i+j)%2==1 to zero you can do (original post):

>>> matrix[::2, 1::2] = 0
>>> matrix[1::2, ::2] = 0
>>> matrix
array([[0, 0, 2, 0, 4, 0],
       [0, 2, 0, 4, 0, 7],
       [2, 0, 4, 0, 7, 0],
       [0, 4, 0, 7, 0, 9]])

You can also set every other value of array to zero then the constructed matrix will have zeros at desired locations:

>>> array[1::2]=0
>>> hankel(array[0:n], array[n:s])
array([[0, 0, 2, 0, 4, 6],
       [0, 2, 0, 4, 6, 0],
       [2, 0, 4, 6, 0, 8],
       [0, 4, 6, 0, 8, 0]])

You can use advanced indexing into array. For efficiency, you can zero odd positions already in the template array.

np.where(np.arange(m+n)&1,0,array)[sum(np.ogrid[:n,:m])]
# array([[0, 0, 2, 0, 4, 0],
#        [0, 2, 0, 4, 0, 6],
#        [2, 0, 4, 0, 6, 0],
#        [0, 4, 0, 6, 0, 8]])

or (faster)

template = np.where(np.arange(m+n)&1,0,array)
np.lib.stride_tricks.as_strided(template,(n,m),2*template.strides)

This is a "compressed" view, if you need to modify the entries you must make a copy (it will still be faster).

Much simpler way to create your table is:

  1. Define a function:

     def tVal(r, c):
         sm = r + c
         return np.where(sm % 2 == 0, sm, 0)
    
  2. Use it as an argument of np.fromfunction:

     arr = np.fromfunction(tVal, (n, m))
    

For your target shape (6 * 4) the result is:

array([[0., 0., 2., 0., 4., 0.],
       [0., 2., 0., 4., 0., 6.],
       [2., 0., 4., 0., 6., 0.],
       [0., 4., 0., 6., 0., 8.]])

Note than tVal is not actually called separately for each array element. It is instead called only once, with 2 arrays (r and c) shaped as the target array, filled with respective arguments for each cell. So this function operates on these arrays (not on single values for each cell index).

This is why this function must contain where, not if for r and c values for particular cell.

And a remark concerning variable names: matrix is a class in Numpy (a subtype of ndarray), so it is a good practice not to use variables with the same name. Use rather other name, as I did in my example.

I would do it directly at the numpy level:

matrix = np.arange(n * m).reshape(n,m)
matrix = matrix // m + matrix % m             # matrix // m is i and matrix % m is j

For n, m = 4, 6 it gives as expected:

array([[0, 1, 2, 3, 4, 5],
       [1, 2, 3, 4, 5, 6],
       [2, 3, 4, 5, 6, 7],
       [3, 4, 5, 6, 7, 8]], dtype=int32)

Your first example:

In [30]: arr=np.arange(24)                                                              
In [31]: [[arr[i+j] for i in range(6)] for j in range(4)]                               
Out[31]: 
[[0, 1, 2, 3, 4, 5],
 [1, 2, 3, 4, 5, 6],
 [2, 3, 4, 5, 6, 7],
 [3, 4, 5, 6, 7, 8]]

To take advantage of 'broadcasting':

In [32]: np.arange(4)[:,None]+np.arange(6)                                              
Out[32]: 
array([[0, 1, 2, 3, 4, 5],
       [1, 2, 3, 4, 5, 6],
       [2, 3, 4, 5, 6, 7],
       [3, 4, 5, 6, 7, 8]])

The outer i loop is replaced by a (n,1) array; the inner j loop is replaced by the (m,) array; together the result is a (n,m) array.

Your more elaborate case:

In [35]: arr = np.arange(24) 
    ...: res = np.zeros((4,6),int) 
    ...: for i in range(4): 
    ...:     for j in range(6): 
    ...:         if (i+j)%2 ==0: 
    ...:             res[i,j] = arr[i+j] 
    ...:                                                                                
In [36]: res                                                                            
Out[36]: 
array([[0, 0, 2, 0, 4, 0],
       [0, 2, 0, 4, 0, 6],
       [2, 0, 4, 0, 6, 0],
       [0, 4, 0, 6, 0, 8]])

So this is the original, with just the even values set.

In [37]: Out[32]                                                                        
Out[37]: 
array([[0, 1, 2, 3, 4, 5],
       [1, 2, 3, 4, 5, 6],
       [2, 3, 4, 5, 6, 7],
       [3, 4, 5, 6, 7, 8]])

Find the odds:

In [38]: Out[32]%2                                                                      
Out[38]: 
array([[0, 1, 0, 1, 0, 1],
       [1, 0, 1, 0, 1, 0],
       [0, 1, 0, 1, 0, 1],
       [1, 0, 1, 0, 1, 0]])

Multiply:

In [39]: Out[32]*(Out[32]%2==0)                                                         
Out[39]: 
array([[0, 0, 2, 0, 4, 0],
       [0, 2, 0, 4, 0, 6],
       [2, 0, 4, 0, 6, 0],
       [0, 4, 0, 6, 0, 8]])

In general to make optimal use of numpy, I try to see overall patterns. That's where small examples are especially valuable.

Related