biopsych_notebook.Rmd

---
title: "biopsych_special_issue"
author: "Nathaniel Haines"
date: "4/24/2022"
output: html_document
---

```{r setup, include=FALSE}
# General rendering options
options(bitmapType = "cairo", mc.cores = 4)
knitr::opts_chunk$set(echo = TRUE, warnings = FALSE, message = FALSE)

# Functions to automatically print nice tables
knit_print.data.frame <- function(x, ...) {
  res <- paste(c("", "", knitr::kable(x)), collapse = "\n")
  knitr::asis_output(res)
}
registerS3method(
  "knit_print", "data.frame", knit_print.data.frame,
  envir = asNamespace("knitr")
)
knit_print.matrix <- function(x, ...) {
  res <- paste(c("", "", knitr::kable(x)), collapse = "\n")
  knitr::asis_output(res)
}
registerS3method(
  "knit_print", "matrix", knit_print.matrix,
  envir = asNamespace("knitr")
)

# register cmdstanr for later
cmdstanr::register_knitr_engine(override = FALSE)
```

```{r, warning=FALSE, message=FALSE}
# Necessary libraries
library(dplyr)
library(foreach)
library(stringr)
library(cmdstanr)
library(ggplot2)
library(bayesplot)
library(posterior)
library(hBayesDM)
library(patchwork)
library(tidyr)

# Color schema
reds <- c("#B97C7C", "#8F2727")
blues <- c("#6497b1", "#03396c")
greens <- c("#66cc99", "#2d8659")

# some convenience functions for working with cmdstanr
get_draws <- function(fit, par) {
  draws_of(as_draws_rvars(fit$draws(par))[[par]])
}

get_ordered_draws <- function(fit, par, ordering) {
  ordered_draws <- as.array(fit$draws(par))[,,ordering]
  
  order_factor <- sprintf("%s_%03d", par, 1:length(ordering))
    
  dimnames(ordered_draws)$variable <- order_factor
  
  return(ordered_draws)
}
```

## The Behavioral Model

### Learning and Choice Rules
```{r}
# Delta learning rule (outcome = R_[i,t]; E = PE_[i,t] in paper)
learning_rule <- function(outcome, V, alpha) {
  E <- outcome - V
  V <- V + alpha * E
  return(c("E" = E, "V" = V))
}

# Softmax function that can cope with numerical underflow
# Mathematically equivalent to exp(V)/sum(exp(V))
choice_rule <- function(V) {
  max_V <- max(V)
  log_sum_exp <- max_V + log(sum(exp(V - max_V)))
  PR_choice <- c(exp((V - log_sum_exp)))
  names(PR_choice) <- c("PR_1", "PR_2")
  return(PR_choice)
}
```

### Simulating Behavioral Data

