ssnt_mete_comparison.py

from __future__ import division
import matplotlib
matplotlib.use('Agg')
import os
import matplotlib.pyplot as plt
import csv
import numpy as np
from numpy.lib.recfunctions import append_fields
from scipy import stats, integrate
import working_functions as wk
import mete
import mete_distributions
import mete_agsne as agsne
import macroecotools as mtools
import macroeco_distributions as md
import multiprocessing

class ssnt_isd_bounded():
    """The individual-size distribution predicted by SSNT.
    
    Diameter is assumed to be lower-bounded at 1, to be consistent
    with METE.
    SSNT is applied on diameter transformed with an arbitrary
    power alpha, i.e., it is assumed that g(D^alpha) is constant.
    The predicted distribution is transformed back to diameter.
    
    """
    def __init__(self, alpha, par):
        """par is the parameter for the lower-truncated exponential
        
        distribution when the scale is D^alpha. 
        The MLE of par is N / (sum(D^alpha) - N) from data.
        
        """
        self.alpha = alpha
        self.par = par
        self.a = 1 # lower bound
        
    def pdf(self, x):
        if x < self.a: return 0
        else: return self.par * self.alpha * np.exp(-self.par * (x ** self.alpha - 1)) * (x ** (self.alpha - 1))
    
    def cdf(self, x): # cdf of D is equal to cdf of D^alpha
        if x < self.a: return 0
        else: return 1 - np.exp(-self.par * (x ** self.alpha - 1))
    
    def ppf(self, q):
        return (1 - np.log(1 - q) / self.par) ** (1 / self.alpha)
    
    def rvs(self, size):
        rand_list = stats.uniform.rvs(size = size)
        out = np.array([self.ppf(x) for x in rand_list])
        return out
    
    def expected(self):  # Note that this is the expected value of D
        ans = integrate.quad(lambda x: x * self.pdf(x), self.a, np.inf)[0]
        return ans
    
    def expected_square(self):
        ans = integrate.quad(lambda x: x ** 2 * self.pdf(x), self.a, np.inf)[0]
        return ans        

def import_likelihood_data(file_name, file_dir = './out_files/'):
    """Import file with likelihood for METE, SSNT, and transformed SSNT"""
    data = np.genfromtxt(file_dir + file_name, dtype = None, 
                         names = ['study', 'site', 'ASNE', 'AGSNE', 'SSNT_N', 'SSNT_M'], delimiter = ' ')
    return data

def import_bootstrap_file(input_filename, Niter = 100):
    """The same function as the one in working_functions, except that dype = None."""
    names_orig = ['dataset', 'site', 'orig']
    names_sim = ['sample'+str(i) for i in xrange(1, Niter + 1)]
    names_orig.extend(names_sim)
    data = np.genfromtxt(input_filename, delimiter = ',', names = names_orig, dtype = None)
    return data

def import_bootstrap_file_incomp(input_filename, Niter = 100):
    """The same as import_bootstrap_file, but accounts for instances where the number of columns are not constant."""
    with open(input_filename) as f:
        content = f.readlines()    
    names_orig = ['dataset', 'site', 'orig']
    names_sim = ['sample' + str(i) for i in xrange(1, Niter + 1)]
    names = names_orig + names_sim
    data_type = ['S15', 'S15'] + ['f8'] * (Niter + 1)
    out_array = np.zeros(len(content), dtype = {'names': names, 'formats': data_type})
    for i, row in enumerate(content):
        row_split = row.strip(' ').split(',')
        row_split = [row_split[j] if j < 2 else float(row_split[j]) for j in range(len(row_split))]
        out_array[i] = tuple(row_split[:(Niter + 3)])
    return out_array

def clean_data_agsne(raw_data_site, cutoff_genera = 4, cutoff_sp = 9, max_removal = 0.1):
    """Further cleanup of data, removing individuals with undefined genus. 
    
    Inputs:
    raw_data_site - structured array generated by wk.import_raw_data(), with three columns 'site', 'sp', and 'dbh', for a single site
    min_genera - minimal number of genera required for analysis
    min_sp - minimal number of species
    max_removal - the maximal proportion of individuals removed with undefined genus
    
    Output:
    a structured array with four columns 'site', 'sp', 'genus', and 'dbh'
    
    """
    counter = 0
    genus_list = []
    row_to_remove = []
    for i, row in enumerate(raw_data_site):
        sp_split = row['sp'].split(' ')
        genus = sp_split[0]
        if len(sp_split) > 1 and genus[0].isupper() and genus[1].islower() and (not any(char.isdigit() for char in genus)):
            genus_list.append(genus)
        else: 
            row_to_remove.append(i)
            counter += 1
    if counter / len(raw_data_site) <= max_removal:
        raw_data_site = np.delete(raw_data_site, np.array(row_to_remove), axis = 0)
        gen_col = np.array(genus_list)
        out = append_fields(raw_data_site, 'genus', gen_col, usemask = False)
        if len(np.unique(out['sp'])) > cutoff_sp and len(np.unique(out['genus'])) > cutoff_genera: 
            return out
        else: return None
    else: return None
  
