R/polar_est_t.R

In jean997/bvpolar: Estimate Distribution of Angles from Bivariate Data

#'Calculate the distribution of theta given that r > 0
#'@param rhat Vector of rhat
#'@param thetahat Vector of thetahat
#'@param fit_r Object containing mixture distribution of r (output by polar_est_r_mixture_uni)
#'@param tvals Grid values of theta
#'@param tprior prior probability of each element of tvals
#'@param orient indicates that data have been oriented so that thetahat is between -pi/2 and pi/2
#'@export
polar_est_t <- function(rhat, thetahat, fit_r, tvals, tprior, orient=FALSE){

  K <- length(tvals)

  p <- length(thetahat)
  stopifnot(length(rhat)==p)
  has_null <- fit_r$fitted_g$b[1]==0
  llmat <- matrix(nrow=p, ncol=K)

  rpost_mix_full <- mixture_posterior(fit_r$llmat, fit_r$fitted_g$pi)
  wts <- rep(1, p)
  if(has_null){
    fit_r$fitted_g$b <- fit_r$fitted_g$b[-1]
    pr_nz <- sum(fit_r$fitted_g$pi[-1])
    fit_r$fitted_g$pi <- fit_r$fitted_g$pi[-1]/pr_nz
    fit_r$fitted_g$a <- fit_r$fitted_g$a[-1]
    rpost_mix <- rpost_mix_full[-1,]
    wts <- colSums(rpost_mix)
    rpost_mix <- apply(rpost_mix, 2, function(x){x/sum(x)})
  }else{
    rpost_mix <- rpost_mix_full
  }

  Kr <- length(fit_r$fitted_g$pi)
  if(orient){
    stopifnot(all(thetahat >= -pi/2 & thetahat <= pi/2))
    llfun <- polar_log_lik_rt2_runi
  }else{
    llfun <- polar_log_lik_rt_runi
  }
  for(i in 1:K){
    cat(i, " ")
    lls <- sapply(1:p, function(j){
      ll1 <- sapply(fit_r$fitted_g$b, function(r){
        llfun(r, tvals[i], rhat[j], thetahat[j])
      })
      #logSumExp(log(rpost_mix[,j]) + ll1-log(fit_r$fitted_g$b))
      logSumExp(log(fit_r$fitted_g$pi) + ll1-log(fit_r$fitted_g$b))
    })
    llmat[,i] <- lls
  }
  matrix_llik <- apply(llmat, 1, function(x){x + log(tprior/sum(tprior))})
  matrix_llik = llmat - apply(llmat, 1, max)
  matrix_lik = exp(matrix_llik)

  postprobs <- apply(matrix_lik, 1, function(x){x/sum(x)})
  mix_overall <- apply(postprobs, 1, function(x){sum(x*wts)})
  mix_overall <- mix_overall/sum(mix_overall)
  return(list("llmat"=llmat, "tvals"=tvals, "tprior"=tprior, "rpost_mix"=rpost_mix,
              "wts"=wts,"postprobs"=postprobs, "mix_overall"=mix_overall))
}