R/nlinla.R

In inlabru: Bayesian Latent Gaussian Modelling using INLA and Extensions

# Linearisation ----

#' Compute inlabru model linearisation information
#'
#' @param \dots Parameters passed on to other methods
#' @export
#' @rdname bru_compute_linearisation
#' @keywords internal
bru_compute_linearisation <- function(...) {
  stopifnot(identical(bru_options_get("bru_method")$autodiff, "pandemic"))
  UseMethod("bru_compute_linearisation")
}

#' @param cmp A [bru_comp] object
#' @param model A [bru_model] object
#' @param lhood_expr A predictor expression
#' @param data Input data
#' @param data_extra Additional data for the predictor
#' @param input Precomputed component inputs from `bru_input()`
#' @param state The state information, as a list of named vectors
#' @param comp_simple Component evaluation information
#' * For `bru_comp`: A [bm_taylor] object
#' * For `bru_obs`: A [bm_list] object
#'   for the components in the likelihood
#' * For `bru_obs_list`: A list of [bm_list] objects
#' @param effects
#' * For `bru_comp`:
#' Precomputed effect list for all components involved in the likelihood
#' expression
#' @param pred0 Precomputed predictor for the given state
#' @param used A [bru_used()] object for the predictor expression
#' @param is_rowwise logical; If `FALSE`, the predictor expression may
#' involve several rows of the input data to influence the same row.
#' @param eps The finite difference step size
#' @param n_pred The length of the predictor expression. If not `NULL`, scalar
#' predictor evaluations are expanded to vectors of length `n_pred`.
#'
#' @export
#' @rdname bru_compute_linearisation
bru_compute_linearisation.bru_comp <- function(
  cmp,
  model,
  lhood_expr,
  data,
  data_extra,
  input,
  state,
  comp_simple,
  effects,
  pred0,
  used,
  is_rowwise,
  eps,
  n_pred = NULL,
  finite_diff = "forward",
  ...
) {
  label <- cmp[["label"]]
  bru_log_message(
    paste0("Linearise with respect to component '", label, "'"),
    verbosity = 5
  )

  if (cmp[["main"]][["type"]] %in% c("offset", "const")) {
    # Zero-column matrix, since a const/offset has no latent state variables.
    return(Matrix::sparseMatrix(
      i = c(),
      j = c(),
      x = c(1),
      dims = c(NROW(pred0), 0)
    ))
  }

  allow_latent <- label %in% used[["latent"]]

  if (is.null(comp_simple)) {
    A <- NULL
    assume_rowwise <- FALSE
  } else {
    A <- ibm_jacobian(
      comp_simple,
      input = input[[label]],
      state = state[[label]]
    )

    assume_rowwise <- !allow_latent && is_rowwise && is.data.frame(data)

    if (assume_rowwise) {
      if (!is.null(n_pred) && (NROW(pred0) != n_pred)) {
        stop(
          glue(
            "The number of values ({NROW(pred0)}) in the predictor for ",
            "observation model ",
            "<unknown> does not match the expected length ({n_pred})."
          )
        )
      }
      if (NROW(A) == 1L) {
        A <- Matrix::kronecker(rep(1, NROW(pred0)), A)
      }
    }
  }

  triplets <- list(
    i = vector("list", NROW(state[[label]])),
    j = vector("list", NROW(state[[label]])),
    x = vector("list", NROW(state[[label]]))
  )

  if (!all(is.finite(pred0))) {
    warning(
      "Non-finite (-Inf/Inf/NaN) entries detected in predictor.\n",
      immediate. = TRUE
    )
  }

  symmetric_diffs <- identical(finite_diff, "central")
  for (k in seq_len(NROW(state[[label]]))) {
    if (is.null(A)) {
      row_subset <- seq_len(NROW(pred0))
    } else {
      Ak <- A[, k, drop = TRUE]
      row_subset <- which(Ak != 0.0)
    }
    if (length(row_subset) > 0) {
      if (symmetric_diffs) {
        state_eps <- list(state, state)
        state_eps[[1]][[label]][k] <- state[[label]][k] - eps
        state_eps[[2]][[label]][k] <- state[[label]][k] + eps
      } else {
        state_eps <- state
        state_eps[[label]][k] <- state[[label]][k] + eps
      }
      # TODO:
      # Option: filter out the data and effect rows for which
      # the rows of A have some non-zeros, or all if !is_rowwise
      # Option: compute predictor for multiple different states. This requires
      # constructing multiple states and corresponding effects before calling
      # evaluate_predictor

      if (symmetric_diffs) {
        effects_eps <- list(effects, effects)
      } else {
        effects_eps <- effects
      }
      if (!is.null(A)) {
        if (assume_rowwise) {
          if (symmetric_diffs) {
            for (label_loop in names(effects)) {
              if (NROW(effects[[label_loop]]) == 1) {
                effects_eps[[1]][[label_loop]] <-
                  rep(effects[[label_loop]], length(row_subset))
                effects_eps[[2]][[label_loop]] <-
                  rep(effects[[label_loop]], length(row_subset))
              } else {
                effects_eps[[1]][[label_loop]] <-
                  effects[[label_loop]][row_subset]
                effects_eps[[2]][[label_loop]] <-
                  effects[[label_loop]][row_subset]
              }
            }
            effects_eps[[1]][[label]] <-
              effects_eps[[1]][[label]] - Ak[row_subset] * eps
            effects_eps[[2]][[label]] <-
              effects_eps[[2]][[label]] + Ak[row_subset] * eps
          } else {
            for (label_loop in names(effects)) {
              if (NROW(effects[[label_loop]]) == 1) {
                effects_eps[[label_loop]] <-
                  rep(effects[[label_loop]], length(row_subset))
              } else {
                effects_eps[[label_loop]] <- effects[[label_loop]][row_subset]
              }
            }
            effects_eps[[label]] <- effects_eps[[label]] + Ak[row_subset] * eps
          }
        } else {
          if (symmetric_diffs) {
            effects_eps <- list(effects, effects)
            effects_eps[[1]][[label]] <- effects_eps[[1]][[label]] - Ak * eps
            effects_eps[[2]][[label]] <- effects_eps[[2]][[label]] + Ak * eps
          } else {
            effects_eps <- effects
            effects_eps[[label]] <- effects_eps[[label]] + Ak * eps
          }
        }
      }
      pred_eps <- evaluate_predictor(
        model,
        state = if (symmetric_diffs) {
          state_eps
        } else {
          list(state_eps)
        },
        data = if (assume_rowwise) {
          data[row_subset, , drop = FALSE]
        } else {
          data
        },
        data_extra = data_extra,
        effects = if (symmetric_diffs) {
          effects_eps
        } else {
          list(effects_eps)
        },
        predictor = lhood_expr,
        used = used,
        format = "matrix",
        n_pred = if (assume_rowwise) {
          length(row_subset)
        } else {
          n_pred
        }
      )
      # Store sparse triplet information
      if (symmetric_diffs) {
        if (assume_rowwise) {
          values <- (pred_eps[, 2] - pred_eps[, 1]) / 2
        } else {
          values <- (pred_eps[, 2] - pred_eps[, 1]) / 2
        }
      } else {
        if (!all(is.finite(pred_eps))) {
          warning(
            "Non-finite (-Inf/Inf/NaN) entries detected in predictor '",
            label,
            "' plus eps.\n",
            immediate. = TRUE
          )
        }
        if (assume_rowwise) {
          values <- (pred_eps - pred0[row_subset])
        } else {
          values <- (pred_eps - pred0)
        }
      }
      nonzero <- is.finite(values)
      if (!all(nonzero)) {
        warning(
          "Non-finite (-Inf/Inf/NaN) entries detected in predictor ",
          "derivatives for '",
          label,
          "'; treated as 0.0.\n",
          immediate. = TRUE
        )
      }
      nonzero[nonzero] <- (values[nonzero] != 0.0) # Detect exact (non)zeros
      if (assume_rowwise) {
        triplets$i[[k]] <- row_subset[nonzero]
      } else {
        triplets$i[[k]] <- which(nonzero)
      }
      triplets$j[[k]] <- rep(k, sum(nonzero))
      triplets$x[[k]] <- values[nonzero] / eps
    }
  }
  B <- Matrix::sparseMatrix(
    i = unlist(triplets$i),
    j = unlist(triplets$j),
    x = unlist(triplets$x),
    dims = c(NROW(pred0), NROW(state[[label]]))
  )
  if (NROW(B) != NROW(pred0)) {
    stop(
      "Jacobian matrix for component '",
      label,
      "' has ",
      NROW(B),
      " rows, but expected ",
      NROW(pred0),
      " rows based on the predictor length."
    )
  }
  B
}

