Numpy apply function elementwise based on position

Question

I have a numpy array A with shape (M,N). I want to create a new array B with shape (M,N,3) where the result would be the same as the following:

import numpy as np

def myfunc(A,sx=1.5,sy=3.5):
    M,N=A.shape
    B=np.zeros((M,N,3))

    for i in range(M):
        for j in range(N):
            B[i,j,0]=i*sx
            B[i,j,1]=j*sy
            B[i,j,2]=A[i,j]
    return B

A=np.array([[1,2,3],[9,8,7]])
print(myfunc(A))

Giving the result:

[[[0.  0.  1. ]
  [0.  3.5 2. ]
  [0.  7.  3. ]]

 [[1.5 0.  9. ]
  [1.5 3.5 8. ]
  [1.5 7.  7. ]]]

Is there a way to do it without the loop? I was thinking whether numpy would be able to apply a function element-wise using the indexes of the array. Something like:

def myfuncEW(indx,value,out,vars):
    out[0]=indx[0]*vars[0]
    out[1]=indx[1]*vars[1]
    out[2]=value

M,N=A.shape
B=np.zeros((M,N,3))
np.applyfunctionelementwise(myfuncEW,A,B,(sx,sy))

Paul Panzer · Accepted Answer · 2018-02-09 17:00:39Z

0

You could use mgrid and moveaxis:

>>> M, N = A.shape
>>> I, J = np.mgrid[:M, :N] * np.array((sx, sy))[:, None, None]
>>> np.moveaxis((I, J, A), 0, -1)
array([[[ 0. ,  0. ,  1. ],
        [ 0. ,  3.5,  2. ],
        [ 0. ,  7. ,  3. ]],

       [[ 1.5,  0. ,  9. ],
        [ 1.5,  3.5,  8. ],
        [ 1.5,  7. ,  7. ]]])
>>>

answered Feb 9, 2018 at 17:00

Paul Panzer

53.3k3 gold badges60 silver badges103 bronze badges

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

Comments

Matthias Fripp · Accepted Answer · 2018-02-09 17:03:10Z

0

You could use meshgrid and dstack, like this:

import numpy as np

def myfunc(A,sx=1.5,sy=3.5):
    M, N = A.shape
    J, I = np.meshgrid(range(N), range(M))
    return np.dstack((I*sx, J*sy, A))

A=np.array([[1,2,3],[9,8,7]])
print(myfunc(A))

# array([[[ 0. ,  0. ,  1. ],
#         [ 0. ,  3.5,  2. ],
#         [ 0. ,  7. ,  3. ]],
#
#        [[ 1.5,  0. ,  9. ],
#         [ 1.5,  3.5,  8. ],
#         [ 1.5,  7. ,  7. ]]])

edited Feb 9, 2018 at 17:03

answered Feb 9, 2018 at 16:57

Matthias Fripp

18.9k5 gold badges36 silver badges49 bronze badges

Comments

kuppern87 · Accepted Answer · 2018-02-09 19:10:33Z

0

By preallocating the 3d array B, you save about half the time compared to stacking I, J and A.

def myfunc(A, sx=1.5, sy=3.5):
    M, N = A.shape
    B = np.zeros((M, N, 3))
    B[:, :, 0] = np.arange(M)[:, None]*sx
    B[:, :, 1] = np.arange(N)[None, :]*sy
    B[:, :, 2] = A
    return B

answered Feb 9, 2018 at 19:10

kuppern87

1,1359 silver badges14 bronze badges

Collectives™ on Stack Overflow

Numpy apply function elementwise based on position

3 Answers 3

Comments

Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related