R/mix.R

In ashr: Methods for Adaptive Shrinkage, using Empirical Bayes

################################## GENERIC FUNCTIONS ############################

#' Generic function of calculating the component densities of the
#'     mixture
#' @param m mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @param y is an n-vector of location
#' @param log whether to use log-scale on densities
#' @return A k by n matrix of densities
#' @export
comp_dens = function(m,y,log=FALSE){
  UseMethod("comp_dens")
}
#' @export
comp_dens.default = function(m,y,log=FALSE){
  stop(paste("Invalid class", class(m), "for first argument in",  match.call()))
}

#' Generic function to extract the standard deviations of
#'     components of the mixture
#' @param m a mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @return it returns a vector of standard deviations
#' @export
comp_sd = function(m){
  UseMethod("comp_sd")
}
#' @export
comp_sd.default = function(m){
  stop("method comp_sd not written for this class")
}

#' Generic function to extract which components of mixture are point mass on 0
#' @param m a mixture of k components generated by normalmix() or unimix() or igmix()
#' @return a boolean vector indicating which components are point mass on 0 
#' @export
pm_on_zero = function(m){
  UseMethod("pm_on_zero")
}
#' @export
pm_on_zero.default = function(m){
  return(comp_mean(m)==0 & comp_sd(m)==0)
}

#' Generic function of calculating the second moment of components of
#'     the mixture
#' @param m a mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @return it returns a vector of second moments.
comp_mean2 = function(m){
  UseMethod("comp_mean2")
}
comp_mean2.default = function(m){
  comp_sd(m)^2 + comp_mean(m)^2
}


#' Generic function of calculating the overall mean of the mixture
#'
#'
#' @param m a mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @return it returns scalar, the mean of the mixture distribution.
calc_mixmean = function(m){
  UseMethod("calc_mixmean")
}
calc_mixmean.default = function(m){
  sum(m$pi * comp_mean(m))
}


#' Generic function of calculating the overall second moment of the
#'     mixture
#' @param m a mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @return it returns scalar
mixmean2 = function(m){
  UseMethod("mixmean2")
}
mixmean2.default = function(m){
  sum(m$pi * comp_mean2(m))
}

#' Generic function of calculating the overall standard deviation of
#'     the mixture
#' @param m a mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @return it returns scalar
#'
#' @export
#' 
calc_mixsd = function(m){
  UseMethod("calc_mixsd")
}
#' @export
calc_mixsd.default = function(m){
  sqrt(mixmean2(m)-calc_mixmean(m)^2)
}

#' Generic function of calculating the first moment of components of
#'     the mixture
#' @param m a mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @return it returns a vector of means.
#' @export
comp_mean = function(m){
  UseMethod("comp_mean")
}
#' @export
comp_mean.default = function(m){
  stop("method comp_mean not written for this class")
}

#number of components
#' ncomp
#' @param m a mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @export
#'
ncomp = function(m){
  UseMethod("ncomp")
}

#' @title ncomp.default
#' @description The default version of \code{\link{ncomp}}.
#' @param m a mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @export
#'
ncomp.default = function(m){
  return(length(m$pi))
}


#' Generic function of extracting the mixture proportions
#' @param m a mixture of k components generated by normalmix() or
#'     unimix() or igmix()
#' @return it returns a vector of component probabilities, summing up
#'     to 1.
#' @export
mixprop = function(m){
  UseMethod("mixprop")
}
#' @export
mixprop.default = function(m){
  m$pi
}

#' @title mixcdf
#'
#' @description Returns cdf for a mixture (generic function)
#'
#' @details None
#'
#' @param x a mixture (eg of type normalmix or unimix)
#' @param y locations at which cdf to be computed
#' @param lower.tail boolean indicating whether to report lower tail
#'
#' @return an object of class normalmix
#'
#' @export
#'
#' @examples mixcdf(normalmix(c(0.5,0.5),c(0,0),c(1,2)),seq(-4,4,length=100))
#'
mixcdf = function(x,y,lower.tail=TRUE){
  UseMethod("mixcdf")
}
#' @title mixcdf.default
#' @description The default version of \code{\link{mixcdf}}.
#' @param x a mixture (eg of type normalmix or unimix)
#' @param y locations at which cdf to be computed
#' @param lower.tail boolean indicating whether to report lower tail
#' @export
#'
mixcdf.default = function(x,y,lower.tail=TRUE){
  x$pi %*% comp_cdf(x,y,lower.tail)
}