def get_GSNE(raw_data_site):
    """Obtain the state variables given data for a single site, returned by clean_data_genera()."""
    G = len(np.unique(raw_data_site['genus']))
    S = len(np.unique(raw_data_site['sp']))
    N = len(raw_data_site)
    E = sum((raw_data_site['dbh'] / min(raw_data_site['dbh'])) ** 2)
    return G, S, N, E
    
def lik_sp_abd_dbh_ssnt(stat_var, beta, model, n, dbh_list, d_list_full, log = True):
    """Probability of a species having abundance n and its individuals having dbh [d1, d2, ..., d_n] in SSNT
    
    Inputs:
    stat_var - [G, S, N, E]
    beta - parameter for SAD
    model - 'ssnt_0', or 'ssnt_1'
    n - abundance
    dbh_list - a list or array of length n with scaled dbh values 
    d_list_full - a list or array of scaled dbh values for all individuals in community
    """
    G, S, N, E = stat_var
    alpha_model = {'ssnt_0': 1, 'ssnt_1': 2/3}
    alpha = alpha_model[model]
    par = N / (sum(np.array(d_list_full)**alpha) - N)
    p_sad_log = stats.logser.logpmf(n, beta)
    isd = ssnt_isd_bounded(alpha, par)
    p_dbh = [isd.pdf(d) for d in dbh_list]
    if log: return p_sad_log + sum([np.log(p_ind) for p_ind in p_dbh])
    else: return np.exp(p_sad_log + sum([np.log(p_ind) for p_ind in dbh_list]))
    
def lik_sp_abd_dbh_asne(stat_var, beta, n, dbh_list, log = True):
    """Probability of a species having abundance n and its individuals having dbh [d1, d2, ..., d_n] in METE
    
    Here unlike SSNT, P(d|n) is not equal to the ISD f(d). 
    Inputs:
    stat_var - [G, S, N, E]
    beta - parameter for SAD
    n - abundance of the species
    dbh_list - a list or array of length n with scaled dbh values
    """
    G, S, N, E = stat_var
    p_sad_log = md.trunc_logser.logpmf(n, beta, N)
    theta = mete_distributions.theta_epsilon(S, N, E)
    p_dbh_log = [theta.logpdf(d ** 2, n) + np.log(2 * d) for d in dbh_list] # Prediction of METE has to be transformed back to distribution of dbh
    if log: return p_sad_log + sum(p_dbh_log)
    else: return np.exp(p_sad_log + sum(p_dbh_log))

def lik_sp_abd_dbh_agsne(stat_var, pars, n, dbh_list, log = True):
    """Probability of a species having abundance n and its individuals having dbh [d1, d2, ..., d_n] in AGSNE.

    Inputs:
    stat_var - [G, S, N, E]
    pars - [lambda1, beta, lambda3, z] in AGSNE
    n - abundance of the species
    dbh_list - a list or array of length n with scaled dbh values
    """
    G, S, N, E = stat_var
    lambda1, beta, lambda3, z = pars
    sad = mete_distributions.sad_agsne(stat_var, pars)
    #logp = 0
    #for d in dbh_list:
        #t = np.exp(-(lambda1 + (beta - lambda1) * n + lambda3 * n * d ** 2))
        #logp_dn = np.log(G / S / z * 2 * d) + np.log(t) - 2 * np.log(1 - t) + np.log(1 - (S + 1) * t ** S + S * t ** (S + 1))
        #logp += logp_dn - np.log(sad.pmf(n))
    #logp += np.log(sad.pmf(n)) # Convert conditional distribution to joint distribution
    logsum = 0
    for d in dbh_list:
        logt = -(lambda1 + (beta - lambda3) * n + lambda3 * n * (d ** 2))
        t = np.exp(logt)
        logsum += logt - 2 * np.log(1 - t) + np.log(1 - (S + 1) * (t ** S) + S * (t ** (S + 1))) + np.log(2 * d)
    logp = logsum + n * np.log(G / S / z) - (n - 1) * sad.logpmf(n)
    if log == True: return logp
    else: return np.exp(logp)
    
def get_obs_pred_sad(raw_data_site, dataset_name, model, out_dir = './out_files/'):
    """Write the observed and predicted RAD to file for a given model.
    
    Inputs:
     raw_data_site - data in the same format as obtained by clean_data_genera(), with
        four columns site, sp, dbh, and genus, and only for one site.
    dataset_name - name of the dataset for raw_data_site.
    model - can take one of three values 'ssnt', 'asne', or 'agsne'. Note that the predicted SAD for SSNT does not 
        change with alternative scaling of D.
    out_dir - directory for output file.
    
    """
    G, S, N, E = get_GSNE(raw_data_site)
    if model == 'ssnt': 
        pred = mete.get_mete_rad(S, N, version = 'untruncated')[0]
    elif model == 'asne': 
        pred = mete.get_mete_rad(S, N)[0]
    elif model == 'agsne': 
        pred = agsne.get_mete_agsne_rad(G, S, N, E)
    obs = np.sort([len(raw_data_site[raw_data_site['sp'] == sp]) for sp in np.unique(raw_data_site['sp'])])[::-1]
    results = np.zeros((S, ), dtype = ('S15, i8, i8'))
    results['f0'] = np.array([raw_data_site['site'][0]] * S)
    results['f1'] = obs
    results['f2'] = pred    
    
    if model == 'ssnt': 
        f1_write = open(out_dir + dataset_name + '_obs_pred_rad_ssnt_0.csv', 'ab')
        f2_write = open(out_dir + dataset_name + '_obs_pred_rad_ssnt_1.csv', 'ab')
        f2 = csv.writer(f2_write)
        f2.writerows(results)
        f2_write.close()
    else: f1_write = open(out_dir + dataset_name + '_obs_pred_rad_' + model + '.csv', 'ab')
    f1 = csv.writer(f1_write)
    f1.writerows(results)
    f1_write.close()

