R/nmfkc.bicv.R

In nmfkc: Non-Negative Matrix Factorization with Kernel Covariates

## ============================================================
## nmfkc.bicv(): Bi-cross-validation for NMF rank selection
## (Owen & Perry 2009, AoAS 3(2):564-594).  PROTOTYPE / experimental.
##
## Idea: split rows into folds and columns into folds.  For a held-out
## row-block I and column-block J, the matrix is partitioned into
##   [ A B ]   A = Y[I, J]   (held-out block, predicted)
##   [ C D ]   B = Y[I, Jc]  C = Y[Ic, J]  D = Y[Ic, Jc]  (retained)
## Fit NMF only on D (D ~ L_D R_D), then "fold in" the held-out rows and
## columns by NON-NEGATIVE REGRESSION onto the fixed D-factors:
##   L_I = argmin_{>=0} ||B - L_I R_D||   (held rows' loadings)
##   R_J = argmin_{>=0} ||C - L_D R_J||   (held cols' scores)
## and predict A_hat = L_I R_J.  The held-out rows AND columns never
## enter the NMF fit, so there is no information leakage.
## ============================================================

#' @title Non-negative least squares by multiplicative updates (Internal)
#' @description Solve \eqn{\min_{L\ge 0}\|M - L F\|^2_F} with \eqn{F}
#'   fixed (\code{side = "left"}), or \eqn{\min_{R\ge 0}\|M - F R\|^2_F}
#'   with \eqn{F} fixed (\code{side = "right"}), via a few multiplicative
#'   updates.  Used by \code{\link{nmfkc.bicv}} to fold held-out rows /
#'   columns onto fixed factors.
#' @param M Target matrix.
#' @param F Fixed factor.
#' @param side \code{"left"} solves for \eqn{L} (\code{nrow(M) x k});
#'   \code{"right"} solves for \eqn{R} (\code{k x ncol(M)}).
#' @param maxit Number of multiplicative updates.
#' @return The non-negative factor.
#' @keywords internal
#' @noRd
.nnls.mu <- function(M, F, side = c("left", "right"), maxit = 100) {
  side <- base::match.arg(side)
  eps <- 1e-10
  if (side == "left") {
    k <- base::nrow(F)
    L <- base::matrix(base::mean(M) / (k * base::mean(F) + eps) + eps,
                      base::nrow(M), k)
    FFt <- F %*% base::t(F); MFt <- M %*% base::t(F)
    for (it in base::seq_len(maxit)) L <- L * MFt / (L %*% FFt + eps)
    L
  } else {
    k <- base::ncol(F)
    R <- base::matrix(base::mean(M) / (k * base::mean(F) + eps) + eps,
                      k, base::ncol(M))
    FtF <- base::t(F) %*% F; FtM <- base::t(F) %*% M
    for (it in base::seq_len(maxit)) R <- R * FtM / (FtF %*% R + eps)
    R
  }
}


