Finding a way to iterate using the input of two xarray dataarrays when chunked

Technical
4

Thank Anton and Jemma,
Those are excellent responses. However, I do not think I can use cumsum. It is because I did not explain well enough in my initial example. So entirely my fault.

After the summing of storage a bunch of conditions need to be met, such as the value of arr_storage not being able to be larger than my_max (see below). Hence, the cumsum solution is too simple.

da = xr.DataArray(
    np.random.randint(low=-5, high=5, size=(10, 100, 100)),
    coords=[range(10), range(100), range(100)],
    dims=["time", "x", "y"],
).chunk(chunks={"x":10, "y":10})

my_max = 10
def sum_storage(arr):
    arr_storage = xr.zeros_like(arr)
    for idx in range(1,len(arr)):
        # print(f'idx: {idx}')
        arr_storage[dict(time=idx)] = arr_storage[dict(time=idx-1)] + da[dict(time=idx)]
        arr_storage[dict(time=idx)] = arr_storage[dict(time=idx)].where(arr_storage[dict(time=idx)] <= my_max, my_max)
    return arr_storage

%time arr_storage = sum_storage(da)

I am now thinking of a solution of manually chunking the parts when they outgrow memory. But that does not sound ideal, i.e. I don’t want to write a lot of manual code when a possible chunked solution is also available…

Thanks for your time, hope you can look at it again?