Drought metric

pcmdi_metrics/drought/README.md

@@ -0,0 +1 @@
+Drought metrics

pcmdi_metrics/drought/drought_metric_driver.py

@@ -0,0 +1,39 @@
+import xarray
+import pandas as pd
+from pcmdi_metrics.drought.lib import spi
+from pcmdi_metrics.drought.lib import cdd
+import matplotlib as mpl
+
+# driver to call SPI and CDD
+
+# PMP PARSER READIN
+
+# LOOP FOR OBS or MODELS
+ # LOOP FOR RUNS
+ # CALL SPI -- calculated for each grid --> output to nc file
+ # Input: monthly precipitation over land (OBS: CPC over CONUS)
+ # CALL CDD -- calculated for each grid --> output to nc file
+ # Input: daily precipitation over land
+
+# OUTPUT TO JSON (which output to JSON??)
+
+
+# The input data should be monthly mean precipitaion. The data used here can be downloaded at https://psl.noaa.gov/data/gridded/data.unified.daily.conus.html
+observe_path = "/Users/lee1043/Documents/Research/DATA/CPC/precip.V1.0.mon.mean.nc"
+observe = xarray.open_dataset(observe_path)
+observe_df = observe.to_dataframe().reset_index()
+observe_regional_mean = observe_df.groupby('time').mean().reset_index()
+date = observe_regional_mean.time
+
+# Build a dataframe to save the SPI3, SPI6, SPI12 and SPI36 of the regional mean monthly precipitation
+observe_regional_spi = pd.DataFrame(columns=['SPI3','SPI6','SPI12','SPI36'], index=date)
+for scale, col in zip([3,6,12,36], observe_regional_spi.columns):
+ observe_regional_spi[col] = spi(observe_regional_mean.precip,observe_regional_mean.precip, scale) 
+observe_regional_spi = observe_regional_spi.reset_index()
+
+# Here I plot the regional mean SPI6 over the total period
+fig, (ax) = plt.subplots(nrows=1, ncols=1, sharex=True,sharey=True,figsize=(20,8))
+ax.fill_between(observe_regional_spi.time, 0, observe_regional_spi.SPI6, where=observe_regional_spi.SPI6 >= 0, facecolor='blue', interpolate=True)
+ax.fill_between(observe_regional_spi.time, 0, observe_regional_spi.SPI6, where=observe_regional_spi.SPI6 < 0, facecolor='red', interpolate=True)
+ax.set_ylabel('Regional Mean SPI6 over CONUS',fontsize=22)
+ax.set_xlabel("Date",fontsize=22)

pcmdi_metrics/drought/lib/__init__.py

@@ -0,0 +1,2 @@
+from .lib_spi import spi #noqa
+from .lib_cdd import cdd #noqa

pcmdi_metrics/drought/lib/lib_cdd.py

@@ -0,0 +1,44 @@
+import numpy as np
+import pandas as pd
+
+#cdd is the function to calculate the CDD, the output is a dataframe grouped by lat lon and time_period.
+
+# data_df should be the input daily precipitation dataframe consists of lat, lon, time_period and daily precipitation.
+# pr_name is the variable name of precipitation
+# period_name is the variable name of period used. like 'year' or 'month'.
+
+def cdd(data_df,pr_name,period_name):
+ data_df['Dryness'] = np.where(data_df[pr_name]>=1,'non-dry','dry')
+ dry_duration = pd.DataFrame(data_df.groupby(['lat','lon',period_name,data_df['Dryness'],(data_df['Dryness']!=data_df['Dryness'].shift()).cumsum()]).size().reset_index(level=4,drop=True).rename('duration'))
+ dry_duration_reset_index = dry_duration.reset_index()
+ dry_duration_reset_index = dry_duration_reset_index[dry_duration_reset_index.Dryness=='dry']
+ CDD_each_grid = pd.DataFrame(dry_duration_reset_index.groupby(['lat','lon',period_name]).duration.max())
+ CDD_each_grid = CDD_each_grid.rename(columns={'duration':"CDD"})
+ return CDD_each_grid.reset_index()
+
+#score_map is a function to plot a variable on map
+
+# data should be the input dataframe.
+# variable_name is varable name of the vairable used like CDD.
+
+def score_map(data,variable_name):
+ val_pivot_df = data.pivot(index='lat', columns='lon', values=variable_name)
+ lons = val_pivot_df.columns.values
+ lats = val_pivot_df.index.values
+ fig, ax = plt.subplots(1, figsize=(8, 8))
+ m = Basemap(projection='merc',
+ llcrnrlat=data.dropna().min().lat-2,
+ urcrnrlat=data.dropna().max().lat+2,
+ llcrnrlon=data.dropna().min().lon-2,
+ urcrnrlon=data.dropna().max().lon+2,
+ resolution='i', area_thresh=10000)
+ m.drawcoastlines()
+ m.drawstates()
+ m.drawcountries()
+ x, y = np.meshgrid(lons,lats) 
+ px, py = m(x,y) 
+ data_values = val_pivot_df.values
+ masked_data = np.ma.masked_invalid(data_values)
+ cmap = plt.cm.viridis
+ m.pcolormesh(px, py, masked_data, cmap=cmap, shading='flat')
+ m.colorbar(label=variable_name)

