# OD migration with Kronecker covariance

## The OD setting

In an origin-destination (OD) model, each observation is a flow between two regions. Migration counts between US states, commuting flows between zip codes, or trade flows between countries all fit this pattern.

The random-effects design has two indicators per observation: one for the origin region and one for the destination region. With \(G\) regions, the random-effects matrix \(Z\) is \(N \times 2G\), and every row sums to 2.

The natural covariance structure on the resulting \(2G\)-vector of random effects is the Kronecker form \[ \Sigma_\alpha \;=\; \underbrace{\Sigma_{2\times 2}}_{\text{origin/dest}} \;\otimes\; \underbrace{\Sigma_{\text{spatial}}}_{\text{between regions}}, \] where \(\Sigma_{2\times 2}\) captures the origin/destination variance plus their covariance, and \(\Sigma_{\text{spatial}}\) (\(G \times G\)) captures spatial dependence between regions (often known a priori from geography).

This reduces \(2G(2G+1)/2\) free covariance parameters to just 3 for \(\Sigma_{2\times 2}\) — a substantial saving when \(G\) is moderate.

## Setup

    library(cevcmm)

## Load the bundled example

The package ships a simulated OD migration dataset of \(N = 3000\) observations: \(G = 10\) regions over 30 years, with all \(10 \times 10\) origin–destination pairs observed annually.

    path <- system.file("extdata", "od_migration.csv", package = "cevcmm")
    od   <- read.csv(path)
    str(od)
    #> 'data.frame':    3000 obs. of  6 variables:
    #>  $ origin   : int  1 1 1 1 1 1 1 1 1 1 ...
    #>  $ dest     : int  1 2 3 4 5 6 7 8 9 10 ...
    #>  $ year     : int  1 1 1 1 1 1 1 1 1 1 ...
    #>  $ t        : num  0 0 0 0 0 0 0 0 0 0 ...
    #>  $ wage_diff: num  0.208 1.409 -1.489 -1.476 0.991 ...
    #>  $ log_flow : num  3.75 4.52 3.97 3.01 3.35 ...
    head(od)
    #>   origin dest year t wage_diff log_flow
    #> 1      1    1    1 0    0.2082   3.7485
    #> 2      1    2    1 0    1.4087   4.5180
    #> 3      1    3    1 0   -1.4891   3.9731
    #> 4      1    4    1 0   -1.4758   3.0113
    #> 5      1    5    1 0    0.9908   3.3466
    #> 6      1    6    1 0   -0.8832   5.2626

The columns:

- `origin`, `dest`: integer region IDs in \(\{1, \dots, 10\}\)
- `year`: 1 through 30
- `t`: year normalised to \([0, 1]\)
- `wage_diff`: a continuous covariate whose effect on flows is allowed to vary smoothly with `t`
- `log_flow`: the response (log migration count)

## Build the random-effects design

    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))
    #>
    #>    2
    #> 3000

## Choose an initial spatial covariance

In a real analysis \(\Sigma_{\text{spatial}}\) would come from geographic distance. Here we use an exponential decay in region-index distance as a reasonable proxy:

    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)
    #>       [,1]  [,2]  [,3]  [,4]
    #> [1,] 1.000 0.717 0.513 0.368
    #> [2,] 0.717 1.000 0.717 0.513
    #> [3,] 0.513 0.717 1.000 0.717
    #> [4,] 0.368 0.513 0.717 1.000

## Fit

A single `vcmm()` call with `re_cov = "kronecker"`. The package detects the constant row-sum automatically and applies the identifiability shift described in `vignette("distributed-fitting")`.

    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
    #> <vcmm_fit>  Varying Coefficient Mixed-Effects Model fit
    #>   method      : CSL
    #>   n_obs       : 3000
    #>   p (fixed)   : 19
    #>   q (random)  : 20
    #>   RE cov      : kronecker
    #>   pilot iter  : 4 (converged)
    #>   newton step : 1
    #>   sigma_eps   : 0.5991
    #>   Sigma_2x2   :
    #>     [0.9444  0.7649]
    #>     [0.7649  0.8597]
    #>   OD corr     : 0.8489
    #>   Sigma_spatial: 10 x 10 (G = 10 groups)
    #>   elapsed     : 0.0010 sec (pilot 0.0010s + newton <0.001s)