#' @title Bi-cross-validation for NMF rank selection
#' @description
#' Owen & Perry's (2009) bi-cross-validation (BCV) for choosing the NMF
#' rank.  A lightweight CV engine in the spirit of \code{\link{nmfkc.ecv}}:
#' it returns the held-out error per rank and nothing more (no plot, no
#' table) -- pass the result to \code{which.min(sigma)}, or build your own
#' diagnostics.
#'
#' Unlike the element-wise CV of \code{\link{nmfkc.ecv}} (which holds out
#' scattered \emph{entries} and refits with weights), BCV holds out a whole
#' \strong{row-block and column-block} simultaneously: the model is fit only
#' on the retained block \eqn{D}, and the held-out block \eqn{A} is predicted
#' by folding the held-out rows/columns onto the fixed \eqn{D}-factors via
#' non-negative regression (\eqn{\hat A = L_I R_J}).  Because the held-out
#' rows and columns never enter the fit, there is no information leakage.
#' Covariates are ignored (plain NMF).  The recommended setting is to leave
#' out roughly half the rows and half the columns (\code{nfolds = 2}).
#'
#' @param Y Observation matrix (\eqn{P \times N}), non-negative.
#' @param rank Integer vector of ranks to evaluate.
#' @details Each fold keeps about \eqn{(1 - 1/\text{nfolds})} of the rows
#'   and columns, so the retained block \eqn{D} must have more than
#'   \code{rank} rows \strong{and} columns.  The largest testable rank is
#'   therefore about \eqn{(1 - 1/\text{nfolds})\min(P, N) - 1}; with
#'   \code{nfolds = 2} this is roughly \eqn{\min(P, N)/2 - 1}.  Ranks above
#'   this return \code{NA} and trigger a \code{\link[base]{warning}} that
#'   names the limit and the \code{nfolds} (or \code{\link{nmfkc.ecv}})
#'   that would reach the requested ranks.  Raising \code{nfolds} lifts the
#'   limit at the cost of a smaller hold-out and more compute
#'   (\eqn{(\text{nfolds} - 1)^2} full fits per rank).
#' @param ... Advanced options, rarely needed (defaults in parentheses):
#'   \code{nfolds} (\code{2}), the number of row and column folds (the grid is
#'   \code{nfolds x nfolds}; \code{2} leaves out half the rows / columns, Owen
#'   & Perry's recommendation); \code{seed} (\code{123}, fold-assignment seed);
#'   and \code{nnls.maxit} (\code{100}, multiplicative-update iterations for the
#'   fold-in non-negative regressions).  Any other arguments are passed to
#'   \code{\link{nmfkc}} for the per-block fits (e.g.\ \code{maxit}).
#' @return A list (cf.\ \code{\link{nmfkc.ecv}}) with:
#'   \item{objfunc}{Held-out mean squared error for each rank.}
#'   \item{sigma}{Its square root (RMSE) for each rank.}
#'   \item{rank}{The evaluated rank vector.}
#'   \item{nfolds}{The number of folds used.}
#' @references A. B. Owen and P. O. Perry (2009).  Bi-cross-validation of
#'   the SVD and the nonnegative matrix factorization.
#'   \emph{Ann. Appl. Stat.} 3(2):564--594. \doi{10.1214/08-AOAS227}.
#' @seealso \code{\link{nmfkc.rank}}, \code{\link{nmfkc.ecv}},
#'   \code{\link{nmfkc.consensus}}, \code{\link{nmfkc.ard}}
#' @export
#' @examples
#' \donttest{
#' ## rank-3 non-negative data; bi-CV needs enough kept rows/cols per
#' ## fold (> rank), so use a matrix with ample dimensions.
#' set.seed(1)
#' X <- matrix(abs(rnorm(30 * 3)), 30, 3)
#' B <- matrix(abs(rnorm(3 * 40)), 3, 40)
#' bv <- nmfkc.bicv(X %*% B, rank = 1:6)   # nfolds = 2 (Owen & Perry) by default
#' bv$sigma                  # held-out RMSE per rank
#' bv$rank[which.min(bv$sigma)]
#' }
nmfkc.bicv <- function(Y, rank = 1:3, ...) {
  Y <- base::as.matrix(Y)
  ## Advanced options live in `...` (the rest passes to nmfkc per block).
  ## nfolds = 2 is Owen & Perry's recommendation and rarely changed.
  dots       <- base::list(...)
  nfolds     <- if (base::is.null(dots$nfolds))     2   else dots$nfolds
  seed       <- if (base::is.null(dots$seed))       123 else dots$seed
  nnls.maxit <- if (base::is.null(dots$nnls.maxit)) 100 else dots$nnls.maxit
  dots$nfolds <- NULL; dots$seed <- NULL; dots$nnls.maxit <- NULL
  P <- base::nrow(Y); N <- base::ncol(Y)
  if (!base::is.null(seed)) base::set.seed(seed)
  row.fold <- base::sample(base::rep_len(1:nfolds, P))
  col.fold <- base::sample(base::rep_len(1:nfolds, N))

  ## Feasibility: the retained block D has min(P - largest row fold,
  ## N - largest col fold) rows/cols, and rank-k needs strictly more than
  ## k of each.  Warn (rather than silently return NA) for ranks the chosen
  ## nfolds cannot reach, and suggest the nfolds that would.
  kept.row <- P - base::max(base::table(row.fold))
  kept.col <- N - base::max(base::table(col.fold))
  rmax <- base::min(kept.row, kept.col) - 1L
  if (base::any(rank > rmax)) {
    g.need <- NA_integer_
    for (g in base::seq.int(nfolds + 1L, base::min(P, N))) {
      kr <- P - base::ceiling(P / g); kc <- N - base::ceiling(N / g)
      if (base::min(kr, kc) - 1L >= base::max(rank)) { g.need <- g; break }
    }
    hint <- if (base::is.na(g.need))
      "no nfolds can reach it (rank exceeds min(P, N) - 1); use a smaller rank"
      else base::sprintf("use nfolds = %d (smaller hold-out, more compute), or nmfkc.ecv", g.need)
    base::warning(base::sprintf(
      "nmfkc.bicv: ranks > %d are infeasible with nfolds = %d (retained block ~%d x %d; rank-k needs > k rows and cols) and return NA. To test up to rank %d, %s.",
      rmax, nfolds, kept.row, kept.col, base::max(rank), hint))
  }

  ## nmfkc args for the (plain-NMF) block fits: drop covariates, skip the
  ## O(N^2) clustering criteria, stay quiet about dimensions.
  ea <- dots; ea$A <- NULL; ea$Q <- NULL
  ea$detail <- "fast"; ea$print.dims <- FALSE

  objfunc <- stats::setNames(base::numeric(base::length(rank)),
                             base::paste0("rank=", rank))
  base::message(base::sprintf(
    "bi-CV: ranks %s, %dx%d fold grid (Owen-Perry 2009)...",
    base::paste(rank, collapse = ","), nfolds, nfolds))

  for (ki in base::seq_along(rank)) {
    k <- rank[ki]
    sse <- 0; cnt <- 0L
    for (fi in 1:nfolds) for (fj in 1:nfolds) {
      I  <- base::which(row.fold == fi); J  <- base::which(col.fold == fj)
      Ic <- base::which(row.fold != fi); Jc <- base::which(col.fold != fj)
      if (base::length(Ic) <= k || base::length(Jc) <= k) next
      D  <- Y[Ic, Jc, drop = FALSE]
      Bm <- Y[I,  Jc, drop = FALSE]   # held rows x kept cols
      Cm <- Y[Ic, J,  drop = FALSE]   # kept rows x held cols
      Am <- Y[I,  J,  drop = FALSE]   # held-out block
      fit <- base::suppressMessages(
        base::do.call("nmfkc", c(base::list(Y = D, rank = k), ea)))
      L_D <- fit$X; R_D <- fit$B
      L_I <- .nnls.mu(Bm, R_D, side = "left",  maxit = nnls.maxit)
      R_J <- .nnls.mu(Cm, L_D, side = "right", maxit = nnls.maxit)
      A_hat <- L_I %*% R_J
      sse <- sse + base::sum((Am - A_hat)^2)
      cnt <- cnt + base::length(Am)
    }
    objfunc[ki] <- if (cnt > 0) sse / cnt else NA_real_
  }

  out <- base::list(objfunc = objfunc, sigma = base::sqrt(objfunc),
                    rank = rank, nfolds = nfolds)
  base::class(out) <- "nmfkc.bicv"
  out
}


## ============================================================
## print / plot methods for the nmfkc CV objects
## ============================================================

## Internal: rank-vs-score curve with a Best (Min) marker (shared by ecv/bicv).
## @keywords internal
## @noRd
.cv.rank.plot <- function(rank, score, main, ylab, ...) {
  score <- base::as.numeric(score)
  graphics::plot(rank, score, type = "b", pch = 19,
                 xlab = "Rank (Q)", ylab = ylab, main = main, ...)
  ok <- base::is.finite(score)
  if (base::any(ok)) {
    bi <- base::which.min(base::replace(score, !ok, Inf))
    graphics::points(rank[bi], score[bi], pch = 19, col = "red", cex = 1.7)
    lp <- if (bi == 1L) 4 else if (bi == base::length(rank)) 2 else 3
    graphics::text(rank[bi], score[bi], labels = "Best (Min)",
                   pos = lp, col = "red", xpd = NA)
  }
  invisible(NULL)
}

## Internal: score to plot/report -- sigma (RMSE) when available, else objfunc.
## @keywords internal
## @noRd
.cv.score <- function(x) {
  if (!base::is.null(x$sigma) && base::any(base::is.finite(x$sigma))) {
    list(v = x$sigma, lab = "sigma (held-out RMSE)")
  } else {
    list(v = x$objfunc, lab = "objfunc (held-out loss)")
  }
}

