<jupyter_start><jupyter_text>Parameter RecoveryThis notebook conducts parameter recovery simulations for modified Rachlin discount function ([Vincent, & Stewart, 2020](https://doi.org/10.1016/j.cognition.2020.104203)).$$V(R, D, k) = R \cdot \frac{1}{1+(k \cdot D)^s}$$where $R$ is a reward, delivered at a delay $D$. The parameters are:- $k$ is the normally interpreted as the discount rate. Although technically in this case it is the product of the discount rate and the constant term in Steven's Power Law.- $s$ is the exponent in Steven's Power Law. Important noteIn order for this to be a meaningful parameter recovery excercise then the data generating model defined in `generate_responses` _must_ be exactly the same model that is used for inference in `infer_parameters`.<jupyter_code>import numpy as np
import pandas as pd
from scipy.stats import norm, bernoulli, uniform
import pymc3 as pm
import math

import matplotlib.pyplot as plt
%config InlineBackend.figure_format = 'retina'
plt.rcParams.update({'font.size': 14})
import matplotlib.ticker as ticker

# Initialize random number generator
np.random.seed(1234)

import sys
print("Python version:\n{}\n".format(sys.version))
print(f"PyMC3 version: {pm.__version__}")
out_dir = 'output/'<jupyter_output><empty_output><jupyter_text>Define parameters of the parameter recovery simulations<jupyter_code>n_simulations = 200
export = True

sample_options = {'tune': 1000, 'draws': 2000,
                  'chains': 4, 'cores': 4,
                  'nuts_kwargs': {'target_accept': 0.95}}<jupyter_output><empty_output><jupyter_text>Define the core `parameter_recovery` function<jupyter_code>def parameter_recovery():
    '''
    Conducts parameter recovery for a single simulated experiment.
    Return a single row of a DataFrame which contains the true and recovered
    parameter values.
    '''
    params_true = generate_true_params()
    expt_data = simulate_experiment(params_true)
    params_inferred = infer_parameters(expt_data)
    row_data = pd.concat([params_true, params_inferred], axis=1)
    return row_data<jupyter_output><empty_output><jupyter_text>Define these functions Our set of true parameters will be generated in `generate_true_params`. These are sampled from normal distributions which equate to our prior beliefs over participant level $\log(k)$ and $\log(s)$ used in our parameter estimation step.<jupyter_code>def generate_true_params():
    '''Sample ONCE from our distribution of plausible parameter values'''
    
    logk_min, logk_max = -5, -1.5
    logk = uniform.rvs(loc=logk_min, scale=logk_max-logk_min) 
    
    logs_min, logs_max = np.log(0.5), np.log(3)
    logs = uniform.rvs(loc=logs_min, scale=logs_max-logs_min) 
    
    return pd.DataFrame({'logk': [logk],
                         'logs': [logs]})
def simulate_experiment(params_true, ϵ=0.01):
    '''Run a simulated experiment, returning simulated behavioural data'''
    designs = generate_designs()
    responses, _ = generate_responses(designs, params_true, ϵ)
    return pd.concat([designs, responses], axis=1)


def generate_designs():
    '''Generate designs (RA, DA, RB, DB). This should precisely match the 
    set of questions we used in the actual experiment.'''
    
    n = 50
    RA_vals = np.array([6, 12, 18, 24, 30, 36, 42, 48, 54, 60])
    DB_vals = np.array([7, 15, 29, 56, 101])
    
    # define constant values
    DA = np.zeros(n)
    RB = np.full(n, 60)

    # shuffle index for DB
    DB_index = np.arange(len(DB_vals))
    np.random.shuffle(DB_index)
    
    # fill remaining design dimensions by iterating over DB (shuffled) and RA
    DB = []
    RA = []
    for db_index in DB_index:
        for ra in RA_vals:
            DB.append(DB_vals[db_index])
            RA.append(ra)
    
    DB = np.array(DB)
    RA = np.array(RA)
    
    designs = pd.DataFrame({'RA': RA, 'DA': DA, 'RB': RB, 'DB': DB})
    return designs


def generate_responses(designs, params_true, ϵ):
    '''Generate simulated responses for the given designs and parameters'''
    
    # unpack designs
    RA = designs['RA'].values
    DA = designs['DA'].values
    RB = designs['RB'].values
    DB = designs['DB'].values
    
    # unpack parameters
    logk = params_true['logk'].values
    logs = params_true['logs'].values
    
    k = np.exp(logk)
    s = np.exp(logs)
    
    VA = RA * (1 / (1 + (k*DA)**s))
    VB = RB * (1 / (1 + (k*DB)**s))
    decision_variable = VB-VA
    p_choose_B = ϵ + (1 - 2 * ϵ) * (1 / (1 + np.exp(-1.7 * decision_variable)))
    responses = bernoulli.rvs(p_choose_B)
    return pd.DataFrame({'R': responses}), p_choose_B
def infer_parameters(data):
    '''Infer parameter values based on response data.
    Return the posterior mean parameter estimates'''
    
    model = generate_model(data)
    
    # do the inference
    with model:
        trace = pm.sample(**sample_options)
    
    return extract_info_from_trace(trace)


def generate_model(data):
    '''Generate a PyMC3 model with the given observed data'''
    
    # decant data
    R = data['R'].values
    RA, DA = data['RA'].values, data['DA'].values
    RB, DB = data['RB'].values, data['DB'].values
    
    with pm.Model() as model:
        # define priors
        logk = pm.Normal('logk', mu=-3.760, sd=3)
        logs = pm.Normal('logs', mu=0, sd=1)
        
        VA = pm.Deterministic('VA', value_function(RA, DA, logk, logs))
        VB = pm.Deterministic('VB', value_function(RB, DB, logk, logs))
        P_chooseB = pm.Deterministic('P_chooseB', choice_psychometric(VB-VA))
        
        R = pm.Bernoulli('R', p=P_chooseB, observed=R)
        
    return model


# helper functions for the model

def value_function(reward, delay, logk, logs):
    '''Calculate the present subjective value of a given prospect'''
    k = pm.math.exp(logk)
    s = pm.math.exp(logs)
    return reward / (1.0+(k*delay)**s)

def choice_psychometric(x, ϵ=0.01):
    # x is the decision variable
    return ϵ + (1.0-2.0*ϵ) * (1/(1+pm.math.exp(-1.7*(x))))


def trace_quantiles(x):
    return pd.DataFrame(pm.quantiles(x, [2.5, 5, 50, 95, 97.5]))

def extract_info_from_trace(trace):
    '''Return a 1-row DataFrame of summary statistics (i.e. means, ranges)
    of the parameters of interest'''
    
    # useful PyMC3 function to get summary statistics
    summary = pm.summary(trace, ['logk', 'logs'])

    logk = summary['mean']['logk']
    logkL = summary['hpd_2.5']['logk']
    logkU = summary['hpd_97.5']['logk']

    logs = summary['mean']['logs']
    logsL = summary['hpd_2.5']['logs']
    logsU = summary['hpd_97.5']['logs']

    return pd.DataFrame({'logk_est': [logk], 'logk_est_L': [logkL], 'logk_est_U': [logkU],
                         'logs_est': [logs], 'logs_est_L': [logsL], 'logs_est_U': [logsU]})<jupyter_output><empty_output><jupyter_text>Conduct the parameter recovery simulations🔥 Warning: this will take some time to compute.<jupyter_code>results = [parameter_recovery() for n in range(n_simulations)]

results = pd.concat(results, ignore_index=True)
results.head()
# export results
results.to_csv(f'{out_dir}parameter_recovery_results.csv')<jupyter_output><empty_output><jupyter_text>Visualisation of parameter recovery resultsPlot the true parameter values (x axis) against the estimated parameter values (y axis).<jupyter_code>results.head()
plt.rcParams.update({'font.size': 14})

fig, ax = plt.subplots(1, 2, figsize=(10, 8))

tick_spacing = 1
params = ['logk', 'logs']

for i, parameter in enumerate(params):
    
    x = results[parameter]
    y = results[f'{parameter}_est']
    
    yerrL = results[f'{parameter}_est'] - results[f'{parameter}_est_L']
    yerrU = results[f'{parameter}_est_U'] - results[f'{parameter}_est']
    
    # line of equality
    ax[i].plot([np.min(x), np.max(x)], [np.min(x), np.max(x)], c='k')
    
    # errorbar
    ax[i].errorbar(x, y, yerr=[yerrL, yerrU], fmt='o', alpha=0.5)
    
    ax[i].set_xlabel(f'true {parameter}')
    ax[i].set_ylabel(f'estimated {parameter}')
    ax[i].grid()
    #ax[i].set_aspect('equal', 'box')
    
    # set same tick spacing for both axes
    ax[i].xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing))
    ax[i].yaxis.set_major_locator(ticker.MultipleLocator(tick_spacing))
    
