ZCA whitening

zca_whiten.R

whitening_matrix <- function(X, eps = 1e-8, zca = FALSE) {
  X <- scale(X, center = TRUE, scale = FALSE)
  s <- svd(X, nu = 0)

  Linvsqr <- sqrt(nrow(X) - 1) / (s$d + eps)

  W <- Linvsqr * t(s$v)
  if (zca) {
    W <- s$v %*% W
  }
  
  W
}

# Perform whitening as Z = X'W where W is the whitening matrix
# X is a matrix with data stored row-wise
# After whitening, column means should be 0
# column standard deviations should be 1
# cov(Z) should be the identity matrix
# with zca = TRUE, Z is closest to X in least-squares sense out of all choices of W
# eps is used to avoid division by zero with small eigenvalues
# values between 0.1 - 1e-5 seem common
whiten <- function(X, eps = 1e-8, zca = FALSE) {
  W <- whitening_matrix(X, eps = eps, zca = zca)
  tcrossprod(X, W)
}