#find cdf for each component, a generic function
#' Generic function of computing the cdf for each component
#' @param m a mixture (eg of type normalmix or unimix)
#' @param y locations at which cdf to be computed
#' @param lower.tail boolean indicating whether to report lower tail
#' @return it returns a vector of probabilities, with length equals to
#'     number of components in m
#' @export
comp_cdf = function(m,y,lower.tail=TRUE){
  UseMethod("comp_cdf")
}
#' @export
comp_cdf.default = function(m,y,lower.tail=TRUE){
  stop("comp_cdf not implemented for this class")
}


#' Find density at y, a generic function
#' @param x A mixture of k components generated by
#'     \code{\link{normalmix}} or \code{\link{unimix}}.
#' @param y An n-vector of the location.
#' @export
dens = function(x,y){
  UseMethod("dens")
}
#' @export
dens.default = function(x,y){
  return (x$pi %*% comp_dens(x, y))
}

#find log likelihood of data in x (a vector) for mixture in m
loglik = function(m,x){
  UseMethod("loglik")
}
loglik.default = function(m,x){
  sum(log(dens(m,x)))
}


#' @title loglik_conv
#' @description find log likelihood of data using convolution of mixture with error distribution
#'
#' @param m mixture distribution with k components
#' @param data details depend on the model
#' @export
#'
loglik_conv = function(m,data){
  UseMethod("loglik_conv")
}
#' @title loglik_conv.default
#'
#' @description The default version of \code{\link{loglik_conv}}.
#'
#' @param m mixture distribution with k components
#' @param data data whose details depend on model
#' @export
#'
loglik_conv.default = function(m,data){
  sum(log(dens_conv(m,data)))
}



#' @title comp_dens_conv
#' @description compute the density of data for each component of mixture when convolved with error distribution 
#' @inheritParams loglik_conv
#' @param \dots other arguments
#' @return a k by n matrix of densities
comp_dens_conv = function(m,data,...){
  UseMethod("comp_dens_conv")
}
comp_dens_conv.default = function(m,data,...){
  stop(paste("Invalid class", class(m), "for first argument in",  match.call()))
}

#' @title log_comp_dens_conv
#' @description compute the log density of the components of the
#'     mixture m when convoluted with a normal with standard deviation
#'     s or a scaled (se) student.t with df v, the density is
#'     evaluated at x
#' @inheritParams loglik_conv
#' 
#' @return a k by n matrix of log densities
log_comp_dens_conv = function(m,data){
  UseMethod("log_comp_dens_conv")
}
log_comp_dens_conv.default = function(m,data){
  log(comp_dens_conv(m,data))
}


#' @title dens_conv
#' @description compute density of mixture m convoluted with normal of
#'     sd (s) or student t with df v at locations x
#' @inheritParams loglik_conv
dens_conv = function(m,data){
  UseMethod("dens_conv")
}
dens_conv.default = function(m,data){
  colSums(m$pi * comp_dens_conv(m,data))
}


#' @title comp_cdf_conv
#' @description compute the cdf of data for each component of mixture when convolved with error distribution 
#' @param m mixture distribution with k components
#' @param data details depend on the model
#' @return a k by n matrix of cdfs
comp_cdf_conv = function (m, data) {
  UseMethod("comp_cdf_conv")
}
comp_cdf_conv.default = function(m, data){
  stop(paste("Invalid class", class(m), "for first argument in",  match.call()))
}


#' @title cdf_conv
#' @description compute cdf of mixture m convoluted with error distribution
#'  either normal of sd (s) or student t with df v at locations x
#' @inheritParams comp_cdf_conv
cdf_conv = function (m, data) {
  UseMethod("cdf_conv")
}
cdf_conv.default = function (m, data) {
  colSums(m$pi * comp_cdf_conv(m, data))
}


#' @title comp_postprob
#' @description compute the posterior prob that each observation came
#'     from each component of the mixture m,output a k by n vector of
#'     probabilities computed by weighting the component densities by
#'     pi and then normalizing
#'
#' @inheritParams loglik_conv
#' @export
comp_postprob=function(m,data){
 UseMethod("comp_postprob")
}


