0

This is closely related to an earlier thread but I am having some trouble reproducing the solution there with a 3d array. My 3d array has one time coordinate and two spatial coordinates. Each spatial coordinate is itself a two-dimensional function of the indices along the spatial dimensions, as might be the case with (longitude,latitude) values when the spatial grid is curvilinear.

Goal: For each time, I want to find the coordinates corresponding to the maximum in space. Here is some code:

# Create cftime object to mimic my real data.
yr1=250; yr2=250
times = xr.cftime_range(start=format(yr1,'04d') + '-01',
                                        end  =format(yr2,'04d') + '-02',
                                        freq='1MS',calendar='noleap')

x=np.arange(6).reshape(2,3); y = x;

nt=len(times); ny=x.shape[0]; nx=x.shape[1]

# Create the dataarray object
dummy = xr.DataArray(data=np.random.random([nt,ny,nx]),dims=['t','iy','ix'], 
                     coords=dict( time=(['t'],times), x=(['iy','ix'],x), y=(['iy','ix'],y) ) )

dummy.where(dummy==dummy.max(dim=['ix','iy']),drop=True ).squeeze()

See the pasted image below to see the output from the last command. It does not look like the coordinates with nan values are being dropped. At each time, the array has size 2 x 2, when it should be 2 x 3 with the max. value and 5 nan values. Finally, if the drop command were working, I should be seeing only two scalar numbers in the output of cell 78.

What am I missing?

enter image description here

1 Answers1

0

This is tricky, I can't figure out what to expect from .where when coordinates are functions of the indices. Have you thought of using .stack as a workaround? If I understand it correctly stacking the array on 'ix' and 'ix' and taking .idxmax() should give you the iy and ix indices you are looking for.

    dummy_stacked=dummy.stack(z=['ix','iy'])
    dummy_stacked.idxmax('z')
Mathi
  • 312
  • 1
  • 8