Scalable Record Linkage with FlexRL

peakRAMNew <- function(...){
  # Capture R expressions or function calls
  args <- c(as.list(match.call(expand.dots = FALSE)$`...`), NULL)
  Function <- sapply(args, function(e) paste0(deparse(e), collapse = ""))
  Function <- sapply(Function, function(e) gsub(" ", "", e))
  # Initialize containers for output
  numfunc <- length(args)
  totaltime <- vector("numeric", numfunc)
  RAMused <- vector("numeric", numfunc)
  RAMpeak <- vector("numeric", numfunc)
  i <- 1
  for(arg in args){
    # Reset garbage collector and save baseline
    start <- gc(verbose = FALSE, reset = TRUE)
    # Evaluate regular and anonymous functions
    evalTime <- system.time(result <- eval.parent(arg))
    if(inherits(result, "function")){
      evalTime <- system.time(output <- result())
      rm(result)
    }else{
      output <- result
      rm(result)
    }
    # Call garbage collector and save post-eval
    end <- gc(verbose = FALSE, reset = FALSE)
    rm(output)
    # Calculate total and peak RAM used
    totaltime[i] <- as.numeric(evalTime["elapsed"])
    get_column_mb = function(object, before_column) {
      ind = which(colnames(object) == before_column)
      stopifnot(colnames(object)[ind+1] == "(Mb)")
      if (length(ind) == 0) {
        stop("No value found for that column")
      }
      # need to ref index and not colname because duplicated colnames for (Mb)
      object[2, ind + 1]
    }
    RAMused[i] <- get_column_mb(end, "used")- get_column_mb(start, "used")
    RAMpeak[i] <- get_column_mb(end, "max used") - get_column_mb(start, "max used")
    i <- i + 1
  }
  data.frame("Function_Call" = Function,
             "Elapsed_Time_sec" = totaltime,
             "Total_RAM_Used_MiB" = RAMused,
             "Peak_RAM_Used_MiB" = RAMpeak,
             row.names = 1:numfunc)
}

wrapper_args <- function(method, prep_data, ...) {
  PIVs <- names(prep_data$PIVs_config)
  args <- c(...)
  switch(method,
         FlexRL    = c(list(data = prep_data, threads = 4),
                       args[intersect(names(args), names(formals(FlexRL::StEM)))]),
         BRL       = c(list(flds = PIVs, types = rep("bi", length(PIVs)), nIter = 1000),
                       args[intersect(names(args), c(names(formals(BRL::compareRecords)), names(formals(BRL::bipartiteGibbs))))]),
         fastLink  = c(list(varnames = PIVs, tol.em = 1e-06, threshold.match = 0.5, return.all = TRUE, n.cores = 4),
                       args[intersect(names(args), names(formals(fastLink::fastLink)))]),
         multilink = c(list(types = rep("bi", length(PIVs)), breaks = stats::setNames(as.list(rep(NA, length(PIVs))), PIVs),
                            duplicates = c(0, 0), n_iter = 1000, max_cc_size = 50),
                       args[intersect(names(args), c(names(formals(multilink::create_comparison_data)),
                                                     names(formals(multilink::reduce_comparison_data)),
                                                     names(formals(multilink::specify_prior)),
                                                     names(formals(multilink::gibbs_sampler)),
                                                     names(formals(multilink::find_bayes_estimate)),
                                                     names(formals(multilink::relabel_bayes_estimate))))]),
         reclin2   = c(list(on = PIVs, formula = stats::as.formula(paste("~", paste(PIVs, collapse = "+"))),
                            type = "mpost", add = TRUE, variable = "selected", score = "mpost", threshold = 0.5),
                       args[intersect(names(args), c(names(formals(reclin2::pair)),
                                                     names(formals(reclin2::compare_pairs)),
                                                     names(formals(reclin2::problink_em)),
                                                     names(formals(stats::predict)),
                                                     names(formals(reclin2::select_threshold))))]),
         diyar     = c(list(probabilistic = TRUE, return_weights = TRUE),
                       args[intersect(names(args), c(names(formals(diyar::prob_score_range)),
                                                     names(formals(diyar::links_wf_probabilistic))))]),
         fedmatch  = c(list(by = PIVs, match_type = "multivar", unique_key_1 = "unique_key_A",
                            unique_key_2 = "unique_key_B", nthread = 4, 
                            multivar_settings = fedmatch::build_multivar_settings(
                              compare_type = rep("indicator", length(PIVs)),
                              wgts = rep(1 / length(PIVs), length(PIVs)))),
                       args[intersect(names(args), names(formals(fedmatch::merge_plus)))])
  )
}

# unix::rlimit_as(20e9) # may be needed on some OS

