plot.Rmd

---
title: "autodiff"
author: "Mathieu"
date: '`r format(Sys.time(), "%d %B, %Y")`'
output:
  pdf_document:
    toc: true
    toc_depth: 3

---

```{r setup, include=FALSE}
library(ggplot2)
library(dplyr)
library(knitr)
library(stringr)
library(tidyr)
library(scales)
library(grid)
library(gtable)
library(cowplot)
library(khroma)
knitr::opts_chunk$set(echo = TRUE)
options(dplyr.summarise.inform = FALSE)
knitr::opts_chunk$set(fig.path='Figs/')
```

# Benchmark description

## Macro benchmark

- flu dataset from treetime_validation
- 5000 iterations
- model:
  - JC69
  - time tree
  - constant population size
  - exp(1000) prior on strict clock
  - OneOnX prior on population size

## Micro benchmark

- flu dataset from treetime_validation
- 1000 replicates
- treelikelihood:
  - input: branch lengths
  - output: log likleihood
  - derivatives wrt branch lengths
- coalescent:
  - input: node heights and pop size
  - output: log likelihood
  - derivatives wrt constant pop size and node heights (not ratios)
- ratio_transform:
  - input: ratios/root height
  - output: internal node heights
  - derivatives wrt ratios
- ratio_transform_jacobian:
  - input: internal node heights
  - output: log determinant of the Jacobian
  - derivatives wrt internal node heights

Programs:

- torchtree: pure pytorch
- bitorch: torchtree with BITO (analytic derivatives of tree likelihood calculated with BEAGLE)
- torchtree SP: torchtree with single precision
- phylojax: pure JAX
- treeflow: pure treeflow
- physher: pure C (analytic derivatives)
- phylostan: Stan language

\newpage
# Results

```{r include=FALSE}
programs = c(
  "physher",
  "torchtree/bito/BEAGLE",
  "torchtree",
  "phylojax",
  "phylostan",
  "treeflow",
  "torchtree SP",
  "jax-analytic",
  "phylojax1"
)
progShapes = setNames(c(4, 4, 16, 16, 16, 16, 17, 16, 1), programs)
bright6 <- colour("bright")(6)
progColors = setNames(c(bright6, c(bright6[3], bright6[4], bright6[4])), programs)

theme_set(theme_bw(16) + theme(strip.background = element_blank()))
```
## Macro benchmark

### Mermory usage
```{r macro-mem}
df.trace = read.csv("results/trace.csv", sep = '\t')
df.trace = df.trace %>%
  group_by(program, size) %>% summarize(max_peak_vmem = min(peak_rss)) %>% ungroup()

df.trace[df.trace$program == 'bitorch',]$program = 'torchtree/bito/BEAGLE'

ggplot(df.trace,
       aes(
         x = size,
         y = max_peak_vmem,
         color = program,
         shape = program
       )) +
  geom_line() +
  geom_point() +
  xlab("Number of sequences") +
  ylab("Peak RSS") +
  scale_color_manual(values = progColors) +
  scale_shape_manual(values = progShapes) +
  theme(
    legend.position = c(0.25, 0.8),
    legend.title = element_blank(),
    legend.key = element_blank(),
    legend.background = element_rect(fill = 'transparent')
  )
```

\newpage

### Runtime against dataset size
```{r include=FALSE}
df.macro.raw = read.csv("results/macro/macro.csv", sep = ',')

df.macro.raw[df.macro.raw$program == 'bitorch', ]$program = 'torchtree/bito/BEAGLE'
max_iter = 5000 # number of iterations in the benchmark
```

