#' Interpolation between baseline and prediction
#'
#' @param baseline_type Baseline sequence, either "zero" for all zeros or "shuffle" for random permutation of input_seq.
#' @param m_steps Number of steps between baseline and original input.
#' @param input_seq Input tensor.
#' @noRd
interpolate_seq <- function(m_steps = 50,
baseline_type = "shuffle",
input_seq) {
stopifnot(baseline_type %in% c("zero", "shuffle", "unif"))
if (is.list(input_seq)) {
baseline <- list()
for (i in 1:length(input_seq)) {
input_dim <- dim(input_seq[[i]])
if (baseline_type == "zero") {
baseline[[i]] <- array(rep(0, prod(input_dim)), dim = input_dim)
}
if (baseline_type == "shuffle") {
input_dim <- dim(input_seq[[i]])
baseline[[i]] <- array(input_seq[[i]][ , sample(input_dim[2]), ], dim = input_dim)
}
if (baseline_type == "unif") {
baseline[[i]] <- array(stats::runif(prod(input_dim)), dim = input_dim)
}
}
} else {
if (baseline_type == "zero") {
baseline <- array(rep(0, prod(dim(input_seq))), dim = dim(input_seq))
}
if (baseline_type == "shuffle") {
baseline <- array(input_seq[ , sample(dim(input_seq)[2]), ], dim = dim(input_seq))
}
if (baseline_type == "unif") {
baseline <- array(stats::runif(prod(dim(input_seq))), dim = dim(input_seq))
}
}
m_steps <- as.integer(m_steps)
alphas <- tensorflow::tf$linspace(start = 0.0, stop = 1.0, num = m_steps + 1L) # Generate m_steps intervals for integral_approximation() below.
alphas_x <- alphas[ , tensorflow::tf$newaxis, tensorflow::tf$newaxis]
if (is.list(baseline)) {
delta <- list()
sequences <- list()
for (i in 1:length(baseline)) {
delta[[i]] <- input_seq[[i]] - baseline[[i]]
sequences[[i]] <- baseline[[i]] + alphas_x * delta[[i]]
}
} else {
delta <- input_seq - baseline
sequences <- baseline + alphas_x * delta
}
return(sequences)
}
#' Compute gradients
#'
#' @param input_idx Input layer to monitor for > 1 input.
#' @param target_class_idx Index of class to compute gradient for.
#' @param model Model to compute gradient for.
#' @param pred_stepwise Whether to do predictions with batch_size 1 rather than all at once. Can be used if
#' input is too big to handle at once.
#' @noRd
compute_gradients <- function(input_tensor, target_class_idx, model, input_idx = NULL, pred_stepwise = FALSE) {
# if (is.list(input_tensor)) {
# stop("Stepwise predictions only supported for single input layer yet")
# }
reticulate::py_run_string("import tensorflow as tf")
py$input_tensor <- input_tensor
py$input_idx <- as.integer(input_idx - 1)
py$target_class_idx <- as.integer(target_class_idx - 1)
py$model <- model
if (!is.null(input_idx)) {
reticulate::py_run_string(
"with tf.GradientTape() as tape:
tape.watch(input_tensor[input_idx])
probs = model(input_tensor)[:, target_class_idx]
")
} else {
reticulate::py_run_string(
"with tf.GradientTape() as tape:
tape.watch(input_tensor)
probs = model(input_tensor)[:, target_class_idx]
")
}
grad <- py$tape$gradient(py$probs, py$input_tensor)
if (!is.null(input_idx)) {
return(grad[[input_idx]])
} else {
return(grad)
}
}
integral_approximation <- function(gradients) {
reticulate::py_run_string("import tensorflow as tf")
py$gradients <- gradients
# riemann_trapezoidal
reticulate::py_run_string("grads = (gradients[:-1] + gradients[1:]) / tf.constant(2.0)")
reticulate::py_run_string("integrated_gradients = tf.math.reduce_mean(grads, axis=0)")
return(py$integrated_gradients)
}
#' Compute integrated gradients
#'
#' Computes integrated gradients scores for model and an input sequence.
#' This can be used to visualize what part of the input is import for the models decision.
#' Code is R implementation of python code from [here](https://www.tensorflow.org/tutorials/interpretability/integrated_gradients).
#' Tensorflow implementation is based on this [paper](https://arxiv.org/abs/1703.01365).
#'
#' @param baseline_type Baseline sequence, either `"zero"` for all zeros or `"shuffle"` for random permutation of `input_seq`.
#' @param m_steps Number of steps between baseline and original input.
#' @param input_seq Input tensor.
#' @param target_class_idx Index of class to compute gradient for
#' @param model Model to compute gradient for.
#' @param pred_stepwise Whether to do predictions with batch size 1 rather than all at once. Can be used if
#' input is too big to handle at once. Only supported for single input layer.
#' @param num_baseline_repeats Number of different baseline estimations if baseline_type is `"shuffle"` (estimate integrated
#' gradient repeatedly for different shuffles). Final result is average of \code{num_baseline} single calculations.
#' @examplesIf reticulate::py_module_available("tensorflow")
#' library(reticulate)
#' model <- create_model_lstm_cnn(layer_lstm = 8, layer_dense = 3, maxlen = 20, verbose = FALSE)
#' random_seq <- sample(0:3, 20, replace = TRUE)
#' input_seq <- array(keras::to_categorical(random_seq), dim = c(1, 20, 4))
#' integrated_gradients(
#' input_seq = input_seq,
#' target_class_idx = 3,
#' model = model)
#'
#' @returns A tensorflow tensor.
#' @export
integrated_gradients <- function(m_steps = 50,
baseline_type = "zero",
input_seq,
target_class_idx,
model,
pred_stepwise = FALSE,
num_baseline_repeats = 1) {
#library(reticulate)
reticulate::py_run_string("import tensorflow as tf")
input_idx <- NULL
if (num_baseline_repeats > 1 & baseline_type == "zero") {
warning('Ignoring num_baseline_repeats if baseline is of type "zero". Did you mean to use baseline_type = "shuffle"?')
}
if (num_baseline_repeats == 1 | baseline_type == "zero") {
baseline_seq <- interpolate_seq(m_steps = m_steps,
baseline_type = baseline_type,
input_seq = input_seq)
if (is.list(baseline_seq)) {
for (i in 1:length(baseline_seq)) {
baseline_seq[[i]] <- tensorflow::tf$cast(baseline_seq[[i]], dtype = "float32")
}
} else {
baseline_seq <- tensorflow::tf$cast(baseline_seq, dtype = "float32")
}
if (is.list(input_seq)) {
path_gradients <- list()
avg_grads <- list()
ig <- list()
if (pred_stepwise) {
path_gradients <- gradients_stepwise(
model = model,
baseline_seq = baseline_seq,
target_class_idx = target_class_idx)
} else {
path_gradients <- compute_gradients(
model = model,
input_tensor = baseline_seq,
target_class_idx = target_class_idx,
input_idx = NULL,
pred_stepwise = pred_stepwise)
}
for (i in 1:length(input_seq)) {
avg_grads[[i]] <- integral_approximation(gradients = path_gradients[[i]])
ig[[i]] <- ((input_seq[[i]] - baseline_seq[[i]][1, , ]) * avg_grads[[i]])[1, , ]
}
} else {
if (pred_stepwise) {
path_gradients <- gradients_stepwise(model = model,
baseline_seq = baseline_seq,
target_class_idx = target_class_idx,
input_idx = NULL)
} else {
path_gradients <- compute_gradients(
model = model,
input_tensor = baseline_seq,
target_class_idx = target_class_idx,
input_idx = NULL,
pred_stepwise = pred_stepwise)
}
avg_grads <- integral_approximation(gradients = path_gradients)
ig <- ((input_seq - baseline_seq[1, , ]) * avg_grads)[1, , ]
}
} else {
ig_list <- list()
for (i in 1:num_baseline_repeats) {
ig_list[[i]] <- integrated_gradients(m_steps = m_steps,
baseline_type = "shuffle",
input_seq = input_seq,
target_class_idx = target_class_idx,
model = model,
pred_stepwise = pred_stepwise,
num_baseline_repeats = 1)
}
ig_stacked <- tensorflow::tf$stack(ig_list, axis = 0L)
ig <- tensorflow::tf$reduce_mean(ig_stacked, axis = 0L)
}
return(ig)
}
#' Compute gradients stepwise (one batch at a time)
#'
#' @noRd
gradients_stepwise <- function(model = model, baseline_seq, target_class_idx,
input_idx = NULL) {
if (is.list(baseline_seq)) {
first_dim <- dim(baseline_seq[[1]])[1]
num_input_layers <- length(baseline_seq)
l <- list()
for (j in 1:first_dim) {
input_list <- list()
for (k in 1:length(baseline_seq)) {
input <- as.array(baseline_seq[[k]][j, , ])
input <- array(input, dim = c(1, dim(baseline_seq[[k]])[-1]))
input <- tensorflow::tf$cast(input, baseline_seq[[k]]$dtype)
input_list[[k]] <- input
}
output <- compute_gradients(
model = model,
input_tensor = input_list,
target_class_idx = target_class_idx,
input_idx = NULL)
for (m in 1:length(output)) {
output[[m]] <- tensorflow::tf$squeeze(output[[m]])
}
l[[j]] <- output
}
path_gradients <- vector("list", num_input_layers)
for (n in 1:num_input_layers) {
temp_list <- vector("list", first_dim)
for (p in 1:first_dim){
temp_list[[p]] <- l[[p]][[n]]
}
path_gradients[[n]] <- tensorflow::tf$stack(temp_list)
}
} else {
l <- list()
for (j in 1:dim(baseline_seq)[1]) {
input <- as.array(baseline_seq[j, , ])
input <- array(input, dim = c(1, dim(baseline_seq)[-1]))
input <- tensorflow::tf$cast(input, baseline_seq$dtype)
output <- compute_gradients(
model = model,
input_tensor = input,
target_class_idx = target_class_idx,
input_idx = NULL)
output <- tensorflow::tf$squeeze(output)
l[[j]] <- output
}
path_gradients <- tensorflow::tf$stack(l)
}
return(path_gradients)
}
#' Heatmap of integrated gradient scores
#'
#' Creates a heatmap from output of \code{\link{integrated_gradients}} function. The first row contains
#' the column-wise absolute sums of IG scores and the second row the sums. Rows 3 to 6 contain the IG scores for each
#' position and each nucleotide. The last row contains nucleotide information.
#'
#' @param integrated_grads Matrix of integrated gradient scores (output of \code{\link{integrated_gradients}} function).
#' @param input_seq Input sequence for model. Should be the same as \code{input_seq} input for corresponding
#' \code{\link{integrated_gradients}} call that computed input for \code{integrated_grads} argument.
#' @examplesIf reticulate::py_module_available("tensorflow") && requireNamespace("ComplexHeatmap", quietly = TRUE)
#' library(reticulate)
#' model <- create_model_lstm_cnn(layer_lstm = 8, layer_dense = 3, maxlen = 20, verbose = FALSE)
#' random_seq <- sample(0:3, 20, replace = TRUE)
#' input_seq <- array(keras::to_categorical(random_seq), dim = c(1, 20, 4))
#' ig <- integrated_gradients(
#' input_seq = input_seq,
#' target_class_idx = 3,
#' model = model)
#' heatmaps_integrated_grad(integrated_grads = ig,
#' input_seq = input_seq)
#'
#' @returns A list of heatmaps.
#' @export
heatmaps_integrated_grad <- function(integrated_grads,
input_seq) {
if (is.list(input_seq)) {
for (i in 1:length(input_seq)) {
input_seq[[i]] <- tensorflow::tf$cast(input_seq[[i]], dtype = "float32")
}
for (i in 1:length(integrated_grads)) {
integrated_grads[[i]] <- tensorflow::tf$cast(integrated_grads[[i]], dtype = "float32")
}
} else {
input_seq <- tensorflow::tf$cast(input_seq, dtype = "float32")
integrated_grads <- tensorflow::tf$cast(integrated_grads, dtype = "float32")
}
if (is.list(input_seq)) {
num_input <- length(input_seq)
attribution_mask <- list()
nuc_matrix <- list()
nuc_seq <- list()
sum_nuc <- list()
for (i in 1:length(integrated_grads)) {
py$integrated_grads <- integrated_grads[[i]]
reticulate::py_run_string("attribution_mask = tf.reduce_sum(tf.math.abs(integrated_grads), axis=-1)")
reticulate::py_run_string("sum_nuc = tf.reduce_sum(integrated_grads, axis=-1)")
attribution_mask[[i]] <- py$attribution_mask
attribution_mask[[i]] <- as.matrix(attribution_mask[[i]], nrow = 1) %>% as.data.frame()
colnames(attribution_mask[[i]]) <- "abs_sum"
sum_nuc[[i]] <- py$sum_nuc
sum_nuc[[i]] <- as.matrix(sum_nuc[[i]], nrow = 1) %>% as.data.frame()
colnames(sum_nuc[[i]]) <- "sum"
if (length(dim(integrated_grads[[i]])) == 3) {
nuc_matrix[[i]] <- as.matrix(integrated_grads[[i]][1, , ])
}
if (length(dim(integrated_grads[[i]])) == 2) {
nuc_matrix[[i]] <- as.matrix(integrated_grads[[i]])
}
amb_nuc <- (apply(input_seq[[i]][1, ,], 1, max) %>% as.character()) != "1"
nuc_seq[[i]] <- apply(input_seq[[i]][1, ,], 1, which.max) %>% as.character()
nuc_seq[[i]] <- nuc_seq[[i]] %>% stringr::str_replace_all("1", "A") %>%
stringr::str_replace_all("2", "C") %>%
stringr::str_replace_all("3", "G") %>%
stringr::str_replace_all("4", "T")
nuc_seq[[i]][amb_nuc] <- "0"
rownames(nuc_matrix[[i]]) <- nuc_seq[[i]]
colnames(nuc_matrix[[i]]) <- c("A", "C", "G", "T")
}
} else {
num_input <- 1
py$integrated_grads <- integrated_grads
reticulate::py_run_string("attribution_mask = tf.reduce_sum(tf.math.abs(integrated_grads), axis=-1)")
reticulate::py_run_string("sum_nuc = tf.reduce_sum(integrated_grads, axis=-1)")
#py_run_string("mean_nuc = tf.reduce_mean(integrated_grads, axis=-1)")
attribution_mask <- py$attribution_mask
attribution_mask <- as.matrix(attribution_mask, nrow = 1) %>% as.data.frame()
colnames(attribution_mask) <- "abs_sum"
sum_nuc <- py$sum_nuc
sum_nuc <- as.matrix(sum_nuc, nrow = 1) %>% as.data.frame()
colnames(sum_nuc) <- "sum"
if (length(dim(integrated_grads)) == 3) {
nuc_matrix <- as.matrix(integrated_grads[1, , ])
}
if (length(dim(integrated_grads)) == 2) {
nuc_matrix <- as.matrix(integrated_grads)
}
amb_nuc <- (apply(input_seq[1, ,], 1, max) %>% as.character()) != "1"
nuc_seq <- apply(input_seq[1, ,], 1, which.max) %>% as.character()
nuc_seq <- nuc_seq %>% stringr::str_replace_all("1", "A") %>%
stringr::str_replace_all("2", "C") %>%
stringr::str_replace_all("3", "G") %>%
stringr::str_replace_all("4", "T")
nuc_seq[amb_nuc] <- "0"
rownames(nuc_matrix) <- nuc_seq
colnames(nuc_matrix) <- c("A", "C", "G", "T")
}
if (num_input == 1) {
ig_min <- keras::k_min(integrated_grads)$numpy()
ig_max <- keras::k_max(integrated_grads)$numpy()
col_fun <- circlize::colorRamp2(c(ig_min, mean(c(ig_max, ig_min)) , ig_max), c("blue", "white", "red"))
} else {
col_fun <- list()
for (i in 1:num_input) {
ig_min <- keras::k_min(integrated_grads[[i]])$numpy()
ig_max <- keras::k_max(integrated_grads[[i]])$numpy()
col_fun[[i]] <- circlize::colorRamp2(c(ig_min, mean(c(ig_max, ig_min)) , ig_max), c("blue", "white", "red"))
}
}
hm_list <- list()
if (num_input == 1) {
row_ha = ComplexHeatmap::columnAnnotation(abs_sum = attribution_mask[,1], sum = sum_nuc[,1]) # mean = mean_nuc[,1]
if (length(unique(row.names(nuc_matrix))) == 4) {
nuc_col <- c("A" = "red", "C" = "green", "G" = "blue", "T" = "yellow")
}
if (length(unique(row.names(nuc_matrix))) == 5) {
nuc_col <- c("A" = "red", "C" = "green", "G" = "blue", "T" = "yellow", "0" = "white")
}
ha <- ComplexHeatmap::HeatmapAnnotation(nuc = row.names(nuc_matrix), col = list(nuc = nuc_col))
hm_list[[1]] <- ComplexHeatmap::Heatmap(matrix = t(nuc_matrix),
name = "hm",
top_annotation = row_ha,
bottom_annotation = ha,
col = col_fun,
cluster_rows = FALSE,
cluster_columns = FALSE,
column_names_rot = 0
)
} else {
for (i in 1:num_input) {
row_ha <- ComplexHeatmap::columnAnnotation(abs_sum = attribution_mask[[i]][,1], sum = sum_nuc[[i]][,1])
if (length(unique(row.names(nuc_matrix[[i]]))) == 4) {
nuc_col <- c("A" = "red", "C" = "green", "G" = "blue", "T" = "yellow")
}
if (length(unique(row.names(nuc_matrix[[i]]))) == 5) {
nuc_col <- c("A" = "red", "C" = "green", "G" = "blue", "T" = "yellow", "0" = "white")
}
ha <- ComplexHeatmap::HeatmapAnnotation(nuc = row.names(nuc_matrix[[i]]), col = list(nuc = nuc_col))
hm_list[[i]] <- ComplexHeatmap::Heatmap(matrix = t(nuc_matrix[[i]]),
name = paste0("hm_", i),
top_annotation = row_ha,
bottom_annotation = ha,
col = col_fun[[i]],
cluster_rows = FALSE,
cluster_columns = FALSE,
column_names_rot = 0
)
}
}
hm_list
}
Datasets
Models
