/* LMM.PRG -- EM program for linear mixed models

  Author: Brian M. Steele
          Dept. of Math. Sci.
          University of Montana
          Missoula MT, 59812
          Email: bsteele@selway.umt.edu

The author gives permission for redistribution of this software for non-
commercial purposes only. */

library glmm;
#include \gauss\src\glmm.dec;
format /m1/rd 4,0;

output file = llm.out reset;

{r,z,q,x,y,name,dname} = DATA1;

lflag = 1;    @ calc log-like. for each iteration.  Not necessary, and sometimes
               time-consuming.  log-lik.  cannot decrease with iteration @
test = 1E-5;  @ convergence criterion @
s = 1 | 1;        @ initial estimate of variance component vector @

{a,s,s0,b} = ML(s,r,z,q,x,y,test,lflag);
{ion,se} = IOML(a,s,s0,b,x,z,q,name,dname);

{a,s0,s,b} = REML(s,r,z,b,q,x,y,test);
{ion,se} = IOREML(a,s0,s,b,x,z,q,name,dname);

/************************  Data Examples  **********************************/

PROC(7) = DATA1;
  @ Turnip data.  See:
     Hemmerle, W.J. and H.O. Hartley. 1973. Technometrics 15, p.819-831
     and Snedecor, G.W. and W.C. Cochran. Statistical Methods, 7 ed. p. 248
     2 random factors @

