Multiply/product along a dimension xarray

Adele Zaini 0
Viewed 25

I was looking for a very easy and elegant way to multiply two DataArrays, that may have different dimensions, along a single axis.

My specific case: a first DataArray A has dimensions (lat, lon, natpft) and the second B (lat, natpft). My purpose is to have the product of A data in (lat, lon) multiplied by B values along lat for each natpft. If this operation could also include the nearest option in choosing which latitude to match, it would be great.

I have some ideas in mind but I guess there could be a single line code that can do this task. I also hope this could help someone else, since I couldn't find any question about this topic.

product python-xarray coordinates multiplication python
1 Answers

Not really a one line solution here, but I think it solves your problem?

First construct some data according to your description (I would suggest you to do this in your post next time).

import numpy as np
import xarray as xr
import pandas as pd

# construct "natpft_1"
natpft_1 = np.random.randn(2, 3)

# create coords
longitude_1 = [1,2]
latitude_1 = [1,2,3]

# put data into a dataset
ds1 = xr.Dataset(data_vars=dict(natpft_1=(["x", "y"], natpft_1)),
                 coords=dict(lon=(["x"], longitude_1),
                             lat=(["y"], latitude_1)),
                 attrs=dict(description="natpft_1"))

# construct "natpft_2" along some dummy latitudes
natpft_2 = np.random.randn(1, 5)[0]

latitude_2 = [0.9,2.2,3.3,4.5,5.3]

# put data into a dataset
ds2 = xr.Dataset(data_vars=dict(natpft_2=(["y"], natpft_2)),
                 coords=dict(lat=(["y"], latitude_2)),
                 attrs=dict(description="natpft_2"))

Then solve your problem:

ds3 = ds1.copy()

for i in range(len(ds3['lon'])):
    for j in range(len(ds3['lat'])):
        nearest_lat_ds2_index = np.argmin(ds2['lat'].values-ds3['lat'].values[j])
        ds3['natpft_1'][i,j] = ds3['natpft_1'][i,j]*ds2['natpft_2'][nearest_lat_ds2_index]


# check results
print(ds3)
userFNkF0
Related