# ------------------------------------------------------------------------------
# Program: ACE_bivariate_4group 
#  Author: Lianne de Vries, adapted from different scripts of Hermine Maes, Lucia Colodro-Conde and Eveline de Zeeuw
#    Date: 16 07 2020
#
# Twin Bivariate ACE model with sex differences to estimate genetic and environmental contribution to the (co)variance
# Continuous data
# -------|---------|---------|---------|---------|---------|---------|---------|

rm(list=ls(all=TRUE)) # Clear working memory

#set working directory
setwd('C:/Users/lvs205/surfdrive/Shared/Team Well-Being/LIANNE/TutorialPaper_bivariatemodel/Scripts')

# Load Libraries & Options
library(OpenMx)
source("miFunctions.R")

#Load your data
load("bivariate_wide2.Rda")
twinData <- bivariate_wide

#Check data, how many rows and columns?
dim(twinData)
describe(twinData, skew=F)

#Rename some variables, easier to work with
colnames(twinData)[52] <- "SAT_1"
colnames(twinData)[53] <- "SAT_2"
colnames(twinData)[9] <- "cito_1"
colnames(twinData)[10] <- "cito_2"
colnames(twinData)[3] <- "aggr_1"
colnames(twinData)[4] <- "aggr_2"
colnames(twinData)[5] <- "anx_1"
colnames(twinData)[6] <- "anx_2"
colnames(twinData)[44] <- "opt3_1"
colnames(twinData)[45] <- "opt3_2"

# Select Variables for Analysis
vars <- c('SAT','anx') # list of variables names
nv <- 2 # number of variables
ntv <- nv*2 # number of total variables
selVars <- paste(vars,c(rep(1,nv),rep(2,nv)),sep="_")

# Select Data for Analysis based on the zygosity groups
mzfData   <- subset(twinData, (zyg5.1==3 | zyg5.2==3), c(selVars))
dzfData   <- subset(twinData, (zyg5.1==4 | zyg5.2==4), c(selVars))
mzmData   <- subset(twinData, (zyg5.1==1 | zyg5.2==1), c(selVars))
dzmData   <- subset(twinData, (zyg5.1==2 | zyg5.2==2), c(selVars))

# Generate Descriptive Statistics
round(colMeans(mzmData,na.rm=TRUE),3) 
round(colMeans(dzmData,na.rm=TRUE),3) 
round(colMeans(mzfData,na.rm=TRUE),3)
round(colMeans(dzfData,na.rm=TRUE),3) 

round(cov(mzmData,use="complete"),3) 
round(cov(dzmData,use="complete"),3)
round(cov(mzfData,use="complete"),3) 
round(cov(dzfData,use="complete"),3)

######################################################################
#The ACE model starts here

# Set Starting Values 
svMe      <- c(27.6, 20)                     # start value for means
svPa      <- .4                    # start value for variance component A
svPe      <- .5                   # start value for variance component E

# Create Algebra for expected Mean Matrices
meanGf      <- mxMatrix( type="Full", nrow=1, ncol=ntv, free=TRUE, values=svMe, 
                         labels=c("meanSATf","meanoptf", "meanSATf","meanoptf"), name="meanGf" )
meanGm      <- mxMatrix( type="Full", nrow=1, ncol=ntv, free=TRUE, values=svMe, 
                         labels=c("meanSATm","meanoptm", "meanSATm","meanoptm"), name="meanGm" )

## Create Matrices for Variance Components
covAf      <- mxMatrix( type="Symm", nrow=nv, ncol=nv, free=TRUE, values=valDiag(svPa,nv),
                        label=labLower("VAf",nv),name="VAf" ) #lbound=1e-4, 
covCf      <- mxMatrix( type="Symm", nrow=nv, ncol=nv, free=TRUE, values=valDiag(0.1,nv),
                        label=labLower("VCf",nv),name="VCf" )
covEf      <- mxMatrix( type="Symm", nrow=nv, ncol=nv, free=TRUE, values=valDiag(svPe,nv),
                        label=labLower("VEf",nv), name="VEf" )

covAm      <- mxMatrix( type="Symm", nrow=nv, ncol=nv, free=TRUE, values=valDiag(svPa,nv),
                        label=labLower("VAm",nv), name="VAm" ) 
