<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Here we will test parameter recovery and model comparison for Rescorla-Wagner (RW), Hierarchical Gaussian Filters (HGF), and Switching Gaussian Filters (SGF) models of the social influence task.
Step1:  Lets start by generating some behavioral data from the social influence task. Here green advice/choice is encoded as 0 and the blue advice/choice is encoded as 1.
Step2:  plot performance of different agents in different blocks
Step3:  Fit simulated behavior
Step4:  Compute fit quality and plot posterior estimates from a hierarchical parameteric model
Step5:  fit HGF agent to simulated data
Step6:  Plot posterior estimates from simulated data for the HGF agent
Step7:  Test model comparison
<ASSISTANT_TASK:>
Python Code:
import numpy as np
from scipy import io
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
sns.set(style = 'white', color_codes = True)
%matplotlib inline
import sys
import os
import os
cwd = os.getcwd()
sys.path.append(cwd[:-len('befit/examples/social_influence')])
Explanation: Here we will test parameter recovery and model comparison for Rescorla-Wagner (RW), Hierarchical Gaussian Filters (HGF), and Switching Gaussian Filters (SGF) models of the social influence task.
End of explanation
import torch
from torch import ones, zeros, tensor
torch.manual_seed(1234)
nsub = 50 #number of subjects
trials = 120 #number of samples
from befit.tasks import SocialInfluence
from befit.simulate import Simulator
from befit.inference import Horseshoe, Normal
from befit.agents import RLSocInf, HGFSocInf, SGFSocInf
# load stimuli (trial offers, advices, and reliability of advices)
reliability = torch.from_numpy(np.load('advice_reliability.npy')).float()
reliability = reliability.reshape(trials, -1, 1).repeat(1, 1, nsub).reshape(trials, -1).unsqueeze(0)
offers = torch.from_numpy(np.load('offers.npy')).reshape(trials, -1, 1).repeat(1, 1, nsub)\
            .reshape(trials, -1).unsqueeze(0)
stimuli = {'offers': offers, 
           'reliability': reliability}
socinfl = SocialInfluence(stimuli, nsub=nsub)
# RL agent
rl_agent = RLSocInf(runs=2*nsub, trials=trials)
trans_pars1 = torch.arange(-.5,.5,1/(2*nsub)).reshape(-1, 1) + tensor([[-2., 4., 0., 0.]]) 
rl_agent.set_parameters(trans_pars1)
sim1 = Simulator(socinfl, rl_agent, runs=2*nsub, trials=trials)
sim1.simulate_experiment()
# HGF agent
hgf_agent = HGFSocInf(runs=2*nsub, trials=trials)
trans_pars2 = torch.arange(-.5, .5, 1/(2*nsub)).reshape(-1, 1) + tensor([[2., 0., 4., 0., 0.]]) 
hgf_agent.set_parameters(trans_pars2)
sim2 = Simulator(socinfl, hgf_agent, runs=2*nsub, trials=trials)
sim2.simulate_experiment()
# SGF agent
sgf_agent = SGFSocInf(runs=2*nsub, trials=trials)
trans_pars3 = torch.arange(-.5, .5, 1/(2*nsub)).reshape(-1, 1) + tensor([[-2., -1., 4., 0., 0.]])
sgf_agent.set_parameters(trans_pars3)
sim3 = Simulator(socinfl, sgf_agent, runs=2*nsub, trials=trials)
sim3.simulate_experiment();
def posterior_accuracy(labels, df, vals):
    for i, lbl in enumerate(labels):
        std = df.loc[df['parameter'] == lbl].groupby(by='subject').std()
        mean = df.loc[df['parameter'] == lbl].groupby(by='subject').mean()
        print(lbl, np.sum(((mean+2*std).values[:, 0] > vals[i])*((mean-2*std).values[:, 0] < vals[i]))/(2*nsub))
Explanation: Lets start by generating some behavioral data from the social influence task. Here green advice/choice is encoded as 0 and the blue advice/choice is encoded as 1.
End of explanation
def compute_mean_performance(outcomes, responses):
    cc1 = (outcomes * responses > 0.).float()  # accept reliable offer
    cc2 = (outcomes * (1 - responses) < 0.).float() # reject unreliable offer
    
    return torch.einsum('ijk->k', cc1 + cc2)/trials
