R/laploglik_products.R

In SSN2: Spatial Modeling on Stream Networks

laploglik_products <- function(params_object, data_object, estmethod) {
  cov_matrix_list <- get_cov_matrix_list(params_object, data_object)
  # cholesky products (no local)
  # if (data_object$parallel) {
  #   cluster_list <- lapply(seq_along(cov_matrix_list), function(l) {
  #     cluster_list_element <- list(
  #       c = cov_matrix_list[[l]],
  #       x = data_object$X_list[[l]],
  #       y = data_object$y_list[[l]]
  #     )
  #   })
  #   cholprods_list <- parallel::parLapply(data_object$cl, cluster_list, get_cholprods_glm_parallel)
  #   names(cholprods_list) <- names(cov_matrix_list)
  # } else {
  #   cholprods_list <- mapply(
  #     c = cov_matrix_list, x = data_object$X_list, y = data_object$y_list,
  #     function(c, x, y) get_cholprods_glm(c, x, y),
  #     SIMPLIFY = FALSE
  #   )
  # }

  cholprods_list <- mapply(
    c = cov_matrix_list, x = data_object$X_list, y = data_object$y_list,
    function(c, x, y) get_cholprods_glm(c, x, y),
    SIMPLIFY = FALSE
  )

  SigInv_list <- lapply(cholprods_list, function(x) x$SigInv)
  SigInv <- Matrix::bdiag(SigInv_list)
  SigInv_X <- do.call("rbind", lapply(cholprods_list, function(x) x$SigInv_X))

  # storing relevant products
  ## lower chol %*% X
  SqrtSigInv_X <- do.call("rbind", lapply(cholprods_list, function(x) x$SqrtSigInv_X))
  ## lower chol %*% y
  SqrtSigInv_y <- do.call("rbind", lapply(cholprods_list, function(x) x$SqrtSigInv_y))
  # covariance of beta hat
  ## t(X) %*% sigma_inverse %*% X
  Xt_SigInv_X <- crossprod(SqrtSigInv_X, SqrtSigInv_X)
  ## t(X) %*% sigma_inverse %*% X)^(-1)
  Xt_SigInv_X_upchol <- chol(Xt_SigInv_X)
  cov_betahat <- chol2inv(Xt_SigInv_X_upchol)

  # find dispersion
  dispersion <- as.vector(params_object$dispersion) # take class away

  # newton rhapson
  w_and_H <- get_w_and_H(
    data_object, dispersion,
    SigInv_list, SigInv_X, cov_betahat, Xt_SigInv_X, estmethod
  )

  w <- w_and_H$w
  # H <- w_and_H$H
  mHldet <- w_and_H$mHldet

  betahat <- tcrossprod(cov_betahat, SigInv_X) %*% w
  # reset w after finding betahat
  if (!is.null(data_object$offset)) {
    w <- w + data_object$offset
  }
  X <- do.call("rbind", data_object$X_list)
  r <- w - X %*% betahat
  rt_SigInv_r <- crossprod(r, SigInv) %*% r

  # get wolfinger objects
  y <- as.vector(do.call("rbind", data_object$y_list))
  l00 <- get_l00(data_object$family, w, y, data_object$size, dispersion)
  l01 <- mHldet
  l1 <- sum(unlist(lapply(cholprods_list, function(x) 2 * sum(log(diag(x$Sig_lowchol))))))
  l2 <- as.numeric(rt_SigInv_r)

  # returning relevant quantities
  if (estmethod == "reml") {
    l3 <- 2 * sum(log(diag(Xt_SigInv_X_upchol)))
    return(list(l00 = l00, l01 = l01, l1 = l1, l2 = l2, l3 = l3))
  }

  if (estmethod == "ml") {
    return(list(l00 = l00, l01 = l01, l1 = l1, l2 = l2))
  }
}

get_minustwolaploglik <- function(laploglik_products, data_object, estmethod) {
  if (estmethod == "reml") {
    minustwolaploglik <- as.numeric(laploglik_products$l00 + laploglik_products$l01 + laploglik_products$l1 +
      laploglik_products$l2 + laploglik_products$l3 +
      (data_object$n - data_object$p) * log(2 * pi))
  } else if (estmethod == "ml") {
    minustwolaploglik <- as.numeric(laploglik_products$l00 + laploglik_products$l01 + laploglik_products$l1 +
      laploglik_products$l2 +
      data_object$n * log(2 * pi))
  }

  minustwolaploglik
}