def get_mete_pred_isd_approx(S, N, E):
    """Obtain the dbh2 for N individuals predicted by METE, using the newly derived approximated ISD."""
    psi_appox = mete_distributions.psi_epsilon_approx(S, N, E)
    scaled_rank = [(x + 0.5) / N for x in range(N)]
    pred = np.array([psi_appox.ppf(q) for q in scaled_rank])
    return np.array(pred)

def get_obs_pred_isd(raw_data_site, dataset_name, model, out_dir = './out_files/'):
    """Write the observed and predicted ISD to file for a given model.
    
    Inputs:
     raw_data_site - data in the same format as obtained by clean_data_genera(), with
        four columns site, sp, dbh, and genus, and only for one site.
    dataset_name - name of the dataset for raw_data_site.
    model - can take one of four values 'ssnt_0' (constant growth of diameter D), 
        'ssnt_1' (constant growth of D^2/3), 'asne', or 'agsne'. 
    out_dir - directory for output file.
    
    """
    G, S, N, E = get_GSNE(raw_data_site)
    if model == 'asne':  # Note both ASNE and AGSNE return values in diameter^2, which needs to be transformed back
        pred = get_mete_pred_isd_approx(range(1, N + 1), S, N, E) ** 0.5
    elif model == 'agsne': 
        pred = np.array(agsne.get_mete_agsne_isd(G, S, N, E)) ** 0.5
    else: 
        dbh_scaled = np.array(raw_data_site['dbh'] / min(raw_data_site['dbh']))
        if model == 'ssnt_0': alpha = 1
        elif model == 'ssnt_1': alpha = 2/3
        par = N / (sum(dbh_scaled ** alpha) - N)
        scaled_rank = [(x + 0.5) / N for x in range(N)]
        isd_ssnt = ssnt_isd_bounded(alpha, par)
        pred = np.array([isd_ssnt.ppf(q) for q in scaled_rank])
        
    obs = np.sort(raw_data_site['dbh'] / min(raw_data_site['dbh']))
    results = np.zeros((N, ), dtype = ('S15, f8, f8'))
    results['f0'] = np.array([raw_data_site['site'][0]] * N)
    results['f1'] = obs
    results['f2'] = pred    
    
    f1_write = open(out_dir + dataset_name + '_obs_pred_isd_' + model + '.csv', 'ab')
    f1 = csv.writer(f1_write)
    f1.writerows(results)
    f1_write.close()

def get_obs_pred_sdr(raw_data_site, dataset_name, model, out_dir = './out_files/'):
    """Write the observed and predicted SDR (in unit of D^2) to file for a given model.
    
    Inputs:
     raw_data_site - data in the same format as obtained by clean_data_genera(), with
        four columns site, sp, dbh, and genus, and only for one site.
    dataset_name - name of the dataset for raw_data_site.
    model - can take one of four values 'ssnt_0' (constant growth of diameter D), 
        'ssnt_1' (constant growth of D^2/3), 'asne', or 'agsne'. 
    out_dir - directory for output file.
    
    """
    scaled_d = raw_data_site['dbh'] / min(raw_data_site['dbh'])
    scaled_d2 = scaled_d **2
    G, S, N, E = get_GSNE(raw_data_site)
    lambda1, beta, lambda3 = agsne.get_agsne_lambdas(G, S, N, E)
    theta_agsne = mete_distributions.theta_agsne([G, S, N, E], [lambda1, beta, lambda3, agsne.agsne_lambda3_z(lambda1, beta, S) / lambda3])
    theta_asne = mete_distributions.theta_epsilon(S, N, E)
    if model == 'ssnt_1': alpha = 2/3
    else: alpha = 1
    par = N / (sum(scaled_d ** alpha) - N)
    iisd_ssnt = ssnt_isd_bounded(alpha, par)
   
    pred, obs = [], []
    for sp in np.unique(raw_data_site['sp']):
        n = len(raw_data_site[raw_data_site['sp'] == sp]) # Number of individuals within species
        if model == 'agsne': 
            genus_sp = raw_data_site['genus'][raw_data_site['sp'] == sp][0]
            m = len(np.unique(raw_data_site['sp'][raw_data_site['genus'] == genus_sp])) # Number of specis within genus
            pred.append(theta_agsne.expected(m, n))
        elif model == 'asne': pred.append(theta_asne.E(n))
        elif model in ['ssnt_0', 'ssnt_1']: pred.append(iisd_ssnt.expected_square())
        obs.append(np.mean(scaled_d2[raw_data_site['sp'] == sp]))
    
    results = np.zeros((S, ), dtype = ('S15, f8, f8'))
    results['f0'] = np.array([raw_data_site['site'][0]] * S)
    results['f1'] = obs
    results['f2'] = pred    
    f1_write = open(out_dir + dataset_name + '_obs_pred_sdr_' + model + '.csv', 'ab')
    f1 = csv.writer(f1_write)
    f1.writerows(results)
    f1_write.close()
                
