stat_models.py 950 Bytes
Newer Older
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
1
2
3
4
"""
function for SNP imputation

"""
hjulienne's avatar
hjulienne committed
5
import numpy as np
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
6

7
def ImpG_model_batch(Zt, Sig_t, Sig_i_t):
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
8
9
10
11
12
13
14
15
16
17
18
19
20
    """
    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())

hjulienne's avatar
hjulienne committed
21
22
    mu = np.dot(Sig_i_t, np.dot(Sig_t_inv, Zt))
    return({"Var":Var, "mu":mu})
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41


def ImpG_model_snp(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.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))
    return({"Var":Var, "mu":mu})