R/Celltype_Annotation_PerCell.R

In SlimR: Adaptive Machine Learning-Powered, Context-Matching Tool for Single-Cell and Spatial Transcriptomics Annotation

#' Annotate Seurat Object with Per-Cell SlimR Predictions
#'
#' @description This function assigns SlimR per-cell predicted cell types directly to 
#'     individual cells in a Seurat object's meta.data slot.
#'
#' @param seurat_obj A Seurat object.
#' @param SlimR_percell_result List generated by Celltype_Calculate_PerCell() containing
#'     Cell_annotations data.frame with Cell_barcode and Predicted_cell_type columns.
#' @param plot_UMAP Logical; if TRUE, plot the UMAP with cell type annotations. Default: TRUE.
#' @param annotation_col Column name to write in meta.data. Default: "Cell_type_PerCell_SlimR".
#' @param plot_confidence Logical; if TRUE, also plot a UMAP colored by confidence scores. Default: FALSE.
#'
#' @return A Seurat object with updated meta.data containing:
#' \itemize{
#'   \item annotation_col: Predicted cell type for each cell
#'   \item paste0(annotation_col, "_score"): Max score for each cell
#'   \item paste0(annotation_col, "_confidence"): Confidence score for each cell
#' }
#'
#' @note If plot_UMAP = TRUE, this function will print UMAP plot(s) as a side effect.
#'
#' @export
#' @family Section_3_Automated_Annotation
#'
#' @importFrom Seurat DimPlot FeaturePlot
#' @importFrom ggplot2 ggtitle theme element_text
#' @importFrom patchwork wrap_plots
#'
#' @examples
#' \dontrun{
#' # Run per-cell annotation
#' result <- Celltype_Calculate_PerCell(
#'     seurat_obj = sce,
#'     gene_list = Markers_list,
#'     species = "Human"
#' )
#' 
#' # Annotate Seurat object
#' sce <- Celltype_Annotation_PerCell(
#'     seurat_obj = sce,
#'     SlimR_percell_result = result,
#'     plot_UMAP = TRUE,
#'     annotation_col = "Cell_type_PerCell_SlimR"
#' )
#' }
#'
Celltype_Annotation_PerCell <- function(
    seurat_obj,
    SlimR_percell_result,
    plot_UMAP = TRUE,
    annotation_col = "Cell_type_PerCell_SlimR",
    plot_confidence = FALSE
) {
  
  # ============================================================================
  # Validate inputs
  # ============================================================================
  if (!inherits(seurat_obj, "Seurat")) {
    stop("seurat_obj must be a Seurat object")
  }
  
  if (!is.list(SlimR_percell_result) || !("Cell_annotations" %in% names(SlimR_percell_result))) {
    stop("SlimR_percell_result must be a list containing 'Cell_annotations' data.frame")
  }
  
  cell_anno <- SlimR_percell_result[["Cell_annotations"]]
  
  if (!is.data.frame(cell_anno)) {
    stop("'Cell_annotations' must be a data.frame")
  }
  
  required_columns <- c("Cell_barcode", "Predicted_cell_type", "Max_score", "Confidence")
  if (!all(required_columns %in% colnames(cell_anno))) {
    missing <- setdiff(required_columns, colnames(cell_anno))
    stop(sprintf("Missing required columns in Cell_annotations: %s", paste(missing, collapse = ", ")))
  }
  
  # ============================================================================
  # Check cell barcode matching
  # ============================================================================
  obj_cells <- colnames(seurat_obj)
  result_cells <- cell_anno$Cell_barcode
  
  if (length(obj_cells) != length(result_cells)) {
    stop("Cell count mismatch: Seurat object has ", length(obj_cells), 
         " cells but SlimR_percell_result has ", length(result_cells), " cells")
  }
  
  if (!all(obj_cells == result_cells)) {
    warning("Cell barcodes are not in the same order. Attempting to match by barcode...")
    
    # Try to match by barcode
    match_idx <- match(obj_cells, result_cells)
    
    if (any(is.na(match_idx))) {
      missing_count <- sum(is.na(match_idx))
      stop(sprintf("%d cells in Seurat object not found in SlimR_percell_result", missing_count))
    }
    
    # Reorder to match Seurat object
    cell_anno <- cell_anno[match_idx, ]
  }
  
  # ============================================================================
  # Add annotations to meta.data
  # ============================================================================
  seurat_obj@meta.data[[annotation_col]] <- cell_anno$Predicted_cell_type
  seurat_obj@meta.data[[paste0(annotation_col, "_score")]] <- cell_anno$Max_score
  seurat_obj@meta.data[[paste0(annotation_col, "_confidence")]] <- cell_anno$Confidence
  
  message(sprintf("SlimR PerCell Annotation: Added '%s' to meta.data", annotation_col))
  message(sprintf("  Also added: '%s_score' and '%s_confidence'", annotation_col, annotation_col))
  
  # ============================================================================
  # Plot UMAP if requested
  # ============================================================================
  if (plot_UMAP) {
    if (!"umap" %in% names(seurat_obj@reductions)) {
      warning("UMAP reduction not found. Skipping plot. Run Seurat::RunUMAP() first.")
    } else {
      
      p1 <- Seurat::DimPlot(
        seurat_obj, 
        reduction = "umap", 
        group.by = annotation_col,
        label = TRUE, 
        pt.size = 0.5,
        repel = TRUE
      ) + 
        Seurat::NoAxes() +
        ggplot2::ggtitle("Per-Cell Annotation | SlimR") +
        ggplot2::theme(plot.title = ggplot2::element_text(hjust = 0.5, face = "bold"))
      
      if (plot_confidence) {
        p2 <- Seurat::FeaturePlot(
          seurat_obj,
          features = paste0(annotation_col, "_confidence"),
          reduction = "umap",
          pt.size = 0.5
        ) +
          ggplot2::ggtitle("Annotation Confidence") +
          ggplot2::theme(plot.title = ggplot2::element_text(hjust = 0.5, face = "bold"))
        
        combined_plot <- patchwork::wrap_plots(p1, p2, ncol = 2)
        print(combined_plot)
      } else {
        print(p1)
      }
    }
  }
  
  # ============================================================================
  # Print summary
  # ============================================================================
  if ("Summary" %in% names(SlimR_percell_result)) {
    message("\nAnnotation Summary:")
    print(SlimR_percell_result$Summary)
  }
  
  return(seurat_obj)
}