LUCID Early Tutorial: Sankey Plot with Balanced Top Omics

1) Goal

A LUCID fit’s Sankey diagram (plot(), in the other vignettes) shows the package’s own default view of exposure -> cluster -> outcome. This tutorial builds a custom Sankey view on top of a fitted model for a situation plot() doesn’t handle out of the box: multiple named omics families concatenated into one early-integration matrix (here, microbiota and metabolite features), where a meaningful figure needs to show a balanced number of top features from each family side by side, not just the overall top-N (which could easily be dominated by one family). It walks through:

  1. Simulating two omics families – microbiota (X0000001, X0000002, …) and metabolites (hilic0001, c180001, …) – concatenated into one Z matrix for early integration.
  2. Fitting a LUCID early model with estimate_lucid().
  3. Ranking and selecting a balanced top-N per family, and plotting the result as a Sankey diagram.
  4. Contrasting that dataset-specific ranking with the package’s general-purpose get_top_omics_features().

2) Setup

required_pkgs <- c("plotly")
missing_pkgs <- required_pkgs[!vapply(required_pkgs, requireNamespace, logical(1), quietly = TRUE)]
if (length(missing_pkgs) > 0) {
  stop("Please install required packages first: ", paste(missing_pkgs, collapse = ", "))
}

library(LUCIDus)
library(plotly)  # for %>% and plot_ly(), used directly by this vignette's own Sankey code

data(sim_data)

3) Simulate a Two-Layer Omics Matrix for Early Integration

As in the other vignettes, the simulation deliberately builds in a signal: a latent state (latent_signal) shifts a subset of features in each family apart by cluster, on top of a shared dependence on one exposure (g_anchor), so the fitted model below has real structure to find and rank. The two metabolite sub-blocks (hilic, c18) are concatenated into one metabolites block before joining microbiota into the single Z matrix early integration expects.

set.seed(1008)

n <- 200

G <- sim_data$G[seq_len(n), , drop = FALSE]
colnames(G) <- paste0("exposure", seq_len(ncol(G)))

Y <- as.numeric(sim_data$Y_normal[seq_len(n), 1])
CoY <- sim_data$Covariate[seq_len(n), , drop = FALSE]
colnames(CoY) <- paste0("covariate", seq_len(ncol(CoY)))

latent_signal <- ifelse(sim_data$X[seq_len(n)] == 1, -1, 1)
g_anchor <- as.numeric(scale(G[, 1]))

simulate_omics_block <- function(feature_names, signal_n, latent_signal, g_anchor,
                                 cluster_scale, noise_sd, offset = 0) {
  n_obs <- length(latent_signal)
  p <- length(feature_names)
  block <- matrix(NA_real_, nrow = n_obs, ncol = p)

  signal_loadings <- seq(cluster_scale, 0.8, length.out = signal_n)

  for (j in seq_len(p)) {
    direction <- if (j %% 2L == 0L) 1 else -1
    cluster_loading <- if (j <= signal_n) {
      direction * signal_loadings[j]
    } else {
      runif(1, -0.20, 0.20)
    }
    exposure_loading <- if (j <= signal_n) {
      direction * runif(1, 0.15, 0.35)
    } else {
      runif(1, -0.10, 0.10)
    }

    block[, j] <- offset +
      cluster_loading * latent_signal +
      exposure_loading * g_anchor +
      rnorm(n_obs, sd = noise_sd)
  }

  block <- scale(block)
  colnames(block) <- feature_names
  block
}

n_microbiota <- 18
n_hilic <- 12
n_c18 <- 8

microbiota_names <- paste0("X", sprintf("%07d", seq_len(n_microbiota)))
hilic_names <- paste0("hilic", sprintf("%04d", seq_len(n_hilic)))
c18_names <- paste0("c18", sprintf("%04d", seq_len(n_c18)))

microbiota <- simulate_omics_block(
  feature_names = microbiota_names,
  signal_n = 6,
  latent_signal = latent_signal,
  g_anchor = g_anchor,
  cluster_scale = 1.80,
  noise_sd = 0.85
)

hilic <- simulate_omics_block(
  feature_names = hilic_names,
  signal_n = 5,
  latent_signal = latent_signal,
  g_anchor = g_anchor,
  cluster_scale = 1.55,
  noise_sd = 0.90,
  offset = 0.15
)

