---
title: "Fitting ML-UMR Models"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Fitting ML-UMR Models}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, include = FALSE, purl = FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  eval = FALSE,
  message = FALSE,
  warning = FALSE
)
```

## Overview

The `mlumr()` function fits Bayesian multilevel unanchored meta-regression
models using Stan. The default backend is `rstan`; `cmdstanr` can be used
when it is installed and configured. Three outcome families are supported:
binary (binomial), continuous (normal), and count (Poisson). Two model
types are available for each family:

- **SPFA** (Shared Prognostic Factor Assumption): shared regression
  coefficients across treatments
- **Relaxed SPFA**: treatment-specific regression coefficients, allowing
  effect modification

## Model specification

### Binary outcomes (binomial)

The SPFA model assumes that covariates affect outcomes equally for both
treatments. The model is:

$$
\text{logit}(P(Y_i = 1 | \text{Index})) = \mu_A + \mathbf{x}_i^T \boldsymbol{\beta}
$$

$$
\text{logit}(P(Y_j = 1 | \text{Comparator})) = \mu_B + \mathbf{x}_j^T \boldsymbol{\beta}
$$

where $\boldsymbol{\beta}$ is **shared** across both treatments. The AgD
likelihood integrates over the comparator population covariate distribution
using QMC points.

The Relaxed model allows treatment-specific effects:

$$
\text{logit}(P(Y_i = 1 | \text{Index})) = \mu_A + \mathbf{x}_i^T \boldsymbol{\beta}_A
$$

$$
\text{logit}(P(Y_j = 1 | \text{Comparator})) = \mu_B + \mathbf{x}_j^T \boldsymbol{\beta}_B
$$

This captures **effect modification** when covariate effects differ between
treatments. The difference $\boldsymbol{\delta} = \boldsymbol{\beta}_A - \boldsymbol{\beta}_B$
quantifies treatment-by-covariate interaction.

### Continuous outcomes (normal)

For continuous outcomes, the SPFA model is:

$$
Y_i | \text{Index} \sim \text{Normal}(\mu_A + \mathbf{x}_i^T \boldsymbol{\beta}, \sigma^2)
$$

The AgD likelihood uses the mean and standard error of the observed outcome.
The generated quantities include mean differences (`delta_index`,
`delta_comparator`) as the primary treatment effect measure.

An additional `prior_sigma` parameter controls the prior on the residual
standard deviation $\sigma$. The default is a half-normal prior induced by
Stan's `<lower=0>` constraint; half-Student-t and exponential priors are also
available through the prior constructors.

### Count outcomes (Poisson)

For count outcomes with exposure, the SPFA model uses a Poisson likelihood
with log link:

$$
Y_i | \text{Index} \sim \text{Poisson}(\exp(\mu_A + \mathbf{x}_i^T \boldsymbol{\beta}) \cdot E_i)
$$

where $E_i$ is the exposure (person-time). The AgD likelihood uses
log-sum-exp for numerical stability. The generated quantities include
rate ratios (`delta_index`, `delta_comparator`).

## Fitting models

The chunks below all use a small toy `dat` built from the same pipeline as
`vignette("data-preparation")`. Run this setup chunk first so the
subsequent fitting, summary, prediction, and comparison calls have
something to operate on.

```{r, eval = TRUE, purl = FALSE}
library(mlumr)
set.seed(2026)

trial_a_data <- data.frame(
  trt      = "Drug_A",
  response = rbinom(300, 1, 0.55),
  age_cat  = rbinom(300, 1, 0.40),
  sex      = rbinom(300, 1, 0.55)
)
trial_b_data <- data.frame(
  trt           = "Drug_B",
  n_total       = 400,
  n_events      = 160,
  age_cat_mean  = 0.35,
  sex_mean      = 0.50
)

ipd <- set_ipd(trial_a_data, treatment = "trt", outcome = "response",
               covariates = c("age_cat", "sex"))
agd <- set_agd(trial_b_data, treatment = "trt",
               outcome_n = "n_total", outcome_r = "n_events",
               cov_means = c("age_cat_mean", "sex_mean"),
               cov_types = c("binary", "binary"))
dat <- combine_data(ipd, agd)
dat <- add_integration(dat, n_int = 64,
                       age_cat = distr(qbern, prob = age_cat_mean),
                       sex = distr(qbern, prob = sex_mean))
```

```{r, eval = TRUE, purl = FALSE}
# SPFA model (default)
fit_spfa <- mlumr(
  dat, model = "spfa",
  chains = 2, iter = 500, warmup = 250,
  seed = 42, refresh = 0, verbose = FALSE
)

