---
title: "speedUpPredCrossVar"
author: "Marnin Wolfe"
date: "2021-May-25"
output: 
  workflowr::wflow_html:
    toc: true
editor_options:
  chunk_output_type: inline
---

```{r test Rcpp and Armadillo}
library(RcppArmadillo); library(Rcpp)

sourceCpp(here::here("code","test.cpp"))
add_two(3)
Rcpp::sourceCpp(here::here("code","helloworld.cpp"))
```

Test Rcpp set-up on both BioHPC and locally.

# Attempts to optimize code

Start with single thread BLAS. 
```{bash}
# module load R/4.0.5;
# module load gcc/10.2.0

cd /home/jj332_cas/marnin/implementGMSinCassava/; 
export OMP_NUM_THREADS=1
R # initiate R session
```

## calcGameticLD
```{r toy inputs}
library(tidyverse); library(Rcpp); library(RcppArmadillo); library(microbenchmark)
haploMat<-sample(0:1,replace = T,size = 20000) %>% 
  matrix(.,nrow=2)
rownames(haploMat)<-paste0("A",c("_HapA","_HapB"))
dim(haploMat)
recombFreqMat<-matrix(0,nrow=ncol(haploMat),ncol=ncol(haploMat))
recombFreqMat<-matrix(runif(length(recombFreqMat)),
                      nrow=ncol(haploMat),
                      ncol=ncol(haploMat))
recombFreqMat[upper.tri(recombFreqMat)]<-t(recombFreqMat[lower.tri(recombFreqMat)])
diag(recombFreqMat)<-1
# dim(recombFreqMat)
# recombFreqMat[1:5,1:5]

parentGID<-"A"
```


```{r}
# X<-haploMat[paste0(parentGID,c("_HapA","_HapB")),,drop=F]
# dim(crossprod(X))
# # 500 x 500
# dim(tcrossprod(colMeans(X)))
# # 500 x 500
# dim((0.5*crossprod(X))-tcrossprod(colMeans(X)))
# # 500 x 500
```


```{r base and cpp functions}
calcGameticLD_r<-function(parentGID,recombFreqMat,haploMat){
  X<-haploMat[paste0(parentGID,c("_HapA","_HapB")),,drop=F]
  D<-recombFreqMat*((0.5*crossprod(X))-tcrossprod(colMeans(X)))
  return(D) } 

calcGameticLD_cpp<-function(parentGID,recombFreqMat,haploMat){
  X<-haploMat[paste0(parentGID,c("_HapA","_HapB")),,drop=F]
  D<-gameticLD(recombFreqMat = recombFreqMat,X = X)
  return(D) } 
```


```{r}
sourceCpp(here::here("code","calcGameticLD.cpp"))
#table(colMeans(X)==col_means(X))
testcpp<-calcGameticLD_cpp(parentGID = "A",
                           recombFreqMat = recombFreqMat,
                           haploMat = haploMat)
testr<-calcGameticLD_r(parentGID = "A",
                       recombFreqMat = recombFreqMat,
                       haploMat = haploMat)
all.equal(testcpp,testr)
```
```{r}
testcpp[1:5,1:5]
```
```{r}
testr[1:5,1:5]
```

```{r}
microbenchmark(cpp=calcGameticLD_cpp(parentGID = "A",
                           recombFreqMat = recombFreqMat,
                           haploMat = haploMat),
               r=calcGameticLD_r(parentGID = "A",
                                 recombFreqMat = recombFreqMat,
                                 haploMat = haploMat),
               times=5) 
```
So my cpp version of `calcGameticLD` is a loser relative to the R version.

What about the next part, dense matrix addition of the two gameticLD matrices?

```{r}
sourceCpp(here::here("code","calcGameticLD.cpp"))
dim(recombFreqMat)
microbenchmark(base=recombFreqMat+recombFreqMat,
               cpp=add_mats(recombFreqMat,recombFreqMat),
               times=5)
```

```{r calcGameticLD - base R optimization}
parentGID<-sireID
X<-haploMat[paste0(parentGID,c("_HapA","_HapB")),,drop=F]
all.equal(recombFreqMat*((0.5*crossprod(X))-tcrossprod(colMeans(X))),
          recombFreqMat*(crossprod(scale(X,scale = F))/2))

p<-colMeans(X)
all.equal(recombFreqMat*((0.5*t(X)%*%X)-p%*%t(p)),
          recombFreqMat*((0.5*crossprod(X))-tcrossprod(colMeans(X))))

calcDv1<-function(X){ recombFreqMat*((0.5*crossprod(X))-tcrossprod(colMeans(X))) }
calcDv2<-function(X) { p<-colMeans(X); return(recombFreqMat*((0.5*t(X)%*%X)-p%*%t(p))); }
calcDv3<-function(X) { p<-colMeans(X); return(recombFreqMat*((0.5*t(X)%*%X)-tcrossprod(p))); }
calcDv4<-function(X) { p<-colMeans(X); return(recombFreqMat*((0.5*crossprod(X))-tcrossprod(p))); }
calcDv5<-function(X) { 
  p<-colMeans(X); 
  XtX<-0.5*crossprod(X); 
  ppt<-tcrossprod(p)
  return(recombFreqMat*(XtX-ppt)); }
```
```{r}
microbenchmark(D1=calcDv1(X),
               D2=calcDv2(X),
               D3=calcDv3(X), 
               D4=calcDv4(X),
               D5=calcDv5(X),
               times=500)
```

Somehow, V2, is fastest!?

Big test on server. 

Ran twice, once with single and once with 56-threaded BLAS.
```{r calcGameticLD - big test on server}
library(microbenchmark)
require(tidyverse); require(magrittr); 
# Parent folds - ONLY 10 CROSSES TO PREDICT!!!!
parentfolds<-readRDS(file=here::here("output","parentfolds.rds")) %>% 
  filter(Repeat=="Repeat1",Fold=="Fold1") %>%   
  mutate(CrossesToPredict=map(CrossesToPredict,~slice(.,1:10)))
parents<-parentfolds %>% 
  select(Repeat,Fold,CrossesToPredict) %>% 
  unnest(CrossesToPredict) %>% 
  distinct(sireID,damID) %$%
  union(sireID,damID)

recombFreqMat<-readRDS(file=here::here("data",
                                       "recombFreqMat_1minus2c_2021May13.rds"))
# Haplotype Matrix
haploMat<-readRDS(file=here::here("data","haps_IITA_filtered_2021May13.rds"))
## keep only haplos for parents-to-be-predicted, saves memory
parenthaps<-sort(c(paste0(parents,"_HapA"),
                   paste0(parents,"_HapB")))
haploMat<-haploMat[parenthaps,colnames(recombFreqMat)]; 
parentGID<-parents[1]
```

 
```{r}
calcGameticLDv1<-function(parentGID,recombFreqMat,haploMat){
  X<-haploMat[paste0(parentGID,c("_HapA","_HapB")),,drop=F]
  D<-recombFreqMat*((0.5*crossprod(X))-tcrossprod(colMeans(X)))
  return(D) } 
calcGameticLDv2<-function(parentGID,recombFreqMat,haploMat){
  X<-haploMat[paste0(parentGID,c("_HapA","_HapB")),,drop=F]
  p<-colMeans(X); 
  D<-recombFreqMat*((0.5*t(X)%*%X)-p%*%t(p));
  return(D) } 
microbenchmark(D1=calcGameticLDv1(parentGID,recombFreqMat,haploMat),
               D2=calcGameticLDv2(parentGID,recombFreqMat,haploMat),
               times=5)
# OMP_NUM_THREADS=1
# Unit: seconds
# expr      min       lq     mean   median       uq      max neval cld
#    D1 49.98993 50.31883 52.24512 52.36281 53.78908 55.13119    10   b
#    D2 27.98526 28.27352 29.09537 28.72401 29.67499 31.39269    10  a
# OMP_NUM_THREADS=56
# Unit: seconds
 # expr      min       lq     mean   median       uq      max neval cld
 #   D1 57.48309 57.98060 61.42607 58.28465 62.35851 71.02349     5   b
 #   D2 15.03936 15.63792 19.27314 16.94444 20.48497 28.25904     5  a
```