def get_isd_lik_three_models(dat_list, out_dir = './out_files/', cutoff = 9):
    """Function to obtain the community-level log-likelihood (standardized by the number of individuals)
    
    as well as AICc values for METE, SSNT on D, and SSNT on D**(2/3) and write to files. 
    
    """
    for dat_name in dat_list:
        dat = wk.import_raw_data('./data/' + dat_name + '.csv')
        for site in np.unique(dat['site']):
            dat_site = dat[dat['site'] == site]
            S0 = len(np.unique(dat_site['sp']))
            if S0 > cutoff:
                N0 = len(dat_site)
                dbh_scaled = dat_site['dbh'] / min(dat_site['dbh'])
                psi = mete_distributions.psi_epsilon(S0, N0, sum(dbh_scaled ** 2))
                ssnt_isd = ssnt_isd_bounded(1, N0 / (sum(dbh_scaled) - N0))
                ssnt_isd_transform = ssnt_isd_bounded(2/3, N0 / (sum(dbh_scaled ** (2/3)) - N0))
                
                lik_mete, lik_ssnt, lik_ssnt_transform = 0, 0, 0
                for dbh in dbh_scaled:
                    lik_mete += np.log(psi.pdf(dbh ** 2) * 2 * dbh) # psi is on dbh**2
                    lik_ssnt += np.log(ssnt_isd.pdf(dbh))
                    lik_ssnt_transform += np.log(ssnt_isd_transform.pdf(dbh))
                out1 = open(out_dir + 'isd_lik_three_models.txt', 'a')
                print>>out1, dat_name, site, str(lik_mete / N0), str(lik_ssnt / N0), str(lik_ssnt_transform / N0)
                out1.close()
                
                out2 = open(out_dir + 'isd_aicc_three_models.txt', 'a')
                # METE has three parameters (S0, N0, E0) for ISD, while SSNT has two (N0 and sum(dbh**alpha))
                print>>out2, dat_name, site, str(mtools.AICc(lik_mete, 3, N0)), str(mtools.AICc(lik_ssnt, 2, N0)), \
                     str(mtools.AICc(lik_ssnt_transform, 2, N0))
                out2.close()
                
def get_lik_sp_abd_dbh_four_models(raw_data_site, dataset_name, out_dir = './out_files/'):
    """Obtain the summed log likelihood of each species having abundance n and its individuals having 
    
    their specific dbh values for the three models METE, SSNT on D, and SSNT on D ** (2/3).
    
    """
    site = raw_data_site['site'][0]
    G, S, N, E = get_GSNE(raw_data_site)
    lambda1, beta, lambda3 = agsne.get_agsne_lambdas(G, S, N, E)
    beta_ssnt = mete.get_beta(S, N, version = 'untruncated')
    beta_asne = mete.get_beta(S, N) 
    d_list = raw_data_site['dbh'] / min(raw_data_site['dbh'])
    lik_asne, lik_agsne, lik_ssnt_0, lik_ssnt_1 = 0, 0, 0, 0
    for sp in np.unique(raw_data_site['sp']):
        sp_dbh = d_list[raw_data_site['sp'] == sp]
        lik_asne += lik_sp_abd_dbh_asne([G, S, N, E], np.exp(-beta_asne), len(sp_dbh), sp_dbh)
        lik_agsne += lik_sp_abd_dbh_agsne([G, S, N, E], 
                                          [lambda1, beta, lambda3, agsne.agsne_lambda3_z(lambda1, beta, S) / lambda3], len(sp_dbh), sp_dbh)
        lik_ssnt_0 += lik_sp_abd_dbh_ssnt([G, S, N, E], np.exp(-beta_ssnt), 'ssnt_0', len(sp_dbh), sp_dbh, d_list)
        lik_ssnt_1 += lik_sp_abd_dbh_ssnt([G, S, N, E], np.exp(-beta_ssnt), 'ssnt_1', len(sp_dbh), sp_dbh, d_list)
    out = open(out_dir + 'lik_sp_abd_dbh_four_models.txt', 'a')
    print>>out, dataset_name, site, str(lik_asne), str(lik_agsne), str(lik_ssnt_0), str(lik_ssnt_1)
    out.close()    