pcmdi_metrics/drought/lib/lib_spi.py

@@ -0,0 +1,92 @@
+import numpy as np
+import scipy.stats
+# Note that the calibration and calculation data can be the same but they should both be 1-d monthly mean precipitation array and start from January.
+
+
+def accumulation(data, temporal_scale): 
+ # calculate the accumulated n months (temporal_scale) precipitation data
+ data_accumulated = np.convolve(data, np.ones(temporal_scale), 'valid')
+ # set the value to NANs for the first n-1 months because there isn't accumulated n months values at the first n-1 months
+ data_accumulated = np.concatenate([[np.nan]*(temporal_scale-1), data_accumulated])
+
+ # If the last year is incomplete, let's set NAN to the months without values in the last year
+ n_year_incomplete = len(data_accumulated)/12 - int(len(data_accumulated)/12)
+ n_month_incomplete = int(n_year_incomplete * 12)
+ data_accumulated = np.concatenate([data_accumulated,[np.nan]*(n_month_incomplete)])
+ # Then reshape the 1-d array to 2-d array with the shape of (the number of years, 12) becasue the SPI's distribution and calculation should be based on each month.
+ data_accumulated_reshape = data_accumulated.reshape(int(len(data_accumulated)/12), 12)
+ return data_accumulated_reshape
+
+
+def gamma_transformation(calibration,calculation, temporal_scale):
+ # calculate the alphas and betas at each month based on calibration data
+ means = np.nanmean(calibration, axis=0)
+ a = np.log(means) - np.nanmean(np.log(calibration), axis=0)
+ alphas = (1 + np.sqrt(1 + 4 * a / 3)) / (4 * a)
+ betas = means / alphas
+ # calculate the probability of 0 values
+ n_zero = (calculation ==0).sum()
+ n_values = np.count_nonzero(~np.isnan(calculation)) + n_zero
+ p_zero = n_zero / len(calculation)
+
+ # set the 0 value as NAN
+ calculation[calculation == 0] = np.NaN
+
+ # build gamma distribution based on the non-zero values of calibration data
+ # calculate the probability and the values of calculation data under gamma distribution built by calibration data
+ p_gamma_non_zero = scipy.stats.gamma.cdf(calculation, a=alphas, scale=betas)
+ p_gamma = p_zero + ( 1 - p_zero ) * p_gamma_non_zero
+ calculation_transformed = scipy.stats.norm.ppf(p_gamma)
+ # change the shape to 1-d array
+ calculation_transformed = calculation_transformed.flatten()
+ return calculation_transformed
+
+
+def spi(precipitation_calibration:np.ndarray,
+ precipitation_calculation:np.ndarray,
+ temporal_scale:int,
+ distribution='gamma',
+ upper_limit=3.09,
+ lower_limit=-3.09):
+
+ # precipitation_calibration: It should be the 1-d monthly precipitation array used to calibrate the distribution. It must start from January and can be the same as precipitation_calculate. 
+ # precipitation_calculate: It should be the 1-d monthly precipitation array used to calculate the SPI_n. It must start from Januaray too.
+ # temporal_scale: It's the temporal scale of SPI_n which indicates the number of n months used to calculate accumulated SPI. For example, to calculate SPI3, it should be 3.
+ # distribution: The distribution of SPI and right now it only supports gamma distribution which is the most commonly used one.
+ # upper_limit, lower_limit: The maximum and minimum SPI values, according to https://onlinelibrary.wiley.com/doi/abs/10.1111/j.1752-1688.1999.tb03592.x , the default values are 3.09 and -3.09.
+
+ """
+ # test if there is missing and negative value 
+ if np.all(np.isnan(precipitation_calibration)):
+ print("calibration data cannot be all NAN!")
+ if np.all(np.isnan(precipitation_calculation)):
+ print("calculation data cannot be all NAN!")
+ """
+
+ if precipitation_calibration[~np.isnan(precipitation_calibration)].min()<0:
+ print("Calibration data must be positive! All negative values will be assign to be 0")
+ precipitation_calibration.clip(min=0)
+
+ if precipitation_calculation[~np.isnan(precipitation_calculation)].min()<0:
+ print("Calculation data must be positive! All negative values will be assign to be 0")
+ precipitation_calculation.clip(min=0)
+
+ # calculate the length of input calculation data
+ length_input = len(precipitation_calibration)
+
+ # calculate the accumulated precipitation of calibration and calculation data, then reshape them to 2-d array.
+ precipitation_calibration_accumulated = accumulation(precipitation_calibration, temporal_scale)
+ precipitation_calculation_accumulated = accumulation(precipitation_calculation, temporal_scale)
+
+ # apply the gamma transformation to accumulated data
+ if distribution == 'gamma':
+ spi_n = gamma_transformation(precipitation_calibration_accumulated,
+ precipitation_calculation_accumulated,
+ temporal_scale)
+
+ # clip the spi values 
+ spi_n = np.clip(spi_n,a_min=lower_limit , a_max=upper_limit)
+
+ # change spi_n to its original length 
+ spi_n = spi_n[0:length_input]
+ return spi_n