perf1 = compute_mean_performance(sim1.stimulus['outcomes'][..., 0],
                                 sim1.responses.float()).numpy().reshape(2, -1) 
print('RL agent: ', np.median(perf1, axis = -1))
fig, ax = plt.subplots(1,2, sharex = True, sharey = True)
ax[0].hist(perf1[0]);
ax[1].hist(perf1[1]);
fig.suptitle('RL agent', fontsize = 20);
ax[0].set_ylim([0, 20]);
ax[0].set_xlim([.5, 1.]);
perf2 = compute_mean_performance(sim2.stimulus['outcomes'][..., 0],
                                 sim2.responses.float()).numpy().reshape(2, -1)
print('HGF agent: ', np.median(perf2, axis = -1))
fig, ax = plt.subplots(1,2, sharex = True, sharey = True)
ax[0].hist(perf2[0]);
ax[1].hist(perf2[1]);
fig.suptitle('HGF agent', fontsize = 20);
ax[0].set_ylim([0, 20]);
ax[0].set_xlim([.5, 1.]);
perf3 = compute_mean_performance(sim3.stimulus['outcomes'][..., 0],
                                 sim3.responses.float()).numpy().reshape(2, -1)
print('SGF agent: ', np.median(perf3, axis = -1))
fig, ax = plt.subplots(1,2, sharex = True, sharey = True)
ax[0].hist(perf3[0]);
ax[1].hist(perf3[1]);
fig.suptitle('SGF agent', fontsize = 20);
ax[0].set_ylim([0, 20]);
ax[0].set_xlim([.5, 1.]);
Explanation: plot performance of different agents in different blocks
End of explanation
stimulus = sim1.stimulus
stimulus['mask'] = torch.ones(1, 120, 100)
rl_infer = Horseshoe(rl_agent, stimulus, sim1.responses)
rl_infer.infer_posterior(iter_steps=200)
labels = [r'$\alpha$', r'$\zeta$', r'$\beta$', r'$\theta$']
tp_df = rl_infer.sample_posterior(labels, n_samples=1000)
sim1.responses.dtype
Explanation: Fit simulated behavior
End of explanation
labels = [r'$\alpha$', r'$\zeta$', r'$\beta$', r'$\theta$']
trans_pars_rl = tp_df.melt(id_vars='subject', var_name='parameter')
vals = [trans_pars1[:,0].numpy(), trans_pars1[:, 1].numpy(), trans_pars1[:, 2].numpy(), trans_pars1[:, 3].numpy()]
posterior_accuracy(labels, trans_pars_rl, vals)
plt.figure()    
#plot convergence of stochasitc ELBO estimates (log-model evidence)
plt.plot(rl_infer2.loss[-400:])
g = sns.FacetGrid(trans_pars_rl, col="parameter", height=3, sharey=False);
g = (g.map(sns.lineplot, 'subject', 'value', ci='sd'));
labels = [r'$\alpha$', r'$\zeta$', r'$\beta$', r'bias']
for i in range(len(labels)):
    g.axes[0,i].plot(np.arange(2*nsub), trans_pars1[:,i].numpy(),'ro', zorder = 0);
Explanation: Compute fit quality and plot posterior estimates from a hierarchical parameteric model
End of explanation
stimulus = sim2.stimulus
stimulus['mask'] = torch.ones(1, 120, 100)
hgf_infer = Horseshoe(hgf_agent, stimulus, sim2.responses)
hgf_infer.infer_posterior(iter_steps=200)
labels = [r'$\mu_0^2$', r'$\eta$', r'$\zeta$', r'$\beta$', r'$\theta$']
hgf_tp_df, hgf_mu_df, hgf_sigma_df = hgf_infer.sample_posterior(labels, n_samples=1000)
labels = [r'$\mu_0^2$', r'$\eta$', r'$\zeta$', r'$\beta$', r'$\theta$']
trans_pars_hgf = hgf_tp_df.melt(id_vars='subject', var_name='parameter')
vals = [trans_pars2[:, i].numpy() for i in range(len(labels))]
posterior_accuracy(labels, trans_pars_hgf, vals)
Explanation: fit HGF agent to simulated data
End of explanation
plt.figure()    
#plot convergence of stochasitc ELBO estimates (log-model evidence)
plt.plot(hgf_infer.loss[-400:])
g = sns.FacetGrid(trans_pars_hgf, col="parameter", height=3, sharey=False);
g = (g.map(sns.lineplot, 'subject', 'value', ci='sd'));
for i in range(len(labels)):
    g.axes[0,i].plot(np.arange(2*nsub), trans_pars2[:,i].numpy(),'ro', zorder = 0);