#' @export
comp_postprob.default = function(m,data){
  lpost = log_comp_dens_conv(m,data) + log(m$pi) # lpost is k by n of log posterior prob (unnormalized)
  lpmax = apply(lpost,2,max) #dmax is of length n
  tmp = exp(t(lpost)-lpmax) #subtracting the max of the logs is just done for numerical stability
  tmp = tmp/rowSums(tmp)
  ismissing = (is.na(data$x) | is.na(data$s))
  tmp[ismissing,]=m$pi
  t(tmp)
}

#' @title comp_cdf_post
#' @description evaluate cdf of posterior distribution of beta at c. m
#'     is the prior on beta, a mixture; c is location of evaluation
#'     assumption is betahat | beta ~ t_v(beta,sebetahat)
#' @param m mixture distribution with k components
#' @param c a scalar
#' @param data details depend on model
#' @return a k by n matrix
#' @examples
#' beta = rnorm(100,0,1)
#' betahat= beta+rnorm(100,0,1)
#' sebetahat=rep(1,100)
#' ash.beta = ash(betahat,1,mixcompdist="normal")
#' comp_cdf_post(get_fitted_g(ash.beta),0,data=set_data(beta,sebetahat))
#' @export
comp_cdf_post=function(m,c,data){
  UseMethod("comp_cdf_post")
}
#' @export
comp_cdf_post.default=function(m,c,data){
  stop("method comp_cdf_post not written for this class")
}

#' @title cdf_post
#' @description evaluate cdf of posterior distribution of beta at c. m
#'     is the prior on beta, a mixture; c is location of evaluation
#'     assumption is betahat | beta ~ t_v(beta,sebetahat)
#' @param m mixture distribution with k components
#' @param c a scalar
#' @param data details depend on model
#' @return an n vector containing the cdf for beta_i at c
#' @examples
#' beta = rnorm(100,0,1)
#' betahat= beta+rnorm(100,0,1)
#' sebetahat=rep(1,100)
#' ash.beta = ash(betahat,1,mixcompdist="normal")
#' cdf0 = cdf_post(ash.beta$fitted_g,0,set_data(betahat,sebetahat))
#' graphics::plot(cdf0,1-get_pp(ash.beta))
#' @export
cdf_post = function(m,c,data){
  UseMethod("cdf_post")
}
#' @export
cdf_post.default=function(m,c,data){
  colSums(comp_postprob(m,data)*comp_cdf_post(m,c,data))
}
#' @title vcdf_post
#' @description vectorized version of \code{\link{cdf_post}}
#' @param m mixture distribution with k components
#' @param c a numeric vector
#' @param data depends on context
#' @return an n vector containing the cdf for beta_i at c
#' @examples
#' beta = rnorm(100,0,1)
#' betahat= beta+rnorm(100,0,1)
#' sebetahat=rep(1,100)
#' ash.beta = ash(betahat,1,mixcompdist="normal")
#' c = vcdf_post(get_fitted_g(ash.beta),seq(-5,5,length=1000),data = set_data(betahat,sebetahat))
#' @export
vcdf_post = function(m,c,data){
  mapply(cdf_post,c,MoreArgs=list(m=m,data=data))
}
#' @title pcdf_post
#' @description ``parallel" vector version of \code{\link{cdf_post}} where c is a vector, of same length as betahat and sebetahat
#' @param m mixture distribution with k components
#' @param c a numeric vector with n elements
#' @param data depends on context
#' @return an n vector, whose ith element is the cdf for beta_i at c_i
#' @examples
#' beta = rnorm(100,0,1)
#' betahat= beta+rnorm(100,0,1)
#' sebetahat=rep(1,100)
#' ash.beta = ash(betahat,1,mixcompdist="normal")
#' c = pcdf_post(get_fitted_g(ash.beta),beta,set_data(betahat,sebetahat))
#' @export
pcdf_post = function(m,c,data){
  vapply(1:length(c),
        function(i){do.call(cdf_post,list(m=m,c=c[i],data=extract_data(data,i)))},FUN.VALUE=0)
}

#output posterior mean for beta for prior mixture m,
#given data
#' postmean
#'
#' @inheritParams loglik_conv
#' @export
postmean = function(m,data){
  UseMethod("postmean")
}

#' @export
postmean.default = function(m,data){
  colSums(comp_postprob(m,data) * comp_postmean(m,data))
}



#' @title postmean2
#' @description output posterior mean-squared value given prior
#'     mixture m and data
#' @inheritParams loglik_conv
#' @export
postmean2 = function(m, data){
  UseMethod("postmean2")
}
#' @export