covCm      <- mxMatrix( type="Symm", nrow=nv, ncol=nv, free=TRUE, values=0.1,
                        label=labLower("VCm",nv), name="VCm" )
covEm      <- mxMatrix( type="Symm", nrow=nv, ncol=nv, free=TRUE, values=valDiag(svPe,nv),
                        label=labLower("VEm",nv) ,name="VEm" )

## Create Algebra for expected Variance/Covariance Matrices in MZ & DZ twins
#Total covariance
covPf      <- mxAlgebra( expression= VAf+VCf+VEf, name="Vf" )
covPm      <- mxAlgebra( expression= VAm+VCm+VEm, name="Vm" )

#Covariance in MZ and DZ twins
covMZf     <- mxAlgebra( expression= VAf+VCf, name="cMZf" )
covDZf     <- mxAlgebra( expression= 0.5%x%VAf+ VCf, name="cDZf" )
covMZm     <- mxAlgebra( expression= VAm+VCm, name="cMZm" )
covDZm     <- mxAlgebra( expression= 0.5%x%VAm+ VCm, name="cDZm" )

#Expected covariance matrices
expCovMZf  <- mxAlgebra( expression= rbind( cbind(Vf, cMZf),
                                            cbind(t(cMZf), Vf)), name="expCovMZf" ) #
expCovDZf  <- mxAlgebra( expression= rbind( cbind(Vf, cDZf),
                                            cbind(t(cDZf), Vf)), name="expCovDZf" )
expCovMZm  <- mxAlgebra( expression= rbind( cbind(Vm, cMZm),
                                            cbind(t(cMZm), Vm)), name="expCovMZm" ) #
expCovDZm  <- mxAlgebra( expression= rbind( cbind(Vm, cDZm),
                                            cbind(t(cDZm), Vm)), name="expCovDZm" )

## Create Data Objects for Multiple Groups
dataMZf    <- mxData( observed=mzfData, type="raw" )
dataDZf    <- mxData( observed=dzfData, type="raw" )
dataMZm    <- mxData( observed=mzmData, type="raw" )
dataDZm    <- mxData( observed=dzmData, type="raw" )

## Create Expectation Objects for Multiple Groups
funML     <- mxFitFunctionML() 
expMZf     <- mxExpectationNormal( covariance="expCovMZf", means="meanGf", dimnames=selVars ) 
expDZf     <- mxExpectationNormal( covariance="expCovDZf", means="meanGf", dimnames=selVars )
expMZm     <- mxExpectationNormal( covariance="expCovMZm", means="meanGm", dimnames=selVars ) 
expDZm     <- mxExpectationNormal( covariance="expCovDZm", means="meanGm", dimnames=selVars )

## Create Model Objects for Multiple Groups
pars      <- list(meanGf, meanGm, covAf, covCf, covEf, covPf,covAm, covCm, covEm, covPm)#, 
modelMZf   <- mxModel(pars, covMZf, expCovMZf, dataMZf, expMZf, funML, name="MZf" ) 
modelDZf   <- mxModel(pars, covDZf, expCovDZf, dataDZf, expDZf, funML, name="DZf" )
modelMZm   <- mxModel(pars, covMZm, expCovMZm, dataMZm, expMZm, funML, name="MZm" ) 
modelDZm   <- mxModel(pars, covDZm, expCovDZm, dataDZm, expDZm, funML, name="DZm" )
multi     <- mxFitFunctionMultigroup( c("MZf","DZf","MZm","DZm"))#, "DZo") ) 

## Create Algebra for Standardization 
matIf      <- mxMatrix( type="Iden", nrow=nv, ncol=nv, name="If" )
matIm      <- mxMatrix( type="Iden", nrow=nv, ncol=nv, name="Im" )

## Calculate genetic,  environmental, and phenotypic correlations
corAf      <- mxAlgebra( expression=solve(sqrt(If*VAf))%&%VAf, name ="rAf" ) 
corCf      <- mxAlgebra( expression=solve(sqrt(If*VCf))%&%VCf, name ="rCf" )
corEf      <- mxAlgebra( expression=solve(sqrt(If*VEf))%&%VEf, name ="rEf" )
corPf      <- mxAlgebra( expression=solve(sqrt(If*Vf))%&%Vf, name ="rPf") 