```{r macro-time}
df.macro = df.macro.raw %>%
  filter(iters == max_iter | program == 'treeflow') %>%
  group_by(program, size) %>% summarize(time = min(time)) %>% ungroup()

ggplot(df.macro, aes(x = size,
                     y = time,
                     color = program,
                     shape = program)) +
  geom_line() +
  geom_point() +
  xlab("Number of sequences") +
  ylab("Time in seconds") +
  scale_color_manual(values = progColors) +
  scale_shape_manual(values = progShapes) +
  theme(
    legend.position = c(0.7, 0.7),
    legend.title = element_blank(),
    legend.key=element_blank(),
    legend.background = element_rect(fill = 'transparent')
  )
```

\newpage
### Runtime against dataset size without phylojax
```{r macro-time-no-jax}
df.macro = df.macro.raw %>%
  filter(iters == max_iter | program == 'treeflow') %>%
  filter(program != 'phylojax') %>%
  group_by(program, size) %>% summarize(time = min(time)) %>% ungroup()

ggplot(df.macro, aes(x = size,
                     y = time,
                     color = program,
                     shape = program)) +
  geom_line() +
  geom_point() +
  xlab("Number of sequences") +
  ylab("Time in seconds") +
  scale_color_manual(values = progColors) +
  scale_shape_manual(values = progShapes) +
  theme(
    legend.position = c(0.25, 0.7),
    legend.title = element_blank(),
    legend.key=element_blank(),
    legend.background = element_rect(fill = 'transparent')
  )
```

\newpage
```{r}
df.macro = df.macro.raw %>%
  filter(program %in%  c('bitorch', 'torchtree/bito/BEAGLE', 'physher')) %>%
  group_by(program, size) %>% summarize(time = min(time)) %>% ungroup()

ggplot(df.macro, aes(x = size,
                     y = time,
                     color = program)) +
  geom_line() +
  geom_point() +
  xlab("Number of sequences") +
  ylab("Time in seconds") +
  scale_color_manual(values = progColors) +
  theme(
    legend.position = c(0.25, 0.7),
    legend.title = element_blank(),
    legend.key=element_blank(),
    legend.background = element_rect(fill = 'transparent')
  )
```

\newpage

### Speed relative to physher
```{r macro-time-rel}
df.macro.physher = df.macro.raw %>% filter(program == 'physher') %>%
  select(!"elbo") %>% select(!"iters")
df.macro.other = df.macro.raw %>% filter(!program %in% c('physher', 'physherpy')) %>%
  select(!"elbo") %>% select(!"iters")
df.relative = inner_join(df.macro.other, df.macro.physher, c("size", "rep")) %>%
  mutate(ratio = time.x / time.y) %>% group_by(program.x, size) %>%
  summarize(minratio = min(ratio)) %>% ungroup()

ggplot(df.relative, aes(x = size,
                        y = minratio,
                        color = program.x,
                        shape = program.x)) +
  geom_line() +
  geom_point() +
  geom_hline(yintercept = 1) +
  scale_y_continuous(trans = 'log2') +
  xlab("Number of sequences") +
  ylab("Ratio program/physher time") +
  scale_color_manual(values = progColors) +
  scale_shape_manual(values = progShapes) +
  theme(
    legend.position = c(0.75, 0.75),
    legend.title = element_blank(),
    legend.key=element_blank(),
    legend.background = element_rect(fill = 'transparent')
  )
```

\newpage

### Runtime against dataset size up fastest
```{r}
df.macro = df.macro.raw %>%
  filter(iters == max_iter) %>%
  filter(!program %in% c('phylostan', 'torchtree', 'treeflow', 'phylojax'))

ggplot(df.macro, aes(x = size,
                     y = time,
                     color = program)) +
  geom_point() +
  xlab("Number of sequences") +
  ylab("Time in seconds") +
  scale_color_manual(values = progColors)
```


