utils.R

# =============== Helper functions =============== #

#'@param sample_contingency data frame generated by as.data.table(xtabs(~, sample))
#'@param sample.margins list of formulas describing sample margins, which must not contain missing values. SAME as survey::rake function, e.g. list(~age_dc, ~opmres_x) 
loading_matrix <- function(sample_contingency, 
                           sample.margins) {
  # sample.margins follow the same rule as survey::rake
  # sample.margins only take formula
  
  sample_vars = colnames(sample_contingency)
  L = c()
  
  for (margin in sample.margins) {
    vars = all.vars(margin)
    stopifnot(all(vars %in% sample_vars))
    form = as.formula(paste0("~-1 + ", paste0(vars, collapse = ":")))
    L = cbind(L, model.matrix(form, sample_contingency))
  }
  return(t(L))
}

#' @param population.margins list of tables describing corresponding population margins MUST be same order as sample.margins
margin_vector <- function(population.margins) { 
  # population.margins follow the same rule as survey::rake
  # population.margins only take table or xtabs
  
  helper <- function(margin_table) {
    tmp = as.data.frame(margin_table)
    vars = colnames(tmp[-NCOL(tmp)])
    N = NROW(tmp)
    rownames(tmp) = 1:N
    for (i in 1:N) {
      tmpname = rep(NA, length(vars))
      for (j in seq_along(vars)) {
        tmpname[j] = paste0(vars[j], tmp[i, vars[j]])
      }
      rownames(tmp)[i] = paste0(tmpname, collapse = ":")
    }
    tmp = tmp[NCOL(tmp)]
    tmp = as.matrix(tmp)
    return(tmp)
  }
  res = c()
  for (i in population.margins) {
    res = rbind(res, helper(i))
  }
  res
}

# =============== Bayes raking wrapper =============== #
#'@param sample_data samples
#'@param vars variable names used for contingency table
#'@param sample.margins list of formulas describing sample margins, which must not contain missing values. SAME as survey::rake function, e.g. list(~age_dc, ~opmres_x)
#'@param desire.margins list of formulas describing which sample margins want to evaluate in the final analysis, which must not contain missing values.
#'@param sreg formula used for inclusion probability model, will pass to model.matrix, should not present any variables outside vars
#'@param template data.table object, used for simulation purpose or if you have a fixed contingency table
#'@return list with basic sample contingency table information
bayes_sample_contingency <- function(sample_data, vars,
                          sample.margins, desire.margins,
                          sreg = NULL,
                          template = NULL) {
  require(magrittr)
  require(data.table)
  
  sample_contingency = as.data.frame(table(sample_data[, vars]))
  if (!is.null(template)) {
    stopifnot("data.table" %in% class(template))
    stopifnot(!is.null(attr(template, "sorted")))
    sample_contingency = template[sample_contingency] %>% filter(!is.na(id)) %>% arrange(id)
  }
  ncell = sample_contingency$Freq
  if (is.null(sreg)) {
    sreg = inclusion_reg_formula(sample.margins) 
  }
  L = loading_matrix(sample_contingency = sample_contingency, 
                     sample.margins = sample.margins)
  D = NROW(L)
  stopifnot(NCOL(L) == length(ncell))
  
  pdesign_J = model.matrix(sreg, sample_contingency)
  L_quant = loading_matrix(sample_contingency = sample_contingency,
                           sample.margins = desire.margins)
  
  return(list(sample_contingency = sample_contingency, 
              ncell = ncell, L = L, D = D, 
              L_quant = L_quant, D_quant = NROW(L_quant),
              pdesign_J = pdesign_J))
}

# ================== NO NEED TO READ ============================== #
inclusion_reg_formula = function(sample.margins) {
  formula_part = rep(NA, length(sample.margins))
  for (i in seq_along(sample.margins)) {
    vars = all.vars(sample.margins[[i]])
    if (length(vars) == 1) {
      formula_part[i] = vars
    } else {
      formula_part[i] = paste0(vars, collapse =':')
    }
  }
  sreg = as.formula(paste("~", paste(formula_part, collapse = "+")))
  return(sreg)
}