```{r}
set.seed(43215)

# Task parameters
n_trials <- 100
task_mu <- c(0,0)
task_mu_base <- c(0,0)
task_sd <- c(.3,.3) 
task_drift_sd <- .1 
task_gamma <- .99 

# Model parameters
V <- c(0,0)
alpha <- .2
xi <- 1

sim_behav <- foreach(t=1:n_trials, .combine = "rbind") %do% {
  # choice = y_[i,t] in paper
  PR_choice <- choice_rule(V * xi)
  choice <- sample(c(1,2), 1, prob=PR_choice)
  
  outcome <- rnorm(1, task_mu[choice], task_sd[choice])
  
  learn <- learning_rule(outcome, V[choice], alpha)
  V[choice] <- learn["V"]
  
  task_mu_prev <- task_mu
  task_mu <- task_mu * task_gamma + rnorm(2, 0, task_drift_sd)
  
  data.frame(trial = t, 
             choice = choice, 
             outcome = outcome, 
             mu_1 = task_mu_prev[1],
             mu_2 = task_mu_prev[2],
             V_1 = V[1],
             V_2 = V[2],
             V = V[choice]-V[3-choice],
             PR_1 = PR_choice[1],
             PR_2 = PR_choice[2],
             E = learn["E"])
}

choice_value <- sim_behav %>%
  ggplot(aes(x = trial)) +
  geom_line(aes(y = mu_1), color = reds[2], alpha = .3, size=2) +
  geom_line(aes(y = mu_2), color = "darkgray", alpha = .3, size=2) +
  geom_line(aes(y = V_1), color = reds[2]) +
  geom_line(aes(y = V_2), color = "darkgray") +
  geom_point(aes(x = trial - 1, y = lag(V_1), alpha = ifelse(choice==1 & E > 0, 1, 0)), color = reds[2]) +
  geom_point(aes(x = trial - 1, y = lag(V_2), alpha = ifelse(choice==2 & E > 0, 1, 0)), color = "darkgray") +
  scale_x_continuous(breaks = seq(0,100,10)) +
  scale_y_continuous(breaks = seq(-1.5,1,.5), limits = c(-1.5,1)) +
  scale_alpha_continuous(range = c(0,1)) +
  geom_text(x = 10, y = .4, label="Option 1", color = reds[2]) +
  geom_text(x = 60, y = .55, label="Option 2", color = "darkgray") +
  xlab("Trial") +
  ylab("Expected Value (V)") +
  theme_minimal(base_size = 15) +
  theme(panel.grid = element_blank(),
        legend.position = "none")

choice_prob <- sim_behav %>%
  ggplot(aes(x = trial)) +
  geom_hline(yintercept = .5, linetype = 2, color = "gray") +
  geom_line(aes(y = PR_1), color = reds[2]) +
  geom_line(aes(y = PR_2), color = "darkgray") +
  geom_text(x = 48, y = .78, label="Option 1", color = reds[2]) +
  geom_text(x = 58, y = .32, label="Option 2", color = "darkgray") +
  scale_x_continuous(breaks = seq(0,100,10)) +
  scale_y_continuous(breaks = seq(0,1,.1)) +
  coord_cartesian(ylim=c(0,1)) +
  xlab("Trial") +
  ylab("Choice Probability") +
  theme_minimal(base_size = 15) +
  theme(panel.grid = element_blank())

choice_value + choice_prob
```

## The Neural Model

### The Hemodynamic Response Function
```{r}
double_gamma <- function(t, onset) {
  # Static parameters per SPM 12
  alpha_1 <- 6
  alpha_2 <- 16
  beta_1 <- beta_2 <- 1
  c_1 <- 1/6
  
  # temporal shift
  t_pr <- t - onset
  
  # Gamma terms 
  term_1 <- t_pr^(alpha_1-1)*(beta_1^alpha_1)*exp(-beta_1*t_pr) /
    gamma(alpha_1)
  term_2 <- t_pr^(alpha_2-1)*(beta_2^alpha_2)*exp(-beta_2*t_pr) / 
    gamma(alpha_2)
  
  dens <-  (term_1 - c_1*term_2)
  return(ifelse(dens == Inf | is.na(dens), 0, dens))
}
```

### Convolution

#### Create an example neural design matrix
```{r}
# examples TRs in seconds
ts_measure <- seq(0, 101, by = 2)
# example stimulus oinsets in seconds
ts_onset <- c(2, 7, 11, 16, 20, 27, 34, 50, 60, 65)
# example beta weights (i.e. BOLD amplitude)
betas <- rnorm(ts_onset, 1, .2)

# design matrix with shifted (but not yet scaled) HRF for each onset
# as a separate row
design_matrix <- foreach(i=seq_along(ts_onset), .combine = "rbind") %do% {
  data.frame(onset = ts_onset[i],
             time = ts_measure,
             value = betas[i]*double_gamma(ts_measure, ts_onset[i])*ifelse(ts_measure - ts_onset[i] < 0, 0, 1)
  )
}

# with beta weights/amplitude
design_matrix %>%
  group_by(time) %>%
  mutate(total = sum(value)) %>%
  ggplot(aes(time, value, group = as.factor(onset))) +
  geom_line(linetype = 2, color = reds[1]) +
  geom_point(aes(x = onset, y = 0), color = reds[2]) +
  geom_line(aes(y = total), size = .8, color = "darkgray") +
  scale_x_continuous(breaks = seq(0,100,10)) +
  scale_y_continuous(breaks = seq(0,.3,.05)) +
  xlab("Time (seconds)") +
  ylab("BOLD Signal") +
  theme_minimal(base_size = 15) +
  theme(panel.grid = element_blank())
```

