pythonnumpy

Optimisation of a numpy script to create links between points in a specific format and order


Starting from a situation like the one below (see image):

starting situation

So, to start with, we have a series of points located on the x axis (the distance between the points can vary and the number of points can vary. In the example, we have 5 starting points). We perform a translation along y in several steps (in the example in 3 steps) to arrive at a final y value (in this case 3 units). The numbers in yellow correspond to their indexes.

My aim is to achieve the following:

final situation

In other words, to create links between the points that follow the translation. In practice, and to keep things simple, I'd like to obtain a final array like this:

final : 
 [[ 0  6] [ 6 12] [12 18]
 [ 1  7] [ 7 13] [13 19]
 [ 2  8] [ 8 14] [14 20]
 [ 3  9] [ 9 15] [15 21]
 [ 4 10] [10 16] [16 22]
 [ 5 11] [11 17] [17 23]]

Here's the script I made that does the job:


import numpy as np
# u: number of points along the x axis
u = 6
# v: number of points along the y axis
v = 4

def x_links(u,v):
    matrice = np.arange(u*v).reshape((v, u))
    links =np.array([])
    for i in range(u):
        a = np.repeat(matrice[:,i],2)[1:-1] 
        links=np.append(links,a).astype(int)
    return links.reshape(int(len(links)/2),2)
    
example = x_links(u,v)
print("final : \n",example)

My questions are:

Thank you very much.


Solution

  • This should work:

    matrix = np.arange(u*v).reshape((v,u))
    np.array(np.stack((matrix, np.roll(matrix, v-1, axis = 0)), axis=-1))[:-1]
    

    Output:

    array([[[ 0,  6],
            [ 1,  7],
            [ 2,  8],
            [ 3,  9],
            [ 4, 10],
            [ 5, 11]],
    
           [[ 6, 12],
            [ 7, 13],
            [ 8, 14],
            [ 9, 15],
            [10, 16],
            [11, 17]],
    
           [[12, 18],
            [13, 19],
            [14, 20],
            [15, 21],
            [16, 22],
            [17, 23]]])