sampling_from_pop = function(population, inclusion_prob) {
  selected = sapply(inclusion_prob, FUN = function(x) {rbinom(1, 1, x) == 1})
  return(population[selected, ])
}

# ================= Summary Utilites ================= #

summary_wrap = function(stan_obj, variable) {
  require(rstan)
  var = c("mean", "sd", "2.5%", "97.5%")
  res = rstan::summary(stan_obj, variable)$summary
  return(res[,var])
}

summary_marginal <- function(estimation, original) {
  bias = estimation[, "mean"] - original
  SquareErr = (bias)^2
  ci = estimation[, 4] - estimation[, 3]
  
  n = length(original)
  coverage = rep(NA, n)
  for (i in 1:n) {
    coverage = 1 * (estimation[i, 3] <= original[i]) * (estimation[i, 4] >= original[i])
  }
  res = data.frame(Bias = bias, Abs.Bias = abs(bias),
                   SquareErr = SquareErr, 
                   StandardErr = estimation[, "sd"], 
                   CIlength = ci, Coverage = coverage, Est = estimation[, "mean"])
  return(res)
}

summary_overall <- function(estimation, original) {
  bias = estimation[1] - original
  SquareErr = bias^2
  ci = estimation[4] - estimation[3]
  coverage = 1 * (estimation[3] <= original) * (estimation[4] >= original)
  res = data.frame(Bias = bias, Abs.Bias = abs(bias), SquareErr = SquareErr,
                   StandardErr = estimation[2],
                   CIlength = ci, Coverage = coverage, Est = estimation[1])
  return(res)
}

summarise_list <- function(summary_list) {
  N = length(summary_list) 
  ymean = 0
  ymarginal = 0
  fit_time = 0
  for (i in 1:N) {
    ymean = ymean + summary_list[[i]][["ymean"]]
    ymarginal = ymarginal + summary_list[[i]][["ymarginal"]]
    fit_time = fit_time + summary_list[[i]][["fit_time"]]
  }
  ymean = ymean / N
  ymarginal = ymarginal / N
  ymarginal['Margin'] = rownames(ymarginal)
  ymarginal['RMSE'] = sqrt(ymarginal['SquareErr'])
  ymarginal$SquareErr = NULL
  fit_time = fit_time / N
  
  return(list(ymean = ymean, ymarginal = ymarginal, fit_time = fit_time[3]))
}

raking_mean = function(raking_design, dependent) {
  
  ymean_raking = as.data.frame(svymean(as.formula(paste0('~', dependent)), raking_design))
  ymean_raking['2.5%'] = ymean_raking[1] - 1.96 * ymean_raking[2]
  ymean_raking['97.5%'] = ymean_raking[1] + 1.96 * ymean_raking[2]
  
  class(ymean_raking) = "numeric"
  names(ymean_raking) = c("mean", "sd", "2.5%", "97.5%")
  return(ymean_raking)
}

raking_marginal = function(raking_design, dependent, subgroup) {
  ymarginal_raking = data.frame()
  for (i in subgroup) {
    vars = all.vars(i)
    test = svyby(as.formula(paste0('~', dependent)), i, raking_design, svymean)
    N = NROW(test)
    tmpname = rep(NA, N)
    for (j in 1:N) {
      varname = rep(N, length(vars))
      for (k in 1:length(vars)) varname[k] = paste0(vars[k], test[j, vars[k]])
      tmpname[j] = paste(varname, collapse = ":")
    }
    rownames(test) = tmpname
    test[vars] = NULL
    ymarginal_raking = rbind(ymarginal_raking, test)
  }
  
  ymarginal_raking['2.5%'] = ymarginal_raking[, DEPENDENT] - 1.96 * ymarginal_raking[, "se"]
  ymarginal_raking['97.5%'] = ymarginal_raking[, DEPENDENT] + 1.96 * ymarginal_raking[, "se"]
  colnames(ymarginal_raking) = c("mean", "sd", "2.5%", "97.5%")
  return(ymarginal_raking)
}