```{r calcCrossLD - additional optimization}
f1<-function(sireID,damID,recombFreqMat,haploMat){
  sireLD<-calcGameticLD(sireID,recombFreqMat,haploMat)
  damLD<-calcGameticLD(damID,recombFreqMat,haploMat)
  progenyLD<-sireLD+damLD
  return(progenyLD)
}
f2<-function(sireID,damID,recombFreqMat,haploMat){
  return(calcGameticLD(sireID,recombFreqMat,haploMat)+calcGameticLD(damID,recombFreqMat,haploMat))
}
microbenchmark(f1=f1(sireID,damID,recombFreqMat,haploMat),
               f2=f2(sireID,damID,recombFreqMat,haploMat),times=1000)
# expr      min       lq     mean   median       uq      max neval cld
# f1	5.976532	6.701085	9.889478	6.960663	7.301757	280.77230	1000
# f2	5.631099	6.463040	8.787175	6.728689	7.003273	56.18071	1000
microbenchmark(f2=f2(sireID,damID,recombFreqMat,haploMat),
               f1=f1(sireID,damID,recombFreqMat,haploMat),times=1000)
# expr      min       lq     mean   median       uq      max neval cld
# f2	5.894926	6.580749	8.988388	6.838795	7.256322	50.56213	1000
# f1	5.788500	6.845660	10.250824	7.103177	7.498478	277.20737	1000
```


## variance prediction - quadratic forms
```{r inputs}
library(tidyverse); library(Rcpp); library(RcppArmadillo); library(microbenchmark)
haploMat<-sample(0:1,replace = T,size = 20000) %>% 
  matrix(.,nrow=2)
rownames(haploMat)<-paste0("A",c("_HapA","_HapB"))
dim(haploMat)
recombFreqMat<-matrix(0,nrow=ncol(haploMat),ncol=ncol(haploMat))
recombFreqMat<-matrix(runif(length(recombFreqMat)),
                      nrow=ncol(haploMat),
                      ncol=ncol(haploMat))
recombFreqMat[upper.tri(recombFreqMat)]<-t(recombFreqMat[lower.tri(recombFreqMat)])
diag(recombFreqMat)<-1
parentGID<-"A"

x<-matrix(rnorm(1000),nrow=1)
xvec<-as.numeric(x)
D<-crossprod(x)
dim(x)
```
```{r}
x%*%D%*%t(x)==xvec%*%D%*%xvec
```
```{r x as mat vs as vec}
library(microbenchmark)
microbenchmark(xasmat=x%*%D%*%t(x),
               xasvec=xvec%*%D%*%xvec)
```
No diff
```{r}
#install.packages("emulator")
library("emulator")
quad.form(D,xvec)==xvec%*%D%*%xvec
## crossprod(crossprod(M, Conj(x)), x)
```
```{r}
microbenchmark(xasmat=x%*%D%*%t(x),
               xasvec=xvec%*%D%*%xvec,
               quad.form(D,xvec))
```
```{r}
colSums(xvec * (D %*% xvec))==as.numeric(xvec%*%D%*%xvec)
```
```{r}
all.equal(colSums(xvec * (D %*% xvec)),as.numeric(xvec%*%D%*%xvec))
```
```{r}
colSums(xvec * (D %*% xvec)) %>% str
```
```{r}
str(as.numeric(xvec%*%D%*%xvec))
```
```{r}
as.numeric(as.matrix(colSums(xvec * (D %*% xvec))))->y
y-as.numeric(xvec%*%D%*%xvec)
## TINY FLOAT POINT DIFFS BETWEEN COMPUTE METHODS
```
```{r fastest quad form}
microbenchmark(xasmat=x%*%D%*%t(x),
               xasvec=xvec%*%D%*%xvec,
               quad_dot_form=quad.form(D,xvec),
               colSums=as.numeric(colSums(xvec * (D %*% xvec))),
               times = 1000)
```
```{r fastest matrix square}
# all.equal(D*D,D^2)
# [1] TRUE
microbenchmark(D*D,D^2,times = 1000)
```

```{r}
sourceCpp(here::here("code","misc.cpp"))
quadraticForm(xvec,D,xvec)
```
```{r}
xvec%*%D%*%xvec
```
```{r}
microbenchmark(cpp=quadraticForm(xvec,D,xvec),
               r=xvec%*%(D%*%xvec))
```

```{r}
sourceCpp(here::here("code","misc.cpp"))
all.equal(D*D,schur(D,D))
microbenchmark(r=D*D,
               cpp=schur(D,D))
```

```{r benchmark dom var pred}
f1<-function(progenyLD,postMeanDomEffects){
  progenyLDsq<-progenyLD^2
  vpm_m2d<-postMeanDomEffects[[Trait1]]%*%progenyLDsq%*%postMeanDomEffects[[Trait2]]
  return(vpm_m2d)
}
f2<-function(progenyLD,postMeanDomEffects){
  progenyLDsq<-progenyLD^2
  vpm_m2d<-quadform(D=progenyLDsq,
                    x=postMeanAddEffects[[Trait1]],
                    y=postMeanAddEffects[[Trait2]])
  return(vpm_m2d)
}
f3<-function(progenyLD,postMeanDomEffects){
  progenyLDsq<-progenyLD*progenyLD
  vpm_m2d<-quadform(D=progenyLDsq,
                    x=postMeanAddEffects[[Trait1]],
                    y=postMeanAddEffects[[Trait2]])
  return(vpm_m2d)
}
f4<-function(progenyLD,postMeanDomEffects){
  return(quadform(D=progenyLD*progenyLD,
                  x=postMeanAddEffects[[Trait1]],
                  y=postMeanAddEffects[[Trait2]])) }
microbenchmark(f1=f1(progenyLD,postMeanDomEffects),
               f2=f2(progenyLD,postMeanDomEffects),
               f3=f3(progenyLD,postMeanDomEffects),
               f4=f4(progenyLD,postMeanDomEffects),times=2000)  
## round 1
# f1	1326.908	1566.399	2267.651	1614.721	1681.483	47055.99	2000
# f2	1146.317	1390.024	2087.710	1436.031	1514.477	46003.58	2000
# f3	957.699	1224.213	1767.830	1278.941	1351.535	45480.06	2000
# f4	938.240	1223.220	1859.626	1282.154	1358.268	256827.98	2000
## round 2
# f1	1318.908	1576.671	2101.041	1629.053	1703.016	46061.93	2000
# f2	1146.777	1397.525	2197.470	1455.725	1537.740	264372.33	2000
# f3	940.414	1231.119	1938.390	1300.162	1383.787	47749.10	2000
# f4	942.878	1235.495	1838.589	1302.143	1386.672	51435.20	2000
```

