---
title: "Randomized matrix computations with matsketch"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Randomized matrix computations with matsketch}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r setup, include = FALSE}
knitr::opts_chunk$set(collapse = TRUE, comment = "#>", fig.width = 7,
                      fig.height = 4.2, dpi = 96, out.width = "100%")
```

Large positive-semidefinite matrices turn up throughout statistics: kernel
matrices in Gaussian-process regression, relationship matrices in
quantitative genetics, covariance matrices in spatial models. Holding one in
memory takes $n^2$ numbers and factorizing it takes about $n^3$ operations,
so exact methods run out of room long before $n$ reaches a hundred thousand.

Recent work in randomized numerical linear algebra shows how much can be done
without either step. matsketch brings four of these methods to R.

| Task | Function | What it needs from the matrix |
|:--|:--|:--|
| Low-rank approximation | `rpchol()` | about $(k + 1)n$ of its entries |
| Trace and diagonal | `trace_est()`, `diag_est()` | $m$ products $A\omega$ |
| Regularized linear systems | `nystrom_precond()`, `pcg()` | products, plus one low-rank approximation |
| Variance components | `reml_sketch()` | all of the above |

This vignette runs each one on a problem where its advantage is easy to
see. `vignette("genomic-reml", package = "matsketch")` develops the last one
in detail.

```{r}
library(matsketch)
```

## Approximating a kernel matrix from a few of its columns

A Nyström approximation rebuilds a positive-semidefinite matrix $A$ from a
subset $S$ of its columns, $\hat A = A_{:,S}\, A_{S,S}^{+}\, A_{S,:}$, and
everything depends on which columns are chosen. Randomly pivoted Cholesky
(Chen, Epperly, Tropp and Webber, 2025) chooses them one at a time, each with
probability proportional to the diagonal of the part of $A$ not yet
explained. Greedy pivoting, the classical rule, always takes the largest
residual instead, and uniform sampling ignores the matrix altogether.

The points below form a large cluster, a small cluster and a scatter of
outliers, a layout that defeats both classical rules.

```{r data}
set.seed(1)
X <- rbind(
  matrix(rnorm(2 * 1300, sd = 0.5), ncol = 2),     # large cluster
  matrix(rnorm(2 * 180, sd = 0.2), ncol = 2) + 4,  # small cluster
  matrix(runif(2 * 20, -6, 10), ncol = 2)          # outliers
)
K <- kernel_matrix(X, "gaussian", bandwidth = 0.5)
K
```

`kernel_matrix()` stores only the points and computes entries when an
algorithm asks for them. Each rule gets a budget of 100 columns, and
`trace_path` records the error after every pivot, so one fit per rule traces
out the whole curve. The best possible error at each rank comes from the
eigenvalues, which only a small example lets us compute:

```{r fits}
rules <- c(rpchol = "accelerated", greedy = "greedy", uniform = "uniform")
fits <- lapply(rules, function(r) rpchol(K, k = 100, method = r))
ev <- eigen(as.matrix(K), symmetric = TRUE, only.values = TRUE)$values
best <- 1 - cumsum(ev[1:100]) / sum(ev)

ranks <- c(20, 40, 100)
tab <- cbind(sapply(fits, function(f) f$trace_path[ranks]), best = best[ranks])
rownames(tab) <- paste("rank", ranks)
signif(tab, 2)
```

```{r error-curves}
cols <- c(rpchol = "#D55E00", greedy = "#0072B2", uniform = "#E69F00")
plot(best, log = "y", type = "l", lwd = 2, lty = 2, col = "grey30",
     ylim = range(best, unlist(lapply(fits, `[[`, "trace_path"))),
     xlab = "rank", ylab = "relative trace error")
for (r in names(fits)) lines(fits[[r]]$trace_path, lwd = 2, col = cols[[r]])
legend("topright", c(names(fits), "best possible"), col = c(cols, "grey30"),
       lty = c(1, 1, 1, 2), lwd = 2, bty = "n")
