R/plot_NBKP.R

In NBKP: Beta Kernel Process Modeling with Negative Binomial Response

#' @rdname plot.NBKP
#' @title Plot Fitted NBKP Models
#'
#' @description Visualizes fitted \code{NBKP} (count) models according to the input
#'   dimensionality. For 1D inputs, it shows predicted mean counts with credible
#'   intervals and observed count data. For 2D inputs, it generates contour plots
#'   of posterior summaries. For higher-dimensional inputs, users must specify
#'   which dimensions to plot.
#'
#' @param x An object of class \code{"NBKP"}, typically returned by
#'   \code{\link{fit_NBKP}}.
#' @param only_mean Logical. If \code{TRUE}, only the predicted mean surface is
#'   plotted for 2D inputs. Default is \code{FALSE}.
#' @param n_grid Positive integer specifying the number of grid points per
#'   dimension for constructing the prediction grid. Larger values produce
#'   smoother and more detailed surfaces, but increase computation time. Default
#'   is \code{80}.
#' @param dims Integer vector indicating which input dimensions to plot. Must
#'   have length 1 (for 1D) or 2 (for 2D). If \code{NULL} (default), all
#'   dimensions are used when their number is <= 2.
#' @param ... Additional arguments passed to internal plotting routines
#'   (currently unused).
#'
#' @return This function is called for its side effects and does not return a
#'   value. It produces plots visualizing the predicted counts, credible
#'   intervals, and posterior summaries.
#'
#' @details
#' The plotting behavior depends on the dimensionality of the input covariates:
#' \itemize{
#'   \item \strong{1D inputs}:
#'     \itemize{
#'       \item The function plots the posterior mean curve with a shaded 95% credible
#'         interval, overlaid with the observed counts.
#'     }
#'
#'   \item \strong{2D inputs}:
#'     \itemize{
#'       \item The function generates contour plots over a 2D prediction grid.
#'       \item Users can choose to plot only the predictive mean surface (\code{only_mean = TRUE})
#'         or a set of four summary plots (\code{only_mean = FALSE}):
#'         \enumerate{
#'           \item Predictive mean
#'           \item 97.5th percentile (upper bound of 95% credible interval)
#'           \item Predictive variance
#'           \item 2.5th percentile (lower bound of 95% credible interval)
#'         }
#'     }
#'
#'   \item \strong{Input dimensions greater than 2}:
#'     \itemize{
#'       \item The function does not automatically support visualization and will
#'         terminate with an error.
#'       \item Users must specify which dimensions to visualize via the \code{dims}
#'         argument (length 1 or 2).
#'     }
#' }
#'
#' @seealso \code{\link{fit_NBKP}} for fitting NBKP models;
#'   \code{\link{predict.NBKP}} for generating predictions from fitted NBKP models.
#'
#' @references Zhao J, Qing K, Xu J (2025). \emph{BKP: An R Package for Beta
#' Kernel Process Modeling}. arXiv. https://doi.org/10.48550/arxiv.2508.10447.
#'
#' @keywords NBKP
#'
#' @examples
#' \donttest{
#' # -------------------------- 1D Example --------------------------
#' set.seed(123)
#'
#' # Define true mean function
#' true_mu_fun <- function(x) {
#'   exp(sin(x) + 0.5)
#' }
#'
#' n <- 30
#' Xbounds <- matrix(c(-2, 2), nrow=1)
#' X <- tgp::lhs(n = n, rect = Xbounds)
#' true_mu <- true_mu_fun(X)
#' y <- rnbinom(n, size = 1, mu = true_mu)
#'
#' # Fit NBKP model
#' model1 <- fit_NBKP(X, y, Xbounds=Xbounds)
#'
#' # Plot results
#' plot(model1)
#'
#' # -------------------------- 2D Example --------------------------
#' set.seed(123)
#'
#' # Define 2D latent function and mean transformation
#' true_mu_fun <- function(X) {
#'   if(is.null(nrow(X))) X <- matrix(X, nrow=1)
#'   x1 <- 4*X[,1] - 2
#'   x2 <- 4*X[,2] - 2
#'   f <- sin(2*pi*x1) * cos(2*pi*x2)
#'   return(exp(f))
#' }
#'
#' n <- 100
#' Xbounds <- matrix(c(0, 0, 1, 1), nrow = 2)
#' X <- tgp::lhs(n = n, rect = Xbounds)
#' true_mu <- true_mu_fun(X)
#' y <- rnbinom(n, size = 0.5, mu = true_mu)
#'
#' # Fit NBKP model
#' model2 <- fit_NBKP(X, y, Xbounds=Xbounds)
#'
#' # Plot results
#' plot(model2, n_grid = 50)
#' }
#'
#' @export
#' @method plot NBKP
plot.NBKP <- function(x, only_mean = FALSE, n_grid = 80, dims = NULL, ...) {
  
  if (!is.logical(only_mean) || length(only_mean) != 1) {
    stop("'only_mean' must be a single logical value (TRUE or FALSE).")
  }
  
  if (!is.numeric(n_grid) || length(n_grid) != 1 || n_grid <= 0) {
    stop("'n_grid' must be a positive integer.")
  }
  n_grid <- as.integer(n_grid)
  
  X <- x$X
  y <- x$y
  Xbounds <- x$Xbounds
  d <- ncol(X)
  
  if (is.null(dims)) {
    if (d > 2) {
      stop("X has more than 2 dimensions. Please specify 'dims' for plotting.")
    }
    dims <- seq_len(d)
  } else {
    if (!is.numeric(dims) || any(dims != as.integer(dims))) {
      stop("'dims' must be an integer vector.")
    }
    if (length(dims) < 1 || length(dims) > 2) {
      stop("'dims' must have length 1 or 2.")
    }
    if (any(dims < 1 || dims > d)) {
      stop(sprintf("'dims' must be within the range [1, %d].", d))
    }
    if (any(duplicated(dims))) {
      stop("'dims' cannot contain duplicate indices.")
    }
  }
  
  X_sub <- X[, dims, drop = FALSE]
  
  if (length(dims) == 1) {
    idx <- order(X_sub[,1])
    x_sorted <- X_sub[idx,1]
    y_sorted <- y[idx]
    fitted_sorted <- (x$alpha_n / x$beta_n)[idx]
    
    plot(x_sorted, y_sorted,
         xlab = "X", ylab = "Count",
         main = "NBKP Model Fit (1D)",
         pch = 19, col = "red")
    lines(x_sorted, fitted_sorted, col = "blue", lwd = 2)
    
  } else {
    x1 <- seq(Xbounds[dims[1], 1], Xbounds[dims[1], 2], length.out = n_grid)
    x2 <- seq(Xbounds[dims[2], 1], Xbounds[dims[2], 2], length.out = n_grid)
    grid <- expand.grid(x1 = x1, x2 = x2)
    
    Xnew_full <- matrix(0, nrow(grid), ncol(X))
    Xnew_full[, dims] <- as.matrix(grid)
    prediction <- predict(x, Xnew_full)
    
    df <- data.frame(
      x1 = grid$x1,
      x2 = grid$x2,
      Mean = prediction$mean
    )
    
    p1 <- my_2D_plot_fun("Mean", "Predictive Mean", df, X_sub, dims = dims)
    print(p1)
  }
}

my_2D_plot_fun <- function(var, title, df, X, dims=c(1,2)){
  filled.plot(df$x1, df$x2, df[[var]],
              main=title, xlab=paste0("x",dims[1]), ylab=paste0("x",dims[2]),
              points=X)
}

filled.plot <- function(x, y, z, main, xlab, ylab, points=NULL){
  n <- sqrt(length(z))
  m <- matrix(z, nrow=n, ncol=n)
  filled.contour(
    seq(min(x),max(x),l=n),
    seq(min(y),max(y),l=n),
    m,
    main=main, xlab=xlab, ylab=ylab,
    plot.axes={
      axis(1)
      axis(2)
      if(!is.null(points)) points(points[,1], points[,2], pch=16, col="red", cex=1)
    }
  )
}