R/ELBO.R

In stephenslab/susiF.alpha: Sum of Single Functions

#' @title Compute KL divergence effect l
#
#' @param obj a susisF object defined by  init_susiF_obj  function
#
#' @param l integer larger or equal to 1. Corresponds to the effect to be accessed
#
#
#' @param X matrix of covariate
#
#' @param D matrix of wavelet D coefficients from the original input data (Y)
#
#' @param C vector of wavelet scaling coefficient from the original input data (Y)
#
#' @param indx_lst list generated by gen_wavelet_indx for the given level of resolution
#
#' @return The KL divergence for effect l
#' @export
#' @keywords internal
cal_KL_l <- function(obj, l, X, D,C , indx_lst, ...)
  UseMethod("cal_KL_l")




#' @rdname cal_KL_l
#
#' @method cal_KL_l susiF
#
#' @export cal_KL_l.susiF
#
#' @export
#' @keywords internal
cal_KL_l.susiF <- function(obj, l, X, D, C , indx_lst, ...)
{

  R_l <- cal_partial_resid(
 
                           obj = obj,
 
                           l         =  (l-1),
                           X         =  X,
                           D         =  D,
                           C         =  C,
                           indx_lst  = indx_lst
                          )


  out <-    - loglik_SFR(obj, l,R_l,X,indx_lst)- loglik_SFR_post(obj, l,R_l,X)


  return(out)
}





#' @title Compute log likelihood of single function regression of effect l
#
#' @param obj a susiF object defined by  init_susiF_obj  function
#
#' @param l effect to update
#
#' @param Y wavelet transformed phenotype
#
#' @param X Matrix of covariates
#' @param indx_lst list generated by gen_wavelet_indx for the given level of resolution
#
#
#' @return The log-likelihood, \eqn{\log p(Y | X, V)}, where V is the prior parameters
#' @export
#' @keywords internal
loglik_SFR <- function    (obj, l,  ...)
  UseMethod("loglik_SFR")

#' @rdname loglik_SFR
#
#' @method loglik_SFR susiF
#
#' @importFrom stats dnorm
#
#' @export loglik_SFR.susiF
#
#' @export
#' @keywords internal
loglik_SFR.susiF <- function (obj, l, Y , X, indx_lst, ...)
{
  lBF            <- get_lBF(obj,l)
  prior_weights  <- rep(1/ncol(X),ncol(X))
  maxlBF         <- max(lBF)
  w              <- exp(lBF - maxlBF)
  w_weighted     <- w * prior_weights
  weighted_sum_w <- sum(w_weighted)
  lBF_model      <- maxlBF + log(weighted_sum_w)
  return(lBF_model + sum(dnorm(Y,0,sd = sqrt(obj$sigma2),log = TRUE)))
}


#' @title Compute posterior expected loglikelihood for single function
#   regression of effect l
#
#' @param obj susiF object, for example created by calling
#'  init_susiF_obj .
#
#' @param l Index of effect to update.
#
#' @param \dots Other arguments.
#
#' @return Posterior expected log-likelihood for the single-function
#   regression of effect l.
#
#' @export
#' @keywords internal
loglik_SFR_post <- function (obj, l, ...)
  UseMethod("loglik_SFR_post")

#' @rdname loglik_SFR_post
#
#' @param Y Wavelet-transformed phenotype.
#
#' @param X Matrix of covariates.
#
#' @method loglik_SFR_post susiF
#
#' @export loglik_SFR_post.susiF
#
#' @export
#' @keywords internal
loglik_SFR_post.susiF <- function (obj, l, Y, X, ...)
{
  n   <- nrow(Y)
  t   <- ncol(Y)
  EF  <- get_post_F(obj,l)
  EF2 <- get_post_F2(obj,l)
  s2  <- obj$sigma2
  return(-n*t/2*log(2*pi*s2)
         - s2/2*(sum(t(Y)%*%Y) - 2*sum(t(Y)%*%X%*%EF) + sum(t(EF2)%*%EF2)))
}




#' @title Expected log likelihood for a   susiF   object
#
#' @param obj a susiF object defined by  init_susiF_obj  function
#
#' @param Y Matrix of outcomes
#
#' @param X Matrix of covariates 
#' 
#' @param \dots Other arguments.
#
#' @return Expected log likelihood
#' @export
#' @keywords internal
Eloglik <- function    (obj,   ...)
  UseMethod("Eloglik")


#' @rdname Eloglik
#
#' @method Eloglik susiF
#
#' @export Eloglik.susiF
#' @export
#' @keywords internal
Eloglik.susiF = function (obj,Y ,X,  ...) {
  n = nrow(Y)
  t = ncol(Y)

  return(-(n*t/2) * log(2*pi*obj$sigma2) - (1/(2*obj$sigma2)) * get_ER2( obj, Y, X))
}
#  @export
#  @keywords internal
#Eloglik.susiF_ss = function (obj,data,...) {
#  n <- data$N
#  t <- ncol(data$Bhat)
#
#  return(-(n*t/2) * log(2*pi*obj$sigma2) - (1/(2*obj$sigma2)) * get_ER2( obj,data))
#}

#' @title Get objective function from data and susiF object
#
#' @param obj a susiF object defined by  init_susiF_obj  function
#
#' @param Y Matrix of outcomes
#
#' @param X matrix of covariate
#
#' @param D matrix of wavelet D coefficients from the original input data (Y)
#
#' @param C vector of wavelet scaling coefficient from the original input data (Y)
#
#' @param indx_lst list generated by gen_wavelet_indx for the given level of resolution
#
#' @return objective function value
#
#' @export
#' @keywords internal
get_objective <- function    (obj,  Y, X, D, C , indx_lst,  ...)
  UseMethod("get_objective")


#' @rdname get_objective
#
#' @method get_objective susiF
#
#' @export get_objective.susiF
#' @export
#' @keywords internal
get_objective.susiF <- function    (obj, Y, X, D, C , indx_lst,  ...)
{
#print(obj$KL)
  out <-  Eloglik(obj, Y, X)  + sum(obj$KL)
  return(out)

}