# ================= Quantities Utilites ================= #


logit <- function(x) {
  return(log(x/(1 - x)))
}

log_inv <- function(x) {
  exp(x)/(1 + exp(x))
}

# ================ Print Utilies ============== #

print_marign = function(list_margin_formula, list_margin_distribution) {
  require(knitr)
  for (i in seq_along(list_margin_formula)) {
    vars = all.vars(list_margin_formula[[i]])
    cat("Variables: ", vars, "\n")
    if (length(vars) == 1) {
      cat("Marginal distribution: ", list_margin_distribution[[i]], "\n")
    } else {
      cat("Joint marginal distribution: \n")
      print(list_margin_distribution[[i]])
    }
    cat("\n")
  }
}

print_summary = function(summary_stat) {
  cat("Estimation: ", summary_stat[,"Est"], "\n",
      "Bias:", summary_stat[,"Bias"], "\n",
      "Average Absolute Bias: ", summary_stat[,"Abs.Bias"], "\n",
      "RMSE: ", sqrt(summary_stat[,"SquareErr"]), "\n",
      "Standard Error: ", summary_stat[,"StandardErr"], "\n", 
      "95% CI length: ", summary_stat[,"CIlength"], "\n")
}

plot_prepare = function(marginal_summary) {
  marginal_summary['Margin'] = rownames(marginal_summary)
  marginal_summary["RMSE"] = sqrt(marginal_summary[, 'SquareErr'])
  return(marginal_summary)
}


# =============== Simulation wrapper ================= #

stan_wrapper = function(stan_model, data_list, chains, iter, seed, ymean_true, ymarginal_true) {
  ptm = proc.time()
  stan_fit = rstan::sampling(stan_model, data = data_list, chains = chains, 
                        iter = iter, seed = seed, open_progress = FALSE, 
                        show_messages = FALSE)
  fit_time = proc.time() - ptm
  
  ymean_fit = summary_overall(summary_wrap(stan_fit, "ymean"), ymean_true)
  ymarginal_fit = summary_marginal(summary_wrap(stan_fit, "ymarginals"), ymarginal_true)
  rownames(ymarginal_fit) = rownames(data_list$L_quant)
  
  return(list(ymean = ymean_fit, ymarginal = ymarginal_fit, fit_time = fit_time))
}

rake_wrapper = function(sample_data, sample.margins, population.margins, 
                        dependent, subgroup, ymean_true, ymarginal_true, ...) {
  ptm = proc.time()
  design = svydesign(id = ~0, probs = NULL, data = sample_data)
  rake_design = rake(design, sample.margins = sample.margins, 
               population.margins = population.margins, ...)
  fit_time = proc.time() - ptm
  ymean_fit = summary_overall(
    raking_mean(rake_design, dependent), 
    ymean_true)
  
  ymarginal_fit = summary_marginal(
    raking_marginal(rake_design, dependent, subgroup), 
    ymarginal_true)
  return(list(ymean = ymean_fit, ymarginal = ymarginal_fit, fit_time = fit_time))
}

# ============== DEBRECATE =========== #
summary_tmp <- function(estimation, original, Method) {
  bias = estimation[, 1] - original
  SquareErr = (bias)^2
  ci = estimation[, 4] - estimation[, 3]
  
  n = length(original)
  coverage = rep(NA, n)
  for (i in 1:n) {
    coverage = 1 * (estimation[i, 3] <= original[i]) * (estimation[i, 4] >= original[i])
  }
  res = data.frame(Bias = bias, 
                   SquareErr = SquareErr, 
                   StandardErr = estimation[, 2], 
                   CIlength = ci, Coverage = coverage, Margin = rownames(estimation), Method = Method)
  return(res)
}