reproduceData.qmd

---
title: "Reproducing the Enpact data and results"
author: "Temi"
description: "This notebook contains steps and codes to reproduce, as much as possible, the Enpact paper results"
date: 'Wed Sep 11 2024'
html:
    self-contained: true
    code-background: true
fig-format: png
---

```{r}
knitr::opts_chunk$set(eval = F, echo = TRUE, message = TRUE, warning = TRUE, cache = TRUE)
```

```{r}
setwd('/beagle3/haky/users/temi/projects/TFXcan/notebooks') # to this repository
```

```{r}
library(data.table)
library(glue)
library(dplyr)
library(RSQLite)
library(rtracklayer)
library(GenomicRanges)
library(magrittr)
library(yaml)
# library(fpeek)
library(purrr)
```

```{r}
# tdate <- '2023-08-17' # modified
tdate <- '2024-04-17'

base_dir <- '/project2/haky/temi/projects/TFXcan/baca_cwas'
project_dir <- '/beagle3/haky/users/temi/projects/TFXcan'
data_dir <- glue('{project_dir}/data')
files_dir <- glue('{project_dir}/files')
if(!(dir.exists(files_dir))){
    dir.create(files_dir, recursive = T)
}

if(!(dir.exists(data_dir))){
    dir.create(data_dir, recursive = T)
}
```

# Introduction
In this notebook, I detail steps to re-create some of the data/results in the Enpact paper. I have set this notebook not to run when rendering. So, it will show you the code and markdown text, but not run the code.

