SRM in Stan

fit.R

#!/usr/bin/env Rscript

library(cmdstanr)
library(dplyr)

options(mc.cores = parallel::detectCores() - 1)

load("./day5.rda")

# Actor list
lvls <- union(trade$Var1, trade$Var2) |> sort()

# Actor level predictors
X_actor <- select(trade, Var1, pop1, gdp1, polity1) |>
    distinct() |>
    mutate(Var1 = factor(Var1, lvls),
           pop1 = scale(pop1),
           gdp1 = scale(gdp1),
           polity1 = scale(polity1)) |>
    arrange(Var1)

# Dyadic level predictors and outcomes
dyads <- combn(lvls, 2) |>
    t() |>
    data.frame() |>
    left_join(select(trade, Var1, Var2, conflicts, distance, shared_igos, y1 = trade),
              by = c("X1" = "Var1", "X2" = "Var2")) |>
    left_join(select(trade, Var1, Var2, y2 = trade),
              by = c("X1" = "Var2", "X2" = "Var1")) |>
    mutate(X1 = factor(X1, lvls) |> as.numeric(),
           X2 = factor(X2, lvls) |> as.numeric(),
           distance = scale(distance),
           shared_igos = scale(shared_igos))

###
# Fit data with stan
data <- list(n_dyads = nrow(dyads),
             n_actors = n_distinct(trade$Var1),
             K_dyad = 3,
             K_actor = 3,
             X_dyad = select(dyads, -X1, -X2, -y1, -y2) |> data.matrix(),
             X_actor = select(X_actor, -Var1) |> data.matrix(),
             actor_id = select(dyads, X1, X2) |> data.matrix(),
             y = select(dyads, y1, y2) |> data.matrix())
str(data)

mod <- cmdstan_model("./srm.stan")
fit <- mod$sample(data = data, chains = 4, max_treedepth = 12)

sim.R

#!/usr/bin/env Rscript
#
# Creates a simulated dataset and checks that we can recover the true
# parameters with Stan.
###

library(bayesplot)
library(cmdstanr)
library(dplyr)
library(extraDistr)
library(MASS, exclude = "select")

options(mc.cores = parallel::detectCores() - 1)

###
# Simulate a directed dyadic dataset
set.seed(4343)

n_actors <- 10
n_dyads <- choose(n_actors, 2)

# List of undirected dyads
dyads <- combn(1:n_actors, 2) |>
    t() |>
    data.frame() |>
    select(Actor1 = X1, Actor2 = X2)

# Dyadic level predictors
X_dyad <- data.frame(X = rnorm(n_dyads))

# Actor level predictors
X_actor <- data.frame(X = rnorm(n_actors))

# Actor level correlation matrix
Omega_actor <- matrix(c(1, 0.8, 0.8, 1), 2, 2)

# Actor level standard deviation
tau <- rexp(2, 10)

# Actor level covariance matrix
Sigma_actor <- diag(tau, 2, 2) %*% Omega_actor %*% diag(tau, 2, 2)

# Actor random effects
gamma <- mvrnorm(n_actors, c(0, 0), Sigma_actor)

# Dyadic level correlation matrix
rho <- 0.6
Omega_dyad <- matrix(c(1, rho, rho, 1), 2, 2)

# Dyadic level variance
sigma <- rexp(1)

# Dyadic level covariance matrix
Sigma_dyad <- sigma * Omega_dyad

# Correlated error terms
epsilon <- mvrnorm(n_dyads, c(0, 0), Sigma_dyad)

# Regression Parameters
alpha <- 0.5
lambda <- 4
beta <- c(-2, 1.5)

# Dyadic level regression terms
nu <- alpha + lambda * X_dyad[, 1]

mu <- matrix(NA, n_dyads, 2)

# Simulate mean for i->j relation w/ Actor1 as sender and Actor2 as receiver
mu[, 1] <- with(dyads, nu + beta[1] * X_actor[Actor1, ] + beta[2] * X_actor[Actor2, ])

# Simulate mean for j->i relation w/ Actor2 as sender and Actor1 as receiver
mu[, 2] <- with(dyads, nu + beta[1] * X_actor[Actor2, ] + beta[2] * X_actor[Actor1, ])

# Add in random effects and error terms. gamma[, 1] captures sender
# heterogeneity and gamma[, 2] receiver heterogeneity.
y <- matrix(NA, n_dyads, 2)
y[, 1] <- with(dyads, mu[, 1] + gamma[Actor1, 1] + gamma[Actor2, 2] + epsilon[, 1])
y[, 2] <- with(dyads, mu[, 2] + gamma[Actor2, 1] + gamma[Actor1, 2] + epsilon[, 2])