## Interpret the estimated \(\Sigma_{2 \times 2}\)

    Sigma_2x2_hat <- fit$re_cov_state$Sigma_left
    round(Sigma_2x2_hat, 3)
    #>       [,1]  [,2]
    #> [1,] 0.944 0.765
    #> [2,] 0.765 0.860

    # 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)
    #> [1] 0.849

The diagonal entries give the origin and destination variance components; the off-diagonal entry summarises whether a region that tends to send many migrants also tends to receive many. A positive correlation says yes — high-traffic regions are high-traffic in both directions.

The true simulation values were \(\Sigma_{2\times2} = \begin{pmatrix} 0.60 & 0.25 \\ 0.25 & 0.50 \end{pmatrix}\) (correlation 0.46).

On small \(G\), the estimated \(\hat\Sigma_{2\times 2}\) will differ from truth. The 3 free parameters of \(\Sigma_{2\times 2}\) are estimated from effectively a single \(G \times 2\) realisation of \(M\), plus an EM correction for the posterior uncertainty in \(\alpha\) given the data. For \(G = 10\), sampling variability in the empirical \(\text{cor}(M)\) has standard error \(\approx 1/\sqrt{G-2} \approx 0.35\), and the EM correction inflates the diagonals further to account for posterior uncertainty. The qualitative pattern — positive variance components, positive OD correlation, “high-traffic” regions that send and receive heavily — is the right takeaway at this \(G\); the exact numerical values rely on a larger network or repeated realisations.

## Recover the per-region random effects

The internal column-stacking convention is \(\alpha = \mathrm{vec}_{\text{col}}(M)\) where \(M\) is \(G \times 2\) — the first column holds origin effects, the second holds destination effects.

    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)
    #>           origin   dest
    #> region_1   0.876  0.629
    #> region_2   1.004  1.201
    #> region_3   0.450 -0.041
    #> region_4   0.109 -0.313
    #> region_5   0.494  0.288
    #> region_6   0.905  0.677
    #> region_7  -0.887 -0.466
    #> region_8  -1.426 -0.911
    #> region_9  -0.903 -0.254
    #> region_10 -0.778 -0.655

A quick visual of the two effect series. Notice how `origin` and `dest` track each other closely across regions — that’s the visual signature of the high empirical OD correlation discussed above.

    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")

Estimated origin and destination effects per region.

## Estimated varying coefficient

    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")

Estimated time-varying effect of wage_diff on log flow.

The wage-difference effect oscillates with time: in the simulation it follows \(1.5\sin(2\pi t)\). Unlike \(\Sigma_{2\times 2}\), the varying-coefficient curve is supported by all \(N = 3000\) observations and tracks the truth essentially perfectly across the 30-year window.

## Distributing this fit

The same OD problem can be fit in the distributed setting. Each node computes a `node_summary()` on its slice and ships the result; the central node aggregates and calls `fit_from_summaries()` with `rowsum_constant = 2` so the identifiability shift matches what `vcmm()` applies automatically:

    # 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))

See `vignette("distributed-fitting", package = "cevcmm")` for the full explanation of the distributed API.

## Where to go next

- Basic usage with diagonal random effects: `vignette("getting-started", package = "cevcmm")`.
- Distributed fitting in detail: `vignette("distributed-fitting", package = "cevcmm")`.

## Reference

Jalili, L. and Lin, L.-H. (2025). Scalable and Communication-Efficient Varying Coefficient Mixed Effect Models: Methodology, Theory, and Applications. arXiv:2511.12732; under review at Journal of the American Statistical Association.