c18 <- simulate_omics_block(
  feature_names = c18_names,
  signal_n = 4,
  latent_signal = latent_signal,
  g_anchor = g_anchor,
  cluster_scale = 1.35,
  noise_sd = 0.95,
  offset = -0.10
)

metabolites <- cbind(hilic, c18)
Z <- cbind(microbiota, metabolites)

omics_dimensions <- data.frame(
  block = c("microbiota", "hilic", "c18", "metabolites_total", "combined_Z"),
  n_features = c(ncol(microbiota), ncol(hilic), ncol(c18), ncol(metabolites), ncol(Z))
)
omics_dimensions
##               block n_features
## 1        microbiota         18
## 2             hilic         12
## 3               c18          8
## 4 metabolites_total         20
## 5        combined_Z         38

combined_Z’s feature count is the sum of microbiota and metabolites_total, confirming the concatenation produced the single matrix early integration expects – one joint cluster variable will be fit from all 38 features together, not one per family.

4) Fit a LUCID Early Model

estimate_lucid(..., lucid_model = "early") fits one latent cluster variable, K = 2 here, from the pooled G/Z/Y. res_Mu (K x M) holds the fitted cluster-specific mean for every omics feature – the ranking used throughout this tutorial reads directly from this matrix.

fit_early <- estimate_lucid(
  G = G,
  Z = Z,
  Y = Y,
  lucid_model = "early",
  CoY = CoY,
  family = "normal",
  K = 2,
  seed = 1008,
  max_itr = 20,
  max_tot.itr = 60,
  tol = 1e-2
)
## Fitting LUCID early model (K = 2)...
## Finished LUCID early model.
class(fit_early)
## [1] "early_lucid"
dim(fit_early$res_Mu)
## [1]  2 38

dim(fit_early$res_Mu) should read 2 x 38: one row per cluster, one column per omics feature (the same 38 confirmed above).

5) Sankey Function with Equal Top N per Omics Family

The Sankey below has three link stages – exposure -> cluster, cluster -> omics feature, cluster -> outcome – built directly from the fitted model’s own parameters (res_Beta for the first, res_Mu for the second, res_Gamma for the third); no new estimation happens here, only reshaping and coloring for the plot. select_top_omics_balanced() is the piece that makes the “balanced” part work: it scores every feature by how far apart its two cluster means are (abs(mu[c1,] - mu[c2,])) and keeps the top top_n_each within each omics family separately, so a family with many features can’t crowd out a smaller one just by having more entries to rank.

detect_omics_family <- function(x) {
  out <- rep("exposure", length(x))
  out[grepl("^Latent Cluster", x)] <- "lc"
  out[grepl("^X\\d{7}$", x)] <- "microbiota"
  out[grepl("^(hilic|c18)", x, ignore.case = TRUE)] <- "metabolite"
  out
}

select_top_omics_balanced <- function(mu_mat, z_names, cluster_pair, top_n_each = NULL) {
  if (is.null(top_n_each)) {
    return(z_names)
  }

  if (!is.numeric(top_n_each) || length(top_n_each) != 1 || is.na(top_n_each) || top_n_each < 1) {
    stop("`top_n_each` must be a single positive integer or NULL.")
  }

  top_n_each <- as.integer(top_n_each)
  c1 <- as.integer(cluster_pair[1])
  c2 <- as.integer(cluster_pair[2])
  score <- abs(mu_mat[c1, ] - mu_mat[c2, ])
  family <- detect_omics_family(z_names)

  ordered_keep <- character(0)
  for (fam in c("microbiota", "metabolite")) {
    fam_idx <- which(family == fam)
    if (!length(fam_idx)) {
      next
    }
    fam_idx <- fam_idx[order(score[fam_idx], decreasing = TRUE)]
    ordered_keep <- c(
      ordered_keep,
      z_names[fam_idx[seq_len(min(top_n_each, length(fam_idx)))]]
    )
  }

  ordered_keep
}

make_link_color <- function(target_group, is_positive) {
  pos_cols <- c(
    lc = "#8E9BAE",
    microbiota = "#1B7F5A",
    metabolite = "#C46A1B",
    outcome = "#B23A48",
    exposure = "#5E6F86"
  )
  neg_cols <- c(
    lc = "#DCE3EA",
    microbiota = "#BFE2D6",
    metabolite = "#F2D4B5",
    outcome = "#E9C0C6",
    exposure = "#CCD6E2"
  )

  ifelse(is_positive, pos_cols[target_group], neg_cols[target_group])
}

