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.
In [1]:
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')])
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.
In [2]:
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))
plot performance of different agents in different blocks
In [5]:
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.]);
Fit simulated behavior
In [ ]:
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)
In [10]:
sim1.responses.dtype
Out[10]:
Compute fit quality and plot posterior estimates from a hierarchical parameteric model
In [ ]:
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)
In [13]:
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);
fit HGF agent to simulated data
In [ ]:
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)
In [ ]:
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)
Plot posterior estimates from simulated data for the HGF agent
In [22]:
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);
In [ ]:
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)
In [24]:
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)
In [26]:
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);
In [27]:
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)
In [29]:
#plt.plot(rl_infer.loss[-400:]);
plt.plot(hgf_infer.loss[-400:]);
plt.plot(sgf_infer.loss[-400:]);
Test model comparison
In [ ]:
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()
In [33]:
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)
Out[33]:
The diagonal elements in the above matrix are not always the lowest values for the corresponding column, which shows that we cannot accuretly infer the correct model over population, and probably not per subject. More detailed analysis of the possible parameteric models is required.
In [ ]: