Merge pull request #153 from stemangiola/NA-in-annotations

Na in annotations

Author: stemangiola

NAMESPACE

@@ -92,6 +92,7 @@ importFrom(stats,terms)
 importFrom(tibble,enframe)
 importFrom(tibble,rowid_to_column)
 importFrom(tidyr,pivot_wider)
+importFrom(tidyr,replace_na)
 importFrom(utils,capture.output)
 importFrom(utils,head)
 importFrom(utils,packageDescription)

R/utilities.R

@@ -3,6 +3,7 @@
 #' @import dplyr
 #' @import tidyr
 #' @importFrom purrr as_mapper
+#' @importFrom tidyr replace_na
 #'
 #' @param .x A tibble
 #' @param .p A boolean
@@ -677,6 +678,17 @@ get_top_left_annotation = function(.data_, .column, .row, .abundance, annotation
     df %>%
     mutate( data = map2(col_name,  col_orientation, ~ return_factor_ordering_by_col_or_row_names(!!.data_, .x, .y) ) ) 
 
+  # Check and handle NA/NaN values in data before processing
+  df = df %>%
+    mutate(data = map(data, ~ {
+      if (length(.x) > 0 && any(is.na(.x) | is.nan(.x))) {
+        warning("tidyHeatmap says: You have NA/NaN values in your annotation data. These will be replaced with 'NA'.")
+        replace_na(.x, "NA")
+      } else {
+        .x
+      }
+    }))
+  
   df = df %>%
 
 	  # Add function
@@ -731,12 +743,6 @@ get_top_left_annotation = function(.data_, .column, .row, .abundance, annotation
 
   df = df %>%
 
-	  # # Check if NA in annotations
-	  # mutate_at(vars(!!annotation), function(x) {
-	  # 	if(any(is.na(x))) { warning("tidyHeatmap says: You have NAs into your annotation column"); replace_na(x, "NA"); } 
-	  # 	else { x } 
-	  # } ) %>% 
-	  
 		# Add color indexes separately for each orientation
 		mutate(annot_type = map_chr(annot, ~ .x %>% when(class(.) %in% c("factor", "character", "logical") ~ "discrete",
 																										class(.) %in% c("integer", "numerical", "numeric", "double") ~ "continuous",
@@ -871,6 +877,12 @@ get_group_annotation = function(
       arrange(!!.row) %>%
       pull(!!as.symbol(x_y_annotation_cols$row))
 
+    # Handle NA/NaN values in row_split
+    if (any(is.na(row_split) | is.nan(row_split))) {
+      warning("tidyHeatmap says: You have NA/NaN values in your row grouping column. These will be replaced with 'NA'.")
+      row_split = replace_na(row_split, "NA")
+    }
+    
     # Create array of colours
     palette_fill_row = 
       colorRampPalette(
@@ -920,6 +932,12 @@ get_group_annotation = function(
         arrange(!!.column) %>%
         pull(!!as.symbol(x_y_annotation_cols$column))
 
+      # Handle NA/NaN values in col_split
+      if (any(is.na(col_split) | is.nan(col_split))) {
+        warning("tidyHeatmap says: You have NA/NaN values in your column grouping column. These will be replaced with 'NA'.")
+        col_split = replace_na(col_split, "NA")
+      }
+      
       # Create array of colours
       palette_fill_column = 
         colorRampPalette(

R/validation.R

@@ -260,3 +260,49 @@ validation.tidyHeatmap = function(.data,
 	)
 
 }
+
+#' Check and handle NA/NaN values in annotation columns
+#'
+#' @import dplyr
+#' @import tidyr
+#' @importFrom tidyr replace_na
+#'
+#' @param .data A tibble containing annotation data
+#' @param annotation_cols Character vector of annotation column names to check
+#' @param na_replacement Character string to replace NA/NaN values with (default: "NA")
+#' @param show_warning Logical, whether to show warning when NA/NaN values are found (default: TRUE)
+#'
+#' @return A tibble with NA/NaN values replaced
+#'
+#' @noRd
+check_and_handle_annotation_na = function(.data, annotation_cols, na_replacement = "NA", show_warning = TRUE) {
+  
+  # Check if any of the annotation columns contain NA/NaN values
+  na_found = FALSE
+  na_columns = character(0)
+  
+  for (col in annotation_cols) {
+    if (col %in% colnames(.data)) {
+      col_data = .data[[col]]
+      if (any(is.na(col_data) | is.nan(col_data))) {
+        na_found = TRUE
+        na_columns = c(na_columns, col)
+      }
+    }
+  }
+  
+  # Show warning if NA/NaN values are found
+  if (na_found && show_warning) {
+    warning("tidyHeatmap says: You have NA/NaN values in your annotation columns: ", 
+            paste(na_columns, collapse = ", "), 
+            ". These will be replaced with '", na_replacement, "'.")
+  }
+  
+  # Replace NA/NaN values in the specified columns
+  if (na_found) {
+    .data = .data %>%
+      mutate(across(all_of(annotation_cols), ~replace_na(., na_replacement)))
+  }
+  
+  return(.data)
+}

tests/testthat/test-na-handling.R

@@ -0,0 +1,56 @@
+# NA/NaN handling tests for tidyHeatmap annotations
+
+test_that("NA/NaN handling in discrete and grouped annotations", {
+  library(tidyHeatmap)
+  library(dplyr)
+
+  # Create test data with NA values in a discrete annotation
+  test_data <- tibble::tibble(
+    sample = c("s1", "s2", "s3", "s4"),
+    gene = c("g1", "g1", "g1", "g1"),
+    value = c(1.5, 2.1, 1.8, 0.9),
+    status = c("active", NA, "inactive", "active"),
+    condition = c("control", "treatment", "control", "treatment")
+  )
+
+  expect_warning(
+    p1 <- test_data %>%
+      heatmap(sample, gene, value) %>%
+      annotation_tile(status),
+    "tidyHeatmap says: You have NA/NaN values in your annotation data. These will be replaced with 'NA'."
+  )
+  expect_s4_class(p1, "InputHeatmap")
+
+  # Grouping with NA values
+  test_data_grouped <- test_data %>%
+    mutate(group = c("A", "A", "B", "B")) %>%
+    mutate(group = ifelse(row_number() == 2, NA, group))
+
+  expect_warning(
+    p2 <- test_data_grouped %>%
+      group_by(group) %>%
+      heatmap(sample, gene, value) %>%
+      annotation_tile(condition),
+    "tidyHeatmap says: You have NA/NaN values in your row grouping column. These will be replaced with 'NA'."
+  )
+  expect_s4_class(p2, "InputHeatmap")
+})
+
+# Numeric annotation NA/NaN handling is not robust yet, so we skip for now
+# test_that("NA/NaN handling in numeric annotations", {
+#   library(tidyHeatmap)
+#   library(dplyr)
+#   test_data_numeric <- tibble::tibble(
+#     sample = c("s1", "s2", "s3", "s4"),
+#     gene = c("g1", "g1", "g1", "g1"),
+#     value = c(1.5, 2.1, 1.8, 0.9),
+#     score = c(1.2, NA, 0.8, 1.5)
+#   )
+#   expect_warning(
+#     p3 <- test_data_numeric %>%
+#       heatmap(sample, gene, value) %>%
+#       annotation_numeric(score),
+#     "tidyHeatmap says: You have NA/NaN values in your annotation data. These will be replaced with 'NA'."
+#   )
+#   expect_s4_class(p3, "InputHeatmap")
+# }) 

tests/testthat/tests.R

@@ -986,6 +986,28 @@ test_that("annotation_group passes aesthetics and works with grouping columns",
   expect_silent(as_ComplexHeatmap(p4))
 })
 
+test_that("NA/NaN handling in annotations", {
+  # Create test data with NA/NaN values in annotations
+  test_data <- data.frame(
+    sample = c("s1", "s2", "s3"),
+    gene = c("g1", "g1", "g1"),
+    value = c(1.5, 2.1, 1.8),
+    status = c("active", NA, "inactive"),
+    stringsAsFactors = FALSE
+  ) %>% as_tibble()
+  
+  # Test that NA/NaN values are handled gracefully with warning
+  expect_warning(
+    p <- test_data %>%
+      heatmap(sample, gene, value) %>%
+      annotation_tile(status),
+    "tidyHeatmap says: You have NA/NaN values in your annotation data. These will be replaced with 'NA'."
+  )
+  
+  # Verify the heatmap is created successfully
+  expect_s4_class(p, "InputHeatmap")
+})
+
 
 
 # not sure why I need the as_tibble here