# ------------------------------------------------------------------------------
# Program: Saturated_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 Saturated model with sex differences to estimate means and (co)variances across multiple groups
# 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, for the means and (co)variances. This might already give an indication for sex differences
round(colMeans(mzfData,na.rm=TRUE),4)
round(colMeans(dzfData,na.rm=TRUE),4)
round(colMeans(mzmData,na.rm=TRUE),4)
round(colMeans(dzmData,na.rm=TRUE),4)

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

# Set Starting Values
svMe <- c(27,21) # start value for means
svVa <- c(.5,.8) # start value for variances

# Create Labels
#Means, per group, per zygosity and overall mean
labMeMZf <- labVars("meanMZf",selVars)
labMeDZf <- labVars("meanDZf",selVars)
labMeMZm <- labVars("meanMZm",selVars)
labMeDZm <- labVars("meanDZm",selVars)

labMeMZ <- labVars("meanMZ",selVars)
labMeDZ <- labVars("meanDZ",selVars)

labMeZ <- labVars("meanZ",selVars)

#Same for the labels for the covariance and variance
labCvMZf <- labLower("covMZf",ntv)
labCvDZf <- labLower("covDZf",ntv)
labCvMZm <- labLower("covMZm",ntv)
labCvDZm <- labLower("covDZm",ntv)

labCvMZ <- labLower("covMZ",ntv)
labCvDZ <- labLower("covDZ",ntv)
labCvZ <- labLower("covZ",ntv)

labVaMZf <- labDiag("covMZf",ntv)
labVaDZf <- labDiag("covDZf",ntv)
labVaMZm <- labDiag("covMZm",ntv)
labVaDZm <- labDiag("covDZm",ntv)

labVaMZ <- labDiag("covMZ",ntv)
labVaDZ <- labDiag("covDZ",ntv)

labVaZ <- labDiag("covZ",ntv)

# ----------------------------------------------------------------------------------------------------------------------
# PREPARE MODEL
# Create Algebra for expected Mean Matrices
meanMZf <- mxMatrix( type="Full", nrow=1, ncol=ntv, free=TRUE, values=svMe, labels=labMeMZf, name="meanMZf" )
meanDZf <- mxMatrix( type="Full", nrow=1, ncol=ntv, free=TRUE, values=svMe, labels=labMeDZf, name="meanDZf" )
meanMZm <- mxMatrix( type="Full", nrow=1, ncol=ntv, free=TRUE, values=svMe, labels=labMeMZm, name="meanMZm" )
meanDZm <- mxMatrix( type="Full", nrow=1, ncol=ntv, free=TRUE, values=svMe, labels=labMeDZm, name="meanDZm" )

# Create Algebra for expected Variance/Covariance Matrices
covMZf <- mxMatrix( type="Symm", nrow=ntv, ncol=ntv, free=TRUE, values=valDiag(svVa,ntv), 
                   labels=labCvMZf, name="covMZf" )
covDZf <- mxMatrix( type="Symm", nrow=ntv, ncol=ntv, free=TRUE, values=valDiag(svVa,ntv), 
                   labels=labCvDZf, name="covDZf" )
covMZm <- mxMatrix( type="Symm", nrow=ntv, ncol=ntv, free=TRUE, values=valDiag(svVa,ntv), 
                   labels=labCvMZm, name="covMZm" )
covDZm <- mxMatrix( type="Symm", nrow=ntv, ncol=ntv, free=TRUE, values=valDiag(svVa,ntv),
                   labels=labCvDZm, name="covDZm" )

# Create correlation matrices cov2cor
corMZf  <- mxAlgebra(cov2cor(covMZf), name="corMZf")
corDZf  <- mxAlgebra(cov2cor(covDZf), name="corDZf")
corMZm  <- mxAlgebra(cov2cor(covMZm), name="corMZm")
corDZm  <- mxAlgebra(cov2cor(covDZm), name="corDZm")

# 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
expMZf <- mxExpectationNormal( covariance="covMZf", means="meanMZf", dimnames=selVars )
expDZf <- mxExpectationNormal( covariance="covDZf", means="meanDZf", dimnames=selVars )
expMZm <- mxExpectationNormal( covariance="covMZm", means="meanMZm", dimnames=selVars )
expDZm <- mxExpectationNormal( covariance="covDZm", means="meanDZm", dimnames=selVars )
funML <- mxFitFunctionML()

# Create Model Objects for Multiple Groups
modelMZf <- mxModel( meanMZf, covMZf, corMZf, dataMZf, expMZf, funML, name="MZf" )
modelDZf <- mxModel( meanDZf, covDZf, corDZf, dataDZf, expDZf, funML, name="DZf" )
modelMZm <- mxModel( meanMZm, covMZm, corMZm, dataMZm, expMZm, funML, name="MZm" )
modelDZm <- mxModel( meanDZm, covDZm, corDZm, dataDZm, expDZm, funML, name="DZm" )
multi <- mxFitFunctionMultigroup( c("MZf","DZf","MZm","DZm") )

# Create Confidence Interval Objects
ciMean <- mxCI( c('MZf.meanMZf','DZf.meanDZf','MZm.meanMZm','DZm.meanDZm') )
ciCor <- mxCI( c("MZf.corMZf", "MZm.corMZm","DZf.corDZf","DZm.corDZm"))