get_w_and_H <- function(data_object, dispersion, SigInv_list, SigInv_X, cov_betahat, cov_betahat_Inv, estmethod, ret_mHInv = FALSE) {
  family <- data_object$family
  SigInv <- Matrix::bdiag(SigInv_list)
  Ptheta <- SigInv - SigInv_X %*% tcrossprod(cov_betahat, SigInv_X)
  y <- as.vector(do.call("rbind", data_object$y_list))
  size <- data_object$size
  w <- get_w_init(family, y, dispersion)
  wdiffmax <- Inf
  iter <- 0



  if (length(SigInv_list) == 1) {
    while (iter < 50 && wdiffmax > 1e-4) {
      iter <- iter + 1
      # if (family %in% c("binomial", "beta")) {
      #   w <- pmax(pmin(w, 8), -8)
      # }
      # compute the d vector
      d <- get_d(family, w, y, size, dispersion)
      # and then the gradient vector
      g <- d - Ptheta %*% w
      # Next, compute H
      D <- get_D(family, w, y, size, dispersion)
      H <- D - Ptheta # not PD but -H is
      solveHg <- solve(H, g)
      wnew <- w - solveHg
      # mH_upchol <- chol(Matrix::forceSymmetric(-H))
      # solveHg <- backsolve(mH_upchol, forwardsolve(t(mH_upchol), g))
      # wnew <- w + solveHg # + because -H is already applied
      # check overshoot on loglik surface
      dnew <- get_d(family, wnew, y, size, dispersion)
      gnew <- dnew - Ptheta %*% wnew
      if (any(is.na(gnew) | is.infinite(gnew))) stop("Convergence problem. Try using a different family, removing extreme observations, rescaling the response variable (if continuous), fixing ie at a known, non-zero value (via spcov_initial), or fixing dispersion at one (via dispersion_initial).", call. = FALSE)
      if (max(abs(gnew)) > max(abs(g))) wnew <- w - 0.1 * solveHg # + because -H is already applied
      # if (max(abs(gnew)) > max(abs(g))) wnew <- w + 0.1 * solveHg
      wdiffmax <- max(abs(wnew - w))
      # update w
      w <- wnew
    }

    # if (family %in% c("binomial", "beta")) {
    #   w <- pmax(pmin(w, 8), -8)
    # }

    mHldet <- as.numeric(determinant(-H, logarithm = TRUE)$modulus)
    # mHldet <- 2 * sum(log(diag(mH_upchol)))
    w_and_H_list <- list(w = w, H = NULL, mHldet = mHldet)
    if (ret_mHInv) {
      # not done above because this is only for model stats and solve(H) slower than solve(H, g)
      HInv <- solve(H)
      w_and_H_list$mHInv <- -HInv
      # mHInv <- chol2inv(mH_upchol)
      # w_and_H_list$mHInv <- mHInv
    }
  } else {
    # # add cov_betahat_Inv stability by same diagonal tolerance as this can have problems too
    # diag(cov_betahat_Inv) <- diag(cov_betahat_Inv) + data_object$diagtol
    #
    # while (iter < 50 && wdiffmax > 1e-4) {
    #   iter <- iter + 1
    #   # compute the d vector
    #   d <- get_d(family, w, y, size, dispersion)
    #   # and then the gradient vector
    #   g <- d - Ptheta %*% w
    #   # Next, compute H
    #   D <- get_D(family, w, y, size, dispersion)
    #   D_diag <- diag(D)
    #   D_list <- lapply(split(D_diag, sort(data_object$local_index)), function(x) Diagonal(x = x))
    #   # cholesky products (while local not implemented)
    #   # if (data_object$parallel) {
    #   #   cluster_list <- lapply(seq_along(D_list), function(l) {
    #   #     cluster_list_element <- list(
    #   #       D = D_list[[l]],
    #   #       S = SigInv_list[[l]]
    #   #     )
    #   #   })
    #   #   DSigInv_list <- parallel::parLapply(data_object$cl, cluster_list, get_DSigInv_parallel)
    #   #   names(DSigInv_list) <- names(D_list)
    #   # } else {
    #   #   DSigInv_list <- mapply(
    #   #     D = D_list, S = SigInv_list,
    #   #     function(D, S) get_DSigInv(D, S),
    #   #     SIMPLIFY = FALSE
    #   #   )
    #   # }
    #
    #   DSigInv_list <- mapply(
    #     D = D_list, S = SigInv_list,
    #     function(D, S) get_DSigInv(D, S),
    #     SIMPLIFY = FALSE
    #   )
    #
    #   # while local not impelmented
    #   # if (data_object$parallel) {
    #   #   cluster_list <- DSigInv_list
    #   #   DSigInv_Inv_list <- parallel::parLapply(data_object$cl, cluster_list, solve)
    #   #   names(DSigInv_Inv_list) <- names(D_list)
    #   # } else {
    #   #   DSigInv_Inv_list <- lapply(DSigInv_list, function(x) solve(x))
    #   # }
    #
    #   DSigInv_Inv_list <- lapply(DSigInv_list, function(x) solve(x))
    #
    #   DSigInv_Inv <- Matrix::bdiag(DSigInv_Inv_list)
    #   HInv <- smw_HInv(AInv = DSigInv_Inv, U = SigInv_X, CInv = cov_betahat_Inv)
    #   solveHg <- HInv %*% g
    #   wnew <- w - solveHg
    #   # check overshoot on loglik surface
    #   dnew <- get_d(family, wnew, y, size, dispersion)
    #   gnew <- dnew - Ptheta %*% wnew
    #   if (any(is.na(gnew) | is.infinite(gnew))) stop("Convergence problem. Try using a different family, removing extreme observations, rescaling the response variable (if continuous), fixing ie at a known, non-zero value (via spcov_initial), or fixing dispersion at one (via dispersion_initial).", call. = FALSE)
    #   if (max(abs(gnew)) > max(abs(g))) wnew <- w - 0.1 * solveHg
    #   wdiffmax <- max(abs(wnew - w))
    #   # update w
    #   w <- wnew
    # }
    #
    # mHldet <- smw_mHldet(A_list = DSigInv_list, AInv = DSigInv_Inv, U = SigInv_X, C = cov_betahat, CInv = cov_betahat_Inv)
    # w_and_H_list <- list(w = w, H = NULL, mHldet = mHldet)
    # if (ret_mHInv) {
    #   w_and_H_list$mHInv <- -HInv
    # }
  }


  # handle offset
  if (!is.null(data_object$offset)) {
    w_and_H_list$w <- w_and_H_list$w - data_object$offset
  }

  w_and_H_list
}