## The Joint Behavior-Brain Model

### Task Design

```{r}
# Generate hypothetical jittered onset times (sec) for prediction errors (i.e. outcome presentation timestamp in the task)
generate_onsets <- function(n_trials) {
  onsets <- vector(length = n_trials)
  onsets[1] <- 10
  for (t in 2:n_trials) {
    onsets[t] <- onsets[t-1] + rlnorm(1, 2, .2)
  }
  return(onsets)
}
onsets <- generate_onsets(n_trials)
```

### Convolving with Behavioral Outcome Presentation Onsets for Design Matrix

```{r}
# TRs in seconds
ts_measure <- seq(0, 780, by = 2)
# Onset of outcome presentation (when prediction error occurs) in seconds
ts_onset <- onsets

# Shifting HRFs to correct position
design_matrix <- foreach(i=ts_onset, .combine = "rbind") %do% {
  double_gamma(ts_measure, i)*ifelse(ts_measure - i < 0, 0, 1)
}
```

### Behavior-brain Correlation Recovery Simulations

#### Generative Model

##### Stan Model

```{cmdstan output.var="generative_model"}
data {
  int N; 
  int N_trials;
  int N_record;
  array[N,N_trials] int y_behav;
  array[N,N_trials] real outcome; 
  array[N] vector[N_record] y_neural;
  array[N] matrix[N_trials,N_record] X; 
}

parameters {
  // Behavioral model parameters
  vector[N] alpha_pr;
  
  // Neural model parameters
  vector<lower=0>[N] sigma_neural;
  vector<lower=0>[N] sigma_beta;
  vector[N] beta_E;
  array[N] row_vector[N_trials] beta_pr;
}

transformed parameters {
  vector[N] alpha;
  
  for (i in 1:N) {
    alpha[i] = Phi_approx(alpha_pr[i]);
  }
}

model {
  array[N] row_vector[N_trials] beta;
  
  // Behavioral model
  alpha_pr ~ normal(-1,.5);
  
  // Neural model
  sigma_neural ~ normal(.5,1);
  sigma_beta ~ normal(1,1);
  beta_E ~ normal(0,2);
  
  vector[2] V;
  array[N] row_vector[N_trials] E_vector;
  real E;
  
  for (i in 1:N) {
    // Behavioral model
    V = rep_vector(0, 2);
    for (t in 1:N_trials) {
      y_behav[i,t] ~ categorical_logit(V);
      E = outcome[i,t] - V[y_behav[i,t]];
      E_vector[i,t] = E;
      //E_vector[i,t] = V[y_behav[i,t]]-V[3-y_behav[i,t]];
      V[y_behav[i,t]] = V[y_behav[i,t]] + alpha[i] * E;
    }
    
    // non-centered parameterization (same as beta ~ normal(E_vector[i] * beta_E[i], sigma_beta[i]))
    beta_pr[i] ~ normal(0,1);
    beta[i] = (E_vector[i] * beta_E[i]) + sigma_beta[i] * beta_pr[i];
    
    // Neural model
    y_neural[i] ~ normal(beta[i] * X[i], sigma_neural[i]);
  }
}

generated quantities {
  vector[N] R;
  array[N] row_vector[N_trials] beta;
  array[N] row_vector[N_trials] E_vector;
  
  {
    vector[2] V;
    real E;
    
    for (i in 1:N) {
      // Behavioral model
      V = rep_vector(0, 2);
      for (t in 1:N_trials) {
        E = outcome[i,t] - V[y_behav[i,t]];
        E_vector[i,t] = E;
        //E_vector[i,t] = V[y_behav[i,t]]-V[3-y_behav[i,t]];
        V[y_behav[i,t]] = V[y_behav[i,t]] + alpha[i] * E;
      }
      beta[i] = (E_vector[i] * beta_E[i]) + sigma_beta[i] * beta_pr[i];
    }
  }
  
  for (i in 1:N) {
    // convert regression weight to correlation
    R[i] = beta_E[i] / (sd(beta[i])/sd(E_vector[i]));
  }
}

```