def plot_likelihood_comp(lik_dir = './out_files/', out_fig_dir = './out_figs/'):
    """Plot the likelihood of the other three models against ASNE."""
    fig = plt.figure(figsize = (3.5, 3.5))
    ax = plt.subplot(1, 1, 1)
    lik_for_sites = import_likelihood_data('lik_sp_abd_dbh_four_models.txt', file_dir = lik_dir)
    lik_asne = lik_for_sites['ASNE']
    other_model_list = ['AGSNE', 'SSNT_N', 'SSNT_M']
    col_list = ['b', '#787878', 'r']
    symbol_list = ['o', 's', '*']
    
    lik_list = list(-np.log(-lik_asne))
    for i, model in enumerate(other_model_list):
        lik_model = lik_for_sites[model]
        plt.scatter(-np.log(-lik_asne), -np.log(-lik_model), s = 20, marker = symbol_list[i], facecolors = col_list[i],
                    edgecolors = 'none', label = model)
        
        lik_list.extend(list(-np.log(-lik_model)))
    
    min_val, max_val = min(lik_list), max(lik_list)
    if min_val < 0: axis_min = 1.1 * min_val
    else: axis_min = 0.9 * min_val
    if max_val < 0: axis_max = 0.9 * max_val
    else: axis_max= 1.1 * max_val    
    plt.plot([axis_min, axis_max], [axis_min, axis_max], 'k-')     
    plt.xlim(axis_min, axis_max)
    plt.ylim(axis_min, axis_max)
    ax.tick_params(axis = 'both', which = 'major', labelsize = 6)
    ax.set_xlabel('ASNE', labelpad = 4, size = 8)
    ax.set_ylabel('Other models', labelpad = 4, size = 8)
    ax.legend(loc = 2, prop = {'size': 8})
    plt.savefig(out_fig_dir + 'lik_comp.png', dpi = 400)

def plot_r2_comp(name_site_combo, dat_dir = './out_files/', out_fig_dir = './out_figs/'):
    """Plot r2 of the three patterns separately for each community."""
    models = ['asne', 'agsne', 'ssnt_0', 'ssnt_1']
    model_names = ['ASNE', 'AGSNE', 'SSNT_N', 'SSNT_M']
    patterns = ['rad', 'isd', 'sdr']
    pattern_names = ['SAD', 'ISD', 'SDR']
    col_list = ['b', '#787878', 'r']
    symbol_list = ['o', 's', '*']
    
    fig = plt.figure(figsize = (10.5, 3.5))
    for i, pattern in enumerate(patterns):
        r2_dic = {'asne':[], 'agsne':[], 'ssnt_0':[], 'ssnt_1':[]}
        r2_list = []
        for j, model in enumerate(models):
            for dat_name, site in name_site_combo:
                pred_obs_model_pattern = wk.import_obs_pred_data(dat_dir + dat_name + '_obs_pred_' + pattern + '_' + model + '.csv')
                pred_obs_site = pred_obs_model_pattern[pred_obs_model_pattern['site'] == site]
                r2 = mtools.obs_pred_rsquare(np.log10(pred_obs_site['obs']), np.log10(pred_obs_site['pred']))
                r2_dic[model].append(r2)
                r2_list.append(r2)
        
        ax = plt.subplot(1, 3, i + 1)
        for j in range(1, 4):
            model = models[j]
            plt.scatter(r2_dic['asne'], r2_dic[model], s = 20, marker = symbol_list[j - 1], facecolors = col_list[j - 1], 
                        edgecolors = 'none', label = model_names[j])
        min_val, max_val = min(r2_list), max(r2_list)
        if min_val < 0: axis_min = 1.1 * min_val
        else: axis_min = 0.9 * min_val
        if max_val < 0: axis_max = 0.9 * max_val
        else: axis_max= 1.1 * max_val    
        plt.plot([axis_min, axis_max], [axis_min, axis_max], 'k-')     
        plt.xlim(axis_min, axis_max)
        plt.ylim(axis_min, axis_max)
        ax.tick_params(axis = 'both', which = 'major', labelsize = 6)
        ax.set_xlabel(r'$R^2$ of ASNE', labelpad = 4, size = 10)
        ax.set_ylabel(r'$R^2$ of the other models', labelpad = 4, size = 10)
        ax.set_title(pattern_names[i], size = 16)
        if i == 0: ax.legend(loc = 2, prop = {'size': 10})
        
    plt.subplots_adjust(left = 0.08, wspace = 0.3)
    plt.tight_layout()
    plt.savefig(out_fig_dir + 'r2_comp.png', dpi = 400)  
         
