R/optimal_model.R

In JMbayes2: Extended Joint Models for Longitudinal and Time-to-Event Data

opt_model <- function (models, newdata, t0, parallel = "snow", cores = 1L) {
    if (!all(sapply(models, function (obj) inherits(obj, "jm")))) {
        stop("all objects in 'models' must inherit from class 'jm'.\n")
    }
    object <- models[[1L]]
    id_var <- object$model_info$var_names$idVar
    time_var <- object$model_info$var_names$time_var
    Time_var <- object$model_info$var_names$Time_var
    TTime_var <- if (length(Time_var) > 1) Time_var[2L] else Time_var
    event_var <- object$model_info$var_names$event_var
    if (!is.data.frame(newdata)) {
        stop("'newdata' must be a data.frame.\n")
    }
    if (!all((c(id_var, time_var) %in% names(newdata)))) {
        stop("'newdata' must contain the variables '", id_var, "' and '",
             time_var, "'.\n")
    }
    ND <- newdata
    ND[[id_var]] <- match(ND[[id_var]], unique(ND[[id_var]]))
    ND_before <- ND[ND[[time_var]] <= t0, ]
    ND_before[[TTime_var]] <- t0
    ND_before[[event_var]] <- 0
    ND_after <- ND[ND[[time_var]] > t0, ]
    data_after_t0 <- nrow(ND_after)
    ###
    MISE_perid <- function (obs, reps, id, times) {
        trapezoid_rule <- function (f, x) {
            sum(0.5 * diff(x) * (f[-length(x)] + f[-1L]))
        }
        loess.smooth2 <- function (x, y) {
            loess.smooth(x, y, span = 0.75, degree = 1,
                         family = "gaussian", evaluation = 200)
        }
        smooth <- function (x, y) {
            n <- length(x)
            if (n >= 4) {
                loess.smooth2(x, y)
            } else if (n > 1 && n < 4) {
                spline(x, y)
            } else {
                list(x = NA_real_, y = NA_real_)
            }
        }
        gof_fun <- function (y, times, id, type) {
            if (type == "variogram") {
                ls <- smooth(times, y)
                ind <- findInterval(times, ls$x)
                rr <- y - ls$y[ind]
                variogram(rr, times, id)[[1L]]
            } else {
                cbind(times, y)
            }
        }
        Obs_ave <- gof_fun(obs, times, id, "average")
        Obs_vario <- gof_fun(obs, times, id, "variogram")
        G_obs_ave <- smooth(Obs_ave[, 1L], Obs_ave[, 2L])
        F_obs_ave <- mapply(smooth, y = split(Obs_ave[, 2L], id),
                            x = split(Obs_ave[, 1L], id), SIMPLIFY = FALSE)
        G_obs_vario <- smooth(Obs_vario[, 1L], Obs_vario[, 2L])
        ni <- tapply(id, id, length)
        id_long <- rep(unique(id), sapply(ni, function (n) if (n < 2) 1 else
            ncol(combn(n, 2))))
        F_obs_vario <- mapply(smooth, y = split(Obs_vario[, 2L], id_long),
                              x = split(Obs_vario[, 1L], id_long),
                              SIMPLIFY = FALSE)
        mise <- function (obs, rep) {
            trapezoid_rule((obs$y - rep$y)^2, obs$x)
        }
        n <- length(unique(id))
        M <- ncol(reps)
        MISE_mod_ave <- MISE_mod_vario <- numeric(M)
        MISE_ave <- MISE_vario <- matrix(0.0, n, M)
        for (m in seq_len(M)) {
            reps_ave <- gof_fun(reps[, m], times, id, "average")
            reps_vario <- gof_fun(reps[, m], times, id, "variogram")
            G_reps_ave <- smooth(reps_ave[, 1L], reps_ave[, 2L])
            F_reps_ave <-
                mapply(smooth, y = split(reps_ave[, 2L], id),
                       x = split(reps_ave[, 1L], id), SIMPLIFY = FALSE)
            G_reps_vario <- smooth(reps_vario[, 1L], reps_vario[, 2L])
            F_reps_vario <-
                mapply(smooth, y = split(reps_vario[, 2L], id_long),
                       x = split(reps_vario[, 1L], id_long), SIMPLIFY = FALSE)
            MISE_mod_ave[m] <- mise(G_obs_ave, G_reps_ave)
            MISE_ave[, m] <- mapply(mise, obs = F_obs_ave, rep = F_reps_ave)
            MISE_mod_vario[m] <- mise(G_obs_vario, G_reps_vario)
            MISE_vario[, m] <- mapply(mise, obs = F_obs_vario, rep = F_reps_vario)
        }
        list(MISE_ave = MISE_ave, MISE_vario = MISE_vario,
             MISE_mod_ave = MISE_mod_ave, MISE_mod_vario = MISE_mod_vario)
    }
    MISE_model <- function (object) {
        prs <- predict(object, newdata = ND_before, newdata2 = ND_after,
                       return_params_mcmc = TRUE)
        sims <- simulate(object, nsim = 200L, newdata = ND_before,
                          include_outcome = TRUE, random_effects = "mcmc",
                          params_mcmc = prs$newdata$params_mcmc)
        n_outcomes <- length(sims[["outcome"]])
        id <- attr(sims$outcome[[1]], "id")
        MISE_mod_ave <- MISE_mod_vario <- matrix(0.0, 200L, n_outcomes)
        MISEs_ave <- MISEs_vario <- array(0, c(length(unique(id)), 200L, n_outcomes))
        dimnames(MISEs_ave) <- dimnames(MISEs_vario) <-
            list(unique(newdata[[id_var]]), colnames(sims[[1]]),
                 names(sims)[seq_len(n_outcomes)])
        for (m in seq_len(n_outcomes)) {
            outcome <- sims$outcome[[m]]
            mises <- MISE_perid(outcome, sims[[m]], attr(outcome, "id"),
                                attr(outcome, "times"))
            MISE_mod_ave[, m] <- mises$MISE_mod_ave
            MISEs_ave[, , m] <- mises$MISE_ave
            MISE_mod_vario[, m] <- mises$MISE_mod_vario
            MISEs_vario[, , m] <- mises$MISE_vario
        }
        list(MISEs_ave = MISEs_ave, MISEs_vario = MISEs_vario, Preds = prs,
             MISE_mod_ave = MISE_mod_ave, MISE_mod_vario = MISE_mod_vario)
    }
    ###
    if (cores > 1L && length(models) > 1L) {
        have_mc <- have_snow <- FALSE
        if (parallel == "multicore") {
            have_mc <- .Platform$OS.type != "windows"
        } else if (parallel == "snow") {
            have_snow <- TRUE
        }
        if (!have_mc && !have_snow) cores <- 1L
        loadNamespace("parallel")
        if (have_mc) {
            out <- parallel::mclapply(models, MISE_model, mc.cores = cores)
        } else {
            cl <- parallel::makePSOCKcluster(rep("localhost", cores))
            invisible(parallel::clusterEvalQ(cl, library("JMbayes2")))
            out <- parallel::parLapply(cl, models, MISE_model)
            parallel::stopCluster(cl)
        }
    } else {
        out <- lapply(models, MISE_model)
    }
    out <- lapply(out, function (x) {
        x$std_MISEs_ave <- x$MISEs_ave
        x$std_MISEs_vario <- x$MISEs_vario
        x
    })
    n_outcomes <- dim(out[[1]][["MISEs_ave"]])[3L]
    for (m in seq_len(n_outcomes)) {
        vals <- do.call("c", lapply(out, function (x) x$MISEs_ave[, , m]))
        out[] <- lapply(out, function (x) {
            x$std_MISEs_ave[, , m] <- (x$MISEs_ave - mean(vals)) / sd(vals)
            x
        })
        vals <- do.call("c", lapply(out, function (x) x$MISEs_vario[, , m]))
        out[] <- lapply(out, function (x) {
            x$std_MISEs_vario[, , m] <- (x$MISEs_vario - mean(vals)) / sd(vals)
            x
        })
    }
    out[] <- lapply(out, function (x) {
        x$MISE_mod_ave <- colMeans(x$MISE_mod_ave, na.rm = TRUE)
        x$MISEs_ave <- apply(x$MISEs_ave, c(1, 3), mean, na.rm = TRUE)
        x$std_MISEs_ave <- apply(x$std_MISEs_ave, c(1, 3), mean, na.rm = TRUE)
        x$MISE_mod_vario <- colMeans(x$MISE_mod_vario, na.rm = TRUE)
        x$MISEs_vario <- apply(x$MISEs_vario, c(1, 3), mean, na.rm = TRUE)
        x$std_MISEs_vario <- apply(x$std_MISEs_vario, c(1, 3), mean, na.rm = TRUE)
        x
    })
    out
}