##### Simulation Procedure

```{r}
# static model parameters
sigma_beta <- 0 # 0 implies true brain-behavior r = 1
beta_E <- 1
betas <- rnorm(n_trials, sim_behav$E*beta_E, sigma_beta)

# grid of sigma_neural parameters to simulate over
sigma_neural <- seq(0, 1, length.out = 15)

# number of simulations to average over for smoother intervals
# increase to 150 to match paper, although this takes a long time!
n_iters <- 1

# stan data list
stan_dat <- list(
  N = 1,
  N_trials = dim(design_matrix)[1],
  N_record = dim(design_matrix)[2],
  X = array(design_matrix, dim = c(1, dim(design_matrix)[1], dim(design_matrix)[2])),
  y_behav = matrix(sim_behav$choice, nrow = 1),
  outcome = matrix(sim_behav$outcome, nrow = 1)
)

sim_error_generative <- foreach(i=1:n_iters, .combine = "rbind") %do% {
  foreach(sig=sigma_neural, .combine = "rbind") %do% {
    # y_neural depends on current value for sigma_neural
    betas <- rnorm(n_trials, sim_behav$E*beta_E, sigma_beta)
    stan_dat$y_neural <- matrix(
      rnorm(length(ts_measure), betas%*%design_matrix, sig),
      nrow = 1
    )

    fit <- generative_model$sample(
      data = stan_dat,
      iter_sampling = 200,
      iter_warmup = 100,
      chains = 8,
      parallel_chains = 8,
      refresh = 0,
    )

    R <- get_draws(fit, "R")[,1]

    data.frame(
      iter = i,
      sigma_neural = sig,
      R_mean = mean(R),
      R_lower = HDIofMCMC(R)[1],
      R_upper = HDIofMCMC(R)[2]
    )
  }
}

error_plot_generative <- sim_error_generative %>%
  group_by(sigma_neural) %>%
  summarize_all(mean) %>%
  ggplot(aes(x = sigma_neural, y = R_mean)) +
  geom_line(color = "#B97C7C") +
  geom_ribbon(aes(ymin = R_lower, ymax = R_upper), alpha = .2, fill = reds[2]) +
  geom_hline(yintercept = 1, linetype = 2) +
  scale_x_continuous(breaks = seq(0, 1, .1)) +
  scale_y_continuous(breaks = seq(-1, 1, .1)) +
  xlab(expression(sigma["neural"])) +
  ylab("Brain-Behavior Correlation") +
  theme_minimal(base_size = 15) +
  theme(panel.grid = element_blank(),
        legend.position = "none")
error_plot_generative
```

#### Post-hoc Model

##### Stan Model

