add Zscore variance normalisation

impute_jass/impute_jass/stat_models.py

@@ -22,6 +22,19 @@ def compute_var(sig_i_t, sig_t_inv, lamb, batch=True):
 def check_inversion(sig_t, sig_t_inv):
      return np.allclose(sig_t, np.dot(sig_t, np.dot(sig_t_inv, sig_t)))
 
+def var_in_boundaries(var,lamb):
+    """
+    Forces the variance to be in the 0 to 1+lambda boundary
+    theoritically we shouldn't have to do that
+    """
+    id_neg = np.where(var < 0)
+    var_norm = var
+    var[id_neg] = 0
+    id_inf = np.where(var > (1+lamb))
+    var[id_inf] = 1+lamb
+
+    return var
+
 def impg_model(zt, sig_t, sig_i_t, lamb=0.01, rcond=0.01,  batch=True):
     """
     Argument:
@@ -39,8 +52,14 @@ def impg_model(zt, sig_t, sig_i_t, lamb=0.01, rcond=0.01,  batch=True):
         correct_inversion = check_inversion(sig_t, sig_t_inv)
 
     var, ld_score = compute_var(sig_i_t, sig_t_inv, lamb, batch)
+
+
     mu = compute_mu(sig_i_t, sig_t_inv, zt)
+
+
     if np.any(mu > 30):
-        print("ABERANT SNP SNiP ")
-    #mu = mu / (((1+lamb)-var)**0.5)
+        print("ABERANT SNP SNiP")
+    var_norm = var_in_boundaries(var, lamb)
+    mu = mu / np.sqrt(var_norm)
+
     return({"var" : var, "mu" : mu, "ld_score" : ld_score, "condition_number" : condition_number, "correct_inversion":correct_inversion })