BayesFda: Bayesian Functional Data Analysis

#'@importFrom Rcpp evalCpp sourceCpp
#'@importFrom splines spline.des
#'@import coda
#'@useDynLib BayesFda

linvgamma <- function (x, shape, scale) {
  return(- (shape + 1) * log(x) - (scale/x))
}

#' @export
Pspline_Gibbs <- function(x, y, lambda = 1, a_sigma=1, b_sigma=1,
                          knots=NULL, inner_knots = min(round(length(y)/4),40), degree=3, dif = 1, 
                          R = 5000, burn_in = 1000, verbose = FALSE) {
  
  if(any(sum(is.na(y)) > 0,sum(is.na(x)) > 0)) stop("There are missing values either on x or y. Please remove them.")
  
  # Verbose setting
  verbose_step = ceiling(R / 10)
  
  
  # Knots placement
  n     <- length(x)
  if(is.null(knots)) {
    xl    <- min(x); xr <- max(x); dx <- (xr - xl) / (inner_knots-1)
    knots <- seq(xl - degree * dx, xr + degree * dx, by = dx)
  }
  # Fixed quantities
  B     <- spline.des(knots, x, degree + 1, 0 * x, outer.ok=TRUE)$design
  BtB   <- crossprod(B)
  # rankB <- NCOL(B) - dif
  p     <- NCOL(B)
  
  if(dif==0) {D <- diag(NCOL(B))} else{D  <- diff(diag(NCOL(B)), dif=dif)}
  DtD <- crossprod(D) 
  Bty <- crossprod(B,y)
  
  # Initialization
  sigma2  <- var(y)/4
  logpost <- -Inf
  
  # Output
  beta_out  <- matrix(0,R,NCOL(B))
  sigma2_out<- numeric(R)
  
  # Start Gibbs sampling
  for (r in 1:(R + burn_in)) {
    
    eig         <- eigen(BtB/sigma2 + lambda*DtD, symmetric = TRUE)
    Sigma_tilde <- crossprod(t(eig$vectors)/sqrt(eig$values))
    mu_tilde    <- Sigma_tilde%*%(Bty/sigma2)
    
    A1        <- t(eig$vectors)/sqrt(eig$values)
    beta      <- mu_tilde + c(matrix(rnorm(1 * p), 1, p) %*% A1)
    
    # OLD version. The above method is fast and reliable.
    
    # Sigma_tilde <- chol2inv(chol(BtB/sigma2 + lambda*DtD))
    # m_tilde     <- Sigma_tilde%*%(Bty/sigma2)
    # beta        <- c(rmvnorm(1,as.numeric(m_tilde),Sigma_tilde))
    
    # Beta quantities
    # mahalanob     <- as.numeric(crossprod(D%*%beta)) 
    pred          <- as.numeric(B%*%beta) 
    
    # Smoothing component
    # a_tau2_tilde  <- a_tau2 + rankB/2
    # b_tau2_tilde  <- b_tau2 + mahalanob/2
    # tau2          <- 1/rgamma(1,a_tau2_tilde,b_tau2_tilde)
    
    # Variance
    a_sigma_tilde <- a_sigma + n/2
    b_sigma_tilde <- b_sigma + sum((y - pred)^2)/2
    sigma2 <-  1/rgamma(1, a_sigma_tilde, b_sigma_tilde)
    
    # Output
    if (r > burn_in) {
        beta_out[r - burn_in, ] <- as.matrix(beta)
        sigma2_out[r - burn_in] <- sigma2
    }
    if (verbose) {
      if(r%%verbose_step==0) cat(paste("Sampling iteration: ", r, " (R + burn in)\n",sep=""))
    }
  }
  out <- list(param=list(beta=beta_out, sigma2 = sigma2_out), data=list(x=x,y=y), fit=list(knots=knots, degree=degree, dif=dif))
  attr(out,"class") <- "PSpline_Gibbs"
  return(out)
}

#' @export
print.PSpline_Gibbs <- function(x,plot=FALSE) {
  param <- as.mcmc(cbind(sigma2=x$param$sigma2))
  if(plot==TRUE) plot(param)
  print(summary(param))
}

#' @export
predict.PSpline_Gibbs <- function(x,newx=NULL,full_matrix=FALSE) {
  if(!is.null(newx)) {x_grid <- newx} 
  else{x_grid <- x$data$x}
  
  B     <- spline.des(x$fit$knots, x_grid, x$fit$degree + 1, 0 * x_grid, outer.ok=TRUE)$design
  pred  <- t(B%*%t(x$param$beta))
  if(full_matrix) return(as.matrix(pred))
  return(colMeans(pred))
}

