xarray data set with two variables

[ghost]

i am attempting to write a python version of the function MBCn so i can use xarray to carry out the computation over a a gridded file.

the R function accepts matrices (size tn by varn) as inputs where tn is the number of timesteps and varn is the number of variables of the dataset. However that same function, great as it is, only accepts point values (inputs for one (lon,lat) pair).

My code is as follows:

import numpy as np
import xarray as xr
from scipy.stats import rankdata
from scipy import linalg


def rotationMatrix(nvars):
    # QR decomposition
    P = np.random.normal(size=(nvars,nvars))
    Q,R = linalg.qr(P)
    diagR = np.diag(R)    
    return Q @ np.diag(np.divide(diagR,np.abs(diagR)))    

def timeParser(DataSet):
    _time_string, _coords = 'time', list(DataSet.coords)
    _time_list = [j for j in _coords if _time_string in j]
    
    if (len(_time_list) > 1): return sorted(_time_list)
    else: return _time_list.pop()
     
def varCollect(DataSet):
    _var_keys, _coords = list(DataSet.variables.keys()), list(DataSet.coords)
    return sorted(list(set(_var_keys).difference(set(_coords))))

def applyRotation(DataSet,mtx):
    _vars = varCollect(DataSet)
    _d_rot = DataSet.to_dataframe()[_vars]
    _d_rot = _d_rot @ mtx
    _d_rot.columns = _vars
    if ( DataSet[_vars[0]][0].shape == () ): 
        return xr.Dataset.from_dataframe(_d_rot).assign_coords({'lon':DataSet.lon.values,'lat':DataSet.lat.values})
    else: return xr.Dataset.from_dataframe(_d_rot)

def reorder(data_a,data_b):
    _a,_b = (rankdata(data_a,method='ordinal')-1).astype(int), sorted(data_b)
    res = [ _b[i] for i in _a ]
    return np.array(res)
      
def nMultiVarBC(obs_h,mod_h,mod_p,adj_proj=False,rot_seq=30):
    _time_sgn,  var_list,  = timeParser(obs_h), varCollect(obs_h)
    nvars = len(var_list)
    
    #Time slices
    hDt_o,hDt_f = mod_h[_time_sgn].dt.year.values[0],mod_h[_time_sgn].dt.year.values[-1]
    pDt_o,pDt_f = mod_p[_time_sgn].dt.year.values[0],mod_p[_time_sgn].dt.year.values[-1]
    
    # Standardize 
    ## observations
    obs_h_scale = xr.apply_ufunc(lambda dat,mu,sig:(dat-mu)/sig, obs_h,obs_h.mean(_time_sgn),obs_h.std(_time_sgn))
    
    ## model
    mod_stitch = xr.concat([mod_h,mod_p],dim=_time_sgn)
    _mu,_sig = mod_stitch.dropna(_time_sgn).mean(_time_sgn), mod_stitch.dropna(_time_sgn).std(_time_sgn)
    mod_scale = xr.apply_ufunc(lambda dat,mu,sig:(dat-mu)/sig, mod_stitch,_mu,_sig)
   
    ROMs = [rotationMatrix(nvars).round(2) for i in np.arange(rot_seq)]
    
    for (k,rot) in zip(np.arange(rot_seq),ROMs):
        _dset_h =  []
        Z_o_h = applyRotation(obs_h_scale,rot)
        if ( k == 0 ): 
            print(f"{k+1}", end = ' ')
            Z_m_h = applyRotation(mod_scale.sel(time=slice(str(hDt_o),str(hDt_f))),rot)
            if adj_proj: 
                _dset_p, Z_m_p = [], applyRotation(mod_scale.sel(time=slice(str(pDt_o),str(pDt_f))),rot)
        elif ( k > 0 ):
            print(f"{k+1}", end = ' ')
            Z_m_h = applyRotation(Z_m_h,rot)
            if adj_proj: 
                _dset_p, Z_m_p = [], applyRotation(Z_m_p,rot)
        for var in var_list:
            Z_h_qdm = quantileDeltaMapping(Z_o_h[var],Z_m_h[var],ratio=False)
            Z_h_qdm = xr.Dataset({var:Z_h_qdm})
            _dset_h.append(Z_h_qdm)
            if adj_proj: 
                Z_p_qdm = quantileDeltaMapping(Z_o_h[var],Z_m_h[var],Z_m_p[var],ratio=False)
                Z_p_qdm = xr.Dataset({var:Z_p_qdm})
                _dset_p.append(Z_p_qdm)
                
        Z_m_h = applyRotation(xr.merge(_dset_h),rot.T)
        if adj_proj:
            Z_m_p = applyRotation(xr.merge(_dset_p),rot.T)
            Z_m = Z_m_p
            
        else: Z_m = Z_m_h
            

    mod_inter = xr.apply_ufunc(lambda dat,mu,sig: (dat*sig)+mu,Z_m,_mu,_sig)
    
    if adj_proj: mod_dat = mod_p
    else: mod_dat = mod_h
    
    _dset = []
    for var in var_list:
        mod_corr = xr.apply_ufunc(reorder,mod_inter[var],mod_dat[var],input_core_dims=[['time'],['time']],
                                   output_core_dims=[['time']],dask='parallelized',vectorize=True,
                                   output_dtypes=[np.float])
        mod_corr = xr.Dataset({var:mod_corr})
        _dset.append(mod_corr)
        
    return xr.merge(_dset).round(2)

The images below is a comparison of the code above with official R version for one iteration:

(top: rainfall, bottom: temp)

However over 30 iterations (the default number of iterations in the R version) the python script deteriorates but R is stable:

(top: rainfall, bottom: temp)

The function quantileDeltaMapping attached below as BA.txt has its own version in R which, after triple-checking, duplicates the R version perfectly so there is no problem there. I am assuming that the problem lies entirely with the piece of code I wrote above.

I should also mention that the R version performs a bias adjustment (called QDM) if the model inputs (mod.c and mod.p in R code) has not been adjusted yet via QDM so the inputs I passed to nMultivarBC has already been adjusted (datasets i used are attached as a zip file)

BA.txt
(attached
data_files.zip
I have specified an extra argument in the function nMultiVarBC above so the user can specify if they wish to adjust historical or model projections need to adjusted (i did this to minimize the computational burden that would be encountered when inputting gridded files)

May someone kindly point out how and why things go awry after several iterations?
Your help would be greatly appreciated.

best,
oxdub