vignettes/stan_model_for_counterfactual_prediction.R

In robertness/bninfo: Queries and information theoritic operations on Bayesian networks.

# Inference on the thing

library(bninfo)
library(tidyverse)
library(rstan)

data(isachs)
.data <- isachs
names(.data) <- c("praf", 'pmek', 'plcg', 'PIP2', 'PIP3', 'perk', 'pakts473', 'PKA', 'PKC', 'P38', 'pjnk', 'INT')
proteins <- setdiff(names(.data), c('plcg', 'PIP2', 'PIP3', 'pakts473', 'INT'))
for(prot in proteins){
  .data[[prot]] <- ifelse(.data[[prot]] == 1, 0, 1)
}
training_raw = .data[(.data$INT == 0) | (.data$INT == 9), ]
test_raw = .data[.data$INT == 2, ]


training_pkc_off = training_raw[training_raw$PKC == 0, ]
training_pkc_on = training_raw[training_raw$PKC == 1, ]
training_pkc_off_upsampled = sample_n(training_pkc_off,  nrow(training_pkc_on), replace=TRUE)
training = rbind(training_pkc_on, training_pkc_off_upsampled)

test_pkc_off = test_raw[test_raw$PKC == 0, ]
test_pkc_on = test_raw[test_raw$PKC == 1, ]
test_pkc_on_upsampled = sample_n(test_pkc_on, nrow(test_pkc_off), replace=TRUE)
test = rbind(test_pkc_off, test_pkc_on_upsampled)

evidence_str <- "
functions {
  real h(real N, real p){
    if(N > p){
      return 1.0;
    }else{
      return 0.0;
    }
  }
  real f_PKC(real N){
    return h(N, .5804);
  }
  real f_PKA(real PKC, real N){
    real p;
    p = PKC == 1.0 ? .9521 : .5423;
    return h(N, p);
  }
  real f_pjnk(real PKC, real PKA, real N){
    real p;
    if(PKC == 1.0){
      p = .46;
    }else{
      p = PKA == 1.0 ? .2696 : .7155;
    }
    return h(N, p);
  }
  real f_P38(real PKC, real PKA, real N){
    real p;
    if(PKC == 1.0){
      p = PKA == 1.0 ? .1946 : .3263;
    }else{
      p = PKA == 1.0 ? .1245 : .7025;
    }
    return h(N, p);
  }
  real f_praf(real PKC, real PKA, real N){
    real p;
    if(PKA == 1.0){
      p = .39;
    }else{
      p = PKC == 1.0 ? .6180: .9379;
    }
    return h(N, p);
  }
  real f_pmek(real PKC, real PKA, real praf, real N){
    real p;
    if(PKC == 1.0 && praf == 1.0){
      p = PKA == 1.0 ? .6822 : .7848;
    }else if(PKC == 1.0 && praf == 0.0){
      p = .2342;
    }else if(PKC == 0.0 && praf == 1.0){
      p = PKA == 1.0 ? .4311 : .8869;
    } else if(PKC == 0.0 && praf == 0.0){
      p = PKA == 1.0 ? .1030 : .3750;
    }
    return h(N, p);
  }
  real f_perk(real PKA, real pmek, real N){
    real p;
    if(pmek == 1.0){
      p = .95;
    } else {
      p = PKA == 1.0 ? .8909 : .1565;
    }
    return h(N, p);
  }
}
data {
  int N; // number of observations
  real<lower=0, upper=1> PKC[N];         // PKC values
  real<lower=0, upper=1> pmek[N];
}
parameters {
  real<lower=0.0, upper=1.0> N_PKC[N];
  real<lower=0.0, upper=1.0> N_PKA[N];
  real<lower=0.0, upper=1.0> N_pjnk[N];
  real<lower=0.0, upper=1.0> N_P38[N];
  real<lower=0.0, upper=1.0> N_praf[N];
  real<lower=0.0, upper=1.0> N_pmek[N];
  real<lower=0.0, upper=1.0> N_perk[N];
}
transformed parameters {
  real PKC_mu[N];
  real PKA_mu[N];
  real pjnk_mu[N];
  real P38_mu[N];
  real praf_mu[N];
  real pmek_mu[N];
  real perk_mu[N];
  for(i in 1:N){
    PKC_mu[i] = f_PKC(N_PKC[i]);
    PKA_mu[i] = f_PKA(PKC_mu[i], N_PKA[i]);
    pjnk_mu[i] = f_pjnk(PKC_mu[i], PKA_mu[i], N_pjnk[i]);
    P38_mu[i] = f_P38(PKC_mu[i], PKA_mu[i], N_P38[i]);
    praf_mu[i] = f_praf(PKC_mu[i], PKA_mu[i], N_praf[i]);
    pmek_mu[i] = f_pmek(PKC_mu[i], PKA_mu[i], praf_mu[i], N_pmek[i]);
    perk_mu[i] = f_perk(PKA_mu[i], pmek_mu[i], N_perk[i]);
  }
}
model {
  target += uniform_lpdf(N_PKC | 0, 1);
  target += normal_lpdf(PKC | PKC_mu, 1e-6);
  target += normal_lpdf(pmek | pmek_mu, 1e-6);
}
"

evidence_model <- stan_model(model_code = evidence_str)
out_test <- optimizing(evidence_model, list(N=nrow(training), PKC=training$PKC, pmek = training$pmek))

noise_posterior <- sampling(
  evidence_model,
  data = list(N=nrow(training), PKC=training$PKC, pmek = training$pmek),
  cores = parallel::detectCores(),
  iter = 2000,
  pars = c('N_PKC', 'N_PKA', 'N_pjnk', 'N_P38', 'N_praf', 'N_pmek', 'N_perk')
)

samples <- rstan::extract(noise_posterior, permuted = T)

as_tibble(samples$N_PKC) %>%
  map_df(summary)

plot(density(samples$N_PKC[, 1]), main = 'PKC = 1')
abline(v = .5804)