postmean2.default = function(m,data){
  postmean2_vals <- colSums(comp_postprob(m,data) * comp_postmean2(m,data))
  postmean2_vals[postmean2_vals < 0] <- 0
  return(postmean2_vals)
}


#' @title postsd
#' @description output posterior sd given prior mixture m and data
#' @inheritParams loglik_conv
#' @export
postsd = function(m,data){
  UseMethod("postsd")
}

#' @export
postsd.default = function(m,data){
  postvar = postmean2(m,data)-postmean(m,data)^2
  postvar[postvar<0]=0 #deal with occasional numerical underflow
  sqrt(postvar)
}

#' @title comp_postmean2
#' @description output posterior mean-squared value given prior
#'     mixture m and data
#' @inheritParams loglik_conv
#' @export
comp_postmean2 = function(m,data){
  UseMethod("comp_postmean2")
}
#' @export
comp_postmean2.default = function(m,data){
  stop("method comp_postmean2 not written for this class")
  #comp_postsd(m,betahat,sebetahat,v)^2 + comp_postmean(m,betahat,sebetahat,v)^2
}



#' @title comp_postmean
#' @description output posterior mean for beta for each component of
#'     prior mixture m,given data
#' @inheritParams loglik_conv
#' @export
comp_postmean = function(m,data){
  UseMethod("comp_postmean")
}
#' @export
comp_postmean.default = function(m,data){
  stop("method comp_postmean not written for this class")
}

#' @title prune
#' @description prunes out mixture components with low weight
#' @param m What is this argument?
#' @param thresh the threshold below which components are removed
#' @export
prune = function(m,thresh=1e-10){
  UseMethod("prune")
}
#' @export
prune.default = function(m,thresh = 1e-10){
  which.comp = (m$pi > thresh)
  for(i in 1:length(m)){
    m[[i]] = m[[i]][which.comp]
  }
  m$pi = m$pi/sum(m$pi) #renormalize
  return(m)
}


#' @title comp_postsd
#' @description output posterior sd for beta for each component of
#'     prior mixture m,given data
#' @inheritParams loglik_conv
#' @examples
#' beta = rnorm(100,0,1)
#' betahat= beta+rnorm(100,0,1)
#' ash.beta = ash(betahat,1,mixcompdist="normal")
#' data= set_data(betahat,rep(1,100))
#' comp_postmean(get_fitted_g(ash.beta),data)
#' comp_postsd(get_fitted_g(ash.beta),data)
#' comp_postprob(get_fitted_g(ash.beta),data)
#' @export
comp_postsd = function(m,data){
  UseMethod("comp_postsd")
}
#' @export
comp_postsd.default = function(m,data){
  stop("method comp_postsd not written for this class")
}

#' @title post_sample
#' @description returns random samples from the posterior, given a prior distribution
#'     m and n observed datapoints. 
#' @details exported, but mostly users will want to use `get_post_sample`
#' @param m prior distribution (eg of type normalmix)
#' @param data a list with components x and s, each vectors of length n, to be interpreted as a 
#'     normally-distributed observations and corresponding standard errors
#' @param nsamp number of random samples to return for each observation
#' @return an nsamp by n matrix
#' @export
post_sample = function(m,data,nsamp){
  UseMethod("post_sample")
}

#' @export
post_sample.default = function(m,data,nsamp){
  stop("method post_sample not written for this class")
}

#find nice limits of mixture m for plotting
min_lim = function(m){
  UseMethod("min_lim")
}
min_lim.default=function(m){
  -5
}

max_lim = function(m){
  UseMethod("max_lim")
}
max_lim.default=function(m){
  5
}


#plot density of mixture
plot_dens = function(m,npoints=100,...){
  UseMethod("plot_dens")
}
plot_dens.default = function(m,npoints=100,...){
  x = seq(min_lim(m),max_lim(m),length=npoints)
  graphics::plot(x,dens(m,x),type="l",xlab="density",ylab="x",...)
}

plot_post_cdf = function(m,data,npoints=100,...){
  UseMethod("plot_post_cdf")
}
plot_post_cdf.default = function(m,data,npoints=100,...){
  x = seq(min_lim(m),max_lim(m),length=npoints)
  x_cdf = vapply(x,cdf_post,FUN.VALUE=data$x,m=m,data=data)
  graphics::plot(x,x_cdf,type="l",xlab="x",ylab="cdf",...)
}