update function

Signed-off-by: Daena Rys <[email protected]>

R/plotScree.R

@@ -7,18 +7,14 @@
 #' @param x a
 #' \code{\link[TreeSummarizedExperiment:TreeSummarizedExperiment-constructor]{TreeSummarizedExperiment}}
 #' or a vector of eigenvalues.
-#' 
 #' @param dimred \code{Character scalar} or \code{integer scalar}. Determines
 #' the reduced dimension to plot. This is used when x is a TreeSummarizedExperiment
 #' to extract the eigenvalues from \code{reducedDim(x, dimred)}.
-#' 
 #' @param cumulative \code{Logical scalar}. Whether to show cumulative explained 
 #' variance. (Default: \code{FALSE}).
-#' 
 #' @param names \code{Character vector}. Optional names for the components 
 #' that will be displayed on the x-axis. If not provided, the components 
 #' are labeled sequentially as 1, 2, 3, etc.
-#' 
 #' @param ... additional parameters for plotting
 #' \describe{
 #'   \item{\code{show.barplot}}{Logical scalar. Whether to show a barplot. 
@@ -88,59 +84,57 @@ setMethod("plotScree", signature = c(x = "SingleCellExperiment"),
         }
         # Get reducedDim
         reduced_dim <- reducedDim(x, dimred)
-
         # Extract eigenvalues
         # Check if data is available
         ind <- names(attributes(reduced_dim)) %in% c("eig", "varExplained")
-        if( any(ind) ){
-            # Add explained variance
-            eig <- attributes(reduced_dim)[ind][[1]]
-        } else{
-            stop("No eigenvalues found in the specified reducedDim.", 
-                 call. = FALSE)
+        if( !any(ind) ){
+          stop("No eigenvalues found in the specified reducedDim.", 
+               call. = FALSE)
         }
-
+        eig <- attributes(reduced_dim)[ind][[1]]
         # Call the vector method
-        plotScree(as.numeric(eig), names(eig), cumulative = cumulative, ...)
+        result <- plotScree(as.numeric(eig), names(eig), cumulative = cumulative, ...)
+        return(result)
     }
 )
-
 #' @rdname plotScree
 #' @export
 setMethod("plotScree", signature = c(x = "vector"),
-    function(x, names = NULL, cumulative = FALSE, ...) {
+    function(x, cumulative = FALSE, ...) {
         # Ensure 'x' is numeric
         if (!is.numeric(x)) {
             stop("'x' must be a numeric vector.", call. = FALSE)
         }
-        plot_data <- .prepare_data(x, names, cumulative)
+        if (anyNA(x)) {
+            warning("NA values found in eigenvalues; 
+                    they will be omitted from the plot.")
+        }
+        plot_data <- .prepare_data(x, cumulative)
         # plot vector
-        .scree_plotter(plot_data, cumulative = cumulative, ...)
+        result <- .scree_plotter(plot_data, cumulative = cumulative, ...)
+        return(result)
     }
 )
-
-.prepare_data <- function(x, names = NULL, cumulative = FALSE) {
+.prepare_data <- function(x, cumulative = FALSE) {
+    # drop NA values in eigenvalues
+    x <- x[!is.na(x)]
     df <- data.frame(
-        Component = if (!is.null(names)) names else seq_along(x),
+        component = if (!is.null(names(x))) names(x) else seq_along(x),
         Eigenvalue = x
     )
-
     # Calculate cumulative proportion if needed
     if (cumulative) {
         df$CumulativeProportion <- cumsum(df$Eigenvalue) / sum(df$Eigenvalue)
     }
-
     return(df)
 }
-
 .scree_plotter <- function(df, show.barplot = TRUE, 
                            show.points = TRUE, 
                            show.line = TRUE, show.labels = FALSE, 
                            cumulative = FALSE, ...) {
     # Create base plot
     p <- ggplot(df, aes(x = Component, y = if (cumulative) 
         CumulativeProportion else Eigenvalue))
-
     # Add layers based on user preferences
     if (show.barplot) {
         p <- p + geom_col(fill = "lightblue", color = "black")
@@ -156,14 +150,10 @@ setMethod("plotScree", signature = c(x = "vector"),
             CumulativeProportion else Eigenvalue, 2)), 
             vjust = -0.5)
     }
-
     # Customize appearance
     p <- p + theme_minimal() +
-        labs(x = "Component", 
-             y = if (cumulative) "Cumulative Proportion of Variance" 
-             else "Eigenvalue",
-             title = if (cumulative) "Cumulative Explained Variance" 
-             else "Scree Plot")
-
+      labs(x = "PC", 
+           y = if (cumulative) "Cumulative Proportion of Variance" 
+           else "Eigenvalue")
     return(p)
 }