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

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

## ----load---------------------------------------------------------------------
path <- system.file("extdata", "od_migration.csv", package = "cevcmm")
od   <- read.csv(path)
str(od)
head(od)

## ----build-z------------------------------------------------------------------
G <- 10L
N <- nrow(od)

Z <- matrix(0, N, 2L * G)
Z[cbind(seq_len(N), od$origin)]    <- 1   # origin indicators in cols 1..G
Z[cbind(seq_len(N), G + od$dest)]  <- 1   # destination indicators in cols (G+1)..2G

# Verify the OD structure: every row sums to 2 (one origin + one dest)
table(rowSums(Z))

## ----sigma-spatial------------------------------------------------------------
Sigma_spatial <- outer(seq_len(G), seq_len(G),
                       function(i, j) exp(-abs(i - j) / 3))
round(Sigma_spatial[1:4, 1:4], 3)

## ----fit----------------------------------------------------------------------
fit <- vcmm(y = od$log_flow,
            X = od$wage_diff,
            Z = Z,
            t = od$t,
            method             = "csl",
            re_cov             = "kronecker",
            n_groups           = G,
            Sigma_spatial_init = Sigma_spatial,
            control            = vcmm_control(
              sigma_eps       = 0.6,
              sigma_alpha     = sqrt(0.5),
              update_variance = TRUE))
fit

## ----sigma-2x2----------------------------------------------------------------
Sigma_2x2_hat <- fit$re_cov_state$Sigma_left
round(Sigma_2x2_hat, 3)

# Correlation between origin and destination effects
corr_OD <- Sigma_2x2_hat[1, 2] /
           sqrt(Sigma_2x2_hat[1, 1] * Sigma_2x2_hat[2, 2])
round(corr_OD, 3)

## ----ranef--------------------------------------------------------------------
alpha_hat <- fit$alpha
M_hat <- matrix(alpha_hat, nrow = G, ncol = 2L)
colnames(M_hat) <- c("origin", "dest")
rownames(M_hat) <- paste0("region_", seq_len(G))
round(M_hat, 3)

## ----ranef-plot, fig.cap = "Estimated origin and destination effects per region."----
matplot(seq_len(G), M_hat, type = "b", pch = 19, lty = 1, lwd = 2,
        col = c("steelblue", "darkorange"),
        xlab = "Region", ylab = "Random effect")
abline(h = 0, lty = 2, col = "grey60")
legend("topright", c("origin", "dest"),
       col = c("steelblue", "darkorange"), pch = 19, lwd = 2, bty = "n")

## ----vcoef, fig.cap = "Estimated time-varying effect of wage_diff on log flow."----
t_grid <- seq(0, 1, length.out = 100L)
vc     <- varying_coef(fit, t_new = t_grid, k = 1L, se.fit = TRUE)

plot(t_grid, vc$fit, type = "l", lwd = 2, col = "steelblue",
     xlab = "t (year, normalised)",
     ylab = expression(hat(beta)[1](t)),
     ylim = range(vc$fit - 2 * vc$se.fit, vc$fit + 2 * vc$se.fit,
                  1.5 * sin(2 * pi * t_grid)))
polygon(c(t_grid, rev(t_grid)),
        c(vc$fit + 2 * vc$se.fit, rev(vc$fit - 2 * vc$se.fit)),
        col = adjustcolor("steelblue", alpha.f = 0.25), border = NA)
lines(t_grid, 1.5 * sin(2 * pi * t_grid),
      col = "red", lty = 2, lwd = 2)
legend("topright", c("estimate", "truth"),
       col = c("steelblue", "red"), lty = c(1, 2), lwd = 2, bty = "n")

## ----distributed-od, eval = FALSE---------------------------------------------
# # Split N obs across 3 nodes
# node_id <- sample.int(3L, N, replace = TRUE)
# splits  <- split(seq_len(N), node_id)
# 
# design <- build_vcmm_design(X = od$wage_diff, t = od$t)
# X_d    <- design$X_design
# 
# summaries <- lapply(splits, function(idx)
#   node_summary(od$log_flow[idx],
#                X_d[idx, , drop = FALSE],
#                Z[idx, , drop = FALSE]))
# 
# fit_dist <- fit_from_summaries(
#   summaries,
#   penalty            = design$penalty,
#   control            = vcmm_control(sigma_eps = 0.6,
#                                     sigma_alpha = sqrt(0.5),
#                                     update_variance = TRUE),
#   method             = "csl",
#   re_cov             = "kronecker",
#   n_groups           = G,
#   Sigma_spatial_init = Sigma_spatial,
#   rowsum_constant    = 2
# )
# 
# # Same answer as the pooled fit above, up to BLAS noise
# max(abs(fit$beta  - fit_dist$beta))
# max(abs(fit$alpha - fit_dist$alpha))