## avoid big objects in workers
"lexical scoping"?
https://furrr.futureverse.org/articles/articles/gotchas.html#package-development-1

```{r}
my_fast_fn <- function() { future_map(1:5, ~.x) }
my_slow_fn <- function() {
  # Massive object - but we don't want it in `.f`
  big <- 1:1e8 + 0L
  future_map(1:5, ~.x) }
plan(multisession, workers = 2)
system.time( my_fast_fn() )
#>    user  system elapsed 
#>   0.120   0.004   0.532

system.time( my_slow_fn() )
#>    user  system elapsed 
#>   7.048   2.080  13.295

fn <- function(x) { x }
my_not_so_slow_fn <- function() {
  big <- 1:1e8 + 0L
  future_map(1:5, fn) }
system.time( my_not_so_slow_fn())
#>    user  system elapsed 
#>   1.405   1.296   3.275
```

## Conclusions and lessons learned

1. `D*D` is faster than `D^2`
2. `as.numeric(colSums(xvec * (D %*% xvec)))` is about 17% faster than alternative quadratic forms `xvec %*% D %*% Xvec` equal up to a tiny ~1e-9 difference due to presumably to numerical precision of functions used.
3. So far, trying matrix operations using `Rcpp` and `RcppArmadillo` **do not** have an impact on performance. R is pretty optimized.
4. Try to use 'lexical scoping' and avoid big objects passed into works without need
5. Avoid grouped data frames, i.e. use `ungroup()` as it may interfere with furrr and prevent it from splitting the job correctly


# predCrossVar() - streamline and optimize

- Write a single function `predCrossVar()` and streamline it for VPM prediction. 
- Compare it to `runMtCrossVarPredsAD()`.
- `runMtCrossVarPredsAD()` is optimized for PMV. 
  - PMV prediction requires the loading-into-memory of huge 3-dimensional arrays of SNP-effects-by-MCMCiteration. 
  - The function must predict all variances and covariances for all crosses.
  - This is done in nested loops. 
  - Choices were made on the ordering of the loops.
  - For PMV, I load the big 3-dim. arrays in an outer-loop
  - In the inner loops compute over crosses and varcomps. 
- For each cross, must compute gameticLD matrices of each parent and sum them.
  - For VPM, no need for the big 3-dim. array.
  - Want to compute the gameticLD matrices _only once_ per cross
  - **Put this step on an outer loop and benchmark.**
```{r inputs for predictCrosses}
require(tidyverse); require(magrittr); 
# Parent folds - ONLY 10 CROSSES TO PREDICT!!!!
parentfolds<-readRDS(file=here::here("output","parentfolds.rds")) %>% 
  filter(Repeat=="Repeat1",Fold=="Fold1") %>%   
  mutate(CrossesToPredict=map(CrossesToPredict,~slice(.,1:10)))

parents<-parentfolds %>% 
  select(Repeat,Fold,CrossesToPredict) %>% 
  unnest(CrossesToPredict) %>% 
  distinct(sireID,damID) %$%
  union(sireID,damID)

# Recomb frequency matrix
recombFreqMat<-readRDS(file=here::here("data",
                                       "recombFreqMat_1minus2c_2021May13.rds"))

# sample snps 2 keep - ONLY 1000
set.seed(1234)
snps2keep<-sample(colnames(recombFreqMat),size = 1000,replace = F)
recombFreqMat %<>% .[snps2keep,snps2keep]
# Haplotype Matrix
haploMat<-readRDS(file=here::here("data","haps_IITA_filtered_2021May13.rds"))
## keep only haplos for parents-to-be-predicted, saves memory
parenthaps<-sort(c(paste0(parents,"_HapA"),
                   paste0(parents,"_HapB")))
haploMat<-haploMat[parenthaps,colnames(recombFreqMat)]; 

# Marker effects - one REPEAT-FOLD only
markEffsAD<-readRDS(here::here("output","markerEffectsAD.rds"))
markEffsAD %<>% 
  filter(Repeat=="Repeat1",Fold=="Fold1") %>%   
  mutate(effects=map(effects,~slice(.,1:2)))
modelType="AD"

# Formatting and tidying inputs for predCrossVar
predvars<-markEffsAD %>%
    unnest(effects) %>% 
    unnest(modelOut) %>% 
    filter(Dataset=="trainset") %>% 
    dplyr::select(Repeat,Fold,Trait,modelType,
                  any_of(c("addsnpeff","domsnpeff"))) %>% 
    nest(EffectList=c(Trait,any_of(c("addsnpeff","domsnpeff")))) %>% 
    mutate(AddEffectList=map(EffectList,
                             function(EffectList){
                               addsnpeff<-map(EffectList$addsnpeff,~t(.))
                               names(addsnpeff)<-EffectList$Trait
                               return(addsnpeff)}))
  if(modelType %in% c("AD")){
    predvars<-predvars %>% 
      mutate(DomEffectList=map(EffectList,
                               function(EffectList){
                                 domsnpeff<-map(EffectList$domsnpeff,~t(.))
                                 names(domsnpeff)<-EffectList$Trait
                                 return(domsnpeff) })) }

predvars %<>%
  left_join(parentfolds %>% 
              dplyr::select(-testparents,-trainset,-testset)) %>% 
  dplyr::select(-EffectList)
predvars$AddEffectList[[1]] %<>% map(.,~.[,snps2keep,drop=F])
predvars$DomEffectList[[1]] %<>% map(.,~.[,snps2keep,drop=F])
```

```{r inputs for predCrossVar}
ncores<-10
library(microbenchmark)

# runMtCrossVarPredsAD
CrossesToPredict<-predvars$CrossesToPredict[[1]]
AddEffectList<-predvars$AddEffectList[[1]]; 
DomEffectList<-predvars$DomEffectList[[1]]
predType="VPM"
# predOneCrossAD
# sireID<-CrossesToPredict$sireID[[1]]
# damID<-CrossesToPredict$damID[[1]]


# predOneVarAD
# Trait1<-varcovars$Trait1[[1]]
# Trait2<-varcovars$Trait2[[1]]
```

```{r calcCrossLD}
calcGameticLD<-function(parentGID,recombFreqMat,haploMat){
  X<-haploMat[paste0(parentGID,c("_HapA","_HapB")),,drop=F]
  p<-colMeans(X); 
  D<-recombFreqMat*((0.5*t(X)%*%X)-p%*%t(p));
  return(D) } 

calcCrossLD<-function(sireID,damID,recombFreqMat,haploMat){
  return(calcGameticLD(sireID,recombFreqMat,haploMat)+
           calcGameticLD(damID,recombFreqMat,haploMat)) }
```

