---
title: "Computational Methods<br/>for GPs"
subtitle: "Lecture 23"
author: "Dr. Colin Rundel"
footer: "Sta 344/644 - Fall 2023"
format:
  revealjs:
    theme: slides.scss
    transition: fade
    slide-number: true
    self-contained: true
    html-math-method: mathjax
execute:
  echo: true
  warning: false
  collapse: true
---

```{r setup}
#| include: false
library(dukestm)
library(sf)
library(raster)

library(tidyverse)
library(patchwork)

knitr::opts_chunk$set(
  fig.align = "center"
)

options(width=70)

set.seed(20221109)

ggplot2::theme_set(ggplot2::theme_bw())

post_summary = function(m, ci_width=0.95) {
  d = data_frame(
    post_mean  = apply(m, 2, mean),
    post_med   = apply(m, 2, median),
    post_lower = apply(m, 2, quantile, probs=(1-ci_width)/2),
    post_upper = apply(m, 2, quantile, probs=1 - (1-ci_width)/2)
  )
  
  if (!is.null(colnames(m)))
    d = d %>% mutate(param = colnames(m)) %>% select(param,post_mean:post_upper)
  
  d
}

```


# GPs and Computational Complexity


## The problem with GPs

Unless you are lucky (or clever), Gaussian process models are difficult to scale to large problems. For a Gaussian process $\underset{n \times 1}{\boldsymbol{y}} \sim \mathcal{N}(\boldsymbol{\mu},\boldsymbol{\Sigma})$:

. . .

Want to sample $\boldsymbol{y}$?
$$
\boldsymbol{\mu} + {\color{red}{\text{Chol}(\boldsymbol{\Sigma})}} \times \boldsymbol{Z} \text{ with } Z_i \sim \mathcal{N}(0,1) \qquad \qquad \color{red}{\mathcal{O}\left({n^3}\right)}
$$
  
. . .


Evaluate the (log) likelihood? 
$$
-\frac{1}{2} \log {\color{red}{\text{det}(\Sigma)}} - \frac{1}{2} (\boldsymbol{x}-\boldsymbol{\mu})'\; {\color{red}{\boldsymbol{\Sigma}^{-1}}} \; (\boldsymbol{x}-\boldsymbol{\mu}) - \frac{n}{2}\log 2\pi \qquad \qquad \color{red}{\mathcal{O}\left(n^3\right)}
$$

. . .

Update covariance parameter?
$$
\color{orange}{\{\Sigma\}_{ij}} = \sigma^2 \exp(-l\,\{d\}_{ij}) + \sigma^2_n \, 1_{i=j} \qquad \qquad \color{orange}{\mathcal{O}\left(n^2\right)}
$$



## A simple guide to computational complexity

<br/>

::: {.center .large}
$\mathcal{O}\left(n\right)$ - Linear complexity 

*Go for it!*
:::
<br/>

. . .

::: {.center .large}
$\color{orange}{\mathcal{O}\left(n^2\right)}$ - Quadratic complexity

*Pray*
:::
<br/>

. . .

::: {.center .large}
$\color{red}{\mathcal{O}\left(n^3\right)}$ - Cubic complexity

*Give up*
:::


## How bad is the problem?

```{r echo=FALSE, message=FALSE}
decomp = readr::read_csv("data/lapack.csv")

ggplot(decomp, aes(y=cpu, x=n, color=method)) + 
  geom_line() +
  geom_point() +
  ylab("time (secs)")
```

## Practice - Migratory Model Prediction

After fitting the GP need to sample from the posterior predictive distribution at $\sim3000$ locations
$$ \boldsymbol{y}_{p} \sim \mathcal{N}\left(\mu_p + \Sigma_{po} \Sigma_o^{-1}(y_o - \mu_o) ,~ \Sigma_p - \Sigma_{po} \Sigma_{o}^{-1} \Sigma_{op}\right) $$

. . .

::: {.small}

| Step            | CPU (secs) |
|--------------------|------------|
| 1. Calc $\Sigma_p$, $\Sigma_{po}$, $\Sigma_{o}$                           | 1.080 |
| 2. Calc $\text{chol}(\Sigma_p - \Sigma_{po} \Sigma_{o}^{-1} \Sigma_{op})$ | 0.467 |
| 3. Calc $\mu_{p|o} + \text{chol}(\Sigma_{p|o}) \times Z$                  | 0.049 |
| 4. Calc Allele Prob                                                       | 0.129 |
|    Total                                                                   | 1.732 |

:::

Total run time for 1000 posterior predictive draws: 28.9 min (CPU)


## A bigger hammer?

::: {.small}

| Step                             | CPU (secs) | CPU+GPU (secs) | Rel. Perf |
|----------------------------------|------------|----------------|-----------|
| 1. Calc. $\Sigma_p$, $\Sigma_{po}$, $\Sigma_{p}$                           | 1.080 | 0.046 | 23.0 |
| 2. Calc. $\text{chol}(\Sigma_p - \Sigma_{po} \Sigma_{o}^{-1} \Sigma_{op})$ | 0.467 | 0.208 |  2.3 |
| 3. Calc. $\mu_{p|o} + \text{chol}(\Sigma_{p|o}) \times Z$                  | 0.049 | 0.052 |  0.9 |
| 4. Calc. Allele Prob                                                       | 0.129 | 0.127 |  1.0 |
|    Total                                                                   | 1.732 | 0.465 |  3.7 |

:::

. . .

<br/>

