From f12aa4a9a68abe88644a8cffed7383208b1a8fe6 Mon Sep 17 00:00:00 2001
From: gufengzhou <gufengzhou@gmail.com>
Date: Wed, 3 May 2023 11:15:10 +0800
Subject: [PATCH] recode: allow convergence status in json

---
 R/R/clusters.R       |  2 ++
 R/R/json.R           | 11 +++++++++--
 R/R/outputs.R        |  3 ++-
 R/man/robyn_write.Rd |  3 +++
 4 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/R/R/clusters.R b/R/R/clusters.R
index 0b46d4886..d0f27601a 100644
--- a/R/R/clusters.R
+++ b/R/R/clusters.R
@@ -236,6 +236,8 @@ confidence_calcs <- function(xDecompAgg, cls, all_paid, dep_var_type, k, boot_n
       boot_ci = sprintf("[%s, %s]", round(.data$ci_low, 2), round(.data$ci_up, 2)),
       ci_low = .data$ci_low,
       ci_up = .data$ci_up,
+      sd = boot_se * sqrt(.data$n - 1),
+      dist100 = (.data$ci_up - .data$ci_low + 2 * boot_se * sqrt(.data$n - 1)) / 99,
       .groups = "drop"
     ) %>%
     ungroup()
diff --git a/R/R/json.R b/R/R/json.R
index 37d4047ae..193f1681a 100644
--- a/R/R/json.R
+++ b/R/R/json.R
@@ -11,6 +11,7 @@
 #'
 #' @inheritParams robyn_outputs
 #' @param InputCollect \code{robyn_inputs()} output.
+#' @param OutputModels \code{robyn_run()} output.
 #' @param select_model Character. Which model ID do you want to export
 #' into the JSON file?
 #' @param dir Character. Existing directory to export JSON file to.
@@ -30,6 +31,7 @@
 #' @export
 robyn_write <- function(InputCollect,
                         OutputCollect = NULL,
+                        OutputModels = NULL,
                         select_model = NULL,
                         dir = OutputCollect$plot_folder,
                         export = TRUE,
@@ -51,7 +53,12 @@ robyn_write <- function(InputCollect,
   ret <- list()
   skip <- which(unlist(lapply(InputCollect, function(x) is.list(x) | is.null(x))))
   skip <- skip[!names(skip) %in% c("calibration_input", "hyperparameters", "custom_params")]
-  ret[["InputCollect"]] <- inputs <- InputCollect[-skip]
+  ret[["InputCollect"]] <- InputCollect[-skip]
+  conv_msg <- mapply(function(x) x[[1]], x = gregexpr(":", OutputModels$convergence$conv_msg),
+                     SIMPLIFY = FALSE)
+  conv_msg <- mapply(function(x, y) substr(x, 1, y-1), x = OutputModels$convergence$conv_msg,
+                     y = conv_msg, USE.NAMES = FALSE)
+  ret[["OutputCollect"]][["conv_msg"]] <- conv_msg
   # toJSON(inputs, pretty = TRUE)
 
   # ExportedModel JSON
@@ -100,7 +107,7 @@ robyn_write <- function(InputCollect,
         (all_sol_json %>% filter(.data$cluster == x))$solID
       })
       names(all_sol_json) <- paste0("cluster", all_c)
-      ret[["all_sols"]] <- all_sol_json
+      ret[["OutputCollect"]][["all_sols"]] <- all_sol_json
     }
     write_json(ret, filename, pretty = TRUE)
   }
diff --git a/R/R/outputs.R b/R/R/outputs.R
index 5c17af6fa..a42f21686 100644
--- a/R/R/outputs.R
+++ b/R/R/outputs.R
@@ -233,7 +233,8 @@ robyn_outputs <- function(InputCollect, OutputModels,
         } else {
           all_sol_json <- NULL
         }
-        robyn_write(InputCollect, dir = OutputCollect$plot_folder, quiet = quiet, all_sol_json = all_sol_json)
+        robyn_write(InputCollect = InputCollect, OutputModels = OutputModels,
+                    dir = OutputCollect$plot_folder, quiet = quiet, all_sol_json = all_sol_json)
 
         # For internal use -> UI Code
         if (ui && plot_pareto) OutputCollect$UI$pareto_onepagers <- pareto_onepagers
diff --git a/R/man/robyn_write.Rd b/R/man/robyn_write.Rd
index 74029e081..22852a35d 100644
--- a/R/man/robyn_write.Rd
+++ b/R/man/robyn_write.Rd
@@ -11,6 +11,7 @@
 robyn_write(
   InputCollect,
   OutputCollect = NULL,
+  OutputModels = NULL,
   select_model = NULL,
   dir = OutputCollect$plot_folder,
   export = TRUE,
@@ -32,6 +33,8 @@ robyn_recreate(json_file, quiet = FALSE, ...)
 
 \item{OutputCollect}{\code{robyn_run(..., export = FALSE)} output.}
 
+\item{OutputModels}{\code{robyn_run()} output.}
+
 \item{select_model}{Character. Which model ID do you want to export
 into the JSON file?}