compare ratxcan with and without loco

#options(error=recover)
#options(error=browser)
options(error=NULL)

## compare observed correlation with null correlation
suppressMessages(devtools::source_gist("a925fea01b365a8c605e")) ## load qqR fn https://gist.github.com/hakyim/a925fea01b365a8c605e
suppressMessages(devtools::source_gist("38431b74c6c0bf90c12f")) ## qqunif https://gist.github.com/hakyim/38431b74c6c0bf90c12f
suppressMessages(devtools::source_gist("115403f16bec0a0e871f3616d552ce9b")) ## source ratxcan functions https://gist.github.com/hakyim/115403f16bec0a0e871f3616d552ce9b 

suppressMessages(library(tidyverse))
suppressMessages(library(glue))
suppressMessages(library(RSQLite))

#suppressMessages(library(expm))
#suppressMessages(library(readxl))
# install.packages("devtools")
# library("devtools")
# install_github("jdstorey/qvalue")
suppressMessages(library(qvalue))
# if (!require("BiocManager", quietly = TRUE))
#     install.packages("BiocManager")
# 
# BiocManager::install("biomaRt")
##suppressMessages(library(biomaRt))
##install.packages("ggrepel")
suppressMessages(library(ggrepel))


recalc=FALSE

WEBDATA="/Users/haekyungim/Library/CloudStorage/Box-Box/LargeFiles/imlab-data/data-Github/web-data"

PLINK="/Users/haekyungim/bin/plink_mac_20231211/plink"
GCTA="/Users/haekyungim/bin/gcta-1.94.2-MacOS-ARM-x86_64/gcta64"

INPUT <- glue("{WEBDATA}/ratxcan-tutorial") ## this has the input data 
OUTPUT <- glue("{WEBDATA}/2023-11-27-ratxcan-tutorial/scratch") ## this has the output data, intermediate results
GENO = glue("{WEBDATA}/2023-11-27-ratxcan-tutorial/data/genotype/")


OUT_TEMPO = glue("{OUTPUT}/testing-2024-08-29")

WEBDATA="/Users/haekyungim/Library/CloudStorage/Box-Box/LargeFiles/imlab-data/data-Github/web-data"

PLINK="/Users/haekyungim/bin/plink_mac_20231211/plink"
GCTA="/Users/haekyungim/bin/gcta-1.94.2-MacOS-ARM-x86_64/gcta64"
  
INPUT=$WEBDATA/ratxcan-tutorial
OUTPUT=$WEBDATA/2023-11-27-ratxcan-tutorial/scratch
GENO=$WEBDATA/2023-11-27-ratxcan-tutorial/data/genotype/

OUT_TEMPO=$OUTPUT/testing-2024-08-29

grm_mat <- read_GRMBin(glue("{OUTPUT}/rat6k_autosome.grm"))

myplot <- function(tempres, post_titulo="",semilla="") {
  # Create a data frame with specific columns
  df <- data.frame(
    p0.01_yes = apply(tempres$pmat_correct, 2, function(x) mean(x < 0.01)),
    p0.01_no = apply(tempres$pmat_raw, 2, function(x) mean(x < 0.01)),
    p0.05_yes = apply(tempres$pmat_correct, 2, function(x) mean(x < 0.05)),
    p0.05_no = apply(tempres$pmat_raw, 2, function(x) mean(x < 0.05)),
    p0.10_yes = apply(tempres$pmat_correct, 2, function(x) mean(x < 0.10)), 
    p0.10_no = apply(tempres$pmat_raw, 2, function(x) mean(x < 0.10))
    # ... [rest of your code for creating df] ...
  )

# Pivot the data frame to long format, specifying the columns to keep
df_long <- pivot_longer(df, cols = starts_with("p"))

df_long <- df_long %>% separate(name,into = c("threshold","corrected"),sep="_") %>% rename(proportion=value)

# Rename the name column to replace p0.xx with p<0.xx
df_long <- df_long %>%
  mutate(threshold = gsub("p0\\.", "p<0.", threshold))

  # Create boxplots with mean
  pp <- ggplot(df_long, aes(x = threshold, y = proportion, fill = corrected)) +
    geom_boxplot(alpha = 0.6) +
    stat_summary(fun = mean, geom = "point", shape = 3, size = 2, stroke = 2, color = "blue",                  #position = position_dodge(width = 0.8)) +
                 position = position_dodge(width = -0.1)) +
    #stat_summary(fun = mean, geom = "crossbar",  size = .5, color = "blue") +
    #stat_summary(fun = mean, geom = "crossbar",  size = .5, color = "darkgray") +
    geom_hline(yintercept = c(0.01, 0.05, 0.10), linetype = "dashed", color = "gray") +
    theme_minimal(base_size = 17) +
    ggtitle(glue("Type I Error Calibration {semilla} {post_titulo}")) +
    xlab("significance") + ylab("false positive rate")

  pp
}

