Create a 3d numpy array whose entries increase in the given direction

Question

Given a shape (n,m,k) and a nonzero vector (u,v,w) with entries in {-1,0,1}, I would like to create a numpy array of shape (n,m,k). The entries of the array should start at 1 and increase in the direction of the vector.

Although I am specifically asking for 3d arrays, let me illustrate with 2d examples:

(n,m) = (3,4) and (u,v) = (-1,0) gives:

4 3 2 1
4 3 2 1
4 3 2 1

(n,m) = (4,3) and (u,v) = (1,-1) gives:

1 2 3
2 3 4
3 4 5
4 5 6

I can create them using nested for loops, but I am wondering if there is a faster solution since I will be working with larger arrays.

I can't understand your rule.. should the second one be (u,v) = (1,1)? and the first on (u,v) = (0,-1)? — ddoGas
– ddoGas, Commented Mar 8, 2020 at 11:15
@ddoGas Ah I see. Sorry about the confusion. Think about the vectors in mathematics. (-1,0) is a vector from right to left, so the array grows in this direction. (1,-1) is a vector from topleft to bottomright, so the array grows in that direction. — ThePortakal
– ThePortakal, Commented Mar 8, 2020 at 11:25

Paul Panzer · Accepted Answer · 2020-03-08 11:52:44Z

Here is a method using stride_tricks. Works for arbitrary number of dimensions.

from numpy.lib.stride_tricks import as_strided                    

def pp(dims,strides):                                                      
    dims,strides = np.asarray(dims),np.asarray(strides)                    
    aux = np.arange(1,(dims-1).sum()+2)                                    
    return as_strided(aux[(dims-1)@(strides==-1):],dims,aux.strides*strides)

Examples:

>>> pp((2,3),(-1,0))
array([[2, 2, 2],
       [1, 1, 1]])
>>> pp((2,3,4),(-1,-1,-1))
array([[[7, 6, 5, 4],
        [6, 5, 4, 3],
        [5, 4, 3, 2]],

       [[6, 5, 4, 3],
        [5, 4, 3, 2],
        [4, 3, 2, 1]]])

Note that the convention for y-axis is it starts at top and goes down. If you want otherwise you'd have to flip it.

Also note that the arrays produced are non-contiguous views, if you want to modify them better make a copy.

Collectives™ on Stack Overflow

Create a 3d numpy array whose entries increase in the given direction

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related