testing stats_model

impute_jass/impute_jass/__init__.py

 import impute_jass.ld_matrix as LD
+import impute_jass.stat_models as model

impute_jass/impute_jass/stat_models.py

+"""
+function for SNP imputation
+
+"""
+import numpy as np
+
+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, Zt))
+    return({"Var":Var, "mu":mu})