---
title: "Filling the gaps: model-based imputation for heatmaps and PCA"
author: "Witold E. Wolski"
date: "`r Sys.Date()`"
output: BiocStyle::html_document
vignette: >
  %\VignetteIndexEntry{7. Filling the gaps: model-based imputation for heatmaps and PCA}
  %\VignetteEncoding{UTF-8}
  %\VignetteEngine{knitr::rmarkdown}
editor_options:
  chunk_output_type: console
---

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

# The problem

Heatmaps, PCA and clustering need a complete data matrix. Label-free data are not complete: a protein is often not observed in some samples, and often not in any sample of one group. While such proteins are frequently below the limit of detection (LOD), complete absence can also reflect true biological on/off expression, which makes them candidates of high biological interest. PCA drops every protein with a missing value, and with it exactly these proteins.

```{r simulate}
library(prolfqua)
library(ggplot2)
sim <- sim_lfq_data_protein_config(Nprot = 100, weight_missing = 0.5, seed = 1234)
lfq <- LFQData$new(sim$data, sim$config)

wide <- lfq$data_wide(as.matrix = TRUE)$data
empty_group <- lfq$data_long() |>
  dplyr::group_by(protein_Id, group_) |>
  dplyr::summarise(empty = all(is.na(abundance)), .groups = "drop") |>
  dplyr::filter(empty)
c(
  proteins = nrow(wide),
  complete = sum(stats::complete.cases(wide)),
  with_empty_group = dplyr::n_distinct(empty_group$protein_Id)
)
```

The proteins with an empty group:

```{r emptyGroups}
lfq$data_long() |>
  dplyr::filter(protein_Id %in% empty_group$protein_Id) |>
  dplyr::select(protein_Id, sampleName, abundance) |>
  tidyr::pivot_wider(names_from = sampleName, values_from = abundance) |>
  knitr::kable(digits = 1)
```

# Fold changes are a separate question

To test for differences you do not need to impute anything yourself: `build_contrast_analysis(method = "lm_impute")` also estimates fold changes for proteins with an empty group, refitting them at the LOD and borrowing their variance from the other proteins. See [Comparing two groups](https://fgcz.github.io/prolfqua/articles/Comparing2Groups.html) and [Comparing modelling methods on one dataset](https://fgcz.github.io/prolfqua/articles/ComparingMethods.html). Here we reuse the models of that fit to fill the data matrix.

# Model-based imputation

`impute_from_model()` fills every missing cell with the prediction of its protein model and keeps the observed values. The summary tells which route each protein took:

- `complete`: no missing values, nothing to fill
- `fitted`: the model fitted on the observed values predicts the missing cells from their group
- `lod_refit`: a protein whose fit failed, typically because a group is empty, is refitted with its missing values set to the LOD and observed values below the LOD raised to it; the empty group gets a value at the LOD

```{r fill}
contrasts <- c(AvsC = "group_A - group_Ctrl", BvsC = "group_B - group_Ctrl")
fa <- build_contrast_analysis(lfq, "~ group_", contrasts, method = "lm_impute")
filled <- impute_from_model(fa$model, lfq)
table(filled$summary$route)
```

`MissingHelpers` estimates the LOD from the data. Groups in which a protein is observed in only one sample, or in the fewest samples, are the closest the data comes to the detection limit; the LOD is the median of their mean intensities. Pass `lod` to `build_contrast_analysis()` to use a LOD of your own.

```{r lod}
MissingHelpers$new(lfq$data_long(), lfq$get_config())$get_lod()
```

The proteins with an empty group after imputation; compare with the table above:

```{r filledCells}
filled$lfqdata$data_long() |>
  dplyr::filter(protein_Id %in% empty_group$protein_Id) |>
  dplyr::select(protein_Id, sampleName, abundance) |>
  tidyr::pivot_wider(names_from = sampleName, values_from = abundance) |>
  knitr::kable(digits = 1)
```

# Heatmap before and after

We cluster the proteins once, on the imputed data, and draw both heatmaps with that row order, so each row shows the same protein before and after.

```{r heatmapOrder}
before <- lfq$data_wide(as.matrix = TRUE)
after <- filled$lfqdata$data_wide(as.matrix = TRUE)
row_order <- rownames(after$data)[hclust(dist(t(scale(t(after$data)))))$order]
```

```{r heatmapBefore, fig.cap = "Before imputation: missing cells are grey."}
ht_before <- plot_heatmap(
  before$data, before$annotation, lfq$factor_keys(), lfq$sample_name(),
  na_fraction = 1, row_order = row_order
)
ht_before
```

```{r heatmapAfter, fig.cap = "After imputation: every cell holds a value."}
ht_after <- plot_heatmap(
  after$data, after$annotation, lfq$factor_keys(), lfq$sample_name(),
  na_fraction = 1, row_order = row_order
)
ht_after
```

# PCA before and after

Before imputation PCA can only use the proteins without missing values; after imputation it uses all proteins.

```{r pca, fig.cap = "PCA before (top) and after (bottom) imputation.", fig.height = 8}
n_complete <- sum(filled$summary$route == "complete")
gridExtra::grid.arrange(
  lfq$get_Plotter()$pca() +
    ggtitle(paste0("Before imputation: ", n_complete, " complete proteins")),
  filled$lfqdata$get_Plotter()$pca() +
    ggtitle(paste0("After imputation: ", nrow(filled$summary), " proteins")),
  ncol = 1
)
```

# Caveats

After imputation the groups separate more cleanly, in the heatmap and in the PCA. Part of that is circular: each filled value is the model's estimate for its group, so imputation adds the group structure of the model to the data. The filled values also carry no noise. Use the filled matrix to look at the data; for tests use the contrasts of `lm_impute`, which account for the imputation.

# Session Info

```{r sessionInfo}
sessionInfo()
```