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