run_wrapper <- function(method, prep_data, arguments) {
  PIVs <- names(prep_data$PIVs_config)
  link_with <- get(paste0("link_with_", method), envir = asNamespace("FlexRL"))
  if (method == "diyar") {
    prep_data$encodedA <- prep_data$encodedA[, c(PIVs, "local_id", "source")]
    prep_data$encodedB <- prep_data$encodedB[, c(PIVs, "local_id", "source")]
  }
  if (Sys.info()["sysname"] == "Darwin") {
    tryCatch(link_with(prep_data$encodedA, prep_data$encodedB, arguments),
           error = function(e) {
             if (grepl("cannot allocate|memory|bad_alloc|too large|long vectors", conditionMessage(e), ignore.case = TRUE)) {
               message(method, ": memory error (", conditionMessage(e), "), results reported as NA")
               return(NULL)
             }
             stop(e)
           })
  } else if (Sys.info()["sysname"] %in% c("Linux","Windows")) {
      tryCatch(callr::r(function(link_with, A, B, arguments) {link_with(A, B, arguments)},
                      args = list(link_with = link_with, A = prep_data$encodedA,
                                  B = prep_data$encodedB, arguments = arguments)),
           error = function(e) {
             if (grepl("cannot allocate|memory|bad_alloc|too large|long vectors|crashed or was killed|signal 9|SIGKILL",
               conditionMessage(e), ignore.case = TRUE)) {
               message(method, ": memory error (", conditionMessage(e), "), results reported as NA")
               return(NULL)
             }
             stop(e)
           })
  } else {
    print("Unknown Operating System.")
  }
}

evaluate_linkage <- function(fit, true_pairs, threshold = 0.5) {
  if (is.null(fit)) return(rep(NA_real_, 5))
  delta_result <- data.frame(idxA = fit$idxA, idxB = fit$idxB)
  if (!is.null(fit$LinkScore)) delta_result <- delta_result[fit$LinkScore > threshold, , drop = FALSE]
  linked_pairs   <- do.call(paste, c(delta_result, list(sep = "_")))
  true_positive  <- length(intersect(linked_pairs, true_pairs))
  false_positive <- length(setdiff(linked_pairs, true_pairs))
  false_negative <- length(setdiff(true_pairs, linked_pairs))
  c(TP = true_positive, FP = false_positive, FN = false_negative,
    sensitivity = true_positive / (true_positive + false_negative),
    FDP = false_positive / (true_positive + false_positive))
}

results_rows <- c("TP               ", "FP               ", "FN               ", 
                  "sensitivity      ", "FDP              ", "est. FDP score   ", 
                  "est. FDP* score  ", "est. FDP* synth  ", "max SMD          ", 
                  "min support IoU  ", "max MMD          ", "time (secs)      ", 
                  "memory (Mb)      ")

methods <- c("BRL", "multilink", "fastLink", "reclin2", "diyar", "fedmatch", "FlexRL")

run_all_wrappers <- function(prep_data, variables_type, methods, FDP_estimation = TRUE,
                             maxIter4CV = 2, n_repeats = 3, ...) {
  args <- list(...)
  op <- options()
  on.exit(options(op))
  PIVs <- names(prep_data$PIVs_config)
  true_pairs <- do.call(paste, c(prep_data$true_pairs, list(sep = "_")))
  df_results <- data.frame(matrix(NA, nrow = length(results_rows), ncol = 0))
  rownames(df_results) <- results_rows
  
  t0 <- Sys.time()
  MeM <- peakRAMNew(naive_fit <- naive_linkage(PIVs = PIVs, encodedA = prep_data$encodedA, encodedB = prep_data$encodedB))
  elapsed <- as.numeric(Sys.time() - t0, units = "secs")
  df_results[1:5, "Naive"] <- evaluate_linkage(list(idxA = naive_fit[[1]], idxB = naive_fit[[2]]), true_pairs)
  df_results[12,"Naive"] <- elapsed
  df_results[13, "Naive"] <- MeM$Peak_RAM_Used_MiB

  for (method in methods) {
    arguments <- wrapper_args(method, prep_data, args)
    
    t0 <- Sys.time()
    MeM <- peakRAMNew(fit <- run_wrapper(method, prep_data, arguments))
    elapsed <- as.numeric(Sys.time() - t0, units = "secs")
    df_results[1:5, method] <- evaluate_linkage(fit, true_pairs)
    if (is.null(fit) || length(fit$idxA) == 0) next
    
    encodedA <- prep_data$encodedA
    encodedB <- prep_data$encodedB
    if (method == "diyar") {
      encodedA <- encodedA[, c(PIVs, "local_id", "source")]
      encodedB <- encodedB[, c(PIVs, "local_id", "source")]
    }
    
    diag <- do.call(RL_diagnostics, c(list(fit, encodedA, encodedB, PIVs, variables_type, PIVs,
                                           true_pairs = prep_data$true_pairs, FDP_estimation = FDP_estimation,
                                           RL_method = method, data = prep_data, maxIter4CV = maxIter4CV, n_repeats = n_repeats),
                                      arguments))
    if (FDP_estimation) {
      df_results[6, method] <- diag$FDP_measures$curve$RL_FDP_score[1]
      df_results[7, method] <- diag$FDP_measures$curve$augmRL_FDP_score[1]
      df_results[8, method] <- diag$FDP_measures$curve$augmRL_FDP_synth[1]
    }
    df_results[9,  method] <- max(diag$discrepancy_measures[1, grep("smd", names(diag$discrepancy_measures))])
    df_results[10, method] <- min(diag$discrepancy_measures[1, grep("iou", names(diag$discrepancy_measures))])
    df_results[11, method] <- max(diag$discrepancy_measures[1, grep("mmd", names(diag$discrepancy_measures))])
    
    df_results[12, method] <- elapsed
    df_results[13, method] <- MeM$Peak_RAM_Used_MiB
  }
  df_results
}

