inst/doc/distributed-fitting.R

In cevcmm: Communication-Efficient Varying Coefficient Mixed-Effects Models

## ----include = FALSE----------------------------------------------------------
knitr::opts_chunk$set(
  collapse  = TRUE,
  comment   = "#>",
  fig.width = 6,
  fig.height = 4,
  fig.align = "center"
)
set.seed(23)

## ----setup--------------------------------------------------------------------
library(cevcmm)

## ----simulate-----------------------------------------------------------------
set.seed(23)
N <- 1200L
q <- 4L

t <- runif(N)
x <- runif(N)
Z <- matrix(rnorm(N * q), N, q)

beta_1_fun <- function(u) sin(2 * pi * u)
alpha_true <- rnorm(q, sd = 0.5)
y <- 2 + beta_1_fun(t) * x + as.vector(Z %*% alpha_true) +
     rnorm(N, sd = 0.5)

## ----split--------------------------------------------------------------------
# Assign each observation to one of 3 nodes
node_id <- sample.int(3L, N, replace = TRUE)
splits  <- split(seq_len(N), node_id)
lengths(splits)

## ----design-------------------------------------------------------------------
design <- build_vcmm_design(X = x, t = t)
X_d    <- design$X_design
# N rows x p columns; p = intercept + spline-basis coefficients,
# with n_basis auto-picked from the sample size.
dim(X_d)

## ----node-summaries-----------------------------------------------------------
summaries <- lapply(splits, function(idx) {
  node_summary(y[idx], X_d[idx, , drop = FALSE], Z[idx, , drop = FALSE])
})
summaries[[1]]

## ----size---------------------------------------------------------------------
object.size(summaries[[1]])
object.size(y[splits[[1]]]) + object.size(X_d[splits[[1]], ])

## ----fit-dist-----------------------------------------------------------------
ctrl <- vcmm_control(sigma_eps       = 0.5,
                     sigma_alpha     = 0.5,
                     update_variance = TRUE)

fit_dist <- fit_from_summaries(summaries,
                               penalty = design$penalty,
                               control = ctrl,
                               method  = "ss",
                               re_cov  = "diag")
fit_dist

## ----fit-pooled---------------------------------------------------------------
fit_pooled <- vcmm(y, X = x, Z = Z, t = t,
                   method = "ss", re_cov = "diag",
                   control = ctrl)

# beta estimates
max(abs(fit_pooled$beta - fit_dist$beta))

# alpha estimates
max(abs(fit_pooled$alpha - fit_dist$alpha))

# variance components
all.equal(fit_pooled$sigma_eps,   fit_dist$sigma_eps)
all.equal(fit_pooled$sigma_alpha, fit_dist$sigma_alpha)

## ----accumulator--------------------------------------------------------------
p <- ncol(X_d)
acc <- init_accumulator(p = p, q = q)

# Pretend chunks arrive one at a time
for (k in seq_along(splits)) {
  idx <- splits[[k]]
  ss_k <- compute_sufficient_stats(y[idx],
                                   X_d[idx, , drop = FALSE],
                                   Z[idx, , drop = FALSE])
  acc <- accumulate_stats(acc, ss_k)
}
acc

## ----fit-acc------------------------------------------------------------------
fit_acc <- fit_from_summaries(acc, penalty = design$penalty,
                              control = ctrl, method = "ss", re_cov = "diag")

# Same answer as Pattern 1
max(abs(fit_acc$beta - fit_dist$beta))

## ----od-shift, eval = FALSE---------------------------------------------------
# # Skeleton — see vignette("od-migration") for a complete runnable example.
# # G and Sigma_spatial come from the OD setup (number of regions and an
# # initial spatial covariance matrix).
# fit_dist_od <- fit_from_summaries(
#   summaries, penalty = design$penalty, control = ctrl,
#   method = "ss", re_cov = "kronecker",
#   n_groups           = G,              # number of OD regions
#   Sigma_spatial_init = Sigma_spatial,  # G x G initial spatial covariance
#   rowsum_constant    = 2               # matches what vcmm() applies internally
# )