diff --git a/impute_jass/impute_jass/stat_models.py b/impute_jass/impute_jass/stat_models.py new file mode 100644 index 0000000000000000000000000000000000000000..9937442f43abac2b7ff1eabae4ec3ff2a99f9a34 --- /dev/null +++ b/impute_jass/impute_jass/stat_models.py @@ -0,0 +1,22 @@ +""" +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))