Fix bugs in instrumental functions

R/adjustment_sets.R

@@ -189,6 +189,7 @@ is_confounder <- function(.tdy_dag, z, x, y, direct = FALSE) {
 #' @param collider_lines logical. Should the plot show paths activated by
 #'   adjusting for a collider?
 #' @param as_factor logical. Should the `adjusted` column be a factor?
+#' @param activate_colliders logical. Include colliders activated by adjustment?
 #'
 #' @return a `tidy_dagitty` with a `adjusted` column for adjusted
 #'   variables, as well as any biasing paths that arise, or a `ggplot`
@@ -202,10 +203,10 @@ is_confounder <- function(.tdy_dag, z, x, y, direct = FALSE) {
 #'
 #' @rdname control_for
 #' @name Adjust for variables
-control_for <- function(.tdy_dag, var, as_factor = TRUE, ...) {
+control_for <- function(.tdy_dag, var, as_factor = TRUE, activate_colliders = TRUE, ...) {
   .tdy_dag <- if_not_tidy_daggity(.tdy_dag, ...)
   dagitty::adjustedNodes(.tdy_dag$dag) <- var
-  .tdy_dag <- activate_collider_paths(.tdy_dag, var)
+  if (isTRUE(activate_colliders)) .tdy_dag <- activate_collider_paths(.tdy_dag, var)
   .tdy_dag$data <- dplyr::mutate(.tdy_dag$data, adjusted = ifelse(name %in% var, "adjusted", "unadjusted"))
   if (as_factor) .tdy_dag$data <- dplyr::mutate(.tdy_dag$data, adjusted = factor(adjusted, exclude = NA))
   .tdy_dag

R/instrumental.R

@@ -57,7 +57,7 @@ node_instrumental <- function(.dag, exposure = NULL, outcome = NULL, ...) {
     conditional_vars <- ifelse(is.null(.z), "", paste("|", paste(.z, collapse = ", ")))
     .dag$data$instrumental_name <- paste(.i, conditional_vars) %>% stringr::str_trim()
     if (!is.null(.z)) {
-      .dag <- .dag %>% control_for(.z)
+      .dag <- .dag %>% control_for(.z, activate_colliders = FALSE)
     } else {
       .dag$data$adjusted <- factor("unadjusted", levels = c("unadjusted", "adjusted"), exclude = NA)
     }
@@ -75,34 +75,53 @@ ggdag_instrumental <- function(.tdy_dag, exposure = NULL, outcome = NULL, ...,
                                node = TRUE, stylized = FALSE, text = TRUE, use_labels = NULL) {
   .tdy_dag <- if_not_tidy_daggity(.tdy_dag) %>%
     node_instrumental(exposure = exposure, outcome = outcome, ...)
+
   p <- .tdy_dag %>%
-    ggplot2::ggplot(ggplot2::aes(x = x, y = y, xend = xend, yend = yend, color = instrumental, shape = adjusted)) +
+    ggplot2::ggplot(ggplot2::aes(x = x, y = y, xend = xend, yend = yend, shape = adjusted)) +
     geom_dag_edges() +
     scale_adjusted() +
     breaks("instrumental")
 
   if (node) {
+    if (all(is.na(.tdy_dag$data$instrumental))) {
+      node_aes <- NULL
+    } else {
+      node_aes <- aes(color = instrumental)
+    }
+
     if (stylized) {
-      p <- p + geom_dag_node(size = node_size)
+      p <- p + geom_dag_node(node_aes, size = node_size)
     } else {
-      p <- p + geom_dag_point(size = node_size)
+      p <- p + geom_dag_point(node_aes, size = node_size)
     }
   }
 
   if (text) p <- p + geom_dag_text(col = text_col, size = text_size)
 
   if (!is.null(use_labels)) {
+    if (all(is.na(.tdy_dag$data$instrumental))) {
+      label_aes <- ggplot2::aes(
+        label = !!rlang::sym(use_labels)
+      )
+    } else {
+      label_aes <- ggplot2::aes(
+        label = !!rlang::sym(use_labels),
+        fill = instrumental
+      )
+    }
+
     p <- p +
       geom_dag_label_repel(
-        ggplot2::aes(
-          label = !!rlang::sym(use_labels),
-          fill = instrumental
-        ),
+        label_aes,
         size = text_size,
         col = label_col,
         show.legend = FALSE
       )
   }
-  if (dplyr::n_distinct(.tdy_dag$data$instrumental_name) > 1) p <- p + ggplot2::facet_wrap(~instrumental_name)
+  if (all(is.na(.tdy_dag$data$instrumental))) {
+    p <- p + ggplot2::facet_wrap(~"{No instrumental variables present}")
+    } else {
+      p <- p + ggplot2::facet_wrap(~instrumental_name)
+    }
   p
 }