R/estimate_cov.R

In kschmaus/gammmbest: Method Of Moments Based Estimation for Generalized Additive Mixed Models

#' Estimate the covariance of the random effect prior distribution.
#'
#' See (5) and (6) of Perry 2015.
#'
#' @param xz_svd_list A list of compact singular vector decomposition objects.
#' @param eta_list A list of numeric vectors encoding fixed and random effects.
#' @param beta A numeric vector estiamte of fixed effects.
#' @param dispersion a scalar estimate of dispersion
#' @param cov_supp A square integer matrix, with possible missing values,
#' see \code{\link{gammmbest_fit}}.
#' @return A non-negative definite matrix.
estimate_cov <- function(xz_svd_list, eta_list, beta, dispersion, cov_supp,
                         cov_weight, w_list) {
    A_vec <- create_A_vec(xz_svd_list, eta_list, beta, w_list)
    B_vec <- create_B_vec(xz_svd_list, w_list)
    omega2 <- create_omega2(xz_svd_list, cov_supp, cov_weight, w_list)

    prior_cov_hat_vec <- MASS::ginv(omega2) %*% (A_vec - dispersion * B_vec)
    prior_cov_hat <- prior_cov_hat_vec[as.vector(cov_supp)] *
                     as.vector(cov_weight)
    prior_cov_hat <- ifelse(is.na(prior_cov_hat), 0, prior_cov_hat)
    dim(prior_cov_hat) <- dim(cov_supp)

    prior_cov <- Matrix::nearPD(prior_cov_hat)$mat
}

create_A_vec <- function(xz_svd_list, eta_list, beta, w_list) {
    V1_list <- Map(getElement, xz_svd_list, "V1")
    V2_list <- Map(getElement, xz_svd_list, "V2")

    f <- function(V1, V2, eta, W, beta)
        tcrossprod(V2 %*% W %*% (eta - t(V1) %*% beta))

    A_mat <-  map_reduce_mean(f, V1_list, V2_list, eta_list, w_list, list(beta))

    as.vector(A_mat)
}

create_B_vec <- function(xz_svd_list, w_list) {
    V2_list <- Map(getElement, xz_svd_list, "V2")
    D_list <- Map(getElement, xz_svd_list, "D")

    f <- function(V2, D, W)
        tcrossprod(V2 %*% W %*% solve(D))

    B_mat <-  map_reduce_mean(f, V2_list, D_list, w_list)

    as.vector(B_mat)
}

create_omega2 <- function(xz_svd_list, cov_supp, cov_weight, w_list) {
    V2_list <- Map(getElement, xz_svd_list, "V2")

    f <- function(V2, W) {
        x <- V2 %*% W %*% t(V2)
        kronecker(x, x)
    }
	omega2 <- map_reduce_mean(f, V2_list, w_list)

	g <- function(x, index, i, weights) {
        x <- x[, which(index == i), drop = FALSE]
        w <- diag(weights[which(index == i)])

        rowSums(x %*% diag(w))
    }

    index <- as.vector(cov_supp)
    unique_index <- na.exclude(unique(index))
    weights <- as.vector(cov_weight)

    mapply(g, list(omega2), list(index), unique_index, list(weights))
}