\newpage
```{r include=FALSE}
df.micro = read.csv("results/micro/micro.csv", sep = ',')

df.micro[df.micro$program == 'treeflow' &
           df.micro$function. == 'constant_coalescent',]$function. = 'coalescent'

df.micro$function.[df.micro$function. == 'treelikelihoodJC69'] = 'treelikelihood'

# Use the minimum time out of the x replicates
df.micro = df.micro %>% group_by(program, size, precision, function., mode, JIT) %>%
  summarize(time = min(time)) %>% ungroup()

df.micro[df.micro$precision == 32, ]$program = 'torchtree SP'

df.micro.naive = df.micro %>% filter(function. == 'ratio_transform_jacobian2') %>%
  mutate(program = "physher-naive")  %>%
  mutate(function. = "Transform Jacobian")

df.physher.jacobian = df.micro %>% filter(grepl('ratio_transform_jacobian', function.) & program=='physher')

df.micro = df.micro %>% filter(function. != 'ratio_transform_jacobian2' &
                                 function. != 'ratio_transform2')

df.micro[df.micro$program == 'treeflow', ]$JIT = 'off'

# add first and other evaluation times
df.micro.jax = df.micro %>% filter(program == 'phylojax')
df.micro.jax.1 = df.micro.jax %>% filter(grepl('1', mode)) %>%
  mutate(mode = str_sub(mode, 1, nchar(mode) - 1))
df.micro.jax.rest = df.micro.jax %>% filter(!grepl('1', mode))
df.micro.jax.all = df.micro.jax.1 %>%
  full_join(df.micro.jax.rest,
            by = c("program", "size", "precision", "function.", "JIT", "mode")) %>%
  mutate(time.x = time.x + time.y) %>%
  rename(time = time.x) %>% select(!time.y)

df.micro = rbind(
  df.micro %>% filter(program != 'phylojax'),
  df.micro.jax.all,
  df.micro.jax.1 %>% mutate(program = 'phylojax1')
)

df.micro.jax.analytic = df.micro %>% filter(function. == 'treelikelihoodAnalytic') %>%
  mutate(program="phylojax analytic") %>% mutate(function. ="Treelikelihood JC69")

df.micro$function.[df.micro$function. == "ratio_transform_jacobian"] = "Transform Jacobian"
df.micro$function.[df.micro$function. == "ratio_transform"] = "Transform"
df.micro$function.[df.micro$function. == "coalescent"] = "Coalescent"
df.micro$function.[df.micro$function. == "treelikelihood"] = "Treelikelihood JC69"
df.micro$function.[df.micro$function. == "treelikelihoodGTR"] = "Treelikelihood GTR"

df.micro.jit = df.micro %>% filter(program=="phylojax") %>% filter(function.!="treelikelihoodAnalytic")
df.micro = df.micro %>% filter(JIT == "off")
df.micro[df.micro$program == 'bitorch',]$program = 'torchtree/bito/\nBEAGLE'
```

\newpage

## Micro benchmark
```{r}
plot.micro <- function(d, ylabel) {
  p <- ggplot(d,
              aes(
                x = size,
                y = time,
                color = program,
                shape = program
              )) +
    geom_line() +
    geom_point() +
    facet_wrap(~ function., scales = "free", nrow = 2) +
    xlab("Number of sequences") +
    ylab(ylabel) +
    scale_color_manual(values = progColors) +
    scale_shape_manual(values = progShapes) +
    theme(
      strip.background = element_blank(),
      legend.title = element_blank(),
      axis.text = element_text(size = rel(0.65)),
      legend.key = element_blank(),
      legend.text = element_text(size = rel(0.8)),
      legend.background = element_rect(fill = 'transparent'),
      legend.position = c(0.85, 0.2)
    )

  p
}
```

\newpage

### Log likelihood evalutation
```{r micro-eval}
programs[programs=='torchtree/bito/BEAGLE']='torchtree/bito/\nBEAGLE'
progColors = setNames(progColors, programs)
progShapes = setNames(progShapes, programs)

micro.like = df.micro %>% filter(mode == 'evaluation')

plot.micro(micro.like, "Time in seconds for evaluation")
```

\newpage

