How can I aggregate a 3D NumPy array over time and zones using index arrays?

Below is one solution that “groups by” time first and then does a spatial bin‐summation using a bincount on the flattened spatial data. In this example we use a for‑loop over the unique time bins in tim_idx. For each time bin we sum over time (so we get a (3,5) array) and then use np.bincount together with zon_arr.ravel() (after flattening) as the “bin labels” to produce three sums (one for each zone). For example:

PYTHONCopied
import numpy as np
import matplotlib.pyplot as plt

# Set up spatial and time binning arrays
zon_arr = np.zeros((3, 5))
tim_idx = np.array([0, 0, 1, 1, 2, 2, 3, 3])

# Set up arbitrary zones
zon_arr[1, :3] = 1
zon_arr[1, 3:] = 2

# Generate arbitrary data array:
# First axis is time (length 8, corresponding to tim_idx)
# Last two axes are spatial (3-by-5, same as zon_arr)
np.random.seed(100)
dat_arr = np.random.rand(8, 3, 5)

# Figure out the unique time bins and unique zones
unique_times = np.unique(tim_idx)        # here: [0 1 2 3]
unique_zones = np.unique(zon_arr)          # here: [0. 1. 2.]
n_zones = unique_zones.size

# Prepare an output array: one row per unique time bin and one column per zone.
result = np.empty((unique_times.size, n_zones))

# Loop over time bins: for a given time, sum along the "time" axis and then aggregate spatially
for i, t in enumerate(unique_times):
    # Get all time indices for this group
    # This gives you a (n_time, 3, 5) subarray.
    subdata = dat_arr[tim_idx == t]
    
    # Sum over the time dimension so we have one 3x5 array of spatial totals for this bin
    spatial_sum = subdata.sum(axis=0)
    
    # Use np.bincount to sum up values for each spatial zone.
    # Since zon_arr (which is 3x5) contains the "zone id" for each spatial location,
    # we flatten both arrays. Note that np.bincount expects integer labels.
    # (In this example the zone values are 0., 1., and 2., and we convert them to int.)
    zone_sums = np.bincount(zon_arr.ravel().astype(int), 
                             weights=spatial_sum.ravel(), 
                             minlength=n_zones)
    result[i, :] = zone_sums

print("Result (each row corresponds to a time bin; columns to zones 0,1,2 respectively):")
print(result)

Explanation

1. Temporal grouping:
   The array tim_idx defines the time bin for each time slice (first axis) of dat_arr. In the loop, we select all time slices (using the boolean index tim_idx == t) that belong to a given time bin and then sum along the time axis.

2. Spatial grouping:
   The zon_arr array defines which spatial cells belong to which zone. Because zon_arr is fixed for all time slices, we simply use it (after flattening) together with our summed 2D (3×5) array. The np.bincount function groups the values according to the integer zone labels.

3. Output:
   The resulting “result” array has one row for each unique time bin and three columns corresponding to zones 0, 1, and 2.

This approach accomplishes exactly what you described: for each unique time value (as given by tim_idx) it aggregates (by summing) the spatial values in dat_arr according to the zones in zon_arr.