## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 4.1,
  dpi = 144,
  dev.args = list(bg = "white")
)
library(drmTMB)

relmat_guide_theme <- function() {
  ggplot2::theme_minimal(base_size = 11) +
    ggplot2::theme(
      panel.grid.minor = ggplot2::element_blank(),
      axis.title = ggplot2::element_text(colour = "grey15"),
      axis.text = ggplot2::element_text(colour = "grey25"),
      plot.title = ggplot2::element_text(
        face = "bold",
        colour = "grey10",
        margin = ggplot2::margin(b = 4)
      ),
      plot.subtitle = ggplot2::element_text(
        colour = "grey30",
        margin = ggplot2::margin(b = 8)
      ),
      plot.background = ggplot2::element_rect(fill = "white", colour = NA),
      panel.background = ggplot2::element_rect(fill = "white", colour = NA),
      legend.position = "bottom"
    )
}

relmat_eye_theme <- function() {
  relmat_guide_theme() +
    ggplot2::theme(
      panel.grid.major.y = ggplot2::element_blank(),
      legend.position = "none"
    )
}

relatedness_heatmap_data <- function(x) {
  out <- as.data.frame(as.table(x), stringsAsFactors = FALSE)
  names(out) <- c("row", "column", "relatedness")
  out$row <- factor(out$row, levels = rev(rownames(x)))
  out$column <- factor(out$column, levels = colnames(x))
  out
}

simulate_relmat_guide_data <- function(seed = 20260524) {
  set.seed(seed)
  n_line <- 24
  n_each <- 5
  line_levels <- paste0("line_", seq_len(n_line))
  line <- factor(rep(line_levels, each = n_each))
  temperature <- runif(length(line), -1, 1)
  treatment <- factor(rep(rep(c("control", "heated"), each = n_each), n_line / 2))

  K <- outer(seq_len(n_line), seq_len(n_line), function(i, j) {
    0.36^abs(i - j)
  })
  diag(K) <- diag(K) + 0.08
  dimnames(K) <- list(line_levels, line_levels)
  Q <- solve(K)

  line_effect <- as.vector(t(chol(K)) %*% rnorm(n_line)) * 0.42
  names(line_effect) <- line_levels
  sigma <- exp(-1.20 + 0.10 * (treatment == "heated"))
  seed_mass <- 2.2 + 0.50 * temperature + 0.25 * (treatment == "heated") +
    line_effect[line] + rnorm(length(line), sd = sigma)

  list(
    data = data.frame(
      seed_mass = seed_mass,
      temperature = temperature,
      treatment = treatment,
      line = line
    ),
    K = K,
    Q = Q
  )
}

simulate_relmat_q2_guide_data <- function(seed = 1) {
  set.seed(seed)
  n_observation_block <- 40
  n_line <- 10
  n_each <- 4
  observation <- factor(rep(paste0("obs_", seq_len(n_observation_block)), each = n_each))
  sex <- factor(rep(rep(c("female", "male"), each = n_each), n_observation_block / 2))
  age <- runif(length(observation), 0, 2)
  line_levels <- paste0("line_", seq_len(n_line))
  line <- factor(
    rep(rep(line_levels, each = 4), each = n_each),
    levels = line_levels
  )

  K <- outer(seq_len(n_line), seq_len(n_line), function(i, j) {
    0.25^abs(i - j)
  })
  diag(K) <- diag(K) + 0.10
  dimnames(K) <- list(line_levels, line_levels)
  Q <- solve(K)

  z_mass <- rnorm(n_line)
  z_height <- 0.55 * z_mass + sqrt(1 - 0.55^2) * rnorm(n_line)
  u_mass <- as.vector(t(chol(K)) %*% z_mass) * 0.35
  u_height <- as.vector(t(chol(K)) %*% z_height) * 0.30
  names(u_mass) <- line_levels
  names(u_height) <- line_levels

  e_mass <- rnorm(length(line))
  e_height <- 0.10 * e_mass + sqrt(1 - 0.10^2) * rnorm(length(line))
  sigma_mass <- exp(-1.0 + 0.15 * (sex == "male"))
  sigma_height <- exp(-1.1 - 0.10 * (sex == "male"))

  list(
    data = data.frame(
      line = line,
      age = age,
      sex = sex,
      seed_mass = 2.1 + 0.45 * age + 0.20 * (sex == "male") +
        u_mass[line] + sigma_mass * e_mass,
      plant_height = 1.4 + 0.30 * age - 0.12 * (sex == "male") +
        u_height[line] + sigma_height * e_height
    ),
    K = K,
    Q = Q
  )
}

## ----relmat-guide-fit---------------------------------------------------------
relmat_example <- simulate_relmat_guide_data()
relmat_dat <- relmat_example$data
K <- relmat_example$K
Q <- relmat_example$Q

fit_relmat <- drmTMB(
  bf(
    seed_mass ~ temperature + treatment + relmat(1 | line, K = K),
    sigma ~ 1
  ),
  family = gaussian(),
  data = relmat_dat
)