dat = {
1 1 3.28,1 1 3.09,
1 2 3.52,1 2 3.48,
1 3 2.88,1 3 2.80,
2 4 2.46,2 4 2.44,
2 5 1.87,2 5 1.92,
2 6 2.19,2 6 2.19,
3 7 2.77,3 7 2.66,
3 8 3.74,3 8 3.44,
3 9 2.55,3 9 2.55,
4 10 3.78,4 10 3.87,
4 11 4.07,4 11 4.12,
4 12 3.31,4 12 3.31};
  y = dat[.,3];
  n = rows(y);
  X = ones(n,1);

  r = rows(q);
  u = unique(dat[.,1],1);
  q = rows(u);
  Z = dat[.,1].==u';

  u = unique(dat[.,2],1);
  q = q|rows(u);
  Z = Z~(dat[.,2].==u');
  r = rows(q);
  let name = const;
  dname = "Turnip data";
RETP(r,z,q,x,y,name,dname);
ENDP;

PROC(7) = DATA2;
  @ Pig data.  See:
     Hemmerle, W.J. and H.O. Hartley. 1973. Technometrics 15, p.819-831
     and Snedecor, G.W. and W.C. Cochran. Statistical Methods, 7 ed. p. 251.
     2 random factors @

dat = {
1 0 0 0 0 1 0 0 0 0 0 0 0 0 0 2.77,
1 0 0 0 0 1 0 0 0 0 0 0 0 0 0 2.38,
1 0 0 0 0 0 1 0 0 0 0 0 0 0 0 2.58,
1 0 0 0 0 0 1 0 0 0 0 0 0 0 0 2.94,
0 1 0 0 0 0 0 1 0 0 0 0 0 0 0 2.28,
0 1 0 0 0 0 0 1 0 0 0 0 0 0 0 2.22,
0 1 0 0 0 0 0 0 1 0 0 0 0 0 0 3.01,
0 1 0 0 0 0 0 0 1 0 0 0 0 0 0 2.61,
0 0 1 0 0 0 0 0 0 1 0 0 0 0 0 2.36,
0 0 1 0 0 0 0 0 0 1 0 0 0 0 0 2.71,
0 0 1 0 0 0 0 0 0 0 1 0 0 0 0 2.72,
0 0 1 0 0 0 0 0 0 0 1 0 0 0 0 2.74,
0 0 0 1 0 0 0 0 0 0 0 1 0 0 0 2.87,
0 0 0 1 0 0 0 0 0 0 0 1 0 0 0 2.46,
0 0 0 1 0 0 0 0 0 0 0 0 1 0 0 2.31,
0 0 0 1 0 0 0 0 0 0 0 0 1 0 0 2.24,
0 0 0 0 1 0 0 0 0 0 0 0 0 1 0 2.74,
0 0 0 0 1 0 0 0 0 0 0 0 0 1 0 2.56,
0 0 0 0 1 0 0 0 0 0 0 0 0 0 1 2.50,
0 0 0 0 1 0 0 0 0 0 0 0 0 0 1 2.48};

  q = 5 | 10;
  r = rows(q);
  X = ones(20,1);
  y = dat[.,cols(dat)];
  Z = dat[.,1:cols(dat)-1];
  let name = const;
  dname = "Pig data";
RETP(r,z,q,x,y,name,dname);
ENDP;

PROC(7) = DATA3;
  @ Oven data.  See:
     Hemmerle, W.J. and H.O. Hartley. 1973. Technometrics 15, p.819-831
     2 random factors @

dat = {
1 0 1 0 0 0 0 0 1 1 0 237,
1 0 1 0 0 0 0 0 1 1 0 254,
1 0 1 0 0 0 0 0 1 1 0 246,
0 1 0 1 0 0 0 0 1 1 0 178,
0 1 0 1 0 0 0 0 1 1 0 179,
1 0 0 0 1 0 0 0 1 0 1 208,
1 0 0 0 1 0 0 0 1 0 1 178,
1 0 0 0 1 0 0 0 1 0 1 187,
0 1 0 0 0 1 0 0 1 0 1 146,
0 1 0 0 0 1 0 0 1 0 1 145,
0 1 0 0 0 1 0 0 1 0 1 141,
1 0 0 0 0 0 1 0 1 -1 -1 186,
1 0 0 0 0 0 1 0 1 -1 -1 183,
0 1 0 0 0 0 0 1 1 -1 -1 142,
0 1 0 0 0 0 0 1 1 -1 -1 125,
0 1 0 0 0 0 0 1 1 -1 -1 136};

  q = {2,6};
  r = 2;
  X = dat[.,9:11];
  y = dat[.,12];
  Z = dat[.,1:8];
  let name = x0 x1 x2;
  dname = "Oven data example";
RETP(r,z,q,x,y,name,dname);
ENDP;

END;
PROC(4) = ML0(s,r,z,q,x,y,test,lflag);
  cls;
  output on;"********** EM algorithm for Linear Mixed Model *************";
  output off;

  if r /= rows(s);
    "s0 is wrong length";
  endif;

  /* setup indexes for random effects */
  qs = sumc(q);
  qidx = cumsumc(q);
  qidx = (trimr(1|qidx+1,0,1))~qidx;

  /* initialize */
  n = rows(y);
  b = zeros(qs,1);
  s0 = vcx(y);
  a = invpd(x'x)'x'y;
  p = cols(x);
  nit = 0;
  iter = {};
  locate 3,1;"iter    llk     s0     a ... s  ";

  /* start of EM cycles.  EM cycles until converge = 1 */
  converge = 0;
  do until converge;

    olpv = pv;

    /* build matrix D and D inverse */
    I = eye(qs);
    d = {};
    j = 1;
    do until j == r + 1;
      t = seqa(qidx[j,1],1,q[j]);
      d = d | s[j]*ones(q[j],1);
      j = j + 1;
    endo;
    din = diagrv(I,1./d);

    /* compute conditional mean (m) */
    m = X*a + Z*b;

    /* Compute b(p) */

    vin = invpd(z'*z/s0 + din);
    b = vin*Z'(y - X*a)/s0;

    /* M-step for alp */
    a = invpd(x'x)*x'(y - Z*b);
    s0 = (1/n)*( (y-m)'(y-m) + tr(vin*Z'Z) );

    /* M-step for s */
    j = 1;
    do until j == r + 1;
      t = seqa(qidx[j,1],1,q[j]);
      s[j] = (b[t]'b[t] + tr(vin[t,t]))/q[j];
      j = j + 1;
    endo;

    /* pv is parameter vector */
    pv = a|s0|s;
    iter = iter|pv';  @ save iterates @
    converge = rows(pv) == sumc(abs(pv-olpv) .< test);

    locate 4,1;
    format 4,0;nit;;format 7,3;
    @ compute log-likehood @
    if lflag*(nit > 1);
        Hin = (1/s0)*(eye(n) - Z*vin*Z'/s0);
        t = -(1/2)*((y-X*a)'Hin*(y-X*a) - ln(det(Hin)));
        e = -(1/2)*((y-X*a-Z*b)'(y-X*a-Z*b)/s0 + n*ln(s0)
                 - ln(det(din)) + b'din*b + ln(det(Z'Z/s0 + din)));
wait;
    else;
      t = 0;
    endif;
    st = st|t;  @ vector of adjusted logL's@
    t;;e;;s0;;a';;s';
    nit = nit + 1;
  endo;
  output on;
  ?;"EM MLEs for linear mixed model ... ";
  "Fixed effects estimates";format 8,4;a';
  "Residual variance";s0;
  "Variance component estimates";format 8,4;s';
  output off;

RETP(a,s,s0,b);
ENDP;