---
title: "Comparing TE annotations across RepeatModeler parameters"
output: html_document
---
  
```{r setup, include=FALSE}
knitr::opts_chunk$set(
  echo = TRUE, message = FALSE, warning = FALSE,
  cache = TRUE, cache.path = '.cache_knitr/compare_RepModel_params/',
  fig.width = 8, fig.height = 5, dpi = 350,
  fig.path='figure/0-compare_RepModel_params/',
  dev = c('png', 'pdf')
)
```

In this document, we are benchmarking the performance of different combination of RepeatModeler on TE annotations. We focus on the *Astatotilapia calliptera* genome.

## Setup
```{r, cache=FALSE}
library(GenomicRanges)
library(scales)
library(reactable)
library(dplyr)
library(tidyr)
library(ggplot2)
library(tibble)
library(readr)
library(glue)

theme_set(theme_classic())
theme_update(
  panel.background = element_rect(fill = "transparent", colour = NA),
  plot.background = element_rect(fill = "transparent", colour = NA),
  legend.background = element_rect(fill = "transparent", colour = NA)
)

setwd(paste0(Sys.getenv("HOME"), "/code/malawi_transposon/script/R_SCRIPT/"))

source("modules/frequently_used_items.R")
source("modules/get_chrom_sizes.R")
source("modules/compare_annotation_overlap.R") # function to computes the overlap between annotations

# lookup table for various directories pertaining to the RepModel and Masker outputs
datadir_lookup <- readr::read_csv(
  paste0(Sys.getenv("HOME"), "/code/malawi_transposon/metadata/datadir_repmodel_parameter_sweep.csv"), 
  comment = '#'
)

# get chrom sizes for astCal
chrom_sizes <- chrom_sizes_list$astCal
chrom_info <- Seqinfo(chrom_sizes[[1]], chrom_sizes[[2]])
```

## Data import and preprocessing
```{r}
read_astCal_repModel_TE_as_GRanges <- function(idx) {
  # idx: row number of the datadir_lookup data frame
  filepath <- paste0(Sys.getenv("HOME"), "/code/malawi_transposon/storage/cloud/", datadir_lookup$simple[idx])
  
  # trim white spaces from repeat_class (named as repeat_class_broad)
  # rename some of the classes that are incorrect, and drop some irrelevant ones
  tmp_df <- readr::read_delim(filepath, delim='\t', col_types='idciiccccccc') %>% 
    mutate(repeat_class_broad = gsub(" ", "", repeat_class_broad)) %>% 
    filter(!repeat_class_broad %in% c('Simple_repeat', 'Low_complexity', 'tRNA', 'snRNA', 'rRNA', 'Satellite')) %>% 
    mutate(repeat_class_broad = if_else(repeat_class_broad == 'RC', 'Helitron', repeat_class_broad)) %>% 
    mutate(repeat_class_broad = if_else(repeat_class_broad == 'SINE?', 'SINE', repeat_class_broad)) %>% 
    arrange(chr, start, end)
  
  # refactor the TE column to have desired order
  tmp_df$repeat_class_broad <- factor(tmp_df$repeat_class_broad, levels = color_schema_te$te_family)
  
  # keep relevant columns and add new ones
  gr_transposons <- GenomicRanges::GRanges(
    seqnames = tmp_df$chr,
    ranges = IRanges(start = tmp_df$start, end = tmp_df$end),
    strand = tmp_df$complement, name = tmp_df$repeat_family, score = tmp_df$perc_div,
    repeat_color = tmp_df$color, 
    repeat_class = tmp_df$repeat_class_broad, repeat_subclass = tmp_df$repeat_class,
    repmask_id = tmp_df$repmask_id, te_id = tmp_df$id, split = tmp_df$split,
    method = datadir_lookup$method[idx], seqinfo = chrom_info
  )
  return(gr_transposons)
}

# read all the RepModel parameter runs as a list of GRanges
gr_te_list <- lapply(seq(nrow(datadir_lookup)), read_astCal_repModel_TE_as_GRanges)
names(gr_te_list) <- datadir_lookup$method

# combine into a data frame
df_te_list_raw <- bind_rows(lapply(gr_te_list, as.data.frame)) %>% 
  mutate(method = factor(method, levels = datadir_lookup$method))
```