## ----relmat-known-matrix, fig.width = 5.3, fig.height = 4.8, fig.cap = "Known relatedness matrix used by the `relmat()` example; this heatmap shows the supplied latent structure, not model-estimated uncertainty.", fig.alt = "Heatmap of the known relatedness matrix for the relmat example. Values are highest on the diagonal and fade as line identifiers are farther apart in the simulated ordering."----
if (requireNamespace("ggplot2", quietly = TRUE)) {
  relmat_matrix <- relatedness_heatmap_data(K)

  ggplot2::ggplot(
    relmat_matrix,
    ggplot2::aes(column, row, fill = relatedness)
  ) +
    ggplot2::geom_tile() +
    ggplot2::coord_equal() +
    ggplot2::scale_fill_gradientn(
      colours = c("#F7FCF5", "#C7E9C0", "#74C476", "#006D2C"),
      name = "Known\nrelatedness"
    ) +
    relmat_guide_theme() +
    ggplot2::theme(
      axis.text = ggplot2::element_blank(),
      axis.ticks = ggplot2::element_blank(),
      panel.grid = ggplot2::element_blank()
    ) +
    ggplot2::labs(
      title = "Known-matrix input structure",
      subtitle = "Validated relatedness among experimental lines",
      x = "Line",
      y = "Line"
    )
}

## ----relmat-sd-figure, fig.width = 6.6, fig.height = 3.4, fig.cap = "Fitted residual `sigma` and level-specific known-matrix marginal SDs `s sqrt(K[i,i])` from a univariate Gaussian `relmat()` model. Points and bars are response-scale estimates with 95% Wald intervals transformed by the known diagonal multipliers.", fig.alt = "Horizontal interval display comparing residual sigma with known-matrix node marginal standard deviations calculated by multiplying the fitted latent scale and interval endpoints by the square root of each known covariance diagonal."----
if (requireNamespace("ggplot2", quietly = TRUE)) {
  relmat_ci <- confint(fit_relmat, parm = "variance_components")
  relmat_parameters <- summary(fit_relmat)$parameters
  node_multiplier <- sqrt(diag(K))
  relmat_sd <- relmat_ci[c(1, rep(2, length(node_multiplier))), , drop = FALSE]
  relmat_sd$estimate <- c(
    unname(exp(coef(fit_relmat, "sigma")["(Intercept)"])),
    relmat_parameters[
      relmat_parameters$parm == "sd:mu:relmat(1 | line)", "estimate"
    ] * node_multiplier
  )
  relmat_sd$lower[-1] <- relmat_sd$lower[-1] * node_multiplier
  relmat_sd$upper[-1] <- relmat_sd$upper[-1] * node_multiplier
  relmat_sd$label <- c(
    "Residual\nsigma",
    rep("relmat node\nmarginal SD", length(node_multiplier))
  )

  ggplot2::ggplot(relmat_sd, ggplot2::aes(y = label)) +
    ggplot2::geom_vline(
      xintercept = 0,
      linewidth = 0.45,
      linetype = "dashed",
      colour = "grey55"
    ) +
    ggplot2::geom_errorbar(
      ggplot2::aes(xmin = lower, xmax = upper),
      width = 0,
      linewidth = 1.1,
      colour = "#009E73"
    ) +
    ggplot2::geom_point(
      ggplot2::aes(x = estimate),
      shape = 21,
      size = 3.5,
      stroke = 1,
      fill = "white",
      colour = "#009E73"
    ) +
    ggplot2::scale_x_continuous(
      expand = ggplot2::expansion(mult = c(0.02, 0.05))
    ) +
    relmat_eye_theme() +
    ggplot2::labs(
      title = "Known-matrix marginal SD is separate from residual sigma",
      subtitle = "Node SD = s sqrt(Kii); bars are transformed 95% Wald intervals",
      x = "Fitted standard deviation",
      y = NULL
    )
}

## ----relmat-q2-correlation-fit------------------------------------------------
relmat_q2_example <- simulate_relmat_q2_guide_data()
relmat_q2_dat <- relmat_q2_example$data
K <- relmat_q2_example$K

fit_relmat_q2_example <- drmTMB(
  bf(
    mu1 = seed_mass ~ age + sex +
      relmat(1 | p | line, K = K),
    mu2 = plant_height ~ age + sex +
      relmat(1 | p | line, K = K),
    sigma1 = ~ 1,
    sigma2 = ~ 1,
    rho12 = ~ 1
  ),
  family = biv_gaussian(),
  data = relmat_q2_dat,
  REML = TRUE
)

relmat_q2_pairs <- corpairs(
  fit_relmat_q2_example,
  level = "relmat"
)
relmat_q2_pairs

## ----relmat-q2-confidence-eye, fig.width = 6.4, fig.height = 2.7, fig.cap = "`relmat()` intercept-only q=2 location-location point estimate from `corpairs()`; the dotted vertical line marks zero correlation and no interval is shown because calibration remains planned.", fig.alt = "Single-row point plot for the known-matrix relmat mean-mean correlation, with a hollow point estimate to the right of the dotted zero reference line."----
if (requireNamespace("ggplot2", quietly = TRUE)) {
  relmat_q2_display <- relmat_q2_pairs
  relmat_q2_display$display_label <- "relmat\nmu1-mu2"

  plot_corpairs(
    relmat_q2_display,
    colour = NULL,
    label = "display_label",
    facet = NULL
  ) +
    relmat_eye_theme() +
    ggplot2::labs(
      title = "Known-matrix latent mean correlation",
      subtitle = "Point estimate only; dotted line marks zero",
      x = "Correlation estimate"
    )
}