Evidence with NAs for imputation

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: arf
 Title: Adversarial Random Forests
-Version: 0.2.0
+Version: 0.2.1
 Date: 2024-01-24
 Authors@R: 
     c(person(given = "Marvin N.",

R/forge.R

@@ -83,22 +83,25 @@ forge <- function(
 
   # Prep evidence
   conj <- FALSE
-  to_sim <- params$meta$variable
   if (!is.null(evidence)) {
     evidence <- prep_evi(params, evidence)
     if (!all(c('f_idx', 'wt') %in% colnames(evidence))) {
       conj <- TRUE
-      to_sim <- setdiff(params$meta$variable, evidence[relation == '==', variable])
+      to_sim <- evidence[relation == '==', setdiff(params$meta$variable, unique(variable)), by = row_idx]
+      setnames(to_sim, 'V1', 'variable')
     }
+  } else {
+    to_sim <- data.table(row_idx = 1, variable = params$meta$variable)
   }
-  factor_cols <- params$meta[variable %in% to_sim, family == 'multinom']
+  factor_cols <- params$meta[variable %in% to_sim$variable, family == 'multinom']
 
   # Prepare the event space
   if (is.null(evidence)) {
     num_trees <- params$forest[, max(tree)]
     omega <- params$forest[, .(f_idx, cvg)]
     omega[, wt := cvg / num_trees]
     omega[, cvg := NULL]
+    omega[, row_idx := 1]
   } else if (isTRUE(conj)) {
     omega <- leaf_posterior(params, evidence)
   } else {
@@ -110,18 +113,20 @@ forge <- function(
     draws <- omega[, .(f_idx)]
   } else {
     # Draw random leaves with probability proportional to weight
-    draws <- data.table(
-      'f_idx' = omega[, sample(f_idx, size = n_synth, replace = TRUE, prob = wt)]
-    )
+    draws <- omega[, sample(f_idx, size = n_synth, replace = TRUE, prob = wt), by = row_idx]
+    setnames(draws, 'V1', 'f_idx')
   }
-  omega <- merge(draws, omega, sort = FALSE)[, idx := .I]
+  omega <- merge(draws, omega, by = c("row_idx", "f_idx"), sort = FALSE)[, idx := .I]
 
   # Simulate continuous data
   synth_cnt <- synth_cat <- NULL
   if (any(!factor_cols)) {
     fam <- params$meta[family != 'multinom', unique(family)]
-    psi <- merge(omega, params$cnt[variable %in% to_sim], by = 'f_idx', 
+    merge(omega, params$cnt, by = 'f_idx', 
+          sort = FALSE, allow.cartesian = TRUE)
+    psi <- merge(omega, params$cnt, by = 'f_idx', 
                  sort = FALSE, allow.cartesian = TRUE)
+    psi <- merge(psi, to_sim, by = c("row_idx", "variable"))
     if (isTRUE(conj)) {
       if (any(evidence$relation %in% c('<', '<=', '>', '>='))) {
         for (k in evidence[, which(grepl('<', relation))]) {
@@ -141,13 +146,14 @@ forge <- function(
     } else if (fam == 'unif') {
       psi[, dat := stats::runif(.N, min = min, max = max)]
     }
-    synth_cnt <- dcast(psi, idx ~ variable, value.var = 'dat')[, idx := NULL]
+    synth_cnt <- dcast(psi, idx + row_idx ~ variable, value.var = 'dat')[, idx := NULL]
   }
 
   # Simulate categorical data
   if (any(factor_cols)) {
-    psi <- merge(omega, params$cat[variable %in% to_sim], by = 'f_idx',
+    psi <- merge(omega, params$cat, by = 'f_idx',
                  sort = FALSE, allow.cartesian = TRUE)
+    psi <- merge(psi, to_sim, by = c("row_idx", "variable"))
     psi[prob == 1, dat := val] 
     if (isTRUE(conj)) {
       if (any(evidence[, relation == '!='])) {
@@ -158,20 +164,45 @@ forge <- function(
       }
     }
     psi[prob < 1, dat := sample(val, 1, prob = prob), by = .(variable, idx)]
-    psi <- unique(psi[, .(idx, variable, dat)])
-    synth_cat <- dcast(psi, idx ~ variable, value.var = 'dat')[, idx := NULL]
+    psi <- unique(psi[, .(idx, row_idx, variable, dat)])
+    synth_cat <- dcast(psi, idx + row_idx ~ variable, value.var = 'dat')[, idx := NULL][, row_idx := NULL]
   }
 
   # Combine, optionally impose constraint(s)
   x_synth <- cbind(synth_cnt, synth_cat)
-  if (length(to_sim) != params$meta[, .N]) {
-    tmp <- evidence[relation == '==']
-    add_on <- setDT(lapply(tmp[, .I], function(k) rep(tmp[k, value], n_synth)))
-    setnames(add_on, tmp[, variable])
-    x_synth <- cbind(x_synth, add_on)
+  if (nrow(to_sim)/to_sim[, max(row_idx)] != params$meta[, .N]) {
+    add_on_cnt <- add_on_cat <- NULL
+    if (any(!params$meta[, family == 'multinom'])) {
+      tmp_cnt <- merge(evidence[relation == '=='], params$meta[family != "multinom", ], by = "variable")[, .(row_idx, variable, value)]
+      if (nrow(tmp_cnt) > 0) {
+        tmp_cnt[, value := as.numeric(value)]
+        add_on_cnt <- dcast(tmp_cnt, row_idx ~ variable, value.var = 'value')
+      }
+    }
+    if (any(params$meta[, family == 'multinom'])) {
+      tmp_cat <- merge(evidence[relation == '=='], params$meta[family == "multinom", ], by = "variable")[, .(row_idx, variable, value)]
+      if (nrow(tmp_cat) > 0) {
+        add_on_cat <- dcast(tmp_cat, row_idx ~ variable, value.var = 'value')
+      }
+    }
+    if (!is.null(add_on_cnt) & !is.null(add_on_cat)) {
+      add_on <- merge(add_on_cnt, add_on_cat, by = "row_idx")
+    } else if (!is.null(add_on_cnt)) {
+      add_on <- add_on_cnt
+    } else if (!is.null(add_on_cat)) {
+      add_on <- add_on_cat
+    }
+    if (any(!(colnames(x_synth) %in% colnames(add_on)))) {
+      x_synth <- cbind(x_synth, add_on)
+    } else {
+      x_synth <- dplyr::rows_patch(add_on, x_synth, by = "row_idx")
+    }
   }
 
   # Clean up, export
+  if (n_synth == 1 | x_synth[, max(row_idx)] == 1) {
+    x_synth[, row_idx := NULL]
+  }
   x_synth <- post_x(x_synth, params)
   return(x_synth)
 }

R/utils.R

@@ -113,15 +113,16 @@ prep_evi <- function(params, evidence) {
       err <- setdiff(colnames(evidence), params$meta$variable)
       stop('Unrecognized feature(s) among colnames: ', err)
     }
+    evidence[, row_idx := .I]
     evidence <- suppressWarnings(
-      melt(evidence, measure.vars = colnames(evidence), variable.factor = FALSE)
+      melt(evidence, id.vars = "row_idx", variable.factor = FALSE, na.rm = TRUE)
     )
     evidence[, relation := '==']
     conj <- TRUE
   }
   if (isTRUE(conj)) {
     evi <- merge(params$meta, evidence, by = 'variable', sort = FALSE)
-    evi[, n := .N, by = variable]
+    evi[, n := .N, by = .(variable, row_idx)]
     if (any(evi[n > 1L, relation == '=='])) {
       culprit <- evi[n > 1L & relation == '==', variable]
       stop(paste('Inconsistent conditioning events for the following variable(s):', 
@@ -156,7 +157,6 @@ prep_evi <- function(params, evidence) {
   ### ALSO: Reduce redundant events to most informative condition
   ###       and check for inconsistencies, e.g. >3 & <2
 
-
   return(evidence)
 }
 
@@ -186,8 +186,8 @@ leaf_posterior <- function(params, evidence) {
   # Continuous features
   if (any(evidence$family != 'multinom')) { 
     evi <- evidence[family != 'multinom']
-    evi[, n := .N, by = variable]
-    psi <- merge(evi, params$cnt, by = 'variable')
+    evi[, n := .N, by = .(variable, row_idx)]
+    psi <- merge(evi, params$cnt, by = 'variable', allow.cartesian = TRUE)
     if (any(evi$relation == '==')) {
       psi[relation == '==', lik := 
             truncnorm::dtruncnorm(value, a = min, b = max, mean = mu, sd = sigma)]
@@ -208,28 +208,32 @@ leaf_posterior <- function(params, evidence) {
       psi[n == 1 & relation %in% c('>', '>='), lik := 1 - lik]
     }
     psi[value == min, lik := 0]
-    psi_cnt <- unique(psi[, .(f_idx, variable, lik)])
+    psi_cnt <- unique(psi[, .(f_idx, row_idx, variable, lik)])
   }
 
   # Categorical features
   psi_eq <- psi_ineq <- NULL
   if (any(evidence$family == 'multinom')) { 
     evi <- evidence[family == 'multinom']
     evi[, value := as.character(value)]
-    grd <- rbindlist(lapply(evi[, variable], function(j) {
+
+
+
+    grd <- rbindlist(lapply(evi[, unique(variable)], function(j) {
       expand.grid('f_idx' = params$forest$f_idx, 'variable' = j,
                   'val' = params$cat[variable == j, unique(val)],
+                  'row_idx' = evi[, unique(row_idx)],
                   stringsAsFactors = FALSE)
     }))
     psi <- merge(params$cat, grd, by = c('f_idx', 'variable', 'val'),
                  sort = FALSE, all.y = TRUE)
     psi[is.na(prob), prob := 0]
     setnames(psi, 'prob', 'lik')
     if (any(evi[, relation == '=='])) {
-      evi_tmp <- evi[relation == '==', .(variable, value)]
+      evi_tmp <- evi[relation == '==', .(variable, value, row_idx)]
       setnames(evi_tmp, 'value', 'val')
-      psi_eq <- merge(psi, evi_tmp, by = c('variable', 'val'), sort = FALSE)
-      psi_eq <- psi_eq[, .(f_idx, variable, lik)]
+      psi_eq <- merge(psi, evi_tmp, by = c('variable', 'val', 'row_idx'), sort = FALSE)
+      psi_eq <- psi_eq[, .(f_idx, row_idx, variable, lik)]
     }
     if (any(evi[, relation == '!='])) {
       evi_tmp <- evi[relation == '!=', .(variable, value)]
@@ -252,18 +256,20 @@ leaf_posterior <- function(params, evidence) {
       exp(mean(log(c(cvg[1], lik))))
     }
   }
-  , by = f_idx]
+  , by = .(f_idx, row_idx)]
 
   # Normalize, export
-  out <- unique(psi[wt > 0, .(f_idx, wt)])
+  out <- unique(psi[wt > 0, .(f_idx, row_idx, wt)])
   if (nrow(out) == 0) {
     # If all leaves have zero weight, choose one randomly
     warning("All leaves have zero likelihood. This is probably because evidence contains an (almost) impossible combination. For categorical data, consider setting alpha>0 in forde().")
     out <- unique(psi[, .(f_idx)])
     out[, wt := 1]
   }
-  out[, wt := (wt / max(wt, na.rm = T))^(nrow(evidence) + 1)][wt > 0, wt := wt / sum(wt)]
-  return(out[])
+  evi_n <- evidence[, .N, by = .(row_idx)]
+  out <- merge(out, evi_n, by = "row_idx")
+  out[, wt := (wt / max(wt, na.rm = T))^(N + 1), by = row_idx][wt > 0, wt := wt / sum(wt), by = row_idx]
+  return(out[, .(f_idx, row_idx, wt)])
 }