---
title: "Mercat optimization analysis"
output: html_document
date: "`r Sys.Date()`"
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
source('../R/mercat.R', echo=F, chdir=T)
source('../R/mercat_exp_metrics.R', echo=F, chdir=T)

library(ggplot2)
library(viridis)
library(ggthemes)
library(reticulate)
library(Seurat)
library(entropy)

cos_sim <- function(x,y)
{
  (x %*% y) / (norm(x, type="2") * norm(y, type="2"))
}
```



## Spectral analysis of angle subsampling matrix

For a better understanding whether the problem is low-rank and hence subsampling is a viable option

Using Paul et al. 2015 hematopoiesis data

```{r spectral analysis of angle subsampling}

exp_data <- anndata::read_h5ad('../data/mousehema.h5ad')
X <- as.matrix(exp_data$X)

## subsample 500 points
X <- X[sample(nrow(X), 500),]

svds <- as.data.frame(do.call(cbind, lapply(1:nrow(X), function(k){
  if (k > 1)
  {
    if (k < nrow(X))
    {
      sub_idcs <- c(1:(k-1), (k+1):nrow(X))
    } else {
       sub_idcs <- 1:(k-1)
    }
  } else {
    sub_idcs <- (k+1):nrow(X)
  }
  X_Xp <- torch_sub(X[sub_idcs,], X[k,])
  angle_subm <- torch_div(torch_matmul(X_Xp, torch_transpose(X_Xp,1,2)),torch_outer(torch_norm(X_Xp, dim=2),torch_norm(X_Xp, dim=2)) )
  svd(as_array(angle_subm))$d
})))
colnames(svds) <- paste0('Sample', 1:nrow(X))
svds$SV <- 1:(nrow(X)-1)
svds_melted <- melt(svds, id.vars=c("SV"), variable.name='Sample', value.name='Singular_value')

gg <- ggplot(svds_melted, aes(x=SV, y=Singular_value, group=Sample)) + theme_minimal() + geom_line(lwd=.5, alpha=.5) + labs(x='singular value index', y='singular value') + scale_y_log10() +
  theme(axis.title = element_text(size=16), axis.text = element_text(size=14))
print(gg)

## Compute effective rank (https://www.eurasip.org/Proceedings/Eusipco/Eusipco2007/Papers/a5p-h05.pdf)
er <- sapply(1:ncol(svds), function(i){ exp(entropy(svds[,i]/norm(as.matrix(svds[,i], ncol=1), type = '1') )) })
gg <- ggplot(data.frame(er=er, lab=rep("", length(er))), aes(x=lab, y=er)) + geom_violin() + geom_boxplot(width=0.2, alpha=.5, outlier.alpha = .5) + theme_minimal() + labs(x="", y="Effective rank") + scale_y_log10() +
  theme(axis.title = element_text(size=16), axis.text = element_text(size=14))
print(gg)
```

## Approximation of gradients through angle subsampling

Compute cosine similarity between full gradient and subsampled gradient

```{r angle sub ana}

## First run full method w/o subsampling
mercat_res <- mercat(X, iter=100, lr=0.001, stepsize = 100)

it_res <- lapply(seq(10, 100, 20), function(it)
{
  angles <- c(2,5,seq(10, nrow(X), 10))
  m_step <- lapply(angles, function(a_i)
  {
    mercat(X, iter=1, lr=0.001, angle_subsample = a_i, Y_prior = mercat_res$Y[[it]], report_progress = F)
  })
  
  err <- sapply(1:(length(m_step)-1), function(i)
  {
    mean(sapply(1:nrow(X), function(j)
    {
      cos_sim(m_step[[i]][j,], m_step[[length(m_step)]][j,])
      #norm(m_step[[i]][j,] - m_step[[length(m_step)]][j,], type="2")
    }))
  })
  data.frame(error=err, angle_sub=angles[-length(angles)])
})

res <- do.call(rbind, it_res)
res$iter <- rep(seq(10, 100, 20), each=nrow(it_res[[1]]))


gg <- ggplot(res, aes(x=iter, y=error, group=angle_sub, color=angle_sub)) +
  geom_point() + geom_line(lwd=.5) + scale_color_viridis() + theme_minimal() + ylim(.99999,1) + labs(x="iteration", y="Cosine similarity full eval") +
  theme(axis.title = element_text(size=16), axis.text = element_text(size=14)) + guides(color=guide_legend(title="#angles\n sampled"))
print(gg)
```

