def test_get_dekad_period(self):
dates = [self.date1, self.date2, self.date3]
periods = dk.get_dekad_period(dates)
assert periods == [4, 5, 6]
def calc_DI(data, inverse=False, interest_period=[6, 12, 24], scaled=False,
scale_zero=False, modf_all=False):
"""
Calculates a Drought Index based on an algorithm developed by
FAO SWALIM.
Parameters
----------
data : pandas.DataFrame
Input data as Pandas DataFrame, must come with column names.
inverse : bool
Inverts the input time series; set True if time series is indirect
proportional to the expected output, e.g. Temperature with output
Temperature Drought Index.
interest_period : list of int, optional
interest periods used to calculate drought index,
defaults to [6, 12, 24]
scaled : boolean, optional
If True values will be scaled between 0 and 1.
scale_zero : boolean, optional
If True values will be shifted around zero, defaults to False.
modf_all : boolean, optional
If True values will be modified, independent of their min.
"""
ts_date = data.index
variables = data.keys()
data['period'] = get_dekad_period(ts_date)
for var in variables:
if inverse is True:
data[var] = ((data[var].max() + 1) - data[var])
if modf_all is True:
data['modf'] = data[var] + 1
del data[var]
elif data[var].min() == 0:
data['modf'] = data[var] + 1
del data[var]
else:
data['modf'] = data[var]
del data[var]
data['modf_avg'] = (data.groupby('period').modf
.transform(lambda x: x.mean()))
# Excess
# Dekads below long term average. If the statement is true the
# program return 1
ex = np.choose((data['modf_avg'] / data['modf']) >= 1, [0., 1.]).values
nanloc = np.where(np.isnan(data['modf']))
ex[nanloc] = np.NAN
data['exc'] = ex
# Run length
# Maximum number of successive dekads below long term average
for ip in interest_period:
data['rlen'] = pd.rolling_apply(data['exc'], ip,
(lambda x:
len(max((''.join(str(j)
for j in map(int,
x)))
.split('0')))),
ip)
# get modified run length
max_rlen = data['rlen'].max()
data['rlen'] = (max_rlen + 1) - data['rlen']
# average run lenghts
rlen_avg = (data.groupby('period').modf
.transform(lambda x: x.mean()))
data['form'] = data['rlen'] / rlen_avg
# sumip matrix
# calculates sum of the values for each interest period
data['sumip'] = pd.rolling_apply(data['modf'], ip,
lambda x: sum(x), ip)
# average values for each interest period over all years
sumip_avg = (data.groupby('period')['sumip']
.transform(lambda x: x.mean()))
data['nrl'] = data['sumip'] / sumip_avg
# calculating PDI/TDI
data['val'] = data['nrl'] * np.sqrt(data['form'])
# scaled index
dkey = var + '_DI_' + str(ip)
if scaled:
data[dkey] = ((data['val'] - data['val'].min()) /
(data['val'].max() - data['val'].min()))
else:
data[dkey] = data['val']
if scale_zero:
data[dkey] = data[dkey] - data[dkey].mean()
del (data['val'], data['nrl'], data['sumip'], data['rlen'],
data['form'])
# deletes not further relevant columns
del data['modf'], data['modf_avg'], data['exc']
del data['period']
return data
