R/likelihood.R

In koekvall/mpredcc: Fully Likelihood-based Principal Components Regression

log_lik <- function(U, d, alpha, Sigma, tau, Y, X)
{
  n <- nrow(Y)
  p <- ncol(X)
  r <- ncol(Y)

  if(min(d) <= -1 || tau <= 0) return(-Inf)
  e_S <- eigen(Sigma, symmetric = TRUE)
  if(min(e_S$values) <= 0) return(-Inf)

  # Replace by square root inverse, and scale inputs
  Sigma <- e_S$vectors %*% diag(1 / sqrt(abs(e_S$values)), r)
  Y <- Y %*% Sigma
  alpha <- alpha %*% Sigma
  # Psi = U diag(d) U'

  ll <- n * sum(log(abs(e_S$values)))

  # Replace Y by Y - X Psi alpha
  Y <- Y - X %*% (U %*% diag(d, p) %*% crossprod(U, alpha))

  ll <- ll + sum(Y^2)
  ll <- ll + n * p * log(tau)
  ll <- ll + n * sum(log(1 + d))

  # Replace X by X U diag(1 / sqrt(1 + d))
  X <- X %*% U %*% diag(1 / sqrt(1 + d), p)

  ll <- ll + sum(X^2) / tau

  # Minus one half in front
  ll <- -0.5 * ll

  return(ll)
}

obj_fun <- function(U, d, alpha, Sigma, tau, Y, X, pen = 0){
  return(-(2 / nrow(X)) * log_lik(U = U, d = d, alpha = alpha,
                                  Sigma = Sigma, tau = tau,
                                  Y = Y, X = X) + pen * sum(alpha^2))
}

grad_obj_psi <- function(U, d, alpha_tilde, tau, Y_tilde, X)
{
  n <- nrow(Y_tilde)
  p <- ncol(X)
  r <- ncol(Y_tilde)

  # This is what _tilde means: scaled by Omega^{1/2}
  # e_omega <- eigen(Omega)
  # Omega <- e_omega$vectors %*% diag(sqrt(e_omega$values), r)
  #
  # alpha <- alpha %*% Omega
  # Y <- Y %*% Omega

  # Constribution from SSE term
  SXY <- crossprod(X, Y_tilde) / n
  SX <- crossprod(X) / n
  G <- tcrossprod(SX %*% (U %*% diag(d, p) %*% crossprod(U, alpha_tilde)), alpha_tilde)
  G <- G - tcrossprod(SXY, alpha_tilde)
  G <- G + t(G)

  # Contribution from marginal dist of X term
  # Replace U by (I + Psi)^{-1}
  U <- U %*% tcrossprod(diag(1 / (1 + d), p), U)
  G <- G + U
  G <- G - (1 / tau) * U %*% SX %*% U

  return(G)
}

sym_psi_obj <- function(Psi, alpha_tilde, tau, Y_tilde, X)
{
  n <- nrow(Y_tilde)
  p <- ncol(X)
  r <- ncol(Y_tilde)
  XtX <- crossprod(X)
  A <- -t(Y_tilde) %*% X %*% Psi %*% alpha_tilde - t(alpha_tilde) %*% Psi %*% t(X) %*% Y_tilde
  A <- A + t(alpha_tilde) %*% Psi %*% XtX %*% Psi %*% alpha_tilde
  obj <- sum(diag(A)) / n
  obj <- obj + determinant(diag(1, p) + Psi, logarithm = TRUE)$modulus
  A <- qr.solve(diag(1, p) + Psi, XtX)
  obj <- obj + sum(diag(A)) / (n * tau)
  return(obj)
}