# Build Saturated Model with Confidence Intervals
modelSAT <- mxModel( "twoSATc", modelMZf, modelDZf,modelMZm, modelDZm, multi, ciCor, ciMean )

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

# Run Saturated Model
fitSAT <- mxRun( modelSAT, intervals=F )
(sumSAT <- summary( fitSAT ))

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

#Print correlation matrices
fitSAT$output$algebras$MZf.corMZf
fitSAT$output$algebras$DZf.corDZf
fitSAT$output$algebras$MZm.corMZm
fitSAT$output$algebras$DZm.corDZm

#------------------------------------------------------------------------------
#Test assumptions
#Test means
#Labels
mMZf <- c('mMZf_SAT','mMZf_oth')
mDZf <- c('mDZf_SAT','mDZf_oth')
mMZm <- c('mMZm_SAT','mMZm_oth')
mDZm <- c('mDZm_SAT','mDZm_oth')

# Constrain expected Means to be equal across twin order
modelEMO  <- mxModel( fitSAT, name="eqMeansTwin" )
modelEMO  <- omxSetParameters( modelEMO, label=labMeMZf, free=TRUE, values=svMe, newlabels= mMZf )
modelEMO  <- omxSetParameters( modelEMO, label=labMeDZf, free=TRUE, values=svMe, newlabels= mDZf )
modelEMO  <- omxSetParameters( modelEMO, label=labMeMZm, free=TRUE, values=svMe, newlabels= mMZm )
modelEMO  <- omxSetParameters( modelEMO, label=labMeDZm, free=TRUE, values=svMe, newlabels= mDZm )
fitEMO    <- mxRun( modelEMO, intervals=F )
mxCompare( fitSAT, fitEMO) 
summary(fitEMO)

# Constrain expected Means to be equal across twin order and zygosity
mZf <- c('mZf_SAT','mZf_oth')
mZm <- c('mZm_SAT','mZm_oth')

modelEMZ <- mxModel( fitEMO, name="eqMZyg" )
modelEMZ <- omxSetParameters( modelEMZ, labels=mMZf, free=TRUE, values=svMe,  newlabels=mZf )
modelEMZ <- omxSetParameters( modelEMZ, labels=mDZf, free=TRUE, values=svMe,  newlabels=mZf )
modelEMZ <- omxSetParameters( modelEMZ, labels=mMZm, free=TRUE, values=svMe,  newlabels=mZm )
modelEMZ <- omxSetParameters( modelEMZ, labels=mDZm, free=TRUE, values=svMe,  newlabels=mZm )

fitEMZ   <- mxRun( modelEMZ, intervals=F )
mxCompare( fitSAT, subs <- list(fitEMO, fitEMZ) )

summary(fitEMZ)

# Constrain expected Means to be equal across twin order and zygosity and  sex
mZ <- c('mZ_SAT','mZ_oth')

modelEMs <- mxModel( fitEMZ, name="eqMsex" )
modelEMs <- omxSetParameters( modelEMs, labels=c("mZf_SAT", "mZm_SAT"), free=TRUE, values=svMe,  newlabels='mZ_SAT' )
modelEMs <- omxSetParameters( modelEMs, labels=c("mZf_oth", "mZm_oth"), free=TRUE, values=svMe,  newlabels='mZ_len' )

modelEMs <- omxAssignFirstParameters(modelEMs)
fitEMs   <- mxRun( modelEMs, intervals=F )
mxCompare(fitEMZ, fitEMs)
mxCompare( fit, subs <- list( fitEMO, fitEMZ, fitEMs) )

summary(fitEMs)

#---------------------------------------------------------------------------------------------------------
#Test difference between female and male MZ covariance
#Constrain the covariance matrices to be equal in MZ males and MZ females
modelCormz  <- mxModel( fitEMZ, name="eqCorrSexmz" )
modelCormz  <- omxSetParameters( modelCormz, label=labCvMZf, free=TRUE, values=svVa, newlabels=labCvMZ) 
modelCormz  <- omxSetParameters( modelCormz, label=labCvMZm, free=TRUE, values=svVa, newlabels=labCvMZ)
fitcormz    <- mxRun( modelCormz, intervals=F )

#Compare this model to the saturated model.
mxCompare( fitEMZ, fitcormz )

#What are the twin correlations constrained?
cov2cor(fitcormz$output$algebras$MZf.covMZf)
cov2cor(fitcormz$output$algebras$MZm.covMZm)

#Test difference between female and male DZ covariance
#Constrain the covariance matrices to be equal in DZ males and DZ females
modelCordz  <- mxModel( fitEMZ, name="eqCorrSexdz" )
modelCordz  <- omxSetParameters( modelCordz, label=labCvDZf, free=TRUE, values=svVa, newlabels=labCvDZ) 
modelCordz  <- omxSetParameters( modelCordz, label=labCvDZm, free=TRUE, values=svVa, newlabels=labCvDZ)
fitcordz    <- mxRun( modelCordz, intervals=F )
mxCompare( fitEMZ, fitcordz )

cov2cor(fitcordz$output$algebras$DZf.covDZf)
cov2cor(fitcordz$output$algebras$DZm.covDZm)
