add plot_pred compatibility with mids object and add test

Also edit verify_data function to work with this.

R/plot_pred.R

@@ -1,6 +1,6 @@
 #' Plot the predictor matrix of an imputation model
 #'
-#' @param data A predictor matrix for `mice`, typically generated with [mice::make.predictorMatrix] or [mice::quickpred].
+#' @param data A predictor matrix for `mice`, typically generated with [mice::make.predictorMatrix] or [mice::quickpred], or an object of class [`mice::mids`].
 #' @param vrb String, vector, or unquoted expression with variable name(s), default is "all".
 #' @param method Character string or vector with imputation methods.
 #' @param label Logical indicating whether predictor matrix values should be displayed.
@@ -20,28 +20,38 @@ plot_pred <-
            label = TRUE,
            square = TRUE,
            rotate = FALSE) {
-    verify_data(data, pred = TRUE)
+    verify_data(data, pred = TRUE, imp = TRUE)
+    if (mice::is.mids(data)) {
+      method <- data$method
+      data <- data$predictorMatrix
+    }
     p <- nrow(data)
-    if (!is.null(method) && is.character(method)) {
-      if (length(method) == 1) {
-        method <- rep(method, p)
+    if (!mice::is.mids(data)) {
+      if (!is.null(method) && is.character(method)) {
+        if (length(method) == 1) {
+          method <- rep(method, p)
+        }
+        if (length(method) == p) {
+          ylabel <- "Imputation method"
+        }
       }
-      if (length(method) == p) {
-        ylabel <- "Imputation method"
+      if (is.null(method)) {
+        method <- rep("", p)
+        ylabel <- ""
+      }
+      if (!is.character(method) || length(method) != p) {
+        cli::cli_abort("Method should be NULL or a character string or vector (of length 1 or `ncol(data)`).")
       }
-    }
-    if (is.null(method)) {
-      method <- rep("", p)
-      ylabel <- ""
-    }
-    if (!is.character(method) || length(method) != p) {
-      cli::cli_abort("Method should be NULL or a character string or vector (of length 1 or `ncol(data)`).")
     }
     vrb <- substitute(vrb)
     if (vrb[1] == "all") {
       vrb <- names(data)
     } else {
-      vrb <- names(dplyr::select(as.data.frame(data), {{vrb}}))
+      vrb <- names(dplyr::select(as.data.frame(data), {
+        {
+          vrb
+        }
+      }))
     }
     vrbs <- row.names(data)
     long <- data.frame(

R/utils.R

@@ -62,7 +62,7 @@ verify_data <- function(data,
       )
     }
   }
-  if (imp && !df) {
+  if (imp && !df && !pred) {
     if (!mice::is.mids(data)) {
       cli::cli_abort(
         c(
@@ -73,7 +73,18 @@ verify_data <- function(data,
       )
     }
   }
-  if (pred) {
+  if (imp && pred){
+    if (!(is.matrix(data) || mice::is.mids(data))) {
+      cli::cli_abort(
+        c(
+          "The 'data' argument requires an object of class 'matrix', or 'mids'.",
+          "i" = "Input object is of class {class(data)}."
+        ),
+        call. = FALSE
+      )
+    }
+  }
+  if (pred && !imp) {
     if (!is.matrix(data)) {
       cli::cli_abort(
         c(

tests/testthat/test-plot_pred.R

@@ -1,6 +1,7 @@
 # create test objects
 dat <- mice::nhanes
 pred <- mice::quickpred(dat)
+imp <- mice::mice(dat, printFlag = FALSE)
 
 # tests
 test_that("plot_pred creates ggplot object", {
@@ -18,6 +19,7 @@ test_that("plot_pred creates ggplot object", {
   expect_s3_class(plot_pred(rbind(
     cbind(pred, "with space" = 0), "with space" = 0
   )), "ggplot")
+  expect_s3_class(plot_pred(imp, vrb = c("age", "bmi")), "ggplot")
 })
 
 test_that("plot_pred with incorrect argument(s)", {