Aggregate data at various levels.
```{r}
# summarise key statistics grouped by individual RepeatModeler families (e.g. ltr-1_family-10)
df_te_list_agg_family <- df_te_list_raw %>% 
  group_by(repmask_id, method, name, repeat_class) %>% 
  summarise(
    width_mean = mean(width),
    width_sd = sd(width),
    width_sum = sum(width),
    score_mean = mean(score),
    score_sd = sd(score),
    count = n() ) %>% 
  ungroup()

# summarise key statistics grouped by TE subclass (e.g. DNA/TcMar2)
df_te_list_agg_subclass <- df_te_list_agg_family %>% 
  group_by(name, repeat_class, method) %>% 
  summarise(count = sum(count), width_sum = sum(width_sum), n_family = n())
```

## Plots 
Count the number of families per subclass, which is also annotated as a number of the edge of the bars.
```{r subclass_count_family, fig.width=16, fig.height=8.5}
# count number of families per subclass
df_te_list_agg_subclass %>% 
  ggplot(aes(x = n_family, y = name, fill = repeat_class)) +
  geom_col(color='black') +
  geom_text(aes(label = n_family), hjust=-0.25, vjust=0.5, size=2.9) +
  facet_wrap(.~method, nrow=1) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  scale_x_continuous('Number of RepeatModeler families by TE subclass', expand = expansion(mult = c(.05, .15))) +
  scale_y_discrete(NULL, limits=rev) +
  theme_bw() +
  theme(legend.position = 'none')
```

Number of fragments per subclass. In the second version, bars are separated into segments to represent individual RepeatModeler families.
```{r subclass_count_fragment, fig.width=16, fig.height=8.5}
# number of TE fragments per subclass 
df_te_list_agg_subclass %>% 
  ggplot(aes(x = count, y = name, fill = repeat_class)) +
  geom_col(color='black') +
  geom_text(aes(label = n_family), hjust=-0.2, vjust=0.5, size=2.9) +
  facet_wrap(.~method, nrow=1) +
  scale_x_continuous(
    'Number of fragments by TE subclass (text: number of RepeatModeler families)', 
    expand = expansion(mult = c(.05, .15))
  ) +
  scale_y_discrete(NULL, limits=rev) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  theme_bw() +
  theme(legend.position = 'none')

# segmented version
df_te_list_agg_family %>% 
  arrange(repeat_class, desc(count)) %>% 
  ggplot(aes(x = count, y = name, fill = repeat_class)) +
  geom_col(color='black', lwd=0.1) +
  geom_text(aes(label = n_family),
            data=df_te_list_agg_subclass, hjust=-0.2, vjust=0.5, size=2.9) +
  facet_wrap(.~method, nrow=1) +
  scale_x_continuous(
    'Number of fragments by TE subclass (text: number of RepeatModeler families)', 
    expand = expansion(mult = c(.05, .15))
  ) +
  scale_y_discrete(NULL, limits=rev) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  theme_bw() +
  theme(legend.position = 'none')
```

Percentage sequence per subclass. In the second version, bars are separated into segments to represent individual RepeatModeler families.
```{r subclass_perc_seq, fig.width=16, fig.height=8.5}
# get percentage sequence per class and subclass
df_te_list_agg_subclass %>% 
  ggplot(aes(x = width_sum /sum(chrom_sizes$length) * 100, y = name, fill = repeat_class )) +
  geom_col(color='black') +
  geom_text(aes(label = n_family), hjust=-0.2, vjust=0.5, size=2.9) +
  facet_wrap(.~method, nrow=1) +
  scale_x_continuous(
    'Percentage sequence of A. calliptera genome annotated as TE subclass (text: number of RepeatModeler families)',
    expand = expansion(mult = c(.05, .15))
  ) +
  scale_y_discrete(NULL, limits=rev) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  theme_bw() +
  theme(legend.position = 'none')

# segmented version
df_te_list_agg_family %>% 
  arrange(repeat_class, desc(width_sum)) %>% 
  ggplot(aes(x = width_sum /sum(chrom_sizes$length) * 100, y = name, fill = repeat_class )) +
  geom_col(color='black', lwd=0.1) +
  geom_text(aes(label = n_family),
            data=df_te_list_agg_subclass, hjust=-0.2, vjust=0.5, size=2.9) +
  facet_wrap(.~method, nrow=1) +
  scale_x_continuous(
    'Percentage sequence of A. calliptera genome annotated as TE subclass (text: number of RepeatModeler families)',
    expand = expansion(mult = c(.05, .15))
  ) +
  scale_y_discrete(NULL, limits=rev) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  theme_bw() +
  theme(legend.position = 'none')
```