sankey_early_integration <- function(
    lucid_fit1,
    top_n_each = 5,
    cluster_pair = c(1, 2),
    text_size = 13,
    drop_covariates = TRUE
) {
  if (!inherits(lucid_fit1, "early_lucid")) {
    stop("`lucid_fit1` must inherit from class 'early_lucid'.")
  }

  needed <- c("K", "var.names", "res_Beta", "res_Mu", "res_Gamma")
  miss <- needed[!needed %in% names(lucid_fit1)]
  if (length(miss) > 0) {
    stop("Missing fields in `lucid_fit1`: ", paste(miss, collapse = ", "))
  }

  K <- lucid_fit1$K
  var.names <- lucid_fit1$var.names
  Gnames <- var.names$Gnames
  Znames <- var.names$Znames
  Yname <- var.names$Ynames

  beta_mat <- as.matrix(lucid_fit1$res_Beta)
  mu_mat <- as.matrix(lucid_fit1$res_Mu)

  if (nrow(beta_mat) != K) stop("Unexpected `res_Beta` shape.")
  if (nrow(mu_mat) != K) stop("Unexpected `res_Mu` shape.")

  if (length(cluster_pair) != 2) stop("`cluster_pair` must be length-2.")
  c1 <- as.integer(cluster_pair[1])
  c2 <- as.integer(cluster_pair[2])
  if (c1 < 1 || c1 > K || c2 < 1 || c2 > K || c1 == c2) {
    stop("`cluster_pair` must reference two distinct cluster indices in 1..K.")
  }

  keep_z <- select_top_omics_balanced(
    mu_mat = mu_mat,
    z_names = Znames,
    cluster_pair = cluster_pair,
    top_n_each = top_n_each
  )

  valueGtoX <- as.vector(t(beta_mat[, -1, drop = FALSE]))

  mu_use <- mu_mat[, Znames %in% keep_z, drop = FALSE]
  mu_use <- mu_use[, keep_z, drop = FALSE]
  z_use <- colnames(mu_use)
  valueXtoZ <- as.vector(t(mu_use))

  gamma_beta <- if (!is.null(lucid_fit1$res_Gamma$beta)) {
    as.numeric(lucid_fit1$res_Gamma$beta)
  } else {
    as.numeric(lucid_fit1$res_Gamma$fit$coefficients)
  }
  if (length(gamma_beta) < K) stop("Could not extract cluster->outcome effects.")

  valueXtoY <- gamma_beta[seq_len(K)]
  valueXtoY[1] <- 0

  GtoX <- data.frame(
    source = rep(Gnames, K),
    target = paste0("Latent Cluster", rep(seq_len(K), each = length(Gnames))),
    value = abs(valueGtoX),
    group = valueGtoX > 0,
    stringsAsFactors = FALSE
  )

  XtoZ <- data.frame(
    source = paste0("Latent Cluster", rep(seq_len(K), each = length(z_use))),
    target = rep(z_use, K),
    value = abs(valueXtoZ),
    group = valueXtoZ > 0,
    stringsAsFactors = FALSE
  )

  XtoY <- data.frame(
    source = paste0("Latent Cluster", seq_len(K)),
    target = rep(Yname, K),
    value = abs(valueXtoY),
    group = valueXtoY > 0,
    stringsAsFactors = FALSE
  )

  links <- rbind(GtoX, XtoZ, XtoY)

  if (drop_covariates) {
    links <- links[!grepl("^covariate", links$source, ignore.case = TRUE), , drop = FALSE]
  }

  nodes <- data.frame(name = unique(c(links$source, links$target)), stringsAsFactors = FALSE)
  nodes$group <- detect_omics_family(nodes$name)
  nodes$group[nodes$name %in% Yname] <- "outcome"

  group_order <- c("exposure", "lc", "microbiota", "metabolite", "outcome")
  nodes$group <- factor(nodes$group, levels = group_order)
  nodes <- nodes[order(nodes$group, nodes$name), , drop = FALSE]
  nodes$group <- as.character(nodes$group)

  node_color_map <- c(
    exposure = "#4C6A92",
    lc = "#B9C6D3",
    microbiota = "#2A9D6F",
    metabolite = "#D98C3F",
    outcome = "#C44536"
  )
  nodes$color <- unname(node_color_map[nodes$group])

  links$IDsource <- match(links$source, nodes$name) - 1L
  links$IDtarget <- match(links$target, nodes$name) - 1L

  target_group <- nodes$group[match(links$target, nodes$name)]
  target_group[grepl("^Latent Cluster", links$target)] <- "lc"
  links$link_color <- make_link_color(target_group, links$group)

  x_map <- c(
    exposure = 0.02,
    lc = 0.34,
    microbiota = 0.68,
    metabolite = 0.80,
    outcome = 0.97
  )
  nodes$x <- unname(x_map[nodes$group])

  nodes$y <- NA_real_
  for (g in group_order) {
    idx <- which(nodes$group == g)
    if (!length(idx)) {
      next
    }
    nodes$y[idx] <- if (length(idx) == 1) 0.5 else seq(0.03, 0.97, length.out = length(idx))
  }

  plotly::plot_ly(
    type = "sankey",
    orientation = "h",
    node = list(
      label = nodes$name,
      color = nodes$color,
      pad = 16,
      thickness = 18,
      line = list(color = "black", width = 0.35),
      x = nodes$x,
      y = nodes$y
    ),
    link = list(
      source = links$IDsource,
      target = links$IDtarget,
      value = pmax(links$value, 1e-12),
      color = links$link_color
    )
  ) %>%
    plotly::layout(
      font = list(size = text_size),
      margin = list(l = 20, r = 20, t = 30, b = 20)
    )
}