# gradient of w (lowcase d)
get_d <- function(family, w, y, size, dispersion) {
  if (family == "poisson") {
    d <- -exp(w) + y
  } else if (family == "nbinomial") {
    d <- dispersion * (y - exp(w)) / (dispersion + exp(w))
  } else if (family == "binomial") {
    d <- y - size * expit(w)
  } else if (family == "Gamma") {
    d <- -dispersion + dispersion * y * exp(-w)
  } else if (family == "inverse.gaussian") {
    # d <- 1 / dispersion * (y - exp(w)) / exp(2 * w)
    d <- dispersion * (y / (2 * exp(w)) - exp(w) / (2 * y)) + 1 / 2
  } else if (family == "beta") {
    one_expw <- 1 + exp(w)
    k0 <- digamma(dispersion * exp(w) / one_expw) - digamma(dispersion / one_expw) + log(1 / y - 1)
    d <- -dispersion * exp(w) * k0 / one_expw^2
  }
  d
}

# Hessian of w (cap D)
get_D <- function(family, w, y, size, dispersion) {
  w <- as.vector(w)

  if (family == "poisson") {
    D_vec <- -exp(w)
  } else if (family == "nbinomial") {
    D_vec <- -(dispersion * exp(w) * (dispersion + y)) / ((dispersion + exp(w))^2)
  } else if (family == "binomial") {
    D_vec <- -size * expit(w) / (1 + exp(w))
  } else if (family == "Gamma") {
    D_vec <- -dispersion * y * exp(-w)
  } else if (family == "inverse.gaussian") {
    # D_vec <- 1 / dispersion * (exp(w) - 2 * y) / exp(2 * w)
    D_vec <- -dispersion * (exp(2 * w) + y^2) / (2 * y * exp(w))
  } else if (family == "beta") {
    one_expw <- 1 + exp(w)
    k0 <- digamma(dispersion * exp(w) / one_expw) - digamma(dispersion / one_expw) + log(1 / y - 1)
    k1 <- dispersion * (trigamma(dispersion * exp(w) / one_expw) + trigamma(dispersion / one_expw)) - 2 * sinh(w) * (k0 + 2 * atanh(1 - 2 * y))
    D_vec <- -dispersion * exp(2 * w) * k1 / one_expw^4
  }
  D <- Diagonal(x = D_vec)
}

