[FIXED] Is it possible to make this function on numpy array more efficient?

Issue

Here a is a 1D array of integer indices. To give some context, a is the atom indices of the 1st molecules. The return is the atom indices for n identical molecules, each of which contains step atoms. This function basically applies the same atom selection to many molecules

def f(a, step, n):
    a.shape=1,-1
    got = np.repeat(a, n, axis=0)
    blocks = np.arange(0, n*step, step)
    blocks.shape = n, -1
    return got + blocks

For example

In [52]: f(np.array([1,3,4]), 10, 4)
Out[52]:
array([[ 1,  3,  4],
       [11, 13, 14],
       [21, 23, 24],
       [31, 33, 34]])

Solution

It looks like broadcasting should be enough:

def f(a, step, n):
    return a + np.arange(0, n*step, step)[:, None]

f(np.array([1,3,4]), 10, 4)

output:

array([[ 1,  3,  4],
       [11, 13, 14],
       [21, 23, 24],
       [31, 33, 34]])

Answered By - mozway