R/CCMnet_Sample_methods.R

In CCMnet: Congruence Class Models for Networks

#' Methods for ccm_sample Objects
#'
#' @description 
#' Printing, summarizing, and plotting methods for results generated by \code{sample_ccm}.
#' 
#' @param x,object An object of class \code{ccm_sample}.
#' @param stats Character vector of statistic names to plot. If \code{NULL}, 
#'   all targeted statistics are plotted.
#' @param type Character string specifying the plot type: \code{"density"}, 
#'   \code{"hist"}, or \code{"trace"}.
#' @param include_theoretical Logical. If \code{TRUE}, overlays the theoretical 
#'   target distribution (requires running \code{sample_theoretical} first).
#' @param ... Additional arguments passed to methods.
#'
#' @details 
#' For \code{type = "trace"}, setting \code{include_theoretical = TRUE} adds a 
#' red dashed line for the theoretical mean and red dotted lines for the 
#' 2.5\% and 97.5\% quantiles.
#'
#' @name plot.ccm_sample
NULL

#' @rdname plot.ccm_sample
#' @import ggplot2
#' @importFrom tidyr pivot_longer
#' @importFrom dplyr mutate row_number group_by summarise select
#' @importFrom stats quantile density
#' @export
plot.ccm_sample <- function(x,
                            stats = NULL,
                            type = c("density", "hist", "trace"),
                            include_theoretical = FALSE,
                            ...) {
  
  type <- match.arg(type)
  fit <- x
  
  # Default: plot all columns if stats is NULL
  if (is.null(stats)) {
    stats <- colnames(fit$mcmc_stats)
  }
  
  # --- 1. TRACEPLOT LOGIC ---
  if (type == "trace") {
    # Prepare MCMC trace data
    df_trace <- fit$mcmc_stats %>%
      as.data.frame() %>%
      mutate(iter = row_number()) %>%
      pivot_longer(cols = all_of(stats), names_to = "stat", values_to = "value")
    
    p <- ggplot(df_trace, aes(x = iter, y = value)) +
      geom_line(color = "gray30") + 
      facet_wrap(~stat, scales = "free_y", labeller = labeller(stat = facet_labeller)) +
      labs(x = "Iteration", y = "Count", title = "") +
      theme_bw() + 
      theme_ccm_plot() # Using a helper for consistency
    
    # Overlay Theoretical H-Lines (Mean + Quantiles)
    if (include_theoretical) {
      if (is.null(fit$theoretical$theory_stats)) {
        warning("No theoretical distribution available for traceplot lines.")
      } else {
        # Calculate summary stats from theory
        df_theory_lines <- fit$theoretical$theory_stats %>%
          as.data.frame() %>%
          select(all_of(stats)) %>%
          pivot_longer(cols = everything(), names_to = "stat", values_to = "val") %>%
          group_by(.data$stat) %>%
          summarise(
            mean_val = mean(.data$val),
            lwr = quantile(.data$val, 0.025),
            upr = quantile(.data$val, 0.975),
            .groups = "drop"
          )
        
        p <- p + 
          geom_hline(data = df_theory_lines, aes(yintercept = .data$mean_val), 
                     color = "red", linetype = "dashed", linewidth = 0.8) +
          geom_hline(data = df_theory_lines, aes(yintercept = .data$lwr), 
                     color = "red", linetype = "dotted", alpha = 0.8) +
          geom_hline(data = df_theory_lines, aes(yintercept = .data$upr), 
                     color = "red", linetype = "dotted", alpha = 0.8)
      }
    }
    return(p)
  }
  
  # --- 2. DENSITY / HIST LOGIC ---
  # Prepare MCMC data
  df_mcmc <- fit$mcmc_stats %>%
    as.data.frame() %>%
    pivot_longer(cols = all_of(stats), names_to = "stat", values_to = "count") %>%
    mutate(source = "MCMC")
  
  df_plot <- df_mcmc
  
  if (include_theoretical && !is.null(fit$theoretical$theory_stats)) {
    df_theory <- fit$theoretical$theory_stats %>%
      as.data.frame() %>%
      pivot_longer(cols = all_of(stats), names_to = "stat", values_to = "count") %>%
      mutate(source = "Theoretical")
    df_plot <- bind_rows(df_mcmc, df_theory)
  }
  
  p <- ggplot(df_plot, aes(x = count, color = source, fill = source)) +
    facet_wrap(~stat, scales = "free_y", labeller = labeller(stat = facet_labeller)) +
    labs(x = "Count", y = ifelse(type == "density", "Density", "Frequency")) +
    theme_bw() + 
    theme_ccm_plot()
  
  if (type == "hist") {
    p <- p + geom_histogram(aes(y = after_stat(density)), 
                            alpha = 0.5, 
                            position = "identity", 
                            binwidth = 1,
                            color = NA)
  } else {
    p <- p + geom_density(alpha = 0.5)
  }
  
  return(p)
}

# Helper to keep theme consistent across types
theme_ccm_plot <- function() {
  theme(
    legend.position = "bottom",
    legend.title = element_blank(),
    panel.grid.major = element_blank(), 
    panel.grid.minor = element_blank(),
    strip.background = element_rect(fill = "lightgray", color = "black"),
    strip.text = element_text(size = 14),
    axis.title = element_text(size = 12)
  )
}

#' @rdname plot.ccm_sample
#' @export
print.ccm_sample <- function(x, ...) {
  cat("Object of class 'ccm_sample'\n")
  cat("-------------------------\n")
  if (!is.null(x$network_stats)) {
    cat("Statistics:       ", paste(unlist(x$network_stats), collapse=", "), "\n")
  }
  if (!is.null(x$prob_distr)) {
    cat("Distribution(s):  ", paste(unlist(x$prob_distr), collapse=", "), "\n")
  }
  if (!is.null(x$population)) {
    cat("Population:       ", x$population, "\n")
  }
  if (!is.null(x$mcmc_stats)) {
    cat("MCMC samples:     ", nrow(x$mcmc_stats), "rows x", ncol(x$mcmc_stats), "cols\n")
  }
  invisible(x)
}

#' @rdname plot.ccm_sample
#' @export
summary.ccm_sample <- function(object, ...) {
  cat("Summary of ccm_sample object\n")
  cat("-------------------------\n")
  if (!is.null(object$mcmc_stats)) {
    for (i in seq_len(ncol(object$mcmc_stats))) {
      nm <- if (!is.null(colnames(object$mcmc_stats))) colnames(object$mcmc_stats)[i] else paste0("V", i)
      cat("\nStatistic: ", nm, "\n")
      vals <- object$mcmc_stats[, i]
      print(summary(vals))
    }
  } else {
    cat("No stats available\n")
  }
  invisible(object)
}