Python Time-Series Calculations in 3D Arrays without using Nested Loops?

Question

I'm have a big 3D array, where each layer is data from one time. I want to compute time-series statistics at each grid cell through time. The code block below is a greatly over-simplified version.

My ancient FORTRAN/C background says that I should use nested loops, but I can't help but think there must be a more "Pythonic" way of doing this (perhaps a lot faster, too?). Any help?

Thanks

# Brute-force time-series stats

import numpy as np

rows = 2
cols = 4
dates = 3

time_series = np.empty([rows, cols, dates])
samples = np.zeros([rows, cols])
means = np.zeros([rows, cols])

for r in range(rows):
    for c in range(cols):
        
        samples[r, c] = np.count_nonzero(~np.isnan(time_series[r, c, :]))
        means[r, c] = np.nanmean(time_series[r, c, :])

print("Means at each grid cell:", means)
print("Samples:", samples)

question from:https://stackoverflow.com/questions/65928977/python-time-series-calculations-in-3d-arrays-without-using-nested-loops

Answer 1

Both numpy.count_nonzero and numpy.nanmean take in axis parameter to specify along which axis it should calculate these values, so if you pass axis = 2 to each (2 being the 3rd axis, i.e.(0,1,2)), you don't need the loops.

samples = np.count_nonzero(~np.isnan(time_series), axis = 2)
means = np.nanmean(time_series, axis = 2)

Will give the same results as your double loop.

Checked both with %%timeit on my jupyter notebook and using numpy axis parameter is indeed at least 5 times faster than the loops (this will depend on the complexity of your arrays, but from my few attempts the gain in performance time gets better with increased complexity).

Results: