R/cmtf_fg.R

In CMTFtoolbox: Create (Advanced) Coupled Matrix and Tensor Factorization Models

#' Function value and gradient calculation for CMTF
#'
#' @param x Vectorized parameters of the CMTF model.
#' @param Z Z object as generated by [setupCMTFdata()].
#'
#' @return A list containing the function ("fn") and the gradient ("gr").
#' @export
#'
#' @examples
#' A = array(rnorm(108*2), c(108, 2))
#' B = array(rnorm(100*2), c(100, 2))
#' C = array(rnorm(10*2), c(10, 2))
#' D = array(rnorm(100*2), c(100,2))
#' E = array(rnorm(10*2), c(10,2))
#'
#' df1 = reinflateTensor(A, B, C)
#' df2 = reinflateTensor(A, D, E)
#' datasets = list(df1, df2)
#' modes = list(c(1,2,3), c(1,4,5))
#' Z = setupCMTFdata(datasets, modes, normalize=FALSE)
#'
#' init = initializeCMTF(Z, 2, output="vect")
#' outcome = cmtf_fg(init, Z)
#' f = outcome$fn
#' g = outcome$gr
cmtf_fg = function(x, Z){

  numDatasets = length(Z$object)
  numModes = max(unlist(Z$modes))
  Fac = vect_to_fac(x, Z, sortComponents=FALSE)
  reinflatedBlocks = reinflateFac(Fac, Z, returnAsTensor=TRUE)

  ## FUN PART ##
  f_per_block = rep(NA, numDatasets)
  for(p in 1:numDatasets){
    modes = Z$modes[[p]]
    reinflatedBlock = reinflatedBlocks[[p]]
    residuals = Z$object[[p]] - reinflatedBlock
    residuals = Z$missing[[p]] * residuals

    Fnorm = rTensor::fnorm(residuals)
    f_per_block[p] = 0.5 * Fnorm^2
  }

  f = sum(f_per_block)

  ## GRADIENT PART ##
  gradient = list()

  # Gradients per mode stored in a list, will be vectorized at the end.
  for(i in 1:numModes){
    gradient[[i]] = array(0L, dim(Fac[[i]]))

    for(p in 1:numDatasets){
      modes = Z$modes[[p]]

      if(i %in% modes){
        idx = which(modes==i)
        otherModes = modes[-idx]

        unfoldedX = rTensor::k_unfold(Z$missing[[p]], idx) * rTensor::k_unfold(Z$object[[p]], idx)
        unfoldedXhat = rTensor::k_unfold(Z$missing[[p]], idx) * rTensor::k_unfold(reinflatedBlocks[[p]], idx)

        if(length(modes) == 3){
          gradientMode = (unfoldedXhat - unfoldedX)@data %*% multiway::krprod(Fac[[otherModes[2]]], Fac[[otherModes[1]]])
        } else if((length(modes) == 2)){
          gradientMode = (unfoldedXhat - unfoldedX)@data %*% Fac[[otherModes[1]]]
        }
        else{
          stop(paste0("Number of modes is incorrect for block ", p))
        }

        gradient[[i]] = gradient[[i]] + gradientMode
      }
    }
  }

  g = fac_to_vect(gradient)
  return(list("fn"=f, "gr"=g))
}