R/check_tpar_depends.R

In fintzij/stemr: Stochastic Epidemic Model Simulation and Estimation

#' Check if time varying parameters depend on initial conditions
#'
#' @param tparam list of time varying parameters
#' @param initdist_objects list with initial distribution objects
#' @param parmat parameter matrix
#'
#' @return time-varying parameter list with indicators for whether each
#'   parameter depends on initial conditions
#' @export
check_tpar_depends = function(tparam, initdist_objects, parmat) {

    # check whether initial conditions are fixed
    if(all(sapply(initdist_objects, "[[", "fixed"))) {

        for(s in seq_along(tparam)) {
            tparam[[s]]$init_dep = FALSE
        }

    } else {
        # copy of parameters from matrix
        pars <- parmat[1,]

        # jitter initial conditions
        for(s in seq_along(initdist_objects)) {
            if(!initdist_objects[[s]]$fixed) {

                # new draws
                draws = rnorm(initdist_objects[[s]]$draws_cur)


                if (stem_object$dynamics$initializer[[s]]$dist == "sbln") {
                  orig <- sbln_normal_to_volume(normal_draws = draws,
                                                stick_means = initdist_objects[[s]]$comp_prior[1:length(initdist_objects[[s]]$draws_cur)],
                                                stick_sds = initdist_objects[[s]]$comp_prior[(length(initdist_objects[[s]]$draws_cur) + 1):length(initdist_objects[[s]]$comp_prior)],
                                                stick_size = initdist_objects[[s]]$comp_size)
                } else {
                  orig <- initdist_objects[[s]]$comp_mean + c(initdist_objects[[s]]$comp_sqrt_cov %*% draws)
                }

                # calculate compartment counts
                copy_vec2(dest = pars,
                          orig = orig,
                          inds = initdist_objects[[s]]$param_inds_Cpp)

                # check boundary conditions
                while(any(pars[initdist_objects[[s]]$param_inds_R] < 0) |
                      any(pars[initdist_objects[[s]]$param_inds_R] >
                          initdist_objects[[s]]$comp_size)) {

                    # redraw and recalculate if needed
                    draws = rnorm(initdist_objects[[s]]$draws_cur)

                    copy_vec2(dest = pars,
                              orig = c(initdist_objects[[s]]$comp_mean +
                                           initdist_objects[[s]]$comp_sqrt_cov %*% draws),
                              inds = initdist_objects[[s]]$param_inds_Cpp)
                }
            }
        }

        # now check whether the new time-varying parameters have changed
        for(s in seq_along(tparam)) {

            tp = tparam[[s]]$draws2par(
                    parameters = pars,
                    draws = tparam[[s]]$draws_cur)[tparam[[s]]$tpar_inds_R]

            tparam[[s]]$init_dep =
                !all.equal(tp, tparam[[s]]$tpar_cur)
        }
    }

    return(tparam)
}