Impg model function

impute_jass/impute_jass/stat_models.py

+"""
+function for SNP imputation
+
+"""
+import numpy
+
+def ImpG_model(Zt, Sig_t, Sig_i_t):
+    """
+    Argument:
+        Zt : (vector) the vector of known Z scores
+
+    """
+
+    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, Zfile.loc[index_known, 'Z']))
+    return((Var, mu))