#object = Models[[4]]
#newdata = testing[testing$time < T0, ]
#newdata2 = testing[testing$time >= T0, ]

dmvnorm <- function (x, mu, Sigma = NULL, invSigma = NULL, log = FALSE,
                     prop = FALSE) {
    if (!is.matrix(x))
        x <- rbind(x)
    if (is.matrix(mu)) {
        if (nrow(mu) != nrow(x))
            stop("incorrect dimensions for 'mu'.")
        p <- ncol(mu)
    } else {
        p <- length(mu)
        mu <- rep(mu, each = nrow(x))
    }
    if (is.null(Sigma) && is.null(invSigma))
        stop("'Sigma' or 'invSigma' must be given.")
    if (!is.null(Sigma)) {
        if (is.list(Sigma)) {
            ev <- Sigma$values
            evec <- Sigma$vectors
        }
        else {
            ed <- eigen(Sigma, symmetric = TRUE)
            ev <- ed$values
            evec <- ed$vectors
        }
        invSigma <- evec %*% (t(evec)/ev)
        if (!prop)
            logdetSigma <- sum(log(ev))
    } else {
        if (!prop)
            logdetSigma <- -determinant(as.matrix(invSigma))$modulus
    }
    ss <- x - mu
    quad <- 0.5 * rowSums((ss %*% invSigma) * ss)
    if (!prop)
        fact <- -0.5 * (p * log(2 * pi) + logdetSigma)
    if (log) {
        if (!prop)
            as.vector(fact - quad)
        else as.vector(-quad)
    } else {
        if (!prop)
            as.vector(exp(fact - quad))
        else as.vector(exp(-quad))
    }
}

