## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)

## ----setup--------------------------------------------------------------------
library(networkscaleup)

## ----simulation---------------------------------------------------------------
set.seed(1998)
N_i = 50
N_k = 5
N = 1e5
sizes = rbinom(N_k, N, prob = runif(N_k, min = 0.01, max = 0.15))
degrees = round(exp(rnorm(N_i, mean = 5, sd = 1)))

ard = matrix(NA, nrow = N_i, ncol = N_k)
for(k in 1:N_k){
  ard[,k] = rbinom(N_i, degrees, sizes[k] / N)
}

## Create some artificial covariates for use later
x = matrix(sample(1:5, size = N_i * N_k, replace = T),
           nrow = N_i,
           ncol = N_k)
z = cbind(rbinom(N_i, 1, 0.3), rnorm(N_i), rnorm(N_i), rnorm(N_i))

## ----prep---------------------------------------------------------------------
x = scale(x)
z = scale(z)
z_subpop = z[,1:2]
z_global = z[,3:4]


## ----pimle--------------------------------------------------------------------
pimle.est = killworth(ard,
  known_sizes = sizes[c(1, 2, 4)],
  known_ind = c(1, 2, 4),
  N = N, model = "PIMLE")


plot(degrees ~ pimle.est$degrees, xlab = "Estimated PIMLE degrees", ylab = "True Degrees")
abline(0, 1, col = "red")

round(data.frame(true = sizes[c(3, 5)],
                 pimle = pimle.est$sizes))

## ----mle----------------------------------------------------------------------
mle.est = killworth(ard,
  known_sizes = sizes[c(1, 2, 4)],
  known_ind = c(1, 2, 4),
  N = N, model = "MLE")


plot(degrees ~ mle.est$degrees, xlab = "Estimated MLE degrees", ylab = "True Degrees")
abline(0, 1, col = "red")

round(data.frame(true = sizes[c(3, 5)],
                 pimle = mle.est$sizes))

## ----overdisp-----------------------------------------------------------------
overdisp_gibbs_metrop_est = overdispersed(
  ard,
  known_sizes = sizes[c(1, 2, 4)],
  known_ind = c(1, 2, 4),
  G1_ind = 1,
  G2_ind = 2,
  B2_ind = 4,
  N = N,
  warmup = 500,
  iter = 1000,
  verbose = TRUE,
  init = "MLE"
)

overdisp_stan = overdispersedStan(
  ard,
  known_sizes = sizes[c(1, 2, 4)],
  known_ind = c(1, 2, 4),
  G1_ind = 1,
  G2_ind = 2,
  B2_ind = 4,
  N = N,
  chains = 2,
  cores = 2,
  warmup = 250,
  iter = 500,
)


round(data.frame(true = sizes,
                 gibbs_est = colMeans(overdisp_gibbs_metrop_est$sizes),
                 stan_est = colMeans(overdisp_stan$sizes)))

plot(degrees ~ colMeans(overdisp_stan$degrees), xlab = "Overdispersed Degree Estimates", ylab = "True Degrees")
abline(0, 1, col = "red")

## ----correlated---------------------------------------------------------------
correlated_cov_stan = correlatedStan(
  ard,
  known_sizes = sizes[c(1, 2, 4)],
  known_ind = c(1, 2, 4),
  model = "correlated",
  scaling = "weighted",
  x = x,
  z_subpop = z_subpop,
  z_global = z_global,
  N = N,
  chains = 2,
  cores = 2,
  warmup = 250,
  iter = 500,
)

correlated_nocov_stan = correlatedStan(
  ard,
  known_sizes = sizes[c(1, 2, 4)],
  known_ind = c(1, 2, 4),
  model = "correlated",
  scaling = "all",
  N = N,
  chains = 2,
  cores = 2,
  warmup = 250,
  iter = 500,
)


uncorrelated_cov_stan = correlatedStan(
  ard,
  known_sizes = sizes[c(1, 2, 4)],
  known_ind = c(1, 2, 4),
  model = "uncorrelated",
  scaling = "all",
  x = x,
  z_subpop = z_subpop,
  z_global = z_global,
  N = N,
  chains = 2,
  cores = 2,
  warmup = 250,
  iter = 500,
)

uncorrelated_x_stan = correlatedStan(
  ard,
  known_sizes = sizes[c(1, 2, 4)],
  known_ind = c(1, 2, 4),
  model = "uncorrelated",
  scaling = "all",
  x = x,
  N = N,
  chains = 2,
  cores = 2,
  warmup = 250,
  iter = 500,
)



round(data.frame(true = sizes,
                 corr_cov_est = colMeans(correlated_cov_stan$sizes),
                 corr_nocov_est = colMeans(correlated_nocov_stan$sizes),
                 uncorr_cov_est = colMeans(uncorrelated_cov_stan$sizes),
                 uncorr_x_est = colMeans(uncorrelated_x_stan$sizes)))

plot(degrees ~ colMeans(correlated_cov_stan$degrees), xlab = "Correlated Covariate Degree Estimates", ylab = "True Degrees")
abline(0, 1, col = "red")

## Examine parameter estimates
colMeans(correlated_cov_stan$alpha)
colMeans(correlated_cov_stan$beta_global)
colMeans(correlated_cov_stan$beta_subpop)

