code: import sys
  from limix.core.old.cobj import *
  from limix.utils.preprocess import regressOut
  import numpy as np
  
  import scipy.linalg as LA
  import copy
  
  def compute_X1KX2(Y, D, X1, X2, A1=None, A2=None):
  
      R,C = Y.shape
      if A1 is None:
          nW_A1 = Y.shape[1]
          #A1 = np.eye(Y.shape[1])	#for now this creates A1 and A2
      else:
          nW_A1 = A1.shape[0]
  
      if A2 is None:
          nW_A2 = Y.shape[1]
          #A2 = np.eye(Y.shape[1])	#for now this creates A1 and A2
      else:
          nW_A2 = A2.shape[0]
  
  
      nW_X1 = X1.shape[1]
      rows_block = nW_A1 * nW_X1
  
      if 0:#independentX2:
          nW_X2 = 1
      else:
          nW_X2 = X2.shape[1]
      cols_block = nW_A2 * nW_X2
  
      block = np.zeros((rows_block,cols_block))
  
  
      if (R>C) or (A1 is None) or (A2 is None):
          for c in range(C):
              X1D = X1 * D[:,c:c+1]
              X1X2 = X1D.T.dot(X2)
              if (A1 is None) and (A2 is None):
                  block[c*X1.shape[1]:(c+1)*X1.shape[1], c*X2.shape[1]:(c+1)*X2.shape[1]] += X1X2
              elif (A1 is None):
                  block[c*X1.shape[1]:(c+1)*X1.shape[1],:] += np.kron(A2[:,c:c+1].T,X1X2)
              elif (A2 is None):
                  block[:,c*X2.shape[1]:(c+1)*X2.shape[1]] += np.kron(A1[:,c:c+1],X1X2)
              else:
                  A1A2 = np.outer(A1[:,c],A2[:,c])
                  block += np.kron(A1A2,X1X2)
      else:
          for r in range(R):
              A1D = A1 * D[r:r+1,:]
              A1A2 = A1D.dot(A2.T)
              X1X2 = X1[r,:][:,np.newaxis].dot(X2[r,:][np.newaxis,:])
              block += np.kron(A1A2,X1X2)
  
      return block
  
  class mean(cObject):
  
      def __init__(self,Y, identity_trick=False):
          """ init data term """
          self.Y = Y
          self.identity_trick=identity_trick
          self.clearFixedEffect()
  
      #########################################
      # Properties
      #########################################
  
      @property
      def A(self):
          return self._A
  
      @property
      def B(self):
          return self._B
  
      @property
      def F(self):
          return self._F
  
      @property
      def A_identity(self):
          return self._A_identity
  
      @property
      def REML_term(self):
          return self._REML_term
  
      @property
      def Y(self):
          return self._Y
  
      @property
      def N(self):
          return self._N
  
      @property
      def P(self):
          return self._P
  
      @property
      def n_fixed_effs(self):
          return self._n_fixed_effs
  
      @property
      def n_terms(self):
          return self._n_terms
  
      @property
      def Lr(self):
          return self._Lr
  
      @property
      def Lc(self):
          return self._Lc
  
      @property
      def d(self):
          return self._d
  
      @property
      def D(self):
          return np.reshape(self.d,(self.N,self.P), order='F')
  
      @property
      def LRLdiag(self):
          return self._LRLdiag
  
      @property
      def LCL(self):
          return self._LCL
  
      #########################################
      # Setters
      #########################################
      def use_identity_trick(self,identity_trick=True):
          self.identity_trick=identity_trick
          self.clear_cache('Fstar','Astar','Xstar','Xhat',
                           'Areml','Areml_eigh','Areml_chol','Areml_inv','beta_hat','B_hat',
                           'LRLdiag_Xhat_tens','Areml_grad',
                           'beta_grad','Xstar_beta_grad','Zstar','DLZ')
  
  
      def clearFixedEffect(self):
          """ erase all fixed effects """
          self._A = []
          self._F = []
          self._B = []
          self._A_identity = []
          self._REML_term = []
          self._n_terms = 0
          self._n_fixed_effs = 0
          self._n_fixed_effs_REML = 0
          self.indicator = {'term':np.array([]),
                              'row':np.array([]),
                              'col':np.array([])}
          self.clear_cache('Fstar','Astar','Xstar','Xhat',
                           'Areml','Areml_eigh','Areml_chol','Areml_inv','beta_hat','B_hat',
                           'LRLdiag_Xhat_tens','Areml_grad',
                           'beta_grad','Xstar_beta_grad','Zstar','DLZ')
  
      def addFixedEffect(self,F=None,A=None, REML=True, index=None):
          """
          set sample and trait designs
          F:      NxK sample design
          A:      LxP sample design
          REML:   REML for this term?
          index:  index of which fixed effect to replace. If None, just append.
          """
          if F is None:   F = np.ones((self.N,1))
          if A is None:
              A = np.eye(self.P)
              A_identity = True
          elif (A.shape == (self.P,self.P)) & (A==np.eye(self.P)).all():
              A_identity = True
          else:
              A_identity = False
  
          assert F.shape[0]==self.N, "F dimension mismatch"
          assert A.shape[1]==self.P, "A dimension mismatch"
          if index is None or index==self.n_terms:
              self.F.append(F)
              self.A.append(A)
              self.A_identity.append(A_identity)
              self.REML_term.append(REML)
              # build B matrix and indicator
              self.B.append(np.zeros((F.shape[1],A.shape[0])))
              self._n_terms+=1
              self._update_indicator(F.shape[1],A.shape[0])
          elif index >self.n_terms:
              raise Exception("index exceeds max index of terms")
          else:
              self._n_fixed_effs-=self.F[index].shape[1]*self.A[index].shape[0]
              if self.REML_term[index]:
                  self._n_fixed_effs_REML-=self.F[index].shape[1]*self.A[index].shape[0]
              self.F[index] = F
              self.A[index] = A
              self.A_identity[index] = A_identity
              self.REML_term[index]=REML
              self.B[index] = np.zeros((F.shape[1],A.shape[0]))
              self._rebuild_indicator()
  
          self._n_fixed_effs+=F.shape[1]*A.shape[0]
          if REML:
              self._n_fixed_effs_REML+=F.shape[1]*A.shape[0]
          self.clear_cache('Fstar','Astar','Xstar','Xhat',
                           'Areml','Areml_eigh','Areml_chol','Areml_inv','beta_hat','B_hat',
                           'LRLdiag_Xhat_tens','Areml_grad',
                           'beta_grad','Xstar_beta_grad','Zstar','DLZ')
  
      def removeFixedEffect(self, index=None):
          """
          set sample and trait designs
          F:      NxK sample design
          A:      LxP sample design
          REML:   REML for this term?
          index:  index of which fixed effect to replace. If None, remove last term.
          """
          if self._n_terms==0:
              pass
          if index is None or index==(self._n_terms-1):
  
              self._n_terms-=1
              F = self._F.pop() #= self.F[:-1]
              A = self._A.pop() #= self.A[:-1]
              self._A_identity.pop() #= self.A_identity[:-1]
              REML_term = self._REML_term.pop()# = self.REML_term[:-1]
              self._B.pop()# = self.B[:-1]
              self._n_fixed_effs-=F.shape[1]*A.shape[0]
              if REML_term:
                  self._n_fixed_effs_REML-=F.shape[1]*A.shape[0]
  
              pass
          elif index >= self.n_terms:
              raise Exception("index exceeds max index of terms")
          else:
              raise NotImplementedError("currently only last term can be removed")
              pass
          self._rebuild_indicator()
          self.clear_cache('Fstar','Astar','Xstar','Xhat',
                           'Areml','Areml_eigh','Areml_chol','Areml_inv','beta_hat','B_hat',
                           'LRLdiag_Xhat_tens','Areml_grad',
                           'beta_grad','Xstar_beta_grad','Zstar','DLZ')
  
      @Y.setter
      def Y(self,value):
          """ set phenotype """
          self._N,self._P = value.shape
          self._Y = value
          self.clear_cache('Ystar1','Ystar','Yhat','LRLdiag_Yhat',
                           'beta_grad','Xstar_beta_grad','Zstar','DLZ')
  
      @Lr.setter
      def Lr(self,value):
          """ set row rotation """
          assert value.shape[0]==self._N, 'dimension mismatch'
          assert value.shape[1]==self._N, 'dimension mismatch'
          self._Lr = value
          self.clear_cache('Fstar','Ystar1','Ystar','Yhat','Xstar','Xhat',
                           'Areml','Areml_eigh','Areml_chol','Areml_inv','beta_hat','B_hat',
                           'LRLdiag_Xhat_tens','LRLdiag_Yhat','Areml_grad',
                           'beta_grad','Xstar_beta_grad',
                           'LRLdiag_Xhat_tens','LRLdiag_Yhat','Areml_grad',
                           'beta_grad','Xstar_beta_grad','Zstar','DLZ')
  
      @Lc.setter
      def Lc(self,value):
          """ set col rotation """
          assert value.shape[0]==self._P, 'Lc dimension mismatch'
          assert value.shape[1]==self._P, 'Lc dimension mismatch'
          self._Lc = value
          self.clear_cache('Astar','Ystar','Yhat','Xstar','Xhat',
                           'Areml','Areml_eigh','Areml_chol','Areml_inv','beta_hat','B_hat',
                           'LRLdiag_Xhat_tens','LRLdiag_Yhat','Areml_grad',
                           'beta_grad','Xstar_beta_grad','Zstar','DLZ')
  
      @d.setter
      def d(self,value):
          """ set anisotropic scaling """
          assert value.shape[0]==self._P*self._N, 'd dimension mismatch'
          self._d = value
          self.clear_cache('Yhat','Xhat','Areml','Areml_eigh','Areml_chol','Areml_inv','beta_hat','B_hat',
                           'LRLdiag_Xhat_tens','LRLdiag_Yhat','Areml_grad',
                           'beta_grad','Xstar_beta_grad','Zstar','DLZ')
  
      @LRLdiag.setter
      def LRLdiag(self,value):
          """ set anisotropic scaling """
          self._LRLdiag = value
          self.clear_cache('LRLdiag_Xhat_tens','LRLdiag_Yhat','Areml_grad',
                           'beta_grad','Xstar_beta_grad')
  
      @LCL.setter
      def LCL(self,value):
          """ set anisotropic scaling """
          self._LCL = value
          self.clear_cache('Areml_grad','beta_grad','Xstar_beta_grad')
  
      #########################################
      # Getters (caching)
      #########################################
  
      @cached
      def Astar(self):
          RV = []
          for term_i in range(self.n_terms):
              RV.append(np.dot(self.A[term_i],self.Lc.T))
          return RV
  
      @cached
      def Fstar(self):
          RV = []
          for term_i in range(self.n_terms):
              RV.append(np.dot(self.Lr,self.F[term_i]))
          return RV
  
      def Ystar1(self):
          return np.dot(self.Lr,self.Y)
  
      @cached
      def Ystar(self):
          return np.dot(self.Ystar1(),self.Lc.T)
  
      @cached
      def Yhat(self):
          return self.D*self.Ystar()
  
      @cached
      def Xstar(self):
          RV = np.zeros((self.N*self.P,self.n_fixed_effs))
          ip = 0
          for i in range(self.n_terms):
              Ki = self.A[i].shape[0]*self.F[i].shape[1]
              RV[:,ip:ip+Ki] = np.kron(self.Astar()[i].T,self.Fstar()[i])
              ip += Ki
          return RV
  
      def var_total(self):
          return (self.Yhat()*self.Ystar()).sum()
  
  
      def var_explained(self):
          XKY = self.compute_XKY(M=self.Yhat())
          beta_hat = self.Areml_solve(XKY)
          return (XKY*beta_hat).sum(), beta_hat
  
  
      @cached
      def Xhat(self):
          RV = self.d[:,np.newaxis]*self.Xstar()
          return RV
  
      @cached
      def Areml(self):
          #A1 = self.XstarT_dot(self.Xhat())
          A2 =  self.compute_XKX()
          return A2
  
      @cached
      def Areml_chol(self):
          return LA.cholesky(self.Areml()).T
  
      @cached
      def Areml_REML_chol(self):
          return LA.cholesky(self.Areml()).T
  
      @cached
      def Areml_inv(self):
          return LA.cho_solve((self.Areml_chol(),True),np.eye(self.n_fixed_effs))
  
      #caching bug:
      #@cached
      def beta_hat(self):
          XKY = self.compute_XKY(M=self.Yhat())
          beta_hat = self.Areml_solve(XKY)
          return beta_hat
  
  
      @cached
      def B_hat(self):
          RV = []
          ip = 0
          for term_i in range(self.n_terms):
              RV.append(np.reshape(self.beta_hat()[ip:ip+self.B[term_i].size],self.B[term_i].shape, order='F'))
              ip += self.B[term_i].size
          return RV
  
      @cached
      def LRLdiag_Xhat_tens(self):
          RV  = np.reshape(self.Xhat(),(self.N,self.P,self.n_fixed_effs),order='F').copy()
          RV *= self.LRLdiag[:,np.newaxis,np.newaxis]
          return RV
  
      @cached
      def LRLdiag_Yhat(self):
          return self.LRLdiag[:,np.newaxis]*self.Yhat()
  
      @cached
      def Areml_grad(self):
          RV = np.einsum('jpk,lp->jlk',self.LRLdiag_Xhat_tens(),self.LCL)
          RV = RV.reshape((self.N*self.P,self.n_fixed_effs),order='F')
          RV*= self.d[:,np.newaxis]
          RV = -self.XstarT_dot(RV)
          return RV
  
      @cached
      def beta_grad(self):
          RV  = np.reshape(np.dot(self.LRLdiag_Yhat(),self.LCL.T),(self.N*self.P),order='F')
          RV *= self.d
          RV  = self.XstarT_dot(RV)
          RV += np.dot(self.Areml_grad(),self.beta_hat())
          RV  = -np.dot(self.Areml_inv(),RV)
          return RV
  
      @cached
      def Xstar_beta_grad(self):
          RV = np.zeros((self.N,self.P))
          ip = 0
          for term_i in range(self.n_terms):
              _Bgrad = np.reshape(self.beta_grad()[ip:ip+self.B[term_i].size],self.B[term_i].shape, order='F')
              RV+=np.dot(self.Fstar()[term_i],np.dot(_Bgrad,self.Astar()[term_i]))
              ip += self.B[term_i].size
          return RV
  
  
      @cached
      def Zstar(self):
          """ predict the value of the fixed effect """
          RV = self.Ystar().copy()
          for term_i in range(self.n_terms):
              if self.identity_trick and self.A_identity[term_i]:
                  RV-=np.dot(self.Fstar()[term_i],self.B_hat()[term_i])
              else:
                  RV-=np.dot(self.Fstar()[term_i],np.dot(self.B_hat()[term_i],self.Astar()[term_i]))
          self.clear_cache('DLZ')
          return RV
  
      @cached
      def Areml_eigh(self):
          """compute the eigenvalue decomposition of Astar"""
          s,U = LA.eigh(self.Areml(),lower=True)
          i_pos = (s>1e-10)
          s = s[i_pos]
          U = U[:,i_pos]
          return s,U
  
      @cached
      def DLZ(self):
          return self.Zstar()*np.reshape(self.D,(self.N,self.P), order='F')
  
      ###############################################
      # Other getters with no caching, should not they have caching somehow?
      ###############################################
  
      def Areml_solve(self, b):
          try:
              res = LA.cho_solve((self.Areml_chol(),True),b)
          except LA.LinAlgError:
  
              s,U = self.Areml_eigh()
              res = U.T.dot(b)
              res /= s[:,np.newaxis]
              res = U.dot(res)
  
          return res
  
  
      def compute_XKY(self, M=None):
          if M is None:
              M = self.Yhat()
          assert M.shape==(self.N,self.P)
          XKY = np.zeros((self.n_fixed_effs))
          n_weights = 0
          for term in range(self.n_terms):
              if self.identity_trick and self.A_identity[term]:
                  XKY_block = compute_XYA(DY=M, X=self.Fstar()[term], A=None)
              else:
                  XKY_block = compute_XYA(DY=M, X=self.Fstar()[term], A=self.Astar()[term])
              XKY[n_weights:n_weights + self.A[term].shape[0] * self.F[term].shape[1]] = XKY_block.ravel(order='F')
              n_weights += self.A[term].shape[0] * self.F[term].shape[1]
          return XKY
  
      def compute_XKX(self):
          #n_weights1 = 0
          #
          #for term1 in xrange(self.n_terms):
          #    n_weights1+=self.Astar()[term1].shape[0] * self.Fstar()[term1].shape[1]
          #cov_beta = np.zeros((n_weights1,n_weights1))
          cov_beta = np.zeros((self.n_fixed_effs,self.n_fixed_effs))
          n_weights1 = 0
          for term1 in range(self.n_terms):
              if self.identity_trick and self.A_identity[term1]:
                  A_term1 = None
              else:
                  A_term1 = self.Astar()[term1]
              n_weights2 = n_weights1
              for term2 in range(term1,self.n_terms):
                  if self.identity_trick and self.A_identity[term2]:
                      A_term2 = None
                  else:
                      A_term2 = self.Astar()[term2]
                  block = compute_X1KX2(Y=self.Ystar(), D=self.D, X1=self.Fstar()[term1], X2=self.Fstar()[term2], A1=A_term1, A2=A_term2)
                  cov_beta[n_weights1:n_weights1 + self.A[term1].shape[0] * self.F[term1].shape[1], n_weights2:n_weights2 + self.A[term2].shape[0] * self.F[term2].shape[1]] = block
                  if term1!=term2:
                      cov_beta[n_weights2:n_weights2 + self.A[term2].shape[0] * self.F[term2].shape[1], n_weights1:n_weights1 + self.A[term1].shape[0] * self.F[term1].shape[1]] = block.T
  
                  n_weights2+=self.A[term2].shape[0] * self.F[term2].shape[1]
  
              n_weights1+=self.A[term1].shape[0] * self.F[term1].shape[1]
          return cov_beta
  
      def predict(self):
          """ predict the value of the fixed effect """
          RV = np.zeros((self.N,self.P))
          for term_i in range(self.n_terms):
              RV+=np.dot(self.Fstar()[term_i],np.dot(self.B()[term_i],self.Astar()[term_i]))
          return RV
  
      def evaluate(self):
          """ predict the value of """
          RV  = -self.predict()
          RV += self.Ystar()
          return RV
  
      def getGradient(self,j):
          """ get rotated gradient for fixed effect i """
          i = int(self.indicator['term'][j])
          r = int(self.indicator['row'][j])
          c = int(self.indicator['col'][j])
          rv = -np.kron(self.Fstar()[i][:,[r]],self.Astar()[i][[c],:])
          return rv
  
      def XstarT_dot(self,M):
          """ get dot product of Xhat and M """
          if 0:
              #TODO: implement this properly
              pass
          else:
              RV = np.dot(self.Xstar().T,M)
          return RV
  
      def getResiduals(self):
          """ regress out fixed effects and results residuals """
          X = np.zeros((self.N*self.P,self.n_fixed_effs))
          ip = 0
          for i in range(self.n_terms):
              Ki = self.A[i].shape[0]*self.F[i].shape[1]
              X[:,ip:ip+Ki] = np.kron(self.A[i].T,self.F[i])
              ip += Ki
          y = np.reshape(self.Y,(self.Y.size,1),order='F')
          RV = regressOut(y,X)
          RV = np.reshape(RV,self.Y.shape,order='F')
          return RV
  
      #########################################
      # Params manipulation
      #########################################
  
      def getParams(self):
          """ get params """
          rv = np.array([])
          if self.n_terms>0:
              rv = np.concatenate([np.reshape(self.B[term_i],self.B[term_i].size, order='F') for term_i in range(self.n_terms)])
          return rv
  
      def setParams(self,params):
          """ set params """
          start = 0
          for i in range(self.n_terms):
              n_effects = self.B[i].size
              self.B[i] = np.reshape(params[start:start+n_effects],self.B[i].shape, order='F')
              start += n_effects
  
      #########################################
      # Utility functions
      #########################################
  
      def getDimensions(self):
          """ get phenotype dimensions """
          return self.N,self.P
  
      def _set_toChange(x):
          """ set variables in list x toChange """
          for key in list(x.keys()):
              self.toChange[key] = True
  
      def _update_indicator(self,K,L):
          """ update the indicator """
          _update = {'term': self.n_terms*np.ones((K,L)).T.ravel(),
                      'row': np.kron(np.arange(K)[:,np.newaxis],np.ones((1,L))).T.ravel(),
                      'col': np.kron(np.ones((K,1)),np.arange(L)[np.newaxis,:]).T.ravel()}
          for key in list(_update.keys()):
              self.indicator[key] = np.concatenate([self.indicator[key],_update[key]])
  
      def _rebuild_indicator(self):
          """ update the indicator """
          indicator = {'term':np.array([]),
                       'row':np.array([]),
                       'col':np.array([])}
  
          for term in range(self.n_terms):
              L = self.A[term].shape[0]
              K = self.F[term].shape[1]
              _update = {'term': (term+1)*np.ones((K,L)).T.ravel(),
                      'row': np.kron(np.arange(K)[:,np.newaxis],np.ones((1,L))).T.ravel(),
                      'col': np.kron(np.ones((K,1)),np.arange(L)[np.newaxis,:]).T.ravel()}
              for key in list(_update.keys()):
                  indicator[key] = np.concatenate([indicator[key],_update[key]])
          self.indicator = indicator
  
apis: numpy.ones numpy.eye numpy.zeros numpy.kron numpy.concatenate numpy.reshape numpy.arange numpy.array numpy.dot numpy.outer