Progress on getting diagnostics into a function

NAMESPACE

@@ -2,6 +2,8 @@
 
 S3method(as_latent_individual,data.frame)
 S3method(epidist,default)
+S3method(epidist_diagnostics,default)
+S3method(epidist_diagnostics,epidist_fit)
 S3method(epidist_family,default)
 S3method(epidist_family,epidist_latent_individual)
 S3method(epidist_formula,default)
@@ -21,6 +23,7 @@ export(correct_primary_censoring_bias)
 export(draws_to_long)
 export(drop_zero)
 export(epidist)
+export(epidist_diagnostics)
 export(epidist_family)
 export(epidist_formula)
 export(epidist_prior)
@@ -45,7 +48,6 @@ export(plot_mean_posterior_pred)
 export(plot_recovery)
 export(plot_relative_recovery)
 export(reverse_obs_at)
-export(sample_model)
 export(simulate_double_censored_pmf)
 export(simulate_exponential_cases)
 export(simulate_gillespie)
@@ -63,6 +65,7 @@ importFrom(checkmate,assert_int)
 importFrom(checkmate,assert_names)
 importFrom(checkmate,assert_numeric)
 importFrom(posterior,as_draws_df)
+importFrom(rstan,lookup)
 importFrom(stats,as.formula)
 importFrom(stats,ecdf)
 importFrom(stats,integrate)

R/diagnostics.R

@@ -4,13 +4,14 @@
 #' @param fit ...
 #' @export
 epidist_diagnostics <- function(fit) {
-  UseMethod("epidist")
+  UseMethod("epidist_diagnostics")
 }
 
 #' Default method for returning diagnostics
 #'
 #' @param fit ...
 #' @family defaults
+#' @method epidist_diagnostics default
 #' @export
 epidist_diagnostics.default <- function(fit) {
   stop(
@@ -23,36 +24,38 @@ epidist_diagnostics.default <- function(fit) {
 #'
 #' @param fit ...
 #' @family defaults
+#' @method epidist_diagnostics epidist_fit
 #' @export
-epidist_diagnostics.default <- function(fit) {
-  stop(
-    "No epidist_diagnostics method implemented for the class ", class(data), "\n",
-    "See methods(epidist_diagnostics) for available methods"
-  )
-}
-
-#' Default method for returning diagnostics
-#'
-#' @param fit ...
-#' @family defaults
-#' @export
-epidist_diagnostics <- function(fit) {
-  # diag <- fit$sampler_diagnostics(format = "df")
-  # diagnostics <- data.table(
-  #   samples = nrow(diag),
-  #   max_rhat = round(max(
-  #     fit$summary(
-  #       variables = NULL, posterior::rhat,
-  #       .args = list(na.rm = TRUE)
-  #     )$`posterior::rhat`,
-  #     na.rm = TRUE
-  #   ), 2),
-  #   divergent_transitions = sum(diag$divergent__),
-  #   per_divergent_transitions = sum(diag$divergent__) / nrow(diag),
-  #   max_treedepth = max(diag$treedepth__)
-  # )
-  # diagnostics[, no_at_max_treedepth := sum(diag$treedepth__ == max_treedepth)]
-  # diagnostics[, per_at_max_treedepth := no_at_max_treedepth / nrow(diag)]
+epidist_diagnostics.epidist_fit <- function(fit) {
+  if (fit$algorithm %in% c("laplace", "meanfield", "fullrank", "pathfinder")) {
+    cli::cli_abort(c(
+      "!" = paste0(
+        "Diagnostics not yet supported for the algorithm: ", fit$algorithm
+      )
+    ))
+  } else if (!fit$algorithm == "sampling") {
+    cli::cli_abort(c( 
+      "!" = paste0(
+        "Unrecognised algorithm: ", fit$algorithm
+      )
+    ))
+  } else if (fit$algorithm == "sampling") {
+    np <- brms::nuts_params(fit)
+    divergent_indices <- np$Parameter == "divergent__"
+    treedepth_indices <- np$Parameter == "treedepth__"
+    diagnostics <- data.table(
+      "samples" = nrow(np) / length(unique(np$Parameter)),
+      "max_rhat" = round(max(brms::rhat(fit)), 3),
+      "divergent_transitions" = sum(np[divergent_indices, ]$Value),
+      "per_divergent_transitions" = mean(np[divergent_indices, ]$Value),
+      "max_treedepth" = max(np[treedepth_indices, ]$Value)
+    )
+    diagnostics[, no_at_max_treedepth :=
+                    sum(np[treedepth_indices, ]$Value == max_treedepth)]
+    diagnostics[, per_at_max_treedepth := no_at_max_treedepth / samples]
+  }
+
+  # rstan::get_elapsed_time(fit$fit)
 
-  # timing <- round(fit$time()$total, 1)
+  return(diagnostics)
 }

R/globals.R

@@ -1,11 +1,6 @@
 # Generated by roxyglobals: do not edit by hand
 
 utils::globalVariables(c(
-  "error", # <sample_model>
-  "no_at_max_treedepth", # <sample_model>
-  "max_treedepth", # <sample_model>
-  "per_at_max_treedepth", # <sample_model>
-  "run_time", # <sample_model>
   "meanlog", # <add_natural_scale_mean_sd>
   "sdlog", # <add_natural_scale_mean_sd>
   "sd", # <add_natural_scale_mean_sd>

tests/testthat/test-unit-diagnostics.R

@@ -0,0 +1,16 @@
+test_that("", { # nolint: line_length_linter.
+  skip_on_cran()
+  set.seed(1)
+  prep_obs <- as_latent_individual(sim_obs)
+  fit_laplace <- epidist(data = prep_obs, seed = 1, algorithm = "laplace")
+  diag <- epidist_diagnostics(fit_laplace)
+})
+
+
+test_that("", { # nolint: line_length_linter.
+  skip_on_cran()
+  set.seed(1)
+  prep_obs <- as_latent_individual(sim_obs)
+  fit <- epidist(data = prep_obs, seed = 1)
+  diag <- epidist_diagnostics(fit)
+})