###
# Fit simulated data
mod <- cmdstan_model("srm.stan")
data <- list(n_dyads = n_dyads,
             n_actors = n_actors,
             K_dyad = 1,
             K_actor = 1,
             X_dyad = data.matrix(X_dyad),
             X_actor = data.matrix(X_actor),
             actor_id = data.matrix(dyads),
             y = y)
str(data)

fit <- mod$sample(data = data, chains = 4)

###
# Check if we recover the true parameter values
theta <- fit$draws(c("alpha", "beta", "lambda", "sigma", "tau", "DyadCorrMat[2,1]"),
                   format = "draws_df") |>
    rename(rho = `DyadCorrMat[2,1]`) |>
    mutate(sigma = sigma^2)

mcmc_recover_intervals(theta, c(alpha, beta, lambda, sigma, tau, rho))

###
# Plot posterior predictions for y
y_hat <- fit$draws("y_hat", format = "draws_matrix")
ppc_dens_overlay(c(y[, 1], y[, 2]), y_hat[1:100, ])

srm.stan

data {
  int n_dyads;                       // Total num unique undirected dyads
  int n_actors;                      // Total num of actors
  int K_dyad;                        // Num of dyadic level predictors
  int K_actor;                       // Num of actor level predicts

  matrix[n_dyads, K_dyad] X_dyad;    // Dyadic level predictors
  matrix[n_actors, K_actor] X_actor; // Actor level predictors

  // Actor IDs comprising each undirected dyad
  array[n_dyads, 2] int<lower=1, upper=n_actors> actor_id;

  array[n_dyads] vector[2] y;        // Two directed outcomes per dyad
}

parameters {
  real alpha;                        // Intercept
  vector[K_dyad] lambda;             // Reg. coefficients for dyadic predictors

  // Reg. coefficients for actor predictors. beta[1:K_actor] are the
  // sender coefficients, and beta[(K_actor+1):(2*K_actor)] are the
  // receiver coefficients.
  vector[2 * K_actor] beta;

  cholesky_factor_corr[2] L_dyad;
  real<lower=0, upper=pi()/2> sigma_unif;

  cholesky_factor_corr[2] L_actor;
  vector<lower=0, upper=pi()/2>[2] tau_unif;
  matrix[2, n_actors] gamma_raw;
}

transformed parameters {
  // Transformed priors
  // sigma ~ Half-Cauchy(0, 1)
  real<lower=0> sigma = tan(sigma_unif);

  // Cholesky factor for dyadic level covariance matrix
  matrix[2, 2] L_Omega = diag_pre_multiply(rep_vector(sigma, 2), L_dyad);

  // tau ~ Half-Cauchy(0, 1)
  vector<lower=0>[2] tau = tan(tau_unif);

  // gamma[actor, ] ~ multi_normal(0, CovMat)
  matrix[n_actors, 2] gamma = (diag_pre_multiply(tau, L_actor) * gamma_raw)';

  // Mean vector for likelihood
  array[n_dyads] vector[2] mu;
  {
    vector[n_dyads] nu = alpha + X_dyad * lambda;
    for (i in 1:n_dyads) {
      // Actor1 as sender and Actor2 as receiver
      mu[i, 1] = nu[i] + gamma[actor_id[i, 1], 1] + gamma[actor_id[i, 2], 2] +
        X_actor[actor_id[i, 1], ] * beta[:K_actor] +
        X_actor[actor_id[i, 2], ] * beta[(K_actor+1):];

      // Actor2 as sender and Actor1 as receiver
      mu[i, 2] = nu[i] + gamma[actor_id[i, 2], 1] + gamma[actor_id[i, 1], 2] +
        X_actor[actor_id[i, 2], ] * beta[:K_actor] +
        X_actor[actor_id[i, 1], ] * beta[(K_actor+1):];
    }
  }
}

model {
  // Priors
  target += normal_lpdf(alpha | 0, 5);
  target += normal_lpdf(lambda | 0, 2.5);
  target += normal_lpdf(beta | 0, 2.5);

  target += lkj_corr_cholesky_lpdf(L_dyad | 2);

  target += lkj_corr_cholesky_lpdf(L_actor | 2);
  target += std_normal_lpdf(to_vector(gamma_raw));

  // Likelihood
  for (i in 1:n_dyads)
    target += multi_normal_cholesky_lpdf(y[i, ] | mu[i, ], L_Omega);
}

generated quantities {
  corr_matrix[2] ActorCorrMat = multiply_lower_tri_self_transpose(L_actor);
  corr_matrix[2] DyadCorrMat = multiply_lower_tri_self_transpose(L_dyad);

  array[n_dyads] vector[2] y_hat;
  for (i in 1:n_dyads)
    y_hat[i, ] = multi_normal_cholesky_rng(mu[i, ], L_Omega);
}