directly pass {g,gw}_idxs

They were previously passed together in a common dataset, but I had a weird bug, not happening all the time, where win_dim dimension would be reduced to the smallest size, the one g_idxs. But gw_idxs is supposed to have a larger dimension. Perhaps the problem is that g_idxs should have the same win_dim as gw_idxs, but fill all irrelevant entries with nans. This would avoid the problem. Perhaps in the future I should try this instead. For now it works
Ouranosinc · Feb 27, 2025 · 90b2ff0 · 90b2ff0
1 parent 5d68d66
commit 90b2ff0
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/src/xsdba/_adjustment.py b/src/xsdba/_adjustment.py
@@ -343,10 +343,12 @@ def _npdft_adjust(sim, af_q, rots, quantiles, method, extrap):
 
 
 def mbcn_adjust(
-    ref: xr.Dataset,
-    hist: xr.Dataset,
-    sim: xr.Dataset,
+    ref: xr.DataArray,
+    hist: xr.DataArray,
+    sim: xr.DataArray,
     ds: xr.Dataset,
+    g_idxs: xr.DataArray,
+    gw_idxs: xr.DataArray,
     pts_dims: Sequence[str],
     interp: str,
     extrapolation: str,
@@ -370,12 +372,14 @@ def mbcn_adjust(
         training data.
     sim : xr.DataArray
         data to adjust (stacked with multivariate dimension).
+    g_idxs : xr.DataArray
+        Indices of the times in each time group.
+    gw_idxs: xr.DataArray
+        Indices of the times in each windowed time group.
     ds : xr.Dataset
         Dataset variables:
             rot_matrices : Rotation matrices used in the training step.
             af_q : Adjustment factors obtained in the training step for the npdf transform
-            g_idxs : Indices of the times in each time group
-            gw_idxs: Indices of the times in each windowed time group
     pts_dims : [str, str]
         The name of the "multivariate" dimension and its primed counterpart. Defaults to "multivar", which
         is the normal case when using :py:func:`xsdba.stack_variables`, and "multivar_prime".
@@ -405,8 +409,6 @@ def mbcn_adjust(
     rot_matrices = ds.rot_matrices
     af_q = ds.af_q
     quantiles = af_q.quantiles
-    g_idxs = ds.g_idxs
-    gw_idxs = ds.gw_idxs
     gr_dim = gw_idxs.attrs["group_dim"]
     win = gw_idxs.attrs["group"][1]
 

diff --git a/src/xsdba/adjustment.py b/src/xsdba/adjustment.py
@@ -1912,14 +1912,14 @@ def _adjust(
 
         g_idxs, gw_idxs = grouped_time_indexes(ref.time, self.group)
         ds = self.ds.copy()
-        ds["g_idxs"] = g_idxs
-        ds["gw_idxs"] = gw_idxs
 
         # adjust (adjust for npft transform, train/adjust for univariate bias correction)
         out = mbcn_adjust(
             ref=ref,
             hist=hist,
             sim=sim,
+            g_idxs=g_idxs,
+            gw_idxs=gw_idxs,
             ds=ds,
             pts_dims=self.pts_dims,
             interp=self.interp,