Plotting all these for the thesis and presentation slides is probably overkill. Here is some code to create subsections of these plots.
```{r subclass_perc_seq_subset, fig.width=12, fig.height=8.5}
# percentage sequence for a subset
subset_to_plot <- unlist(strsplit("repbase + LTR,no repbase + LTR,repbase + no LTR", ","))
df_te_list_agg_family %>% 
  filter(method %in% subset_to_plot) %>% 
  arrange(repeat_class, desc(width_sum)) %>%  
  ggplot(aes(x = width_sum /sum(chrom_sizes$length) * 100, y = name, fill = repeat_class )) +
  geom_col(color='black', lwd=0.1) +
  geom_text(aes(label = n_family),
            data=df_te_list_agg_subclass %>% filter(method %in% subset_to_plot), 
            hjust=-0.25, vjust=0.5, size=2.9) +
  facet_wrap(.~method, nrow=1) +
  scale_x_continuous(
    'Percentage sequence of A. calliptera genome annotated as TE subclass (text: number of RepeatModeler families)',
    expand = expansion(mult = c(.05, .15))
  ) +
  scale_y_discrete(NULL, limits=rev) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  theme_bw() +
  theme(legend.position = 'none')
```

## Annotation overlap comparisons
Run it through several pairs that are meaningful to compare.
```{r overlap_comparisons, fig.width=3.8, fig.height=3.4}
pairs_to_plot <- list(
  c(2,1), c(1,2),  # contribution of LTR module
  c(3,1), c(1,3),  # contribution of RepBase
  c(4,1), c(1,4),  # both of them
  c(1,5), c(2,5),  # greg as subject (show that we capture most that was found)
  c(4,3), c(3,4)   # what LTR module does (probably don't include)
)
for (pair in pairs_to_plot) {
  compare_annotation_overlap(gr_te_list[pair[1]], gr_te_list[pair[2]])
}
```

## Some class-level statistics
Print out tables.
```{r}
# summarise key statistics grouped by TE class (e.g. DNA)
df_te_list_agg_class <- df_te_list_agg_subclass %>% 
  group_by(repeat_class, method) %>% 
  summarise(count = sum(count), width_sum = sum(width_sum), n_family = sum(n_family)) %>% 
  mutate(perc_sum = round(width_sum / sum(chrom_sizes$length) * 100, 3))

for (variable in c("perc_sum", "n_family", "count")) {
  print(variable)
  df_te_list_agg_class %>% 
    select(method, repeat_class, !!variable) %>% 
    pivot_wider(id_cols = "method", names_from = "repeat_class", values_from = !!variable) %>% 
    mutate(Total = rowSums(across(where(is.numeric)), na.rm=TRUE)) %>% 
    as.data.frame() %>% 
    print()
}
```

Percentage sequence.
```{r class_perc_seq, fig.width=10, fig.height=3.5}
df_te_list_agg_class %>% 
  ggplot(aes(x = perc_sum, y = repeat_class, fill = repeat_class)) +
  geom_col(color='black') +
  geom_text(aes(label = paste0(format(round(perc_sum, 2), n_small=2), "%")), 
            hjust=-0.1, vjust=0.5, size=2.9) +
  facet_wrap(.~method, nrow=1) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  scale_x_continuous('Percentage sequence of A. calliptera genome annotated as TE class', 
                     expand = expansion(mult = c(.05, .35))) +
  scale_y_discrete(NULL, limits=rev) +
  theme_bw() +
  theme(legend.position = 'none')
```

Fragment count.
```{r class_count_fragment, fig.width=10, fig.height=3.5}
df_te_list_agg_class %>% 
  ggplot(aes(x = count/10^3, y = repeat_class, fill = repeat_class)) +
  geom_col(color='black') +
  geom_text(aes(label = format(count/10^3, nsmall=3)), hjust=-0.1, vjust=0.5, size=2.9) +
  facet_wrap(.~method, nrow=1) +
  scale_x_continuous(
    expression('Number of fragments by TE class,' %*% 10^{3}~'annotations'), 
    expand = expansion(mult = c(.05, .4))) +
  scale_y_discrete(NULL, limits=rev) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  theme_bw() +
  theme(legend.position = 'none')
```

Number of families.
```{r class_n_family, fig.width=10, fig.height=3.5}
df_te_list_agg_class %>% 
  ggplot(aes(x = n_family, y = repeat_class, fill = repeat_class)) +
  geom_col(color='black') +
  geom_text(aes(label = n_family), hjust=-0.1, vjust=0.5, size=2.9) +
  facet_wrap(.~method, nrow=1) +
  scale_fill_manual(values = color_schema_te$color, breaks = color_schema_te$te_family) +
  scale_x_continuous('Number of RepeatModeler families by TE class', 
                     expand = expansion(mult = c(.05, .3))) +
  scale_y_discrete(NULL, limits=rev) +
  theme_bw() +
  theme(legend.position = 'none')
```
