Numpy next element minus previous

Question

I would like to subtract the next element from current element in a specific axis from numpy array. But, I know how to do that using a lot of loops. My question is: How to do that in the most efficient way? Maybe using numpy?

My Python code below:

import numpy as np

np.random.seed(0)

myarr = np.random.rand(20, 7, 11, 151, 161)

newarr = np.full((20, 6, 11, 151, 161), np.nan)

for iyr in range(20):
    for iwk in range(6):
        for imb in range(11):
            for ilat in range(151):
                for ilon in range(161):
                     newarr[iyr, iwk, imb, ilat, ilon] = myarr[iyr, iwk + 1, imb, ilat, ilon] - myarr[iyr, iwk, imb, ilat, ilon]

Answer 1

There are a few good ways of doing this. If you did not care about the last element being NaN, you could use np.diff

myarr = np.random.rand(20, 7, 11, 151, 161)
newarr = np.diff(myarr, axis=1)

The result will have shape (20, 6, 11, 151, 161).

If you really want to keep those NaNs, I would recommend using np.empty_like and np.subtract. Allocating with np.full is somewhat wasteful since you are definitely setting almost all the elements. The only NaNs you have are in the last index along the second dimension, which you can initialize very cheaply yourself:

myarr = np.random.rand(20, 7, 11, 151, 161)
newarr = np.empty_like(myarr)  # Don't repeat shape
newarr[:, -1, ...] = np.nan
np.subtract(myarr[:, 1:, ...], myarr[:, :-1, ...], out=newarr[:, :-1, ...])

Since myarr[:, 1:, ...], myarr[:, :-1, ...] and newarr[:, :-1, ...] are views, this operation avoids temporary arrays and unnecessary initialization almost completely.