```{r predCrossVar - top level}
predCrossVar<-function(CrossesToPredict,AddEffectList,DomEffectList,
                       haploMat,recombFreqMat,ncores=1,predType="VPM"){
  starttime<-proc.time()[3]
  # Center posterior distribution of effects
  ## on posterior mean across MCMC samples
  AddEffectList<-purrr::map(AddEffectList,~scale(.,center = T, scale = F));
  DomEffectList<-purrr::map(DomEffectList,~scale(.,center = T, scale = F))
  
  ## Extract the posterior mean effects vectors
  ## If predType="VPM" this just recovers the original effects
  postMeanAddEffects<-purrr::map(AddEffectList,~attr(.,which = "scaled:center"))
  postMeanDomEffects<-purrr::map(DomEffectList,~attr(.,which = "scaled:center"))
  
  parents<-union(CrossesToPredict$sireID,
                 CrossesToPredict$damID)
  haploMat<-haploMat[sort(c(paste0(parents,"_HapA"),paste0(parents,"_HapB"))),]
  
  if(predType!="PMV"){ AddEffectList<-NULL; DomEffectList<-NULL; }
  
  # Set-up a loop over the crosses
  require(furrr); options(mc.cores=ncores); plan(multicore)
  #require(furrr); options(mc.cores=ncores); plan(multisession)
  crossespredicted<-CrossesToPredict %>% 
    mutate(predVars=future_pmap(.,
                                predOneCrossAD,
                                haploMat=haploMat,
                                recombFreqMat=recombFreqMat,
                                predType=predType,
                                postMeanAddEffects=postMeanAddEffects,
                                postMeanDomEffects=postMeanDomEffects)) %>% 
    unnest(predVars)
  totcomputetime<-proc.time()[3]-starttime
  print(paste0("Done predicting fam vars. ",
               "Took ",round((totcomputetime)/60,2),
               " mins for ",nrow(crossespredicted)," crosses"))
  return(crossespredicted) 
}
```
```{r}
test<-microbenchmark(onecore=predCrossVar(CrossesToPredict,AddEffectList,DomEffectList,
                                    haploMat,recombFreqMat,ncores=1,predType="VPM"),
               multicore=predCrossVar(CrossesToPredict,AddEffectList,DomEffectList,
                                      haploMat,recombFreqMat,ncores=10,predType="VPM"),
               times=5)
test
# On server, with OMP_NUM_THREADS=56
# Unit: milliseconds
#       expr      min       lq     mean   median       uq      max neval cld
#    onecore  805.877 1107.004 1233.274 1245.842 1416.614 1591.035     5  a 
#  multicore 1851.174 2082.563 2158.426 2215.975 2229.953 2412.465     5   b
```

runMtCrossVarPreds with 1 vs. 10 threads

On server. `OMP_NUM_THREADS=56`
```{r}
microbenchmark(runMt_10core=predCrossVar::runMtCrossVarPredsAD(CrossesToPredict = CrossesToPredict, 
                                                               AddEffectList = AddEffectList,
                                                               DomEffectList = DomEffectList,
                                                               haploMat = haploMat,
                                                               recombFreqMat = recombFreqMat,ncores=10),
               runMt_1core=predCrossVar::runMtCrossVarPredsAD(CrossesToPredict = CrossesToPredict, 
                                                              AddEffectList = AddEffectList,
                                                              DomEffectList = DomEffectList,
                                                              haploMat = haploMat,
                                                              recombFreqMat = recombFreqMat,ncores=1),
               times=5)
# locally 
# Unit: seconds
#          expr       min        lq      mean    median        uq       max neval
#  runMt_10core 19.876403 19.880891 20.295269 20.139956 20.152783 21.426310     5
#   runMt_1core  4.956376  5.007242  5.028677  5.042166  5.053157  5.084445     5
```

Local.
```{r test predCrossVar vs runMtCrossVar 1 - single core}
test<-microbenchmark(predCrossVar=predCrossVar(CrossesToPredict,AddEffectList,DomEffectList,
                                               haploMat,recombFreqMat,ncores=1,predType="VPM"),
                     runMtCrossVarPreds=predCrossVar::runMtCrossVarPredsAD(CrossesToPredict = CrossesToPredict, 
                                                                           AddEffectList = AddEffectList,
                                                                           DomEffectList = DomEffectList, 
                                                                           haploMat = haploMat,
                                                                           recombFreqMat = recombFreqMat,ncores=1),
                     times = 10)
test
# run 1
# predCrossVar        	364.6201	372.6661	399.4438	383.7957	385.6236	490.5137	5
# runMtCrossVarPreds	6092.3644	6106.3892	6193.0670	6172.5342	6192.1642	6401.8829	5
# run 2
# predCrossVar	       356.8027	370.7798	384.5027	378.4938	384.4371	459.2163	10
# runMtCrossVarPreds	6182.1222	6201.7970	6232.8530	6216.2808	6261.3899	6320.9898	10
```

```{r test predCrossVar vs runMtCrossVar 2 - multicore}
test2<-microbenchmark(predCrossVar=predCrossVar(CrossesToPredict,
                                               AddEffectList,
                                               DomEffectList,
                                               haploMat,recombFreqMat,
                                               ncores=10,predType="VPM"),
                     runMtCrossVarPreds=predCrossVar::runMtCrossVarPredsAD(CrossesToPredict = CrossesToPredict, 
                                                                           AddEffectList = AddEffectList,
                                                                           DomEffectList = DomEffectList, 
                                                                           haploMat = haploMat,
                                                                           recombFreqMat = recombFreqMat,
                                                                           ncores=10),
                     times = 5)
test2
# run 1
# Unit: milliseconds
#                expr       min         lq       mean     median         uq
#        predCrossVar   377.562   403.2724   405.9821   403.6371   414.4919
#  runMtCrossVarPreds 19484.025 19793.0929 20156.7465 20172.3361 20580.0936
```
```{r test predCrossVar - multicore vs single core}
test3<-microbenchmark(predCrossVar10core=predCrossVar(CrossesToPredict,
                                               AddEffectList,
                                               DomEffectList,
                                               haploMat,recombFreqMat,
                                               ncores=10,predType="VPM"),
                     predCrossVar1core=predCrossVar(CrossesToPredict,
                                               AddEffectList,
                                               DomEffectList,
                                               haploMat,recombFreqMat,
                                               ncores=1,predType="VPM"),
                     times = 5)
test3
# Unit: milliseconds
#                expr      min       lq     mean   median       uq      max neval
#  predCrossVar10core 339.9820 347.8085 377.7624 354.1135 403.4701 443.4378     5
#   predCrossVar1core 400.3883 444.9101 455.7601 445.8505 478.5584 509.0932     5
```


