---
title: "Speed Comparisons"
output: 
  rmarkdown::html_vignette:
    toc: yes
vignette: >
  %\VignetteIndexEntry{Speed Comparisons}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

\newcommand{\QQ}{\mathbf{Q}}
\newcommand{\RR}{\mathbf{R}}
\newcommand{\xx}{\mathbf{x}}

```{r setup, include = FALSE}
# knitr options
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)
run.comparisons <- FALSE
set.seed(1234)
```

Here are speed and estimate comparisons of EM, BFGS with the Kalmon filter (no gradient functions), TMB and Tim Clines original code.

## Example data

```{r}
library(MARSS)
library(dplyr)
data(lakeWAplankton, package = "MARSS")
phytoplankton <- c("Cryptomonas", "Diatoms", "Greens", "Unicells", "Other.algae")
dat <- as.data.frame(lakeWAplanktonTrans) |>
  subset(Year >= 1980 & Year <= 1989) |>
  subset(select=phytoplankton) |>
  t() |>
  MARSS::zscore()
```

## Fit models without covariates

```{r}
R <- "unconstrained"
mod.list <- list(R=R, m=1, tinitx=1)
```

Fit with MARSS with EM or optim and BFGS.
```{r}
m1 <- MARSS(dat, model=mod.list, form='dfa', z.score=FALSE, silent = TRUE)
m2 <- MARSS(dat, model=mod.list, form='dfa', z.score=FALSE, silent = TRUE, method="BFGS")
```

Fit with TMB.
```{r}
library(marssTMB)
m3 <- dfaTMB(dat, model=list(m=1, R=R))
m4 <- MARSS(dat, model=mod.list, form='dfa', method="TMB", silent=TRUE)
m5 <- MARSS(dat, model=mod.list, method="BFGS_TMB", form='dfa', silent=TRUE)
```

### Log likelihoods

```{r echo=FALSE}
df <- data.frame(name=c("MARSS-EM", "MARSS-BFGS", "dfaTMB-nlminb", "MARSS_tmb-nlminb", "MARSS_tmb-optim-BFGS"),
                 logLik=c(m1$logLik, m2$logLik, m3$logLik, m4$logLik, m5$logLik))
knitr::kable(df)
```

### Compare parameter estimates

```{r echo=FALSE}
library(tidyr)
library(ggplot2)
# BFGS was swapped in sign so just mult by -1
adj <- 1
if(sign(coefficients(m1)$Z[1]) != sign(coefficients(m2)$Z[1])) adj <- -1
pars <- data.frame(
  name = c(paste0("R", rownames(coefficients(m1)$R)), paste0("Z", rownames(coefficients(m1)$Z))),
  EM = c(coefficients(m1)$R, coefficients(m1)$Z),
  BFGS = c(coefficients(m2)$R, adj*coefficients(m2)$Z),
  TMB1 = c(as.vector(m3$Estimates$R[lower.tri(m3$Estimates$R,diag=TRUE)]), as.vector(m3$Estimates$Z)),
  TMB2 = c(coefficients(m4)$R, coefficients(m4)$Z))
pars <- pars %>% tidyr::pivot_longer(!name, names_to = "model")

dodge <- position_dodge(width=0.5)
ggplot(pars, aes(x=name, y=value, col=model)) +
  geom_point(position=dodge) +
  ggtitle("same estimates")
```

## Add example with covariates

For `form="dfa"`, pass in covariates with `covariates=xyz`. If using the default model form (not dfa), then pass in covariates with `model$d` or `model$c`.

### Fit model

```{r}
# use a simpler R
R <- "diagonal and unequal"
mod.list2 <- list(m=3, R=R, tinitx=1)
# add a temperature covariate
temp <- as.data.frame(lakeWAplanktonTrans) |>
    subset(Year >= 1980 & Year <= 1989) |>
    subset(select=Temp)
covar <- t(temp) |> zscore()
m <- MARSS(dat, model=mod.list2, form="dfa", covariates=covar, silent = TRUE, z.score = FALSE, method="TMB", fit=FALSE)
t6 <- system.time(m6 <- marssTMB:::MARSSfit.TMB(m, fun=1))
t6c <- system.time(m6c <- marssTMB:::MARSSfit.TMB(m, fun=2))
t7 <- system.time(m7 <- MARSS(dat, model=mod.list2, form="dfa", covariates=covar, silent = TRUE, control=list(maxit=10000), z.score = FALSE))
t6b <- system.time(m6b <- dfaTMB(dat, model=list(m=3, R=R), 
                    EstCovar = TRUE, Covars = covar))
```

Add a 2nd covariate

```{r}
TP <- as.data.frame(lakeWAplanktonTrans) |>
    subset(Year >= 1980 & Year <= 1989) |>
    subset(select=TP)
covar <- rbind(covar, t(TP)) |> zscore()
m <- MARSS(dat, model=mod.list2, form="dfa", covariates=covar, silent = TRUE, z.score=FALSE, method="TMB", fit=FALSE)
t8 <- system.time(m8 <- marssTMB:::MARSSfit.TMB(m, fun=1))
t8c <- system.time(m8c <- marssTMB:::MARSSfit.TMB(m, fun=2))
t9 <- system.time(m9 <- MARSS(dat, model=mod.list2, form="dfa", covariates=covar, silent = TRUE, control=list(maxit=10000), z.score=FALSE))
t8b <- system.time(m8b <- dfaTMB(dat, model=list(m=3, R=R), 
                    EstCovar = TRUE, Covars = covar))
```