Total run time for 1000 posterior predictive draws: 

* CPU (28.9 min)
* CPU+GPU (7.8 min)

::: {.aside}
Benchmarks based on decade old consumer hardware
:::

## Cholesky CPU vs GPU (P100)

```{r echo=FALSE}
decomp %>%
  tidyr::gather(comp, time, cpu:gpu) %>%
  ggplot(aes(y=time, x=n, color=method, linetype=comp)) + 
    geom_line() +
    geom_point() +
  ylab("time (secs)")
```

::: {.aside}
Benchmarks based on ~5 year old server hardware
:::

##

```{r echo=FALSE}
decomp %>%
  tidyr::gather(comp, time, cpu:gpu) %>%
  ggplot(aes(y=time, x=n, color=method, linetype=comp)) + 
    geom_line() +
    geom_point() +
    ylab("time (secs)") +
    scale_y_log10()
```

## Relative Performance

```{r echo=FALSE}
decomp %>%
  ggplot(aes(y=cpu/gpu, x=n, color=method)) + 
    geom_line() +
    geom_point() +
    ylab("Relative performance") +
    scale_y_log10()
```

## Aside (1) - Matrix Multiplication (P100)

```{r echo=FALSE, message=FALSE}
mat_mult = readr::read_csv("data/mat_mult.csv")

mat_mult %>%
  tidyr::gather(comp, time, -n) %>%
  ggplot(aes(x=n, y=time/1000, linetype=comp)) +
    geom_line() +
    geom_point() +
    ylab("time (sec)") + 
    labs(title="Matrix Multiplication")
```

##

```{r echo=FALSE, message=FALSE}
mat_mult %>%
  ggplot(aes(x=n, y=cpu/gpu)) +
    geom_line() +
    geom_point() +
    ylab("rel. perf") + 
    labs(title="Matrix Multiplication - Relative Performance")
```

## Aside (2) - Memory Limitations

A general covariance is a dense $n \times n$ matrix, meaning it will require $n^2 \times$ 64-bits to store.

```{r echo=FALSE}
size = data_frame(
  n = seq(1000, 50000, by=1000)
) %>%
  mutate(size_gb = n^2 * 64/8 / 10^9)

ggplot(size, aes(x=n, y=size_gb)) +
  geom_line(size=1.5) +
  geom_hline(yintercept = c(4,8,16), color="red", alpha=0.5, linetype="dashed") +
  labs(y="Cov Martrix Size (GB)")
```



## Other big hammers

::: {.small}
`bigGP` is an R package written by Chris Paciorek (UC Berkeley), et al.

*  Specialized distributed implementation of linear algebra operation for GPs

*  Designed to run on large super computer clusters

*  Uses both shared and distributed memory

*  Able to fit models on the order of $n = 65$k (32 GB Cov. matrix)
:::

![](imgs/Lec22/Paciorek.png){fig-align="center" width="70%"}


## More scalable solutions?

* Spectral domain / basis functions

* Covariance tapering 

* GMRF approximations 

* Low-rank approximations

* Nearest-neighbor models




# Low Rank Approximations

## Low rank approximations in general

Lets look at the example of the singular value decomposition of a matrix,

$$
\underset{n \times m}{M} = \underset{n \times n}{U}\,\underset{n \times m}{\text{diag}(S)}\,\underset{m \times m}{V^{\,t}} 
$$

where $U$ are the left singular vectors, $V$ the right singular vectors, and $S$ the singular values. Usually the singular values and vectors are ordered such that the singular values are in descending order. 

. . .

The Eckart–Young theorem states that we can construct an approximatation of $M$ with rank $k$ by setting $\tilde S$ to contain only the $k$ largest singular values and all other values set to zero.

$$ 
\begin{aligned}
\underset{n \times m}{\tilde M} 
  &= \underset{n \times n}{U}\,\underset{n \times m}{\text{diag}(\tilde S)}\,\underset{m \times m}{V^{\,t}} 
  = \underset{n \times k}{\tilde U}\,\underset{k \times k}{\text{diag}(\tilde S)}\,\underset{k \times m}{\tilde{V}^{\,t}} 
\end{aligned}
$$

## Example

```{r echo=FALSE, message=FALSE, eval=FALSE}

print_mat = function(m, digits=3) {
  print(
    xtable::xtable(m, digits=digits), floating=FALSE, tabular.environment="pmatrix", 
    hline.after=NULL, include.rownames=FALSE, include.colnames=FALSE
  )  
}

hilbert = function(n) { i <- 1:n; 1 / outer(i - 1, i, "+") }
M = hilbert(4)
s = svd(M)
D = diag(s$d)

svd(M)

print_mat(M)
print_mat(s$u)
print_mat(D)
print_mat(t(s$v))

D_tilde = D = diag(c(s$d[1:2],0,0))
M_tilde = s$u %*% D_tilde %*% t(s$v)
print_mat(M_tilde)
```

::: {.small}
$$ 
\begin{aligned}
M 
&= \begin{pmatrix}
  1.000 & 0.500 & 0.333 & 0.250 \\ 
  0.500 & 0.333 & 0.250 & 0.200 \\ 
  0.333 & 0.250 & 0.200 & 0.167 \\ 
  0.250 & 0.200 & 0.167 & 0.143 \\ 
\end{pmatrix} 
  = U \, \text{diag}(S) \, V^{\,t} \\