#' @title Print/plot methods for nmfkc cross-validation objects
#' @description Compact print and (for rank sweeps) a score-vs-rank plot with a
#'   \dQuote{Best (Min)} marker, for the objects returned by
#'   \code{\link{nmfkc.ecv}}, \code{\link{nmfkc.cv}} and \code{\link{nmfkc.bicv}}.
#' @param x A \code{"nmfkc.ecv"} / \code{"nmfkc.cv"} / \code{"nmfkc.bicv"} object.
#' @param main,ylab Plot labels.
#' @param ... Passed to \code{\link[graphics]{plot}}.
#' @return \code{x}, invisibly.
#' @name nmfkc.cv.methods
#' @export
print.nmfkc.ecv <- function(x, ...) {
  sc <- .cv.score(x)
  base::cat(sprintf("nmfkc element-wise CV (%s-fold)\n",
                    if (!base::is.null(x$nfolds)) x$nfolds else "?"))
  base::print(base::data.frame(rank = x$rank, sigma = base::as.numeric(x$sigma),
                               objfunc = base::as.numeric(x$objfunc)),
              row.names = FALSE)
  best <- x$rank[base::which.min(base::replace(sc$v, !base::is.finite(sc$v), Inf))]
  base::cat(sprintf("Best rank (min %s): %d\n",
                    if (grepl("sigma", sc$lab)) "sigma" else "objfunc", best))
  invisible(x)
}

#' @rdname nmfkc.cv.methods
#' @export
plot.nmfkc.ecv <- function(x, main = "nmfkc element-wise CV", ylab = NULL, ...) {
  sc <- .cv.score(x)
  .cv.rank.plot(x$rank, sc$v, main = main,
                ylab = if (!base::is.null(ylab)) ylab else sc$lab, ...)
  invisible(x)
}

#' @rdname nmfkc.cv.methods
#' @export
print.nmfkc.bicv <- function(x, ...) {
  base::cat(sprintf("nmfkc bi-cross-validation (%s-fold, Owen & Perry)\n",
                    if (!base::is.null(x$nfolds)) x$nfolds else "?"))
  base::print(base::data.frame(rank = x$rank, sigma = base::as.numeric(x$sigma),
                               objfunc = base::as.numeric(x$objfunc)),
              row.names = FALSE)
  best <- x$rank[base::which.min(base::replace(x$sigma, !base::is.finite(x$sigma), Inf))]
  base::cat(sprintf("Best rank (min sigma): %d\n", best))
  invisible(x)
}

#' @rdname nmfkc.cv.methods
#' @export
plot.nmfkc.bicv <- function(x, main = "nmfkc bi-cross-validation", ylab = "sigma (held-out RMSE)", ...) {
  .cv.rank.plot(x$rank, x$sigma, main = main, ylab = ylab, ...)
  invisible(x)
}

#' @rdname nmfkc.cv.methods
#' @export
print.nmfkc.cv <- function(x, ...) {
  base::cat(sprintf("nmfkc sample-wise CV (rank Q=%s, %s-fold)\n",
                    base::paste(x$rank, collapse = ","),
                    if (!base::is.null(x$nfolds)) x$nfolds else "?"))
  base::cat(sprintf("  objfunc (mean held-out loss): %.6g\n", x$objfunc))
  base::cat(sprintf("  sigma   (held-out RMSE):      %s\n",
                    if (base::is.finite(x$sigma)) sprintf("%.6g", x$sigma) else "NA (non-EU)"))
  invisible(x)
}

#' @rdname nmfkc.cv.methods
#' @export
plot.nmfkc.cv <- function(x, main = "nmfkc CV: per-fold held-out loss", ylab = "objfunc (block)", ...) {
  graphics::barplot(base::as.numeric(x$objfunc.block),
                    names.arg = base::seq_along(x$objfunc.block),
                    xlab = "fold", ylab = ylab, main = main, ...)
  invisible(x)
}