### Compare time and log likelihoods

The MARSS_tmb-1 uses the splitting approach for the diag versus the chol, while tmb-2 is using just the chol. Clearly splitting is faster.

```{r echo=FALSE}
df <- data.frame(name=rep(c("MARSS-EM", "MARSS_tmb-1", "MARSS_tmb-2", "dfaTMB"), 2),
                 num_covar = rep(1:2, each = 4),
                 time=c(t7[1], t6[1], t6c[1], t6b[1], t9[1], t8[1], t8c[1], t8b[1]),
                 logLik=c(m7$logLik, m6$logLik, m6c$logLik, m6b$logLik, m9$logLik, m8$logLik, m8c$logLik, m8b$logLik))
knitr::kable(df)
```

### Compare parameter estimates

```{r echo=FALSE}
library(tidyr)
n1 <- length(coef(m6, type="vector"))
n2 <- length(coef(m8, type="vector"))
pars <- data.frame(
  name = c(rep(names(coef(m6, type="vector")), 2), rep(names(coef(m8, type="vector")), 2)),
  ncovar = as.factor(c(rep(c(1, 1), each=n1), rep(c(2, 2), each = n2))),
  model = c(rep(c("EM", "TMB"), each=n1), rep(c("EM", "TMB"), each = n2)),
  value = c(coef(m6, type="vector"), coef(m7, type="vector"), coef(m8, type="vector"), coef(m9, type="vector")))
```

```{r echo=FALSE}
library(ggplot2)
dodge <- position_dodge(width=1)
ggplot(pars, aes(x=name, y=value, col=model)) +
  geom_point(position=dodge) +
  coord_flip() + 
  facet_wrap(~ncovar) +
  ggtitle("estimates should be the same (roughly)")
```

## More MARSS models

```{r echo = FALSE}
dat <- t(harborSealWA)
dat <- dat[2:4, ] # remove the year row
contrl <- list(maxit = 10000)
mod <- list(Z = matrix(1, 3, 1), R = "diagonal and equal")
fit1 <- MARSS(dat, model = mod, control = contrl, silent = TRUE)
fit2 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="BFGS")
fit3 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="TMB")
```
```{r}
cat("mod <- list(Z = matrix(1, 3, 1), R = 'diagonal and equal')\n")
c(EM=fit1$logLik, BFGS=fit2$logLik, TMB=fit2$logLik)
```
```{r echo=FALSE}
fit1 <- MARSS(dat, control = contrl, silent = TRUE)
fit2 <- MARSS(dat, control = contrl, silent = TRUE, method="BFGS")
fit3 <- MARSS(dat, control = contrl, silent = TRUE, method="TMB")
```
```{r}
cat("MARSS(dat)\n")
c(EM=fit1$logLik, BFGS=fit2$logLik, TMB=fit2$logLik)
```

```{r echo=FALSE}
mod <- list(Q = "unconstrained")
fit1 <- MARSS(dat, model = mod, control = contrl, silent = TRUE)
fit2 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="BFGS")
fit3 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="TMB")
```
```{r}
cat("mod <- list(Q = 'unconstrained')\n")
c(EM=fit1$logLik, BFGS=fit2$logLik, TMB=fit2$logLik)
```

```{r echo=FALSE}
mod <- list(U = matrix(c("N", "S", "S")))
fit1 <- MARSS(dat, model = mod, control = contrl, silent = TRUE)
fit2 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="BFGS")
fit3 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="TMB")
```
```{r}
cat('list(U = matrix(c("N", "S", "S")))\n')
c(EM=fit1$logLik, BFGS=fit2$logLik, TMB=fit2$logLik)
```

```{r echo=FALSE}
mod <- list(U = matrix(list("N", "N", 0), 3, 1), R = diag(0.01, 3))
fit1 <- MARSS(dat, model = mod, control = contrl, silent = TRUE)
fit2 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="BFGS")
fit3 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="TMB")
```
```{r}
cat('list(U = matrix(list("N", "N", 0), 3, 1), R = diag(0.01, 3))\n')
c(EM=fit1$logLik, BFGS=fit2$logLik, TMB=fit2$logLik)
```

```{r echo=FALSE}
Z <- matrix(c(1, 1, 0, 0, 0, 1), 3, 2)
colnames(Z) <- c("N", "S")
mod <- list(Z = Z, Q = "diagonal and equal", R = "diagonal and unequal", U = "equal")
fit1 <- MARSS(dat, model = mod, control = contrl, silent = TRUE)
fit2 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="BFGS")
fit3 <- MARSS(dat, model = mod, control = contrl, silent = TRUE, method="TMB")
```
```{r}
cat('list(Z = Z, Q = "diagonal and equal", R = "diagonal and unequal", U = "equal")\n')
c(EM=fit1$logLik, BFGS=fit2$logLik, TMB=fit2$logLik)
```