#gene_annotation <- readRDS(glue("{INPUT}/data/expression/gene_annotation.RDS"))
gene_annotation <- readRDS(glue("{WEBDATA}/2023-11-27-ratxcan-tutorial/data/expression/gene_annotation.RDS"))

read_pred_expr = function(filename)
{
  ##usage: Br_pred_expr = read_pred_expr(glue("{OUTPUT}/Br-hki-rat6k__predict.txt"))
  pred_expr <- vroom::vroom(filename) %>% 
  select(-FID) %>%  # Remove the FID column
  mutate(IID = str_split(IID, "_", simplify = TRUE)[, 1])  # Keep the first part of IID
  pred_expr
}

# pred_expr = read_pred_expr(glue("{OUTPUT}/Br-hki-rat6k__predict.txt"))
# ## WHYYYY br-hki-rat6k__predict.txt has no chr 1????
# df1= tibble(gene=names(pred_expr)) %>% mutate(isinac=TRUE) %>% left_join(gene_annotation,by=c("gene"="gene_id")) 
# df1 %>% count(chr)
# cat("{OUTPUT}/Br_rat6k__predict.txt has no chr 1 genes!!! Where did I lose them?\n")

pred_expr = read_pred_expr(glue("{OUTPUT}/AC-filtered__predict.txt"))

## checked that I have chr1 genes
df1= tibble(gene=names(pred_expr)[-1]) %>% mutate(isinac=TRUE) %>% left_join(gene_annotation,by=c("gene"="gene_id")) 
df1 %>% count(chr)

cat("{OUTPUT}/AC-filtered__predict.txt has genes in ",length(unique(df1$chr)), "chr")

## HERE WE USE THE FULL GRM MATRIX AND CALCULATE THE INVERSE OF THE SIGMA MATRIX
## define lmm association function 
lmmGRM = function(pheno, grm_mat, h2, pred_expr,pheno_id_col=1,pheno_value_cols=6:6,out=NULL)
{
  ## input pheno is a data frame with id column pheno_id_col=1 by default
  ## phenotype values are in pheno_value_cols, 6:6 by default (SCORE column location in plink output), it can have more than one phenotype
  ## but h2 has to be the same, this is useful when running simulations with different h2
  ## call lmmXcan(pheno %>% select(IID,SCORE))
  
  ## format pheno to matrix form
  phenomat <- as.matrix(pheno[,pheno_value_cols])
  rownames(phenomat) <- pheno[[pheno_id_col]]
  
  ## turn pred_expr into matrix with rownames =IID, keep only IIDs in ymat
  exp_mat = as.matrix(pred_expr %>% select(-IID))
  rownames(exp_mat) = pred_expr$IID

  ## align pheno and expr matrices
  idlist = intersect(rownames(phenomat), rownames(exp_mat))
  
  nsam = length(idlist)
  
  ## CALCULATE SIGMA
  ID_mat = diag(rep(1,nsam))
  
  #testing_scale_grm = TRUE
  #if(testing_scale_grm) grm_mat = sweep( sweep(grm_mat,2, 1/sqrt(diag(grm_mat)),"*"), 1, 1/sqrt(diag(grm_mat)),"*")    
  
  Sigma = grm_mat[idlist,idlist] * h2 + (1 - h2) * ID_mat
  
  Sig_eigen = eigen(Sigma)
  rownames(Sig_eigen$vectors) = rownames(Sigma)
  
  isighalf = Sig_eigen$vectors %*% diag( 1 / sqrt(  Sig_eigen$values  ) ) %*% t(Sig_eigen$vectors)
  
  ## perform raw association
  cormat_raw = matrix_lm(phenomat[idlist,, drop = FALSE], exp_mat[idlist,])
  pmat_raw = cor2pval(cormat_raw,nsam)
  colnames(pmat_raw) <- gsub("cor_", "pval_", colnames(pmat_raw))
  
  ## perform corrected association
  cormat_correct = matrix_lm(isighalf%*% phenomat[idlist,, drop = FALSE], isighalf %*% exp_mat[idlist,])
  pmat_correct = cor2pval(cormat_correct,nsam)
  colnames(pmat_correct) <- gsub("cor_", "pval_", colnames(pmat_correct))
  
  if(!is.null(out))
  {
    saveRDS(cormat_correct,file = glue("{out}_cormat_correct.RDS"))
    saveRDS(pmat_correct,  file = glue("{out}_pmat_correct.RDS"))
    saveRDS(cormat_raw,    file = glue("{out}_cormat_raw.RDS"))
    saveRDS(pmat_raw,      file = glue("{out}_pmat_raw.RDS"))
  }
  res = list(
    cormat_correct=cormat_correct, 
    pmat_correct=pmat_correct, 
    cormat_raw=cormat_raw, 
    pmat_raw=pmat_raw)
  res
  
}

## read simulated Y
pheno = read_table(glue("{WEBDATA}/2023-11-27-ratxcan-tutorial/scratch/sim/tempo/PRS_output_100-32240.profile"))