corAm      <- mxAlgebra( expression=solve(sqrt(Im*VAm))%&%VAm, name ="rAm" ) 
corCm      <- mxAlgebra( expression=solve(sqrt(Im*VCm))%&%VCm, name ="rCm" )
corEm      <- mxAlgebra( expression=solve(sqrt(Im*VEm))%&%VEm, name ="rEm" )
corPm      <- mxAlgebra( expression=solve(sqrt(Im*Vm))%&%Vm, name ="rPm") 

## Create Algebra for Variance Components ( formatting output)
#Unstandardized variance components
rowUV     <- rep('UV',nv)
colUV     <- rep(c('VA','VC','VE'),each=nv)
estUVf     <- mxAlgebra( expression=cbind(VAf,VCf,VEf), name="UVf", dimnames=list(rowUV,colUV) )
estUVm     <- mxAlgebra( expression=cbind(VAm,VCm,VEm), name="UVm", dimnames=list(rowUV,colUV) )

#Standardized variance components
rowSV     <- rep('SV',nv)
colSV     <- rep(c('SA','SC','SE'),each=nv)
estSVf     <- mxAlgebra( expression=cbind(VAf/Vf,VCf/Vf,VEf/Vf), name="SVf", dimnames=list(rowSV,colSV) )
estSVm     <- mxAlgebra( expression=cbind(VAm/Vm,VCm/Vm,VEm/Vm), name="SVm", dimnames=list(rowSV,colSV) )

#Correlations
rowcorr     <- rep('corr',nv)
colcorr     <- rep(c('rA','rC','rE','rP'),each=nv)
estcorrf     <- mxAlgebra( expression=cbind(rAf,rCf,rEf,rPf), name="corrf", dimnames=list(rowcorr,colcorr) )
estcorrm     <- mxAlgebra( expression=cbind(rAm,rCm,rEm,rPm), name="corrm", dimnames=list(rowcorr,colcorr) )

## Create Confidence Interval Objects
ciUV      <- mxCI( c("UV") )
ciSV      <- mxCI( c("SVf", "SVm") )
cicorr      <- mxCI( c("corrf", "corrm") )

## Build Model with Confidence Intervals
calc      <- list( matIf,  corAf, corCf, corEf,corPf, estUVf, estSVf,estcorrf,
                   matIm,  corAm, corCm, corEm,corPm, estUVm, estSVm,estcorrm, ciSV, cicorr)
modelACE  <- mxModel( "twoACEvc", pars, modelMZf, modelDZf, modelMZm, modelDZm,multi, calc ) #

# ----------------------------------------------------------------------------------------------------------------------
# RUN MODEL

# Run ACE Model
fitACE    <- mxRun( modelACE, intervals=T)
(sumACE    <- summary( fitACE ) )

# Print Goodness-of-fit Statistics & Parameter Estimates
fitGofs(fitACE)
fitEsts(fitACE)

# Print Covariance & Correlation Matrices
fitACE$algebras$UVf
fitACE$algebras$UVm
fitACE$algebras$SVf
fitACE$algebras$SVm
fitACE$algebras$corrf
fitACE$algebras$corrm

#For females:
#For males, change the SVf to SVm
# prop% phenotypic covariance due to A #
fitACE$SVf$result[1,2]
# And now with its confidence interval (will only work if intervals=T in mxRun)
sumACE$CI[grep(pattern = mxEval(SVf[1,2], fitACE),sumACE$CI$estimate),]
# prop% phenotypic covariance due to C #
fitACE$SVf$result[1,4]
# And now with its confidence interval (will only work if intervals=T in mxRun)
sumACE$CI[grep(pattern = mxEval(SVf[1,4], fitACE),sumACE$CI$estimate),]
# prop% phenotypic covariance due to E #
fitACE$SVf$result[1,6]
# And now with its confidence interval (will only work if intervals=T in mxRun)
sumACE$CI[grep(pattern = mxEval(SVf[1,6], fitACE),sumACE$CI$estimate),]
# It should add up to 1
mxEval(SVf[1,2]+SVf[1,4]+SVf[1,6], fitACE)

