R/sm_index.R

In ammiBayes: Bayesian Ammi Model for Continuous Data with or without Additive and Dominance Effect

#' Bayesian Stability Mahalanobis Distance (SM_i) Index
#'
#' Computes the Bayesian Stability Mahalanobis Distance index for genotypes
#' based on MCMC samples of interaction principal component scores from ammiBayes.
#'
#' Nascimento, A. C. C., Nascimento, M., Sagae, V. S., & Jarquín, D. (2025). 
#' Bayesian AMMI-based indexes for genotype selection: Integrating novel stability 
#' measures for enhanced G x E inference. 
#' \emph{Crop Science}, e20732.

# Plot Method for Bayesian Stability Mahalanobis Distance (SM_i) with Quadrants



sm.index <- function(object, prob = 0.95) {
  if (!inherits(object, "ammiBayes")) {
    stop("Object must be of class 'ammiBayes'")
  }
  
  # Obtém os BLUPs e seus intervalos HPD prontos
  gen_ef <- gen.effects(object)
  n.gen <- nrow(gen_ef)
  
  chains <- object$output
  ch <- chains[[1]]
  
  atu_names <- grep("^atu[0-9]+$", names(ch), value = TRUE)
  n.axes <- length(atu_names)
  
  if (n.axes < 2) {
    stop("At least 2 AMMI axes are required to compute the Mahalanobis distance index.")
  }
  
  # Extração dos escores médios posteriores ponderados
  scores_mean <- matrix(NA, nrow = n.gen, ncol = n.axes)
  for (k in seq_len(n.axes)) {
    atu_mat <- ch[[paste0("atu", k)]]
    lambda_k_mean <- mean(ch[["L"]][, k])
    scores_mean[, k] <- colMeans(atu_mat) * sqrt(lambda_k_mean)
  }
  
  # Matriz de covariância e inversa baseada nos escores médios
  Sigma_mean <- stats::cov(scores_mean)
  inv_Sigma <- tryCatch(solve(Sigma_mean), error = function(e) MASS::ginv(Sigma_mean))
  
  # Cálculo do SM_i para cada genótipo
  sm_values <- numeric(n.gen)
  for (i in seq_len(n.gen)) {
    alpha_i <- scores_mean[i, , drop = FALSE]
    sm_values[i] <- sqrt(as.numeric(alpha_i %*% inv_Sigma %*% t(alpha_i)))
  }
  
  gen_labels <- rownames(gen_ef)
  if (is.null(gen_labels)) gen_labels <- paste0("G", seq_len(n.gen))
  
  gen_means <- if (!is.null(object$means)) {
    rowMeans(object$means, na.rm = TRUE)
  } else {
    gen_ef$Mean
  }
  
  # Margem de erro proporcional baseada na amplitude dos efeitos genotípicos
  error_margin <- abs(gen_ef$`97.5%` - gen_ef$`2.5%`) / 2
  
  # Montagem final integrando os limites ajustados à escala do SM_i
  results_mat <- data.frame(
    Genotype = gen_labels,
    YieldMean = as.numeric(gen_means),
    Mean = sm_values,
    Median = sm_values,
    Lower = pmax(0, sm_values - error_margin), 
    Upper = sm_values + error_margin,
    stringsAsFactors = FALSE
  )
  
  rownames(results_mat) <- gen_labels
  
  out <- list(
    statistics = results_mat,
    gen.effects = gen_ef,
    prob = prob,
    call = match.call()
  )
  class(out) <- "ammiBayes.sm"
  return(out)
}

ammiBayes.sm.plot <- function(x, 
                               xlab = "Genotype Posterior Mean",
                               ylab = "Bayesian Stability Index (SM_i)",
                               line.col = "darkblue", main = NULL, text.col = "red",
                               pch = 16, cex = 1.5, lwd = 1.5, lty = 2,
                               ylimits = 0.1, grid.col = "gray50")
 {
  if (!inherits(x, "ammiBayes.sm")) {
    stop("Object must be of class 'ammiBayes.sm'")
  }
  
  df <- x$statistics

	limits <- c(min(df$Lower)-ylimits, max(df$Upper)+ylimits)
  
  xyplot(Mean ~ YieldMean, data = df,
         xlab = xlab,
         ylab = ylab,
         main = main,
         pch = pch,
         line.col = line.col,
				 text.col = text.col,
				 ylim=limits,
         panel = function(x, y, subscripts, ...) {

           panel.abline(v = median(x, na.rm = TRUE), lty = lty, col = grid.col)
           panel.abline(h = median(y, na.rm = TRUE), lty = lty, col = grid.col)
           

           lower <- df$Lower[subscripts]
           upper <- df$Upper[subscripts]
           

           panel.segments(x, lower, x, upper, col = line.col, lwd = lwd)

           panel.xyplot(x, y, subscripts = subscripts)
           

           gen_names <- df$Genotype[subscripts]
           panel.text(x, y, labels = gen_names, pos = 3, cex = cex, col = text.col)
         })
}