---
title: "Custom spatial models with RStan and geostan"
date: October 2, 2023
author: Connor Donegan
output: 
  rmarkdown::html_vignette:
    toc: true 
    toc_depth: 3  # upto three depths of headings (specified by #, ## and ###)
    number_sections: true    
header-includes:
   - \usepackage{amsmath}
vignette: >
  %\VignetteIndexEntry{Custom spatial models with RStan and geostan}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
bibliography: spatial-me.bib
link-citations: yes
---

```{r, include=FALSE}
knitr::opts_chunk$set(eval = FALSE, echo = TRUE)
```

The spatial models in geostan use custom Stan functions that are far more efficient than using built-in functions, including the conditional (CAR) and simultaneous spatial autoregressive (SAR) models (both are particular specifications of the multivariate normal distribution). This vignette shows how you can use those functions together with some R functions in geostan to start building custom spatial models.

This tutorial is written with the assumption that the reader is already familiar with [RStan](https://mc-stan.org/users/documentation/). Users are expected to make adjustments to the code as needed, including to prior distributions. For more details and an explanation of the computational approach, see @donegan_2022. For more background on spatial modeling see @haining_2020. 

# CAR models

The CAR model is a multivariate normal distribution with covariance matrix $$\Sigma = (I - \rho \cdot C)^{-1} M,$$ where $I$ is the identity matrix, $\rho$ a spatial autocorrelation parameter which may be positive or negative, $C$ is a sparse connectivity matrix, and $M$ is a diagonal matrix with scale parameters $\tau_i^2$. There is typically a single scale parameter $\tau$ multiplied by weights $\delta_i$. The term $\tau^2 \cdot \delta_i$ is the conditional variance pertaining to $y_i | y_{n(i)}$ where $n(i)$ lists the areas that neighbor the $i^{th}$ one.

The defining feature of the CAR model is that, similar to time-series autoregression, the expectation for the value at the $i^{th}$ site given any or all other values of $y$ besides $i$ ($y_{-i}$) are known, is a function of the neighboring values: $$\mathbb{E}[y_i | y_{-i}] = \mu_j + \sum_{j \in n(i)} \rho \cdot c_{ij} (y_j - \mu_j).$$

The CAR function presented here is valid for what is sometimes called the WCAR specification. This is the most commonly employed CAR specification (see @donegan_2022 for CAR models without this restriction). The connectivity matrix is row-standardized. Let $A$ be a symmetric binary connectivity matrix where $a_{ij}= 1$ only if the $i^{th}$ and $j^{th}$ sites are neighbors, all other entries are zero (including the diagonal entries $i=j$). The elements of $C$ are
$$c_{ij} = \frac{a_{ij}}{\sum_j^n a_{ij}} = \frac{a_{ij}}{|n(i)|}.$$

The diagonal of matrix $M$ then is specified by the n-length vector of weights $\delta_i = \tau^2 \cdot \frac{1}{|n(i)|}$.

The `geostan::prep_car_data` function will complete all of this for you. The user passes in a binary connectivity matrix $A$ and and specifies `style = "WCAR"`, then the function returns a list of data inputs for the Stan model. 

## Autonormal model

This first example is an autonormal model: $$y \sim Normal(\alpha + x \beta, \Sigma).$$ Here you have $n$ observations of a continuously measured outcome $y$, possibly $k=1$ or more covariates $x$, and you're accounting for spatial autocorrelation in $y$ using the covariance matrix of an $n$-dimensional multivariate normal distribution.

The following Stan code implements the above model; if you're following along, copy the Stan code and save it in a file inside your working directory; I'll name it "autonormal.stan" and I'll save the name as an R variable:

```{r}
autonormal_file <- "autonormal.stan"
```

Here's the Stan code:

```{stan output.var= "auto_normal"}
functions {
#include wcar-lpdf.stan
}

data {
  // data
  int<lower=1> n;
  int<lower=1> k;
  vector[n] y;
  matrix[n, k] x;

    // CAR parts
  int nAx_w;
  int nC;
  vector[nAx_w] Ax_w;
  array[nAx_w] int Ax_v;
  array[n + 1] int Ax_u;
  array[nC] int Cidx;
  vector[n] Delta_inv;
  real log_det_Delta_inv;
  vector[n] lambda;
}

parameters {
  // spatial autocorrelation (SA) parameter
  real<lower=1/min(lambda), upper=1/max(lambda)> rho;
  
  // scale parameter
  real<lower=0> tau;
  
  // intercept
  real alpha;
  
  // coefficients
  vector[k] beta;  
}

model {
  vector[n] mu = alpha + x * beta;

  // Likelihood: y ~ Normal(Mu, Sigma)
  target += wcar_normal_lpdf(y |
			     mu, tau, rho, // mean, scale, SA
			     Ax_w, Ax_v, Ax_u, // stuff from prep_car_data
			     Delta_inv, 
			     log_det_Delta_inv,
			     lambda, 
			     n);    
  
  // prior for scale parameter
  target += student_t_lpdf(tau | 10, 0, 5);

  // prior for beta
  target += normal_lpdf(beta | 0, 5);

  // prior for intercept
  target += normal_lpdf(alpha | 0, 5);
}

```

The input file "wcar-lpdf.stan" contains the following code (save this code inside your working directory and name it "wcar-lpdf.stan"):

```{stan output.var = 'wcar_fun', eval = FALSE}
/**
 * Log probability density of the conditional autoregressive (CAR) model: WCAR specifications only
 *
 * @param y Process to model
 * @param mu Mean vector
 * @param tau Scale parameter
 * @param rho Spatial dependence parameter
 * @param A_w Sparse representation of the symmetric connectivity matrix, A
 * @param A_v Column indices for values in A_w
 * @param A_u Row starting indices for values in A_u
 * @param D_inv The row sums of A; i.e., the diagonal elements from the inverse of Delta, where M = Delta * tau^2 is a diagonal matrix containing the conditional variances.
 * @param log_det_D_inv Log determinant of Delta inverse.
 * @param lambda Eigenvalues of C (or of the symmetric, scaled matrix Delta^{-1/2}*C*Delta^{1/2}); for the WCAR specification, C is the row-standardized version of W.
 * @param n Length of y
 *
 * @return Log probability density of CAR prior up to additive constant
 */
real wcar_normal_lpdf(vector y, vector mu,
		      real tau, real rho,
		      vector A_w, array[] int A_v, array[] int A_u,
		      vector D_inv, real log_det_D_inv,
		      vector lambda,
		      int n) {
  vector[n] z = y - mu;
  real ztDz; 
  real ztAz; 
  vector[n] ldet_ImrhoC;
  ztDz = (z .* D_inv)' * z;
  ztAz = z' * csr_matrix_times_vector(n, n, A_w, A_v, A_u, z);
  for (i in 1:n) ldet_ImrhoC[i] = log1m(rho * lambda[i]);
  return 0.5 * (
		-n * log( 2 * pi() )
		-2 * n * log(tau)
		+ log_det_D_inv
		+ sum(ldet_ImrhoC)
		- (1 / tau^2) * (ztDz - rho * ztAz));
}
```

From R, you can use the following code to prepare the input data ($y$, $x$, etc.). I use `prep_car_data` to get a list of parts for the CAR model, then I append the outcome data to the same list and pass it all to a Stan model to draw samples.

```{r}
library(rstan)
library(geostan)
data(georgia)

A <- shape2mat(georgia, "B")
car_list <- prep_car_data(A, style = "WCAR")

# add data
car_list$y <- log(georgia$income / 10e3)
car_list$x <- matrix(log(georgia$population / 10e3), ncol = 1)
car_list$k <- ncol(car_list$x)

# compile Stan model from file
car_model <- stan_model(autonormal_file)

# sample from model
samples <- sampling(car_model, data = car_list, iter = 1e3)
```

The same results can be obtained using `geostan::stan_car`:

```{r}
fit <- stan_car(log(income / 10e3) ~ log(population / 10e3),
    data = georgia, car = car_list, iter = 1e3)
```


## Poisson models

Disease mapping is a common use for CAR models, though these models have many applications. In this class of models, the CAR model is assigned as the prior distribution to a parameter vector $\phi$, which is used to model disease incidence rates across small areas like counties. The disease data consist of counts $y$ together with the size of population at risk $p$, or possibly a different denominator such as the expected number of cases $E$ (which occurs when using indirect age-standardization). These models have the form
\begin{equation}
\begin{aligned}
	&y \sim Poisson(e^\phi) \\
	&\phi \sim Normal(\mu, \Sigma) \\
	&\mu = log(p) + \alpha + x \beta.
\end{aligned}
\end{equation}

Here you see that the linear predictor $\mu$ is embedded with the CAR model.

Here is another way of writing down the very same model:
\begin{equation}
\begin{aligned}
	&y \sim Poisson(e^\eta) \\
	&\eta = log(p) + \alpha + x \beta + \phi \\
	&\phi \sim Normal(0, \Sigma) \\
\end{aligned}
\end{equation}

It is possible to build a Stan modeling using either one of these two formulations. They are substantively equivalent, *but* you may find that one is far more amenable to Stan's MCMC algorithm. The former specification ($\mu$ inside the CAR model) is less commonly presented in papers, but in this author's experience it is often far more stable computationally than forcing the CAR model to have zero mean. (You may very well encounter a case where the opposite is true.)

The purpose of the offset term is sometimes unclear in introductory texts. The offset term $p$ is from the denominator of the rates $\frac{y}{p}$. The expectation or mean of the model is
\begin{equation}
\begin{aligned}
	&\mathbb{E}[y] = e^{log(p) + \mu} = e^{log(p)} \cdot e^{\mu} = p \cdot e^\mu \\
	&\mathbb{E}\Bigl[ \frac{y}{p} \Bigr] = e^\mu
\end{aligned}
\end{equation}
The smaller is the denominator, the less informative is the rate with respect to the characteristic level of risk bearing upon the population of the given period and place. With small denominators, chance renders the rates uninformative (and unstable over time and space), and the Poisson model accounts for this.

The following Stan code provides a template for a simple disease mapping model using the CAR model as a prior for $\phi$. As before, you can save the code in your working directory and give it a name like "car_poisson.stan". I'll store the file name in my R environment:

```{r}
car_poisson_file <- "car_poisson.stan"
```

```{stan output.var = "car_poisson"}
functions {
#include wcar-lpdf.stan
}

data {
  // data
  int<lower=1> n;
  int<lower=1> k;
  array[n] int<lower=0> y;
  matrix[n, k] x;
  vector[n] const_offset;

    // CAR parts
  int nAx_w;
  int nC;
  vector[nAx_w] Ax_w;
  array[nAx_w] int Ax_v;
  array[n + 1] int Ax_u;
  array[nC] int Cidx;
  vector[n] Delta_inv;
  real log_det_Delta_inv;
  vector[n] lambda;
}

parameters {
  // spatial autocorrelation (SA) parameter
  real<lower=1/min(lambda), upper=1/max(lambda)> rho;
  
  // scale parameter
  real<lower=0> tau;
  
  // intercept
  real alpha;
  
  // coefficients
  vector[k] beta;
  
  // SA trend component
  vector[n] phi;
}


model {
  vector[n] mu = const_offset + alpha + x * beta;
  
  // Likelihood: y ~ Poisson(e^[phi])
  target += poisson_lpmf(y | exp(phi));
    
  // phi ~ Normal(Mu, Sigma)
  target += wcar_normal_lpdf(phi |
			     mu, tau, rho, // mean, scale, SA
			     Ax_w, Ax_v, Ax_u, // stuff from prep_car_data
			     Delta_inv, 
			     log_det_Delta_inv,
			     lambda, 
			     n);

  // prior for scale parameter
  target += student_t_lpdf(tau | 10, 0, 1);

  // prior for beta
  target += normal_lpdf(beta | 0, 5);

  // prior for intercept
  target += normal_lpdf(alpha | 0, 5);
}

```

Again, you can use `geostan::prep_car_data` to easily convert the spatial weights matrix into the list of required inputs for the CAR model:

```{r}
library(rstan)
library(geostan)
data(georgia)

A <- shape2mat(georgia, "B")
car_list <- prep_car_data(A, style = "WCAR")

# add data
car_list$y <- georgia$deaths.male
car_list$const_offset <- log(georgia$pop.at.risk.male)
car_list$x <- matrix(log(georgia$income / 10e3), ncol = 1)
car_list$k <- ncol(car_list$x)

# compile Stan model from file
car_poisson <- stan_model(car_poisson_file)

# sample from model
samples <- sampling(car_poisson,
	data = car_list,
	iter = 1e3)
```

The same results can be obtained using `geostan::stan_car`:

```{r}
fit <- stan_car(deaths.male ~ offset(log(pop.at.risk.male)) + log(income / 10e3),
    data = georgia, car = car_list, family = poisson(), iter = 1e3)
```

# SAR models

The simultaneously-specified spatial autoregression (SAR) is written as
\begin{equation}
  \begin{aligned}
    y = \mu + (I - \rho \cdot W)^{-1} \epsilon \\
    \epsilon \sim Normal(0, \sigma^2 \cdot I)
    \end{aligned}
\end{equation}	
where $W$ is a row-standardized spatial weights matrix, $I$ is the n-by-n identity matrix, and $\rho$ is a spatial autocorrelation parameter. This is also a multivariate normal distribution but with a different covariance matrix than the CAR model:
\begin{equation}
  \Sigma = \sigma^2 \cdot (I - \rho \cdot W)^{-1} (I - \rho \cdot W^T)^{-1},
\end{equation}
where $T$ is the matrix transpose operator. The SA parameter $\rho$ for the SAR model has a more intuitive connection to the degree of SA than the CAR model (which is usually above .95 for moderately high SA).

The Stan function that is used by `geostan::stan_sar` is as follows (the R code below will assume that you have saved this as "sar-lpdf.stan"):

```{stan output.var = "sar_fun", eval = FALSE}
/**
 * Log probability density of the simultaneous autoregressive (SAR) model (spatial error model)
 *
 * @param y Process to model
 * @param mu Mean vector
 * @param sigma Scale parameter
 * @param rho Spatial dependence parameter
 * @param ImW Sparse representation of (I - W): non-zero values only
 * @param ImW_v Column indices for values in ImW
 * @param ImW_u Row starting indices for values in ImW
 * @param Widx Indices for the off-diagonal elements in ImC
 * @param lambda Eigenvalues of W
 * @param n Length of y
 *
 * @return Log probability density of SAR model up to additive constant
*/
  real  sar_normal_lpdf(vector y,
		      vector mu,
		      real sigma,
		      real rho,
		      vector ImW,
		      array[] int ImW_v,
		      array[] int ImW_u,
		      array[] int Widx,
		      vector lambda,
		      int n) {
  vector[n] z = y - mu;
  real tau = 1 / sigma^2;
  vector[num_elements(ImW)] ImrhoW = ImW;  // (I - rho W)
  vector[n] ImrhoWz;                       // (I - rho * W) * z
  real zVz;
  real ldet_V = 2 * sum(log1m(rho * lambda)) - 2 * n * log(sigma);
  ImrhoW[Widx] = rho * ImW[Widx];
  ImrhoWz = csr_matrix_times_vector(n, n, ImrhoW, ImW_v , ImW_u , z);
  zVz = tau * dot_self(ImrhoWz);
  return  0.5 * (-n * log(2 * pi()) + ldet_V - zVz);		 
 }    

```

The following Stan model provides an example of how to use this function to build a SAR model. Again, its an autonormal model for continuous outcome variable with $k$ covariates. 

```{r}
sar_model_file <- "sar_model.stan"
```

```{stan output.var = "auto_sar"}
functions {
#include sar-lpdf.stan
}

data {
  // data
  int<lower=1> n;
  int<lower=1> k;
  vector[n] y;
  matrix[n, k] x;


// SAR
  int nImW_w;
  int nW;
  vector[nImW_w] ImW_w;
  array[nImW_w] int ImW_v;
  array[n + 1] int ImW_u;
  array[nW] int Widx;
  vector[n] eigenvalues_w;
}

parameters {
  // SA parameter
  real<lower=1/min(eigenvalues_w), upper=1/max(eigenvalues_w)> rho;     

  // scale parameter
  real<lower=0> sigma;

  // intercept
  real alpha;

  // coefficients
  vector[k] beta;
}

model{
  vector[n] mu = alpha + x * beta;

  // Likelihood: Y ~ Normal(Mu, Sigma)
  target += sar_normal_lpdf(y |
                  mu, sigma, rho,
                  ImW_w,
                  ImW_v,
                  ImW_u,
                  Widx,
                  eigenvalues_w,
                  n);

  // prior for scale parameter
  target += student_t_lpdf(sigma | 10, 0, 5);

  // prior for beta
  target += normal_lpdf(beta | 0, 5);

  // prior for intercept
  target += normal_lpdf(alpha | 0, 5);
}

```


```{r}
library(geostan)
library(rstan)
data(georgia)

W <- shape2mat(georgia, "W")
sar_list <- prep_sar_data(W)

# add data
sar_list$y <- log(georgia$income / 10e3)
sar_list$x <- matrix(log(georgia$population / 10e3), ncol = 1)
sar_list$k <- ncol(sar_list$x)

# compile Stan model from file
sar_model <- stan_model(sar_model_file)

# sample from model
samples <- sampling(sar_model, data = sar_list, iter = 1e3)
```

We can draw samples from the same model using `geostan::stan_sar`:

```{r}
sar_dl <- prep_sar_data(W)
fit <- stan_sar(log(income / 10e3) ~ log(population / 10e3), data = georgia,
    sar_parts = sar_dl, iter = 1e3)
```

One can also use the SAR model as a prior distribution for a parameter vector, just as was done above with the CAR model.


# References