def bootstrap_SAD(name_site_combo, model, in_dir = './data/', out_dir = './out_files/', Niter = 200):
    """A general function of bootstrapping for SAD applying to all four models. 
    
    Inputs:
    name_site_combo: a list with dat_name and site
    model - takes one of four values 'ssnt_0', 'ssnt_1', 'asne', or 'agsne'
    in_dir - directory of raw data
    out_dir - directory used both in input (obs_pred.csv file) and output 
    Niter - number of bootstrap samples
    
    Output:
    Writes to disk, with one file for R^2 and one for KS statistic.
    
    """
    dat_name, site = name_site_combo
    dat = wk.import_raw_data(in_dir + dat_name + '.csv')
    dat_site = dat[dat['site'] == site]
    dat_clean = clean_data_agsne(dat_site)    
    G, S, N, E = get_GSNE(dat_clean)
    beta_ssnt = mete.get_beta(S, N, version = 'untruncated')
    beta_asne = mete.get_beta(S, N)
    lambda1, beta, lambda3 = agsne.get_agsne_lambdas(G, S, N, E)
    sad_agsne = mete_distributions.sad_agsne([G, S, N, E], [lambda1, beta, lambda3, agsne.agsne_lambda3_z(lambda1, beta, S) / lambda3])
    dist_for_model = {'ssnt_0': stats.logser(np.exp(-beta_ssnt)), 
                      'ssnt_1': stats.logser(np.exp(-beta_ssnt)), 
                      'asne': md.trunc_logser(np.exp(-beta_asne), N),
                      'agsne': sad_agsne}
    dist = dist_for_model[model]
    pred_obs = wk.import_obs_pred_data(out_dir + dat_name + '_obs_pred_rad_' + model + '.csv')
    pred = pred_obs[pred_obs['site'] == site]['pred'][::-1]
    obs = pred_obs[pred_obs['site'] == site]['obs'][::-1]
    
    out_list_rsquare = [dat_name, site, str(mtools.obs_pred_rsquare(np.log10(obs), np.log10(pred)))]
    emp_cdf = mtools.get_emp_cdf(obs)
    out_list_ks = [dat_name, site, str(max(abs(emp_cdf - np.array([dist.cdf(x) for x in obs]))))]
    
    for i in range(Niter):
        obs_boot = np.array(sorted(dist.rvs(S)))
        cdf_boot = np.array([dist.cdf(x) for x in obs_boot])
        emp_cdf_boot = mtools.get_emp_cdf(obs_boot)
        out_list_rsquare.append(str(mtools.obs_pred_rsquare(np.log10(obs_boot), np.log10(pred))))
        out_list_ks.append(str(max(abs(emp_cdf_boot - np.array(cdf_boot)))))
    
    wk.write_to_file(out_dir + 'SAD_bootstrap_' + model + '_rsquare.txt', ",".join(str(x) for x in out_list_rsquare))
    wk.write_to_file(out_dir + 'SAD_bootstrap_' + model + '_ks.txt', ",".join(str(x) for x in out_list_ks))

def bootstrap_ISD(name_site_combo, model, in_dir = './data/', out_dir = './out_files/', Niter = 200):
    """A general function of bootstrapping for ISD applying to all four models. 
    
    Inputs:
    name_site_combo: a list with dat_name and site
    model - takes one of four values 'ssnt_0', 'ssnt_1', 'asne', or 'agsne'
    in_dir - directory of raw data
    out_dir - directory used both in input (obs_pred.csv file) and output 
    Niter - number of bootstrap samples
    
    Output:
    Writes to disk, with one file for R^2 and one for KS statistic.
    
    """
    dat_name, site = name_site_combo
    dat = wk.import_raw_data(in_dir + dat_name + '.csv')
    dat_site = dat[dat['site'] == site]
    dat_clean = clean_data_agsne(dat_site)    
    G, S, N, E = get_GSNE(dat_clean)
    lambda1, beta, lambda3 = agsne.get_agsne_lambdas(G, S, N, E)
    isd_agsne = mete_distributions.psi_agsne([G, S, N, E], [lambda1, beta, lambda3, agsne.agsne_lambda3_z(lambda1, beta, S) / lambda3])
    isd_asne = mete_distributions.psi_epsilon_approx(S, N, E)
    dbh_scaled = np.array(dat_clean['dbh'] / min(dat_clean['dbh']))
    isd_ssnt_0 = ssnt_isd_bounded(1, N / (sum(dbh_scaled ** 1) - N))
    isd_ssnt_1 = ssnt_isd_bounded(2/3, N / (sum(dbh_scaled ** (2/3)) - N))
    dist_for_model = {'ssnt_0': isd_ssnt_0, 'ssnt_1': isd_ssnt_1, 'asne': isd_asne, 'agsne': isd_agsne}
    dist = dist_for_model[model]
    pred_obs = wk.import_obs_pred_data(out_dir + dat_name + '_obs_pred_isd_' + model + '.csv')
    pred = pred_obs[pred_obs['site'] == site]['pred']
    obs = pred_obs[pred_obs['site'] == site]['obs']
    
    out_list_rsquare = [dat_name, site, str(mtools.obs_pred_rsquare(np.log10(obs), np.log10(pred)))]
    wk.write_to_file(out_dir + 'ISD_bootstrap_' + model + '_rsquare.txt', ",".join(str(x) for x in out_list_rsquare), new_line = False)
    emp_cdf = mtools.get_emp_cdf(obs)
    out_list_ks = [dat_name, site, str(max(abs(emp_cdf - np.array([dist.cdf(x) for x in obs]))))]
    wk.write_to_file(out_dir + 'ISD_bootstrap_' + model + '_ks.txt', ",".join(str(x) for x in out_list_ks), new_line = False)
    
    num_pools = 8  # Assuming that 8 pools are to be created
    for i in xrange(Niter):
        obs_boot = []
        cdf_boot = []
        while len(obs_boot) < N:
            pool = multiprocessing.Pool(num_pools)
            out_sample = pool.map(wk.generate_isd_sample, [dist for j in xrange(num_pools)])
            for combo in out_sample:
                cdf_sublist, sample_sublist = combo
                obs_boot.extend(sample_sublist)
                cdf_boot.extend(cdf_sublist)
            pool.close()
            pool.join()
        if model in ['asne', 'agsne']: obs_boot = np.sort(obs_boot[:N]) ** 0.5 # Convert to diameter
        else: obs_boot = np.sort(obs_boot[:N])
        sample_rsquare = mtools.obs_pred_rsquare(np.log10(obs_boot), np.log10(pred))
        sample_ks = max(abs(emp_cdf - np.sort(cdf_boot[:N])))
        
        wk.write_to_file(out_dir + 'ISD_bootstrap_' + model + '_rsquare.txt', "".join([',', str(sample_rsquare)]), new_line = False)
        wk.write_to_file(out_dir + 'ISD_bootstrap_' + model + '_ks.txt', "".join([',', str(sample_ks)]), new_line = False)
    
    wk.write_to_file(out_dir + 'ISD_bootstrap_' + model + '_rsquare.txt', '\t')
    wk.write_to_file(out_dir + 'ISD_bootstrap_' + model + '_ks.txt', '\t')
    wk.write_to_file(out_dir + 'ISD_bootstrap_' + model + '_ks.txt', ",".join(str(x) for x in out_list_ks))
    
