Add ego_color_* and ego_label_var parameters to plot_egograms() for homophily visualization

DESCRIPTION

@@ -55,7 +55,7 @@ Suggests:
     network,
     haven
 VignetteBuilder: knitr
-RoxygenNote: 7.3.3
+RoxygenNote: 7.3.3.9000
 Roxygen: list(markdown = TRUE)
 LazyData: true
 Encoding: UTF-8

R/plot_ego_gram.R

@@ -102,6 +102,10 @@ plot_egograms <- function(x,
                           vertex_color_palette = "Heat Colors",
                           vertex_color_legend_label = vertex_color_var,
                           vertex_label_var = "name",
+                          ego_color_var = vertex_color_var,
+                          ego_color_palette = vertex_color_palette,
+                          ego_color_legend_label = ego_color_var,
+                          ego_label_var = vertex_label_var,
                           edge_width_var = NULL,
                           edge_color_var = NULL,
                           edge_color_palette = "Heat Colors",
@@ -153,6 +157,10 @@ plot_egograms <- function(x,
         font_size = font_size,
         show_venn_labels = show_venn_labels,
         include_ego = include_ego,
+        ego_color_var = ego_color_var,
+        ego_color_palette = ego_color_palette,
+        ego_color_legend_label = ego_color_legend_label,
+        ego_label_var = ego_label_var,
         ...
       )
     }
@@ -170,6 +178,10 @@ plot_egogram <-
            vertex_color_palette = "Heat Colors",
            vertex_color_legend_label = vertex_color_var,
            vertex_label_var = "name",
+           ego_color_var = vertex_color_var,
+           ego_color_palette = vertex_color_palette,
+           ego_color_legend_label = ego_color_var,
+           ego_label_var = vertex_label_var,
            edge_width_var = NULL,
            edge_color_var = NULL,
            edge_color_palette = "Heat Colors",
@@ -251,23 +263,27 @@ plot_egogram <-
 
     # Pieces of the pie
     plot.new()
-    pie_add(
-      rep(1, piece_n),
-      labels = levels(pie_var),
-      radius = 1,
-      clockwise = TRUE,
-      border = FALSE,
-      add = TRUE,
-      col = pie_colors,
-      cex = font_size
-    )
+    if (piece_n > 0) {
+      pie_add(
+        rep(1, piece_n),
+        labels = levels(pie_var),
+        radius = 1,
+        clockwise = TRUE,
+        border = FALSE,
+        add = TRUE,
+        col = pie_colors,
+        cex = font_size
+      )
+    }
 
-    # Venns
-    radi <- c(1:(venn_n + 1) / (venn_n + 1))
-    cols <- paste0("#ffffff", format(as.hexmode(round(seq(0, 220,  220 / venn_n)))))
-    if(venn_gradient_reverse) cols <- rev(cols)
-    for(i in 1:venn_n) {
-      ring(0, 0, radi[i+1], radi[i], col = cols[i], border = "grey70")
+    # Venns - only draw if there are alters
+    if (nrow(ego_object$alter) > 0 && venn_n > 0) {
+      radi <- c(1:(venn_n + 1) / (venn_n + 1))
+      cols <- paste0("#ffffff", format(as.hexmode(round(seq(0, 220,  220 / venn_n)))))
+      if(venn_gradient_reverse) cols <- rev(cols)
+      for(i in 1:venn_n) {
+        ring(0, 0, radi[i+1], radi[i], col = cols[i], border = "grey70")
+      }
     }
 
     # plotrix::draw.circle(0, 0, c(1:(venn_n + 1) / (venn_n + 1)),
@@ -350,11 +366,36 @@ plot_egogram <-
 
     if(include_ego) {
       # Place ego in middle of plot
-      lay <- rbind(lay, c(0, 0, 0))
-      # Set curvature of ego-alter ties to zero
-      igraph::E(g)$curved[is.na(igraph::E(g)$curved)] <- 0
-      # Set ego-alter weights to a dummy value
-      igraph::E(g)$weight[is.na(igraph::E(g)$weight)] <- min(igraph::E(g)$weight, na.rm = TRUE)
+      if (nrow(lay) > 0) {
+        # Add ego to layout using proper tibble row binding
+        ego_row <- tibble::tibble(.altID = factor(".ego", levels = c(levels(lay$.altID), ".ego")), 
+                                   x = 0, y = 0)
+        lay <- rbind(lay, ego_row)
+      } else {
+        # If no alters, create a layout with just the ego
+        lay <- tibble::tibble(.altID = factor(".ego"), x = 0, y = 0)
+      }
+      # Only set edge attributes if there are edges
+      if (igraph::ecount(g) > 0) {
+        # Set curvature of ego-alter ties to zero
+        curved_vals <- igraph::E(g)$curved
+        if (is.null(curved_vals) || length(curved_vals) == 0 || all(is.na(curved_vals))) {
+          igraph::E(g)$curved <- rep(0, igraph::ecount(g))
+        } else {
+          curved_vals[is.na(curved_vals)] <- 0
+          igraph::E(g)$curved <- curved_vals
+        }
+        # Set ego-alter weights to a dummy value
+        if (any(!is.na(igraph::E(g)$weight))) {
+          # Set to min of other weights, so scale of weights is comparable
+          weight_vals <- igraph::E(g)$weight
+          weight_vals[is.na(weight_vals)] <- min(weight_vals, na.rm = TRUE)
+          igraph::E(g)$weight <- weight_vals
+        } else {
+          # If all weights are NA (no aaties), set them all to 1
+          igraph::E(g)$weight <- 1
+        }
+      }
     }
 
     # Plot
@@ -380,6 +421,10 @@ plot_egogram <-
       edge.curved = igraph::E(g)$curved,
       include_ego = include_ego, 
       font_size = font_size,
+      ego_color_var = ego_color_var,
+      ego_color_palette = ego_color_palette,
+      ego_color_legend_label = ego_color_legend_label,
+      ego_label_var = ego_label_var,
       ...
     )
   }