```{cmdstan output.var="posthoc_model"}
data {
  int N; 
  int N_trials;
  int N_record;
  array[N,N_trials] int y_behav;
  array[N,N_trials] real outcome; 
  array[N] vector[N_record] y_neural;
  array[N] matrix[N_trials,N_record] X; 
}

parameters {
  // Behavioral model parameters
  vector[N] alpha_pr; 
  
  // Neural model parameters
  vector[N] mu_beta;
  vector<lower=0>[N] sigma_neural;
  vector<lower=0>[N] sigma_beta;
  array[N] row_vector[N_trials] beta_pr;
}

transformed parameters {
  vector[N] alpha;
  
  for (i in 1:N) {
    alpha[i] = Phi_approx(alpha_pr[i]);
  }
}

model {
  array[N] row_vector[N_trials] beta;
  
  // Behavioral model
  alpha_pr ~ normal(-1,.5);
  
  // Neural model
  sigma_neural ~ normal(.5,1);
  sigma_beta ~ normal(1,1);
  mu_beta ~ normal(0,2);
  
  vector[2] V;
  array[N] row_vector[N_trials] E_vector;
  real E;
  
  for (i in 1:N) {
    beta_pr[i] ~ normal(0,1);
    beta[i] = mu_beta[i] + sigma_beta[i] * beta_pr[i];
    
    // Behavioral model
    V = rep_vector(0, 2);
    for (t in 1:N_trials) {
      y_behav[i,t] ~ categorical_logit(V);
      E = outcome[i,t] - V[y_behav[i,t]];
      E_vector[i,t] = E;
      //E_vector[i,t] = V[y_behav[i,t]]-V[3-y_behav[i,t]];
      V[y_behav[i,t]] = V[y_behav[i,t]] + alpha[i] * E;
    }
    
    // Neural model
    y_neural[i] ~ normal(beta[i] * X[i], sigma_neural[i]);
  }
}

generated quantities {
  array[N] row_vector[N_trials] E_vector;
  array[N] row_vector[N_trials] beta;
  
  {
    vector[2] V;
    real E;
    
    for (i in 1:N) {
      beta[i] = mu_beta[i] + sigma_beta[i] * beta_pr[i];
      // Behavioral model
      V = rep_vector(0, 2);
      for (t in 1:N_trials) {
        E = outcome[i,t] - V[y_behav[i,t]];
        E_vector[i,t] = E;
        //E_vector[i,t] = V[y_behav[i,t]]-V[3-y_behav[i,t]];
        V[y_behav[i,t]] = V[y_behav[i,t]] + alpha[i] * E;
      }
    }
  }
}

```

##### Simulation Procedure

```{r}
# static model parameters and stan data already set above

# number of simulations (needs fewer than generative model for stability)
# increase to 30 to match paper
n_iters <- 1

sim_error_posthoc <- foreach(i=1:n_iters, .combine = "rbind") %do% {
  foreach(sig=sigma_neural, .combine = "rbind") %do% {
    betas <- rnorm(n_trials, sim_behav$E*beta_E, sigma_beta)
    stan_dat$y_neural <- matrix(
      rnorm(length(ts_measure), betas%*%design_matrix, sig), 
      nrow = 1
    )
    
    fit <- posthoc_model$sample(
      data = stan_dat,
      iter_sampling = 200,
      iter_warmup = 100,
      chains = 8,
      parallel_chains = 8,
      refresh = 0,
    )
    
    beta <- apply(get_draws(fit, "beta"), c(2,3), mean)[1,]
    E <- apply(get_draws(fit, "E_vector"), c(2,3), mean)[1,]
    R <- cor.test(beta, E)
    
    data.frame(
      iter = i, 
      sigma_neural = sig,
      R_mean = R$estimate,
      R_lower = R$conf.int[1],
      R_upper = R$conf.int[2]
    )
  }
}

error_plot_posthoc <- sim_error_posthoc %>%
  group_by(sigma_neural) %>%
  summarize_all(mean) %>%
  ggplot(aes(x = sigma_neural, y = R_mean)) +
  geom_line(color = "darkgray") +
  geom_ribbon(aes(ymin = R_lower, ymax = R_upper), alpha = .2, fill = "gray") +
  geom_hline(yintercept = 1, linetype = 2) +
  scale_x_continuous(breaks = seq(0, 1, .1)) +
  scale_y_continuous(breaks = seq(-1, 1, .1)) +
  xlab(expression(sigma["neural"])) +
  ylab("Brain-Behavior Correlation") +
  theme_minimal(base_size = 15) +
  theme(panel.grid = element_blank(),
        legend.position = "none")
error_plot_posthoc
```

#### Sequential Model

##### Stan Model

