diff --git a/xclim/indices/fire/_cffwis.py b/xclim/indices/fire/_cffwis.py index 8012e7ebc..7d1e99db5 100644 --- a/xclim/indices/fire/_cffwis.py +++ b/xclim/indices/fire/_cffwis.py @@ -158,9 +158,9 @@ "overwintering_drought_code", ] -default_params = dict( - temp_start_thresh=(12, "degC"), - temp_end_thresh=(5, "degC"), +default_params: dict[str, int | float | tuple[float, str]] = dict( + temp_start_thresh=(12.0, "degC"), + temp_end_thresh=(5.0, "degC"), snow_thresh=(0.01, "m"), temp_condition_days=3, snow_condition_days=3, @@ -388,9 +388,13 @@ def _duff_moisture_code( @vectorize(nopython=True) -def _drought_code( - t: np.ndarray, p: np.ndarray, mth: np.ndarray, lat: float, dc0: float -) -> np.ndarray: # pragma: no cover +def _drought_code( # pragma: no cover + t: np.ndarray, + p: np.ndarray, + mth: np.ndarray, + lat: float, + dc0: float, +) -> np.ndarray: """Compute the drought code over one time step. Parameters @@ -411,10 +415,10 @@ def _drought_code( array_like Drought code at the current timestep """ - fl = _day_length_factor(lat, mth) + fl = _day_length_factor(lat, mth) # type: ignore if t < -2.8: - t = -2.8 + t = -2.8 # type: ignore pe = (0.36 * (t + 2.8) + fl) / 2 # *Eq.22*# pe = max(pe, 0.0) @@ -431,7 +435,7 @@ def _drought_code( dc = pe else: # f p <= 2.8: dc = dc0 + pe - return dc + return dc # type: ignore def initial_spread_index(ws: np.ndarray, ffmc: np.ndarray) -> np.ndarray: @@ -451,7 +455,7 @@ def initial_spread_index(ws: np.ndarray, ffmc: np.ndarray) -> np.ndarray: """ mo = 147.2 * (101.0 - ffmc) / (59.5 + ffmc) # *Eq.1*# ff = 19.1152 * np.exp(mo * -0.1386) * (1.0 + (mo**5.31) / 49300000.0) # *Eq.25*# - isi = ff * np.exp(0.05039 * ws) # *Eq.26*# + isi: np.ndarray = ff * np.exp(0.05039 * ws) # *Eq.26*# return isi @@ -503,7 +507,7 @@ def fire_weather_index(isi, bui): return fwi -def daily_severity_rating(fwi: np.ndarray) -> np.ndarry: +def daily_severity_rating(fwi: np.ndarray) -> np.ndarray: """Daily severity rating. Parameters @@ -548,6 +552,7 @@ def _overwintering_drought_code(DCf, wpr, a, b, minDC): # pragma: no cover # SECTION 2 : Iterators +# FIXME: default_params should be supplied within the logic of the function. def _fire_season( tas: np.ndarray, snd: np.ndarray | None = None, @@ -1056,15 +1061,16 @@ def fire_weather_ufunc( # noqa: C901 ) # Arg order : tas, pr, hurs, sfcWind, snd, mth, lat, season_mask, dc0, dmc0, ffmc0, winter_pr # 0 1 2 3 4 5 6 7 8 9 10 11 - args = [None] * 12 - input_core_dims = [[]] * 12 + args: list[xr.DataArray | None] = [None] * 12 + input_core_dims: list[list[str | None]] = [[]] * 12 # Verification of all arguments for i, (arg, name, usedby, has_time_dim) in enumerate(needed_args): if any([ind in indexes + [season_method] for ind in usedby]): if arg is None: raise TypeError( - f"Missing input argument {name} for index combination {indexes} with fire season method '{season_method}'" + f"Missing input argument {name} for index combination {indexes} " + f"with fire season method '{season_method}'." ) args[i] = arg input_core_dims[i] = ["time"] if has_time_dim else [] @@ -1078,17 +1084,14 @@ def fire_weather_ufunc( # noqa: C901 raise ValueError("'dry_start' must be one of None, 'CFS' or 'GFWED'.") # Always pass the previous codes. - if dc0 is None: - dc0 = xr.full_like(tas.isel(time=0), np.nan) - if dmc0 is None: - dmc0 = xr.full_like(tas.isel(time=0), np.nan) - if ffmc0 is None: - ffmc0 = xr.full_like(tas.isel(time=0), np.nan) - args[8:11] = [dc0, dmc0, ffmc0] + _dc0 = xr.full_like(tas.isel(time=0), np.nan) if dc0 is None else dc0 + _dmc0 = xr.full_like(tas.isel(time=0), np.nan) if dmc0 is None else dmc0 + _ffmc0 = xr.full_like(tas.isel(time=0), np.nan) if ffmc0 is None else ffmc0 + args[8:11] = [_dc0, _dmc0, _ffmc0] # Output config from the current indexes list outputs = indexes - output_dtypes = [tas.dtype] * len(indexes) + output_dtypes: list[np.dtype] = [tas.dtype] * len(indexes) output_core_dims = len(indexes) * [("time",)] if season_mask is not None: diff --git a/xclim/indices/fire/_ffdi.py b/xclim/indices/fire/_ffdi.py index b986fe2e3..2acc607f6 100644 --- a/xclim/indices/fire/_ffdi.py +++ b/xclim/indices/fire/_ffdi.py @@ -41,7 +41,7 @@ nopython=True, cache=True, ) -def _keetch_byram_drought_index(p, t, pa, kbdi0, kbdi: float): # pragma: no cover +def _keetch_byram_drought_index(p, t, pa, kbdi0, kbdi): # pragma: no cover """Compute the Keetch-Byram drought (KBDI) index. Parameters @@ -239,7 +239,7 @@ def keetch_byram_drought_index( :cite:cts:`ffdi-keetch_1968,ffdi-finkele_2006,ffdi-holgate_2017,ffdi-dolling_2005` """ - def _keetch_byram_drought_index_pass(pr, tasmax, pr_annual, kbdi0): + def _keetch_byram_drought_index_pass(_pr, _tasmax, _pr_annual, _kbdi0): """Pass inputs on to guvectorized function `_keetch_byram_drought_index`. This function is actually only required as `xr.apply_ufunc` will not receive @@ -249,7 +249,7 @@ def _keetch_byram_drought_index_pass(pr, tasmax, pr_annual, kbdi0): -------- DO NOT CALL DIRECTLY, use `keetch_byram_drought_index` instead. """ - return _keetch_byram_drought_index(pr, tasmax, pr_annual, kbdi0) + return _keetch_byram_drought_index(_pr, _tasmax, _pr_annual, _kbdi0) pr = convert_units_to(pr, "mm/day", context="hydro") tasmax = convert_units_to(tasmax, "C") @@ -259,7 +259,7 @@ def _keetch_byram_drought_index_pass(pr, tasmax, pr_annual, kbdi0): else: kbdi0 = xr.full_like(pr.isel(time=0), 0) - kbdi = xr.apply_ufunc( + kbdi: xr.DataArray = xr.apply_ufunc( _keetch_byram_drought_index_pass, pr, tasmax, @@ -270,7 +270,7 @@ def _keetch_byram_drought_index_pass(pr, tasmax, pr_annual, kbdi0): dask="parallelized", output_dtypes=[pr.dtype], ) - kbdi.attrs["units"] = "mm/day" + kbdi = kbdi.assign_attrs(units="mm/day") return kbdi @@ -317,7 +317,7 @@ def griffiths_drought_factor( :cite:cts:`ffdi-griffiths_1999,ffdi-finkele_2006,ffdi-holgate_2017` """ - def _griffiths_drought_factor_pass(pr, smd, lim): + def _griffiths_drought_factor_pass(_pr, _smd, _lim): """Pass inputs on to guvectorized function `_griffiths_drought_factor`. This function is actually only required as xr.apply_ufunc will not receive @@ -327,7 +327,7 @@ def _griffiths_drought_factor_pass(pr, smd, lim): -------- DO NOT CALL DIRECTLY, use `griffiths_drought_factor` instead. """ - return _griffiths_drought_factor(pr, smd, lim) + return _griffiths_drought_factor(_pr, _smd, _lim) pr = convert_units_to(pr, "mm/day", context="hydro") smd = convert_units_to(smd, "mm/day") @@ -339,7 +339,7 @@ def _griffiths_drought_factor_pass(pr, smd, lim): else: raise ValueError(f"{limiting_func} is not a valid input for `limiting_func`") - df = xr.apply_ufunc( + df: xr.DataArray = xr.apply_ufunc( _griffiths_drought_factor_pass, pr, smd, @@ -349,7 +349,7 @@ def _griffiths_drought_factor_pass(pr, smd, lim): dask="parallelized", output_dtypes=[pr.dtype], ) - df.attrs["units"] = "" + df = df.assign_attrs(units="") # First non-zero entry is at the 19th time point since df is calculated # from a 20-day rolling window. Make prior points NaNs. diff --git a/xclim/indices/generic.py b/xclim/indices/generic.py index f594b6b3e..76d987024 100644 --- a/xclim/indices/generic.py +++ b/xclim/indices/generic.py @@ -1061,7 +1061,9 @@ def _get_zone(_da): return zones -def detrend(ds, dim="time", deg=1) -> xr.Dataset | xr.DataArray: +def detrend( + ds: xr.DataArray | xr.Dataset, dim="time", deg=1 +) -> xr.DataArray | xr.Dataset: """Detrend data along a given dimension computing a polynomial trend of a given order. Parameters @@ -1075,7 +1077,7 @@ def detrend(ds, dim="time", deg=1) -> xr.Dataset | xr.DataArray: Returns ------- - detrended : xr.Dataset or xr.DataArray + xr.Dataset or xr.DataArray Same as `ds`, but with its trend removed (subtracted). """ if isinstance(ds, xr.Dataset):