Optional profiling

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: epinowcast
 Title: Hierarchical Nowcasting of Right Censored Epidemiological Counts
-Version: 0.0.4.10000
+Version: 0.0.6.2000
 Authors@R:
   c(person(given = "Sam Abbott",
            role = c("aut", "cre"),

NEWS.md

@@ -1,4 +1,7 @@
-# epinowcast 0.0.4.10000
+# epinowcast 0.0.6
+* Add profiling switch to model compilation, allowing to toggle profiling (https://mc-stan.org/cmdstanr/articles/profiling.html) on/off in the same model (see [PR#41](https://github.com/epiforecasts/epinowcast/pull/41) by [@adrian-lison](https://github.com/adrian-lison)).
+
+# epinowcast 0.0.5
 
 * Convert retrospective data date fields to class of `IDate` when utilising `enw_retrospective_data` to solve esoteric error.
 * Added full argument name for `include_paths` to avoid console chatter

R/model.R

@@ -229,6 +229,8 @@ enw_inits <- function(data) {
 #' @param verbose Logical, defaults to `TRUE`. Should verbose
 #' messages be shown.
 #'
+#' @param profile Logical, defaults to `FALSE`. Should the model be profiled (see https://mc-stan.org/cmdstanr/articles/profiling.html)?
+#'
 #' @param stanc_options A list of options to pass to the `stanc_options` of [cmdstanr::cmdstan_model()]
 #' by default "01" is passed which specifies simple optimisations should be done by the prior to compilation
 #'
@@ -242,7 +244,7 @@ enw_inits <- function(data) {
 #' @examplesIf interactive()
 #' mod <- enw_model()
 enw_model <- function(model, include,
-                      compile = TRUE, threads = FALSE, stanc_options = list("O1"), verbose = TRUE, ...) {
+                      compile = TRUE, threads = FALSE, stanc_options = list("O1"), verbose = TRUE, profile = FALSE, ...) {
   if (missing(model)) {
     model <- "stan/epinowcast.stan"
     model <- system.file(model, package = "epinowcast")
@@ -252,6 +254,12 @@ enw_model <- function(model, include,
   }
 
   if (compile) {
+    if (!profile) {
+      code <- paste(readLines(model), collapse = "\n")
+      code_no_profile <- remove_profiling(code)
+      model <- cmdstanr::write_stan_file(code_no_profile)
+    }
+
     if (verbose) {
       model <- cmdstanr::cmdstan_model(model,
         include_paths = include,
@@ -276,6 +284,17 @@ enw_model <- function(model, include,
   return(model)
 }
 
+#' Remove profiling statements from a character vector representing stan code
+#'
+#' @param s Character vector representing stan code
+#' @return A `character` vector of the stan code without profiling statements
+remove_profiling <- function(s) {
+  while (grepl("profile\\(.+\\)\\{", s, perl = T)) {
+    s <- gsub("profile\\(.+\\)\\{((?:[^{}]++|\\{(?1)\\})++)\\}", "\\1", s, perl = T)
+  }
+  return(s)
+}
+
 #' Fit a CmdStan model using NUTS
 #'
 #' @param data A list of data as produced by [enw_as_data_list()].

inst/stan/epinowcast.stan

@@ -84,12 +84,16 @@ transformed parameters{
   vector[t] imp_obs[g]; // Expected final observations
   real phi; // Transformed overdispersion (joint across all observations)
   // calculate log mean and sd parameters for each dataset from design matrices
+  profile("transformed_delay_reference_date_total") {
+  profile("transformed_delay_reference_date_effects") {
   logmean = combine_effects(logmean_int, logmean_eff, d_fixed, logmean_sd,
                             d_random);
   logsd = combine_effects(log(logsd_int), logsd_eff, d_fixed, logsd_sd,
                           d_random); 
   logsd = exp(logsd);
+  }
   // calculate pmfs
+  profile("transformed_delay_reference_date_pmfs") {
   for (i in 1:npmfs) {
     pmfs[, i] = calculate_pmf(logmean[i], logsd[i], dmax, dist);
   }
@@ -101,11 +105,16 @@ transformed parameters{
   }else{
     ref_lh = pmfs;
   }
+  }
+  }
   // calculate sparse report date effects with forced 0 intercept
+  profile("transformed_delay_reporting_date_effects") {
   srdlh = combine_effects(0, rd_eff, rd_fixed, rd_eff_sd, rd_random);
+  }
   // estimate unobserved expected final reported cases for each group
   // this could be any forecasting model but here its a 
   // first order random walk for each group on the log scale.
+  profile("transformed_expected_final_observations") {
   for (k in 1:g) {
     real llast_obs;
     imp_obs[k][1] = leobs_init[k];
@@ -114,6 +123,7 @@ transformed parameters{
     }
     imp_obs[k] = exp(imp_obs[k]);
   }
+  }
   // transform phi to overdispersion scale
   phi = 1 / sqrt(sqrt_phi);
   // debug issues in truncated data if/when they appear
@@ -123,6 +133,7 @@ transformed parameters{
 }
 
 model {
+  profile("model_priors") {
   // priors for unobserved expected reported cases
   leobs_init ~ normal(eobs_init, 1);
   for (i in 1:g) {
@@ -154,14 +165,18 @@ model {
   }
   // reporting overdispersion (1/sqrt)
   sqrt_phi ~ normal(sqrt_phi_p[1], sqrt_phi_p[2]) T[0,];
+  }
   // log density: observed vs model
   if (likelihood) {
+    profile("model_likelihood") {
     target += reduce_sum(obs_lupmf, st, 1, obs, sl, imp_obs, sg, st, rdlurd,
                          srdlh, ref_lh, dpmfs, ref_p, phi);
+    }
   }
 }
 
 generated quantities {
+  profile("generated_total") {
   int pp_obs[pp ? sum(sl) : 0];
   vector[ologlik ? s : 0] log_lik;
   int pp_inf_obs[cast ? dmax : 0,cast ? g : 0];
@@ -172,18 +187,23 @@ generated quantities {
     int pp_obs_tmp[s, dmax];
     // Posterior predictions for observations
     for (i in 1:s) {
+      profile("generated_obs") {
       tar_obs = imp_obs[sg[i]][st[i]];
       rdlh = srdlh[rdlurd[st[i]:(st[i] + dmax - 1), sg[i]]];
       exp_obs = expected_obs(tar_obs, ref_lh[1:dmax, dpmfs[i]], rdlh, ref_p);
       pp_obs_tmp[i, 1:dmax] = neg_binomial_2_rng(exp_obs, phi);
+      }
+      profile("generated_loglik") {
       if (ologlik) {
         log_lik[i] = 0;
         for (j in 1:sl[i]) {
           log_lik[i] += neg_binomial_2_lpmf(obs[i, j] | exp_obs[j], phi);
         }
       }
+      }
     }
     // Posterior prediction for final reported data (i.e at t = dmax + 1)
+    profile("generated_obs") {
     for (k in 1:g) {
       int start_t = t - dmax;
       for (i in 1:dmax) {
@@ -204,5 +224,7 @@ generated quantities {
         start_t += sl[i];
       }
     }
+    }
   }
 }
+}