U = V &= \begin{pmatrix}
  -0.79 & 0.58 & -0.18 & -0.03 \\ 
  -0.45 & -0.37 & 0.74 & 0.33 \\ 
  -0.32 & -0.51 & -0.10 & -0.79 \\ 
  -0.25 & -0.51 & -0.64 & 0.51 \\ 
  \end{pmatrix} \\
S &= 
\begin{pmatrix} 1.50 & 0.17  & 0.01 & 0.00 \end{pmatrix}
\end{aligned}
$$
:::

## Rank 2 approximation

::: {.small}
$$ 
\begin{aligned}
\tilde M
&= \begin{pmatrix}
  -0.79 &  0.58 \\ 
  -0.45 & -0.37 \\ 
  -0.32 & -0.51 \\ 
  -0.25 & -0.51 \\ 
\end{pmatrix}
\begin{pmatrix}
  1.50 & 0.00 \\ 
  0.00 & 0.17 \\ 
\end{pmatrix}
\begin{pmatrix}
  -0.79 & -0.45 & -0.32 & -0.25 \\ 
  0.58 & -0.37 & -0.51 & -0.51 \\ 
\end{pmatrix} \\
&= 
\begin{pmatrix}
  1.000 & 0.501 & 0.333 & 0.249 \\ 
  0.501 & 0.330 & 0.251 & 0.203 \\ 
  0.333 & 0.251 & 0.200 & 0.166 \\ 
  0.249 & 0.203 & 0.166 & 0.140 \\ 
\end{pmatrix}
\end{aligned} 
$$
:::

. . .

::: {.small}
$$ 
\begin{aligned}
M 
&= \begin{pmatrix}
  1.000 & 0.500 & 0.333 & 0.250 \\ 
  0.500 & 0.333 & 0.250 & 0.200 \\ 
  0.333 & 0.250 & 0.200 & 0.167 \\ 
  0.250 & 0.200 & 0.167 & 0.143 \\ 
\end{pmatrix}
\end{aligned}
$$
:::



## Approximation Error

We can measure the error of the approximation using the Frobenius norm,
$$ \lVert M-\tilde M\rVert_F = \left( \sum_{i=1}^m\sum_{j=1}^n (M_{ij}-\tilde M_{ij})^2\right)^{1/2} $$

. . .

$$  M-\tilde M = 
\begin{pmatrix}
  0.00022 & -0.00090 & 0.00012 & 0.00077 \\ 
  -0.00090 & 0.00372 & -0.00053 & -0.00317 \\ 
  0.00012 & -0.00053 & 0.00013 & 0.00039 \\ 
  0.00077 & -0.00317 & 0.00039 & 0.00277 \\ 
\end{pmatrix}
$$

$$ \lVert M-\tilde M\rVert_F = 0.00674 $$


## Strong dependence

For a $50 \times 50$ covariance matrix with a *large* effective range,

```{r echo=FALSE, fig.height=4}
d = runif(2*50) %>% matrix(ncol=2) %>% dist() %>% as.matrix()
cov = exp(-d * 3)
svd_m = svd(cov)

sing_values = svd_m$d

res = data_frame(rank=50:0, frob=NA)

for(i in 1:nrow(res)) {
  res$frob[i] = (cov - svd_m$u %*% diag(svd_m$d) %*% t(svd_m$v))^2 %>% sum() %>% sqrt()
  svd_m$d[50-i+1] = 0
}
```


::: {.panel-tabset}

### Singular values

```{r}
#| echo: false
par(mar=c(4,4,0.1,0.1))
plot(sing_values, type='b', xlab="k", ylab="Singular Values", ylim=c(0,15))
```

### Approximation error

```{r}
#| echo: false
par(mar=c(4,4,0.1,0.1))
plot(res$rank, res$frob, type='b', xlab="Rank", ylab="Error (Frob. norm)", ylim=c(0,17.5))
```

:::




## Weak dependence

For a $50 \times 50$ covariance matrix with a *short* effective range,

```{r echo=FALSE, fig.height=4}
d = runif(2*50) %>% matrix(ncol=2) %>% dist() %>% as.matrix()
cov = exp(-d * 9)
svd_m = svd(cov)

sing_values = svd_m$d


res = data_frame(rank=50:0, frob=NA)

for(i in 1:nrow(res)) {
  res$frob[i] = (cov - svd_m$u %*% diag(svd_m$d) %*% t(svd_m$v))^2 %>% sum() %>% sqrt()
  svd_m$d[50-i+1] = 0
}
```

::: {.panel-tabset}

### Singular values

```{r}
#| echo: false
par(mar=c(4,4,0.1,0.1))
plot(sing_values, type='b', xlab="", ylab="Singular Values", ylim=c(0,15))
```

### Approximation error

```{r}
#| echo: false
par(mar=c(4,4,0.1,0.1))
plot(res$rank, res$frob, type='b', xlab="Rank", ylab="Error (Frob. norm)", ylim=c(0,17.5))
```
:::



## How does this help?

There is an immensely useful linear algebra identity, the Sherman-Morrison-*Woodbury* formula, for the inverse (and determinant) of a decomposed matrix,

::: {.small}
$$
\begin{aligned}
\underset{n \times m}{M}^{-1} 
&= \left(\underset{n \times m}{A} + \underset{n \times k}{U} ~ \underset{k \times k}{S} ~ \underset{k \times m}{V^t}\right)^{-1} \\
&= A^{-1} - A^{-1} U \left(S^{-1}+V^{\,t} A^{-1} U\right)^{-1}V^{\,t} A^{-1}.
\end{aligned}
$$
:::

