Process 3d array with 2d zonal array and get statistics through time-CodePudding

I would like to generate statistics for each 'zone' in the example 3D array below. Zones are defined by the 2D array offered below. If the first index in the 3d array represents a time index (and the other two indices represent spatial coordinates), how might I generate a time series of the mean of each zone through time? The final result would be 4 times series (since there are 4 zones - 0's represent an inactive area) of length 10.

Starter code:

import numpy as np

# Create example 3d array
sampl = np.random.uniform(low=0.1, high=10, size=(10,6,6))

# Generate an example zone array
zn = np.array([
    0,1,1,2,2,0,
    1,1,1,2,2,2,
    1,1,1,2,2,2,
    3,3,3,4,4,4,
    3,3,3,4,4,4,
    0,0,3,4,0,0,
])

CodePudding user response：

Since the zones are of different sizes, your target data are jagged and the solution can only be partially vectorised. If the zones were guaranteed to be the same size this would be improved.

import numpy as np
from numpy.random import default_rng

rand = default_rng(seed=0)

zones = np.array((
    (0,1,1,2,2,0),
    (1,1,1,2,2,2),
    (1,1,1,2,2,2),
    (3,3,3,4,4,4),
    (3,3,3,4,4,4),
    (0,0,3,4,0,0),
))
sample = rand.uniform(low=0.1, high=10, size=(10, *zones.shape))

zone_values = np.unique(zones)[1:]
mask_y, mask_x, mask_zone = (
    zones[..., np.newaxis] == zone_values[np.newaxis, np.newaxis, :]
).nonzero()
slices = sample[:, mask_y, mask_x]

for zone in range(len(zone_values)):
    slice = slices[:, mask_zone == zone]
    print(slice.shape, slice.mean())

(10, 8) 5.356025823305826
(10, 8) 4.998502146907645
(10, 7) 5.568793552270735
(10, 7) 5.641964094593679