efficient, pythonic way to zig-zag through meshgrid data

Question

I want to make a trajectory of (x, y) points that zig-zag through an even grid. I want to be able to specify which axis to start on. Is there an efficient way to do this with numpy?

For instance:

xv, yv = np.meshgrid([1,2,3], [4,5])

produces

xv: [[1 2 3]
     [1 2 3]]
yv: [[4 4 4]
     [5 5 5]]

I want my trajectory to be:

(1,4), (2,4), (3,4), (3,5), (2,5), (1,5) 

or

(1,4), (1,5), (2,5), (2,4), (3,4), (3,5)

Answer 1

You might flip the coordinate every other row by using:

xv[1::2] = xv[1::2,::-1]

Then stack xv and yv .

---

xv, yv = np.meshgrid([1,2,3], [4,5])
xv[1::2] = xv[1::2,::-1]

xv
#array([[1, 2, 3],
#       [3, 2, 1]])

np.stack((xv, yv), axis=-1).reshape(-1, 2)

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

Or flip yv :

xv, yv = np.meshgrid([1,2,3], [4,5])
yv[:,1::2] = yv[::-1, 1::2]

yv
#array([[4, 5, 4],
#       [5, 4, 5]])

np.stack((xv.T, yv.T), axis=-1).reshape(-1, 2)
#array([[1, 4],
#       [1, 5],
#       [2, 5],
#       [2, 4],
#       [3, 4],
#       [3, 5]])

Answer 2

We can avoid creating the heavy meshgrids by directly working on the 1D arrays as a performance measure with array-initialization for the two output arrays and then broadcasted assignments for assigning the input arrays once as they are along the even-numbered places and flipped along the odd numbered ones. This gives us the zig-zag ordering.

The implementation would be -

def create_zigzag_grids(a,b):
    m,n = len(b), len(a)

    out1 = np.empty((m,n,2),dtype=int)
    out1[:,:,1] = b[:,None]
    out1[::2,:,0] = a
    out1[1::2,:,0] = a[::-1]

    out2 = np.empty((n,m,2),dtype=int)
    out2[:,:,0] = a[:,None]
    out2[::2,:,1] = b
    out2[1::2,:,1] = b[::-1]
    return out1.reshape(-1,2), out2.reshape(-1,2)

Sample run -

1) Input and output :

In [748]: a,b = np.array([1,2,3,4]), np.array([5,6,7])
     ...: out1, out2 = create_zigzag_grids(a,b)
     ...: 

2) Create meshgrid to verify our results :

In [749]: xv, yv = np.meshgrid(a,b)

In [750]: xv
Out[750]: 
array([[1, 2, 3, 4],
       [1, 2, 3, 4],
       [1, 2, 3, 4]])

In [751]: yv
Out[751]: 
array([[5, 5, 5, 5],
       [6, 6, 6, 6],
       [7, 7, 7, 7]])

3) Verify results :

In [752]: out1
Out[752]: 
array([[1, 5],
       [2, 5],
       [3, 5],
       [4, 5],
       [4, 6],
       [3, 6],
       [2, 6],
       [1, 6],
       [1, 7],
       [2, 7],
       [3, 7],
       [4, 7]])

In [753]: out2
Out[753]: 
array([[1, 5],
       [1, 6],
       [1, 7],
       [2, 7],
       [2, 6],
       [2, 5],
       [3, 5],
       [3, 6],
       [3, 7],
       [4, 7],
       [4, 6],
       [4, 5]])

Runtime test

Other approach(es) -

def Psidom(a,b): # @Psidom's stack based soln
    xv, yv = np.meshgrid(a,b)
    xv[1::2] = xv[1::2,::-1]
    out1 = np.stack((xv, yv), axis=-1).reshape(-1, 2)
    yv[:,1::2] = yv[::-1, 1::2]
    out2 = np.stack((xv.T, yv.T), axis=-1).reshape(-1, 2)
    return out1, out2 

Timings -

In [829]: a = np.arange(100)
     ...: b = np.arange(100)+100

In [830]: %timeit Psidom(a,b)
     ...: %timeit create_zigzag_grids(a,b)
10000 loops, best of 3: 133 µs per loop
10000 loops, best of 3: 25.1 µs per loop

In [831]: a = np.arange(1000)
     ...: b = np.arange(1000)+1000

In [832]: %timeit Psidom(a,b)
     ...: %timeit create_zigzag_grids(a,b)
10 loops, best of 3: 23.8 ms per loop
100 loops, best of 3: 5.2 ms per loop