diff --git a/lshmm/core.py b/lshmm/core.py index 95b106c..dd822a1 100644 --- a/lshmm/core.py +++ b/lshmm/core.py @@ -71,8 +71,4 @@ def get_index_in_emission_prob_matrix_diploid(ref_allele, query_allele): is_allele_match = ref_allele == query_allele is_ref_one = ref_allele == 1 is_query_one = query_allele == 1 - return ( - 4 * is_allele_match - + 2 * is_ref_one - + is_query_one - ) + return 4 * is_allele_match + 2 * is_ref_one + is_query_one diff --git a/lshmm/vit_diploid.py b/lshmm/vit_diploid.py index 4937b02..e1c0304 100644 --- a/lshmm/vit_diploid.py +++ b/lshmm/vit_diploid.py @@ -21,8 +21,7 @@ def forwards_viterbi_dip_naive(n, m, G, s, e, r): for j1 in range(n): for j2 in range(n): emission_index = core.get_index_in_emission_prob_matrix_diploid( - ref_allele=G[0, j1, j2], - query_allele=s[0, 0] + ref_allele=G[0, j1, j2], query_allele=s[0, 0] ) V[0, j1, j2] = 1 / (n**2) * e[0, emission_index] @@ -74,8 +73,7 @@ def forwards_viterbi_dip_naive_low_mem(n, m, G, s, e, r): for j1 in range(n): for j2 in range(n): emission_index = core.get_index_in_emission_prob_matrix_diploid( - ref_allele=G[0, j1, j2], - query_allele=s[0, 0] + ref_allele=G[0, j1, j2], query_allele=s[0, 0] ) V_prev[j1, j2] = 1 / (n**2) * e[0, emission_index] @@ -129,8 +127,7 @@ def forwards_viterbi_dip_low_mem(n, m, G, s, e, r): for j1 in range(n): for j2 in range(n): emission_index = core.get_index_in_emission_prob_matrix_diploid( - ref_allele=G[0, j1, j2], - query_allele=s[0, 0] + ref_allele=G[0, j1, j2], query_allele=s[0, 0] ) V_prev[j1, j2] = 1 / (n**2) * e[0, emission_index] @@ -218,8 +215,7 @@ def forwards_viterbi_dip_low_mem_no_pointer(n, m, G, s, e, r): for j1 in range(n): for j2 in range(n): emission_index = core.get_index_in_emission_prob_matrix_diploid( - ref_allele=G[0, j1, j2], - query_allele=s[0, 0] + ref_allele=G[0, j1, j2], query_allele=s[0, 0] ) V_prev[j1, j2] = 1 / (n**2) * e[0, emission_index] @@ -300,8 +296,7 @@ def forwards_viterbi_dip_naive_vec(n, m, G, s, e, r): for j1 in range(n): for j2 in range(n): emission_index = core.get_index_in_emission_prob_matrix_diploid( - ref_allele=G[0, j1, j2], - query_allele=s[0, 0] + ref_allele=G[0, j1, j2], query_allele=s[0, 0] ) V[0, j1, j2] = 1 / (n**2) * e[0, emission_index] @@ -453,8 +448,7 @@ def get_phased_path(n, path): def path_ll_dip(n, m, G, phased_path, s, e, r): """Evaluate log-likelihood path through a reference panel which results in sequence s.""" emission_index = core.get_index_in_emission_prob_matrix_diploid( - ref_allele=G[0, phased_path[0][0], phased_path[1][0]], - query_allele=s[0, 0] + ref_allele=G[0, phased_path[0][0], phased_path[1][0]], query_allele=s[0, 0] ) log_prob_path = np.log10(1 / (n**2) * e[0, emission_index]) old_phase = np.array([phased_path[0][0], phased_path[1][0]]) diff --git a/lshmm/vit_haploid.py b/lshmm/vit_haploid.py index 0904601..f0d537d 100644 --- a/lshmm/vit_haploid.py +++ b/lshmm/vit_haploid.py @@ -18,8 +18,7 @@ def viterbi_naive_init(n, m, H, s, e, r): for i in range(n): emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[0, i], - query_allele=s[0, 0] + ref_allele=H[0, i], query_allele=s[0, 0] ) V[0, i] = 1 / n * e[0, emission_idx] @@ -36,8 +35,7 @@ def viterbi_init(n, m, H, s, e, r): for i in range(n): emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[0, i], - query_allele=s[0, 0] + ref_allele=H[0, i], query_allele=s[0, 0] ) V_prev[i] = 1 / n * e[0, emission_idx] @@ -54,8 +52,7 @@ def forwards_viterbi_hap_naive(n, m, H, s, e, r): v = np.zeros(n) for k in range(n): emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[j, i], - query_allele=s[0, j] + ref_allele=H[j, i], query_allele=s[0, j] ) v[k] = V[j - 1, k] * e[j, emission_idx] if k == i: @@ -81,8 +78,7 @@ def forwards_viterbi_hap_naive_vec(n, m, H, s, e, r): v = np.copy(v_tmp) v[i] += V[j - 1, i] * (1 - r[j]) emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[j, i], - query_allele=s[0, j] + ref_allele=H[j, i], query_allele=s[0, j] ) v *= e[j, emission_idx] P[j, i] = np.argmax(v) @@ -103,8 +99,7 @@ def forwards_viterbi_hap_naive_low_mem(n, m, H, s, e, r): v = np.zeros(n) for k in range(n): emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[j, i], - query_allele=s[0, j] + ref_allele=H[j, i], query_allele=s[0, j] ) v[k] = V_prev[k] * e[j, emission_idx] if k == i: @@ -133,8 +128,7 @@ def forwards_viterbi_hap_naive_low_mem_rescaling(n, m, H, s, e, r): v = np.zeros(n) for k in range(n): emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[j, i], - query_allele=s[0, j] + ref_allele=H[j, i], query_allele=s[0, j] ) v[k] = V_prev[k] * e[j, emission_idx] if k == i: @@ -168,8 +162,7 @@ def forwards_viterbi_hap_low_mem_rescaling(n, m, H, s, e, r): V[i] = r_n[j] P[j, i] = argmax emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[j, i], - query_allele=s[0, j] + ref_allele=H[j, i], query_allele=s[0, j] ) V[i] *= e[j, emission_idx] V_prev = np.copy(V) @@ -185,8 +178,7 @@ def forwards_viterbi_hap_lower_mem_rescaling(n, m, H, s, e, r): V = np.zeros(n) for i in range(n): emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[0, i], - query_allele=s[0, 0] + ref_allele=H[0, i], query_allele=s[0, 0] ) V[i] = 1 / n * e[0, emission_idx] P = np.zeros((m, n), dtype=np.int64) @@ -204,8 +196,7 @@ def forwards_viterbi_hap_lower_mem_rescaling(n, m, H, s, e, r): V[i] = r_n[j] P[j, i] = argmax emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[j, i], - query_allele=s[0, j] + ref_allele=H[j, i], query_allele=s[0, j] ) V[i] *= e[j, emission_idx] @@ -220,8 +211,7 @@ def forwards_viterbi_hap_lower_mem_rescaling_no_pointer(n, m, H, s, e, r): V = np.zeros(n) for i in range(n): emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[0, i], - query_allele=s[0, 0] + ref_allele=H[0, i], query_allele=s[0, 0] ) V[i] = 1 / n * e[0, emission_idx] r_n = r / n @@ -286,8 +276,7 @@ def backwards_viterbi_hap_no_pointer(m, V_argmaxes, recombs): def path_ll_hap(n, m, H, path, s, e, r): """Evaluate the log-likelihood of a path through a reference panel resulting in a sequence.""" emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[0, path[0]], - query_allele=s[0, 0] + ref_allele=H[0, path[0]], query_allele=s[0, 0] ) log_prob_path = np.log10((1 / n) * e[0, emission_idx]) old = path[0] @@ -295,8 +284,7 @@ def path_ll_hap(n, m, H, path, s, e, r): for l in range(1, m): emission_idx = core.get_index_in_emission_prob_matrix( - ref_allele=H[l, path[l]], - query_allele=s[0, l] + ref_allele=H[l, path[l]], query_allele=s[0, l] ) current = path[l] same = old == current