```{r predCrossVar - predOneCrossAD}
predOneCrossAD<-function(sireID,damID,
                         haploMat,recombFreqMat,
                         predType,
                         postMeanAddEffects,postMeanDomEffects,...){
  starttime<-proc.time()[3]
  # Before predicting variances
  # check for and remove SNPs that
  # won't segregate, i.e. are fixed in parents
  ### hopes to save time / mem
  x<-colSums(rbind(haploMat[grep(paste0("^",sireID,"_"),rownames(haploMat)),],
                   haploMat[grep(paste0("^",damID,"_"),rownames(haploMat)),]))
  segsnps2keep<-names(x[x>0 & x<4])
  
  if(length(segsnps2keep)>0){
    recombFreqMat<-recombFreqMat[segsnps2keep,segsnps2keep,drop=F]
    haploMat<-haploMat[,segsnps2keep,drop=F]
    postMeanAddEffects<-purrr::map(postMeanAddEffects,~.[segsnps2keep])
    postMeanDomEffects<-purrr::map(postMeanDomEffects,~.[segsnps2keep])
    if(predType=="PMV"){
      postVarCovarOfAddEffects<-postVarCovarOfAddEffects[segsnps2keep,
                                                         segsnps2keep,drop=F]
      postVarCovarOfDomEffects<-postVarCovarOfDomEffects[segsnps2keep,
                                                         segsnps2keep,drop=F]
    }
    # calc cross LD matrix
    progenyLD<-calcCrossLD(sireID,damID,recombFreqMat,haploMat)
    
    # Set-up loop over variance and covarance parameters
    ## Trait variances to-be-predicted
    traits<-names(postMeanAddEffects)
    varcovars<-tibble::tibble(Trait1=traits,
                              Trait2=traits) 
    ## If multiple traits
    if(length(traits)>1){
      ## Add covariances to-be-predicted
      varcovars<-dplyr::bind_rows(varcovars, # trait variances
                                  combn(traits,2,simplify = T) %>% # covariances
                                    t(.) %>% #
                                    `colnames<-`(.,c("Trait1","Trait2")) %>%
                                    tibble::as_tibble(.)) }
    
    varcovars<-varcovars %>%
      dplyr::mutate(predVars=purrr::pmap(.,
                                         predOneCrossVarAD,
                                         progenyLD=progenyLD,
                                         haploMat=haploMat,
                                         recombFreqMat=recombFreqMat,
                                         predType=predType,
                                         postMeanAddEffects=postMeanAddEffects,
                                         postMeanDomEffects=postMeanDomEffects,
                                         postVarCovarOfAddEffects=postVarCovarOfAddEffects,
                                         postVarCovarOfDomEffects=postVarCovarOfDomEffects)) %>% 
      unnest(predVars)
    computetime<-proc.time()[3]-starttime
    out_thiscross<-tibble(Nsegsnps=length(segsnps2keep),
                          ComputeTime=computetime,
                          predVars=list(varcovars)) 
  } else { 
    out_thiscross<-tibble(Nsegsnps=length(segsnps2keep),
                          ComputeTime=computetime,
                          predVars=list()) }
  return(out_thiscross)
}
# predOneCrossAD(sireID,damID,
#                haploMat,recombFreqMat,
#                predType="VPM",
#                postMeanAddEffects,postMeanDomEffects)
```

```{r predCrossVar - predOneVarAD}
predOneCrossVarAD<-function(Trait1,Trait2,progenyLD,
                            haploMat,recombFreqMat,predType,
                            postMeanAddEffects,postMeanDomEffects,
                            postVarCovarOfAddEffects=NULL,postVarCovarOfDomEffects=NULL,...){
  ## Predicts variance for one cross and
  ### one variance or covariance paramater (Trait1-Trait2)
  
  ## Predict cross additive variance
  #### posterior mean (VPM)
  quadform<-function(D,x,y){ return(as.numeric(colSums(x*(D%*%y)))) } # does x %*% D %*% t(x) faster
  predVarA_vpm<-quadform(D=progenyLD,
                         x=postMeanAddEffects[[Trait1]],
                         y=postMeanAddEffects[[Trait2]])
  
  #### posterior mean (co)variance (PMV)
  if(predType=="PMV"){ predVarA_pmv<-predVarA_vpm+sum(diag(progenyLD%*%postVarCovarOfAddEffects)) }
  
  ## Predict cross dominance variance
  #### VPM
  progenyLDsq<-progenyLD*progenyLD
  predVarD_vpm<-quadform(D=progenyLDsq,
                         x=postMeanAddEffects[[Trait1]],
                         y=postMeanAddEffects[[Trait2]])
  #### PMV
  if(predType=="PMV"){ predVarD_pmv<-predVarD_vpm+sum(diag(progenyLDsq%*%postVarCovarOfDomEffects)) }
  
  # Tidy the results
  out<-bind_rows(tibble(predOf="VarA",predVar=predVarA_vpm),
                 tibble(predOf="VarD",predVar=predVarD_vpm)) %>% 
    mutate(predMethod="VPM")
  if(predType=="PMV"){
    out<-out %>% 
      bind_rows(bind_rows(tibble(VarComp="VarA",predVar=predVarA_pmv),
                          tibble(VarComp="VarD",predVar=predVarD_pmv)) %>% 
                  mutate(predMethod="PMV")) }
  return(out)
}
```

# Big test - big data

## Input
```{bash}
# module load R/4.0.5;
# module load gcc/10.2.0

cd /home/jj332_cas/marnin/implementGMSinCassava/; 
export OMP_NUM_THREADS=5
R # initiate R session
```

```{r inputs for predictCrosses}
require(tidyverse); require(magrittr); library(microbenchmark)
# Parent folds - ONLY 20 CROSSES TO PREDICT!!!!
parentfolds<-readRDS(file=here::here("output","parentfolds.rds")) %>% 
  filter(Repeat=="Repeat1",Fold=="Fold1") %>%   
  mutate(CrossesToPredict=map(CrossesToPredict,~slice(.,1:20)))

parents<-parentfolds %>% 
  select(Repeat,Fold,CrossesToPredict) %>% 
  unnest(CrossesToPredict) %>% 
  distinct(sireID,damID) %$%
  union(sireID,damID)

# Recomb frequency matrix
recombFreqMat<-readRDS(file=here::here("data",
                                       "recombFreqMat_1minus2c_2021May13.rds"))

# sample snps 2 keep - ONLY 5000
set.seed(1234)
snps2keep<-sample(colnames(recombFreqMat),size = 5000,replace = F)
recombFreqMat %<>% .[snps2keep,snps2keep]
# Haplotype Matrix
haploMat<-readRDS(file=here::here("data","haps_IITA_filtered_2021May13.rds"))
## keep only haplos for parents-to-be-predicted, saves memory
parenthaps<-sort(c(paste0(parents,"_HapA"),
                   paste0(parents,"_HapB")))
haploMat<-haploMat[parenthaps,colnames(recombFreqMat)]; 

# Marker effects - one REPEAT-FOLD only
markEffsAD<-readRDS(here::here("output","markerEffectsAD.rds"))
markEffsAD %<>% 
  filter(Repeat=="Repeat1",Fold=="Fold1") %>%   
  mutate(effects=map(effects,~slice(.,1:2)))
modelType="AD"

# Formatting and tidying inputs for predCrossVar
predvars<-markEffsAD %>%
    unnest(effects) %>% 
    unnest(modelOut) %>% 
    filter(Dataset=="trainset") %>% 
    dplyr::select(Repeat,Fold,Trait,modelType,
                  any_of(c("addsnpeff","domsnpeff"))) %>% 
    nest(EffectList=c(Trait,any_of(c("addsnpeff","domsnpeff")))) %>% 
    mutate(AddEffectList=map(EffectList,
                             function(EffectList){
                               addsnpeff<-map(EffectList$addsnpeff,~t(.))
                               names(addsnpeff)<-EffectList$Trait
                               return(addsnpeff)}))
  if(modelType %in% c("AD")){
    predvars<-predvars %>% 
      mutate(DomEffectList=map(EffectList,
                               function(EffectList){
                                 domsnpeff<-map(EffectList$domsnpeff,~t(.))
                                 names(domsnpeff)<-EffectList$Trait
                                 return(domsnpeff) })) }

predvars %<>%
  left_join(parentfolds %>% 
              dplyr::select(-testparents,-trainset,-testset)) %>% 
  dplyr::select(-EffectList)
predvars$AddEffectList[[1]] %<>% map(.,~.[,snps2keep,drop=F])
predvars$DomEffectList[[1]] %<>% map(.,~.[,snps2keep,drop=F])

# runMtCrossVarPredsAD
CrossesToPredict<-predvars$CrossesToPredict[[1]]
AddEffectList<-predvars$AddEffectList[[1]]; 
DomEffectList<-predvars$DomEffectList[[1]]
predType="VPM"

```