plt.tight_layout()

if export:
    plt.savefig(f'{out_dir}parameter_recovery.pdf', bbox_inches='tight')<jupyter_output><empty_output><jupyter_text>Visualise set of true parameters used in the parameter recovery<jupyter_code>fig, ax = plt.subplots(1, 1, figsize=(10, 8))

tick_spacing = 1
params = ['logk', 'logs']

ax.scatter(x=results['logk'], y=results['logs'])
ax.set(xlabel='true $\log(k)$', ylabel='true $\log(s)$')

ax.grid()<jupyter_output><empty_output><jupyter_text>Visualise a grid of discount functions spanning this parameter range<jupyter_code>logk_res, logs_res = 6, 6

logk_min, logk_max = -5, -1.5
logk_vec = np.linspace(logk_min,logk_max, logk_res)

logs_min, logs_max = np.log(0.2), np.log(3)
logs_vec = np.linspace(logs_min,logs_max, logs_res)

logk_matrix, logs_matrix = np.meshgrid(logk_vec, logs_vec)<jupyter_output><empty_output><jupyter_text>Flip the row ordering so that we plot high values of s in the top rows<jupyter_code>logk_matrix = np.flipud(logk_matrix)
logs_matrix = np.flipud(logs_matrix)
np.exp(logs_matrix)
logs_vec = np.flip(logs_vec)
np.exp(logs_vec)
def discount_function(delay, k, s):
    ''' This is the MODIFIED Rachlin discount function. This is outlined
    in Vincent & Stewart (2020).
    <NAME>., & <NAME>. (2020). The case of muddled 
    units in temporal discounting. https://doi.org/10.1016/j.cognition.2020.104203
    '''
    return 1 / (1.0+(k*delay)**s)