```{cmdstan output.var="sequential_model"}
data {
  int N; 
  int N_trials;
  int N_record;
  array[N] row_vector[N_trials] E_hat;
  array[N] vector[N_record] y_neural;
  array[N] matrix[N_trials,N_record] X; 
}

parameters {
  // Neural model parameters
  vector[N] beta_E;
  vector<lower=0>[N] sigma_neural;
  vector<lower=0>[N] sigma_beta;
  array[N] row_vector[N_trials] beta_pr;
}

model {
  array[N] row_vector[N_trials] beta;
  
  // Neural model
  sigma_neural ~ normal(.5,1);
  sigma_beta ~ normal(1,1);
  beta_E ~ normal(0,2);
  
  for (i in 1:N) {
    // point-estimates of prediction errors as input (non-centered parameterization)
    beta_pr[i] ~ normal(0,1);
    beta[i] = (beta_E[i]*E_hat[i]) + sigma_beta[i] * beta_pr[i];
    
    // Neural model
    y_neural[i] ~ normal(beta[i] * X[i], sigma_neural[i]);
  }
}

generated quantities {
  vector[N] R;
  array[N] row_vector[N_trials] beta;
  
  for (i in 1:N) {
    beta[i] = (beta_E[i]*E_hat[i]) + sigma_beta[i] * beta_pr[i];
    
    // convert regression weight to correlation
    R[i] = beta_E[i] / (sd(beta[i])/sd(E_hat[i]));
  }
}

```

##### Simulation Procedure

```{r}
# static model parameters already set above

# E_hat as estimated from the behavior-only model (in post-hoc above)
stan_dat$E_hat <- matrix(E, nrow = 1)

# number of simulations (needs fewer than generative model for stability)
# increase to 30 to match paper
n_iters <- 1

sim_error_sequential <- foreach(i=1:n_iters, .combine = "rbind") %do% {
  foreach(sig=sigma_neural, .combine = "rbind") %do% {
    betas <- rnorm(n_trials, sim_behav$E*beta_E, sigma_beta)
    stan_dat$y_neural <- matrix(
      rnorm(length(ts_measure), betas%*%design_matrix, sig), 
      nrow = 1
    )
    
    fit <- sequential_model$sample(
      data = stan_dat,
      iter_sampling = 200,
      iter_warmup = 100,
      chains = 8,
      parallel_chains = 8,
      refresh = 0,
    )
    
    R <- get_draws(fit, "R")[,1]
    
    data.frame(
      iter = i, 
      sigma_neural = sig,
      R_mean = mean(R),
      R_lower = HDIofMCMC(R)[1],
      R_upper = HDIofMCMC(R)[2]
    )
  }
}

error_plot_sequential <- sim_error_sequential %>% 
  group_by(sigma_neural) %>%
  summarize_all(mean) %>%
  ggplot(aes(x = sigma_neural, y = R_mean)) + 
  geom_line(color = "darkgray") +
  geom_ribbon(aes(ymin = R_lower, ymax = R_upper), alpha = .2, fill = "gray") + 
  geom_hline(yintercept = 1, linetype = 2) +
  scale_x_continuous(breaks = seq(0, 1, .1)) +
  scale_y_continuous(breaks = seq(-1, 1, .1)) +
  xlab(expression(sigma["neural"])) +
  ylab("Brain-Behavior Correlation") +
  theme_minimal(base_size = 15) +
  theme(panel.grid = element_blank(),
        legend.position = "none")
error_plot_sequential
```


#### Plot Models Together

```{r}
summary_posthoc <- sim_error_posthoc %>% 
  group_by(sigma_neural) %>%
  summarize_all(mean)

summary_sequential <- sim_error_sequential %>% 
  group_by(sigma_neural) %>%
  summarize_all(mean)

error_plot_generative + 
  geom_line(aes(x = sigma_neural, y = R_mean), data = summary_posthoc, color = "darkgray") +
  geom_ribbon(aes(ymin = R_lower, ymax = R_upper), data = summary_posthoc, alpha = .5, fill = "gray") +
  geom_line(aes(x = sigma_neural, y = R_mean), data = summary_sequential, color = blues[1]) +
  geom_ribbon(aes(ymin = R_lower, ymax = R_upper), data = summary_sequential, alpha = .4, fill = blues[2]) +
  geom_text(x = .35, y = -.1, label="Generative", color = reds[1], size = 5) +
  geom_text(x = .55, y = -.32, label="Sequential", color = blues[1], size = 5) +
  geom_text(x = .15, y = .1, label="Post-hoc", color = "darkgray", size = 5)
```