## Function

Export these functions to `code/predCrossVar.R`. 

Complete the function by allowing user choice between modelType = "A" and "B"

I want to test these, and let peer review on the MS complete before including in **predCrossVar**.

```{r}
predCrossVar<-function(CrossesToPredict,AddEffectList,DomEffectList,
                       haploMat,recombFreqMat,ncores=1,predType="VPM"){
  starttime<-proc.time()[3]
  # Center posterior distribution of effects
  ## on posterior mean across MCMC samples
  AddEffectList<-purrr::map(AddEffectList,~scale(.,center = T, scale = F));
  DomEffectList<-purrr::map(DomEffectList,~scale(.,center = T, scale = F))
  
  ## Extract the posterior mean effects vectors
  ## If predType="VPM" this just recovers the original effects
  postMeanAddEffects<-purrr::map(AddEffectList,~attr(.,which = "scaled:center"))
  postMeanDomEffects<-purrr::map(DomEffectList,~attr(.,which = "scaled:center"))
  
  parents<-union(CrossesToPredict$sireID,
                 CrossesToPredict$damID)
  haploMat<-haploMat[sort(c(paste0(parents,"_HapA"),paste0(parents,"_HapB"))),]
  
  if(predType!="PMV"){ AddEffectList<-NULL; DomEffectList<-NULL; }
  
  # Set-up a loop over the crosses
  require(furrr); options(mc.cores=ncores); plan(multicore)
  #require(furrr); options(mc.cores=ncores); plan(multisession)
  crossespredicted<-CrossesToPredict %>% 
    mutate(predVars=future_pmap(.,
                                predOneCrossAD,
                                haploMat=haploMat,
                                recombFreqMat=recombFreqMat,
                                predType=predType,
                                postMeanAddEffects=postMeanAddEffects,
                                postMeanDomEffects=postMeanDomEffects)) %>% 
    unnest(predVars)
  totcomputetime<-proc.time()[3]-starttime
  print(paste0("Done predicting fam vars. ",
               "Took ",round((totcomputetime)/60,2),
               " mins for ",nrow(crossespredicted)," crosses"))
  return(crossespredicted) 
}

predOneCrossVarAD<-function(Trait1,Trait2,progenyLD,
                            haploMat,recombFreqMat,predType,
                            postMeanAddEffects,postMeanDomEffects,
                            postVarCovarOfAddEffects=NULL,postVarCovarOfDomEffects=NULL,...){
  ## Predicts variance for one cross and
  ### one variance or covariance paramater (Trait1-Trait2)
  
  ## Predict cross additive variance
  #### posterior mean (VPM)
  quadform<-function(D,x,y){ return(as.numeric(colSums(x*(D%*%y)))) } # does x %*% D %*% t(x) faster
  predVarA_vpm<-quadform(D=progenyLD,
                         x=postMeanAddEffects[[Trait1]],
                         y=postMeanAddEffects[[Trait2]])
  
  #### posterior mean (co)variance (PMV)
  if(predType=="PMV"){ predVarA_pmv<-predVarA_vpm+sum(diag(progenyLD%*%postVarCovarOfAddEffects)) }
  
  ## Predict cross dominance variance
  #### VPM
  progenyLDsq<-progenyLD*progenyLD
  predVarD_vpm<-quadform(D=progenyLDsq,
                         x=postMeanAddEffects[[Trait1]],
                         y=postMeanAddEffects[[Trait2]])
  #### PMV
  if(predType=="PMV"){ predVarD_pmv<-predVarD_vpm+sum(diag(progenyLDsq%*%postVarCovarOfDomEffects)) }
  
  # Tidy the results
  out<-bind_rows(tibble(predOf="VarA",predVar=predVarA_vpm),
                 tibble(predOf="VarD",predVar=predVarD_vpm)) %>% 
    mutate(predMethod="VPM")
  if(predType=="PMV"){
    out<-out %>% 
      bind_rows(bind_rows(tibble(VarComp="VarA",predVar=predVarA_pmv),
                          tibble(VarComp="VarD",predVar=predVarD_pmv)) %>% 
                  mutate(predMethod="PMV")) }
  return(out)
}

predOneCrossAD<-function(sireID,damID,
                         haploMat,recombFreqMat,
                         predType,
                         postMeanAddEffects,postMeanDomEffects,...){
  starttime<-proc.time()[3]
  # Before predicting variances
  # check for and remove SNPs that
  # won't segregate, i.e. are fixed in parents
  ### hopes to save time / mem
  x<-colSums(rbind(haploMat[grep(paste0("^",sireID,"_"),rownames(haploMat)),],
                   haploMat[grep(paste0("^",damID,"_"),rownames(haploMat)),]))
  segsnps2keep<-names(x[x>0 & x<4])
  
  if(length(segsnps2keep)>0){
    recombFreqMat<-recombFreqMat[segsnps2keep,segsnps2keep,drop=F]
    haploMat<-haploMat[,segsnps2keep,drop=F]
    postMeanAddEffects<-purrr::map(postMeanAddEffects,~.[segsnps2keep])
    postMeanDomEffects<-purrr::map(postMeanDomEffects,~.[segsnps2keep])
    if(predType=="PMV"){
      postVarCovarOfAddEffects<-postVarCovarOfAddEffects[segsnps2keep,
                                                         segsnps2keep,drop=F]
      postVarCovarOfDomEffects<-postVarCovarOfDomEffects[segsnps2keep,
                                                         segsnps2keep,drop=F]
    }
    # calc cross LD matrix
    progenyLD<-calcCrossLD(sireID,damID,recombFreqMat,haploMat)
    
    # Set-up loop over variance and covarance parameters
    ## Trait variances to-be-predicted
    traits<-names(postMeanAddEffects)
    varcovars<-tibble::tibble(Trait1=traits,
                              Trait2=traits) 
    ## If multiple traits
    if(length(traits)>1){
      ## Add covariances to-be-predicted
      varcovars<-dplyr::bind_rows(varcovars, # trait variances
                                  combn(traits,2,simplify = T) %>% # covariances
                                    t(.) %>% #
                                    `colnames<-`(.,c("Trait1","Trait2")) %>%
                                    tibble::as_tibble(.)) }
    
    varcovars<-varcovars %>%
      dplyr::mutate(predVars=purrr::pmap(.,
                                         predOneCrossVarAD,
                                         progenyLD=progenyLD,
                                         haploMat=haploMat,
                                         recombFreqMat=recombFreqMat,
                                         predType=predType,
                                         postMeanAddEffects=postMeanAddEffects,
                                         postMeanDomEffects=postMeanDomEffects,
                                         postVarCovarOfAddEffects=postVarCovarOfAddEffects,
                                         postVarCovarOfDomEffects=postVarCovarOfDomEffects)) %>% 
      unnest(predVars)
    computetime<-proc.time()[3]-starttime
    out_thiscross<-tibble(Nsegsnps=length(segsnps2keep),
                          ComputeTime=computetime,
                          predVars=list(varcovars)) 
  } else { 
    out_thiscross<-tibble(Nsegsnps=length(segsnps2keep),
                          ComputeTime=computetime,
                          predVars=list()) }
  return(out_thiscross)
}

calcGameticLD<-function(parentGID,recombFreqMat,haploMat){
  X<-haploMat[paste0(parentGID,c("_HapA","_HapB")),,drop=F]
  p<-colMeans(X); 
  D<-recombFreqMat*((0.5*t(X)%*%X)-p%*%t(p));
  return(D) } 

calcCrossLD<-function(sireID,damID,recombFreqMat,haploMat){
  return(calcGameticLD(sireID,recombFreqMat,haploMat)+
           calcGameticLD(damID,recombFreqMat,haploMat)) }

```