## calculate h2 of pheno
## mpheno 4 means fourth column after FID and IID
## $GCTA --grm $OUTPUT/rat6k_autosome --reml --pheno $OUTPUT/sim/tempo/PRS_output_100-32240.profile --mpheno 4 --out $OUTPUT/test-2024-08-29

## create phenotype with h2
h2=0.5
pheno$trait = (pheno$SCORE - mean(pheno$SCORE))/sd(pheno$SCORE) * sqrt(h2) + sqrt(1 - h2)*rnorm(nrow(pheno))

if(recalc)
{res = lmmGRM(pheno, grm_mat, h2, pred_expr,pheno_id_col=1,pheno_value_cols=7:7,out=NULL)
saveRDS(res,glue("{OUT_TEMPO}/ratxcan-ac-res.RDS"))
} else
  res = readRDS(glue("{OUT_TEMPO}/ratxcan-ac-res.RDS"))

$PLINK --bfile $GENO/rat6k_autosome --chr 1 --write-snplist --out $OUT_TEMPO/chr1_snps

$GCTA --bfile $GENO/rat6k_autosome --exclude $OUT_TEMPO/chr1_snps.snplist --make-grm-bin --out $OUT_TEMPO/rat6k_autosome_loco_chr01 --thread-num 8
## that took 

grm_mat_no_chr01 = read_GRMBin(glue("{OUT_TEMPO}/rat6k_autosome_loco_chr01.grm"))


## need to recalculate h2
# $GCTA --reml --grm $OUT_TEMPO/rat6k_autosome  --out $OUT_TEMPO/rat6k_autosome_loco_chr01 --thread-num 8

## save the phenotype with error added
write_tsv(pheno %>% select(IID, FID, trait),file=glue("{OUT_TEMPO}/pheno_PRS_output_100-32240.txt"))

# $GCTA --grm $OUT_TEMPO/rat6k_autosome_loco_chr01 --reml --pheno $OUT_TEMPO/pheno_PRS_output_100-32240.txt --mpheno 1 --out $OUT_TEMPO/loco_chr01 --thread-num 8

h2_loco_chr01 = read_tsv(glue("{OUT_TEMPO}/loco_chr01.hsq")) %>% filter(Source=="V(G)/Vp") %>% pull(Variance)

if(recalc)
{
  res_no_chr01 = lmmGRM(pheno, grm_mat_no_chr01, h2= h2_loco_chr01, pred_expr,pheno_id_col=1,pheno_value_cols=7:7,out=NULL)
  saveRDS(res_no_chr01,file=glue("{OUT_TEMPO}/ratxcan-ac-res-loco-chr01.RDS"))
} else res_no_chr01=readRDS(glue("{OUT_TEMPO}/ratxcan-ac-res-loco-chr01.RDS"))

pmat2df = function(pmat)
{
  df = data.frame(pmat)
  df$gene = rownames(pmat)
  rownames(df) = NULL
  df %>% left_join(gene_annotation %>% select(gene_id,chr,gene_name),by=c("gene"="gene_id") )
}

df_loco_chr01 = pmat2df(res_no_chr01$pmat_correct) 
df = pmat2df(res$pmat_correct) 

par(mfrow=c(2,2))
hist(df_loco_chr01 %>% filter(chr=="1") %>% pull(trait),main="ratxcan GRM no chr1 - genes in chr1")
hist(df_loco_chr01 %>% filter(chr!="1") %>% pull(trait),main="ratxcan GRM no chr1 - genes chr!=1")
hist(df %>% filter(chr=="1") %>% pull(trait),main="ratxcan GRM all - genes in chr1")
hist(df %>% filter(chr!="1") %>% pull(trait),main="ratxcan GRM all - genes chr!=1")

par(mfrow=c(1,1))

par(mfrow=c(2,2))
rango = range(-log10(c(df_loco_chr01$trait, df$trait)))
qqunif(df_loco_chr01 %>% filter(chr=="1") %>% pull(trait),main="ratxcan GRM no chr1 - genes in chr1",BH=FALSE,CI=FALSE,xlim=rango,ylim=rango)
qqunif(df_loco_chr01 %>% filter(chr!="1") %>% pull(trait),main="ratxcan GRM no chr1 - genes chr!=1",BH=FALSE,CI=FALSE,xlim=rango,ylim=rango)
qqunif(df %>% filter(chr=="1") %>% pull(trait),main="ratxcan genes in chr1",BH=FALSE,CI=FALSE,xlim=rango,ylim=rango)
qqunif(df %>% filter(chr!="1") %>% pull(trait),main="ratxcan genes chr!=1",BH=FALSE,CI=FALSE,xlim=rango,ylim=rango)

par(mfrow=c(1,1))


qqunif(df_loco_chr01 %>% filter(chr=="1") %>% pull(trait),main="ratxcan GRM no chr1 - genes in chr1")