. . .

How does this help?

* Imagine that $A = \text{diag}(A)$, then it is trivial to find $A^{-1}$.

* $S^{-1}$ is $k \times k$ which is hopefully small, or even better $S = \text{diag}(S)$.

* $\left(S^{-1}+V^{\,t} A^{-1} U\right)$ is $k \times k$ which is also small.


## Aside - Determinant

Remember for any MVN distribution when evaluating the likelihood

::: {.small}
$$
-\frac{1}{2} \log {|\Sigma|} - \frac{1}{2} (\boldsymbol{x}-\boldsymbol{\mu})' {\boldsymbol{\Sigma}^{-1}} (\boldsymbol{x}-\boldsymbol{\mu}) - \frac{n}{2}\log 2\pi
$$
:::

we need the inverse of $\Sigma$ as well as its *determinant*.

. . .

* For a full rank Cholesky decomposition we get the determinant for ``free''.
$$|M| = |LL^t| = \prod_{i=1}^n \left(\text{diag}(L)_i\right)^2$$

* The Sherman-Morrison-Woodbury Determinant lemma gives us,
$$
\begin{aligned}
\det(M) 
  &= \det({A} + {U} {S} {V^t})
   = \det(S^{-1} + V^t A^{-1} U) ~ \det(S) ~ \det(A)
\end{aligned}
$$



## Low rank approximations for GPs

For a standard spatial random effects model, 

$$
y(\boldsymbol{s}) = x(\boldsymbol{s}) \, \boldsymbol{\beta} + w(\boldsymbol{s}) + \epsilon, \quad \epsilon \sim N(0,~\tau^2 I)
$$

$$
w(\boldsymbol{s}) \sim \mathcal{N}(0,~\boldsymbol{\Sigma}(\boldsymbol{s})), \quad \boldsymbol{\Sigma}(\boldsymbol{s},\boldsymbol{s}')=\sigma^2\;\rho(\boldsymbol{s},\boldsymbol{s}'|\theta)
$$

if we can replace $\boldsymbol{\Sigma}(\boldsymbol{s})$ with a low rank approximation of the form $\boldsymbol{\Sigma}(\boldsymbol{s}) \approx \boldsymbol{U}\,\boldsymbol{S}\,\boldsymbol{U}^t$ where 

* $\boldsymbol{U}$ is $n \times k$, 

* $\boldsymbol{S}$ is $k \times k$, and

* $A = \tau^2 I$ or a similar diagonal matrix




# Predictive Processes

## Gaussian Predictive Processes

For a rank $k$ approximation,

*  Pick $k$ knot locations $\boldsymbol{s}^\star$

. . .

*  Calculate knot covariance, $\boldsymbol{\Sigma}(\boldsymbol{s}^\star)$, and knot cross-covariance, $\boldsymbol{\Sigma}(\boldsymbol{s}, \boldsymbol{s}^\star)$

. . .

*  Approximate full covariance using

$$ 
\boldsymbol{\Sigma}(\boldsymbol{s}) \approx \underset{n \times k}{\boldsymbol{\Sigma}(\boldsymbol{s},\boldsymbol{s}^\star)} \, \underset{k \times k}{\boldsymbol{\Sigma}(\boldsymbol{s}^\star)^{-1}} \, \underset{k \times n}{\boldsymbol{\Sigma}(\boldsymbol{s}^\star,\boldsymbol{s})}.
$$


::: {.aside}
These are also called inducing points in other non-spatial GP literature.
:::


##

*  PPs systematically underestimates variance ($\sigma^2$) and inflate $\tau^2$, Modified predictive processs corrects this using 

$$
\begin{aligned}
\boldsymbol{\Sigma}(\boldsymbol{s}) \approx &
\boldsymbol{\Sigma}(\boldsymbol{s},\boldsymbol{s}^\star) \, \boldsymbol{\Sigma}(\boldsymbol{s}^\star)^{-1} \, \boldsymbol{\Sigma}(\boldsymbol{s}^\star,\boldsymbol{s}) \\
&+ \text{diag}\Big(\boldsymbol{\Sigma}(\boldsymbol{s}) - \boldsymbol{\Sigma}(\boldsymbol{s},\boldsymbol{s}^\star) \, \boldsymbol{\Sigma}(\boldsymbol{s}^\star)^{-1} \, \boldsymbol{\Sigma}(\boldsymbol{s}^\star,\boldsymbol{s})\Big).
\end{aligned}
$$



::: {.aside}
Banerjee, Gelfand, Finley, Sang (2008); Finley, Sang, Banerjee, Gelfand (2008)
:::


## Example

Below we have a surface generate from a squared exponential Gaussian Process where

::: {.small}
$$ \{\Sigma\}_{ij} = \sigma^2 \exp\left(-(\phi\,d)^2\right) + \tau^2 I $$
$$ \sigma^2 = 1 \quad \phi=9 \quad \tau^2 = 0.1 $$
:::