npar <- function (object, newdata, type = c("marginal", "conditional")) {
    type <- match.arg(type)
    if (type == "marginal") {
        TermsX <- object$terms
        mfX <- model.frame(TermsX, object$data)
        X <- model.matrix(TermsX, mfX)
        formYz <- formula(object$modelStruct$reStruct[[1]])
        mfZ <- model.frame(terms(formYz), data = object$data)
        TermsZ <- attr(mfZ, "terms")
        Z <- model.matrix(TermsZ, mfZ)
        nx <- qr(model.matrix(TermsX, mfX))$rank
        sigma <- object$sigma
        D <- lapply(pdMatrix(object$modelStruct$reStruct), "*", sigma^2)[[1]]
        D <- unique(D[lower.tri(D, TRUE)])
        D <- D[D > 1e-06]
        nz <- length(D)
        nx + nz + 1
        #qr(cbind(X, Z))$rank + 1
    } else {
        TermsX <- object$terms
        formYz <- formula(object$modelStruct$reStruct[[1]])
        mfZ <- model.frame(terms(formYz), data = object$data)
        TermsZ <- attr(mfZ, "terms")
        sigma <- object$sigma
        D <- lapply(pdMatrix(object$modelStruct$reStruct), "*", sigma^2)[[1]]
        #######
        all_vars <- unique(c(all.vars(TermsX), all.vars(TermsZ)))
        newdata_nomiss <- newdata[complete.cases(newdata[all_vars]), ]
        mfX_new <- model.frame(TermsX, data = newdata_nomiss)
        X_new <- model.matrix(TermsX, mfX_new)
        mfZ_new <- model.frame(TermsZ, data = newdata_nomiss)
        Z_new <- model.matrix(TermsZ, mfZ_new)
        XZ <- cbind(X_new, Z_new)
        V11 <- crossprod(X_new)
        V22 <- crossprod(Z_new) + sigma^2 * solve(D)
        VV <- nearPD(
            rbind(cbind(V11, crossprod(X_new, Z_new)),
                  cbind(crossprod(Z_new, X_new), V22)))
        sum(diag(XZ %*% solve(VV, t(XZ))))
    }
}

