R/gammmbest_fit.R

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

#' @title Title
#'
#' @description
#' Description.
#'
#' @param x A design matrix.
#' @param y A vector of response observation.
#' @param group A character vector describing grouping of random effects.
#' @param fixef_index An integer vector specifying which columns of \code{x}
#'        have fixed effects coefficients.
#' @param ranef_index An integer vector specifying which columns of \code{x}
#'        have random effects coefficients.
#' @param cov_supp An integer matrix encoding weight sharing in the prior
#'        prior's covariance matrix.
#' @param alpha The elasticnet mixing parameter, with
#'        \eqn{0 \le \alpha \le 1}. The penalty is defined as
#'        \deqn{(1-\alpha)/2 ||\Gamma \beta||_2^2 + \alpha ||\Gamma||_1}
#'        where \eqn{\Sigma^{-1} = \Gamma^T \Gamma} and \eqn{\Sigma} is the
#'        prior covariance.
#' @return An object of class \code{gammmbest} containing
#'   \item{beta}{The estimated coefficients for the fixed effects.}
#'   \item{dispersion}{The estimated dispersion of the model.}
#'   \item{prior_cov}{The estimated covariance of the random effects prior.}
#'   \item{ranef}{The posterior estimate of the random effects coefficients.}
#'
#' @details
#' hello said kyle
#'
#' @examples
#' print("hello world")
#' @export
gammmbest_fit <- function(x, y, group, family, fixef_index, ranef_index,
		cov_supp = NULL, cov_weight = NULL, eta_precision = 0.01, alpha = 0,
		n_loops = 2) {

    assertthat::are_equal(length(x), length(y))
    assertthat::are_equal(length(x), length(group))

    len_ran <- length(ranef_index)
    if (is.null(cov_supp)) {
        cov_supp <- matrix(NA, len_ran, len_ran)
        cov_supp[lower.tri(cov_supp, diag = TRUE)] <- seq_len(choose(len_ran + 1, 2))
        cov_supp[upper.tri(cov_supp)] <- cov_supp[lower.tri(cov_supp)]
    }

    if (is.null(cov_weight)) {
        cov_weight <- matrix(1, len_ran, len_ran)
    }

	x <- split.data.frame(x, group)
	y <- split(y, group)

	xz_svd_list <- compact_svd_list(x, fixef_index, ranef_index)
	x <- NULL

	eta_list <- estimate_eta_list(xz_svd_list, y, family, eta_precision)
	dispersion <- estimate_dispersion(xz_svd_list, y, eta_list, family)
    prior_cov <- diag(length(ranef_index))

    for (i in seq_len(n_loops)) {
    	w_list <- initialize_weights(xz_svd_list, prior_cov)
    	beta <- estimate_beta(xz_svd_list, eta_list, w_list)
    	prior_cov <- estimate_cov(xz_svd_list, eta_list, beta, dispersion,
    	                          cov_supp, cov_weight, w_list)
    }
	ranef <- estimate_ranef(xz_svd_list, y, beta, prior_cov, family, alpha)

	structure(
	    list(beta       = beta,
    	     dispersion = dispersion,
    	     prior_cov  = prior_cov,
    	     ranef      = ranef
    	), class = "gammmbest")
}