#-------------------------------------
# Bivariate
#-------------------------------------

########################
# Clear Workspace
rm(list = ls())

########################
# Switch off scientific notation 
options(scipen=999)

########################
# install packages
#install.packages("tidyverse")
#install.packages("OpenMx")
#install.packages("umx")
#install.packages("haven")
#install.packages("lavaan")
#devtools::install_github('mirkoruks/twinflex')

########################
# load packages
library(tidyverse)
library(umx)
library(OpenMx)
library(twinflex)
library(haven)
library(lavaan)

# load functions
source("functions.R") 

# run recode script
source("Session 4/Scripts/recode.R")

View(df_wide)

df_wide <- as.data.frame(df_wide)

########################
# prepare MZ and DZ data frames
########################
mz_df <- df_wide %>% filter(zyg == 1)
dz_df <- df_wide %>% filter(zyg == 2)

# bivariate ACE modeling
################################################################################################

# bivariate ACE model with umx
ACE1_umx <- umxACE(name = "Bivariate ACE 1",
                  selDVs = c("CA","math"),
                  sep = "_",
                  mzData = mz_df,
                  dzData = dz_df,
                  boundDiag = 0.001)

summary(ACE1_umx)
umxSummary(ACE1_umx)
ACE1_umxCI <- umxCI(ACE1_umx, run = "yes") # these are likelihood-based confidence intervals. Sometimes, they have some issues
# try to fix it by dropping boundaries
ACE1a_umx <- umxACE(name = "Bivariate ACE 1",
                   selDVs = c("CA","math"),
                   sep = "_",
                   mzData = mz_df,
                   dzData = dz_df,
                   boundDiag = NULL)
ACE1a_umxCI <- umxCI(ACE1a_umx, run = "yes")

# get genetic correlation from genetic covariance (formula: cor(x,y) = cov(x,y)/(sd(x)*sd(y)))
#0.334767951/(sqrt(0.702606073)*sqrt(0.342539459))

# get p values
makeOutput(summary(ACE1_umx), rounded = TRUE)

# calculate genetic/environmental correlations
umxSummaryACE(ACE1_umx, showRg = TRUE) # strong genetic correlation between cognitive abilities and math, but also strong shared environmental correlation
# however, (to my knowledge) umx does not give you a test statistic for the correlations, since they are not treated as a model parameters
# see the code below to calculate the standard errors for the correlations

# we can calculate the correlations manually:
#
# rG = (a21*a11)/sqrt((a11^2)*((a21^2)+(a22^2)))
#
summary(ACE1_umx)
View(ACE1_umx)
a11 <- ACE1_umx@output[["matrices"]][["top.a"]][1,1]
a21 <- ACE1_umx@output[["matrices"]][["top.a"]][2,1]
a22 <- ACE1_umx@output[["matrices"]][["top.a"]][2,2]
(rg <- (a21*a11)/sqrt((a11^2)*((a21^2)+(a22^2)))) # check! 

# calculate heritabilities by hand 
  # Trait 1: (h2_1 = (a11^2)/(a11^2 + c11^2 + e11^2))
  # Trait 2: (h2_2 = (a21^2 + a22^2)/(a21^2 + a22^2 + c21^2 + c22^2 + e21^2 + e22^2))
  # problem: no inference stats
summary(ACE1_umx)
VarT1 <- 0.838214977**2 + 0.652303595**2 + 0.812551327**2 
VarA1 <- 0.838214977**2
(h2_1 <- VarA1/VarT1)
VarT2 <- 0.399381990**2 + 0.427824758**2 + 0.182499899**2 + 0.427984602**2 + 0.072746870**2 + 0.523645329**2
VarA2 <- 0.399381990**2 + 0.427824758**2
VarA2_1 <- 0.399381990**2
VarA2_2 <- 0.427824758**2
(h2_2 <- VarA2/VarT2)
VarA2_1/VarT2 
VarA2_2/VarT2

# inspect fit object
View(ACE1_umx)