def bootstrap_SDR(name_site_combo, model, in_dir = './data/', out_dir = './out_files/', Niter = 200):
    """A general function of bootstrapping for ISD applying to all four models. 
    
    Inputs:
    name_site_combo: a list with dat_name and site
    model - takes one of four values 'ssnt_0', 'ssnt_1', 'asne', or 'agsne'
    in_dir - directory of raw data
    out_dir - directory used both in input (obs_pred.csv file) and output 
    Niter - number of bootstrap samples
    
    Output:
    Writes to one file on disk for R^2.
    
    """
    dat_name, site = name_site_combo
    dat = wk.import_raw_data(in_dir + dat_name + '.csv')
    dat_site = dat[dat['site'] == site]
    dat_clean = clean_data_agsne(dat_site)    
    G, S, N, E = get_GSNE(dat_clean)
    lambda1, beta, lambda3 = agsne.get_agsne_lambdas(G, S, N, E)
    
    par_list = []
    for sp in np.unique(dat_clean['sp']):
        dat_sp = dat_clean[dat_clean['sp'] == sp]
        n = len(dat_sp)
        genus_sp = dat_sp['genus'][0]
        m = len(np.unique(dat_clean[dat_clean['genus'] == genus_sp]['sp']))
        par_list.append([m, n])
        
    pred_obs = wk.import_obs_pred_data(out_dir + dat_name + '_obs_pred_sdr_' + model + '.csv')
    pred = pred_obs[pred_obs['site'] == site]['pred']
    obs = pred_obs[pred_obs['site'] == site]['obs'] 
    out_list_rsquare = [dat_name, site, str(mtools.obs_pred_rsquare(np.log10(obs), np.log10(pred)))]
    
    iisd_agsne = mete_distributions.theta_agsne([G, S, N, E], [lambda1, beta, lambda3, agsne.agsne_lambda3_z(lambda1, beta, S) / lambda3])
    iisd_asne = mete_distributions.theta_epsilon(S, N, E)
    dbh_scaled = np.array(dat_clean['dbh'] / min(dat_clean['dbh']))
    iisd_ssnt_0 = ssnt_isd_bounded(1, N / (sum(dbh_scaled ** 1) - N))
    iisd_ssnt_1 = ssnt_isd_bounded(2/3, N / (sum(dbh_scaled ** (2/3)) - N))
    dist_for_model = {'ssnt_0': iisd_ssnt_0, 'ssnt_1': iisd_ssnt_1, 'asne': iisd_asne, 'agsne': iisd_agsne}
    dist = dist_for_model[model]
        
    for i in range(Niter):
        if model in ['ssnt_0', 'ssnt_1']: obs_boot = np.array([np.mean((dist.rvs(par[1])) ** 2) for par in par_list]) # Here par[1] is n for each species
        elif model == 'asne': 
            obs_boot = np.array([np.mean(np.array(dist.rvs(par[1], par[1]))) for par in par_list])
        else:
            obs_boot = np.array([np.mean(np.array(dist.rvs(par[1], par[1], par[0]))) for par in par_list])
        out_list_rsquare.append(str(mtools.obs_pred_rsquare(np.log10(obs_boot), np.log10(pred))))
    
    wk.write_to_file(out_dir + 'SDR_bootstrap_' + model + '_rsquare.txt', ",".join(str(x) for x in out_list_rsquare))
            