6) Top Features Selected per Family

selected_top <- select_top_omics_balanced(
  mu_mat = fit_early$res_Mu,
  z_names = fit_early$var.names$Znames,
  cluster_pair = c(1, 2),
  top_n_each = 5
)

data.frame(
  feature = selected_top,
  family = detect_omics_family(selected_top),
  stringsAsFactors = FALSE
)
##      feature     family
## 1   X0000001 microbiota
## 2   X0000006 microbiota
## 3   X0000003 microbiota
## 4   X0000002 microbiota
## 5   X0000005 microbiota
## 6    c180003 metabolite
## 7  hilic0003 metabolite
## 8    c180002 metabolite
## 9  hilic0002 metabolite
## 10 hilic0001 metabolite

select_top_omics_balanced() above is specific to this dataset: it scores features by their difference between one pair of clusters and balances the count kept from each omics family by name pattern. The package’s general-purpose ranking function, get_top_omics_features(), instead scores every feature by its separation across all clusters (the same criterion plot_cluster_omic_profile() uses) with no family balancing or pairwise comparison:

top_general <- get_top_omics_features(fit_early, top_n = 10)
top_general
## $Omics
##  X0000001   c180003 hilic0003  X0000002  X0000003 hilic0001 hilic0002   c180002 
## 1.1267944 1.0120169 0.9964746 0.9627817 0.9452440 0.9064645 0.9024558 0.8979317 
##  X0000004  X0000005 
## 0.8701005 0.8416784

7) Display the Plot

Read either figure left to right: exposure nodes flow into the two latent cluster nodes (link width proportional to |G -> X coefficient|, color by sign), the clusters flow into omics feature nodes (width by |cluster mean|), and the clusters flow into the outcome node. A feature whose incoming link is wide is one the model estimates a large cluster-mean difference for; a thin link means that feature barely differs between clusters despite appearing in the diagram.

7.1 All omics in the concatenated early-integration matrix

top_n_each = NULL keeps every one of the 38 features, so this is a complete but visually dense view – useful as a first look, less useful for picking out which features actually matter.

fig_all <- sankey_early_integration(
  lucid_fit1 = fit_early,
  top_n_each = NULL,
  cluster_pair = c(1, 2),
  text_size = 11
)
fig_all

7.2 Top 5 microbiota and top 5 metabolites

Setting top_n_each = 5 trims the diagram to the 5 most cluster-differentiating features from each family (10 total omics nodes), which is where the “balanced” ranking in section 5 matters: without it, a simple overall top-10 could plausibly be all microbiota (or all metabolites) if one family’s effects happen to be larger in scale, hiding the other family’s real signal entirely from the figure.

fig_top_balanced <- sankey_early_integration(
  lucid_fit1 = fit_early,
  top_n_each = 5,
  cluster_pair = c(1, 2),
  text_size = 13
)
fig_top_balanced

8) Notes