## Run some time comparisons

```{r include=FALSE}
# no covars
if(!run.comparisons){
  load("dfa-time-comparisons.rda")
}else{
df <- c()
mods <- list()
for(R in c("diagonal and equal", "diagonal and unequal", "unconstrained")){
  for(m in 1:3){
    mod <- list(m = m, R = R, tinitx=1)
    tfit <- system.time(fit <- MARSS(dat, model=mod, form="dfa", control=list(maxit=10000)))
    df <- rbind(df, data.frame(fun="MARSS", opt.function="EM", m=m, R=R, ncovar = 0, time=tfit[1], logLik=fit$logLik, convergence = fit$convergence))
    mods <- c(mods, list(fit))
    tfit <- system.time(fit <- MARSS(dat, model=mod, form="dfa", method="BFGS"))
    df <- rbind(df, data.frame(fun="MARSS", opt.function="BFGS", m=m, R=R, ncovar = 0, time=tfit[1], logLik=fit$logLik, convergence = fit$convergence))
    mods <- c(mods, list(fit))
    tfit <- system.time(fit <- MARSS(dat, model=mod, form="dfa", method="TMB"))
    df <- rbind(df, data.frame(fun="TMB", opt.function="nlminb", m=m, R=R, ncovar = 0, time=tfit[1], logLik=fit$logLik, convergence = fit$convergence))
    mods <- c(mods, list(fit))
  }
}
}
```


```{r include=FALSE}
# with covars
if(!run.comparisons){
  load("dfa-time-comparisons.rda")
}else{
for(R in c("diagonal and equal", "diagonal and unequal", "unconstrained")){
  for(m in 1:3){
    mod <- list(m = m, R = R, tinitx=1)
    tfit <- system.time(fit <- MARSS(dat, model=mod, form="dfa", covariates=covar, control=list(maxit=10000)))
    df <- rbind(df, data.frame(fun="MARSS", opt.function="EM", m=m, R=R, ncovar=2, time=tfit[1], logLik=fit$logLik, convergence = fit$convergence))
    mods <- c(mods, list(fit))
    tfit <- system.time(fit <- MARSS(dat, model=mod, form="dfa", covariates=covar, method="BFGS"))
    df <- rbind(df, data.frame(fun="MARSS", opt.function="BFGS", m=m, R=R, ncovar=2, time=tfit[1], logLik=fit$logLik, convergence = fit$convergence))
    mods <- c(mods, list(fit))
    tfit <- system.time(fit <- MARSS(dat, model=mod, form="dfa", covariates=covar, method="TMB"))
    df <- rbind(df, data.frame(fun="TMB", opt.function="nlminb", m=m, R=R, ncovar=2, time=tfit[1], logLik=fit$logLik, convergence = fit$convergence))
    mods <- c(mods, list(fit))
  }
}
}
save(df, mods, file="dfa-time-comparisons2.rda")
```

```{r include=FALSE}
df2 <- df |> mutate(mod = paste0(fun, "-", opt.function)) |>
  mutate(ncovar = as.factor(ncovar)) |>
  group_by(ncovar)
ggplot(df2, aes(alpha=ncovar, fill=mod, y=time, x=m)) + 
    geom_bar(stat="identity", position="dodge", color="black") +
  facet_wrap(~R, scales = "free_y") +
  scale_y_continuous() +
  ggtitle("TMB is faster than MARSS EM")
```
```{r}
df2 <- df %>% mutate(mod = paste0(fun, "_", opt.function)) |>
  select(mod, time, m, R, ncovar) %>%
  pivot_wider(id_cols=c(m, R, ncovar), names_from = mod, values_from = time) |>
  mutate(
    MARSS_EM = MARSS_EM/TMB_nlminb,
    MARSS_BFGS = MARSS_BFGS/TMB_nlminb,
    TMB_nlminb = TMB_nlminb/TMB_nlminb
  ) |>
  mutate(ncovar = as.factor(ncovar)) |>
  pivot_longer(cols = 4:6)
mean.improv <- df2 |> subset(name != "TMB_nlminb") |> summarize(mean = mean(value))
ggplot(df2, aes(alpha=ncovar, fill=name, y=value, x=m)) + 
    geom_bar(stat="identity", position="dodge", color="black") +
  facet_wrap(~R, scales = "free_y") +
  ylab("Time relative to TMB") +
  scale_y_continuous() +
  ggtitle(paste("speed relative to TMB.", round(mean.improv, digits=2), "fold improvement"))
```

```{r}
df2 <- df |> mutate(mod = paste0(fun, "-", opt.function))
df2$ncovar <- as.factor(df2$ncovar)
ggplot(df2, aes(col=ncovar, y=logLik, x=m, shape=mod)) + 
    geom_point(position=position_dodge(width=0.3)) +
  facet_wrap(~R, scales = "free_y") +
  scale_y_continuous() +
  ggtitle("logLik comparison")
```