em_vs_mixsqp.Rmd

---
title: Comparing two different ways of estimating the mixture weights in mr-ash
author: Peter Carbonetto
date: February 20, 2020
site: workflowr::wflow_site
output: workflowr::wflow_html
---

Here we compare two different ways of updating the mixture weights
in mr-ash: using EM, and using mix-SQP.

```{r knitr, echo=FALSE}
knitr::opts_chunk$set(comment = "#",results = "hold",collapse = TRUE,
                      fig.align = "center")
```

## Script parameters

These are the data simulation settings.

I suggest trying this example with `s = 0` (all predictors are
independent) and `s = 0.4` (the correlation between all predictors
is 0.4).

```{r datasim-params}
n  <- 500
p  <- 100
sd <- c(0,   1,    2)
w  <- c(0.9, 0.05, 0.05)
s  <- 0
```

This specifies the variances for the mixture-of-normals prior on the
regression coefficients.

```{r set-s0}
s0 <- 10^seq(-4,0,length.out = 12)
```

## Load functions

These are the packages used in this analysis.

```{r load-pkgs, message=FALSE, warning=FALSE}
library(ggplot2)
library(cowplot)
library(MASS)
library(mixsqp)
```

This R code provides a very simple implementation of the mr-ash
algorithm.

```{r load-functions}
source("../code/misc.R")
source("../code/mr_ash.R")
```

## Simulate data

The predictors are drawn from the multivariate normal with zero mean
and covariance matrix S, in which all diagonal entries are 1, and
all off-diagonal entries are s. Setting `s = 0.5` reproduces the
simulation of the predictors used in Example 3 of Zou & Hastie
(2005).

```{r simulate-data}
set.seed(2)
S       <- matrix(s,p,p)
diag(S) <- 1
X       <- mvrnorm(n,rep(0,p),S)
k       <- sample(length(w),p,replace = TRUE,prob = w)
beta    <- sd[k] * rnorm(p)
y       <- drop(X %*% beta + rnorm(n))
```

## Fit model

These are the initial estimates of residual variance (`s`), mixture
weights (`w0`), and posterior mean estimates of the regression
coefficients (b).

```{r init-model}
k  <- length(s0)
se <- 1
w0 <- rep(1/k,k)
b  <- rep(0,p)
```

Fit the model by running 200 EM updates for the mixture weights.

```{r fit-model-em}
fit1 <- mr_ash(X,y,se,s0,w0,b,maxiter = 200,verbose = FALSE)
```

Fit the model a second time using the mix-SQP updates for the mixture
weights.

```{r fit-model-mixsqp}
fit2 <- mr_ash_with_mixsqp(X,y,se,s0,w0,b,numiter = 10,maxiter.inner = 50,
                           tol.inner = 1e-8)
```

## Review model fit

Compare the posterior mean estimates from the two fits.

```{r, fig.height=3, fig.width=3}
ggplot(data.frame(em = fit1$b,mixsqp = fit2$b),
             aes(x = em,y = mixsqp)) +
  geom_point(color = "darkblue") +
  geom_abline(intercept = 0,slope = 1,col = "magenta",lty = "dotted") +
  labs(title = "coefs") +
  theme_cowplot()
```

Compare the posterior mean estimates against the values used to
simulate the data.

```{r, fig.height=3, fig.width=3}
ggplot(data.frame(true = beta,mixsqp = fit2$b),
             aes(x = true,y = mixsqp)) +
  geom_point(color = "darkblue") +
  geom_abline(intercept = 0,slope = 1,col = "magenta",lty = "dotted") +
  labs(title = "coefs") +
  theme_cowplot()
```

Plot the improvement in the solution over time. The EM updates are
shown in blue, and the mix-SQP updates are shown in orange.

```{r fig.height=4, fig.width=4}
elbo.best <- max(c(fit1$elbo,fit2$elbo))
pdat      <- rbind(data.frame(update = "em",
                              iter   = 1:length(fit1$elbo),
                              elbo   = fit1$elbo),
                   data.frame(update = "mixsqp",
                              iter   = cumsum(fit2$niter),
                              elbo   = fit2$elbo))
pdat$elbo <- elbo.best - pdat$elbo + 1e-4
ggplot(pdat,aes(x = iter,y = elbo,color = update)) +
  geom_line() +
  geom_point() +
  scale_y_log10() +
  scale_color_manual(values = c("royalblue","darkorange")) +
  labs(y = "distance to \"best\" elbo",title = paste("s =",s)) +
  theme_cowplot()
```