streamline S3 classes

[sbfnk]

Currently the three models included in the package return S3 class with similar but slightly different contents. Whilst to some degree that's expected because of their different use cases it might be nice to streamline this a bit to simplify the user experience and allow common access functions.

I'm thinking that perhaps as a minimum they should all return args (the result of create_stan_args), fit (the stanfit object) and observations as everything else, I think, could be constructed from these three elements using appropriate summary or extract_... methods.

If doing this changes would be likely to break existing workflows, but as we're introducing breaking changes towards 2.0.0 anyway this might be a good time to do this if at all.

Current returns are:

library("EpiNow2")
library("data.table")

reported_cases <- example_confirmed[1:60]
est_inf <- estimate_infections(reported_cases)
#> Warning: There were 32 divergent transitions after warmup. See
#> https://mc-stan.org/misc/warnings.html#divergent-transitions-after-warmup
#> to find out why this is a problem and how to eliminate them.
#> Warning: Examine the pairs() plot to diagnose sampling problems

cases <- as.data.table(reported_cases)[, primary := confirm]
cases[, scaling := 0.4]
cases[, meanlog := 1.8][, sdlog := 0.5]
cases <- simulate_secondary(cases, type = "incidence")
est_sec <- estimate_secondary(cases,
  obs = obs_opts(scale = list(mean = 0.2, sd = 0.2), week_effect = FALSE)
)

trunc <- dist_spec(
  mean = convert_to_logmean(3, 2),
  mean_sd = 0.1,
  sd = convert_to_logsd(3, 2),
  sd_sd = 0.1,
  max = 10
)
construct_truncation <- function(index, cases, dist) {
  set.seed(index)
  if (dist$dist == 0) {
    dfunc <- dlnorm
  } else {
    dfunc <- dgamma
  }
  cmf <- cumsum(
    dfunc(
      1:(dist$max + 1),
      rnorm(1, dist$mean_mean, dist$mean_sd),
      rnorm(1, dist$sd_mean, dist$sd_sd)
    )
  )
  cmf <- cmf / cmf[dist$max + 1]
  cmf <- rev(cmf)[-1]
  trunc_cases <- data.table::copy(cases)[1:(.N - index)]
  trunc_cases[
    (.N - length(cmf) + 1):.N, confirm := as.integer(confirm * cmf)
  ]
  return(trunc_cases)
}
example_data <- purrr::map(c(20, 15, 10, 0),
  construct_truncation,
  cases = reported_cases,
  dist = trunc
)
est_trunc <- estimate_truncation(example_data,
  verbose = interactive(),
  chains = 2, iter = 2000
)

names(est_inf)
#> [1] "samples"      "summarised"   "fit"          "args"         "observations"
names(est_sec)
#> [1] "predictions" "posterior"   "data"        "fit"
names(est_trunc)
#> [1] "dist"     "obs"      "last_obs" "cmf"      "data"     "fit"

^{Created on 2023-09-08 with reprex v2.0.2}

[seabbs]

Yes, I agree. This is a nice idea. Helps make a point that the models are really subsets/supersets of each other as well.

[jamesmbaazam]

I'm working on this and have a few questions:

and observations as everything else, I think, could be constructed from these three elements using appropriate summary or extract_... methods.

Are all three functions going to return the unformatted stanfit object, which will then be formatted downstream using the extract() and summary() methods?

Is observations supposed to be only the raw data passed through the reported_cases/reports/obs arguments?

estimate_truncation() returns a dist object that's passed on to the other two estimate_*() functions. Should this function be returning that as a fourth element? Alternatively, the truncation distribution can be constructed via a construct_trunc_dist() helper using the returned fit object.