### Gradient evaluation
```{r micro-grad}
micro.grad = df.micro %>% filter(mode == 'gradient')

plot.micro(micro.grad, "Time in seconds for gradient")
```

\newpage

### Log likelihood evaluation wihtout phylojax
```{r micro-eval-no-jax}
micro.like = df.micro %>% filter(mode == 'evaluation') %>%
  filter(!grepl('phylojax', program) & program != 'torchtree SP')

plot.micro(micro.like, "Time in seconds for evaluation")
```

\newpage

### Gradient evaluation wihtout phylojax
```{r micro-grad-no-jax}
micro.grad = df.micro %>% filter(mode == 'gradient') %>%
  filter(JIT != 'on' |
           program != 'torchtree' & program != 'torchtree SP') %>%
  filter(!grepl('phylojax', program) & program != 'torchtree SP')

plot.micro(micro.grad, "Time in seconds for gradient")
```

\newpage
### Gradient evaluation phylojax only
```{r micro-grad-jax-only}
df.micro.jit2 = df.micro.jit %>% 
  mutate(program = ifelse(program == 'phylojax',
                           'phylojax AD',
                           program))
micro.jax.grad = rbind(df.micro.jax.analytic, df.micro.jit2) %>%
  filter(mode == 'gradient') %>%
  filter(program %in% c('phylojax AD', "phylojax analytic"))

ggplot(micro.jax.grad,
       aes(
         x = size,
         y = time,
         color = program,
         linetype = JIT
       )) +
  geom_line() +
  geom_point() +
  facet_wrap( ~ function., scales = "free") +
  xlab("Number of sequences") +
  ylab("Time in seconds for gradient") +
  theme(
    strip.background = element_blank(),
    panel.border = element_blank(),
    strip.placement = "outside",
    axis.text = element_text(size = rel(0.5)),
    legend.background = element_rect(fill = 'transparent'),
    legend.box.background = element_rect(fill = 'transparent')
  )
```

\newpage

### Gradient evaluation log scale
```{r micro-grad-log}
dd = df.micro %>% filter(mode == 'gradient') %>% filter(!program %in% c('phylojax1','torchtree SP'))

ggplot(dd, aes(
  x = size,
  y = time,
  color = program,
  shape = program
)) +
  geom_point() +
  scale_x_continuous(
    trans = log10_trans()
  ) +
  scale_y_continuous(
    trans = log10_trans(),
    breaks = trans_breaks("log10", function(x)
      10 ^ x),
    labels = trans_format("log10", math_format(10 ^ .x))
  ) +
  facet_wrap( ~ function., scales = "free", nrow = 2) +
  xlab("Number of sequences (log scale)") +
  ylab("Time in seconds for gradient (log scale)") +
  scale_color_manual(values = progColors) +
  scale_shape_manual(values = progShapes) +
  theme(strip.background = element_blank(),
        #panel.border = element_blank(),
        #strip.placement = "outside",
        legend.title = element_blank(),
        axis.text = element_text(size = rel(0.65)),
legend.key=element_blank(),
    legend.background = element_rect(fill = 'transparent'),
        legend.position = c(0.85, 0.2))
```

\newpage
### Only physher and bitorch
```{r}
dd = df.micro %>% filter(program %in% c('physher', 'torchtree/bito/\nBEAGLE')) %>% filter(grepl('ikeli', function.))

ggplot(dd, aes(x = size,
                    y = time,
                    color = program)) +
  geom_point() +
  scale_x_continuous(
    trans = log10_trans(),
  ) +
  scale_y_continuous(
    trans = log10_trans()
  ) +
  
  xlab("Number of sequences (log scale)") +
  ylab("Time in seconds for gradient (log scale)") +
  facet_wrap(mode ~ function., scales = "free", nrow = 2) +
  scale_color_manual(values = progColors) +
  theme(
    strip.background = element_blank(),
    panel.border = element_blank(),
    strip.placement = "outside",
    legend.title = element_blank(),
    axis.text = element_text(size = rel(0.5)),
    legend.background = element_rect(fill = 'transparent'),
    legend.box.background = element_rect(fill = 'transparent')
  )
```
\newpage

