2121from cmethods .scaling import linear_scaling as __linear_scaling
2222from cmethods .scaling import variance_scaling as __variance_scaling
2323from cmethods .static import SCALING_METHODS
24- from cmethods .utils import UnknownMethodError , check_xr_types
24+ from cmethods .utils import UnknownMethodError , ensure_xr_dataarray
2525
2626if TYPE_CHECKING :
2727 from cmethods .types import XRData
3737
3838def apply_ufunc (
3939 method : str ,
40- obs : XRData ,
41- simh : XRData ,
42- simp : XRData ,
40+ obs : xr . xarray . core . dataarray . DataArray ,
41+ simh : xr . xarray . core . dataarray . DataArray ,
42+ simp : xr . xarray . core . dataarray . DataArray ,
4343 ** kwargs : dict ,
44- ) -> XRData :
44+ ) -> xr . xarray . core . dataarray . DataArray :
4545 """
4646 Internal function used to apply the bias correction technique to the
4747 passed input data.
4848 """
49- check_xr_types (obs = obs , simh = simh , simp = simp )
49+ ensure_xr_dataarray (obs = obs , simh = simh , simp = simp )
5050 if method not in __METHODS_FUNC__ :
5151 raise UnknownMethodError (method , __METHODS_FUNC__ .keys ())
5252
@@ -96,11 +96,11 @@ def apply_ufunc(
9696
9797def adjust (
9898 method : str ,
99- obs : XRData ,
100- simh : XRData ,
101- simp : XRData ,
99+ obs : xr . xarray . core . dataarray . DataArray ,
100+ simh : xr . xarray . core . dataarray . DataArray ,
101+ simp : xr . xarray . core . dataarray . DataArray ,
102102 ** kwargs ,
103- ) -> XRData :
103+ ) -> xr . xarray . core . dataarray . DataArray | xr . xarray . core . dataarray . Dataset :
104104 """
105105 Function to apply a bias correction technique on single and multidimensional
106106 data sets. For more information please refer to the method specific
@@ -119,19 +119,19 @@ def adjust(
119119 :param method: Technique to apply
120120 :type method: str
121121 :param obs: The reference/observational data set
122- :type obs: XRData
122+ :type obs: xr.xarray.core.dataarray.DataArray
123123 :param simh: The modeled data of the control period
124- :type simh: XRData
124+ :type simh: xr.xarray.core.dataarray.DataArray
125125 :param simp: The modeled data of the period to adjust
126- :type simp: XRData
126+ :type simp: xr.xarray.core.dataarray.DataArray
127127 :param kwargs: Any other method-specific parameter (like
128128 ``n_quantiles`` and ``kind``)
129129 :type kwargs: dict
130130 :return: The bias corrected/adjusted data set
131- :rtype: XRData
131+ :rtype: xr.xarray.core.dataarray.DataArray | xr.xarray.core.dataarray.Dataset
132132 """
133133 kwargs ["adjust_called" ] = True
134- check_xr_types (obs = obs , simh = simh , simp = simp )
134+ ensure_xr_dataarray (obs = obs , simh = simh , simp = simp )
135135
136136 if method == "detrended_quantile_mapping" : # noqa: PLR2004
137137 raise ValueError (
@@ -169,6 +169,8 @@ def adjust(
169169 obs_group = group ["obs" ]
170170 simh_group = group ["simh" ]
171171 simp_group = group ["simp" ]
172+ else :
173+ raise ValueError ("'group' must be a string or a dict!" )
172174
173175 del kwargs ["group" ]
174176
0 commit comments