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

Question

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]])

How big is your array in practice (ie. shape)? The current solution seems quite fine for a Numpy code. — Jérôme Richard
– Jérôme Richard, Commented Aug 30, 2022 at 17:41

mozway · Accepted Answer · 2022-08-30 17:34:09Z

3

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 Aug 30, 2022 at 17:34

mozway

267k13 gold badges55 silver badges106 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Quang Hoang Over a year ago

I was going for: return a + (np.arange(n) * step)[:,None]

mozway Over a year ago

@Quang it looks like it's similar in terms of timing. I naively thought using directly the correct step in arange would be faster.

Collectives™ on Stack Overflow

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

1 Answer 1

2 Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Related