```{r}
fit_log_log <- function(df, m) {
  fits = df %>% filter(mode == m) %>%
    mutate(size = log(size), time = log(time)) %>%
    group_by(program, JIT, function., mode) %>%
    do(model = lm(time ~ size, data = .)) %>% ungroup()
  
  fits$intercept = sapply(fits$model, function(x)
    x$coefficients[[1]])
  fits$slope = sapply(fits$model, function(x)
    x$coefficients[[2]])
  fits$r_squared = sapply(fits$model, function(x)
    summary(x)$r.squared)
  fits$cook = lapply(fits$model, function(x)
    cooks.distance(x))
  fits = fits %>% select(!model)
  fits
}
```

\newpage

### Log likelihood evaluation log scale
```{r}
dd.log = fit_log_log(df.micro, 'evaluation')

kable(dd.log %>% select(-cook) %>% arrange(function., slope, r_squared))
# cook=dd.log%>%filter(program =='treeflow') %>%  filter(mode=='evaluation') %>% filter(function.=='treelikelihood')
# plot(cook$cook[[1]])
```

\newpage

### Gradient evaluation log scale
```{r}
dd.log.grad = fit_log_log(df.micro, 'gradient')

kable(dd.log.grad%>% select(-cook) %>% select(-mode)%>% select(-JIT) %>% arrange(function., slope))
```

\newpage

### Only bitorch and physher
```{r}
dd = df.micro %>%  filter(program %in% c('physher', 'bitorch', 'torchtree/bito/\nBEAGLE')) %>%
  filter(function. != 'coalescent')

ggplot(dd, aes(
  x = size,
  y = time,
  color = program,
  linetype = mode
)) +
  geom_line() +
  geom_point() +
  facet_wrap( ~ function., scales = "free") +
  xlab("Number of sequences") +
  ylab("Time in seconds") +
  scale_color_manual(values = progColors)
```

```{r}
dd = df.micro %>%  filter(program %in% c('physher', 'torchtree', 'jax-analytic'))#%>% filter(function. == 'treelikelihood')%>% filter(mode != 'gradient1')%>% filter(mode != 'evaluation1')

ggplot(dd, aes(
  x = size,
  y = time,
  color = program,
  linetype = mode
)) +
  geom_line() +
  geom_point() +
  facet_wrap( ~ function., scales = "free") +
  xlab("Number of sequences") +
  ylab("Time in seconds") +
  scale_color_manual(values = progColors)
```

\newpage

### Ratio of gradient and likelihood execution time
```{r}
micro.join = df.micro  %>% filter(mode == 'evaluation') %>%
  full_join(df.micro  %>% filter(mode == 'gradient'),
            by = c("program", "size", "precision", "function.", "JIT")) %>%
  mutate(ratio = time.y / time.x) %>% filter(!grepl('phylojax', program))

ggplot(micro.join, aes(
  x = size,
  y = ratio,
  color = program
)) +
  geom_line() +
  geom_point() +
  facet_wrap( ~ function., scales = "free") +
  xlab("Number of sequences") +
  ylab("Ratio gradient/evaluation execution times") +
  scale_color_manual(values = progColors)
```

\newpage
### ELBO convergence
```{r elbos}
df.elbos = read.csv("results/elbos.csv")
df.elbos = df.elbos %>% filter(iteration <= 25000)
ggplot(df.elbos, aes(
  x = iteration/1000,
  y = elbo,
  color = program
)) +
  geom_line() +
  facet_wrap( ~ model) +
  xlab("Number of iterations") +
  ylab("ELBO") +
  theme(
    legend.position = c(0.2, 0.7),
    legend.title = element_blank(),
    legend.box.background = element_rect(fill = 'transparent')
  )
```