df1982 <- read.csv("ex-reg-NLTCS-82.csv", row.names = 1)
df1994 <- read.csv("ex-reg-NLTCS-94.csv", row.names = 1)

PIVs_config <- list(
  SEX       = list(dynamics = "stable",   bound_mistakes = c(0.10, 0.10), fix_mistakes = c(NA, NA)),
  DOB_MONTH = list(dynamics = "stable",   bound_mistakes = c(0.10, 0.10), fix_mistakes = c(NA, NA)),
  DOB_YEAR  = list(dynamics = "stable",   bound_mistakes = c(0.10, 0.10), fix_mistakes = c(NA, NA)),
  STATE     = list(dynamics = "flexible", bound_mistakes = c(NA, NA),     fix_mistakes = c(NA, NA)))

PIVs_type <- list(SEX = FALSE, DOB_MONTH = TRUE, DOB_YEAR = TRUE, STATE = FALSE)

prep_data <- prepare_data(df1982, df1994, "1982", "1994", PIVs_config, same_mistakes = TRUE, uniq_id = "ID")
#> '1982' is the larger source, saved as B; '1994' saved as A.

c(n_A = nrow(prep_data$encodedA), n_B = nrow(prep_data$encodedB), n_links = nrow(prep_data$true_pairs))
#>     n_A     n_B n_links 
#>     491     901     332

prep_data$n_values
#>       SEX DOB_MONTH  DOB_YEAR     STATE 
#>         2        12        28         9

df_results_NLTCS <- run_all_wrappers(prep_data, PIVs_type, methods)

cat("Record Linkage\n")
#> Record Linkage
knitr::kable(df_results_NLTCS[1:5, ], digits = 2, align=rep('c', ncol(df_results_NLTCS)))

cat("FDP estimation\n")
#> FDP estimation
knitr::kable(df_results_NLTCS[6:8, ], digits = 2, align=rep('c', ncol(df_results_NLTCS)))

cat("Linked divergence\n")
#> Linked divergence
knitr::kable(df_results_NLTCS[9:11, ], digits = 2, align=rep('c', ncol(df_results_NLTCS)))

cat("Computational time and memory\n")
#> Computational time and memory
knitr::kable(df_results_NLTCS[12:13, ], digits = 2, align=rep('c', ncol(df_results_NLTCS)))

df2012 <- comp[comp$ANNO == 2012, var_of_interest]
df2014 <- comp[comp$ANNO == 2014, var_of_interest]
df2012 <- df2012[!df2012$ID %in% df2012$ID[duplicated(df2012$ID)], ]
df2014 <- df2014[!df2014$ID %in% df2014$ID[duplicated(df2014$ID)], ]
rownames(df2012) <- 1:nrow(df2012)
rownames(df2014) <- 1:nrow(df2014)

PIVs_config <- list(
  ANASC          = list(dynamics = "stable", bound_mistakes = c(0.10, 0.10), fix_mistakes = c(NA, NA)),
  SESSO          = list(dynamics = "stable", bound_mistakes = c(0.10, 0.10), fix_mistakes = c(NA, NA)),
  STACIV         = list(dynamics = "flexible", bound_mistakes = c(NA, NA), fix_mistakes = c(NA, NA)),
  NASCREG        = list(dynamics = "stable", bound_mistakes = c(0.10, 0.10), fix_mistakes = c(NA, NA)),
  SETTP11QUALP10 = list(dynamics = "flexible", bound_mistakes = c(NA, NA), fix_mistakes = c(NA, NA)))

