<jupyter_start><jupyter_text>To do* Add f_IExcess back in for this time period and see if it still runs  * If it doesn't, decrease the maximum value to 0.1 and try again  * If it does, run for a longer period (as before when didn't work), but with the maximum value at 0.1<jupyter_code>%matplotlib inline
import matplotlib.pyplot as plt, seaborn as sn, emcee, corner, mpld3
from scipy import optimize
import numpy as np, pandas as pd
from scipy.integrate import odeint
from scipy.stats import norm

sn.set_context('notebook')
# FUNCTION USED BY SCIPY.INTEGRATE.ODEINT TO SOLVE THE MODEL EQUATIONS AT EACH TIME STEP

def f(y, t, ode_params):
    """ Define ODE system.
            y is list [Vs, Qs, Qg, Ds, Dg]
            t is an array of time points of interest
            params is a tuple of input values & model params
            (P, E, f_IExcess, alpha, beta, T_s, T_g, fc)
    """
    # Unpack incremental values for Qs and Qg
    Vs_i = y[0]
    Qs_i = y[1]
    Qg_i = y[2]
    
    # Unpack params
    P, E, f_IExcess, alpha, beta, T_s, T_g, fc = ode_params

    # Model equations (see section 2.2)
    ##dQq_dt = (f_IExcess*P - Qq_i)/T_q
    dQs_dV = (((Vs_i - fc)*np.exp(fc - Vs_i))/(T_s*((np.exp(fc-Vs_i) + 1)**2))) + (1/(T_s*(np.exp(fc-Vs_i) + 1)))
    dVs_dt = P*(1-f_IExcess) - alpha*E*(1 - np.exp(-0.02*Vs_i)) - Qs_i
    dQs_dt = dQs_dV*dVs_dt
    dQg_dt = (beta*Qs_i - Qg_i)/T_g
    dDs_dt = (1 - beta)*Qs_i
    dDg_dt = Qg_i
    
    # Add results of equations to an array
    res = np.array([dVs_dt, dQs_dt, dQg_dt, dDs_dt, dDg_dt])
    
    return res
# FUNCTION TO WRAP THE MODEL EQUATIONS IN A LOOP TO EVALUATE DRAINAGE VOLUME AT EACH STEP IN THE INPUT TIME SERIES

def hydro_model(met_df, ics, p, period, step_len=1):
    """ The hydrological model

            met_df         Dataframe containing columns 'Rainfall_mm' and 'PET_mm', with datetime index
            ics            Vector of initial conditions [Vs0, Vg0]
            p              Series of parameter values (index = param name)
            period         Vector of [start, end] dates [yyyy-mm-dd, yyyy-mm-dd]
            step_len       Length of each step in the input dataset (days)

        Returns a dataframe with column headings
        [Vs, Qs, Qg, Ds, Dg, Sim_Runoff, Obs_Runoff]
    """
    # Truncate the met data to the desired period
    input_df = met_df.truncate(before=period[0], after=period[1])

    # Unpack initial conditions
    Vs0, Vg0 = ics

    # Unpack model parameters
    #     f_IExcess, alpha, beta, T_s, T_g, fc = mod_params

    # Time points to evaluate ODEs at. We're only interested in the start and the end of each step
    ti = [0, step_len]

    # Lists to store output
    output_ODEs = []
    output_rest = []

    # Loop over met data
    for idx in range(len(input_df)):

        # Get P and E for this day
        P = input_df.ix[idx, 'Rainfall_mm']
        E = input_df.ix[idx, 'PET_mm']

        # Calculate infiltration excess and add to results
        Qq = p['f_IExcess']*P
        output_rest.append(Qq)

        # Calculate Qs0 and Qg0 from Vs0 and Vg0
        Qs0 = (Vs0 - p['fc'])/(p['T_s']*(1 + np.exp(p['fc'] - Vs0)))
        Qg0 = Vg0/p['T_g']

        # Vector of initial conditions
        y0 = [Vs0, Qs0, Qg0, 0., 0.]

        # Model parameters plus rainfall and ET, for input to solver
        ode_params = np.array([P, E, p['f_IExcess'], p['alpha'], p['beta'], p['T_s'],
                               p['T_g'], p['fc']])

        # Solve
        y = odeint(f, y0, ti, args=(ode_params,))

        # Extract values for end of step
        res = y[1]

        # Numerical errors may result in very tiny values <0
        # set these back to 0
        res[res<0] = 0
        output_ODEs.append(res)

        # Update initial conditions for next step
        Vs0 = res[0]
        Vg0 = res[2]*p['T_g']

    # Build a dataframe of ODE results
    df1 = pd.DataFrame(data=np.vstack(output_ODEs),
                      columns=['Vs', 'S', 'G', 'Ds', 'Dg'],
                      index=input_df.index)

    # Dataframe of non ODE results
    df2 = pd.DataFrame(data=np.vstack(output_rest), columns=['Qq'],
                     index=input_df.index)

    # Concatenate results dataframes
    df = pd.concat([df1,df2], axis=1)

    # Estimate runoff as Ds + Dg
    df['Sim_Runoff_mm_IE'] = df['Ds'] + df['Dg'] + df['Qq']
    df['Sim_Runoff_mm'] = df['Ds'] + df['Dg']

    # Add observed runoff to df
    df['Obs_Runoff_mm'] = input_df['Runoff_mm']

    return df