#' @param lhood A `bru_obs` object
#' @param model A `bru_model` object
#' @export
#' @rdname bru_compute_linearisation
bru_compute_linearisation.bru_obs <- function(
  lhood,
  model,
  data,
  input,
  state,
  comp_simple,
  eps,
  ...
) {
  comp_input <- input[["comp"]]
  used <- bru_used(lhood)
  pred_expr <- bru_pred_expr(lhood)
  is_rowwise <- bru_is_rowwise(pred_expr)
  effects <- evaluate_effect_single_state(
    comp_simple[used[["effect"]]],
    input = comp_input[used[["effect"]]],
    state = state[used[["effect"]]]
  )

  lhood_expr <- bru_pred_expr(pred_expr, format = "expr")
  n_pred <- bru_response_size(lhood)

  pred0 <- evaluate_predictor(
    model,
    state = list(state),
    data = data,
    data_extra = lhood[["data_extra"]],
    effects = list(effects),
    predictor = lhood_expr,
    used = used,
    format = "matrix",
    n_pred = n_pred
  )

  # Compute derivatives for each non-const/offset component
  B <- list()
  offset <- pred0
  # Either this loop or the internal bru_comp specific loop
  # can in principle be parallelised.
  comp_lst <- as_bru_comp_list(model)
  additive_and_rowwise <- bru_is_additive(lhood) && is_rowwise
  for (label in union(used[["effect"]], used[["latent"]])) {
    if (ibm_n(comp_lst[[label]][["mapper"]]) > 0) {
      if (additive_and_rowwise) {
        # If additive and no combinations allowed, just need to copy the
        # non-offset A matrix, and possibly expand to full size
        A <- ibm_jacobian(
          comp_simple[[label]],
          comp_input[[label]],
          state[[label]]
        )
        if (NROW(A) == 1) {
          if (NROW(offset) > 1) {
            B[[label]] <- Matrix::kronecker(rep(1, NROW(offset)), A)
          } else {
            B[[label]] <- A
          }
        } else {
          B[[label]] <- A
        }
      } else {
        B[[label]] <-
          bru_compute_linearisation(
            comp_lst[[label]],
            model = model,
            lhood_expr = lhood_expr,
            data = data,
            data_extra = lhood[["data_extra"]],
            input = comp_input,
            state = state,
            comp_simple = comp_simple[[label]],
            effects = effects,
            pred0 = pred0,
            used = used,
            is_rowwise = is_rowwise,
            eps = eps,
            n_pred = n_pred,
            ...
          )
      }
      if ((NROW(offset) == 1L) && (NROW(B[[label]]) > 1L)) {
        offset <- matrix(offset, NROW(B[[label]]), 1)
      }
      offset <- offset - B[[label]] %*% state[[label]]
    }
  }

  bm_taylor(offset = offset, jacobian = B, state0 = NULL)
}

#' @param lhoods A `bru_obs_list` object
#' @export
#' @rdname bru_compute_linearisation
bru_compute_linearisation.bru_obs_list <- function(
  lhoods,
  model,
  input,
  state,
  comp_simple,
  eps = 1e-5,
  ...
) {
  # TODO: set the eps default more intelligently
  lapply(seq_along(lhoods), function(idx) {
    x <- lhoods[[idx]]
    bru_compute_linearisation(
      x,
      model = model,
      data = x[["data"]],
      input = input[[idx]],
      state = state,
      comp_simple = comp_simple[[idx]],
      eps = eps,
      ...
    )
  })
}

#' @export
#' @rdname bru_compute_linearisation
bru_compute_linearisation.bru_model <- function(
  model,
  lhoods,
  input,
  state,
  comp_simple,
  ...
) {
  bru_compute_linearisation(
    lhoods,
    model = model,
    input = input,
    state = state,
    comp_simple = comp_simple,
    ...
  )
}