stat_models.py 1.52 KB
Newer Older
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
1
2
3
"""
function for SNP imputation
"""
hjulienne's avatar
hjulienne committed
4
import numpy as np
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
5
import scipy as sc
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
6

Hanna  JULIENNE's avatar
Hanna JULIENNE committed
7
8
def compute_mu(sig_i_t, sig_t_inv, zt):
    return np.dot(sig_i_t, np.dot(sig_t_inv, zt))
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
9

Hanna  JULIENNE's avatar
Hanna JULIENNE committed
10
11
12
def compute_var(sig_i_t, sig_t_inv, lamb, batch=True):
    if batch:
        var = (1 + lamb) - np.einsum('ij,jk,ki->i', sig_i_t, sig_t_inv ,sig_i_t.transpose())
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
13
14
        ld_score = (sig_i_t**2).sum(1)

Hanna  JULIENNE's avatar
Hanna JULIENNE committed
15
16
    else:
        var = (1 + lamb) - np.dot(sig_i_t, np.dot(sig_t_inv, sig_i_t.transpose()))
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
17
18
19
20
21
22
        ld_score = (sig_i_t**2).sum()

    return var, ld_score

def check_inversion(sig_t, sig_t_inv):
     return np.allclose(sig_t, np.dot(sig_t, np.dot(sig_t_inv, sig_t)))
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
23

Hanna  JULIENNE's avatar
Hanna JULIENNE committed
24
def impg_model(zt, sig_t, sig_i_t, lamb=0.01, rcond=0.01,  batch=True):
25
26
    """
    Argument:
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
27
        zt : (vector) the vector of known Z scores
28
    """
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
29
30
    sig_t = sig_t.values
    np.fill_diagonal(sig_t, (1+lamb))
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
31
    sig_t_inv = sc.linalg.pinv(sig_t, rcond=rcond)
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
32
33
34
35
36
37
38
39
40

    if batch:
        condition_number = np.array([np.linalg.cond(sig_t)]*sig_i_t.shape[0])
        correct_inversion = np.array([check_inversion(sig_t, sig_t_inv)]*sig_i_t.shape[0])
    else:
        condition_number = np.linalg.cond(sig_t)
        correct_inversion = check_inversion(sig_t, sig_t_inv)

    var, ld_score = compute_var(sig_i_t, sig_t_inv, lamb, batch)
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
41
    mu = compute_mu(sig_i_t, sig_t_inv, zt)
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
42
    if np.any(mu > 30):
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
43
        print("ABERANT SNP SNiP ")
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
44
    #mu = mu / (((1+lamb)-var)**0.5)
Hanna  JULIENNE's avatar
Hanna JULIENNE committed
45
    return({"var" : var, "mu" : mu, "ld_score" : ld_score, "condition_number" : condition_number, "correct_inversion":correct_inversion })