R/tmle3_Spec_SET.R

In yulun-rayn/tmle3trans: What the Package Does (One Line, Title Case)

#' Defines a TMLE with Survival Effect Transportation
#'
#'
#' @importFrom R6 R6Class
#' @importFrom assertthat assert_that
#' @importFrom tmle3 tmle3_Spec define_lf tmle3_Update Targeted_Likelihood
#'
#' @export
#

tmle3_Spec_SET <- R6Class(
  classname = "tmle3_Spec_SET",
  portable = TRUE,
  class = TRUE,
  inherit = tmle3_Spec,
  public = list(
    initialize = function(intervention, onsite = 1, offsite = 0, target_times = NULL, 
                          fit_s_marginal = "empirical", ...) {
      super$initialize(intervention = intervention, 
                       onsite = onsite, offsite = offsite, 
                       target_times = target_times,
                       fit_s_marginal = fit_s_marginal, ...)
    },
    make_tmle_task = function(data, node_list, ...) {
      variable_types <- self$options$variable_types
      
      tmle_task <- survival_task(data, node_list, survival_e_npsem, variable_types)
      
      return(tmle_task)
    },
    
    make_initial_likelihood = function(tmle_task, learner_list = NULL) {
      likelihood <- survival_e_likelihood(tmle_task, learner_list)
      return(likelihood)
    },
    
    make_params = function(tmle_task, likelihood) {
      tmle_params <- define_param(Param_SET, likelihood, 
                                  intervention = self$options$intervention,
                                  onsite = self$options$onsite,
                                  offsite = self$options$offsite,
                                  target_times = self$options$target_times,
                                  fit_s_marginal = self$options$fit_s_marginal)
      return(tmle_params)
    }
  ),
  active = list(),
  private = list()
)

#' Wrapper
#'
#' @importFrom sl3 make_learner Lrnr_mean
#' @param onsite value for onsite
#' @param offsite value for offsite
#' @export
tmle_SET <- function(intervention, onsite = 1, offsite = 0, target_times = NULL, 
                     fit_s_marginal = "empirical", ...) {
  tmle3_Spec_SET$new(intervention, onsite = onsite, offsite = offsite,
                     target_times = target_times,
                     fit_s_marginal = fit_s_marginal, ...)
}