# FUNCTION TO DEFINE LOG LIKELIHOOD

def log_likelihood(p_cal_values, p_list, p, met_df, ics, period):
    """ p_cal_values   Row vector of value for parameters being calibrated
        p_list         List of parameters to be calibrated/varied
        p              Series of all parameter values, including those to be calibrated (index = param names)
        met_df         Dataframe including columns 'Rainfall_mm' and 'PET_mm', with datetime index
        ics            List of initial conditions for Vs and Vg, e.g. [200., 100.]
        
        Returns the log likelihood.
    """   

    # Update parameter values being calibrated in the parameter Series, p
    for idx, param in enumerate(p_list):
        p[param] = p_cal_values[idx]

    # Run deterministic model with these parameters
    df = hydro_model(met_df=met_df, ics=ics, period=period, p=p)

    # Extract arrays for simulated and observed runoff
    sim = df['Sim_Runoff_mm']
    obs = df['Obs_Runoff_mm']

    # Calculate sigma_e for each step
    sigma_e = p['m']*sim

    # Calculate log likelihood. For each element in the arrays sim, sigma_e and obs,
    # this code calculates the log prob of drawing 'obs' from a Gaussian centred on 
    # 'sim' with std. dev. 'sigma_e'
    likes = norm(sim, sigma_e).logpdf(obs)

    # If flow is zero, sigma_e is zero and scipy.norm returns NaN
    # Set these to -inf instead
    likes[np.isnan(likes)] = -np.inf

    # Sum log likes
    ll = np.sum(likes)
    
    return ll
# FUNCTION TO DEFINE LOG PRIOR

def log_prior(p_cal_values, priors_df):
    """ p_cal_values is an array of parameter values for calibrating params
        priors_df is a dataframe of min and max values for priors for each calibrated param
    
        Returns the log prior probability of p
    """
    # Add the current parameter values to the priors dataframe
    priors_df['value'] = p_cal_values

    # Determine whether the current parameter values lie within the priors
    # (add a boolean column to the dataframe, value 0 if param not in range)
    priors_df['valid'] = np.where((priors_df['value']>= priors_df['min']) &
                                  (priors_df['value']<priors_df['max']), 1, 0)

    # If all parameters are within allowed ranges, return a constant
    # Otherwise, the parameter set is invalid. In which case, it has
    # probability 0, i.e. log prob = -inf
    if sum(priors_df['valid']) == len(priors_df['valid']):
        return 0
    else:
        return -np.inf
# FUNCTION TO DEFINE LOG POSTERIOR

def log_posterior(p_cal_values, p_list, p, met_df, priors_df, ics, period):
    """ mcmc_params    Vector of parameters (model params for calibration + error variance)
        met_df         Dataframe containing columns 'Rainfall_mm' and 'PET_mm', with datetime index
        
        Returns the log posterior.
        The log posterior distribution is (proportional to) the sum of the log prior and the log likelihood
    """   
    # Get log prior prob
    log_pri = log_prior(p_cal_values, priors_df)

    # Evaluate log likelihood if necessary
    if np.isfinite(log_pri):
        log_like = log_likelihood(p_cal_values, p_list, p, met_df, ics, period)
        
        # Calculate log posterior
        return log_pri + log_like
    
    else:
        # Log prior is -inf, so log posterior is -inf too
        return -np.inf
# FUNCTION TO DECIDE ON STARTION LOCATIONS FOR EACH OF THE MCMC WALKERS
# To do this: (1) use an optimiser to estimate the maximum of the posterior
# (2) add a small amount of random noise so each walker slights from a slightly different location

def neg_log_posterior(p_cal_values, p_list, p, met_df, priors_df, ics, period):
    """ Negative of log posterior.
    """
    return -log_posterior(p_cal_values, p_list, p, met_df, priors_df, ics, period)