```{r echo=FALSE}
set.seed(20170410)

if(!file.exists("Lec22_pp_data.Rdata")) {
  n=4900
  n_samp = 1000
  
  r = raster::raster(xmn=0, xmx=1, ymn=0, ymx=1, nrow=sqrt(n), ncol=sqrt(n))
  
  coords = raster::xyFromCell(r, 1:length(r))
  
  cov_func = function(d) exp(-(9*d)^2) + ifelse(d==0, 0.1, 0) 
  Sigma = coords %>% dist() %>% as.matrix() %>% cov_func()
  
  r[] = t(chol(Sigma)) %*% rnorm(n)
  
  obs_coords = runif(2*n_samp) %>% matrix(ncol=2)
  
  Sigma_2 = obs_coords %>% rdist() %>% cov_func()
  Sigma_21 = rdist(obs_coords, coords) %>% cov_func()
  
  obs = Sigma_21 %*% solve(Sigma) %*% r[] + t(chol(Sigma_2 - Sigma_21 %*% solve(Sigma, t(Sigma_21)))) %*% rnorm(n_samp)
  
  
  d = data_frame(z=c(obs), x=obs_coords[,1], y=obs_coords[,2])
  
  r_obs = r
  r_obs[] = NA
  r_obs[raster::cellFromXY(r_obs, d[,c("x","y")] %>% as.matrix())] = d$z

  save(d, r, coords, r_obs, cov_func, n, n_samp, file = "Lec22_pp_data.Rdata")
} else {
  load("Lec22_pp_data.Rdata")
}
```

```{r}
#| echo: false
par(mfrow=c(1,2))
plot(r, main="True Surface", axes=FALSE, asp=0)
plot(r_obs, main="Observed Data", axes=FALSE, asp=0)
```

## Predictive Process Model Results

```{r echo=FALSE, message=FALSE}
library(spBayes)

if (!file.exists("Lec22_pp_models.Rdata"))
{
  n.samples = 20000
  starting = list("phi"=3/0.3, "sigma.sq"=1, "tau.sq"=0.1)

  tuning = list("phi"=0.1, "sigma.sq"=0.1, "tau.sq"=0.1)
  
  priors = list("beta.Norm"=list(0, 100),
                  "phi.Unif"=c(3/1, 3/0.1), 
                  "sigma.sq.IG"=c(2, 2),
                  "tau.sq.IG"=c(2, 2))

  cov.model = "gaussian"

  m = spLM(z~1, data=d, coords=d[,c("x","y")] %>% as.matrix(), starting=starting,
           tuning=tuning, priors=priors, cov.model=cov.model,
           n.samples=n.samples, verbose=TRUE, n.report=n.samples/2+1)

  pp_5 = spLM(z~1, data=d, coords=d[,c("x","y")] %>% as.matrix(), 
              knots=c(5,5,0.1), modified.pp = FALSE,
              starting=starting, tuning=tuning, priors=priors, cov.model=cov.model,
              n.samples=n.samples, verbose=FALSE, n.report=n.samples/2+1)
  pp_10 = spLM(z~1, data=d, coords=d[,c("x","y")] %>% as.matrix(), 
               knots=c(10,10,0.05), modified.pp = FALSE,
               starting=starting, tuning=tuning, priors=priors, cov.model=cov.model,
               n.samples=n.samples, verbose=FALSE, n.report=n.samples/2+1)
  pp_15 = spLM(z~1, data=d, coords=d[,c("x","y")] %>% as.matrix(), 
               knots=c(15,15,0.05), modified.pp = FALSE,
               starting=starting, tuning=tuning, priors=priors, cov.model=cov.model,
               n.samples=n.samples, verbose=FALSE, n.report=n.samples/2+1)
  
  mpp_5 = spLM(z~1, data=d, coords=d[,c("x","y")] %>% as.matrix(), 
              knots=c(5,5,0.1), modified.pp = TRUE,
              starting=starting, tuning=tuning, priors=priors, cov.model=cov.model,
              n.samples=n.samples, verbose=FALSE, n.report=n.samples/2+1)
  mpp_10 = spLM(z~1, data=d, coords=d[,c("x","y")] %>% as.matrix(), 
               knots=c(10,10,0.05), modified.pp = TRUE,
               starting=starting, tuning=tuning, priors=priors, cov.model=cov.model,
               n.samples=n.samples, verbose=FALSE, n.report=n.samples/2+1)
  mpp_15 = spLM(z~1, data=d, coords=d[,c("x","y")] %>% as.matrix(), 
               knots=c(15,15,0.05), modified.pp = TRUE,
               starting=starting, tuning=tuning, priors=priors, cov.model=cov.model,
               n.samples=n.samples, verbose=FALSE, n.report=n.samples/2+1)
  
  models = list(m=m, pp_5=pp_5, pp_10=pp_10, pp_15=pp_15, mpp_5=mpp_5, mpp_10=mpp_10, mpp_15=mpp_15)
  
  save(models, file="Lec22_pp_models.Rdata")
} else {
  load("Lec22_pp_models.Rdata")
}
```

```{r echo=FALSE}
if (!file.exists("Lec22_pp_predict.Rdata"))
{
  predictions = lapply(
    models,
    function(m)
    {
      pred = spPredict(m, coords, matrix(1, nrow=nrow(coords)), start = 10001, thin=100, n.report=10)
      rast = r
      rast[] = pred$p.y.predictive.samples %>% t() %>% post_summary() %>% .$post_mean
      rast
    }
  ) %>% setNames(names(models)) 
  
  save(predictions, file="Lec22_pp_predict.Rdata")
} else {
  load("Lec22_pp_predict.Rdata")
}
```