get_w_init <- function(family, y, dispersion) {
  if (family == "poisson") {
    w_init <- 0.5 * log(y + 1)
  } else if (family == "nbinomial") {
    w_init <- 0.5 * log(y + 1)
  } else if (family == "binomial") {
    w_init <- rep(0, times = length(y))
  } else if (family == "Gamma") {
    w_init <- 0.5 * log(y + 1)
  } else if (family == "inverse.gaussian") {
    w_init <- 0.5 * log(y + 1)
  } else if (family == "beta") {
    w_init <- rep(0, times = length(y))
  }
  w_init
}

get_l00 <- function(family, w, y, size, dispersion) {
  w <- as.vector(w)
  y <- as.vector(y)
  # -2 is for -2ll constant
  if (family == "poisson") {
    mu <- exp(w)
    l00 <- -2 * sum(dpois(y, lambda = mu, log = TRUE))
  } else if (family == "nbinomial") {
    mu <- exp(w)
    l00 <- -2 * sum(dnbinom(x = y, mu = mu, size = dispersion, log = TRUE))
  } else if (family == "binomial") {
    mu <- expit(w)
    l00 <- -2 * sum(dbinom(y, size, mu, log = TRUE))
  } else if (family == "Gamma") {
    mu <- exp(w)
    # disp_recip <- 1 / dispersion
    # l00 <- -2 * sum(dgamma(y, shape = disp_recip, scale = dispersion * mu, log = TRUE))
    l00 <- -2 * sum(dgamma(y, shape = dispersion, scale = mu / dispersion, log = TRUE))
  } else if (family == "inverse.gaussian") {
    mu <- exp(w)
    # disp_recip <- 1 / dispersion
    # l00 <- -2 * sum((log(disp_recip) - log(2 * pi) - 3 * log(y)) / 2 - (disp_recip * (y - mu)^2 / (2 * y * mu^2)))
    l00 <- -2 * sum(1 / 2 * (log(dispersion) + log(exp(w)) - log(2 * pi) - log(y^3)) - dispersion * (y - exp(w))^2 / (2 * exp(w) * y))
  } else if (family == "beta") {
    mu <- expit(w)
    a <- mu * dispersion
    b <- (1 - mu) * dispersion
    l00 <- -2 * sum(dbeta(x = y, shape1 = a, shape2 = b, log = TRUE))
  }
  l00
}

# smw_HInv <- function(AInv, U, CInv) {
#   mid <- CInv + t(U) %*% AInv %*% U
#   # solve_mid <- tryCatch(solve(mid), error = function(e) {
#   #   diag(mid) <- diag(mid) + 1e-4 # inverse stability
#   #   solve(mid)
#   # })
#   # diag(mid) <- diag(mid) + 1e-4
#   # if (all(mid == 0)) diag(mid) <- diag(mid) + 1e-4
#   AInv - (AInv %*% U) %*% solve(mid) %*% (t(U) %*% AInv)
# }
#
# smw_mHldet <- function(A_list, AInv, U, C, CInv) {
#   Aldet <- sum(unlist(lapply(A_list, function(x) determinant(x, logarithm = TRUE)$modulus))) # must be positive det for -H
#   Cldet <- 2 * sum(log(diag(t(chol(C)))))
#   mid <- CInv + t(U) %*% AInv %*% U
#   # diag(mid) <- diag(mid) + 1e-4
#   midldet <- determinant(mid, logarithm = TRUE)$modulus
#   as.numeric(Aldet + Cldet + midldet)
# }
#
# get_DSigInv <- function(D, SigInv) {
#   D - SigInv
# }
#
# get_DSigInv_parallel <- function(cluster_list) {
#   D <- cluster_list$D
#   S <- cluster_list$S
#   get_DSigInv(D, S)
# }