def find_map(init_guess, p_list, p, met_df, priors_df, ics, period):
    """ Estimate the location of the maximum of the posterior density.
    
            init_guess   Initial guess for starting optimiser
            met_df       Data frame of meteorological data
    """
    # Run optimiser
    param_est = optimize.fmin(neg_log_posterior, init_guess, args=(p_list, p, met_df, priors_df, ics, period))

    return param_est<jupyter_output><empty_output><jupyter_text>Call functions<jupyter_code># SET UP

# USER INPUT

# Input data. Download into a dataframe
data_url = r'https://drive.google.com/uc?export=&id=0BximeC_RweaecHNIZF9GMHkwaWc'

# Simulation period
st_dt = '2004-04-01'  # Start date
end_dt = '2004-08-31' # End date

# Catchment area (m2)
cat_area = 51.7E6

# Model parameters, including starting guesses for those being calibrated.
# Include parameters that define the error variance (e.g. sigma or m)
param_dict = {'fc':200., 'beta':0.6, 'f_IExcess':0.02, 'alpha':0.75,
              'T_s':10.,'T_g':100.,'m':0.5}
# beta = 0.6            # BFI (dimensionless)
# fc = 200.             # Field capacity (mm) #SSKIB: 290
# f_IExcess = 0.02

# Initial conditions
Vs0_init = param_dict['fc']       # Initial soil volume (mm)
Vg0_init = 90.       # Initial groundwater volume (mm)

# list of params to calibrate
p_list = ['f_IExcess', 'alpha','T_s','T_g','m']

# Define priors for each of the params being calibrated. Just involves defining upper limits
# (for now, assume uniform; lower limits all 0). Can change lower limits by changing input
# to priors_df (below)
max_dict = {'f_IExcess':1, 'alpha':2,'T_s':20,'T_g':500,'m':1}

# ADMIN
ics=[Vs0_init, Vg0_init]  # Initial conditions
period=[st_dt, end_dt]    # Simulation period

# # Dictionary of parameter values to be calibrated
# p_cal = p.copy()
# for param in p.keys():
#     if param not in p_list:
#         del p_cal[param]

# Store parameter values in series, referenced by their names (the row indices)
p = pd.Series(param_dict)  # All params; cols = param names, one row with values
# p.sort_index(axis=1, inplace=True)

# Array of values for parameters to be calibrated. Note, this has the same ordering as 
# p_list
p_cal = p[p_list] # Same as p, but only with parameters to be calibrated
p_cal_values = p_cal.values

# Create a dataframe of priors for parameters being calibrated (to be input to the 'prior' function)
cols = ['min','max','value','valid']
priors_df = pd.DataFrame(columns=cols, index=p_list)
priors_df['min'] = np.zeros(len(p_list))
priors_df['max'] = pd.Series(max_dict)

# READ IN INPUT DATA
met_df = pd.read_csv(data_url, parse_dates=True, dayfirst=True, index_col=0)

# Convert cumecs to mm
met_df['Runoff_mm'] = met_df['Q_Cumecs']*60*60*24*1000/cat_area
del met_df['Q_Cumecs']

# Linear interpolation of any missing values
met_df.interpolate(method='linear', inplace=True)
# RUN THE OPTIMISER TO FIND THE MAP

LL = log_likelihood(p_cal_values, p_list, p, met_df, ics, period)
L_prior = log_prior(p_cal, priors_df)
L_post = log_posterior(p_cal_values, p_list, p, met_df, priors_df, ics, period)
neg_L_post = neg_log_posterior(p_cal_values, p_list, p, met_df, priors_df, ics, period)

# print LL
# print L_prior
# print L_post
# print neg_L_post

# Run optimiser
param_est = find_map(p_cal_values, p_list, p, met_df, priors_df, ics, period)

# Print results
print '\n'
for idx, param in enumerate(p_cal.index):
    print 'Estimated %s: %.2f' % (param, param_est[idx])<jupyter_output>Optimization terminated successfully.
         Current function value: -45.265279
         Iterations: 342
         Function evaluations: 574


Estimated f_IExcess: 0.35
Estimated alpha: 0.53
Estimated T_s: 2.44
Estimated T_g: 87.99
Estimated m: 0.22
{'fraction_non_alphanumeric': 0.1018748962999834, 'fraction_numerical': 0.015098722415795587, 'mean_word_length': 2.9576493762311227, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 9, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '115858', 'license': 'MIT', 'n_tokens_mistral': 4382, 'n_tokens_neox': 4120, 'n_words': 1348, 'path': 'Development/2016/MAP.ipynb', 'repo_name': 'LeahJB/Simple_P_model', 'size': '12054'}