(rp <- 0.5129239 / sqrt(1.7883440 * 0.8385136))
# from the model parameters
(covxy <- 0.838214977*0.399381990 + 0.652303595*0.182499899 + 0.812551327*0.072746870)
(varx <- 0.838214977**2 + 0.652303595**2 + 0.812551327**2)
(vary <- 0.399381990**2 + 0.427824758**2 + 0.182499899**2 + 0.427984602**2 + 0.072746870**2 +0.523645329**2) 
(covxy / sqrt(varx * vary))

# Calculate bivariate heritability
(biv_h2 <- (sqrt(h2_1) * rg * sqrt(h2_2)) / rp)


# I wrote a small function called "exph" (explained heritability). It involves three arguments:
# 1) The fit object of the model for which you want to calculate the heritability
# 2) The position of the outcome 
# 3) If you want to use tryHard
ACE1_umx_h2 <- exph(ACE1_umx, 2, tryHard = FALSE)
ACE1_umx_h2[[2]]
View(ACE1_umx_h2)

# Different parametrization (variances instead of paths)
ACE1v_umx <- umxACEv(name = "Bivariate ACE 1 Variances",
                     selDVs = c("CA","math"),
                     sep = "_",
                     mzData = mz_df,
                     dzData = dz_df,
                     boundDiag = 0.001)
summary(ACE1v_umx)

# calculate rg
(rg_v <- 0.334767951/sqrt(0.702606073*0.342539459))

# bivariate heritability
View(ACE1v_umx)
# expected correlation:
rp_v <- 0.5129238 / sqrt(1.7883439 * 0.8385134) # (see View(ACE1v_umx) -> submodels -> top -> algebras -> Vtot)
h2_1_v <- 0.3928808 # (see View(ACE1v_umx) -> submodels -> top -> algebras -> A_std)
h2_2_v <- 0.4085080 # (see View(ACE1v_umx) -> submodels -> top -> algebras -> A_std)

sqrt(h2_1_v)*rg_v*sqrt(h2_2_v) / rp_v

# Bivariate ACE model with twinflex
ACE2_tf <- twinflex(acevars = c("CA","math"), sep = "_", data = df_wide, lboundACE = TRUE)
summary(ACE2_tf)
makeOutput(summary(ACE2_tf), rounded = TRUE)


# Multivariate models umx
ACE2_umx <- umxACE(name = "Bivariate ACE 2",
                   selDVs = c("selfeff","CA","math"),
                   sep = "_",
                   mzData = mz_df,
                   dzData = dz_df,
                   boundDiag = 0.001)
summary(ACE2_umx)
makeOutput(summary(ACE2_umx), rounded = TRUE)

# Multivariate models twinflex
ACE2_tf <- twinflex(acevars = c("selfeff","CA","math"), sep = "_", data = df_wide, lboundACE = TRUE)
summary(ACE2_tf)
makeOutput(summary(ACE2_tf), rounded = TRUE)

# Inspect heritabilities
ACE2_umx_h2 <- exph(ACE2_umx, 3, tryHard = FALSE)
ACE2_umx_h2[[2]]

# Inspect genetic correlations
umxSummaryACE(ACE2_umx, showRg = TRUE)

# model comparisons with umx (alternative omxSetParameters())

# inspect fit object
View(ACE2_umx)

# drop some paths
makeOutput(summary(ACE2_umx), rounded = TRUE)
ACE2_umx_reduced1 <- umxModify(ACE2_umx, 
                             update = c("c_r1c1","c_r2c1","c_r3c1","e_r2c1"), # specify parameters you want to fix/drop
                             free = FALSE, # default
                             value = 0, # default
                             comparison = TRUE, 
                             name = "Bivariate ACE 2 Reduced")
plot(ACE2_umx)
# see model comparison
umxCompare(ACE2_umx,ACE2_umx_reduced1)
mxCompare(ACE2_umx,ACE2_umx_reduced1) # question of personal preference what to use...
makeOutput(summary(ACE2_umx_reduced1), rounded = TRUE)

