创建一个 3d numpy 数组，其条目在给定方向上增加

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.

Answer 1

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.