diff --git a/hdc/algo/_version.py b/hdc/algo/_version.py index b0187ac..7dd4300 100644 --- a/hdc/algo/_version.py +++ b/hdc/algo/_version.py @@ -1,2 +1,2 @@ """Version only in this file.""" -__version__ = "0.2.1" +__version__ = "0.3.0" diff --git a/hdc/algo/accessors.py b/hdc/algo/accessors.py index 8f8f35c..51a91bb 100644 --- a/hdc/algo/accessors.py +++ b/hdc/algo/accessors.py @@ -1,5 +1,5 @@ """Xarray Accesor classes.""" -from typing import List, Optional, Union +from typing import Iterable, List, Optional, Union from warnings import warn from dask import is_dask_collection @@ -100,10 +100,26 @@ def midx(self): def yidx(self): return self._tseries.apply(lambda x: self._period_cls(x).yidx).to_xarray() + @property + def ndays(self): + return self._tseries.apply(lambda x: self._period_cls(x).ndays).to_xarray() + @property def label(self): return self._tseries.apply(lambda x: str(self._period_cls(x))).to_xarray() + @property + def start_date(self): + return self._tseries.apply(lambda x: self._period_cls(x).start_date).to_xarray() + + @property + def end_date(self): + return self._tseries.apply(lambda x: self._period_cls(x).end_date).to_xarray() + + @property + def raw(self): + return self._tseries.apply(lambda x: self._period_cls(x).raw).to_xarray() + @xarray.register_dataset_accessor("dekad") @xarray.register_dataarray_accessor("dekad") @@ -444,12 +460,35 @@ def spi( self, calibration_start: Optional[str] = None, calibration_stop: Optional[str] = None, + nodata: Optional[Union[float, int]] = None, + groups: Optional[Iterable[int]] = None, + dtype="int16", ): - """Calculate the SPI along the time dimension.""" + """Calculate the SPI along the time dimension. + + Calculates the Standardized Precipitation Index along the time dimension. + Optionally, a calibration start and / or stop date can be provided which + determine the part of the timeseries used to fit the gamma distribution. + + `groups` can be supplied as list of group labels (attention, they are required + to be in format {0..n-1} where n is the number of unique groups. + If `groups` is supplied, the SPI will be computed for each individual group. + This is intended to be used when SPI should be calculated for specific timesteps. + """ if not self._check_for_timedim(): raise MissingTimeError("SPI requires a time dimension!") - from .ops.stats import gammastd_yxt # pylint: disable=import-outside-toplevel + if nodata is None: + if (nodata := self._obj.attrs.get("nodata")) is None: + raise ValueError( + "Need nodata attribute defined, or nodata argument provided." + ) + + # pylint: disable=import-outside-toplevel + from .ops.stats import ( + gammastd_yxt, + gammastd_grp, + ) tix = self._obj.get_index("time") @@ -479,17 +518,44 @@ def spi( "Timeseries too short for calculating SPI. Please adjust calibration period!" ) - res = xarray.apply_ufunc( - gammastd_yxt, - self._obj, - kwargs={ - "cal_start": calstart_ix, - "cal_stop": calstop_ix, - }, - input_core_dims=[["time"]], - output_core_dims=[["time"]], - dask="parallelized", - ) + if groups is None: + res = xarray.apply_ufunc( + gammastd_yxt, + self._obj, + kwargs={ + "nodata": nodata, + "cal_start": calstart_ix, + "cal_stop": calstop_ix, + }, + input_core_dims=[["time"]], + output_core_dims=[["time"]], + keep_attrs=True, + dask="parallelized", + dask_gufunc_kwargs={"meta": self._obj.data.astype(dtype)}, + ) + + else: + groups = np.array(groups) if not isinstance(groups, np.ndarray) else groups + num_groups = np.unique(groups).size + + if not groups.dtype.name == "int16": + warn("Casting groups to int16!") + groups = groups.astype("int16") + + res = xarray.apply_ufunc( + gammastd_grp, + self._obj, + groups, + num_groups, + nodata, + calstart_ix, + calstop_ix, + input_core_dims=[["time"], ["grps"], [], [], [], []], + output_core_dims=[["time"]], + keep_attrs=True, + dask="parallelized", + dask_gufunc_kwargs={"meta": self._obj.data.astype(dtype)}, + ) res.attrs.update( { @@ -523,6 +589,7 @@ def lroo(self): input_core_dims=[["time"]], dask="parallelized", output_dtypes=["uint8"], + keep_attrs=True, ) def autocorr(self): @@ -533,8 +600,7 @@ def autocorr(self): xarray.DataArray with lag1 autocorrelation """ xx = self._obj - nodata = xx.attrs.get("nodata", None) - if nodata is None: + if (nodata := xx.attrs.get("nodata", None)) is None: warn("Calculating autocorr without nodata value defined!") if xx.dims[0] == "time": # I don't know how to tell xarray's map_blocks about @@ -606,6 +672,90 @@ def mktrend(self): x.trend.attrs["nodata"] = -2 return x + def mean_grp( + self, + groups: Iterable[int], + nodata: Optional[Union[int, float]] = None, + ): + """Calculate mean over groups along time dimension. + + This calculates a simple average over groups along the time + dimension. The groups are identified by an int16 array, and + the **need to be in ascending order from 0 to n-1**! + + The function will return an array of original size with averages + at the respective positions. + """ + if not self._check_for_timedim(): + raise MissingTimeError("Grouped mean requires a time dimension!") + + # pylint: disable=import-outside-toplevel + from .ops.stats import mean_grp + + if nodata is None: + nodata = self._obj.attrs.get("nodata", None) + + if nodata is None: + raise ValueError("Need to define nodata value!") + + groups = ( + np.array(groups, dtype="int16") + if not isinstance(groups, np.ndarray) + else groups + ) + + if groups.size != self._obj.time.size: + raise ValueError("Need array of groups same length as time dimension!") + + num_groups = np.unique(groups).size + + return xarray.apply_ufunc( + mean_grp, + self._obj, + groups, + num_groups, + nodata, + input_core_dims=[["time"], ["grps"], [], []], + output_core_dims=[["time"]], + keep_attrs=True, + dask="parallelized", + dask_gufunc_kwargs={"meta": self._obj.data}, + ) + + +class RollingWindowAlgos(AccessorBase): + """Class to calculate rolling window algos on dimenson.""" + + def sum( + self, + window_size: int, + dtype: str = "float32", + dimension: str = "time", + nodata: Optional[Union[int, float]] = None, + ): + # pylint: disable=import-outside-toplevel + from .ops.stats import rolling_sum + + if nodata is None: + if (nodata := self._obj.attrs.get("nodata")) is None: + raise ValueError( + "Need nodata attribute defined, or nodata argument provided." + ) + + xx = xarray.apply_ufunc( + rolling_sum, + self._obj, + window_size, + nodata, + input_core_dims=[[dimension], [], []], + output_core_dims=[[dimension]], + keep_attrs=True, + dask="parallelized", + dask_gufunc_kwargs={"meta": self._obj.astype(dtype).data}, + ) + xx = xx[..., window_size - 1 :] + return xx + class ZonalStatistics(AccessorBase): """Class to claculate zonal statistics.""" @@ -614,7 +764,7 @@ def mean( self, zones: xarray.DataArray, zone_ids: Union[List, np.ndarray], - dtype: str = "float64", + dtype: str = "float32", dim_name: str = "zones", name: Optional[str] = None, ) -> xarray.DataArray: @@ -660,6 +810,9 @@ def mean( "stat": ["mean", "valid"], } + # convert str datatype to type + dtype = np.dtype(dtype).type + if is_dask_collection(xx): dask_name = name if isinstance(dask_name, str): @@ -677,7 +830,7 @@ def mean( drop_axis=[1, 2], new_axis=[1, 2], chunks=chunks, - dtype=dtype, + out_dtype=dtype, name=dask_name, ) else: @@ -687,6 +840,7 @@ def mean( num_zones, xx.nodata, zones.nodata, + out_dtype=dtype, ) return xarray.DataArray( @@ -703,5 +857,6 @@ def __init__(self, xarray_obj): self.algo = PixelAlgorithms(xarray_obj) self.anom = Anomalies(xarray_obj) self.iteragg = IterativeAggregation(xarray_obj) + self.rolling = RollingWindowAlgos(xarray_obj) self.whit = WhittakerSmoother(xarray_obj) self.zonal = ZonalStatistics(xarray_obj) diff --git a/hdc/algo/ops/stats.py b/hdc/algo/ops/stats.py index 48d7880..5d81f1a 100644 --- a/hdc/algo/ops/stats.py +++ b/hdc/algo/ops/stats.py @@ -1,7 +1,8 @@ """Numba accelerated statistical funtions.""" from math import erf, log, sqrt -import numba + +from numba import guvectorize, njit import numpy as np import scipy.special as sc @@ -11,7 +12,7 @@ _init_extension() -@numba.njit +@njit def brentq(xa, xb, s): """ Root finding optimization using Brent's method. @@ -108,7 +109,7 @@ def brentq(xa, xb, s): return xcur -@numba.njit +@njit def gammafit(x): # pylint: disable=line-too-long """ @@ -150,8 +151,8 @@ def gammafit(x): return (a, b) -@numba.njit -def gammastd(x, cal_start, cal_stop, a=0, b=0): +@njit +def gammastd(x, nodata, cal_start, cal_stop, a=0, b=0): """Calculate a standardized index for observations based on fitted gamma distribution. The standardized index are anomalies relative to a variable's statistics @@ -175,7 +176,7 @@ def gammastd(x, cal_start, cal_stop, a=0, b=0): p_zero = 1 - (n_valid / t) if p_zero > 0.9: - return np.full_like(x, -9999, dtype="float64") + return np.full_like(x, nodata, dtype="float64") if (a == 0) and (b == 0): alpha, beta = gammafit(x[cal_start:cal_stop]) @@ -183,7 +184,7 @@ def gammastd(x, cal_start, cal_stop, a=0, b=0): alpha, beta = (a, b) if alpha == 0 or beta == 0: - return np.full_like(x, -9999, dtype="float64") + return np.full_like(x, nodata, dtype="float64") y = np.full(t, p_zero, dtype="float64") # type: ignore @@ -198,9 +199,10 @@ def gammastd(x, cal_start, cal_stop, a=0, b=0): return y -@numba.njit +@njit def gammastd_yxt( x, + nodata, cal_start=None, cal_stop=None, ): @@ -210,7 +212,6 @@ def gammastd_yxt( of the array with the assumption the time dimension is the last one. The outputs are scaled by 1000 and cast to in16. """ - y = np.full_like(x, -9999, dtype="int16") r, c, t = x.shape if cal_start is None: @@ -219,11 +220,16 @@ def gammastd_yxt( if cal_stop is None: cal_stop = t + y = np.full_like(x, nodata, dtype="int16") + for ri in range(r): for ci in range(c): xt = x[ri, ci, :] - s = gammastd(xt, cal_start, cal_stop) - if (s != -9999).sum() > 0: + if (xt != nodata).sum() == 0: + y[ri, ci, :] = nodata + continue + s = gammastd(xt, nodata, cal_start, cal_stop) + if (s != nodata).sum() > 0: s = s * 1000 np.round_(s, 0, s) y[ri, ci, :] = s[:] @@ -231,7 +237,37 @@ def gammastd_yxt( return y -@numba.njit +@lazycompile( + guvectorize( + [ + "(int16[:], int16[:], float64, float64, float64, float64, int16[:])", + "(float32[:], int16[:], float64, float64, float64, float64, int16[:])", + ], + "(n),(m),(),(),(),() -> (n)", + ) +) +def gammastd_grp(xx, groups, num_groups, nodata, cal_start, cal_stop, yy): + """Calculate the gammastd for specific groups. + + This calculates gammastd across xx for indivual groups + defined in `groups`. These need to be in ascending order from + 0 to num_groups - 1. + """ + for grp in range(num_groups): + grp_ix = groups == grp + pix = xx[grp_ix] + if (pix != nodata).sum() == 0: + yy[grp_ix] = nodata + continue + res = gammastd(pix, nodata, cal_start, cal_stop) + if (res != nodata).sum() > 0: + valid_ix = res != nodata + res[valid_ix] = res[valid_ix] * 1000 + np.round_(res, 0, res) + yy[grp_ix] = res[:] + + +@njit def mk_score(x): """Calculate MK score (S) and Kendall's Tau. @@ -257,7 +293,7 @@ def mk_score(x): return s, tau -@numba.njit +@njit def mk_variance_s(x): """Mann-Kendall's variance S. @@ -285,7 +321,7 @@ def mk_variance_s(x): return ((n * (n - 1) * (2 * n + 5)) - tp) / 18 -@numba.njit +@njit def mk_z_score(s, vs): """Calculate Mann-Kendall's Z (MKZ). @@ -304,7 +340,7 @@ def mk_z_score(s, vs): return 0 -@numba.njit +@njit def mk_p_value(z, alpha=0.05): """Calculate Mann-Kendall's p value and significance. @@ -320,7 +356,7 @@ def mk_p_value(z, alpha=0.05): return p, h -@numba.njit +@njit def mk_sens_slope(x): """Calculate Sen's slope and intercept for Mann-Kendall test. @@ -343,7 +379,7 @@ def mk_sens_slope(x): return slope, intercept -@numba.njit +@njit def mann_kendall_trend_yxt(x): """Calculate Mann-Kendall trend over y, x, t array. @@ -382,7 +418,7 @@ def mann_kendall_trend_yxt(x): return r -@numba.njit +@njit def mann_kendall_trend_1d(x): """Caluclate Mann-Kendall trend metric on 1-d array. @@ -411,7 +447,7 @@ def mann_kendall_trend_1d(x): @lazycompile( - numba.guvectorize( + guvectorize( [ "(int16[:], float64, float32[:], float32[:], float32[:], int8[:])", "(float32[:], float64, float32[:], float32[:], float32[:], int8[:])", @@ -433,7 +469,7 @@ def _mann_kendall_trend_gu_nd(x, nodata, tau, p, slope, trend): @lazycompile( - numba.guvectorize( + guvectorize( [ "(int16[:], float32[:], float32[:], float32[:], int8[:])", "(float32[:], float32[:], float32[:], float32[:], int8[:])", @@ -445,3 +481,62 @@ def _mann_kendall_trend_gu_nd(x, nodata, tau, p, slope, trend): def _mann_kendall_trend_gu(x, tau, p, slope, trend): """Guvectorize wrapper for mann_kendall_trend_1d.""" tau[0], p[0], slope[0], trend[0] = mann_kendall_trend_1d(x) + + +@lazycompile( + guvectorize( + [ + "(float32[:], int16[:], float64, float64, float32[:])", + "(int16[:], int16[:], float64, float64, float32[:])", + "(int32[:], int16[:], float64, float64, float32[:])", + "(int64[:], int16[:], float64, float64, float32[:])", + ], + "(n),(m),(),() -> (n)", + ) +) +def mean_grp(xx, groups, num_groups, nodata, yy): + """Calculate grouped mean.""" + for grp in range(num_groups): + grp_ix = groups == grp + pix = xx[grp_ix] + n = 0 + for pixv in pix: + if pixv == nodata: + continue + if n == 0: + avg = pixv + else: + avg += pixv + n += 1 + if n == 0: + avg = nodata + else: + avg = avg / n + + yy[grp_ix] = avg + + +@lazycompile( + guvectorize( + [ + "(float32[:], float64, float64, float32[:])", + "(int16[:], float64, float64, float32[:])", + "(int64[:], float64, float64, float32[:])", + ], + "(n),(),() -> (n)", + ) +) +def rolling_sum(xx, window_size, nodata, yy): + """Calculate moving window sum over specified size.""" + n = xx.size + yy[:] = 0 + for ii in range(n): + if ii - window_size + 1 < 0: + yy[ii] = nodata + continue + + for jj in range(ii - window_size + 1, ii + 1): + if xx[jj] == nodata: + yy[ii] = nodata + continue + yy[ii] += xx[jj] diff --git a/hdc/algo/ops/zonal.py b/hdc/algo/ops/zonal.py index 7a5ba3b..dce7789 100644 --- a/hdc/algo/ops/zonal.py +++ b/hdc/algo/ops/zonal.py @@ -6,7 +6,7 @@ @lazycompile(njit) -def do_mean(pixels, z_pixels, num_zones, nodata, z_nodata, dtype="float64"): +def do_mean(pixels, z_pixels, num_zones, nodata, z_nodata, out_dtype=np.float32): """Calculate the zonal mean. The mean for each pixel of `pixels` is calculated for each zone in @@ -22,11 +22,11 @@ def do_mean(pixels, z_pixels, num_zones, nodata, z_nodata, dtype="float64"): z_pixels: zonal pixels (Y,X) nodata: nodata value in values num_zones: number of zones in z_pixels - dtype: datatype + out_dtype: datatype """ t, nr, nc = pixels.shape + result = np.zeros((t, num_zones, 2), dtype=out_dtype) - result = np.zeros((t, num_zones, 2), dtype=dtype) # 0 mean # 1 valids diff --git a/tests/test_accessors.py b/tests/test_accessors.py index bbc5e27..24dc9b5 100644 --- a/tests/test_accessors.py +++ b/tests/test_accessors.py @@ -73,6 +73,54 @@ def test_period_yidx_dekad(darr): np.testing.assert_array_equal(darr.time.dekad.yidx, [1, 2, 3, 3, 4]) +def test_period_ndays_dekad(darr): + assert isinstance(darr.time.dekad.ndays, xr.DataArray) + np.testing.assert_equal( + darr.time.dekad.ndays.values, np.array([10, 10, 11, 11, 10]) + ) + + +def test_period_start_date_dekad(darr): + assert isinstance(darr.time.dekad.start_date, xr.DataArray) + np.testing.assert_equal( + darr.time.dekad.start_date.values, + np.array( + [ + "2000-01-01T00:00:00.000000000", + "2000-01-11T00:00:00.000000000", + "2000-01-21T00:00:00.000000000", + "2000-01-21T00:00:00.000000000", + "2000-02-01T00:00:00.000000000", + ], + dtype="datetime64[ns]", + ), + ) + + +def test_period_end_date_dekad(darr): + assert isinstance(darr.time.dekad.end_date, xr.DataArray) + np.testing.assert_equal( + darr.time.dekad.end_date.values, + np.array( + [ + "2000-01-10T23:59:59.999999000", + "2000-01-20T23:59:59.999999000", + "2000-01-31T23:59:59.999999000", + "2000-01-31T23:59:59.999999000", + "2000-02-10T23:59:59.999999000", + ], + dtype="datetime64[ns]", + ), + ) + + +def test_period_raw_dekad(darr): + assert isinstance(darr.time.dekad.raw, xr.DataArray) + np.testing.assert_array_equal( + darr.time.dekad.raw, [72000, 72001, 72002, 72002, 72003] + ) + + def test_period_labels_dekad(darr): assert isinstance(darr.time.dekad.label, xr.DataArray) np.testing.assert_array_equal( @@ -145,6 +193,12 @@ def test_algo_spi(darr, res_spi): np.testing.assert_array_equal(_res, res_spi) +def test_algo_spi_grouped(darr, res_spi): + _res = darr.astype("float32").hdc.algo.spi(groups=np.array([0] * 5, dtype="int16")) + assert isinstance(_res, xr.DataArray) + np.testing.assert_array_equal(_res, res_spi) + + def test_algo_spi_transp(darr, res_spi): _darr = darr.transpose(..., "time") _res = _darr.hdc.algo.spi() @@ -216,6 +270,13 @@ def test_algo_spi_decoupled_3(darr): assert _res.attrs["spi_calibration_stop"] == "2000-02-10" +def test_algo_spi_nodata(darr): + _ = darr.attrs.pop("nodata") + with pytest.raises(ValueError): + _ = darr.hdc.algo.spi() + _ = darr.hdc.algo.spi(nodata=-9999) + + def test_algo_spi_decoupled_err_1(darr): with pytest.raises(ValueError): _res = darr.hdc.algo.spi( @@ -888,7 +949,7 @@ def test_zonal_mean(darr, zones): [[63.0, 2.0], [16.0, 2.0]], ] ), - dtype="float64", + dtype="float32", ) z_ids = np.unique(zones.data) @@ -897,6 +958,7 @@ def test_zonal_mean(darr, zones): np.testing.assert_equal(x.coords["zones"].data, [0, 1]) assert list(x.coords["stat"].values) == ["mean", "valid"] np.testing.assert_almost_equal(x, res) + assert x.dtype == res.dtype def test_zonal_mean_nodata(darr, zones): @@ -908,7 +970,7 @@ def test_zonal_mean_nodata(darr, zones): [[28.0, 2.0], [12.0, 2.0]], [[63.0, 2.0], [16.0, 2.0]], ], - dtype="float64", + dtype="float32", ) darr[0, 0, 0] = darr.nodata @@ -916,6 +978,7 @@ def test_zonal_mean_nodata(darr, zones): z_ids = np.unique(zones.data) x = darr.hdc.zonal.mean(zones, z_ids) np.testing.assert_almost_equal(x, res) + assert x.dtype == res.dtype def test_zonal_mean_nodata_nan(darr, zones): @@ -943,8 +1006,9 @@ def test_zonal_mean_nodata_nan_float(darr, zones): darr[0, 0, 0] = "nan" z_ids = np.unique(zones.data) - x = darr.hdc.zonal.mean(zones, z_ids) + x = darr.hdc.zonal.mean(zones, z_ids, dtype="float64") np.testing.assert_almost_equal(x, res) + assert x.dtype == res.dtype def test_zonal_zone_nodata_nan(darr, zones): @@ -961,13 +1025,14 @@ def test_zonal_zone_nodata_nan(darr, zones): zones.attrs = {"nodata": 0} z_ids = np.unique(zones.data) - x = darr.hdc.zonal.mean(zones, z_ids, dim_name="foo") + x = darr.hdc.zonal.mean(zones, z_ids, dim_name="foo", dtype="float64") np.testing.assert_almost_equal(x, res) + assert x.dtype == res.dtype def test_zonal_dimname(darr, zones): z_ids = np.unique(zones.data) - x = darr.hdc.zonal.mean(zones, z_ids, dim_name="foo") + x = darr.hdc.zonal.mean(zones, z_ids, dim_name="foo", dtype="float64") assert x.dims == ("time", "foo", "stat") @@ -989,3 +1054,41 @@ def test_zonal_type_exc(darr, zones): z_ids = np.unique(zones.data) with pytest.raises(ValueError): _ = darr.hdc.zonal.mean(zones.data, z_ids) + + +def test_rolling_sum(darr): + res = darr.hdc.rolling.sum(window_size=3).transpose("time", ...) + xr.testing.assert_equal( + res, + darr.rolling(time=3).sum().dropna("time"), + ) + np.testing.assert_equal(darr[-3:].sum("time").data, res[-1].data) + + +def test_rolling_sum_nodata(darr): + darr[:, 0, 0] = -9999 + xr.testing.assert_equal( + darr.hdc.rolling.sum(3, nodata=-9999).transpose("time", ...), + darr.rolling(time=3).sum().where(darr != -9999, -9999).dropna("time"), + ) + + +def test_mean_grp(darr): + grps = np.array([0, 0, 1, 1, 2], dtype="int16") + tst = darr.hdc.algo.mean_grp(grps)[..., [0, 2, 4]] + cntrl = ( + darr.groupby(xr.DataArray(grps, dims=("time",))).mean().transpose(..., "group") + ) + + np.testing.assert_allclose(tst.data, cntrl.data) + + +def test_mean_grp_exc(darr): + grps = np.array([0, 0, 1, 1, 2], dtype="int16") + _darr = darr.copy() + _ = _darr.attrs.pop("nodata") + with pytest.raises(ValueError): + _darr.hdc.algo.mean_grp(grps) + + with pytest.raises(MissingTimeError): + darr.rename(time="foo").hdc.algo.mean_grp(grps) diff --git a/tests/test_algos.py b/tests/test_algos.py index d237740..2a631d1 100644 --- a/tests/test_algos.py +++ b/tests/test_algos.py @@ -22,12 +22,15 @@ gammafit, gammastd, gammastd_yxt, + gammastd_grp, + mean_grp, mk_score, mk_p_value, mk_z_score, mk_variance_s, mk_sens_slope, mann_kendall_trend_1d, + rolling_sum, ) @@ -92,8 +95,8 @@ def test_gammafit(ts): assert parameters == pytest.approx((1.083449238500003, 0.9478709674697126)) -def test_spi(ts): - xspi = gammastd_yxt.py_func(ts.reshape(1, 1, -1)) +def test_gammastd(ts): + xspi = gammastd_yxt.py_func(ts.reshape(1, 1, -1), -9999) assert xspi.shape == (1, 1, 10) np.testing.assert_array_equal( xspi[0, 0, :], @@ -101,8 +104,17 @@ def test_spi(ts): ) -def test_spi_nofit(ts): - xspi = gammastd.py_func(ts, 0, len(ts), a=1, b=2) +def test_gammastd(ts): + xspi = gammastd.py_func(ts, -9999, 0, 10) + + np.testing.assert_array_equal( + (xspi * 1000).round(), + [-382.0, 1654.0, 588.0, 207.0, -1097.0, -1098.0, -1677.0, 1094.0, 213.0, 514.0], + ) + + +def test_gammastd_nofit(ts): + xspi = gammastd.py_func(ts, -9999, 0, len(ts), a=1, b=2) xspi np.testing.assert_array_equal( (xspi * 1000).round(), @@ -121,8 +133,8 @@ def test_spi_nofit(ts): ) -def test_spi_selfit(ts): - xspi = gammastd_yxt.py_func(ts.reshape(1, 1, -1), cal_start=0, cal_stop=3) +def test_gammastd_selfit(ts): + xspi = gammastd_yxt.py_func(ts.reshape(1, 1, -1), -9999, cal_start=0, cal_stop=3) assert xspi.shape == (1, 1, 10) np.testing.assert_array_equal( xspi[0, 0, :], @@ -141,15 +153,56 @@ def test_spi_selfit(ts): ) -def test_spi_selfit_2(ts): +def test_gammastd_selfit_2(ts): cal_start = 2 cal_stop = 8 + nodata = -9999 a, b = gammafit.py_func(ts[cal_start:cal_stop]) - xspi_ref = gammastd.py_func(ts, cal_start=cal_start, cal_stop=cal_stop, a=a, b=b) - xspi = gammastd.py_func(ts, cal_start=cal_start, cal_stop=cal_stop) + xspi_ref = gammastd.py_func( + ts, nodata, cal_start=cal_start, cal_stop=cal_stop, a=a, b=b + ) + xspi = gammastd.py_func(ts, nodata, cal_start=cal_start, cal_stop=cal_stop) np.testing.assert_equal(xspi, xspi_ref) +def test_gammastd_grp(ts): + tts = np.repeat(ts, 5).astype("float32") + grps = np.tile(np.arange(5), 10).astype("int16") + xspi = gammastd_grp(tts, grps, np.unique(grps).size, 0, 0, 10) + + res = [ + -0.38238713, + 1.6544422, + 0.5879236, + 0.20665395, + -1.0974495, + -1.0975538, + -1.6773673, + 1.093621, + 0.21322519, + 0.5144766, + ] + + np.testing.assert_almost_equal(xspi, (np.repeat(res, 5) * 1000).round()) + + +def test_mean_grp(ts): + tts = np.repeat(ts, 5).astype("float32") + grps = np.tile(np.arange(5), 10).astype("int16") + res = mean_grp(tts, grps, np.unique(grps).size, 0) + assert (res == 1.02697).all() + + +def test_rolling_sum(): + res = rolling_sum(np.arange(10).astype("float32"), 3, 0) + np.testing.assert_almost_equal( + res, + np.array( + [0.0, 0.0, 3.0, 6.0, 9.0, 12.0, 15.0, 18.0, 21.0, 24.0], dtype="float32" + ), + ) + + def test_ws2dgu(ts): _ts = ts * 10 z = ws2dgu(_ts, 10, 0)