2

I have two numpy ndarrays, array1 and array 2, with array1.shape = array2.shape = (n, l, m).

A 3rd ndarray is initialized as array3 = np.nan * np.zeros((n-1, l, m + 1)) and is then computed using the following for loop:

for i in range(m):
   array3[:n - i - 1, :, i] = array1[i + 1:, :, i] - array2[:n - i - 1, :, i]

Is there a simple way to vectorize this and avoid the for loop ?

Here is a simple example:

import numpy as np

a = np.ones((6, 4, 4)) * np.arange(1, 5)
b = np.ones((6, 4, 4))
c = np.nan * np.zeros((5, 4, 4))
n = a.shape[0]
m = a.shape[2]

for i in range(m):
   c[:n - i - 1, :, i] = a[i + 1:, :, i] - b[:n - i - 1, :, i]
WinstonSmith
  • 21
  • 3
  • 1
    No. Some times loops are unavoidable. In case you need to improve speed, you can use something like numba. –  Jul 26 '22 at 07:13

1 Answers1

0

I tried rewriting array a the following way:

a = np.concatenate((a, np.nan * np.zeros((4, 4, 4))), axis=0)

row_idx, column_idx, slice_idx = np.ogrid[:a.shape[0], :a.shape[1], :a.shape[2]]

r = -1.0 * np.arange(1, 5) + a.shape[0]

row_idx = (row_idx - r[np.newaxis, np.newaxis, :]).astype(int)

a = a[row_idx, column_idx, slice_idx]
a = a[:6, :, :]

and then subtract array b directly but it was only marginally faster for arrays of larger size.

WinstonSmith
  • 21
  • 3