#' @export
plot.PSpline_Gibbs <- function(x,ngrid=500,ncurves=200) {
  x_grid <- seq(from=min(x$data$x),to=max(x$data$x),length=ngrid)
  fit <- predict(x,newx=x_grid,full_matrix=TRUE)
  plot(x$data$x,x$data$y,xlab="x",ylab="y",type="n");
  for(r in 1:ncurves) lines(x_grid,fit[r,],col="deepskyblue")
  points(x$data$x,x$data$y,cex=0.7)
}


#' @export
Pspline_CM <- function(x, y, lambda = 1, a_sigma=1, b_sigma=1, knots=NULL, inner_knots = min(length(y)/4,40), degree=3, dif = 1, maxiter = 50, verbose = FALSE) {
  
  if(any(sum(is.na(y)) > 0,sum(is.na(x)) > 0)) stop("There are missing values either on x or y. Please remove them.")
  
  # Convergence tolerance
  tol=1e-5
  
  # Knots placement
  n     <- length(x)
  if(is.null(knots)) {
    xl    <- min(x); xr <- max(x); dx <- (xr - xl) / (inner_knots-1)
    knots <- seq(xl - degree * dx, xr + degree * dx, by = dx)
  }
  
  
  # Fixed quantities
  B     <- spline.des(knots, x, degree + 1, 0 * x, outer.ok=TRUE)$design
  BtB   <- crossprod(B)
  rankB <- NCOL(B) - dif
  
  if(dif==0) {D <- diag(NCOL(B))} else{D   <- diff(diag(NCOL(B)),dif=dif)}
  DtD <- crossprod(D)
  Bty <- crossprod(B,y)
  
  # Initialization
  sigma2 <- var(y)/4
  logpost <- -Inf
  
  # Start CM algorithm
  for(r in 1:maxiter){
    
    beta       <- solve(BtB/sigma2 + lambda*DtD, Bty/sigma2)

    # Beta quantities
    mahalanob     <- as.numeric(crossprod(D%*%beta)) 
    pred          <- as.numeric(B%*%beta) 
    
    # Smoothing component
    # a_tau2_tilde  <- a_tau2 + rankB/2
    # b_tau2_tilde  <- b_tau2 + mahalanob/2
    # tau2          <- b_tau2_tilde/ (a_tau2_tilde+1)

    # Variance
    a_sigma_tilde <- a_sigma + n/2
    b_sigma_tilde <- b_sigma + sum((y-pred)^2)/2
    sigma2        <- b_sigma_tilde / (a_sigma_tilde+1)
    
    # Convergence checks
    loglik         <-  sum(dnorm(y,pred,sqrt(sigma2),log=TRUE))
    logpriors      <- -0.5*mahalanob*lambda + linvgamma(sigma2,a_sigma, b_sigma) 
    logpost_new    <- loglik + logpriors
    
    # Break the loop at convergence
    if((logpost_new - logpost) < tol) {
      if(verbose) cat(paste("Convergence reached after",r,"iterations."))
      break
    }
    logpost <- logpost_new
    
    # Display status
    if (verbose) {
      cat(paste("log-posterior:",round(logpost,4),", iteration:", r, "\n",sep=""))
    }
  }
  out <- list(param=list(beta=beta, sigma2=sigma2), data=list(x=x,y=y), fit=list(knots=knots, degree=degree, dif=dif, iter=r,logpost=logpost))
  attr(out,"class") <- "PSpline_CM"
  return(out)
}

#' @export
print.PSpline_CM <- function(x) {
  cat(paste("Variance parameter: sigma2=",round(x$param$sigma2,6),".\nConvergence reached after ",x$fit$iter," iterations.", sep=""))
}


#' @export
predict.PSpline_CM <- function(x,newx=NULL) {
  if(!is.null(newx)) {x_grid <- newx} 
  else{x_grid <- x$data$x}
  
  B     <- spline.des(x$fit$knots, x_grid, x$fit$degree + 1, 0 * x_grid, outer.ok=TRUE)$design
  pred  <- as.numeric(B %*% x$param$beta)
  return(pred)
}

#' @export
plot.PSpline_CM <- function(x,ngrid=500) {
  x_grid <- seq(from=min(x$data$x),to=max(x$data$x),length=ngrid)
  fit <- predict(x,newx=x_grid)
  plot(x$data$x,x$data$y,xlab="x",ylab="y",cex=.6);
  lines(x_grid,fit,col="deepskyblue",lwd=2)
}