Multiply xarray variables with a numpy array-CodePudding

I have an xarray.DataArray:

<xarray.DataArray 'eofs' (mode: 3, lat: 41, lon: 109)>
array([[[ 0.0436546 ,  0.04423243,  0.04479113, ...,  0.02626665,
          0.0257012 ,  0.02513798],
        [ 0.04490738,  0.04549405,  0.04606867, ...,  0.02734869,
          0.026775  ,  0.02619634],
        [ 0.04613635,  0.04672573,  0.0473114 , ...,  0.02844015,
          0.02785003,  0.02725326],
        ...,
        [ 0.05721894,  0.05787764,  0.05852705, ...,  0.0474642 ,
          0.04676448,  0.04606556],
        [ 0.05649778,  0.05714934,  0.05779039, ...,  0.04714761,
          0.04645796,  0.04576794],
        [ 0.05573076,  0.05637409,  0.05700592, ...,  0.04678733,
          0.04610882,  0.04542908]],
]]])
Coordinates:
  * mode     (mode) int32 0 1 2
  * lat      (lat) float64 50.0 50.25 50.5 50.75 51.0 ... 59.25 59.5 59.75 60.0
  * lon      (lon) float64 -2.0 -1.75 -1.5 -1.25 -1.0 ... 24.25 24.5 24.75 25.0

and I want ro multiply each mode with an element of the array array([68.8005905 , 17.8959575 , 8.46729004]). Any idea how to do that?

CodePudding user response：

This is exactly what xarray is intended to help with you're on the right track!

Unlike numpy, which relies on dimension order and positional indices to align arrays, xarray relies on dimension names and coordinate labels. So the key to making sure xarray knows how to align your np.array with your xr.DataArray is to assign the appropriate dimension names and coordinates to the array.

In this case:

multiplier = xr.DataArray(
    [68.8005905 , 17.8959575 ,  8.46729004],
    dims=['mode'],
    coords=[[0, 1, 2]],
)

Now you can simply multiply them! Assuming your DataArray is called da:

# these will be aligned automatically
da * multiplier

See the xarray docs on computation: automatic alignment for more on this topic.

Alternatively, you can still do this operation in numpy by accessing the DataArray.data property:

da.data = da.data * np.array([68.8005905 , 17.8959575 ,  8.46729004]).reshape(-1, 1, 1)

When using direct array manipulation like this, normal numpy broadcasting rules apply. This will in almost all cases be faster than using xarray's automatic alignment, but there are downsides. You cannot assign the result to the same DataArray's data attribute if you change the dimensions, and be careful not to re-order the values - assigning directly to .data like this is playing with fire and removes xarray's ability to ensure data alignment for you - you've been warned! Note that the result of an operation using da.data will be a numpy array (or whatever array backend is in use, e.g. a dask array), so if you do want to dip into numpy and then come back to xarray, you can always just create a new DataArray and assign the dimensions and coordinates using the DataArray constructor as I did above.