# Inspect heritabilities
ACE2_umx_reduced1_h2 <- exph(ACE2_umx_reduced1, 3, tryHard = FALSE)
ACE2_umx_reduced1_h2[[2]]
plot(ACE2_umx_reduced1)
# set equality constraint
makeOutput(summary(ACE2_umx_reduced1), rounded = TRUE)
ACE2_umx_equality <- umxModify(ACE2_umx_reduced1, 
                             update = c("a_r2c2","a_r3c3"),
                             newlabels = c("a3132"),
                             name = "Bivariate ACE 2 Equal")

makeOutput(summary(ACE2_umx_equality), rounded = TRUE) 
mxCompare(ACE2_umx_reduced1, ACE2_umx_equality)

# Tasks:
# 1. At the moment we can explain ~53% of the heritability of math grades with cognitive ability and self-efficacy. Can you 
# increase the proportion of explained heritability of math grades by adding more covariates? Use the function exph to check
# whether the added covariates explain a part of the heritability

# 2. Calculate a multivariate model with different personality items. Estimate Cholesky or a variance components model. 
# Are there any correlations? Do the correlations are due to common genes or due to common (non-)shared environmental influences?
  # does the pattern differ across domains? 
    # nervous, worries, notrelaxed -> neuroticism
    # quiet, talkactive, outgoing -> extraversion
    # notlazy, thorough, efficient -> conscientiousness


# General --> Missing "build-in" functions for inference of second order statistics, e.g. heritability, genetic correlation or bi-
# variate heritability --> use OpenMx to build second order statistics as a model parameter and then use mxSE() to get standard error
########################################################################################################################
########################################################################################################################
# optional!!
# ok... now, what about the test statistic for the correlations?
# this is somewhat more advanced... you don't need to read this. But you can read it if you are interested :-)
# -> we need openMx, the powerful big brother of umx..

# let's check the model fit objects
View(ACE1_umx)

# we need to build an algebra object that defines a new model parameter
# these are our inputs:
ACE1_umx@submodels[["top"]]@matrices[["I"]] # identity matrix
ACE1_umx@submodels[["top"]]@algebras[["A"]] # unstandardized genetic covariance matrix

# this is our algebra object
rA <- mxAlgebra(expression = solve(sqrt(top.I * top.A)) %*% top.A %*% solve(sqrt(top.I * top.A)), name = "rA") # this is just the formula of a correlation in matrix notation (see answer of Brian B: https://math.stackexchange.com/questions/186959/correlation-matrix-from-covariance-matrix/300775)
# you can read a book on matrix algebra or you can check the umx function to see the formula. For the latter go with the cursor on the function "umxSummaryACE" and press control/strg to see the code, in line 84-85 you can see the formula for the genetic correlation
rA <- mxAlgebra(expression = cov2cor(top.A), name = "rA") # this is also an option (less maths :-))
# we define the model with the new parameter
ACE1_umxb <- mxModel(ACE1_umx,rA)
# we run the model with the new parameter  
ACE1_umxb <- umxRun(ACE1_umxb)
# let's see where the parameter is stored
View(ACE1_umxb)
# what is the genetic correlation?
ACE1_umxb@algebras[["rA"]] # tada! the genetic correlations
# are they equal?
umxSummaryACE(ACE1_umx, showRg = TRUE) # yes!
# with the genetic correlation being a (second order) model parameter we can calculate the SE
mxSE("rA",ACE1_umxb) # and here the Wald-based SE! You can use them to calculate CIs 

Est <- ACE1_umxb@algebras[["rA"]]$result[2,1]
SE <- mxSE("rA",ACE1_umxb)[2,1]
# example: 
rAdf <- as.data.frame(cbind(Est,SE))
rAdf <- rAdf %>% 
  mutate(lowerCI = Est - 1.96*SE) %>% 
  mutate(upperCI = Est + 1.96*SE) 
print(rAdf)



