text: import numpy as NP
  import scipy as SP
  from numpy import dot
  import scipy.stats as ST
  from scipy.linalg import cholesky,solve_triangular
  from fastlmmhpc.external.util.math import check_definite_positiveness,check_symmetry,mvnormpdf,ddot,trace2
  from fastlmmhpc.external.util.math import stl, stu, dotd
  from fastlmmhpc.inference.glmm import GLMM_N1K3, GLMM_N3K1
  from fastlmmhpc.inference.likelihood import LogitLikelihood, ProbitLikelihood
  import likelihood as LH
  
  
  '''
      Important! Always run test.py in the current folder for unit testing after
      changes have been made.
  '''
  
  class EPGLMM(object):
      def __init__(self, link):
          self._debugUACalls = []
          self._link = link
          if link == "logistic":
              self._likelihood = LogitLikelihood()
          elif link == "erf":
              self._likelihood = ProbitLikelihood()
          else:
              assert False, "Unknown link function."
  
      def _updateApproximation(self):
          '''
          Calculates the Laplace approximation for the posterior.
          It can be defined by two variables: f mode and W at f mode.
          '''
          if self._updateApproximationCount == 0:
              return
  
          if self._is_kernel_zero():
              self._updateApproximationCount = 0
              return
  
          self._updateApproximationBegin()
  
          self._mean = self._calculateMean()
          m = self._mean
  
          ttau = NP.zeros(self._N)
          tnu = NP.zeros(self._N)
          sig2_ = NP.zeros(self._N)
          mu_ = NP.zeros(self._N)
  
          prevsig2 = self._dKn()
  
          prevmu = m.copy()
  
          converged = False
          outeriter = 1
          iterMax = 1000
          
          while outeriter <= iterMax and not converged:
  
              tau_ = 1.0/prevsig2 - ttau
              nu_ = prevmu/prevsig2 - tnu
  
              tt = tau_**2 + tau_
              stt = NP.sqrt(tt)
              c = (self._y*nu_) / stt
  
              # dan: using _cdf and _pdf instead of cdf and pdf I avoid
              # a lot of overhead due to error checking and other things
              nc_hz = NP.exp( ST.norm._logpdf(c) - ST.norm._logcdf(c) )
  
              hmu = nu_/tau_ + nc_hz*self._y/stt
  
              hsig2 = 1.0/tau_ - (nc_hz/tt) * (c + nc_hz)
  
              tempttau = 1.0/hsig2 - tau_
              ok = NP.bitwise_and(~NP.isnan(tempttau), abs(tempttau)>1e-8)
              ttau[ok] = tempttau[ok]
  
              tnu[ok] = hmu[ok]/hsig2[ok] - nu_[ok]
  
              (sig2, mu) = self._calculateSig2Mu(ttau, tnu, m)
  
              lastMuDiff = abs(mu-prevmu).max()
              lastSig2Diff = abs(sig2-prevsig2).max()
              if lastMuDiff  < 1e-6 and lastSig2Diff < 1e-6:
                  converged = True
  
              prevmu = mu
              prevsig2 = sig2
  
              outeriter += 1
  
          if outeriter > iterMax:
              Pr.prin('EP did not converge. Printing debug information...')
              Pr.prin('beta ' + str(self.beta))
              Pr.prin('sig02 '+str(self.sig02)+' sig12 '+str(self.sig12)+' sign2 '+str(self.sign2))
              Pr.prin('abs(mu-prevmu).max() '+str(lastMuDiff))
              Pr.prin('abs(sig2-prevsig2).max() '+str(lastSig2Diff))
  
          self._updateApproximationEnd(ttau, tnu, tau_, nu_)
          self._updateApproximationCount = 0
  
  class EPGLMM_N1K3(GLMM_N1K3, EPGLMM):
      def __init__(self, link, penalty=None, penalizeBias=False, debug=False):
          GLMM_N1K3.__init__(self, penalty=penalty, penalizeBias=penalizeBias, debug=debug)
          EPGLMM.__init__(self, link)
          self._link = link
  
      def _updateApproximationBegin(self):
          self._G01 = self._calculateG01()
  
      def _updateApproximationEnd(self, ttau, tnu, tau_, nu_):
          self._ttau = ttau
          self._tnu = tnu
          self._tau_ = tau_
          self._nu_ = nu_
          self._V = self._calculateV(ttau, self._sign2)
          self._Lk = self._calculateLk(self._G01, self._V)
          self._H = stl(self._Lk, ddot(self._G01.T, self._V, left=False))
  
      def _regular_marginal_loglikelihood(self):
          self._updateConstants()
          self._updateApproximation()
  
          if self._is_kernel_zero():
              return self._likelihood.log(self._mean, self._y)
  
          m = self._mean
          sign2 = self._sign2
          tnu = self._tnu
          ttau = self._ttau
          tau_ = self._tau_
          nu_ = self._nu_
          G01 = self._G01
          V = self._V
          Lk = self._Lk
          H = self._H
  
          A = ttau*sign2 + 1.0
          Ktnu = self._rdotK(tnu)
          HKtnu = dot(H,Ktnu)
          Hm = dot(H,m)
          ts_1 = 1.0/(tau_+ttau)
          mV = m*V
  
          r = 0.5*NP.sum(NP.log(1.0+ttau/tau_)) - 0.5*NP.sum(NP.log(A)) - NP.sum(NP.log(NP.diag(Lk)))\
              + NP.sum( ST.norm._logcdf((self._y*nu_)/NP.sqrt(tau_**2+tau_)) )\
              + 0.5*(dot(tnu, Ktnu) - dot(Ktnu.T, V*Ktnu) + dot(HKtnu.T, HKtnu) - dot(tnu, ts_1*tnu))\
              + dot(m,tnu) - dot(mV, Ktnu) + dot(Hm.T, HKtnu) - 0.5*(dot(mV, m) - dot(Hm.T, Hm))\
              + 0.5*dot(nu_*ts_1, (ttau/tau_)*nu_ - 2.0*tnu)
  
          return r
  
      def _rmll_gradient(self, optSig02=True, optSig12=True, optSign2=True, optBeta=True):
          self._updateConstants()
          self._updateApproximation()
          
          m = self._mean
          H = self._H
          V = self._V
          ttau = self._ttau
          tnu = self._tnu
          G0 = self._G0
          G1 = self._G1
  
          Smtnu = ttau*m - tnu
          KSmtnu = self._rdotK(Smtnu)
  
          b = Smtnu - V*KSmtnu + dot(H.T, dot(H, KSmtnu))
  
          ret = []
          if optSig02:
              r = 0.5*(dot(b, dot(G0, dot(G0.T, b))) - trace2(ddot(V, G0, left=True), G0.T)\
                  + trace2(H.T, dot(dot(H, G0), G0.T)))
              
              ret.append(r)
  
          if optSig12:
              r = 0.5*(dot(b, dot(G1, dot(G1.T, b))) - trace2(ddot(V, G1, left=True), G1.T)\
                  + trace2(H.T, dot(dot(H, G1), G1.T)))
              
              ret.append(r)
  
          if optSign2:
              r = 0.5*(dot(b, b) - NP.sum(V)\
                  + trace2(H.T, H))
              
              ret.append(r)
  
          if optBeta:
              ret += list(-dot(b, self._X))
  
          return NP.array(ret)
  
      def _updateApproximation(self):
          EPGLMM._updateApproximation(self)
  
      def _calculateV(self, ttau, sign2):
          V = ttau/(ttau*sign2 + 1.0)
          return V
  
      def _calculateSig2Mu(self, ttau, tnu, mean):
          sign2 = self._sign2
  
          V = self._calculateV(ttau, sign2)
          G01 = self._G01
      
          Lk = self._calculateLk(G01, V)
      
          G01tV = ddot(G01.T, V, left=False)
          H = stl(Lk, G01tV)
  
          HK = self._ldotK(H)
      
          sig2 = self._dKn() - sign2**2*V - dotd(G01, dot(dot(G01tV, G01), G01.T))\
              - 2.0*sign2*dotd(G01, G01tV) + dotd(HK.T, HK)
  
          assert NP.all(NP.isfinite(sig2)), 'sig2 should be finite.'
  
          u = self._mean + self._rdotK(tnu)
      
          mu = u - self._rdotK(V*u) + self._rdotK(H.T.dot(H.dot(u)))
          assert NP.all(NP.isfinite(mu)), 'mu should be finite.'
  
          return (sig2, mu)
  
      def _predict(self, meanstar, kstar, kstarstar, prob):
          self._updateConstants()
          self._updateApproximation()
  
          m = self._mean
          tnu = self._tnu
          Lk = self._Lk
          H = self._H
          V = self._V
  
          Ktnu = self._rdotK(tnu)
          mKtnu = m + Ktnu
          Vkstar = ddot(V, kstar, left=True)
          Hkstar = dot(H, kstar)
  
          mustar = meanstar + dot(kstar.T,tnu) - dot(Vkstar.T, mKtnu) + dot(Hkstar.T, dot(H, mKtnu))
  
          if prob is False:
              if nom > 0.0:
                  return +1.0
              return -1.0
  
          sig2star = kstarstar - dotd(kstar.T,Vkstar)\
              + dotd(dot(Hkstar.T, H), kstar) + self._sign2
  
          return LH.probit_sigmoid(mustar/NP.sqrt(1.0 + sig2star))
  
  class EPGLMM_N3K1(GLMM_N3K1, EPGLMM):
      def __init__(self, link, penalty=None, penalizeBias=False, debug=False):
          GLMM_N3K1.__init__(self, penalty=penalty, penalizeBias=penalizeBias, debug=debug)
          EPGLMM.__init__(self, link)
          self._link = link
  
      def _updateApproximationBegin(self):
          self._K = NP.eye(self._N) * self._sign2
          if self._isK0Set:
              self._K += self._sig02*(dot(self._G0, self._G0.T))
          if self._isK1Set:
              self._K += self._sig12*(dot(self._G1, self._G1.T))
  
      def _regular_marginal_loglikelihood(self):
          self._updateConstants()
          self._updateApproximation()
  
          if self._is_kernel_zero():
              return self._likelihood.log(self._mean, self._y)
  
          m = self._mean
          K = self._K
          Ln = self._Ln
          ttau = self._ttau
          tnu = self._tnu
          tau_ = self._tau_
          nu_ = self._nu_
          LnSsq = self._LnSsq
          LnSsqK = self._LnSsqK
  
          LnSsqKtnu = dot(LnSsqK, tnu)
          ts_1 = 1.0/(tau_ + ttau)
          LnSsqm = dot(LnSsq, m)
  
          r = 0.5*NP.sum(NP.log(1.0 + ttau/tau_)) - NP.sum(NP.log(NP.diag(Ln)))\
              + NP.sum( ST.norm._logcdf( (self._y*nu_) / NP.sqrt(tau_**2+tau_) ) )\
              + 0.5*(dot(tnu, dot(K, tnu)) - dot(LnSsqKtnu,LnSsqKtnu) - dot(tnu*ts_1, tnu))\
              + dot(m, tnu) - dot(LnSsqm.T, LnSsqKtnu) - 0.5*dot(LnSsqm.T, LnSsqm)\
              + 0.5*dot(nu_*ts_1, (ttau/tau_)*nu_ - 2.0*tnu)
  
          return r
  
      def _rmll_gradient(self, optSig02=True, optSig12=True, optSign2=True, optBeta=True):
          self._updateConstants()
          self._updateApproximation()
  
          K0 = self._K0
          K1 = self._K1
          m = self._mean
          ttau = self._ttau
          tnu = self._tnu
          Ssq = self._Ssq
          LnSsq = self._LnSsq
          LnSsqK = self._LnSsqK
  
          Smtnu = ttau*m - tnu
  
          b = Smtnu - dot(LnSsq.T, dot(LnSsqK, Smtnu) )
  
          ret = []
  
          if optSig02:
              r = 0.5*(dot(b,dot(K0,b)) - trace2(LnSsq.T, dot(LnSsq, K0)))
              ret.append(r)
  
          if optSig12:
              r = 0.5*(dot(b,dot(K1,b)) - trace2(LnSsq.T, dot(LnSsq, K1)))
              ret.append(r)
  
          if optSign2:
              r = 0.5*(dot(b,b) - trace2(LnSsq.T, LnSsq))
              ret.append(r)
  
          if optBeta:
              r = -dot(b, self._X)
              ret += list(r)
              
          ret = NP.array(ret)
          assert NP.all(NP.isfinite(ret)), 'Not finite regular marginal loglikelihood gradient.'
          
          return ret
  
      def _updateApproximation(self):
          EPGLMM._updateApproximation(self)
  
      def _updateApproximationEnd(self, ttau, tnu, tau_, nu_):
          self._ttau = ttau
          self._tnu = tnu
          self._tau_ = tau_
          self._nu_ = nu_
          self._Ssq = NP.sqrt(ttau)
          self._Ln = self._calculateLn(self._K, self._Ssq)
          self._LnSsq = stl(self._Ln, NP.diag(self._Ssq))
          self._LnSsqK = dot(self._LnSsq, self._K)
  
      def _calculateSig2Mu(self, ttau, tnu, mean):
          Ssq = NP.sqrt(ttau)
  
          Ln = self._calculateLn(self._K, Ssq)
  
          LnSsq = stl(Ln, NP.diag(Ssq))
          V = dot(LnSsq,self._K)
          sig2 = self._dKn() - dotd(V.T,V)
          mu = mean + dot(self._K,tnu) - dot(V.T,dot(LnSsq,mean)) - dot(V.T,dot(V,tnu))
  
          return (sig2, mu)
  
      def _predict(self, meanstar, kstar, kstarstar, prob):
          self._updateConstants()
          self._updateApproximation()
  
          m = self._mean
          tnu = self._tnu
          K = self._K
          Ln = self._Ln
          Ssq = self._Ssq
          LnSsq = self._LnSsq
          LnSsqkstar = dot(LnSsq, kstar)
          LnSsqK = self._LnSsqK
  
          mustar = meanstar + dot(kstar.T,tnu) - dot(LnSsqkstar.T, dot(LnSsq,m) + dot(LnSsqK,tnu))
  
          if prob is False:
              if nom > 0.0:
                  return +1.0
              return -1.0
  
          sig2star = kstarstar - dotd(dot(LnSsqkstar.T, LnSsq), kstar) + self._sign2
  
          return LH.probit_sigmoid(mustar/NP.sqrt(1.0 + sig2star))
  