# Train DL-based Enpact models from scratch
Use [this](https://github.com/hakyimlab/TFPred-snakemake/tree/main) repository to train the DL-based Enpact models from scratch.

To train all 700+ TF-tissue binding predictors, 

```{bash}
snakemake -s snakefile.smk --configfile minimal/pipeline.734.yaml --profile profiles/simple/ -np > dryrun.out
```

Steps required to train Enpact models are [here](https://github.com/hakyimlab/TFPred-snakemake). There is a minimal example to follow, as well as the config needed to train all 734 TF-tissue binding predictors.

# Train DL-based Enpact models using prepared matrix
Alternatively, you may also train the AR-Prostate Enpact model, as well as the others, using the prepared train and test data here:
-- Train AR_Prostate model
```{r}
train_script <- file.path(project_dir, 'src', 'trainEnpactModel.R')
evaluate_script <- file.path(project_dir, 'src', 'evaluateEnpactModel.R')
train_data <- file.path(project_dir, 'data', 'enpact/training', 'AR_Prostate.train_epigenome.csv.gz')
eval_data <- file.path(project_dir, 'data', 'enpact/training', 'AR_Prostate.test_epigenome.csv.gz')
model_rds <- file.path(project_dir, 'models', 'enpact', 'AR_Prostate_logistic.enpact.rds')
basename_eval_output <- file.path(project_dir, 'data', 'enpact/evaluations', 'AR_Prostate_logistic.enpact.eval_output')
```

```{r}
cmd <- glue('sbatch {project_dir}/src/trainAndEvaluateEnpactModel.sbatch {train_script} {evaluate_script} {train_data} {eval_data} {model_rds} {basename_eval_output}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```

-- Train FOXA1_Prostate model
```{r}
train_script <- file.path(project_dir, 'src', 'trainEnpactModel.R')
evaluate_script <- file.path(project_dir, 'src', 'evaluateEnpactModel.R')
train_data <- file.path(project_dir, 'data', 'enpact/training', 'FOXA1_Prostate.train_epigenome.csv.gz')
eval_data <- file.path(project_dir, 'data', 'enpact/training', 'FOXA1_Prostate.test_epigenome.csv.gz')
model_rds <- file.path(project_dir, 'models', 'enpact', 'FOXA1_Prostate_logistic.enpact.rds')
basename_eval_output <- file.path(project_dir, 'data', 'enpact/evaluations', 'FOXA1_Prostate_logistic.enpact.eval_output')
```

```{r}
cmd <- glue('sbatch {project_dir}/src/trainAndEvaluateEnpactModel.sbatch {train_script} {evaluate_script} {train_data} {eval_data} {model_rds} {basename_eval_output}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```

-- Train GATA2_Prostate model
```{r}
train_script <- file.path(project_dir, 'src', 'trainEnpactModel.R')
evaluate_script <- file.path(project_dir, 'src', 'evaluateEnpactModel.R')
train_data <- file.path(project_dir, 'data', 'enpact/training', 'GATA2_Prostate.train_epigenome.csv.gz')
eval_data <- file.path(project_dir, 'data', 'enpact/training', 'GATA2_Prostate.test_epigenome.csv.gz')
model_rds <- file.path(project_dir, 'models', 'enpact', 'GATA2_Prostate_logistic.enpact.rds')
basename_eval_output <- file.path(project_dir, 'data', 'enpact/evaluations', 'GATA2_Prostate_logistic.enpact.eval_output')
```

```{r}
cmd <- glue('sbatch {project_dir}/src/trainAndEvaluateEnpactModel.sbatch {train_script} {evaluate_script} {train_data} {eval_data} {model_rds} {basename_eval_output}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```

-- Train HOXB13_Prostate model
```{r}
train_script <- file.path(project_dir, 'src', 'trainEnpactModel.R')
evaluate_script <- file.path(project_dir, 'src', 'evaluateEnpactModel.R')
train_data <- file.path(project_dir, 'data', 'enpact/training', 'HOXB13_Prostate.train_epigenome.csv.gz')
eval_data <- file.path(project_dir, 'data', 'enpact/training', 'HOXB13_Prostate.test_epigenome.csv.gz')
model_rds <- file.path(project_dir, 'models', 'enpact', 'HOXB13_Prostate_logistic.enpact.rds')
basename_eval_output <- file.path(project_dir, 'data', 'enpact/evaluations', 'HOXB13_Prostate_logistic.enpact.eval_output')
```

```{r}
cmd <- glue('sbatch {project_dir}/src/trainAndEvaluateEnpactModel.sbatch {train_script} {evaluate_script} {train_data} {eval_data} {model_rds} {basename_eval_output}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```


# Prepare Baca's CWAS models in sqlite and create the necessary databases
The AR weights in the CWAS paper were trained using the Fusion pipeline in mind. Here, we prepared the AR weights such that they are compatible with the PredictDB pipeline.

These weights were shared by our collaborators. So, you can either find them online yourself through the [CWAS paper](https://www.nature.com/articles/s41588-022-01168-y) or you can ask the authors for the weights.
-- the CWAS weights for AR
```{r}
transcription_factor <- 'AR'
ar_zip <- '/project2/haky/Data/baca_cwas/cwas_weights/AR.zip' # use the weights zip file here.
output_dir <- '/beagle3/haky/users/temi/projects/TFXcan/data/baca_cwas_weights'
print(file.exists(ar_zip))
```

-- First unzip the file
```{r}

if(!dir.exists(glue('{output_dir}/{transcription_factor}'))){
    file_names <- unzip(ar_zip, list=T)$Name
    files_to_read <- grep(pattern='^\\bAR\\b.*\\bRDat\\b$', x=file_names, value=T)
    #files_to_read[1:5]

    # unzip the file
    zip::unzip(ar_zip, files=files_to_read, exdir=output_dir)

} 

ar_files <- list.files(glue('{output_dir}/{transcription_factor}'))
ar_files_locus <- sapply(strsplit(x=ar_files, split='\\.'), getElement, 1)
ar_files_locus[1:5]
```

```{r}
length(ar_files_locus)
```

```{r}
hsq.out <- purrr::map(.x=seq_along(ar_files_locus), .f=function(i){
    locus <- ar_files_locus[i]
    #print(file.exists(glue('{output_dir}/{transcription_factor}/{locus}.wgt.RDat')))
    rdt <- new.env(parent = emptyenv())
    load(glue('{output_dir}/{transcription_factor}/{locus}.wgt.RDat'), envir=rdt)
    hsqs <- data.frame(rdt$hsq, rdt$hsq.as, rdt$hsq.qtl, rdt$hsq.pv) %>% 
        dplyr::mutate(locus=locus)
    return(hsqs)
}, .progress=T)

hsq.dt <- do.call('rbind', hsq.out)

# dim(cwas_db) ; cwas_db[1:5, ]
```

```{r}
range(hsq.dt[, 1])
```

-- read weights; Next, read the files `.wgt` files
```{r}
out <- purrr::map(.x=seq_along(ar_files_locus), .f=function(i){
    locus <- ar_files_locus[i]
    #print(file.exists(glue('{output_dir}/{transcription_factor}/{locus}.wgt.RDat')))
    rdt <- new.env(parent = emptyenv())
    load(glue('{output_dir}/{transcription_factor}/{locus}.wgt.RDat'), envir=rdt)
    wgts <- as.data.frame(rdt$wgt.matrix) %>% 
        tibble::rownames_to_column('snp_id') %>% 
        dplyr::mutate(locus=locus)
    
    snp_info <- rdt$snps %>% 
        as.data.frame() %>% 
        dplyr::select(all_of(c('V1', 'V3', 'V2', 'V4', 'V5'))) 

    colnames(snp_info) <- c('chr', 'snp_id', 'position', 'a1', 'a2')

    dt <- base::merge(wgts, snp_info, by='snp_id') %>% 
        dplyr::relocate(all_of(c('locus', 'chr', 'position', 'a1', 'a2')), .after=snp_id)
    return(dt)
}, .progress=T)

cwas_db <- do.call('rbind', out)

dim(cwas_db) ; cwas_db[1:5, ]
```

-- write out the weights
```{r}
data.table::fwrite(cwas_db, file=glue('{files_dir}/{transcription_factor}_baca_cwas_weights.hg19.{tdate}.txt.gz'), col.names=T, row.names=F, quote=F, compress='gzip',sep = '\t')
```
```{r}
cwas_db <- data.table::fread(glue('{files_dir}/{transcription_factor}_baca_cwas_weights.hg19.{tdate}.txt.gz'))

cwas_db$locus |> unique() |> length()
```


-- read and write out the extras

```{r}
out <- purrr::map(.x=seq_along(ar_files_locus), .f=function(i){
    locus <- ar_files_locus[i]
    #print(file.exists(glue('{output_dir}/{transcription_factor}/{locus}.wgt.RDat')))
    rdt <- new.env(parent = emptyenv())
    load(glue('{output_dir}/{transcription_factor}/{locus}.wgt.RDat'), envir=rdt)

    cv_perf <- rbind(rdt$cv.performance['pval', ], rdt$cv.performance['rsq', ]) %>%
        as.data.frame() %>%
        dplyr::mutate(measure = c('pval', 'rsq'), locus) %>%
        dplyr::relocate(locus, measure) %>%
        dplyr::mutate(locus = locus, transcription_factor = transcription_factor, n_snps_in_window = rdt$N.tot, n.snps.in.model = rdt$N.as)

    return(cv_perf)
}, .progress=T)

cvperf_dt <- do.call('rbind', out)
dim(cvperf_dt) ; cvperf_dt[1:5, ]
```

```{r}
data.table::fwrite(cvperf_dt, file=glue('{files_dir}/{transcription_factor}_baca_cwas_extras.hg19.{tdate}.txt.gz'), col.names=T, row.names=F, quote=F, compress='gzip',sep = '\t')
```

```{r}
cvperf_dt <- data.table::fread(glue('{files_dir}/{transcription_factor}_baca_cwas_extras.hg19.{tdate}.txt.gz'))

cvperf_dt$locus |> unique() |> length()
```

```{r}
emptyperf <- setdiff(cvperf_dt$locus |> unique(), cwas_db$locus |> unique())
sapply(emptyperf, function(locus){
    rdt <- new.env(parent = emptyenv())
    load(glue('{output_dir}/{transcription_factor}/{locus}.wgt.RDat'), envir=rdt)
    if(all(is.na(rdt$cv.performance))){
        return(NA)
    } else {
        return(1)
    }
}) 
```

-- map the loci to hg38 from hg19

### Bed mappings i.e. liftover the arbs' coordinates from hg19 to hg38
```{r}
cvperf_dt <- data.table::fread(glue('{files_dir}/{transcription_factor}_baca_cwas_extras.hg19.{tdate}.txt.gz'))
```

```{r}
cvperf_dt$locus |> unique() |> length()
```
```{r}
#write out the bed files
# make sure to remove the empty ones
cvperf_dt %>%
    dplyr::filter(!locus %in% emptyperf) %>%
    tidyr::separate_wider_delim(col = locus, names = c('chr', 'start', 'end'), delim = stringr::regex(':|-')) %>%
    dplyr::mutate(across(c(start, end), as.numeric)) %>%
    dplyr::select(chr, start, end) %>%
    dplyr::distinct() %>%
    dplyr::mutate(id = 1:nrow(.)) %>%
    data.table::fwrite(file=glue('{project_dir}/files/baca_cwas_loci_hg19.bed'), col.names=F, row.names=F, quote=F, sep='\t')
```

```{r}
# lift over these files
# download the liftover command : https://hgdownload.cse.ucsc.edu/admin/exe/linux.x86_64/liftOver
if(!file.exists(file.path(project_dir, 'software', 'liftOver'))){
    download.file('https://hgdownload.cse.ucsc.edu/admin/exe/linux.x86_64/liftOver', destfile = file.path(project_dir, 'software', 'liftOver'))
}
```

```{r}
liftover_sh <-  file.path(project_dir, 'src', 'liftoverSingleBed.sbatch')
liftover_exe <- file.path(project_dir, 'software', 'liftOver')
input_bed <- glue('{project_dir}/files/baca_cwas_loci_hg19.bed')
chain_file <- file.path(project_dir, 'helpers', 'hg19ToHg38.over.chain.gz')
output_bed <- glue('{project_dir}/files/baca_cwas_loci_hg38.bed')
unmapped_bed <- glue('{project_dir}/files/baca_cwas_loci_hg19.unmapped.bed')

file.exists(liftover_sh) ; file.exists(liftover_exe) ; file.exists(chain_file) ; file.exists(input_bed)
```

```{r}
cmd <- glue('sbatch {liftover_sh} {liftover_exe} {chain_file} {input_bed} {output_bed} {unmapped_bed}')
cmd
```

```{r}
system(cmd)
```

- read in, merge and save

```{r}
hg19_bed <- data.table::fread(input_bed, col.names=c('chr', 'start.hg19', 'end.19', 'id'))
hg38_bed <- data.table::fread(output_bed, col.names=c('chr', 'start.hg38', 'end.hg39', 'id'))
bed_mappings <- dplyr::inner_join(hg19_bed, hg38_bed, by=c('chr' = 'chr', 'id' = 'id')) %>%
    dplyr::select(-id) %>% 
    dplyr::rename(chrom = chr)

data.table::fwrite(bed_mappings, file=glue('{project_dir}/files/baca_cwas_arbs_mappings.txt'), col.names=T, row.names=F, quote=F, sep='\t')
```

```{r}
dim(bed_mappings)
```


### SNP mappings i.e. liftover the snps' coordinates from hg19 to hg38

```{r}
# write out the cwas database snps
cwas_db %>%
    dplyr::select(chrom = chr, start = position, locus, rsid=snp_id, a1, a2) %>%
    dplyr::mutate(chrom = paste0('chr', chrom, sep=''), end = start + 1) %>%
    dplyr::relocate(end, .after = start) %>%
    data.table::fwrite(file=glue('{project_dir}/files/baca_cwas_snps_hg19.bed'), col.names=F, row.names=F, quote=F, sep='\t')

```

```{r}
liftover_sh <-  file.path(project_dir, 'src', 'liftoverSingleBed.sbatch')
liftover_exe <- file.path(project_dir, 'software', 'liftOver')
input_bed <- glue('{project_dir}/files/baca_cwas_snps_hg19.bed')
chain_file <- file.path(project_dir, 'helpers', 'hg19ToHg38.over.chain.gz')
output_bed <- glue('{project_dir}/files/baca_cwas_snps_hg38.bed')
unmapped_bed <- glue('{project_dir}/files/baca_cwas_snps_hg19.unmapped.bed')

file.exists(liftover_sh) ; file.exists(liftover_exe) ; file.exists(chain_file) ; file.exists(input_bed)
```

```{r}
cmd <- glue('sbatch {liftover_sh} {liftover_exe} {chain_file} {input_bed} {output_bed} {unmapped_bed}')
cmd
```

```{r}
system(cmd)
```

- read in, merge and save

```{r}
hg19_bed <- data.table::fread(input_bed, col.names=c('chrom', 'start.hg19', 'end.19', 'arbs.hg19', 'rsid', 'a1', 'a2'))
hg38_bed <- data.table::fread(output_bed, col.names=c('chr', 'start.hg38', 'end.hg39', 'arbs.hg19', 'rsid', 'a1', 'a2'))
snp_mappings <- dplyr::inner_join(hg19_bed, hg38_bed, by = c('arbs.hg19' = 'arbs.hg19', 'rsid' = 'rsid', 'a1' = 'a1', 'a2' = 'a2'))
data.table::fwrite(snp_mappings, file=glue('{project_dir}/files/baca_cwas_snp_mappings.txt'), col.names=T, row.names=F, quote=F, sep='\t')
```

-- Now, you can write out the db and save

```{r}
baca_models <- c('lasso', 'lasso.as', 'lasso.plasma', 'top1.as', 'top1.qtl', 'top1')
db_folder <- glue('{project_dir}/models/cwas/db_folder')
if(!dir.exists(db_folder)){dir.create(db_folder)}

db_folder_chr <- glue('{project_dir}/models/cwas/db_folder_chr')
if(!dir.exists(db_folder_chr)){dir.create(db_folder_chr)}
```

```{r}
snp_mappings <- data.table::fread(glue('{project_dir}/files/baca_cwas_snp_mappings.txt'))
weights_dt <- data.table::fread(file.path(project_dir, 'files', 'AR_baca_cwas_weights.hg19.2024-04-17.txt.gz'))
weights_dt <- dplyr::inner_join(weights_dt, snp_mappings, by=c('locus' = 'arbs.hg19', 'snp_id' = 'rsid', 'a1' = 'a1', 'a2' = 'a2')) %>%
    dplyr::select(rsid = snp_id, locus, chrom, position = start.hg38, a1, a2, all_of(baca_models))
head(weights_dt)
```

```{r}
# next, map the locus to hg38
bed_mappings <- data.table::fread(glue('{project_dir}/files/baca_cwas_arbs_mappings.txt')) %>%
    dplyr::mutate(locus.hg19 = paste0(chrom, ':', start.hg19, '-', end.19),
        locus.hg38 = paste0(chrom, '_', start.hg38, '_', end.hg39)) %>%
        dplyr::select(locus.hg19, locus.hg38)

weights_dt <- dplyr::inner_join(weights_dt, bed_mappings, by=c('locus' = 'locus.hg19')) %>%
    dplyr::select(-locus) %>%
    dplyr::rename(locus = locus.hg38) %>%
    dplyr::relocate(locus, .before = rsid)

head(weights_dt)
```

```{r}
# read in and map the extras

baca_extra <- data.table::fread(glue('{files_dir}/{transcription_factor}_baca_cwas_extras.hg19.{tdate}.txt.gz'))
baca_extra <- dplyr::inner_join(baca_extra, bed_mappings, by=c('locus' = 'locus.hg19')) %>%
    dplyr::select(-locus) %>%
    dplyr::rename(locus = locus.hg38) %>%
    dplyr::relocate(locus, .before = measure)
baca_extra$pred.perf.qval <- NA
baca_extra[1:5, ]
```

```{r}
baca_extra %>% dplyr::select(locus, measure, transcription_factor, n_snps_in_window, n.snps.in.model, pred.perf.qval, as.symbol('lasso')) %>% 
    tidyr::pivot_wider(names_from = 'measure', values_from = 'lasso') %>% 
    dplyr::relocate(c(pval, rsq), .after = locus)
```

```{r}
weights_dt %>% 
    dplyr::mutate(varID = paste0(chrom, '_', position, '_', a1, '_', a2, sep=''), gene = locus) %>%
    dplyr::select(gene, rsid, varID, ref_allele=a1, eff_allele=a2, weight=all_of('lasso')) %>%
    dplyr::mutate(varID = gsub("chr", '', varID))
```

```{r}
baca_weights_list <- purrr::map(.x=baca_models, function(each_m){
    model_weights <- weights_dt %>% 
        dplyr::mutate(varID = paste0(chrom, '_', position, '_', a1, '_', a2, sep=''), gene = locus) %>%
        dplyr::select(gene, rsid, varID, ref_allele=a1, eff_allele=a2, weight=all_of(each_m)) %>% 
        dplyr::mutate(varID = gsub("chr", '', varID))

    ## hg38
    edt <- baca_extra %>% 
        dplyr::select(locus, measure, transcription_factor, n_snps_in_window, n.snps.in.model, pred.perf.qval, all_of(each_m)) %>% 
        tidyr::pivot_wider(names_from = 'measure', values_from = each_m) %>% 
        dplyr::relocate(c(pval, rsq), .after = locus) %>%
        dplyr::filter(locus %in% weights_dt$locus) %>%
        dplyr::rename(gene = locus, genename = transcription_factor, pred.perf.R2 = rsq, pred.perf.pval = pval)

    each_db <- DBI::dbConnect(RSQLite::SQLite(), glue('{db_folder}/baca_cwas_{each_m}.hg38.db'))
    dbWriteTable(each_db, "extra", edt, overwrite=T)
    dbWriteTable(each_db, "weights", model_weights, overwrite=T)
    dbDisconnect(each_db)

    return(0)
})

# names(baca_weights_list) <- baca_models
```
-- see an example
```{r}
# '/project2/haky/temi/projects/TFXcan/baca_cwas/db_folder/baca_cwas_lasso.db'
mydb <- dbConnect(SQLite(), glue('{db_folder}/baca_cwas_lasso.hg38.db'))
ex <- dbGetQuery(mydb, 'SELECT * FROM extra')
wt <- dbGetQuery(mydb, 'SELECT * FROM weights')
ex |> head(); wt |> head()

dbDisconnect(mydb)
```

# Predict CWAS scores in Baca individuals using the CWAS models 

Above, we prepared the CWAS models for AR. We can now predict CWAS scores using these models.

```{r}
db_folder <- '/project/haky/users/temi/projects/Enpact/models/cwas/db_folder'
txt_genotypes <- '/project2/haky/Data/baca_cwas/vcfs/hg38/formatted_geno/all_chrs.text_dosages.txt.gz'
txt_samples <- '/project2/haky/Data/baca_cwas/vcfs/hg38/formatted_geno/samples.text_dosages.txt'
exec_file <- '/beagle3/haky/users/temi/software/MetaXcan/software/Predict.py'
output_folder <- "/project/haky/users/temi/projects/Enpact/data/baca_cwas/output"
if(!dir.exists(output_folder)){dir.create(output_folder, recursive = T)}
```

```{r}
# if needed, you should edit the sbatch file: {project_dir}/src/predictCWASscores.sbatch
cmd <- glue('sbatch {project_dir}/src/predictCWASscores.sbatch {db_folder} {txt_genotypes} {txt_samples} {exec_file} {output_folder}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```

- prepare the file in a simple matrix

```{r}
# read in the individual mappings
samples_metadata <- data.table::fread(file.path(project_dir, 'metadata', 'baca_samples_mappings.metadata.txt'))

cwas_scores <- data.table::fread(glue('{output_folder}/top1.hg38/baca_cwas_predict.txt')) %>%
    dplyr::select(-FID) %>%
    tibble::column_to_rownames('IID') %>%
    t() %>% as.data.frame()

cwas_mat <- cwas_scores[,]

colnames(cwas_mat) <- gsub('UW_PDX_172', 'UW_PDX_170_2', colnames(cwas_mat))
colnames(cwas_mat) <- gsub('UW_PDX_173', 'UW_PDX_170_3', colnames(cwas_mat))
colnames(cwas_mat) <- samples_metadata[match(colnames(cwas_mat), samples_metadata$vcf), ]$id
cwas_mat <- cwas_mat %>% tibble::rownames_to_column('locus')
```

```{r}
data.table::fwrite(cwas_mat, file=glue('{project_dir}/files/baca_cwas_scores.hg38.txt.gz'), col.names=T, row.names=F, quote=F, compress='gzip', sep='\t')
```

# Predict Enpact scores using the DL-based Enpact model

Here, I just need the AR_Prostate model tbf

The steps here are involved. 

1. Get epigenomic features for all Baca individuals with Enformer. In this paper, we used [this pipeline]().
2. With the DL-based Enpact model, predict the Enpact scores using [this pipeline](https://github.com/hakyimlab/enpact-predict-snakemake)


## For the 521 1KG EUR individuals at the 18,000 ARBS
```{r}
exec_file <- '/beagle3/haky/users/temi/projects/TFXcan/src/calculateEnpactScores.sbatch'
data_directory <- '/beagle3/haky/users/temi/data/1KG_AR_prostate' 
individuals_list <- '/beagle3/haky/users/temi/projects/TFXcan-snakemake/metadata/EUR_individuals.1KG.txt' 
enpact_weights <- '/beagle3/haky/users/temi/projects/Enpact/data/enpact/weights/ENPACT_734_2025-04-24.compiled_weights.with_intercept.lambda.1se.txt.gz' 
output_file <-'/beagle3/haky/users/temi/projects/TFXcan/data/enpact_scores/AR_Prostate.1KG.CWAS_ARBS.predictions.2025-10-16.rds.gz' 
files_pattern <- '_aggByCollect_AR_Prostate.csv'
md <- 'AR_Prostate'
```

```{r}
cmd <- glue('sbatch {exec_file} {data_directory} {individuals_list} {enpact_weights} {output_file} {files_pattern} {md}')
cmd
```

## For the 131 samples in Baca

```{r}
exec_file <- '/beagle3/haky/users/temi/projects/TFXcan/src/calculateEnpactScores.sbatch'
data_directory <- '/beagle3/haky/users/temi/data/baca_AR_prostate' 
individuals_list <- '/project2/haky/temi/projects/enpact-predict-snakemake/metadata/cwas_individuals.txt' 
enpact_weights <- '/beagle3/haky/users/temi/projects/Enpact/data/enpact/weights/ENPACT_734_2025-04-24.compiled_weights.with_intercept.lambda.1se.txt.gz' 
output_file <- '/beagle3/haky/users/temi/projects/TFXcan/data/enpact_scores/AR_Prostate.BACA.CWAS_ARBS.predictions.2025-10-16.rds.gz' 
files_pattern <- '_aggByMeanCenter_AR_Prostate.csv'
md <- 'AR_Prostate'
```

```{r}
cmd <- glue('sbatch {exec_file} {data_directory} {individuals_list} {enpact_weights} {output_file} {files_pattern} {md}')
cmd
```


# Train a SNP-based Enpact model

1. Get epigenomic features across 521 EUR individuals with Enformer; In this paper, we used [this pipeline]().
     Alternatively, these features can be downloaded from here: [epigenomic features](https://zenodo.org/record/5520737/files/epigenomic_features.tar.gz)
2. With the DL-based Enpact model, predict the Enpact scores using [this pipeline](https://github.com/hakyimlab/enpact-predict-snakemake)
3. Prepare the matrix and data for the SNP-based Enpact model such that they are compatible with predictDB format

## Format 1KG predictions
```{r}
output_directory <- '/scratch/beagle3/temi/enpact_scores/cwas_arbs_1000G'
model <- 'AR_Prostate'
exec_file <- '/beagle3/haky/users/temi/projects/Enpact/src/formatForPredictDB.R' #file.path(project_dir, "src", 'formatForPredictDB.R')
enpact_scores <- file.path(output_directory, glue("1KG.{model}.enpact_scores.txt"))
enpact_annot <- file.path(output_directory, glue("1KG.{model}.enpact_annot.txt"))

# enpact_scores <- file.path(output_directory, "lenpact", glue("1KG_logistic.{model}.enpact_scores.txt"))
# enpact_annot <- file.path(output_directory, "lenpact", glue("1KG_logistic.{model}.enpact_annot.txt"))
rds_file <- '/beagle3/haky/users/temi/projects/TFXcan/data/enpact_scores/AR_Prostate.1KG.CWAS_ARBS.predictions.2025-10-16.rds.gz'

#'/beagle3/haky/users/temi/projects/TFXcan/data/enpact_scores/AR_Prostate.1KG.CWAS_ARBS.predictions.2025-10-16.rds.gz'


#'/beagle3/haky/users/temi/projects/TFXcan/misc/reruns/enpact_predictions/ENPACT_48.1KG.CWAS_ARBS.predictions.2025-04-28.rds.gz'

#endregion"/project2/haky/temi/projects/enpact-predict-snakemake/output/1KG_linear_2025-02-11/1KG_linear_2025-02-11.enpact_scores.array.rds.gz"
# if(!dir.exists(output_folder)){dir.create(output_folder, recursive = T)}
```


```{r}
# if needed, you should edit the sbatch file: {project_dir}/src/predictCWASscores.sbatch
cmd <- glue('sbatch {project_dir}/src/formatForPredictDB.sbatch {exec_file} {rds_file} {enpact_scores} {enpact_annot} {model}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```


## Train the SNP-based Enpact model

```{r}
phenotype <- 'AR_Prostate.EUR.CWAS_ARBS'
outputpath <- "/scratch/beagle3/temi/lEnpact/AR_Prostate.2025-10-23"
annotfile <- file.path(output_directory, glue("1KG.{model}.enpact_annot.txt")) #file.path(data_dir, "lenpact", "1KG.enpact_annot.txt")
enpactfile <- file.path(output_directory, glue("1KG.{model}.enpact_scores.txt")) #file.path(data_dir, "lenpact", "1KG.enpact_scores.txt")
geno_dir <- "/project2/haky/Data/1000G/population_data/EUR/annot_files"
conda_env <- "/beagle3/haky/users/temi/software/conda_envs/predictdb-env"
nextflow_command <- "/beagle3/haky/users/temi/projects/TFXcan-snakemake/workflow/PredictDb-nextflow/main.nf"
```

```{r}
# if needed, you should edit the sbatch file: {project_dir}/src/predictCWASscores.sbatch
cmd <- glue('sbatch {project_dir}/src/generateLinearModels.sbatch {phenotype} {outputpath} {annotfile} {enpactfile} {geno_dir} {conda_env} {nextflow_command}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```

```{bash}
# format the covariances
gzip -dk predict_db_AR_Prostate_filtered.txt.gz
sed -e 's/:/_/g' predict_db_AR_Prostate_filtered.txt > predict_db_AR_Prostate_filtered.formatted.txt

# or
#unpigz -ck predict_db_AR_Prostate.logistic_filtered.txt.gz | sed -e 's/:/_/g' | pigz > predict_db_AR_Prostate.logistic_filtered.formatted.txt.gz

unpigz -ck predict_db_AR_Prostate.EUR.CWAS_ARBS_filtered.txt.gz | sed -e 's/:/_/g' | pigz > predict_db_AR_Prostate.EUR.CWAS_ARBS_filtered.formatted.txt.gz
```


# Predict Enpact scores on Baca individuals using the SNP-based Enpact model
## Format Baca predictions
```{r}
output_directory <- '/scratch/beagle3/temi/enpact_scores/cwas_arbs_baca'
model <- 'AR_Prostate'
exec_file <- '/beagle3/haky/users/temi/projects/Enpact/src/formatForPredictDB.R'
enpact_scores <- file.path(output_directory, glue("Baca.{model}.enpact_scores.txt"))
enpact_annot <- file.path(output_directory, glue("Baca.{model}.enpact_annot.txt"))
rds_file <- '/beagle3/haky/users/temi/projects/TFXcan/data/enpact_scores/AR_Prostate.BACA.CWAS_ARBS.predictions.2025-10-16.rds.gz'

#"/project2/haky/temi/projects/enpact-predict-snakemake/output/baca_2024-11-15/baca_2024-11-15.enpact_scores.array.rds.gz"
# if(!dir.exists(output_folder)){dir.create(output_folder, recursive = T)}
```

```{r}
# if needed, you should edit the sbatch file: {project_dir}/src/predictCWASscores.sbatch
cmd <- glue('sbatch {project_dir}/src/formatForPredictDB.sbatch {exec_file} {rds_file} {enpact_scores} {enpact_annot} {model}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```

## Predict on Baca individuals
```{r}
db_name <- 'AR_Prostate.EUR.CWAS_ARBS'
model_db <- '/scratch/beagle3/temi/lEnpact/AR_Prostate.2025-10-23/AR_Prostate.EUR.CWAS_ARBS/models/filtered_db/predict_db_AR_Prostate.EUR.CWAS_ARBS_filtered.db'
# '/beagle3/haky/users/temi/projects/TFXcan/models/lenpact/AR_Prostate.2025-04-28/AR_Prostate.logistic/models/filtered_db/predict_db_AR_Prostate.logistic_filtered.db'
txt_genotypes <- '/project2/haky/Data/baca_cwas/vcfs/hg38/formatted_geno/all_chrs.text_dosages.txt.gz'
txt_samples <- '/project2/haky/Data/baca_cwas/vcfs/hg38/formatted_geno/samples.text_dosages.txt'
exec_file <- '/beagle3/haky/users/temi/software/MetaXcan/software/Predict.py'
output_folder <- "/beagle3/haky/users/temi/projects/TFXcan/data/lenpact_scores" #"/beagle3/haky/users/temi/projects/TFXcan/data/baca_lenpact/output"
output_basename <- "baca"
if(!dir.exists(output_folder)){dir.create(output_folder, recursive = T)}
```

```{r}
# if needed, you should edit the sbatch file: {project_dir}/src/predictCWASscores.sbatch
cmd <- glue('sbatch {project_dir}/src/predictENPACTscores.sbatch {db_name} {model_db} {txt_genotypes} {txt_samples} {exec_file} {output_folder} {output_basename}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```

# Running TFXcan on prostate cancer data (GWAS summary statistics)
You can download the GWAS summary statistic from the CWAS paper repo [here](https://github.com/scbaca/cwas/blob/master/gwas_data/ProstateCancer_Meta_Schumacher2018.nodup.sumstats.gz) or, better still, from the [GWAS catalog](https://www.ebi.ac.uk/gwas/publications/29892016). This is [the direct link](https://ftp.ebi.ac.uk/pub/databases/gwas/summary_statistics/GCST006001-GCST007000/GCST006085/harmonised/29892016-GCST006085-EFO_0001663.h.tsv.gz)

-- If you have/want to download from the GWAS catalog, this following scripts should help with processing
```{r}
gwas_data <- file.path(project_dir, 'data', 'sumstats', '29892016-GCST006085-EFO_0001663.h.tsv.gz')
dir.create(dirname(gwas_data), recursive = F)
if(!file.exists(gwas_data)){
    download.file('https://ftp.ebi.ac.uk/pub/databases/gwas/summary_statistics/GCST006001-GCST007000/GCST006085/harmonised/29892016-GCST006085-EFO_0001663.h.tsv.gz', destfile = gwas_data)
}
```

-- process the downloaded file. I have processed this file in accordance with the weights data from the SNP-based Enpact model. So, you will need to read in the models and use the data to filter for compatible SNPs.

```{r}
pcr_sumstat <- data.table::fread(gwas_data)
dim(pcr_sumstat); pcr_sumstat[1:5, 1:10]
```
-- read in the lEnpact weights
```{r}
# model_db <- file.path(project_dir, 'models', 'lenpact', 'predict_db_EUR_AR_Prostate_logistic_filtered.db')
enpactdb <- dbConnect(SQLite(), model_db)
weights <- dbGetQuery(enpactdb, 'SELECT * FROM weights')
dbDisconnect(enpactdb)
```

reformat the file and filter for only those with variant Id in the weights data
```{r}
pcgss <- pcr_sumstat %>% 
    dplyr::select(chrom=hm_chrom, variant_id=hm_variant_id, rsid=hm_rsid, pos=hm_pos, A2=hm_other_allele, A1=hm_effect_allele, beta=hm_beta, p_value=p_value, se=standard_error, maf=hm_effect_allele_frequency) %>% 
    dplyr::filter(!(is.na(variant_id) | is.na(rsid) | is.na(pos))) %>%
    dplyr::mutate(zscore=beta/se) %>% 
    dplyr::filter(variant_id %in% weights$varID)

pcgss[1:5, 1:5]; dim(pcgss)
```

```{r}
output_folder <- file.path(project_dir, 'data', 'sumstats')
if(!dir.exists(output_folder)){
    dir.create(output_folder, recursive=T)
}

pcgss %>% split(.$chrom) %>% imap(~data.table::fwrite(.x, glue('{output_folder}/chr{.y}_Schumacher.gwas_ss.txt.gz'), compress='gzip', row.names=F, quote=F, sep = '\t'))
```

-- Now, you can run TFXcan on the prostate cancer data

```{r}
db_name <- 'AR_Prostate.EUR.CWAS_ARBS'
model_db <- "/scratch/beagle3/temi/lEnpact/AR_Prostate.2025-10-23/AR_Prostate.EUR.CWAS_ARBS/models/filtered_db/predict_db_AR_Prostate.EUR.CWAS_ARBS_filtered.db"

#file.path(project_dir, 'models/lenpact/AR_Prostate.2025-04-28/AR_Prostate.logistic/models/filtered_db/predict_db_AR_Prostate.logistic_filtered.db')
covariances <- "/scratch/beagle3/temi/lEnpact/AR_Prostate.2025-10-23/AR_Prostate.EUR.CWAS_ARBS/models/filtered_db/predict_db_AR_Prostate.EUR.CWAS_ARBS_filtered.formatted.txt.gz"
#file.path(project_dir, 'models/lenpact/AR_Prostate.2025-04-28/AR_Prostate.logistic/models/filtered_db/predict_db_AR_Prostate.logistic_filtered.formatted.txt.gz')
gwas_folder <- "/beagle3/haky/users/temi/projects/Enpact/data/sumstats" #file.path(project_dir, 'data/sumstats')
gwas_file_pattern <- ".*_Schumacher.gwas_ss.txt.gz"
exec_file <- '/beagle3/haky/users/temi/software/MetaXcan/software/SPrediXcan.py'
output_file <- file.path(project_dir, 'data', 'tfxcan', 'AR_Prostate.EUR.CWAS_ARBS.TFXcan.prostate_cancer_risk.csv')
if(!dir.exists(dirname(output_file))){dir.create(dirname(output_file), recursive = T)}
```

```{r}
cmd <- glue('sbatch {project_dir}/src/sTFXcan.sbatch {db_name} {model_db} {covariances} {gwas_folder} {gwas_file_pattern} {exec_file} {output_file}')
cmd
```

```{r}
system(cmd) ; system('squeue -u temi')
```

# Running TWAS on prostate cancer data

This is very similar to using the PrediXcan pipeline. You will need to download the TWAS weights from the [PrediXcan website](https://predictdb.org/). You can use the [GTEx v8 weights](https://predictdb.org/download/weights/GTEx_V8_HapMap-2017-11-29.tar.gz) or the [GTEx v7 weights](https://predictdb.org/download/weights/GTEx_V7_HapMap-2017-11-29.tar.gz).

# Consensus matrix factorization 
```{r warning=F, message=F}
devtools::source_gist('https://gist.github.com/TemiPete/250d9922b9516691f83bb1fd999a3ccc')
devtools::source_gist('https://gist.github.com/hakyim/38431b74c6c0bf90c12f')
devtools::source_gist('https://gist.github.com/hakyim/5d2251ea1a86009499e4ffdf47fe2735')
devtools::source_gist('https://gist.github.com/TemiPete/e7214eb308c4116e2f10ee96401be0c2')
devtools::source_gist('https://gist.github.com/TemiPete/b309a46e25ecec93127fb6756e68fb14')

# analyze tenerife
devtools::source_gist("https://gist.github.com/TemiPete/d303781c7ddd9b6d8b4d0163804e80c2")
```

### Step 1: First prepare the z-ratio matrices

We want to remove some tissues 
```{r}
# these weird tissues should be removed 
weird_tissues <- c("endometrioidadenocarcinoma", "HCT116", "HeLacontaminant", "Headandneck", "LNCaPcells", 'PeritonealEffusion')
#performance_file <- '/beagle3/haky/users/temi/projects/TFPred-snakemake/data/ENPACT_734_2024-07-26/statistics/ENPACT_734_2024-07-26.compiled_stats.txt'
performance_file <- '/beagle3/haky/users/temi/projects/Enpact/data/enpact/weights/ENPACT_734_2025-04-24.compiled_stats.with_auprc.tsv'

#'/beagle3/haky/users/temi/projects/Enpact/data/enpact/weights/ENPACT_734_2025-04-24.compiled_stats.txt'
performance <- data.table::fread(performance_file) %>%
    dplyr::filter(type == 'test') %>%
    dplyr::filter((!is.na(auroc) & auroc >= 0.7) & !context %in% weird_tissues)
dim(performance)
```

You need the summary statistics

```{r}
# GWAS information
ss <- data.table::fread("/beagle3/haky/users/temi/projects/TFXcan-snakemake/data/prostate_cancer_risk_2024-09-30/collection/prostate_cancer_risk.filteredGWAS.topSNPs.txt.gz") %>%
    dplyr::mutate(loci = paste('chr', chr, sep='') %>% paste(., pos, sep = ':')) 
    
    # %>%
    # dplyr::mutate(locus = paste('chr', chr, sep = '') %>% paste(., pos, pos + 1, sep='_'))
```

You need the TFXcan results
```{r}
dt <- data.table::fread("/beagle3/haky/users/temi/projects/Enpact/misc/reruns/prca_tfxcan/prostate_cancer_risk.enpactScores.2025-10-21.complete.spredixcan.tsv") %>%
    dplyr::filter(!is.na(pvalue), !is.na(zscore)) %>%
    tidyr::separate_wider_delim(tfbs, delim = "_", names = c("tf", 'tissue', 'chrom', 'start', 'end')) %>%
    dplyr::mutate(locus = paste(chrom, start, end, sep = "_"), tf_tissue = paste(tf, tissue, sep = "_")) %>%
    dplyr::filter(tf_tissue %in% performance$model) %>%
    dplyr::select(locus, tf, tissue, tf_tissue, chrom, start, end, zscore, pvalue) %>%
    dplyr::filter(!grepl("HCT116|HeLacontaminant|PeritonealEffusion|LNCaPcells|endometrioidadenocarcinoma|Headandneck", tissue)) %>%
    dplyr::mutate(snp = paste(chrom, start, sep = ':'))
```

```{r}
process_tfxcan_results <- function(summary_results, gwas_summary_stats, combine_at, retain_at, aggregate_on = c('tf', 'tissue', 'tf_tissue'), how = 'weights'){
    # analysis_factor <- switch(analysis_level, 
    #     'tf' = 'tf', 
    #     'tissue' = 'tissue', 
    #     'tf_tissue' = 'tf_tissue'
    # )
    aggregate_on <- switch(aggregate_on, 
        'tf' = 'tf', 
        'tissue' = 'tissue', 
        'tf_tissue' = 'tf_tissue'
    )

    # process the sumstats to make it easier to work with
    res_sumstats <- gwas_summary_stats %>%
        dplyr::rename(locus = loci, zscore = zscore, pvalue = pval) %>%
        dplyr::filter(nchar(a0) == 1) %>%
        dplyr::mutate(bind_level = 'GWAS') %>%
        dplyr::select(locus, bind_level, zscore, pvalue) %>% 
        dplyr::filter(locus %in% summary_results$snp)

    print(head(res_sumstats))

    if(aggregate_on %in% c('tf', 'tissue')){

        dt_gwas <- res_sumstats %>% 
            dplyr::select(locus, bind_level, zscore, pvalue)
        dt_tfxcan <- summary_results %>% 
            dplyr::select(locus, all_of(c(combine_at, aggregate_on)), zscore, pvalue) 
        mtx <- dt_tfxcan %>%
            tidyr::separate(locus, c('chrom', 'start'), sep = ':', remove = FALSE) %>%
            dplyr::mutate(start = as.numeric(start), end = as.numeric(start + 1)) 
            
        if(how == 'weights'){
            mtx <- mtx %>%
                dplyr::group_split(!!rlang::sym(retain_at), chrom, start, end) %>%
                purrr::map_dfr(merge_zscores, level = retain_at, how = 'weights') %>%
                do.call(cbind, .) %>% as.data.table() %>% 
                dplyr::mutate(locus = paste0(chrom, ":", start), zscore = as.numeric(zscore)) %>% 
                dplyr::arrange(desc(zscore)) %>%
                dplyr::select(locus, all_of(c('level', 'zscore'))) %>%
                dplyr::rename(bind_level = 'level') 
        } else if (how == 'acat'){
            # sum and convert to zscores
            # dplyr::group_by(!!rlang::sym(retain_at))
            mtx <- mtx %>%
                dplyr::group_by(!!rlang::sym(aggregate_on), chrom, start, end) %>%
                dplyr::summarize(nv = acat(pvalue) + 2.2e-16) %>%
                dplyr::mutate(locus = paste0(chrom, ":", start), zscore = qnorm(nv, lower.tail = FALSE)) %>%
                dplyr::ungroup() %>%
                dplyr::group_by(!!rlang::sym(aggregate_on)) %>% 
                dplyr::arrange(desc(zscore)) %>%
                dplyr::select(locus, all_of(c(aggregate_on, 'zscore'))) %>%
                dplyr::rename(bind_level = aggregate_on)
            
        } 
    } else if (aggregate_on %in% c('tf_tissue')) {
        # gg <- res_sumstats %>% 
        #     dplyr::select(locus, bind_level, zscore, pvalue)
        mtx <- summary_results %>% 
            dplyr::select(locus = snp, bind_level = !!rlang::sym(aggregate_on), zscore, pvalue) 
        dt_gwas <- res_sumstats %>% 
            dplyr::select(locus, bind_level, zscore, pvalue)

        # mtx <- dplyr::bind_rows(gg, xx)
    }

    print(head(mtx))
    print("===")
    print(head(dt_gwas))

    mat_tfxcan <- dplyr::bind_rows(mtx, dt_gwas) %>%
        dplyr::select(locus, bind_level, zscore) %>%
        tidyr::pivot_wider(id_cols = 'locus', values_from = zscore, names_from = bind_level) %>% 
        tibble::column_to_rownames('locus') %>%
        as.matrix() %>%
        t()
    print(mat_tfxcan[1:3, 1:3])
    #mat_tfxcan <- mat_tfxcan[complete.cases(mat_tfxcan), ]
    return(mat_tfxcan)
}
```

```{r}
# ss %>%
#         dplyr::rename(locus = loci, zscore = zscore, pvalue = pval) %>%
#         dplyr::filter(nchar(a0) == 1) %>%
#         dplyr::mutate(bind_level = 'GWAS') %>%
#         dplyr::select(locus, bind_level, zscore, pvalue) %>% head()
#         dplyr::filter(locus %in% summary_results$locus) %>% head()
```

```{r}
# collect the Z-scores into a matrix of loci x TF/tissue pairs
tf_tissue_processed <- process_tfxcan_results(dt, ss, combine_at = 'tf', retain_at = 'tissue', aggregate_on = 'tf_tissue')

# collect the Z-scores into two matrices for TFXcan and GWAS results
data_list <- list()
data_list[['zratios']] <- collect_matrices(tf_tissue_processed)

# divide or get eh Z-ratios
zratio_matdev <- (data_list[["zratios"]][["tfxcan"]])^2/(data_list[["zratios"]][["gwas"]][,,drop = TRUE])^2

# there are missing values because I removed associations with pvalue at fdr < 0.05

dim(zratio_matdev); zratio_matdev[1:5, 1:5]
```

```{r}
saveRDS(data_list, '/beagle3/haky/users/temi/projects/TFXcan/data/tenerife/PrCa.tfxcan.zscores.matrices.rds')
saveRDS(zratio_matdev, '/beagle3/haky/users/temi/projects/TFXcan/data/tenerife/PrCa.tfxcan.zratios.matrix.rds')
```

```{r}
sum(complete.cases(zratio_matdev))
```

### Step 2: Split the z-ratio matrix into random subsets
```{r}
X <- readRDS('/beagle3/haky/users/temi/projects/TFXcan/data/tenerife/PrCa.tfxcan.zratios.matrix.rds') # loci by TF_tissues
```

```{r}
# subsets directory
consensus_subsets_dir <- file.path(project_dir, 'data', 'tenerife', 'repeat_flash', 'prca_risk')
dir.create(consensus_subsets_dir, recursive = TRUE)
```

```{r}
# create random subsets; here I use 1000 random subsets, each with 80% of the loci
set.seed(2025)
random_loci <- purrr::map(1:1000, function(k){
    jj <- sample(nrow(X), round(nrow(X) * 0.8))
    return(jj)
})

# confirm length
lapply(random_loci, length) %>% unlist() %>% unique() 

# split into 100 batches
batches_loci <- purrr::map(1:100, function(k){
    jj <- random_loci[seq(k, 1000, by = 100) ] #|> unlist() #seq(k, 1, 1) 
    return(jj)
})

names(batches_loci) <- paste0('batch', 1:100)

lapply(batches_loci, length) |> unlist() %>% sum()

# save these as a list
lapply(names(batches_loci), function(ni) {
    saveRDS(batches_loci[ni], glue::glue('{consensus_subsets_dir}/prca_risk.random_subsets.{ni}.rds'))
})

# save this as a list
# saveRDS(random_loci, '/beagle3/haky/users/temi/projects/Enpact/data/tenerife/prca_risk.random_subsets_1000.rds')
```

```{r}
write(names(batches_loci), file = glue::glue('{consensus_subsets_dir}/prca_risk.random_subsets.txt'))
```

### Step 3: Run Flashier on the random subsets
```{r}
script <- file.path(project_dir, 'src/repeat_flash.sbatch')
datafile <- file.path(project_dir, 'data/tenerife/PrCa.tfxcan.zratios.matrix.rds')
output_basename <- file.path('/beagle3/haky/users/temi/projects/TFXcan/data/tenerife/flash_results', 'prca_risk') # not a directory but a prefix for the output files
priorL <- 'ebnm_point_exponential'
priorF <- 'ebnm_point_exponential'
greedy_Kmax <- 40L
batch_list <- glue::glue('{consensus_subsets_dir}/prca_risk.random_subsets.txt')
splits_directory <- consensus_subsets_dir
merged_output_basename <- '/beagle3/haky/users/temi/projects/TFXcan/data/tenerife/prca_risk'
```

```{r}
glue::glue("sbatch {script} {datafile} {output_basename} {priorL} {priorF} {greedy_Kmax} {batch_list} {splits_directory} {merged_output_basename}")
```

### Step 4: Run consensus factor analysis on the outputs

```{r}
consensus_script <- file.path(project_dir, 'src/cluster_programs.sbatch')
datafile <- glue::glue("{merged_output_basename}.ebnm_point_exponential-ebnm_point_exponential.1000Iters.rds.gz")

#file.path(splits_directory, 'prca_risk.retrain_repeats.ebnm_point_exponential-ebnm_point_exponential.1000Iters.rds.gz')
outputbasename <- file.path(project_dir, 'data/tenerife/consensus/prca_risk')
normalization <- 'L2'
```

```{r}
glue::glue("sbatch {consensus_script} {datafile} {outputbasename} {normalization}")
```

```{r}
knitr::knit_exit()
```



# You can also run some codes individually

## develop snp-predictors

For this one, I have a snakemake pipeline ready to go. 

## run summary TFXcan

```{r}
mms <- list.files("/scratch/beagle3/temi/lEnpact/enpact_models_linearization_2025-10-16")
mms_models <- glue::glue("/scratch/beagle3/temi/lEnpact/enpact_models_linearization_2025-10-16/{mms}/models/filtered_db/predict_db_{mms}_filtered.db")
mms_cov <- glue::glue("/scratch/beagle3/temi/lEnpact/enpact_models_linearization_2025-10-16/{mms}/models/filtered_db/Covariances.varID.txt.gz")
```

```{r}
mms_metadata <- data.table(tf_tissue = mms, model = mms_models, covariances = mms_cov)

mms_metadata %>% data.table::fwrite("/beagle3/haky/users/temi/projects/TFXcan/files/enpact_linearized_models.tsv", sep = '\t', quote = F, row.names = F, col.names = T)
```

```{r}
phenotype = "prostate_cancer_risk"
models_list = "/beagle3/haky/users/temi/projects/TFXcan/files/enpact_linearized_models.tsv"
output_directory = '/scratch/beagle3/temi/summary_tfxcan/prostate_cancer_risk'
gwas_folder = "/beagle3/haky/users/temi/projects/TFXcan-snakemake/data/prostate_cancer_risk_2024-09-30/processed_sumstats/prostate_cancer_risk"
gwas_pattern = '.*.sumstats.txt.gz'
exec_file = "/beagle3/haky/users/temi/software/MetaXcan/software/SPrediXcan.py"
conda_environment = "/beagle3/haky/users/temi/software/conda_envs/imlabtools"

glue::glue("sbatch /beagle3/haky/users/temi/projects/TFXcan/src/run_summary_tfxcan.sbatch {phenotype} {models_list} {output_directory} {gwas_folder} {gwas_pattern} {exec_file} {conda_environment}")
```

```{r}
# # I have to re-run mammary gland
# mms_metadata %>% dplyr::filter(tf_tissue == 'AR_MammaryGland') %>% data.table::fwrite("/beagle3/haky/users/temi/projects/TFXcan/files/enpact_linearized_models.AR_MammaryGland.tsv", sep = '\t', quote = F, row.names = F, col.names = T)
```

```{r}
# phenotype = "prostate_cancer_risk"
# models_list = "/beagle3/haky/users/temi/projects/TFXcan/files/enpact_linearized_models.AR_MammaryGland.tsv"
# output_directory = '/scratch/beagle3/temi/summary_tfxcan/prostate_cancer_risk'
# gwas_folder = "/beagle3/haky/users/temi/projects/TFXcan-snakemake/data/prostate_cancer_risk_2024-09-30/processed_sumstats/prostate_cancer_risk"
# gwas_pattern = '.*.sumstats.txt.gz'
# exec_file = "/beagle3/haky/users/temi/software/MetaXcan/software/SPrediXcan.py"
# conda_environment = "/beagle3/haky/users/temi/software/conda_envs/imlabtools"

# glue::glue("sbatch /beagle3/haky/users/temi/projects/TFXcan/src/run_summary_tfxcan.sbatch {phenotype} {models_list} {output_directory} {gwas_folder} {gwas_pattern} {exec_file} {conda_environment}")
```

## collect all summary TFXcan into one

```{r}
fpattern <- glue::glue("/scratch/beagle3/temi/summary_tfxcan/prostate_cancer_risk/\\*.prostate_cancer_risk.summaryTFXcan.csv")
output_file <- "/beagle3/haky/users/temi/projects/Enpact/misc/reruns/prca_tfxcan/prostate_cancer_risk.enpactScores.2025-10-21.complete.spredixcan.tsv"
phenotype <- "prostate_cancer_risk"

glue::glue("sbatch /beagle3/haky/users/temi/projects/TFXcan/src/gather_sTFXcan.sbatch {fpattern} {phenotype} {output_file}")
```





# Collect weights

```{r}
date_2025 <- '2025-04-24'
directory_2025 <- '/beagle3/haky/users/temi/projects/TFPred-snakemake/data/ENPACT_734_2025-04-24/models'
models_2025 <- list.files(directory_2025)
metadata_2025 <- data.table(tf_tissue = models_2025, path = glue::glue("{directory_2025}/{models_2025}/{models_2025}_{date_2025}.logistic.rds")) %>%
    tidyr::separate_wider_delim(col = tf_tissue, delim = '_', names = c('tf', 'tissue'), cols_remove = FALSE)
```

```{r}
date_2024 <- '2024-07-26'
directory_2024 <- '/beagle3/haky/users/temi/projects/TFPred-snakemake/data/ENPACT_734_2024-07-26/models'
models_2024 <- list.files(directory_2024)
metadata_2024 <- data.table(tf_tissue = models_2024, path = glue::glue("{directory_2024}/{models_2024}/{models_2024}_{date_2024}.logistic.rds")) %>%
    tidyr::separate_wider_delim(col = tf_tissue, delim = '_', names = c('tf', 'tissue'), cols_remove = FALSE) %>%
    dplyr::filter(!tf_tissue %in% metadata_2025$tf_tissue)
```

```{r}
models_metadata <- dplyr::bind_rows(metadata_2025, metadata_2024) %>%
    dplyr::rename(model = tf_tissue) %>%
    dplyr::filter(file.size(path) != 0)
```

```{r}
data.table::fwrite(models_metadata, '/beagle3/haky/users/temi/projects/TFXcan/files/enpact_models.metadata.tsv', sep = '\t', quote = F, col.names = TRUE, row.names = FALSE)
```

```{r}
"sbatch /beagle3/haky/users/temi/projects/TFXcan/src/collect_enpact_weights.sbatch /beagle3/haky/users/temi/projects/TFXcan/files/enpact_models.metadata.tsv /beagle3/haky/users/temi/projects/Enpact/data/enpact/weights/ENPACT_734_2025-04-24.compiled_weights.with_intercept"
```

### Are they the same, or at least, really well-correlated?
```{r}
old_stfxcan <- data.table::fread("/beagle3/haky/users/temi/projects/TFXcan-snakemake/data/prostate_cancer_risk_2024-09-30/important/prostate_cancer_risk.enpactScores.2024-09-30.spredixcan.complete.txt") %>% dplyr::rename(tfbs = gene)
# data.table::fread("/beagle3/haky/users/temi/projects/Enpact/misc/reruns/prca_tfxcan/prostate_cancer_risk.enpactScores.2025-04-25.complete.spredixcan.txt")
new_stfxcan <- data.table::fread("/beagle3/haky/users/temi/projects/Enpact/misc/reruns/prca_tfxcan/prostate_cancer_risk.enpactScores.2025-10-21.complete.spredixcan.tsv")
common_loci <- intersect(old_stfxcan$tfbs, new_stfxcan$tfbs)
dt_old <- subset(old_stfxcan, tfbs %in% common_loci) %>% dplyr::filter(tfbs %in% common_loci) %>% dplyr::arrange(match(tfbs, common_loci))
dt_new <- subset(new_stfxcan, tfbs %in% common_loci) %>% dplyr::filter(tfbs %in% common_loci) %>% dplyr::arrange(match(tfbs, common_loci))
```

```{r}
plot(dt_old$zscore, dt_new$zscore)
```

```{r}
cor.test(dt_old$zscore, dt_new$zscore)
```

### Create hdf5 database

#### For prca loci

```{r}
folder <- '/beagle3/haky/users/temi/projects/TFXcan-snakemake/data/prostate_cancer_risk_2024-09-30/aggregated_predictions/prostate_cancer_risk'
individual_enformer_predictions <- list.files(folder, pattern = "*.aggByCollect_prostate_cancer_risk.csv.gz")
```


```{r}
individuals <- gsub('_aggByCollect_prostate_cancer_risk.csv.gz', '', individual_enformer_predictions)
individual_metadata <- data.table(individual = individuals, path = glue::glue("{folder}/{individual_enformer_predictions}"))

individual_metadata %>%
    data.table::fwrite('/beagle3/haky/users/temi/projects/TFXcan/files/individual_metadata.eur_1000G.metadata.tsv', sep = '\t', quote = F, col.names = TRUE, row.names = FALSE)
```


#### For Baca ARBS - 1000G predictions

```{r}
folder <- '/beagle3/haky/users/temi/data/1KG_AR_prostate'
individual_enformer_predictions <- list.files(folder, pattern = "*_aggByCollect_AR_Prostate.csv")
```

```{r}
individuals <- gsub('_aggByCollect_AR_Prostate.csv', '', individual_enformer_predictions)
individual_metadata <- data.table(individual = individuals, path = glue::glue("{folder}/{individual_enformer_predictions}"))

individual_metadata %>%
    data.table::fwrite('/beagle3/haky/users/temi/projects/TFXcan/files/individual_metadata.1000G.baca_arbs.metadata.tsv', sep = '\t', quote = F, col.names = TRUE, row.names = FALSE)
```

```{r}
glue::glue("sbatch /beagle3/haky/users/temi/projects/TFXcan/src/create_h5database_from_csvs.sbatch /beagle3/haky/users/temi/projects/TFXcan/files/individual_metadata.1000G.baca_arbs.metadata.tsv /scratch/beagle3/temi/epigenomes/1KG_BACA_ARBS/BACA_ARBS.1KG.2025-10-21")
```


#### For Baca ARBS - Baca predictions

```{r}
folder <- '/beagle3/haky/users/temi/data/baca_AR_prostate'
individual_enformer_predictions <- list.files(folder, pattern = "*_aggByMeanCenter_AR_Prostate.csv")
```

```{r}
individuals <- gsub('_aggByMeanCenter_AR_Prostate.csv', '', individual_enformer_predictions)
individual_metadata <- data.table(individual = individuals, path = glue::glue("{folder}/{individual_enformer_predictions}"))

individual_metadata %>%
    data.table::fwrite('/beagle3/haky/users/temi/projects/TFXcan/files/individual_metadata.BACA.baca_arbs.metadata.tsv', sep = '\t', quote = F, col.names = TRUE, row.names = FALSE)
```

```{r}
glue::glue("sbatch /beagle3/haky/users/temi/projects/TFXcan/src/create_h5database_from_csvs.sbatch /beagle3/haky/users/temi/projects/TFXcan/files/individual_metadata.BACA.baca_arbs.metadata.tsv /scratch/beagle3/temi/epigenomes/baca_BACA_ARBS/BACA_ARBS.baca.2025-10-21")
```

## Predict with Enpact weights
--matrix {input.matrix} --weights {params.enpact_weights} --metadata {input.metadata} --split --output_basename {params.output_basename} --subset_of_loci {params.loci_subset}

```{r}
input_matrix <- "/beagle3/haky/users/temi/projects/TFXcan/files/prostate_cancer_risk.2024-09-30.processed.matrix.h5.gz"
input_metadata <- "/beagle3/haky/users/temi/projects/TFXcan/files/prostate_cancer_risk.2024-09-30.processed.metadata.tsv"
enpact_weights <- "/beagle3/haky/users/temi/projects/Enpact/data/enpact/weights/ENPACT_734_2025-04-24.compiled_weights.with_intercept.lambda.1se.txt.gz"
output_basename <- "/scratch/beagle3/temi/enpact_scores/prostate_cancer_risk/prostate_cancer_risk/prostate_cancer_risk"

glue::glue("sbatch /beagle3/haky/users/temi/projects/TFXcan/src/predict_enpact_scores.sbatch {input_matrix} {enpact_weights} {input_metadata} {output_basename}")
```



## Checks

```{r}
old_db <- readRDS("/beagle3/haky/users/temi/projects/TFXcan-snakemake/data/prostate_cancer_risk_2024-09-30/enpactdb/prostate_cancer_risk.enpact_scores.array.rds.gz")
dim(old_db)
old_scores <- as.data.table(old_db[, , 'ATF2_Embryo'], keep.rownames = 'locus')
old_scores[1:5, 1:5]
```

```{r}
new_scores <- data.table::fread('/scratch/beagle3/temi/enpact_scores/prostate_cancer_risk/prostate_cancer_risk.ATF2_Embryo.enpact_scores.txt') %>%
    dplyr::mutate(NAME = gsub('ATF2_Embryo_', '', NAME)) %>% dplyr::rename(locus = NAME)
new_scores[1:5, 1:5]
```


```{r}
dim(old_scores); dim(new_scores)
```


```{r}
common_loci <- intersect(old_scores$locus, new_scores$locus)
length(common_loci)
```


```{r}
new_scores <- new_scores %>% dplyr::arrange(locus = match(locus, common_loci))
new_scores[1:5, 1:5]
```

```{r}
old_scores <- old_scores %>% dplyr::arrange(locus = match(locus, common_loci))
old_scores[1:5, 1:5]
```


```{r}
cor(old_scores$HG01771, new_scores$HG01771)
```


```{r}
purrr::map(colnames(old_scores)[-1], function(each_name){
    cor(old_scores[[each_name]], new_scores[[each_name]])
}) %>% unlist() %>% hist()
```


```{r}
wgts <- data.table::fread("/beagle3/haky/users/temi/projects/Enpact/data/enpact/weights/ENPACT_734_2025-04-24.compiled_weights.with_intercept.lambda.1se.txt.gz")

wgts %>% dplyr::select("ATF2_Embryo")
```



```{r}
list.files('/scratch/beagle3/temi/enpact_scores/prostate_cancer_risk') %>%
    gsub("prostate_cancer_risk.|.annotation.txt|.enpact_scores.txt", "", .) %>% unique() %>%
    as.data.table() %>%
    data.table:::fwrite("/beagle3/haky/users/temi/projects/TFXcan/files/enpact_models.linearize.list.tsv", col.names = F, row.names = F, quote = F, sep = '\t')
```


```{r}
xenformer <- data.table::fread(base::subset(individual_metadata, individual == 'HG01771')$path) %>%
    dplyr::arrange(id = match(id, common_loci))
locs <- xenformer$id
Xmat <- xenformer %>% dplyr::select(-id) %>% as.matrix()
dim(Xmat)
```


```{r}
wgts <- data.table::fread("/beagle3/haky/users/temi/projects/Enpact/data/enpact/weights/ENPACT_734_2025-04-24.compiled_weights.with_intercept.lambda.1se.txt.gz")
wgts %>% dplyr::select("AR_Prostate")
```

```{r}
wgts <- wgts %>% dplyr::select(-feature)  %>% as.matrix()
wgts_intercept <- wgts[1, , drop = F]
wgts_features <- wgts[-1, , drop = F]
dim(wgts_intercept); dim(wgts_features)
```


```{r}
Ymat_noIntercept <- (Xmat %*% wgts_features)
Ymat <- apply(Ymat_noIntercept, 1, function(each_row){
    wgts_intercept + each_row
}) %>% t()

colnames(Ymat) <- colnames(wgts_features)

dim(Ymat)
```

```{r}
Ymat[1:5, 1:5]
```

```{r}
dt1 <- data.table::fread("/scratch/midway3/temi/enpact_predictions/prostate_cancer_risk/HG01771.prostate_cancer_risk.aggByCollect.2024-09-30.csv.gz") %>%
    dplyr::arrange(locus = match(locus, common_loci)) %>% dplyr::select("AR_Prostate")
dt1[1:10, 1]
```

```{r}
old_scores$HG01771[1:10]; new_scores$HG01771[1:10]
```


```{r}
Ymat[1:10, 'AR_Prostate']
```


```{r}
Ymat_noIntercept[1:10, 'AR_Prostate']
```


```{r}
new_scores$HG01771[1:10] - old_scores$HG01771[1:10]
```


```{r}
wgts_intercept[1:5]
```


```{r}
each_db <- DBI::dbConnect(RSQLite::SQLite(), "/scratch/beagle3/temi/lEnpact/enpact_models_linearization_2025-10-16/AR_MammaryGland/models/filtered_db/predict_db_AR_MammaryGland_filtered.db")
dbListTables(each_db)
extras <- dbGetQuery(each_db, 'SELECT * FROM extra LIMIT 5')
dbDisconnect(each_db)
```


```{r}
dbGetQuery(each_db, 'SELECT * FROM weights LIMIT 5')
```

```{r}
dbDisconnect(each_db)
```