# Relaxed SPFA
fit_relaxed <- mlumr(
  dat, model = "relaxed",
  chains = 2, iter = 500, warmup = 250,
  seed = 43, refresh = 0, verbose = FALSE
)
```

### Controlling the sampler

```{r, eval = FALSE, purl = FALSE}
fit <- mlumr(
  dat,
  model = "spfa",
  chains = 4,               # Number of MCMC chains
  iter = 4000,               # Total iterations per chain
  warmup = 2000,             # Warmup iterations
  seed = 42,                 # For reproducibility
  adapt_delta = 0.99,        # Increase for divergences
  max_treedepth = 15,        # Maximum tree depth
  refresh = 500              # Print progress every 500 iterations
)
```

The backend can be selected per fit. The code below runs the same small
demonstration model with each available engine and records elapsed wall-clock
time. `rstan` uses the model compiled into the installed R package. The first
`cmdstanr` call may include external executable compilation or cache setup, so
the table reports cold and warm `cmdstanr` runs separately. For repeated
analyses, the warm-cache `cmdstanr` row is the more relevant comparison.

```{r, eval = TRUE, purl = FALSE}
run_backend_timed <- function(engine, run, note, seed = 2026, ...) {
  start_time <- proc.time()[["elapsed"]]
  fit <- tryCatch(
    mlumr(
      dat,
      model = "spfa",
      engine = engine,
      chains = 2,
      iter = 300,
      warmup = 150,
      seed = seed,
      refresh = 0,
      verbose = FALSE,
      ...
    ),
    error = function(e) e
  )
  elapsed <- proc.time()[["elapsed"]] - start_time

  if (inherits(fit, "error")) {
    return(data.frame(
      run = run,
      engine = engine,
      status = "failed",
      elapsed_seconds = round(elapsed, 2),
      lor_comparator_mean = NA_real_,
      max_rhat = NA_real_,
      min_ess = NA_real_,
      note = conditionMessage(fit),
      stringsAsFactors = FALSE
    ))
  }

  lor_comp <- marginal_effects(fit, effect = "lor", population = "comparator")
  data.frame(
    run = run,
    engine = engine,
    status = "ran",
    elapsed_seconds = round(elapsed, 2),
    lor_comparator_mean = round(lor_comp$mean, 3),
    max_rhat = round(max(fit$summary$Rhat, na.rm = TRUE), 3),
    min_ess = round(min(fit$summary$n_eff, na.rm = TRUE), 1),
    note = note,
    stringsAsFactors = FALSE
  )
}

cmdstanr_ready <- requireNamespace("cmdstanr", quietly = TRUE) &&
  isTRUE(tryCatch({
    nzchar(cmdstanr::cmdstan_path())
  }, error = function(e) FALSE))

backend_speed <- run_backend_timed(
  engine = "rstan",
  run = "rstan",
  note = "Package-compiled model, already used above"
)

if (cmdstanr_ready) {
  backend_speed <- rbind(
    backend_speed,
    run_backend_timed(
      engine = "cmdstanr",
      run = "cmdstanr_cold",
      note = "First call; may include compile/cache setup"
    ),
    run_backend_timed(
      engine = "cmdstanr",
      run = "cmdstanr_warm",
      note = "Second call; executable/cache already available"
    ),
    run_backend_timed(
      engine = "cmdstanr",
      run = "cmdstanr_warm_parallel",
      note = "Warm cache with two parallel chains",
      parallel_chains = 2
    )
  )
} else {
  backend_speed <- rbind(
    backend_speed,
    data.frame(
      run = "cmdstanr",
      engine = "cmdstanr",
      status = "skipped",
      elapsed_seconds = NA_real_,
      lor_comparator_mean = NA_real_,
      max_rhat = NA_real_,
      min_ess = NA_real_,
      note = "cmdstanr or CmdStan is not available in this R session",
      stringsAsFactors = FALSE
    )
  )
}

backend_speed
```

For production analyses, choose the backend based on your local toolchain,
diagnostics, and workflow. `cmdstanr` is often faster after the executable is
compiled, especially with parallel chains and longer runs, but this is not
guaranteed for tiny demonstration fits. In short examples, compilation, process
startup, and CSV readback can dominate the timing. Posterior estimates should be
similar up to Monte Carlo error when the same model, data, priors, and sampler
settings are used.

### Setting priors

By default, weakly informative `Normal(0, 10)` priors are used for
intercepts and `Normal(0, 2.5)` priors are used for regression coefficients.
Normal, Student-t, and Cauchy priors are available for intercepts and
regression coefficients. For normal-family residual standard deviations,
`prior_sigma` also supports an exponential prior.

```{r, eval = FALSE, purl = FALSE}
fit <- mlumr(
  dat,
  model = "spfa",
  prior_intercept = prior_normal(0, 10),   # Default
  prior_beta = prior_normal(0, 2.5)        # More informative for betas
)
```

Per-covariate beta priors and autoscaling are supported:

```{r, eval = FALSE, purl = FALSE}
fit <- mlumr(
  dat,
  model = "spfa",
  prior_beta = list(
    prior_normal(0, 2.5, autoscale = TRUE),
    prior_normal(0, 1.5, autoscale = TRUE)
  )
)
```

Inspect the priors actually passed to Stan, including autoscaled beta scales:

```{r, eval = FALSE, purl = FALSE}
prior_summary(fit)
```

## Inspecting results

### Print and summary

```{r, eval = TRUE, purl = FALSE}
# Quick overview
print(fit_spfa)

