@@ -60,21 +60,21 @@ def apply_ufunc(
6060 'input_core_dims must have three key-value pairs like: {"obs": "time", "simh": "time", "simp": "time"}' ,
6161 )
6262
63- input_core_dims = kwargs [ "input_core_dims" ]
63+ input_core_dims = kwargs . pop ( "input_core_dims" )
6464 else :
6565 input_core_dims = {"obs" : "time" , "simh" : "time" , "simp" : "time" }
6666
6767 result : XRData = xr .apply_ufunc (
6868 __METHODS_FUNC__ [method ],
6969 obs ,
7070 simh ,
71- # Need to spoof a fake time axis since 'time' coord on full dataset is different
72- # than 'time' coord on training dataset.
71+ # Need to spoof a fake time axis since 'time' coord on full dataset is
72+ # different than 'time' coord on training dataset.
7373 simp .rename ({input_core_dims ["simp" ]: "__t_simp__" }),
7474 dask = "parallelized" ,
7575 vectorize = True ,
76- # This will vectorize over the time dimension, so will submit each grid cell
77- # independently
76+ # This will vectorize over the time dimension, so will submit each grid
77+ # cell independently
7878 input_core_dims = [
7979 [input_core_dims ["obs" ]],
8080 [input_core_dims ["simh" ]],
@@ -89,9 +89,9 @@ def apply_ufunc(
8989 # Rename to proper coordinate name.
9090 result = result .rename ({"__t_simp__" : input_core_dims ["simp" ]})
9191
92- # ufunc will put the core dimension to the end (time), so want to preserve original
93- # order where time is commonly first.
94- return result .transpose (* obs .dims )
92+ # ufunc will put the core dimension to the end (time), so want to preserve
93+ # original order where time is commonly first.
94+ return result .transpose (* obs .rename ({ input_core_dims [ "obs" ]: input_core_dims [ "simp" ]}). dims )
9595
9696
9797def adjust (
0 commit comments