library(mvtnorm)
library(OpenMx)
mxOption(NULL,"Default optimizer","SLSQP")

#True MZ correlation matrix:
Rmz <- matrix(c(
	1.0, 0.5, 0.5, 0.4, 0.2, 0.2,
	0.5, 1.0, 0.5, 0.2, 0.4, 0.2,
	0.5, 0.5, 1.0, 0.2, 0.2, 0.4,
	0.4, 0.2, 0.2, 1.0, 0.5, 0.5,
	0.2, 0.4, 0.2, 0.5, 1.0, 0.5,
	0.2, 0.2, 0.4, 0.5, 0.5, 1.0
), nrow=6, ncol=6)

#True DZ correlation matrix:
Rdz <- matrix(c(
	1.0, 0.5, 0.5, 0.3, 0.15, 0.15,
	0.5, 1.0, 0.5, 0.15, 0.3, 0.15,
	0.5, 0.5, 1.0, 0.15, 0.15, 0.3,
	0.3, 0.15, 0.15, 1.0, 0.5, 0.5,
	0.15, 0.3, 0.15, 0.5, 1.0, 0.5,
	0.15, 0.15, 0.3, 0.5, 0.5, 1.0
), nrow=6, ncol=6)

#Set RNG seed and generate multivariate-normal data:
set.seed(180313)
mzdat <- rmvnorm(n=2000,mean=rep(0,6),sigma=Rmz)
dzdat <- rmvnorm(n=2000,mean=rep(0,6),sigma=Rdz)
colnames(mzdat) <- c("t1y1","t1y2","t1y3","t2y1","t2y2","t2y3")
colnames(dzdat) <- c("t1y1","t1y2","t1y3","t2y1","t2y2","t2y3")
#Trait #1 will have a mean of 10 and a variance of 4:
mzdat[,1] <- (mzdat[,1]*2)+10
mzdat[,4] <- (mzdat[,4]*2)+10
dzdat[,1] <- (dzdat[,1]*2)+10
dzdat[,4] <- (dzdat[,4]*2)+10
#Trait #2 will have a mean of 20 and a variance of 9:
mzdat[,2] <- (mzdat[,2]*3)+20
mzdat[,5] <- (mzdat[,5]*3)+20
dzdat[,2] <- (dzdat[,2]*3)+20
dzdat[,5] <- (dzdat[,5]*3)+20
#Trait #3 will be negative-binomial with size=1 and mu=1:
mzdat[,3] <- qnbinom(pnorm(mzdat[,3]),size=1,mu=1)
mzdat[,6] <- qnbinom(pnorm(mzdat[,6]),size=1,mu=1)
dzdat[,3] <- qnbinom(pnorm(dzdat[,3]),size=1,mu=1)
dzdat[,6] <- qnbinom(pnorm(dzdat[,6]),size=1,mu=1)

#Largest and smallest counts observed on trait #3:
maxct <- max(c(mzdat[,3],mzdat[,6],dzdat[,3],dzdat[,6]),na.rm=T)
minct <- min(c(mzdat[,3],mzdat[,6],dzdat[,3],dzdat[,6]),na.rm=T)

#Datasets actually used will store trait #3 as an ordered factor:
mzdat2 <- data.frame(
	t1y1=mzdat[,1], t1y2=mzdat[,2], t1y3=mxFactor(mzdat[,3],levels=minct:maxct), 
	t2y1=mzdat[,4], t2y2=mzdat[,5], t2y3=mxFactor(mzdat[,6],levels=minct:maxct))
dzdat2 <- data.frame(
	t1y1=dzdat[,1], t1y2=dzdat[,2], t1y3=mxFactor(dzdat[,3],levels=minct:maxct), 
	t2y1=dzdat[,4], t2y2=dzdat[,5], t2y3=mxFactor(dzdat[,6],levels=minct:maxct))

#Start values for some parameters:
sv11 <- var(mzdat[,1],na.rm=T)/3
sv12 <- cov(mzdat[,1],mzdat[,2],use="pair")/3
sv22 <- var(mzdat[,2],na.rm=T)/3
svmu <- mean(mzdat[,3],na.rm=T)
svnu <- (svmu^2)/(var(mzdat[,3],na.rm=T)-svmu)

#Additive-genetic covariance matrix:
A <- mxMatrix(
	type="Symm",nrow=3,ncol=3,free=T,
	values=c(sv11, sv12, 0.1, sv22, 0.1, 0.1),
	labels=c("a11","a21","a31","a22","a32","a33"),
	lbound=c(rep(NA,5),-1),
	ubound=c(rep(NA,5),1),
	name="A")