```

Randomly pivoted Cholesky stays within a small factor of the best possible
error at every rank. Greedy pivoting falls far behind at first, because it
spends its early pivots on the outliers and on isolated points at the edges
of the clusters, each of which explains little beyond itself; it catches up
only after working through them. Uniform sampling keeps pace early on and
then stalls. It almost never draws an outlier, and every outlier it misses
leaves a whole diagonal entry unexplained: the 20 outliers hold 1.3% of the
trace, which is about where its error levels off. The first 40 pivots of
each rule show where the budget goes:

```{r pivots, fig.height = 3}
op <- par(mfrow = c(1, 3), mar = c(0.5, 0.5, 2, 0.5))
for (r in names(fits)) {
  plot(X, pch = 16, cex = 0.35, col = "grey75", asp = 1, axes = FALSE,
       xlab = "", ylab = "", main = r)
  points(X[fits[[r]]$pivots[1:40], ], pch = 16, cex = 0.7, col = "#D55E00")
  box(col = "grey85")
}
par(op)
```

All three rules read the same small share of the matrix; they differ only
in which entries they read.

```{r entries}
sapply(fits, function(f) f$entries / K$n^2)
```

The default `method = "accelerated"` draws pivots from the same distribution
as one-at-a-time random pivoting but proposes them in blocks and accepts
each by rejection sampling (Epperly, Tropp and Webber, 2025), so most of its
work is matrix-matrix arithmetic.

## Estimating a trace from matrix-vector products

Traces of matrices that are expensive to form are everywhere: the effective
degrees of freedom of a smoother, the derivative of a log-determinant, the
score equations of REML. The Girard-Hutchinson estimator averages
$\omega^\top A \omega$ over random vectors $\omega$ and needs only products
with $A$, but its error falls only like $1/\sqrt{m}$ in the number of
products $m$. Hutch++ (Meyer, Musco, Musco and Woodruff, 2021) first removes
a low-rank approximation of $A$ and estimates only what is left. XTrace
(Epperly, Tropp and Webber, 2024) does the same, but uses every product for
both jobs through a leave-one-out construction, and XNysTrace is its
counterpart for positive-semidefinite matrices.

A test matrix with eigenvalues $1, 1/4, 1/9, \dots$:

```{r trace-setup}
set.seed(2)
n <- 500
U <- qr.Q(qr(matrix(rnorm(n * n), n)))
A <- U %*% ((1:n)^-2 * t(U))
trA <- sum(diag(A))
trA
trace_est(A, m = 60)
```

XTrace reports a standard error from the spread of its leave-one-out
estimates. The median relative error of each estimator over ten repetitions,
as the budget of products grows:

```{r trace-curves}
budgets <- c(12, 24, 48, 96, 192)
methods <- c(hutchinson = "#999999", hutchpp = "#E69F00",
             xtrace = "#D55E00", xnystrace = "#0072B2")
err <- sapply(names(methods), function(meth) {
  sapply(budgets, function(m) {
    median(replicate(10, abs(trace_est(A, m, meth)$estimate - trA) / trA))
  })
})
matplot(budgets, err, log = "xy", type = "b", pch = 19, lty = 1, lwd = 2,
        col = methods, xlab = "matrix-vector products",
        ylab = "relative error")
legend("bottomleft", names(methods), col = methods, lwd = 2, pch = 19,
       bty = "n")
```

The Girard-Hutchinson line has slope $-1/2$; the others fall far faster,
because each extra product also sharpens the low-rank part. `diag_est()`
applies the same construction to the diagonal:

```{r diag}
cor(diag_est(A, m = 100), diag(A))
```

## Solving regularized linear systems

Systems $(K + \mu I) x = b$ arise in kernel ridge regression, in
Gaussian-process prediction, and at every step of a REML fit. Conjugate
gradients need only products with $K$, but when $\mu$ is small they need many
of them. Frangella, Tropp and Udell (2023) precondition with a low-rank
approximation $K \approx U \hat\Lambda U^\top$, and show that a rank of about
three times the effective dimension
$d_{\mathrm{eff}}(\mu) = \mathrm{tr}\{K (K + \mu I)^{-1}\}$ keeps the
iteration count small regardless of $n$.

`effective_dim()` estimates $d_{\mathrm{eff}}$ from any low-rank
approximation and returns the recommended rank:

```{r pcg}
set.seed(3)
Z <- matrix(rnorm(3 * 1000), ncol = 3)
K2 <- as.matrix(kernel_matrix(Z, "gaussian", bandwidth = 2))
b <- rnorm(1000)
mu <- 1e-4

