Root finding for extended cases

Emily C. Voldal, Fan Xia, Avi Kenny, Patrick J. Heagerty, and James P. Hughes

This document demonstrates how to find closed-form roots for misspecified parameters in cases not covered by the primary manuscript ‘Model misspecification in stepped wedge trials: Random effects for time or treatment’.  Perhaps most importantly, this document covers the case where the true data distribution is a ‘full model’ with random intercepts, random time effects, and random treatment effects.  See supplemental R files for a demonstration of how to calculate validity and efficiency after finding these roots.
For updates and any corrections, see https://github.com/voldal/SWT_model_misspecification

Defining true model

To demonstrate, we’ll be using a classic two-sequence design with the same fixed effects used in the primary manuscript and a full random effects model.  These inputs can be modified for other designs and models.

globalJ = 3 (*Total number of time points in the design*)

globalK = 2 (*Number of observations per cluster per period*)

globalI = 2 (*Total number of sequences*)

To specify the true distribution of the data, I’m using a matrix, ordered by sequence (rows) and then time (columns).  Here, μT is the intercept, βT is the slope for the linear time trend, and θT is the effect of treatment.  The random effects are denoted by u for intercept, v for treatment, and w for time, indexed by sequence and time as appropriate.  As in the main manuscript, this setup collapses clusters within sequences; however, users may expand this model to list each cluster individually if desired (e.g. modifying to include some cluster-specific fixed effects or varying cluster sizes).

globalTrueMeans = {{μT+u1T+w11T,μT+βT+θT+u1T+v1T+w12T,μT+2*βT+θT+u1T+v1T+w13T},{μT+u2T+w21T,μT+βT+u2T+w22T,μT+2*βT+θT+u2T+v2T+w23T}} (*{{Sequence 1 Time 1, Sequence 1 Time 2,...},{Sequence 2 Time 1, Sequence 2 Time 2,...}}*)

Now, we specify the distribution of the random effects in the model above for each cluster, where τT is the true standard deviation of the random intercepts and γT is the true standard deviation of the random time effect.  In this example, we continue to assume that all random effects are independent but this assumption may be relaxed by specifying a joint distribution.

globalTrueDist = {{u1T \[Distributed] NormalDistribution[0,τT],v1T \[Distributed] NormalDistribution[0,ηT],w11T \[Distributed] NormalDistribution[0,γT],w12T \[Distributed] NormalDistribution[0,γT],w13T \[Distributed] NormalDistribution[0,γT]},{u2T \[Distributed] NormalDistribution[0,τT],v2T \[Distributed] NormalDistribution[0,ηT],w21T \[Distributed] NormalDistribution[0,γT],w22T \[Distributed] NormalDistribution[0,γT],w23T \[Distributed] NormalDistribution[0,γT]}}(*{{Sequence 1},{Sequence 2}}*)

Defining misspecified model

Using the same style as above, now we will define the model that will be used to fit the data.  In this example, we are fitting a random time model.

globalMisspecMeans = {{μ+u1+w11,μ+β+θ+u1+w12,μ+2*β+θ+u1+w13},{μ+u2+w21,μ+β+u2+w22,μ+2*β+θ+u2+w23}}

globalMisspecDist = {{u1 \[Distributed] NormalDistribution[0,τ],w11 \[Distributed] NormalDistribution[0,γ],w12 \[Distributed] NormalDistribution[0,γ],w13 \[Distributed] NormalDistribution[0,γ]},{u2 \[Distributed] NormalDistribution[0,τ],w21 \[Distributed] NormalDistribution[0,γ],w22 \[Distributed] NormalDistribution[0,γ],w23 \[Distributed] NormalDistribution[0,γ]}}

Functions

If new parameters are being added to the true or misspecified models above, these functions will need to be modified.

Function: replicate each element in a list

