inst/scripts/d_zibb_6.R

In snaketron/IgGeneUsage: Differential gene usage in immune repertoires

require(rstan)

# Stan generative model
sim_stan <- "
functions {
  int zibb_rng(int y, int n, real mu, real phi, real kappa) {
    if (bernoulli_rng(kappa) == 1) {
      return (0);
    } else {
      return (beta_binomial_rng(n, mu * phi, (1 - mu) * phi));
    }
  }
}
data {
  int<lower=0> N_rep;
  int<lower=0> N_sample;                   // number of repertoires
  int<lower=0> N_gene;                     // gene
  int<lower=0> N_individual;               // number of individuals
  int<lower=0> N_condition;                // number of conditions
  array [N_sample] int N;                  // repertoire size
  array [N_sample] int condition_id;       // id of conditions
  array [N_sample] int individual_id;      // id of replicate
  real <lower=0> phi;
  real <lower=0, upper=1> kappa;
  vector <lower=0> [N_condition] sigma_condition;
  real <lower=0> sigma_alpha;
}
generated quantities {
  vector [N_gene] alpha;
  array [N_individual] vector [N_gene] alpha_individual;
  array [N_sample] vector [N_gene] beta_individual;
  // generate usage
  array [N_sample] vector <lower=0, upper=1> [N_gene] theta;
  array [N_gene, N_sample, N_rep] int Y;
  array [N_condition] vector [N_gene] beta_condition;
  
  for(i in 1:N_gene) {
    alpha[i] = normal_rng(-3, 0.5);
  }
  
  for(i in 1:N_condition) {
    for(j in 1:N_gene) {
      beta_condition[i][j] = normal_rng(0, sigma_condition[i]);
    }
  }
  
  for(i in 1:N_sample) {
    for(j in 1:N_gene) {
      alpha_individual[individual_id[i]][j] = normal_rng(alpha[j], sigma_alpha);
      beta_individual[i][j] = normal_rng(beta_condition[condition_id[i]][j], sigma_condition[condition_id[i]]);
      theta[i][j] = inv_logit(alpha_individual[individual_id[i]][j] + beta_individual[i][j]);
      for(r in 1:N_rep) {
        Y[j, i, r] = zibb_rng(1, N[i], theta[i][j], phi, kappa);
      }
    }
  }
}
"

# compile model
model <- rstan::stan_model(model_code = sim_stan)



# generate data based on the following parameters parameters
set.seed(1113)
N_gene <- 10
N_individual <- 9
N_condition <- 3
N_sample <- N_individual*N_condition
N_rep <- 4

condition_id <- rep(x = 1:N_condition, each = N_individual) 

N <- rep(x = 1000, times = N_sample)

individual_id <- rep(x = 1:N_individual, times = N_condition)

phi <- 200
kappa <- 0.02

sigma_condition <- c(1, 0.5, 0.25)
sigma_alpha <- 0.25


l <- list(N_rep = N_rep,
          N_sample = N_sample, 
          N_gene = N_gene, 
          N_individual = N_individual,
          N_condition = N_condition,
          N = N,
          condition_id = condition_id,
          individual_id = individual_id,
          phi = phi,
          kappa = kappa,
          sigma_condition = sigma_condition,
          sigma_alpha = sigma_alpha)

# simulate
sim <- rstan::sampling(object = model,
                       data = l, 
                       iter = 1, 
                       chains = 1, 
                       algorithm="Fixed_param")

# extract simulation and convert into data frame which can 
# be used as input of IgGeneUsage
ysim <- rstan::extract(object = sim, par = "Y")$Y
ysim <- ysim[1,,,]

ysim_df <- reshape2::melt(ysim)
colnames(ysim_df) <- c("gene_name", "sample_id", "replicate", "gene_usage_count")

m <- data.frame(sample_id = 1:l$N_sample,
                individual_id = l$individual_id,
                condition_id = l$condition_id)

ysim_df <- merge(x = ysim_df, y = m, by = "sample_id", all.x = TRUE)

ysim_df$condition <- paste0("C_", ysim_df$condition_id)
ysim_df$gene_name <- paste0("G_", ysim_df$gene_name)
ysim_df$individual_id <- paste0("I_", ysim_df$individual_id)
ysim_df$replicate <- paste0("R_", ysim_df$replicate)

ysim_df$condition_id <- NULL
ysim_df$sample_id <- NULL
ysim_df <- ysim_df[, c("individual_id", "condition", "gene_name",
                       "replicate", "gene_usage_count")]

d_zibb_6 <- ysim_df

# save
save(d_zibb_6, file = "data/d_zibb_6.RData", compress = TRUE)