ed <- effective_dim(rpchol(K2, k = 200), mu)
ed
```

The estimate from a rank-200 approximation already asks for more than rank
200, so the preconditioner is built at the recommended rank. For
comparison, a preconditioner of rank 50 and no preconditioner at all:

```{r pcg-runs}
precond_of_rank <- function(k) nystrom_precond(rpchol(K2, k), mu)
runs <- list(
  plain = pcg(K2, b, mu = mu, maxit = 5000),
  `rank 50` = pcg(K2, b, mu = mu, precond = precond_of_rank(50)),
  recommended = pcg(K2, b, mu = mu,
                    precond = precond_of_rank(ed$recommended_rank))
)
sapply(runs, `[[`, "iterations")
```

```{r pcg-plot}
cols <- c("#999999", "#E69F00", "#D55E00")
plot(runs$plain$residuals, log = "y", type = "l", lwd = 2, col = cols[1],
     xlab = "iteration", ylab = "relative residual")
for (i in 2:3) {
  lines(runs[[i]]$residuals, type = "o", pch = 19, cex = 0.4, lwd = 2,
        col = cols[i])
}
legend("topright", names(runs), col = cols, lwd = 2, bty = "n")
```

Even a rank of 50 cuts the count from `r runs$plain$iterations` iterations
to `r runs[["rank 50"]]$iterations`. At the recommended rank the
approximation captures nearly all of the matrix above the level $\mu$, and
the preconditioned system is so well conditioned that
`r runs$recommended$iterations` iterations are enough.

The preconditioner depends on $\mu$ only through a diagonal, so one
approximation serves every value of $\mu$ an application visits. `pcg()`
also accepts a matrix of right-hand sides and advances all of them together,
which turns each step into one matrix-matrix product.

## Variance components on a genomic relationship matrix

`reml_sketch()` combines the three tools to fit
$y = X\beta + g + e$ with $g \sim N(0, \sigma^2_g G)$ by average-information
REML, without ever forming $V = \sigma^2_g G + \sigma^2_e I$. With
`grm_matrix()` it does not form $G$ either, only the genotypes:

```{r reml}
set.seed(4)
dat <- sim_genomic(n = 1000, p = 2000, h2 = 0.5, pops = 4)
G <- grm_matrix(dat$M)
G
fit <- reml_sketch(dat$y, G)
fit
exact <- reml_exact(dat$y, as.matrix(G))
rbind(sketched = c(fit$sigma2, h2 = fit$h2),
      exact = c(exact$sigma2, h2 = exact$h2))
```

At this size the exact fit is still quick. How the two scale, and how close
they stay, is the subject of
`vignette("genomic-reml", package = "matsketch")`.

## References

Chen, Y., Epperly, E. N., Tropp, J. A. and Webber, R. J. (2025). Randomly
pivoted Cholesky: practical approximation of a kernel matrix with few entry
evaluations. *Communications on Pure and Applied Mathematics* 78, 995--1041.
[doi:10.1002/cpa.22234](https://doi.org/10.1002/cpa.22234)

Epperly, E. N., Tropp, J. A. and Webber, R. J. (2024). XTrace: making the
most of every sample in stochastic trace estimation. *SIAM Journal on Matrix
Analysis and Applications* 45, 1--23.
[doi:10.1137/23m1548323](https://doi.org/10.1137/23m1548323)

Epperly, E. N., Tropp, J. A. and Webber, R. J. (2025). Embrace rejection:
kernel matrix approximation by accelerated randomly pivoted Cholesky. *SIAM
Journal on Matrix Analysis and Applications* 46, 2527--2557.
[doi:10.1137/24m1699048](https://doi.org/10.1137/24m1699048)

Frangella, Z., Tropp, J. A. and Udell, M. (2023). Randomized Nyström
preconditioning. *SIAM Journal on Matrix Analysis and Applications* 44,
718--752. [doi:10.1137/21m1466244](https://doi.org/10.1137/21m1466244)

Meyer, R. A., Musco, C., Musco, C. and Woodruff, D. P. (2021). Hutch++:
optimal stochastic trace estimation. *Symposium on Simplicity in Algorithms*,
142--155.
[doi:10.1137/1.9781611976496.16](https://doi.org/10.1137/1.9781611976496.16)