total_var <- function (ys) {
    n <- length(ys)
    s <- 0
    count <- 0
    f <- function(x, y) 0.5 * (x - y)^2
    for (i in seq_len(n-1)) {
        for (j in seq(i + 1, n)) {
            r <- outer(ys[[i]], ys[[j]], f)
            s <- s + sum(r)
            count <- count + length(r)
        }
    }
    s / count
}

#object = tt[[2]]
#newdata = training[training$time <= 3, ]
#newdata2 = training[training$time > 3, ]

AIC_lme <- function (object, newdata) {
    TermsX <- object$terms
    formYz <- formula(object$modelStruct$reStruct[[1]])
    mfZ <- model.frame(terms(formYz), data = object$data)
    TermsZ <- attr(mfZ, "terms")
    id_var <- names(object$modelStruct$reStruct)
    betas <- fixef(object)
    sigma <- object$sigma
    D <- lapply(pdMatrix(object$modelStruct$reStruct), "*", sigma^2)[[1]]
    all_vars <- unique(c(all.vars(TermsX), all.vars(TermsZ)))
    newdata <- newdata[complete.cases(newdata[all_vars]), ]
    mfX_new <- model.frame(TermsX, data = newdata)
    X_new <- model.matrix(TermsX, mfX_new)
    mfZ_new <- model.frame(TermsZ, data = newdata)
    Z_new <- model.matrix(TermsZ, mfZ_new)
    y_new <- model.response(mfX_new, "numeric")
    id_new <- newdata[[id_var]]
    unq_id <- unique(id_new)
    id_new <- match(id_new, unq_id)
    n <- length(unq_id)
    lL <- numeric(n)
    for (i in seq_len(n)) {
        id_i <- id_new == i
        y_new_i <- y_new[id_i]
        X_new_i <- X_new[id_i, , drop = FALSE]
        Z_new_i <- Z_new[id_i, , drop = FALSE]
        mu_i <- c(X_new_i %*% betas)
        V_i <- Z_new_i %*% tcrossprod(D, Z_new_i) + sigma^2 * diag(sum(id_i))
        lL[i] <- dmvnorm(y_new_i, mu_i, V_i, log = TRUE)
    }
    nx <- qr(X_new)$rank
    D <- unique(D[lower.tri(D, TRUE)])
    D <- D[D > 1e-06]
    nz <- length(D)
    - 2 * sum(lL, na.rm = TRUE) + 2 * (nx + nz + 1)
}