## Multi-vs-single (Ncore vs. BLAS)

```{r}
microbenchmark(ompthreads1=predCrossVar(CrossesToPredict,
                                        AddEffectList,
                                        DomEffectList,
                                        haploMat,recombFreqMat,
                                        ncores=20,predType="VPM"),
               times=5)
# Unit: seconds
#         expr      min       lq     mean   median      uq      max neval
#  ompthreads1 3.307141 3.498848 3.913035 3.659239 4.09278 5.007166     5

microbenchmark(ncore20=predCrossVar(CrossesToPredict,
                                        AddEffectList,
                                        DomEffectList,
                                        haploMat,recombFreqMat,
                                        ncores=20,predType="VPM"),
               ncore1=predCrossVar(CrossesToPredict,
                                        AddEffectList,
                                        DomEffectList,
                                        haploMat,recombFreqMat,
                                        ncores=1,predType="VPM"),
               times=5)
# Unit: seconds
#     expr       min        lq      mean    median        uq      max neval cld
#  ncore20  3.827271  3.941241  4.392303  4.503967  4.652848  5.03619     5  a
#   ncore1 19.861448 19.876688 20.472155 20.498971 20.825927 21.29774     5   b

microbenchmark(ompthreads=predCrossVar(CrossesToPredict,
                                        AddEffectList,
                                        DomEffectList,
                                        haploMat,recombFreqMat,
                                        ncores=20,predType="VPM"),
               times=5)
# export OMP_NUM_THREADS=5
# Unit: seconds
#         expr      min       lq     mean   median       uq      max neval
#  ompthreads5 1.319202 1.364892 1.660863 1.396969 1.442215 2.781038     5

# export OMP_NUM_THREADS=10
# Unit: seconds
#        expr      min       lq    mean   median       uq      max neval
#  ompthreads10 1.688315 1.693673 1.78982 1.732201 1.744468 2.090443     5

# export OMP_NUM_THREADS=56
# Unit: seconds
#         expr     min       lq     mean   median       uq      max neval
#  ompthreads56 4.85575 4.884644 5.021363 4.927789 4.968646 5.469988     5

# export OMP_NUM_THREADS=10
microbenchmark(ompthreads=predCrossVar(CrossesToPredict,
                                        AddEffectList,
                                        DomEffectList,
                                        haploMat,recombFreqMat,
                                        ncores=10,predType="VPM"),
               times=20)
# Unit: seconds
#        expr      min       lq     mean   median       uq      max neval
#  ompthreads 1.683265 1.704042 1.759613 1.719433 1.780987 1.910336     5
 #       expr      min       lq     mean   median       uq      max neval
 # ompthreads 1.738217 1.769593 1.824076 1.828887 1.862609 1.918279    20

# export OMP_NUM_THREADS=56
microbenchmark(ompthreads=predCrossVar(CrossesToPredict,
                                        AddEffectList,
                                        DomEffectList,
                                        haploMat,recombFreqMat,
                                        ncores=2,predType="VPM"),
               times=5)
# Unit: seconds
#        expr      min       lq     mean   median       uq      max neval
#  ompthreads 8.319719 8.343258 8.446788 8.386593 8.581757 8.602615     5
### OUCH :( X_X

# export OMP_NUM_THREADS=56
microbenchmark(ompthreads=predCrossVar(CrossesToPredict,
                                        AddEffectList,
                                        DomEffectList,
                                        haploMat,recombFreqMat,
                                        ncores=1,predType="VPM"),
               times=5)
# Unit: seconds
#        expr     min       lq     mean   median       uq      max neval
#  ompthreads 11.0073 11.33078 11.38762 11.34614 11.36676 11.88711     5
# export OMP_NUM_THREADS=56
microbenchmark(ompthreads=predCrossVar(CrossesToPredict,
                                        AddEffectList,
                                        DomEffectList,
                                        haploMat,recombFreqMat,
                                        ncores=10,predType="VPM"),
               times=5)
# Unit: seconds
#        expr      min       lq     mean   median       uq      max neval
#  ompthreads 5.056288 5.086857 5.131601 5.153773 5.160996 5.200093     5
microbenchmark(ompthreads=predCrossVar(CrossesToPredict,
                                        AddEffectList,
                                        DomEffectList,
                                        haploMat,recombFreqMat,
                                        ncores=20,predType="VPM"),
               times=5)
# Unit: seconds
#        expr      min       lq     mean  median      uq      max neval
#  ompthreads 4.862031 5.155379 5.128031 5.15584 5.19865 5.268255     5
```

## Conclusion

Best performance when balancing ncores (parallel across families) with OMP_NTHREADS. 

WORST performance using OMP_NTHREADS=56!

New function is decisively faster than old function for cross-variance prediction... at least for predType="VPM".

# Final? debugging 

```{bash}
cd /home/jj332_cas/marnin/implementGMSinCassava/; 
export OMP_NUM_THREADS=5 # <-- for a 112 core machine. Use ncores=20 below
screen; 
R # initiate R session
```

```{r inputs for predictCrosses}
require(tidyverse); require(magrittr); 
parentfolds<-readRDS(file=here::here("output","parentfolds.rds"))
parents<-parentfolds %>% 
  select(Repeat,Fold,CrossesToPredict) %>% 
  unnest(CrossesToPredict) %>% 
  distinct(sireID,damID) %$%
  union(sireID,damID)

# Recomb frequency matrix
recombFreqMat<-readRDS(file=here::here("data",
                                       "recombFreqMat_1minus2c_2021May13.rds"))
# Haplotype Matrix
haploMat<-readRDS(file=here::here("data","haps_IITA_filtered_2021May13.rds"))
## keep only haplos for parents-to-be-predicted, saves memory
parenthaps<-sort(c(paste0(parents,"_HapA"),
                   paste0(parents,"_HapB")))
haploMat<-haploMat[parenthaps,colnames(recombFreqMat)]; 
markEffsAD<-readRDS(here::here("output","markerEffectsAD.rds"))

```
```{r debug params}
ncores<-1
modelType="AD"
snpeffs<-markEffsAD
parentfolds=parentfolds
haploMat=haploMat
recombFreqMat=recombFreqMat
```