```{r echo=FALSE}
par(mfrow=c(2,4), mar=c(1,1,3,1))

plot(r,                  axes=FALSE, asp=0, legend=FALSE, main="True Field")
plot(predictions$pp_5,   axes=FALSE, asp=0, legend=FALSE, main="PP - 5 x 5 knots")
plot(predictions$pp_10,  axes=FALSE, asp=0, legend=FALSE, main="PP - 10 x 10 knots")
plot(predictions$pp_15,  axes=FALSE, asp=0, legend=FALSE, main="PP - 15 x 15 knots")
plot(predictions$m,      axes=FALSE, asp=0, legend=FALSE, main="Full GP")
plot(predictions$mpp_5,  axes=FALSE, asp=0, legend=FALSE, main="Mod. PP - 5 x 5 knots")
plot(predictions$mpp_10, axes=FALSE, asp=0, legend=FALSE, main="Mod. PP - 10 x 10 knots")
plot(predictions$mpp_15, axes=FALSE, asp=0, legend=FALSE, main="Mod. PP - 15 x 15 knots")
```

## Performance

```{r echo=FALSE}
res = data_frame(
  time = purrr::map_dbl(models, ~ .$run.time[3]),
  error = purrr::map_dbl(predictions, ~ (.[] - r[])^2 %>% sum() %>% sqrt()),
  model = c("Full GP", "PP", "PP","PP", "Mod. PP", "Mod. PP", "Mod. PP") %>% forcats::as_factor(),
  knots = c("-","5x5", "10x10", "15x15","5x5", "10x10", "15x15") %>% forcats::as_factor()
)

ggplot(res, aes(x=time, y=error, color=knots, shape=model)) +
  geom_point(size=4)
```

## Parameter Estimates

<br/>

```{r}
#| echo: false
post_mean = purrr::map(models, ~ .$p.theta.samples %>% post_summary() %>% select(param, post_mean) %>% tibble::deframe()) %>% 
  purrr::transpose() %>% 
  purrr::simplify_all() %>% 
  bind_cols() %>% 
  cbind(model = names(models),.) %>%
  mutate(model = as.character(model)) %>%
  rbind(list("true",1,0.1,9)) %>%
  mutate(
    model = c("Full GP", "PP", "PP","PP", "Mod. PP", "Mod. PP", "Mod. PP","True") %>% forcats::as_factor(),
    knots = c("-",25, 100, 225,25, 100, 225,"-") %>% forcats::as_factor()
  ) %>%
  tidyr::gather(param, value, -model, -knots)

ggplot(post_mean, aes(x=value, y=as.integer(forcats::as_factor(model)), col=knots, shape=model)) +
  geom_point(size=3) +
  facet_wrap(~param, scale="free_x", ncol = 2) +
  labs(x="Parameter Value", y="") +
  theme(axis.title.y=element_blank(),
        axis.text.y=element_blank(),
        axis.ticks.y=element_blank())
```


# Random Projections

## Low Rank via Random Projections

::: {.small}
1.  Starting with an matrix $\underset{m \times n}{\boldsymbol{A}}$.
:::

. . .

::: {.small}
2.  Draw a Gaussian random matrix $\underset{n \times k+p}{\boldsymbol{\Omega}}$.
:::

. . .

::: {.small}
3.  Form $\boldsymbol{Y} = \boldsymbol{A}\,\boldsymbol{\Omega}$ and compute its QR factorization $\boldsymbol{Y} = \boldsymbol{Q}\,\boldsymbol{R}$
:::

. . .

<!--
%*  For $j=1,2,\ldots,q$ where q is integer power
%   Form Y􏱷 = AT Q and compute its QR factorization Y􏱷 = Q􏱷 R􏱷 j j−1 j jj
%   Form Yj = AQ􏱷j and compute its QR factorization Yj = Qj Rj
%   End
%*  $Q = Q_q$
-->

::: {.small}
4.  Form $\boldsymbol{B}=\boldsymbol{Q}'\,\boldsymbol{A}$.
:::

. . .

::: {.small}
5.  Compute the SVD of $\boldsymbol{B} = \boldsymbol{\hat{U}}\,\boldsymbol{S}\,\boldsymbol{V}'$.
:::

. . .

::: {.small}
6.  Form the matrix $\boldsymbol{U} = \boldsymbol{Q} \, \boldsymbol{\hat{U}}$.
:::

. . .

::: {.small}
7.  Form  $\boldsymbol{\tilde{A}} = \boldsymbol{U}\boldsymbol{S}\boldsymbol{V}'$
:::

##

Resulting approximation has a bounded expected error,

$$
E| \boldsymbol{A} - \boldsymbol{U}\boldsymbol{S}\boldsymbol{V}'\|_F \leq \left[1 + \frac{4\sqrt{k+p}}{p-1} \sqrt{\min(m,n)} \right] \sigma_{k+1}. 
$$


::: {.aside}
Halko, Martinsson, Tropp (2011)
:::



## Random Matrix Low Rank Approxs and GPs

::: {.small}
The preceding algorithm can be modified slightly to take advantage of the positive definite structure of a covariance matrix.

1.  Starting with an $n \times n$ covariance matrix $\boldsymbol{A}$.

2. Draw Gaussian random matrix $\underset{n \times k+p}{\boldsymbol{\Omega}}$.

3.  Form $\boldsymbol{Y} = \boldsymbol{A}\,\boldsymbol{\Omega}$ and compute its QR factorization $\boldsymbol{Y} = \boldsymbol{Q}\,\boldsymbol{R}$

4.  Form the $\boldsymbol{B}=\boldsymbol{Q}'\,\boldsymbol{A} \, \boldsymbol{Q}$.