indv_pred_lme <- function (object, newdata, newdata2,
                          interval = c("none", "confidence", "prediction"),
                          M = 200, level = 0.95) {
    if (!inherits(object, "lme") && !inherits(object, "lmeComponents")) {
        stop("Use only with 'lme' or 'lmeComponents' objects.\n")
    }
    interval <- match.arg(interval)
    if (inherits(object, "lme")) {
        TermsX <- object$terms
        formYz <- formula(object$modelStruct$reStruct[[1]])
        mfZ <- model.frame(terms(formYz), data = object$data)
        TermsZ <- attr(mfZ, "terms")
        id_var <- names(object$modelStruct$reStruct)
        betas <- fixef(object)
        sigma <- object$sigma
        D <- lapply(pdMatrix(object$modelStruct$reStruct), "*", sigma^2)[[1]]
        V <- vcov(object)
    } else {
        TermsX <- object$TermsX
        TermsZ <- object$TermsZ
        id_var <- object$idVar
        betas <- object$betas
        sigma <- object$sigma
        D <- object$D
        V <- object$V
    }
    all_vars <- unique(c(all.vars(TermsX), all.vars(TermsZ)))
    newdata_nomiss <- newdata[complete.cases(newdata[all_vars]), ]
    mfX_new <- model.frame(TermsX, data = newdata_nomiss)
    X_new <- model.matrix(TermsX, mfX_new)
    mfZ_new <- model.frame(TermsZ, data = newdata_nomiss)
    Z_new <- model.matrix(TermsZ, mfZ_new)
    y_new <- model.response(mfX_new, "numeric")
    if (length(id_var) > 1)
        stop("the current version of the function only works with a single grouping variable.\n")
    if (is.null(newdata[[id_var]]))
        stop("subject id variable not in newdata.")
    id_nomiss <- match(newdata_nomiss[[id_var]], unique(newdata_nomiss[[id_var]]))
    n <- length(unique(id_nomiss))
    modes <- matrix(0, n, ncol(Z_new))
    chol_post_vars <- DZtVinv <- vector("list", n)
    for (i in seq_len(n)) {
        id_i <- id_nomiss == i
        y_new_id <- y_new[id_i]
        X_new_id <- X_new[id_i, , drop = FALSE]
        Z_new_id <- Z_new[id_i, , drop = FALSE]
        Vi_inv <- solve(Z_new_id %*% tcrossprod(D, Z_new_id) +
                            sigma^2 * diag(sum(id_i)))
        DZtVinv[[i]] <- tcrossprod(D, Z_new_id) %*% Vi_inv
        mu_id <- c(X_new_id %*% betas)
        modes[i, ] <- c(DZtVinv[[i]] %*% (y_new_id - mu_id))
        t1 <- DZtVinv[[i]] %*% Z_new_id %*% D
        t2 <- DZtVinv[[i]] %*% X_new_id %*% V %*%
            crossprod(X_new_id, Vi_inv) %*% Z_new_id %*% D
        chol_post_vars[[i]] <- chol(D - t1 + t2)
    }
    fitted_y <- c(X_new %*% betas) +
        rowSums(Z_new * modes[id_nomiss, , drop = FALSE])
    #####
    newdata2_nomiss <- newdata2[complete.cases(newdata2[all_vars]), ]
    mfX_new2 <- model.frame(TermsX, data = newdata2_nomiss)
    X_new2 <- model.matrix(TermsX, mfX_new2)
    mfZ_new2 <- model.frame(TermsZ, data = newdata2_nomiss)
    Z_new2 <- model.matrix(TermsZ, mfZ_new2)
    id_nomiss2 <- match(newdata2_nomiss[[id_var]], unique(newdata2_nomiss[[id_var]]))
    predicted_y <- c(X_new2 %*% betas) +
        rowSums(Z_new2 * modes[id_nomiss2, , drop = FALSE])
    out <- list(fitted_y = fitted_y, predicted_y = predicted_y)
    if (interval != "none") {
        betas_M <- MASS::mvrnorm(M, betas, V)
        modes_fun <- function (betas) {
            v <- mapply("%*%", DZtVinv, split(y_new - X_new %*% betas, id_nomiss))
            if (ncol(Z_new) > 1L) t(v) else as.matrix(v)
        }
        modes_M <- lapply(split(betas_M, row(betas_M)), modes_fun)
        n_pred <- length(predicted_y)
        sampled_y <- matrix(0, n_pred, M)
        ncz <- ncol(Z_new)
        for (m in seq_len(M)) {
            betas_m <- betas_M[m, ]
            modes_m <- modes_M[[m]]
            b_m <- matrix(rnorm(n * ncz), n, ncz)
            for (i in seq_len(n)) {
                b_m[i, ] <- b_m[i, ] %*% chol_post_vars[[i]]
            }
            b_m <- modes_m + b_m
            mean_m <- c(X_new2 %*% betas_m) +
                rowSums(Z_new2 * b_m[id_nomiss2, , drop = FALSE])
            sampled_y[, m] <- if (interval == "confidence") {
                mean_m
            } else {
                rnorm(n_pred, mean_m, sigma)
            }
        }
        low <- apply(sampled_y, 1L, quantile, probs = (1 - level) / 2)
        upp <- apply(sampled_y, 1L, quantile, probs = 1 - (1 - level) / 2)
        out <- c(out, list(low = low, upp = upp))
    }
    out
}

cor2cov <- function (R, vars, sds = NULL) {
    p <- nrow(R)
    if (is.null(sds)) sds <- sqrt(vars)
    sds * R * rep(sds, each = p)
}