R/diri.reg.irls.R

In Compositional: Compositional Data Analysis

diri.reg.irls <- function(y, x, xnew = NULL, tol = 1e-6, maxit = 100) {

  runtime <- proc.time()
  n <- dim(y)[1]   ;  D <- dim(y)[2]  ;  p <- D - 1
  x <- model.matrix( y~., data = as.data.frame(x) )
  K <- dim(x)[2]
  ly <- log(y)

  beta <- matrix(0, nrow = K, ncol = p)   # columns = components 2..D
  beta <- Compositional::kl.compreg(y, x[, -1], tol = 1e-5)$be

  phi <- 1.0
  loglik_old <-  -Inf

  txi <- list()
  for ( i in 1:n ) txi[[ i ]] <- tcrossprod(x[i, ])

  for ( iter in 1:maxit ) {
    eta <- x %*% beta                     # n x p
    exp_eta <- exp(eta)
    mu <- cbind(1, exp_eta) / (1 + Rfast::rowsums(exp_eta))   # n x D
    psi1_phi <- trigamma(phi)
    alpha <- phi * mu
    loglik <- sum( lgamma(phi) - Rfast::rowsums( lgamma(alpha) ) + Rfast::rowsums( (alpha - 1) * ly) )

    if ( abs(loglik - loglik_old) < tol ) {
      break
    }

    contribs <- lapply(1:n, function(i) {
      x_i <- x[i, ]
      mu_i <- mu[i, ]
      eta_i <- eta[i, ]
      y_i <- y[i, ]
      ly_i <- ly[i, ]
      alpha_i <- alpha[i, ]

      J <- matrix(0, nrow = D, ncol = p)
      mu_sub <- mu_i[2:D]   # components 2..D
      for ( d in 1:D ) {
        for ( k in 1:p ) {
          J[d, k] <- mu_i[d] * ( (d == (k + 1) ) - mu_sub[k] )
        }
      }

      h_i <- ly_i - digamma(alpha_i)
      psi1_alpha <- trigamma(alpha_i)
      C <- diag(psi1_alpha) - psi1_phi
      JtCJ <- crossprod(J, C) %*% J
      Jt_h <- crossprod(J, h_i)
      S_i <- as.vector( phi * tcrossprod(x_i, Jt_h) )
      I_i <- phi^2 * kronecker(JtCJ, txi[[ i ]] )

      list(S = S_i, I = I_i)
    })

    S_vec <- Reduce(`+`, lapply(contribs, `[[`, "S"))
    I_mat <- Reduce(`+`, lapply(contribs, `[[`, "I"))

    beta_vec <- as.vector(beta)
    beta_vec_new <- beta_vec + solve(I_mat, S_vec)
    beta <- matrix(beta_vec_new, nrow = K, ncol = p)

    # ---- 5. Update phi (direct Newton, no link) ----
    S_phi <- sum( digamma(phi) - Rfast::rowsums(mu * digamma(alpha)) + Rfast::rowsums(mu * ly) )
    H_phi <- sum( trigamma(phi) - Rfast::rowsums( mu^2 * trigamma(alpha) ) )
    phi_new <- phi - S_phi / H_phi
    if ( phi_new <= 0 ) phi_new <- 1e-4
    phi <- phi_new

    loglik_old <- loglik
  }

  colnames(beta) <- paste0("Y", 2:D)
  rownames(beta) <- colnames(x)

  est <- NULL
  if ( !is.null(xnew) ) {
    xnew <- model.matrix(~., data.frame(xnew) )
    mu <- cbind( 1, exp(xnew %*% beta) )
    est <- mu / Rfast::rowsums(mu)
  }

  runtime <- proc.time() - runtime

  list( runtime = runtime, iters = iter, loglik = loglik, phi = phi, be = beta, est = est)
}