PIVs_type <- list(ANASCI = TRUE, SESSO = FALSE, STACIV = FALSE, NASCREG = FALSE, SETTP11QUALP10 = FALSE)

PIVs_config_full <- c(PIVs_config,
  list(IREG = list(dynamics = "flexible", bound_mistakes = c(NA, NA), fix_mistakes = c(NA, NA))))
PIVs_type_full <- c(PIVs_type, IREG = FALSE)

prep_data_full <- prepare_data(df2012, df2014, "2012", "2014", PIVs_config_full, same_mistakes = TRUE, uniq_id = "ID")
#> '2012' is the larger source, saved as B; '2014' saved as A.
true_pairs_full <- do.call(paste, c(prep_data_full$true_pairs, list(sep = "_")))

df_results_full <- data.frame(matrix(NA, nrow = 7, ncol = 0))
rownames(df_results_full) <- c("TP", "FP", "FN", "sensitivity", "FDP", "time (mins)", "memory (Mb)")
for (method in methods) {
  t0 <- Sys.time()           
  args <- list(n_iter = 100)
  MeM <- peakRAMNew(fit <- run_wrapper(method, prep_data_full, wrapper_args(method, prep_data_full, args)))
  df_results_full[1:5, method] <- evaluate_linkage(fit, true_pairs_full)
  df_results_full[6, method] <- if (is.null(fit)) NA else round(as.numeric(Sys.time() - t0, units = "mins"))
  df_results_full[7, method] <- if (is.null(fit)) NA else MeM$Peak_RAM_Used_MiB
}

cat("Record Linkage performance on large files")
#> Record Linkage performance on large files
knitr::kable(df_results_full, digits = 2, align=rep('c', ncol(df_results_full)))

regions <- sort(unique(c(df2012$IREG, df2014$IREG)))
results_IREG <- list()
for (r in regions) {
  block2012 <- df2012[df2012$IREG == r, ]
  block2014 <- df2014[df2014$IREG == r, ]
  prep_data_r <- prepare_data(block2012, block2014, "2012", "2014", PIVs_config, same_mistakes = TRUE, uniq_id = "ID")
  true_pairs_r <- do.call(paste, c(prep_data_r$true_pairs, list(sep = "_")))
  for (method in methods) {
    t0 <- Sys.time()
    MeM <- peakRAMNew(fit <- run_wrapper(method, prep_data_r, wrapper_args(method, prep_data_r)))
    results_IREG[[length(results_IREG) + 1]] <- data.frame(
      IREG = r, n_A = nrow(prep_data_r$encodedA),
      n_B = nrow(prep_data_r$encodedB), n_links = length(true_pairs_r),
      method = method, round(t(evaluate_linkage(fit, true_pairs_r)),2),
      seconds = round(as.numeric(Sys.time() - t0, units = "secs")),
      memory = MeM$Peak_RAM_Used_MiB)
  }
}

results_IREG <- do.call(rbind, results_IREG)

cat("Record Linkage on regional subsets")
#> Record Linkage on regional subsets
knitr::kable(results_IREG, digits = 2, align=rep('c', ncol(results_IREG)))

summary_IREG <- do.call(rbind, lapply(split(results_IREG, results_IREG$method), function(d) {
  data.frame(TP = mean(d$TP, na.rm=TRUE), FP = mean(d$FP, na.rm=TRUE), 
             FN = mean(d$FN, na.rm=TRUE), sensitivity = mean(d$sensitivity, na.rm=TRUE),
             FDP = mean(d$FDP, na.rm=TRUE), method = d$method[1],
             minutes = sum(d$seconds) / 60,
             memory = mean(d$memory))
}))
rownames(summary_IREG) <- summary_IREG$method

newdf <- as.data.frame(apply(t(summary_IREG)[c("TP","FP","FN","sensitivity","FDP","minutes","memory"),methods], 2, function(x)  sprintf("%.2f",as.numeric(as.character(x)))))
names(newdf) <- names(df_results_full)
rownames(newdf) <- rownames(df_results_full)
knitr::kable(newdf, align=rep('c', ncol(newdf)))

colours <- grDevices::hcl.colors(length(methods), "Dark 3")
names(colours) <- methods
plot(NA, xlim = c(0, 1), ylim = c(0, 1), xlab = "sensitivity", ylab = "FDP",
     main = "SHIW linked region by region (one point per region)")
for (method in methods) {
  d <- results_IREG[results_IREG$method == method, ]
  points(d$sensitivity, d$FDP, col = colours[method], pch = 19)
}
legend("bottomright", legend = methods, col = colours, pch = 19, bty = "n")