plt.rcParams.update({'font.size': 14})

fig, ax = plt.subplots(logs_res, logk_res,
                       gridspec_kw = {'wspace':0, 'hspace':0},
                       figsize=(14, 14))

max_delay = 101
delays = np.linspace(0, max_delay, 1000)

for row in range(logs_res):
    for col in range(logk_res):
        
        logk = logk_matrix[row, col]
        logs = logs_matrix[row, col]
        
        ax[row,col].plot(delays, 
                         discount_function(delays, np.exp(logk), np.exp(logs)),
                         c='k')
        ax[row,col].set(xlim=[0, 101], ylim=[0, 1])

        if col > 0 or row < logs_res-1:
            ax[row,col].set_yticklabels([])
            ax[row,col].set_xticklabels([])
        else:
            ax[row,col].set_xlabel("delay [sec]")
            ax[row,col].set_ylabel("discount fraction")
            
        ax[row,col].set_xticks(np.arange(0,101,20))

# Create row titles ------------------------------------------------------------
row_headings = [f's={np.exp(logs):.2f}' for logs in logs_vec]
                            
pad = 13 # in points
for axis, row_title in zip(ax[:,0], row_headings):
    axis.annotate(row_title, xy=(0, 0.5), xytext=(-axis.yaxis.labelpad - pad, 0),
                  xycoords=axis.yaxis.label, textcoords='offset points',
                  size='large', ha='center', va='center', rotation=90)
    
# Create col titles ------------------------------------------------------------  
col_headings = [f'ln(k)=\n{logk:.2f}' for logk in logk_vec]

for axis, col in zip(ax[0], col_headings):
    axis.annotate(col, xy=(0.5, 1), xytext=(0, pad),
                  xycoords='axes fraction', textcoords='offset points',
                  size='large', ha='center', va='baseline')
            
        
plt.savefig(f'{out_dir}df_grid.pdf', bbox_inches='tight')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14729894042035782, 'fraction_numerical': 0.020844189682126108, 'mean_word_length': 3.343643907959261, 'pattern_counts': {'":': 0, '<': 40, '<?xml version=': 0, '>': 40, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '332196', 'license': 'MIT', 'n_tokens_mistral': 4148, 'n_tokens_neox': 3896, 'n_words': 967, 'path': '1_parameter_recovery.ipynb', 'repo_name': 'drbenvincent/Xu-GonzalezVallejo-Vincent', 'size': '11525'}