CODEWITHSUNDEEP
arraysnumpyvariable-assignmentdaskpython-xarray
Rens Vermeltfoort
Rens Vermeltfoort

Reputation: 21

Workaround for Item Assignment in Dask

I have a piece of code in a model that transforms arrays based on the value of another array one day prior. If the maximum temperature for the current day is lower than the minimum temperature from the day before, the value should be changed to the value of the minimum temperature of the day before the current date. Likewise when the minimum temperature of the before the current day for the current day is higher than the maximum temperature from the day before, the value should be changed to the value of the maximum temperature of the day before the current date. I have two versions that both give the same error, as the Dask arrays do not support item assignment:

for ens in range(0,num_ens):
    for lat in range(0, num_lats):
        for lon in range(0, num_lons):
            for day in range(1, 240):
                if max_ens[ens, lat, lon, day] < min_ens[ens, lat, lon, day-1]:
                    max_ens[ens, lat, lon, day] = min_ens[ens, lat, lon, day-1]
                if min_ens[ens, lat, lon, day] > max_ens[ens, lat, lon, day-1]:
                    min_ens[ens, lat, lon, day] = max_ens[ens, lat, lon, day-1]

Vectorized version:

max_ens[:, :, :, day] = np.maximum(max_ens[:, :, :, day], min_ens[:, :, :, day - 1])
min_ens[:, :, :, day] = np.minimum(min_ens[:, :, :, day], max_ens[:, :, :, day - 1])

Does anyone know how to efficiently do this in Dask?

Thanks in advance!

EDIT:

I have been trying to implement the where function but the shape of the arrays is adjusted when the where function is applied.

Upvotes: 0

Views: 83

Answers (1)

MRocklin
MRocklin

Reputation: 57281

I would try to solve this problem with just numpy without using for loops or in-place assignment. My guess is that it's doable with np.where, np.roll, and np.stack, or maybe clever broadcasting

Upvotes: 1

Related Questions