# Detailed summary
summary(fit_spfa)
```

### Marginal treatment effects

The main quantities of interest are population-level treatment effects,
marginalized over the covariate distribution:

```{r, eval = TRUE, purl = FALSE}
# All effects in both populations
marginal_effects(fit_spfa)

# Log odds ratios only
marginal_effects(fit_spfa, effect = "lor")

# In index population only
marginal_effects(fit_spfa, effect = "lor", population = "index")

# Full posterior draws (for custom summaries)
lor_draws <- marginal_effects(fit_spfa, effect = "lor", summary = FALSE)
head(lor_draws)
```

### Absolute predictions

```{r, eval = TRUE, purl = FALSE}
# Predicted probabilities P(Y=1) for each treatment in each population
predict(fit_spfa, population = "both")

# On logit scale
predict(fit_spfa, type = "link")

# Full posterior draws
pred_draws <- predict(fit_spfa, summary = FALSE)
head(pred_draws)
```

### Conditional predictions and effects

`predict()` and `marginal_effects()` report population-averaged quantities.
Use `conditional_predict()` and `conditional_effects()` when the estimand is a
specific covariate profile rather than the index or comparator population.

```{r, eval = TRUE, purl = FALSE}
profiles <- data.frame(
  age_cat = c(0, 1),
  sex = c(0, 1)
)

# Absolute event probabilities at each profile
conditional_predict(fit_spfa, newdata = profiles)

# Conditional link-scale, risk-difference, and risk-ratio effects
conditional_effects(fit_spfa, newdata = profiles)
```

For SPFA binary models, the conditional link-scale effect is constant across
profiles because the shared beta cancels on the fitted link scale. Risk
differences and risk ratios can still vary by profile because they depend on
the absolute baseline probabilities.

## Model comparison

Compare SPFA and Relaxed models using DIC:

```{r, eval = TRUE, purl = FALSE}
compare_models(fit_spfa, fit_relaxed)
```

You can also compute DIC individually:

```{r, eval = TRUE, purl = FALSE}
dic_spfa <- calculate_dic(fit_spfa)
print(dic_spfa)
```

## MCMC diagnostics

The package automatically warns about:

- Divergent transitions (increase `adapt_delta`)
- Maximum treedepth hits (increase `max_treedepth`)
- High Rhat values (> 1.01, run more iterations)
- Low ESS (< 400, run more iterations)

Check diagnostics manually:

```{r, eval = TRUE, purl = FALSE}
# Diagnostics stored in the fit object
fit_spfa$diagnostics$n_divergent
fit_spfa$diagnostics$n_max_treedepth

# Rhat and ESS from the summary table
max(fit_spfa$summary$Rhat, na.rm = TRUE)
min(fit_spfa$summary$n_eff, na.rm = TRUE)
```

For visual diagnostics, use the bayesplot package:

```{r, eval = TRUE, fig.width = 7, fig.height = 4, purl = FALSE}
library(bayesplot)
pars <- c("mu_index", "mu_comparator", "lor_index")
draws_array <- rstan::extract(fit_spfa$stanfit, pars = pars, permuted = FALSE)
mcmc_dens_overlay(draws_array, pars = pars)
mcmc_intervals(as.matrix(fit_spfa$draws[, c("lor_index", "lor_comparator")]))
```

## Troubleshooting

### Divergent transitions

Increase `adapt_delta` closer to 1:

```{r, eval = FALSE, purl = FALSE}
fit <- mlumr(dat, adapt_delta = 0.99)
```

### Slow convergence

Increase iterations and warmup:

```{r, eval = FALSE, purl = FALSE}
fit <- mlumr(dat, iter = 8000, warmup = 4000)
```

### Identifiability in Relaxed model

The Relaxed model has more parameters. With only one AgD row, the
comparator-specific coefficients may be weakly identified. The package
warns about this. Consider:

- Using the SPFA model if effect modification is not expected
- Providing multiple AgD subgroups
- Using more informative priors for betas
- Running `prior_sensitivity()` to assess how strongly relaxed-model
  conclusions depend on the beta prior

```{r, eval = FALSE, purl = FALSE}
prior_sensitivity(fit_relaxed, prior_beta_scales = c(1, 2.5, 5, 10))
```