R/plot_ego_graph.R

@@ -9,6 +9,10 @@ plot_ego_graphs <- function(x,
                             vertex_color_palette = "Heat Colors",
                             vertex_color_legend_label = vertex_color_var,
                             vertex_label_var = "name",
+                            ego_color_var = vertex_color_var,
+                            ego_color_palette = vertex_color_palette,
+                            ego_color_legend_label = ego_color_var,
+                            ego_label_var = vertex_label_var,
                             edge_width_var = NULL,
                             ego_alter_edge_width_var = 
                               if(!is.null(edge_width_var) & include_ego) edge_width_var,
@@ -57,6 +61,10 @@ plot_ego_graphs <- function(x,
         font_size = font_size,
         include_ego = include_ego,
         ego_attrs = ego_attrs,
+        ego_color_var = ego_color_var,
+        ego_color_palette = ego_color_palette,
+        ego_color_legend_label = ego_color_legend_label,
+        ego_label_var = ego_label_var,
         ...
       )
     }
@@ -71,6 +79,10 @@ plot_one_ego_graph <- function(x,
                                vertex_color_palette = "Heat Colors",
                                vertex_color_legend_label = vertex_color_var,
                                vertex_label_var = "name",
+                               ego_color_var = vertex_color_var,
+                               ego_color_palette = vertex_color_palette,
+                               ego_color_legend_label = ego_color_var,
+                               ego_label_var = vertex_label_var,
                                edge_width_var = NULL,
                                ego_alter_edge_width_var = edge_width_var,
                                edge_color_var = NULL,
@@ -92,6 +104,22 @@ plot_one_ego_graph <- function(x,
     if (vertex_label_var %in% names(x$ego)) {
       ego_attrs <- c(ego_attrs, vertex_label_var)
     }
+    # Add ego_label_var to ego_attrs if it's different from vertex_label_var
+    if (!is.null(ego_label_var) && 
+        ego_label_var %in% names(x$ego) && 
+        !identical(ego_label_var, vertex_label_var)) {
+      ego_attrs <- c(ego_attrs, ego_label_var)
+    }
+    if (!is.null(vertex_color_var) && vertex_color_var %in% names(x$ego)) {
+      ego_attrs <- c(ego_attrs, vertex_color_var)
+    }
+    # Add ego_color_var to ego_attrs if it's different from vertex_color_var
+    ego_needs_separate_color_var <- !is.null(ego_color_var) && 
+                                     ego_color_var %in% names(x$ego) && 
+                                     !identical(ego_color_var, vertex_color_var)
+    if (ego_needs_separate_color_var) {
+      ego_attrs <- c(ego_attrs, ego_color_var)
+    }
   }
 
   gr <- as_igraph(x, 
@@ -138,17 +166,53 @@ plot_one_ego_graph <- function(x,
     #vertex.color[is.na(vertex.color)] <- 0
     vertex.color <- factor(vertex.color)
     colors_ <- egor_col_pal(vertex_color_palette,
-                            length(levels(
-                              factor(igraph::vertex_attr(gr,
-                                                                  vertex_color_var))
-                            )))
+                            length(levels(vertex.color)))
     clrs <- colors_[vertex.color]
     clrs[is.na(clrs)] <- "#ffffff"
   } else {
     vertex.color <- 1
     clrs <- "coral"
   }
 
+  # Ego Color (if include_ego is TRUE and ego_color_var is specified)
+  # Note: When include_ego=TRUE, the ego vertex is always added as the last vertex in the igraph
+  if (include_ego && !is.null(ego_color_var)) {
+    # Determine if ego needs separate coloring
+    ego_has_diff_color_config <- !identical(ego_color_var, vertex_color_var) || 
+                                  !identical(ego_color_palette, vertex_color_palette)
+
+    if (ego_has_diff_color_config) {
+      # If ego_color_var and vertex_color_var are the same variable but different palettes
+      if (identical(ego_color_var, vertex_color_var)) {
+        # Same variable, different palette - use the combined levels but apply ego_color_palette to ego
+        ego_colors_ <- egor_col_pal(ego_color_palette,
+                                    length(levels(vertex.color)))
+        # The last vertex is always ego when include_ego=TRUE
+        clrs[length(clrs)] <- ego_colors_[vertex.color[length(vertex.color)]]
+      } else {
+        # Different variables - ego_color_var should be an ego-level attribute
+        # Check if the attribute exists in the graph (it should for the ego vertex)
+        if (ego_color_var %in% igraph::vertex_attr_names(gr)) {
+          ego_color_values <- igraph::vertex_attr(gr, ego_color_var)
+          # Get unique non-NA values to determine factor levels
+          unique_ego_values <- unique(ego_color_values[!is.na(ego_color_values)])
+          ego_color_factor <- factor(ego_color_values, levels = unique_ego_values)
+          ego_colors_ <- egor_col_pal(ego_color_palette,
+                                      length(levels(ego_color_factor)))
+          # Apply ego color only to the last vertex (ego)
+          ego_idx <- length(clrs)
+          if (!is.na(ego_color_factor[ego_idx])) {
+            clrs[ego_idx] <- ego_colors_[ego_color_factor[ego_idx]]
+          } else {
+            clrs[ego_idx] <- "#ffffff"
+          }
+        }
+      }
+    }
+    # If ego_color_var and vertex_color_var are identical (including palette),
+    # the ego color is already set correctly by the vertex color logic above
+  }
+
   # Edge Width
   if (!is.null(edge_width_var)) {
     edge.width <-
@@ -185,6 +249,18 @@ plot_one_ego_graph <- function(x,
     vertex.label <- ""
   }
 
+  # Ego Label (if include_ego is TRUE and ego_label_var is different)
+  if (include_ego && !is.null(ego_label_var) && !identical(ego_label_var, vertex_label_var)) {
+    # Check if the ego_label_var attribute exists in the graph
+    if (ego_label_var %in% igraph::vertex_attr_names(gr)) {
+      ego_label_value <- igraph::vertex_attr(gr, ego_label_var)[length(igraph::V(gr))]
+      # Only set if not NA
+      if (!is.na(ego_label_value)) {
+        vertex.label[length(vertex.label)] <- ego_label_value
+      }
+    }
+  }
+
   par(mar = c(0.5, 0.5, 0.5, 0.5))
   if (!is.null(vertex_color_var))
     par(mar = c(0.5, 5, 0.5, 0.5))
@@ -200,17 +276,20 @@ plot_one_ego_graph <- function(x,
     # Set curvature of ego-alter ties to zero
     # igraph::E(gr)$curved[is.na(igraph::E(gr)$curved)] <- 0
     # Set ego-alter weights to a dummy value
-    if (any(!is.na(igraph::E(gr)$weight))) {
-      # Set to min of other weights, so scale of weights is comparable
-      igraph::E(gr)$weight[is.na(igraph::E(gr)$weight)] <- min(igraph::E(gr)$weight, na.rm = TRUE)
-    } else {
-      # no other weights in the graph, so just set a hardwired dummy value
-      # if there is no weight variable at all, E(gr)$weight will be NULL rather than a vector,
-      # so the syntax igraph::E(gr)$weight[is.na(igraph::E(gr)$weight)] will fail. Since all weights
-      # are NULL or NA at this point, set them all to 1. This will change missing aatie weights to 1,
-      # which may not be desirable, but overwriting missing aatie weights is also the behavior of the
-      # code above when there is at least one nonmissing aatie weight
-      igraph::E(gr)$weight <- 1
+    # Only set edge attributes if there are edges
+    if (igraph::ecount(gr) > 0) {
+      if (any(!is.na(igraph::E(gr)$weight))) {
+        # Set to min of other weights, so scale of weights is comparable
+        igraph::E(gr)$weight[is.na(igraph::E(gr)$weight)] <- min(igraph::E(gr)$weight, na.rm = TRUE)
+      } else {
+        # no other weights in the graph, so just set a hardwired dummy value
+        # if there is no weight variable at all, E(gr)$weight will be NULL rather than a vector,
+        # so the syntax igraph::E(gr)$weight[is.na(igraph::E(gr)$weight)] will fail. Since all weights
+        # are NULL or NA at this point, set them all to 1. This will change missing aatie weights to 1,
+        # which may not be desirable, but overwriting missing aatie weights is also the behavior of the
+        # code above when there is at least one nonmissing aatie weight
+        igraph::E(gr)$weight <- 1
+      }
     }
   }
 
@@ -255,19 +334,79 @@ plot_one_ego_graph <- function(x,
       ifelse(vertex_color_legend_label == "",
              vertex_color_var,
              vertex_color_legend_label)
-    legend(
-      x = -1.9,
-      y = 1.1,
-      legend = levels(factor(color_var)),
-      pt.bg = colors_,
-      pt.cex = 1.5,
-      pch = 22,
-      bty = "n",
-      y.intersp = 1,
-      title = title_,
-      xpd = TRUE,
-      cex = font_size
-    )
+
+    # Determine if we need a separate ego legend
+    # Reuse the same condition as in ego color logic
+    ego_has_diff_color_config <- include_ego && !is.null(ego_color_var) && 
+                                 (!identical(ego_color_var, vertex_color_var) || 
+                                  !identical(ego_color_palette, vertex_color_palette))
+
+    if (ego_has_diff_color_config && !identical(ego_color_var, vertex_color_var)) {
+      # Different variables: show two legends stacked
+      # Constants for legend spacing
+      LEGEND_ITEM_HEIGHT <- 0.15
+      LEGEND_VERTICAL_GAP <- 0.3
+
+      # First show alter/vertex legend
+      # Note: ego is always the last vertex when include_ego=TRUE
+      alter_color_var <- color_var[-length(color_var)]  # Exclude ego
+      legend(
+        x = -1.9,
+        y = 1.1,
+        legend = levels(factor(alter_color_var)),
+        pt.bg = colors_,
+        pt.cex = 1.5,
+        pch = 22,
+        bty = "n",
+        y.intersp = 1,
+        title = paste0("Alter: ", title_),
+        xpd = TRUE,
+        cex = font_size
+      )
+
+      # Now show ego legend below
+      ego_color_var_values <- igraph::vertex_attr(gr, ego_color_var)
+      ego_color_factor <- factor(ego_color_var_values)
+      ego_colors_ <- egor_col_pal(ego_color_palette,
+                                  length(levels(ego_color_factor)))
+      ego_title_ <-
+        ifelse(ego_color_legend_label == "",
+               ego_color_var,
+               ego_color_legend_label)
+
+      # Calculate y position for second legend
+      n_alter_levels <- length(levels(factor(alter_color_var)))
+      y_offset <- 1.1 - (n_alter_levels * LEGEND_ITEM_HEIGHT) - LEGEND_VERTICAL_GAP
+
+      legend(
+        x = -1.9,
+        y = y_offset,
+        legend = levels(ego_color_factor),
+        pt.bg = ego_colors_,
+        pt.cex = 1.5,
+        pch = 22,
+        bty = "n",
+        y.intersp = 1,
+        title = paste0("Ego: ", ego_title_),
+        xpd = TRUE,
+        cex = font_size
+      )
+    } else {
+      # Same variable or same variable with different palette: show single legend
+      legend(
+        x = -1.9,
+        y = 1.1,
+        legend = levels(factor(color_var)),
+        pt.bg = colors_,
+        pt.cex = 1.5,
+        pch = 22,
+        bty = "n",
+        y.intersp = 1,
+        title = title_,
+        xpd = TRUE,
+        cex = font_size
+      )
+    }
   }
   par(mar = c(0.5, 0.5, 0.5, 0.5))
   graphics::box(lty = 'solid', col = highlight_box_col, lwd = 5)