How to apply a for loop function to a dataset, but with a different number of loops along a dimension?

[XPR23]

Hi, I need help.
I'm not an expert in xarray, however everything was going fine until I came to the following problem.

I am trying to apply a function to a dataset. The function contains a loop, however, the number of cycles is different for each of the elements of a dimension (nodeID) en el dataset Qhbv. The value of maxbas is the number of cycles, plus there are three values of maxbas for each nodeID.

These are the datasets:

---

<xarray.Dataset> Size: 147GB
Dimensions:  (time: 251304, number: 10, map: 10, nodeID: 1462)
Coordinates:
  * map      (map) int64 80B 0 1 2 3 4 5 6 7 8 9
  * nodeID   (nodeID) int64 12kB 0 1 2 3 4 5 6 ... 1456 1457 1458 1459 1460 1461
  * number   (number) int32 40B 0 1 2 3 4 5 6 7 8 9
  * time     (time) datetime64[ns] 2MB 1940-01-01 ... 2024-12-31T21:00:00
Data variables:
    Qhbv    (time, number, map, nodeID) float32 147GB

---

<xarray.Dataset> Size: 21kB
Dimensions:  (nodeID: 1462, m: 3)
Coordinates:
  * nodeID   (nodeID) int64 12kB 734 703 715 725 704 678 702 ... 25 19 43 18 8 9
  * m        (m) <U3 36B 'max' 'med' 'min'
Data variables:
    maxbas   (nodeID, m) int16 9kB ...

---

FUNCTION

def routinw(Qhbv,maxbas)
    Qhbv_mxb= xr.zeros_like(Qhbv*maxbas)
    for i in range(1,maxbas+1):
        c = -((abs(maxbas-((i-1)*2))-abs(maxbas-(i*2))-4)*maxbas-...  
        Qhbv_i=Qhbv.shift(time=(i-1), fill_value=0)
        Qhbv_mxb+=Qhbv_i*c

I have tried several ways to solve this problem, including apply_ufunc, but I have not been successful.
Do you have any idea?

[dcherian]

Can you make a minimal reproducible example please? https://matthewrocklin.com/minimal-bug-reports

[XPR23]

Hello
Sorry, this is my first time asking a question. Here is a small example of what I need to do.

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

#DATA:
Q = xr.DataArray(
    data=[[[0.1 , 0.15],[0.12, 0.09]],
       [[1.5 , 2.25],[1.8 , 1.3 ]],
       [[3.  , 4.5 ],[3.6 , 2.7 ]],
       [[1.5 , 2.25],[1.8 , 1.3 ]],
       [[0.1 , 0.15],[0.12, 0.09]]],
    dims=["time","a", "b"],
    coords=dict(time=("time", pd.date_range("2001-01-01", freq="D", periods=5)),
                a=("a",[0,1]),b=("b",[0,1]), ),name="q", )

M = xr.DataArray(
    data= [3, 4],
    dims=["b"],
    coords=dict(b=("b",[0,1]), 
    ),name="m", )

#FUNCTION:
def routing(Q,m):
    Qi_ = []
    for i in range(1,m+1):     
        c = -((abs(m-((i-1)*2))-abs(m-(i*2))-4)*m-
                ((i-1)*2)*abs(m-((i-1)*2))+(i*2)*abs(m-(i*2)))/(2*m**2)
        Qi=np.roll(Q,i-1,axis=0)*c
        Qi[0:i-1] = 0        
        Qi_.append(Qi)
    Qm=sum(Qi_)
    return Qm

Q_m=xr.apply_ufunc(routing,Q,M)

If "M" has a single value, this works, for example:
Q_m=xr.apply_ufunc(routing, Q, 3)

However, I have different values ​​of "M" for each element of dimension b,
Q_m=xr.apply_ufunc(routing,Q,M)
In this case, I get an error.

    for i in range(1,m+1):
             ~~~~~^^^^^^^
TypeError: only integer scalar arrays can be converted to a scalar index

[mhungen]

i'm a xarray noob, but you could do this:

done = []

for b_cord in M.b:
    q_sub = Q.sel(b=b_cord)
    m_sub = M.sel(b=b_cord)
    ds_res = xr.apply_ufunc(routing, q_sub, int(m_sub))
    done.append(ds_res)

Q_m = xr.concat(done, dim='b')

at least this is how i would do it. maybe there is a more elegant solution...

[XPR23]

Thanks mhungen

it also worked like this

Qhbv_mxb_=xr.apply_ufunc(routing,Q,M,
                         input_core_dims=[["time"], []],
                         output_core_dims=[["time"]],
                         exclude_dims=set(("time",)),
                         vectorize=True,)