5.  Compute the eigen decomposition of $\boldsymbol{B} = \boldsymbol{\hat{U}}\,\boldsymbol{S}\,\boldsymbol{\hat{U}}'$.

6.  Form the matrix $\boldsymbol{U} = \boldsymbol{Q} \, \boldsymbol{\hat{U}}$.

Once again we have a bound on the error,

$$
  E \|\boldsymbol{A} - \boldsymbol{U}\boldsymbol{S}\boldsymbol{U}'\|_F 
\lesssim c \cdot \sigma_{k+1}. 
$$
:::

::: {.aside}
Halko, Martinsson, Tropp (2011), Banerjee, Dunson, Tokdar (2012)
:::


## Low Rank Approximations and GPUs

Both predictive process and random matrix low rank approximations are good candidates for acceleration using GPUs.

* Both use Sherman-Woodbury-Morrison to calculate the inverse (involves matrix multiplication, addition, and a small matrix inverse).

* Predictive processes involves several covariance matrix calculations (knots and cross-covariance) and a small matrix inverse.

* Random matrix low rank approximations involves a large matrix multiplication ($\boldsymbol{A}\,\boldsymbol{\Omega}$) and several small matrix decompositions (QR, eigen).




## Comparison $n=15,000,\;k=\{100,\ldots,4900\}$

```{r fig.height=4, echo=FALSE}
load("data/res3.Rdata")
strong_cpu = res$time[res$method=="cpu"]
strong_gpu = res$time[res$method=="gpu"] 
strong = res %>%
  filter(!method %in% c("cpu","gpu")) %>%
  filter(!method %in% c("lr2", "lr2 mod", "lr3", "lr3 mod"))

load("data/res12.Rdata")
weak_cpu = res$time[res$method=="cpu"]
weak_gpu = res$time[res$method=="gpu"] 
weak = res %>%
  filter(!method %in% c("cpu","gpu")) %>%
  filter(!method %in% c("lr2", "lr2 mod", "lr3", "lr3 mod"))

( 
  ggplot(strong, aes(x=time, y=error, color=method)) +
    geom_line() +
    geom_point() + 
    geom_vline(xintercept = strong_cpu, color="red") +
    geom_vline(xintercept = strong_gpu, color="orange") + 
    labs(title="Strong Dependence")
) + (
  ggplot(weak, aes(x=time, y=error, color=method)) +
    geom_line() +
    geom_point() + 
    geom_vline(xintercept = weak_cpu, color="red") +
    geom_vline(xintercept = weak_gpu, color="orange") + 
    labs(title="Weak Dependence")
)
```

## Rand. Projection LR Depositions for Prediction



This approach can also be used for prediction, if we want to sample 

$$
\boldsymbol{y} \sim \mathcal{N}(0,\boldsymbol{\Sigma})
$$
$$
\Sigma \approx \boldsymbol{U} \boldsymbol{S} \boldsymbol{U}^t = (\boldsymbol{U} \boldsymbol{S}^{1/2} \boldsymbol{U}^t)(\boldsymbol{U} \boldsymbol{S}^{1/2} \boldsymbol{U}^t)^t 
$$

then 

$$
y_{\text{pred}} = (\boldsymbol{U}\, \boldsymbol{S}^{1/2}\,\boldsymbol{U}^t) \times \boldsymbol{Z} \text{ where } Z_i \sim \mathcal{N}(0,1)
$$

because $\boldsymbol{U}^t \, \boldsymbol{U} = I$ since $\boldsymbol{U}$ is an orthogonal matrix.

::: {.aside}
Dehdari, Deutsch (2012)
:::



##

![](imgs/Lec22/RandLRPred.png){fig-align="center" width="100%"}

$$ n=1000, \quad p=10000 $$

# Nearest-neighbor models

## Vecchia’s approximation

::: {.medium}
Another approach for simplifying the computational complexity is simplify / approximate the likelihood we are evaluating. As we saw with both time series models and the areal models we can rewrite our joint likelihood as a product of conditional likelihoods.
:::

::: {.small}
$$
\begin{aligned}
p(\boldsymbol{y}) &= p(y_1,y_2,\ldots,y_n) \\
                  &= p(y_1) \, p(y_2|y_1) \, p(y_3|y_1,y_2) \, \cdots \, p(y_n|y_1,\ldots,y_{n-1}) \\
                  &= p(y_1) \prod_{i=1}^n p(y_i|y_1,\ldots,y_{i-1})
\end{aligned}
$$
:::

. . .

::: {.medium}
The Vecchia approach is to approximate the conditional likelihoods by using only the $k$ nearest neighbors of the $i$th observation. This is appealing in the spatial context as we expect correlation to depend on distance, so nearer observations should be more relevant than distant observations.
:::

::: {.small}
$$
\begin{aligned}
\overset{\sim}{p}(\boldsymbol{y}) 
= \prod_{i=1}^n p(y_i|\boldsymbol{y}_{N(y_i)})
\end{aligned}
$$
:::

## Choosing neighbors

::: {.medium}
In order to choose neighbors we need to first defined an ordering of the observations as each $y_i$'s neighbors may only be chosen from the observations that precede it in the ordering. With time series this is easy, but in space it is non-obvious.
:::

. . .


:::: {.columns .small}
::: {.column width='25%'}
One approach is to order the observations by their x *or* y coordinates or some combination thereof - note that this is a (mostly) arbitrary choice and can potentially affect the results.
:::

