Checked statistic computation on correctly indexed matrices

impute_jass/impute_jass/stat_models.py

@@ -11,14 +11,15 @@ def ImpG_model_batch(Zt, Sig_t, Sig_i_t):
 
     """
 
-    np.fill_diagonal(Sig_t.values, 1.01)
-    Sig_t.fillna(0, inplace=True)
+    #np.fill_diagonal(Sig_t.values, 1.01)
+    #Sig_t.fillna(0, inplace=True)
 
     Sig_t_inv =np.linalg.inv(Sig_t)
 
     Var = np.diag(Sig_t)[0] - np.einsum('ij,jk,ki->i', Sig_i_t, Sig_t_inv ,Sig_i_t.transpose())
 
     mu = np.dot(Sig_i_t, np.dot(Sig_t_inv, Zt))
+    
     return({"Var":Var, "mu":mu})
 
 
@@ -28,14 +29,13 @@ def ImpG_model_snp(Zt, Sig_t, Sig_i_t):
         Zt : (vector) the vector of known Z scores
 
     """
-
-    np.fill_diagonal(Sig_t.values, 1.01)
-    Sig_t.fillna(0, inplace=True)
+    #np.fill_diagonal(Sig_t.values, 1.01)
+    #Sig_t.fillna(0, inplace=True)
 
     Sig_t_inv =np.linalg.inv(Sig_t)
-
     Var = np.diag(Sig_t)[0] - np.dot(Sig_i_t, np.dot(Sig_t_inv, Sig_i_t.transpose()))
     #np.einsum('ij,jk,ki->i', Sig_i_t, Sig_t_inv ,Sig_i_t.transpose())
 
     mu = np.dot(Sig_i_t, np.dot(Sig_t_inv, Zt))
+    mu = mu / ((1-Var)**0.5)
     return({"Var":Var, "mu":mu})