myReplicate[list_,n_] := Sequence @@ ConstantArray[#,n]&/@list

Function: check for valid roots

checkThis[myList_] :=
Module[{firstCheck,secondCheck,thirdCheck,fourthCheck,fifthCheck,sixthCheck},
firstCheck=Element[myList[[1]][[2]],Reals];(*Intercept*)
secondCheck=Element[myList[[2]][[2]],Reals];(*Treatment effect*)
thirdCheck=Element[myList[[3]][[2]],Reals];(*Time slope*)
fourthCheck = NonNegative[myList[[4]][[2]]];(*Residual variance*)
fifthCheck = NonNegative[myList[[5]][[2]]];(*Random time or treatment variance*)
sixthCheck = NonNegative[myList[[6]][[2]]];(*Random intercept variance*)
AllTrue[{firstCheck,secondCheck,thirdCheck,fourthCheck,fifthCheck,sixthCheck},TrueQ]
]

Function: get elements of system of equations for one cluster when fitting random time

This function returns a vector of six elements, which are the relevant portions of the system of equations for a specific cluster.  Inputs come from the findRoots function, where the misspecified model has random intercepts and random time effects.

oneClusterFittedTime[myK_,myJ_,myMisspecMeans_,myMisspecDist_,myTrueMeans_,myTrueDist_]:=
Module[{myMatrixY,myConditionalLikelihood,myMarginalLikelihood,myPartialMu,myPartialTheta,myPartialBeta,myPartialSigma,myPartialGamma,myPartialTau,myEPartialMu,myEPartialTheta,myEPartialBeta,myEPartialSigma,myEPartialGamma,myEPartialTau},
(*Make names of y's: for each time point (1,...,J), need K unique y's.  Arrange in a matrix so rows can be pulled by time index later.*)
myMatrixY = Table[StringJoin["yj",TextString[j],"k",TextString[k]],{j,myJ},{k,myK}];(*Each row corresponds to a time point, pull with myMatrixY[[j]]*)
myMatrixY = ToExpression[myMatrixY];
(*Build initial likelihood by looping through each time point*)
myConditionalLikelihood = 1;
Do[myConditionalLikelihood=myConditionalLikelihood * Likelihood[NormalDistribution[myMisspecMeans[[j]],σ],myMatrixY[[j]]],{j,myJ}];
(*Find marginal likelihood*)
myMarginalLikelihood= Expectation[myConditionalLikelihood, myMisspecDist];
(*Take derivatives; note that these depend on consistent notation in the MisspecMeans*)
myPartialMu= Simplify[D[Log[myMarginalLikelihood],μ]];
myPartialTheta= Simplify[D[Log[myMarginalLikelihood],θ]];
myPartialBeta= Simplify[D[Log[myMarginalLikelihood],β]];
myPartialSigma= Simplify[D[Log[myMarginalLikelihood],σ]];
myPartialGamma= Simplify[D[Log[myMarginalLikelihood],γ]];
myPartialTau= Simplify[D[Log[myMarginalLikelihood],τ]];
(*Take expectation over true distribution of Y's*)
myEPartialMu=Expectation[Expectation[myPartialMu ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialTheta=Expectation[Expectation[myPartialTheta ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialBeta=Expectation[Expectation[myPartialBeta,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialSigma=Expectation[Expectation[myPartialSigma ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialGamma=Expectation[Expectation[myPartialGamma ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialTau=Expectation[Expectation[myPartialTau ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
{myEPartialMu,myEPartialTheta,myEPartialBeta,myEPartialSigma,myEPartialGamma,myEPartialTau}
]

Function: construct and solve system of equations when fitting random time

This function returns solutions to the system of equations that pass some basic screening.  Inputs are the true and misspecified models as described above, where the misspecified model has random intercepts and random time effects.

findRootsFittedTime[myK_,myJ_,myI_,myMisspecMeansList_,myMisspecDistList_,myTrueMeansList_,myTrueDistList_]:=
Module[{allclusterresults,systemsolutions,systemsolutionsplugin,systemsolutionspluginlist},
allclusterresults = Table[oneClusterFittedTime[myK,myJ,myMisspecMeansList[[i]],myMisspecDistList[[i]],myTrueMeansList[[i]],myTrueDistList[[i]]],{i,myI}];
systemsolutions = Solve[{Simplify[Total[allclusterresults[[All,1]]]]==0,Simplify[Total[allclusterresults[[All,2]]]]==0,Simplify[Total[allclusterresults[[All,3]]]]==0,Simplify[Total[allclusterresults[[All,4]]]]==0,Simplify[Total[allclusterresults[[All,5]]]]==0,Simplify[Total[allclusterresults[[All,6]]]]==0},{μ,θ,β,σ,γ,τ}];
systemsolutionsplugin = systemsolutions/.{μT -> 0, θT -> 0, βT -> 0,σT -> 1, ηT -> 1, τT -> 1, γT -> 1};
systemsolutionspluginlist = systemsolutionsplugin/. Rule->List;
{Do[If[checkThis[systemsolutionspluginlist[[i]]],Print[systemsolutions[[i]]],],{i,Length[systemsolutions]}]}
]

Function: get elements of system of equations for one cluster when fitting random treatment

This function returns a vector of six elements, which are the relevant portions of the system of equations for a specific cluster.  Inputs come from the findRoots function, where the misspecified model has random intercepts and random treatment effects.

oneClusterFittedTreatment[myK_,myJ_,myMisspecMeans_,myMisspecDist_,myTrueMeans_,myTrueDist_]:=
Module[{myMatrixY,myConditionalLikelihood,myMarginalLikelihood,myPartialMu,myPartialTheta,myPartialBeta,myPartialSigma,myPartialEta,myPartialTau,myEPartialMu,myEPartialTheta,myEPartialBeta,myEPartialSigma,myEPartialEta,myEPartialTau},
(*Make names of y's: for each time point (1,...,J), need K unique y's.  Arrange in a matrix so rows can be pulled by time index later.*)
myMatrixY = Table[StringJoin["yj",TextString[j],"k",TextString[k]],{j,myJ},{k,myK}];(*Each row corresponds to a time point, pull with myMatrixY[[j]]*)
myMatrixY = ToExpression[myMatrixY];
(*Build initial likelihood by looping through each time point*)
myConditionalLikelihood = 1;
Do[myConditionalLikelihood=myConditionalLikelihood * Likelihood[NormalDistribution[myMisspecMeans[[j]],σ],myMatrixY[[j]]],{j,myJ}];
(*Find marginal likelihood*)
myMarginalLikelihood= Expectation[myConditionalLikelihood, myMisspecDist];
(*Take derivatives; note that these depend on consistent notation in the MisspecMeans*)
myPartialMu= Simplify[D[Log[myMarginalLikelihood],μ]];
myPartialTheta= Simplify[D[Log[myMarginalLikelihood],θ]];
myPartialBeta= Simplify[D[Log[myMarginalLikelihood],β]];
myPartialSigma= Simplify[D[Log[myMarginalLikelihood],σ]];
myPartialEta= Simplify[D[Log[myMarginalLikelihood],η]];
myPartialTau= Simplify[D[Log[myMarginalLikelihood],τ]];
(*Take expectation over true distribution of Y's*)
myEPartialMu=Expectation[Expectation[myPartialMu ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialTheta=Expectation[Expectation[myPartialTheta ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialBeta=Expectation[Expectation[myPartialBeta,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialSigma=Expectation[Expectation[myPartialSigma ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialEta=Expectation[Expectation[myPartialEta ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
myEPartialTau=Expectation[Expectation[myPartialTau ,{Flatten[myMatrixY] \[Distributed]MultinormalDistribution[myReplicate[myTrueMeans,myK],σT^2*IdentityMatrix[myJ*myK]]}],myTrueDist];
{myEPartialMu,myEPartialTheta,myEPartialBeta,myEPartialSigma,myEPartialEta,myEPartialTau}
]

Function: construct and solve system of equations when fitting random treatment

This function returns solutions to the system of equations that pass some basic screening.  Inputs are the true and misspecified models as described above, where the misspecified model has random intercepts and random treatment effects.

findRootsFittedTreatment[myK_,myJ_,myI_,myMisspecMeansList_,myMisspecDistList_,myTrueMeansList_,myTrueDistList_]:=
Module[{allclusterresults,systemsolutions,systemsolutionsplugin,systemsolutionspluginlist},
allclusterresults = Table[oneClusterFittedTreatment[myK,myJ,myMisspecMeansList[[i]],myMisspecDistList[[i]],myTrueMeansList[[i]],myTrueDistList[[i]]],{i,myI}];
systemsolutions = Solve[{Simplify[Total[allclusterresults[[All,1]]]]==0,Simplify[Total[allclusterresults[[All,2]]]]==0,Simplify[Total[allclusterresults[[All,3]]]]==0,Simplify[Total[allclusterresults[[All,4]]]]==0,Simplify[Total[allclusterresults[[All,5]]]]==0,Simplify[Total[allclusterresults[[All,6]]]]==0},{μ,θ,β,σ,η,τ}];
systemsolutionsplugin = systemsolutions/.{μT -> 0, θT -> 0, βT -> 0,σT -> 1, γT -> 1, τT -> 1, ηT -> 1};
systemsolutionspluginlist = systemsolutionsplugin/. Rule->List;
Do[If[checkThis[systemsolutionspluginlist[[i]]],Print[systemsolutions[[i]]],],{i,Length[systemsolutions]}]
]

Examples

Using the design set up above, here are the roots.  Note that this can be a little time-consuming to run, especially for designs with many time points or sequences.  These roots are very similar in form to the roots identified for the simpler time-fitted random treatment case in the main manuscript.  As we might have expected, the only difference in Root 1 compared to the main manuscript is that the fitted random time effect now involves the true time effect.  Note that plugging in 0 for the true random treatment variance provides results for a correctly specified random time model.

findRootsFittedTime[globalK,globalJ,globalI,globalMisspecMeans,globalMisspecDist,globalTrueMeans,globalTrueDist]

In theory, you can do the same thing when the misspecified model has a random treatment effect instead of time (see example below, for a SWT with only two time points where the second sequence remains in control); however, this is much more difficult computationally than the fitted time case.  It may be more feasible to find numeric solutions, using the same system of equations.

globalJ = 2

globalK = 5

globalI = 2

globalTrueMeans = {{μT+u1T+w11T,μT+βT+θT+u1T+v1T+w12T},{μT+u2T+w21T,μT+βT+u2T+w22T}}

globalTrueDist = {{u1T \[Distributed] NormalDistribution[0,τT],v1T \[Distributed] NormalDistribution[0,ηT],w11T \[Distributed] NormalDistribution[0,γT],w12T \[Distributed] NormalDistribution[0,γT]},{u2T \[Distributed] NormalDistribution[0,τT],w21T \[Distributed] NormalDistribution[0,γT],w22T \[Distributed] NormalDistribution[0,γT]}}

globalMisspecMeans = {{μ+u1,μ+β+θ+u1+v1},{μ+u2,μ+β+u2}}

globalMisspecDist = {{u1 \[Distributed] NormalDistribution[0,τ],v1 \[Distributed] NormalDistribution[0,η]},{u2 \[Distributed] NormalDistribution[0,τ]}}

(*findRootsFittedTreatment[globalK,globalJ,globalI,globalMisspecMeans,globalMisspecDist,globalTrueMeans,globalTrueDist]*)(*Commented out because it takes so long to run*)

Here is an example of how you might try to solve the problem above using numeric solutions, for a particular set of true parameter settings (all fixed effects are 0, all standard deviations are 1).

allclusterresults = Table[oneClusterFittedTreatment[globalK,globalJ,globalMisspecMeans[[i]],globalMisspecDist[[i]],globalTrueMeans[[i]],globalTrueDist[[i]]],{i,globalI}];

systemsolutions = NSolve[{Simplify[Total[allclusterresults[[All,1]]]]==0,Simplify[Total[allclusterresults[[All,2]]]]==0,Simplify[Total[allclusterresults[[All,3]]]]==0,Simplify[Total[allclusterresults[[All,4]]]]==0,Simplify[Total[allclusterresults[[All,5]]]]==0,Simplify[Total[allclusterresults[[All,6]]]]==0}/.{μT -> 0, θT -> 0, βT -> 0,σT -> 1, γT -> 1, τT -> 1, ηT -> 1},{μ,θ,β,σ,η,τ},Reals];

As in the primary cases covered in the main manuscript, the solutions below include four sets of values for which all the variance components are non-negative, following the same pattern of Root 1, Root 2, Root 3, and Root 4.

systemsolutions

Below, we find another set of solutions for a different set of true parameters, but we set the true random time standard deviation to 0 so that the model is correctly fitted.  With these numeric solutions the number of roots returned varies with different settings, but seems to usually include Root 1. For example, running this design with a smaller K produces all four roots.  This may be a result of the algorithm used to find numeric solutions, and perhaps using a different algorithm or better settings would resolve this inconsistency.

systemsolutionsCorrectModel = NSolve[{Simplify[Total[allclusterresults[[All,1]]]]==0,Simplify[Total[allclusterresults[[All,2]]]]==0,Simplify[Total[allclusterresults[[All,3]]]]==0,Simplify[Total[allclusterresults[[All,4]]]]==0,Simplify[Total[allclusterresults[[All,5]]]]==0,Simplify[Total[allclusterresults[[All,6]]]]==0}/.{μT -> 1, θT -> 0.5, βT -> 2,σT -> 1, γT -> 0, τT -> 0.5, ηT -> 0.1},{μ,θ,β,σ,η,τ},Reals];

systemsolutionsCorrectModel