::: {.column width='75%'}
```{r}
#| echo: false
set.seed(1234)
d = tibble(
  x = runif(30), 
  y = runif(30)
) |>
  arrange(x,y)
  
d |>
  mutate(i = row_number()) |>
  ggplot(aes(x,y)) +
  geom_point() +
  ggrepel::geom_text_repel(aes(label=i), size=5)
```
:::
::::





## Neighbors in practice

```{r}
#| echo: false

library(gganimate)

cur = d |>
  mutate(i = row_number())

nn = RANN::nn2(d, d, k=30)$nn.idx |>
  apply(1, c, simplify = FALSE) |>
  tibble(nn = _) |>
  mutate(
    i = row_number(),
    nn = map2(i, nn, ~ .y[.y < .x][1:5] |> na.omit())
  ) |> 
  unnest_longer(nn) |>
  left_join(
    cur, by = c("nn" = "i")
  )


g = ggplot(mapping = aes(x,y)) +
  geom_point(data = d) +
  geom_point(data = cur, color="red", size=3) +
  geom_point(data = nn, color="cyan", size=2) +
  transition_states(i,transition_length = 0.5, state_length=5, wrap=FALSE) 

animate(g, nframes = 300, fps = 10)
```


## Benefits?

- The NN approximation is a **huge** computational savings over the full likelihood.

  - Dealing with $n$ $k \times k$ matrices is much easier than dealing with a single $n \times n$ matrix.
  
  - The NN approximation is $O(nk^3)$ while the full likelihood is $O(n^3)$.

- The NN approximation is also a **huge** memory savings over the full likelihood.

- The NNGP is a well-defined spatial process




## NNGP Model Results

```{r echo=FALSE, message=FALSE}
#| echp: false
#| message: false

library(spNNGP)

load("Lec22_pp_data.Rdata")

if (!file.exists("Lec22_nn_models.rds")) {
  n.samples = 10000
  starting = list("phi"=3/0.3, "sigma.sq"=1, "tau.sq"=0.1)

  tuning = list("phi"=0.1, "sigma.sq"=0.1, "tau.sq"=0.1)
  
  priors = list(
    "phi.Unif"=c(3/1, 3/0.1), 
    "sigma.sq.IG"=c(2, 2),
    "tau.sq.IG"=c(2, 2)
  )
  
  cov.model = "gaussian"
  
  nn = c(3,5,10,15)
  nn_models = lapply(
    nn, 
    function(n) {
      spNNGP::spNNGP(z~1, data=d, coords=d[,c("x","y")] %>% as.matrix(), starting=starting,
           tuning=tuning, priors=priors, cov.model=cov.model,
           n.samples=n.samples, verbose=TRUE, n.neighbors = n, 
           n.report = 1000)
    }  
  ) |> set_names(paste0("nn_",nn))
  
  nn_models$full = spBayes::spLM(
    z~1, data=d, coords=d[,c("x","y")] %>% as.matrix(), starting=starting,
    tuning=tuning, priors=priors, cov.model=cov.model,
    verbose=TRUE, amcmc = list(n.batch=200, batch.length=50, accept.rate=0.42)
  )
  
  saveRDS(nn_models, file="Lec22_nn_models.rds")
} else {
  nn_models = readRDS("Lec22_nn_models.rds")
}
```

```{r echo=FALSE}
if (!file.exists("Lec22_nn_predict.rds")) {
  predictions = lapply(
    names(nn_models),
    function(m) {
      rast = r
      if (m == "full") {
        pred = spBayes::spPredict(nn_models[[m]], coords, matrix(1, nrow=nrow(coords)), start = 5001, thin=50, n.report=10)
        rast[] = pred$p.y.predictive.samples |> t() |> post_summary() %>% .$post_mean
      } else {
        pred = predict(nn_models[[m]], matrix(1, nrow = nrow(coords)), coords, sub.sample = list( start = 5001, thin=50), n.report = 1000)
        rast[] = pred$p.y.0 |> t() |> post_summary() %>% .$post_mean
      }
      
      
      rast
    }
  ) |> setNames(names(nn_models)) 
  
  saveRDS(predictions, file="Lec22_nn_predict.rds")
} else {
  predictions = readRDS("Lec22_nn_predict.rds")
}
```

```{r echo=FALSE}
par(mfrow=c(2,3), mar=c(1,1,3,1))

plot(r,                  axes=FALSE, asp=0, legend=FALSE, main="True Field")
plot(predictions$nn_3,   axes=FALSE, asp=0, legend=FALSE, main="NNGP - 3 neighbors")
plot(predictions$nn_5,  axes=FALSE, asp=0, legend=FALSE, main="NNGP - 5 neighbors")
plot(predictions$full,      axes=FALSE, asp=0, legend=FALSE, main="Full GP")
plot(predictions$nn_10,  axes=FALSE, asp=0, legend=FALSE, main="NNGP - 10 neighbors")
plot(predictions$nn_15, axes=FALSE, asp=0, legend=FALSE, main="NNGP - 15 neighbors")
```

## Performance

```{r echo=FALSE}
res = data_frame(
  time = purrr::map_dbl(nn_models, ~ .$run.time[3]),
  error = purrr::map_dbl(predictions, ~ (.[] - r[])^2 %>% sum() %>% sqrt()),
  model = c("NNGP - 3", "NNGP - 5","NNGP - 10", "NNGP - 15", "Full GP") %>% forcats::as_factor()
)

ggplot(res, aes(x=time, y=error, color=model, shape=model)) +
  geom_point(size=4)
```