```{r}
# crossValPredsAD<-predictCrossVars(modelType="AD",ncores=10,
#                                   snpeffs=markEffsAD,parentfolds=parentfolds,
#                                   haploMat=haploMat,recombFreqMat=recombFreqMat)
```
```{r}
  predvars<-snpeffs %>%
    unnest(effects) %>%
    unnest(modelOut) %>%
    filter(Dataset=="trainset") %>%
    dplyr::select(Repeat,Fold,Trait,modelType,
                  any_of(c("addsnpeff","domsnpeff"))) %>%
    nest(EffectList=c(Trait,any_of(c("addsnpeff","domsnpeff")))) %>%
    mutate(AddEffectList=map(EffectList,
                             function(EffectList){
                               addsnpeff<-map(EffectList$addsnpeff,~t(.))
                               names(addsnpeff)<-EffectList$Trait
                               return(addsnpeff)}))
  if(modelType %in% c("AD")){
    predvars<-predvars %>%
      mutate(DomEffectList=map(EffectList,
                               function(EffectList){
                                 domsnpeff<-map(EffectList$domsnpeff,~t(.))
                                 names(domsnpeff)<-EffectList$Trait
                                 return(domsnpeff) })) }

  predvars %<>%
    left_join(parentfolds %>%
                dplyr::select(-testparents,-trainset,-testset)) %>%
    dplyr::select(-EffectList)

  require(furrr); options(future.globals.maxSize=40000*1024^2)

  if(modelType=="A"){
    predvars<-predvars %>%
      mutate(predVars=map2(CrossesToPredict,AddEffectList,
                           ~predCrossVar(CrossesToPredict=.x,
                                         AddEffectList=.y,
                                         modelType=modelType,
                                         haploMat=haploMat,
                                         recombFreqMat=recombFreqMat,
                                         ncores=ncores))) }
  if(modelType=="AD"){
    predvars<-predvars %>%
      mutate(predVars=pmap(.,function(CrossesToPredict,
                                      AddEffectList,DomEffectList,...){
        out<-predCrossVar(CrossesToPredict=CrossesToPredict,
                            AddEffectList=AddEffectList,
                            DomEffectList=DomEffectList,
                            modelType=modelType,
                            haploMat=haploMat,
                            recombFreqMat=recombFreqMat,
                            ncores=ncores)
        return(out) })) }
#   return(predvars)
# }

```
```{r debug inputs for predCrossVar}
modelType=="AD"
CrossesToPredict<-predvars$CrossesToPredict[[1]]
AddEffectList<-predvars$AddEffectList[[1]]; 
DomEffectList<-predvars$DomEffectList[[1]]
predType="VPM"
ncores=1
```

```{r}
  AddEffectList<-purrr::map(AddEffectList,~scale(.,center = T, scale = F));
  if(modelType=="AD"){
    DomEffectList<-purrr::map(DomEffectList,~scale(.,center = T, scale = F)) }

  ## Extract the posterior mean effects vectors
  ## If predType="VPM" this just recovers the original effects
  postMeanAddEffects<-purrr::map(AddEffectList,~attr(.,which = "scaled:center"))
  if(modelType=="AD"){
    postMeanDomEffects<-purrr::map(DomEffectList,~attr(.,which = "scaled:center"))
  } else {
    postMeanDomEffects<-NULL
    }
  parents<-union(CrossesToPredict$sireID,
                 CrossesToPredict$damID)
  haploMat<-haploMat[sort(c(paste0(parents,"_HapA"),paste0(parents,"_HapB"))),]

  if(predType!="PMV"){
    AddEffectList<-NULL;
    if(predType!="PMV" & modelType=="AD"){ DomEffectList<-NULL; } }
```

```{r debug predOneCross}
# predOneCross
sireID<-CrossesToPredict$sireID[[1]]
damID<-CrossesToPredict$damID[[1]]


# predOneCrossVar
Trait1<-varcovars$Trait1[[1]]
Trait2<-varcovars$Trait2[[1]]
progenyLD=progenyLD

```

# An extra (final?) local debugging test

### New predCrossVar functions

```{bash}
cd /home/jj332_cas/marnin/implementGMSinCassava/; 
export OMP_NUM_THREADS=5 # <-- for a 112 core machine. Use ncores=20 below
screen; 
R # initiate R session
```

```{r small inputs}
require(tidyverse); require(magrittr); 
parentfolds<-readRDS(file=here::here("output","parentfolds.rds")) %>% 
  filter(Repeat=="Repeat1",Fold=="Fold1") %>%   
  mutate(CrossesToPredict=map(CrossesToPredict,~slice(.,1:10)))

parents<-parentfolds %>% 
  select(Repeat,Fold,CrossesToPredict) %>% 
  unnest(CrossesToPredict) %>% 
  distinct(sireID,damID) %$%
  union(sireID,damID)

# Recomb frequency matrix
recombFreqMat<-readRDS(file=here::here("data",
                                       "recombFreqMat_1minus2c_2021May13.rds"))
# sample snps 2 keep - ONLY 1000
set.seed(1234)
snps2keep<-sample(colnames(recombFreqMat),size = 1000,replace = F)
recombFreqMat %<>% .[snps2keep,snps2keep]

# Haplotype Matrix
haploMat<-readRDS(file=here::here("data","haps_IITA_filtered_2021May13.rds"))
## keep only haplos for parents-to-be-predicted, saves memory
parenthaps<-sort(c(paste0(parents,"_HapA"),
                   paste0(parents,"_HapB")))
haploMat<-haploMat[parenthaps,colnames(recombFreqMat)]; 

# NEW predCrossVar function
source(here::here("code","predCrossVar.R"))


source(here::here("code","parentWiseCrossVal.R"))

# Marker effects - one REPEAT-FOLD only
markEffsAD<-readRDS(here::here("output","markerEffectsAD.rds"))
markEffsAD %<>% 
  filter(Repeat=="Repeat1",Fold=="Fold1") %>%   
  mutate(effects=map(effects,~slice(.,1:4)))
modelType="AD"

markEffsAD

  mutate(effects=map(effects,~mutate(.data = effects,
                                     modelOut=map(modelOut,function(modelOut,snps2keep){
                                       out<-modelOut
                                       out[[1]]$addsnpeff[[1]]<-out[[1]]$addsnpeff[[1]][snps2keep,,drop=F]
                                       out[[1]]$domsnpeff[[1]]<-out[[1]]$domsnpeff[[1]][snps2keep,,drop=F]
                                       return(out)
                                     },snps2keep=snps2keep))))
markEffsAD$effects[[1]]$modelOut %<>% map(.,~.$addsnpeff[[1]][snps2keep,,drop=F])
markEffsAD$effects[[1]]
#%>% map(modelOut,)



starttime<-proc.time()
crossValPredsAD<-predictCrossVars(modelType="AD",ncores=20,
                                  snpeffs=markEffsAD,parentfolds=parentfolds,
                                  haploMat=haploMat,recombFreqMat=recombFreqMat)
saveRDS(crossValPredsAD,here::here("output","crossValPredsAD.rds"))
timeelapsed<-proc.time()[3] - starttime; 
print(paste0("Elapsed: ",timeelapsed/60/60," hrs"))

```