# ----------------------------------------------------------------------------------------------------------------------
# RUN SUBMODELS
#Step 1: Sex diff
#Test sex difference in parameter estimates
modelACEs   <- mxModel( fitACE, name="mulAEsex" )
modelACEs   <- omxSetParameters( modelACEs, labels=labLower("VAf",nv), free=T,  newlabels=labLower("VA",nv) )
modelACEs   <- omxSetParameters( modelACEs, labels=labLower("VAm",nv), free=T,  newlabels=labLower("VA",nv) )
modelACEs   <- omxSetParameters( modelACEs, labels=labLower("VCf",nv), free=T, newlabels=labLower("VC",nv) )
modelACEs   <- omxSetParameters( modelACEs, labels=labLower("VCm",nv), free=T,  newlabels=labLower("VC",nv) )
modelACEs   <- omxSetParameters( modelACEs, labels=labLower("VEf",nv), free=T,  newlabels=labLower("VE",nv) )
modelACEs   <- omxSetParameters( modelACEs, labels=labLower("VEm",nv), free=T,  newlabels=labLower("VE",nv) )
modelACEs <- omxAssignFirstParameters(modelACEs)

fitACEs     <- mxTryHard( modelACEs, intervals=F)
mxCompare( fitACE, fitACEs )

#Step 2: AE model. Drop all C components
modelAE   <- mxModel( fitACE, name="twoAEvc" )
modelAE   <- omxSetParameters( modelAE, labels=labLower("VCf",nv), free=FALSE, values=0 )
modelAE   <- omxSetParameters( modelAE, labels=labLower("VCm",nv), free=FALSE, values=0 )
fitAE     <- mxTryHard( modelAE, intervals= T)
mxCompare(fitACE, fitAE)

fitGofs(fitAE); fitEsts(fitAE)
(sumAE <- summary(fitAE))

#Standardized covariance matrices and correlation matrices for females and males
fitAE$algebras$SVf
fitAE$algebras$SVm
fitAE$algebras$corrf
fitAE$algebras$corrm

#Same as above:
#Females
# prop% phenotypic covariance due to A #
sumAE$CI[grep(pattern = mxEval(SVf[1,2], fitAE),sumAE$CI$estimate),]
# prop% phenotypic covariance due to E #
sumAE$CI[grep(pattern = mxEval(SVf[1,6], fitAE),sumAE$CI$estimate),]
# It should add up to 1
mxEval(SVf[1,2]+SVf[1,4]+SVf[1,6], fitAE)

#Males
# prop% phenotypic covariance due to A #
sumAE$CI[grep(pattern = mxEval(SVm[1,2], fitAE),sumAE$CI$estimate),]
# prop% phenotypic covariance due to E #
sumAE$CI[grep(pattern = mxEval(SVm[1,6], fitAE),sumAE$CI$estimate),]

# It should add up to 1
mxEval(SVm[1,2]+SVm[1,6], fitAE)


#################

# Run CE model
modelCE   <- mxModel( fitACE, name="twoCEvc" )
modelCE   <- omxSetParameters( modelCE, labels=labLower("VAf",nv), free=FALSE, values=0 )
modelCE   <- omxSetParameters( modelCE, labels=labLower("VAm",nv), free=FALSE, values=0 )
modelCE   <- omxSetParameters( modelCE, labels=labLower("VCf",nv), free=TRUE, values=.6 )
modelCE   <- omxSetParameters( modelCE, labels=labLower("VCm",nv), free=TRUE, values=.6 )
fitCE     <- mxRun( modelCE, intervals=F )
mxCompare(fitACE,fitCE)
fitGofs(fitCE); fitEsts(fitCE)

# Run E model
modelE    <- mxModel( fitAE, name="twoEvc" )
modelE    <- omxSetParameters( modelE, labels=labLower("VAf",nv), free=FALSE, values=0 )
modelE    <- omxSetParameters( modelE, labels=labLower("VAm",nv), free=FALSE, values=0 )
fitE      <- mxRun( modelE, intervals=F )
mxCompare(fitAE,fitE)
fitGofs(fitE); fitEsts(fitE)

# Print Comparative Fit Statistics
mxCompare( fitACE, nested <- list(fitAE, fitCE, fitE) ) #generate likelihood ratio test