#Shared-environmental covariance matrix:
C <- mxMatrix(
	type="Symm",nrow=3,ncol=3,free=T,
	values=c(sv11, sv12, 0.1, sv22, 0.1, 0.1),
	labels=c("c11","c21","c31","c22","c32","c33"),
	lbound=c(rep(NA,5),-1),
	ubound=c(rep(NA,5),1),
	name="C")

#Nonshared-environmental covariance matrix:
E <- mxMatrix(
	type="Symm",nrow=3,ncol=3,free=T,
	values=c(sv11, sv12, 0.1, sv22, 0.1, 0.8),
	labels=c("e11","e21","e31","e22","e32","e33"),
	lbound=c(1e-4, NA, NA, 1e-4, NA, 1e-4),
	ubound=c(rep(NA,5),1),
	name="E")

#Vector of Gaussian means:
M <- mxMatrix(
	type="Full",nrow=1,ncol=6,
	free=c(T,T,F,T,T,F),
	values=c(
		mean(c(mzdat[,1],mzdat[,4],dzdat[,1],dzdat[,4])),
		mean(c(mzdat[,2],mzdat[,5],dzdat[,2],dzdat[,5])),
		0,
		mean(c(mzdat[,1],mzdat[,4],dzdat[,1],dzdat[,4])),
		mean(c(mzdat[,2],mzdat[,5],dzdat[,2],dzdat[,5])),
		0),
	labels=c("m1","m2","m3","m1","m2","m3"),
	name="M")

#Mean of trait #3:
Mu <- mxMatrix(type="Full",nrow=1,ncol=1,free=T,values=svmu,labels="mu",lbound=1e-4,name="Mu")

#Index parameter of trait #3:
Nu <- mxMatrix(type="Full",nrow=1,ncol=1,free=T,values=svnu,labels="nu",lbound=1e-4,name="Nu")

#This is how we use thresholds to kludge a variable that has a parametric discrete marginal distribution.  Note that trait #3 has
#an ordered category for every unique count observed in the sample, but the threshold vector--no matter how many elements it has--is always
#a function of only two parameters:
allCounts <- mxMatrix(type="Full",nrow=maxct-minct+1,ncol=1,free=F,values=minct:maxct,name="allCounts")
Tau <- mxAlgebra( cbind(p2z(omxPnbinom(allCounts,Nu,-1,Mu,1,0)), p2z(omxPnbinom(allCounts,Nu,-1,Mu,1,0))), name="Tau")

#Constraint to identify the scale of the Gaussian continuum presumed to underlie trait #3:
ic <- mxConstraint(A[3,3]+C[3,3]+E[3,3]-1 == 0, name="identifying")

#MxModel object:
m1 <- mxModel(
	"GaussCopula",
	mxModel(
		"MZ",
		mxData(mzdat2,type="raw"),
		A, C, E, M, Mu, Nu, allCounts, Tau, ic,
		mxAlgebra(rbind(
			cbind(A+C+E, A+C),
			cbind(A+C, A+C+E)
		), name="SigmaMZ"),
		mxAlgebra(cov2cor(SigmaMZ),name="RhoMZ"),
		mxFitFunctionML(),
		mxExpectationNormal(covariance="SigmaMZ",means="M",dimnames=colnames(mzdat2),thresholds="Tau",threshnames=c("t1y3","t2y3"))
	),
	mxModel(
		"DZ",
		mxData(dzdat2,type="raw"),
		A, C, E, M, Mu, Nu, allCounts, Tau, ic,
		mxAlgebra(rbind(
			cbind(A+C+E, 0.5%x%A+C),
			cbind(0.5%x%A+C, A+C+E)
		), name="SigmaDZ"),
		mxAlgebra(cov2cor(SigmaDZ),name="RhoDZ"),
		mxFitFunctionML(),
		mxExpectationNormal(covariance="SigmaDZ",means="M",dimnames=colnames(dzdat2),thresholds="Tau",threshnames=c("t1y3","t2y3"))
	),
	mxFitFunctionMultigroup(groups=c("MZ","DZ"))
)

#Run it and see results:
m2 <- mxRun(m1)
summ <- summary(m2)
summ
mxEval(RhoMZ,m2$MZ,T)
mxEval(RhoDZ,m2$DZ,T)