stimulus = sim3.stimulus
stimulus['mask'] = torch.ones(1, 120, 100)
sgf_infer = Horseshoe(sgf_agent, stimulus, sim3.responses)
sgf_infer.infer_posterior(iter_steps=200)
labels = [r'$\rho_1$', r'$h$',  r'$\zeta$', r'$\beta$', r'$\theta$']
sgf_tp_df, sgf_mu_df, sgf_sigma_df = sgf_infer.sample_posterior(labels, n_samples=1000)
labels = [r'$\rho_1$', r'$h$',  r'$\zeta$', r'$\beta$', r'$\theta$']
trans_pars_sgf = sgf_tp_df.melt(id_vars='subject', var_name='parameter')
vals = [trans_pars3[:, i].numpy() for i in range(len(labels))]
posterior_accuracy(labels, trans_pars_sgf, vals)
plt.figure()    
#plot convergence of stochasitc ELBO estimates (log-model evidence)
plt.plot(sgf_infer.loss[-400:])
g = sns.FacetGrid(trans_pars_sgf, col="parameter", height=3, sharey=False);
g = (g.map(sns.lineplot, 'subject', 'value', ci='sd'));
for i in range(len(labels)):
    g.axes[0,i].plot(np.arange(2*nsub), trans_pars3[:,i].numpy(),'ro', zorder = 0);
g = sns.PairGrid(sgf_mu_df)
g = g.map_diag(sns.kdeplot)
g = g.map_offdiag(plt.scatter)
g = sns.PairGrid(sgf_sigma_df)
g = g.map_diag(sns.kdeplot)
g = g.map_offdiag(plt.scatter)
#plt.plot(rl_infer.loss[-400:]);
plt.plot(hgf_infer.loss[-400:]);
plt.plot(sgf_infer.loss[-400:]);
Explanation: Plot posterior estimates from simulated data for the HGF agent
End of explanation
stimulus = sim1.stimulus
stimulus['mask'] = torch.ones(1, 120, 100)
rl_infer = [Horseshoe(rl_agent, stimulus, sim1.responses),
            Horseshoe(rl_agent, stimulus, sim2.responses),
            Horseshoe(rl_agent, stimulus, sim3.responses)]
evidences = torch.zeros(3, 3, 2*nsub)
for i in range(3):
    rl_infer[i].infer_posterior(iter_steps = 500)
    evidences[0, i] = rl_infer[i].get_log_evidence_per_subject()
    
hgf_infer = [Horseshoe(hgf_agent, stimulus, sim1.responses),
             Horseshoe(hgf_agent, stimulus, sim2.responses),
             Horseshoe(hgf_agent, stimulus, sim3.responses)]
for i in range(3):
    hgf_infer[i].infer_posterior(iter_steps = 500)
    evidences[1, i] = hgf_infer[i].get_log_evidence_per_subject()
    
sgf_infer = [Horseshoe(sgf_agent, stimulus, sim1.responses),
             Horseshoe(sgf_agent, stimulus, sim2.responses),
             Horseshoe(sgf_agent, stimulus, sim3.responses)]
for i in range(3):
    sgf_infer[i].infer_posterior(iter_steps = 500)
    evidences[2, i] = sgf_infer[i].get_log_evidence_per_subject()
print((evidences[:, 0].argmax(dim=0) == 0).sum().float()/(2*nsub))
print((evidences[:, 1].argmax(dim=0) == 1).sum().float()/(2*nsub))
print((evidences[:, 2].argmax(dim=0) == 2).sum().float()/(2*nsub))
evidences.sum(-1)
Explanation: Test model comparison
End of explanation
<END_TASK>