def plot_hist_quan(dat_file, dat_type = 'r2', ax = None):    
    """Similar to the function under the same name in working_functions,
    
    only with different text.
    
    """
    if not ax:
        fig = plt.figure(figsize = (3.5, 3.5))
        ax = plt.subplot(111)
  
    quan_list = []
    num_row = len(dat_file['dataset'])
    for i in range(num_row):
        dat_row = list(dat_file[i])
        stat_orig = dat_row[2]
        stat_sim = dat_row[3:]
        if dat_type == 'r2':
            quan_row = len([x for x in stat_sim if x < stat_orig]) / len(stat_sim)
        else: quan_row = len([x for x in stat_sim if x > stat_orig]) / len(stat_sim)
        quan_list.append(quan_row)
    
    n, bins, patches = plt.hist(quan_list,  facecolor='grey', alpha=0.5, range = (0, 1))
    ax.annotate('Zero quantile: ' + str(len([x for x in quan_list if x == 0])) + '/' + str(len(quan_list)), \
                xy = (0.05, 0.92), xycoords = 'axes fraction', fontsize = 8)
    ax.tick_params(axis = 'both', which = 'major', labelsize = 6)
    plt.ylim(0, max(n) * 1.2)
    return ax
      
def plot_bootstrap(model, Niter = 500, out_file_dir = './out_files/', out_fig_dir = './out_figs/'):
    """Plot the bootstrap results for the a given model and the three patterns (SAD, ISD, SDR). 
    
    The output is a 3*2 plot (with the last subplot missing) with name bootstrap_model.pdf.
    
    """
    patterns = ['SAD', 'ISD', 'SDR']
    stats = ['rsquare', 'ks']
    titles = [r'$R^2$', 'K-S Statistic']
    fig = plt.figure(figsize = (5, 8))
    iplot = 1
    for pattern in patterns:
        for stat in stats:
            if iplot < 6:
                boot_dir = out_file_dir + pattern + '_bootstrap_' + model + '_' + stat + '.txt'
                boot_out = import_bootstrap_file_incomp(boot_dir, Niter = Niter)
                ax = plt.subplot(3, 2, iplot)
                if stat == 'ks': plot_hist_quan(boot_out, dat_type = 'ks', ax = ax)
                else: plot_hist_quan(boot_out, ax = ax)
                plt.xlabel('Samples closer to prediction', fontsize = 8)
                plt.ylabel('Number of communities', fontsize = 8)
                if iplot in [1, 2]: ax.set_title(titles[iplot - 1], size = 14,y = 1.1)
                if iplot in [1, 3, 5]: plt.figtext(0.01, 0.8 - int(np.floor(iplot / 2)) * 0.32, patterns[int(np.floor(iplot / 2))], 
                                                                   size = 14, rotation = 'horizontal')
                iplot += 1
    plt.subplots_adjust(left = 0.17, top = 0.95, bottom = 0.05, right = 0.95, wspace = 0.3, hspace = 0.3)
    plt.savefig(out_fig_dir + 'bootstrap_' + model + '.png', dpi = 400)
    
def plot_obs_pred_four_models(dat_list, out_file_dir = './out_files/', out_fig_dir = './out_figs/'):
    """Create the obs-pred plots for the three patterns (SAD, ISD, and SDR) and four models.
    
    The output is a 4*3 plot with name obs_pred_3patterns_4models.pdf.
    
    """
    dat_list_exist = [x for x in dat_list if os.path.isfile(out_file_dir + x + '_obs_pred_rad_asne.csv')]
    model_list = ['asne', 'agsne', 'ssnt_0', 'ssnt_1']
    pattern_list = ['rad', 'isd', 'sdr']
    model_names = ['ASNE', 'AGSNE', 'SSNT_N', 'SSNT_M']
    pattern_names = ['SAD', 'ISD', 'SDR']
    xylabel = {'rad': ['Predicted abundance', 'Observed abundance'], 
               'isd': ['Predicted diameter', 'Observed diameter'],
               'sdr': ['Predicted average metabolic rate', 'Observed average metabolic rate']}
    fig = plt.figure(figsize = (8, 10))
    iplot = 1
    for model in model_list:
        for pattern in pattern_list:
            filename = '_obs_pred_' + pattern + '_' + model + '.csv'
            sites, obs, pred = wk.get_obs_pred_from_file(dat_list_exist,out_file_dir, filename)
            ax = plt.subplot(4, 3, iplot)
            ax = wk.plot_obs_pred(obs, pred, 2, True, ax = ax)
            xlab, ylab = xylabel[pattern]
            ax.set_xlabel(xlab, labelpad = 4, size = 8)
            ax.set_ylabel(ylab, labelpad = 4, size = 8)
            if iplot in [1, 2, 3]:  ax.set_title(pattern_names[iplot - 1], size = 14,y = 1.03)
            if iplot in [1, 4, 7, 10]: plt.figtext(0.01, 0.85 - int(iplot / 3) * 0.24, model_names[int(iplot / 3)], 
                                                   size = 14, rotation = 'horizontal')
            iplot += 1
    plt.subplots_adjust(left = 0.17, top = 0.95, bottom = 0.05, right = 0.95, wspace = 0.3, hspace = 0.3)
    plt.savefig(out_fig_dir + 'obs_pred_3patterns_4models.png', dpi = 400)