Skip to content

Commit

Permalink
Reformat
Browse files Browse the repository at this point in the history
  • Loading branch information
szhan committed Apr 19, 2024
1 parent fd011ca commit 1cbfc64
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 41 deletions.
6 changes: 1 addition & 5 deletions lshmm/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
18 changes: 6 additions & 12 deletions lshmm/vit_diploid.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down Expand Up @@ -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]

Expand Down Expand Up @@ -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]

Expand Down Expand Up @@ -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]

Expand Down Expand Up @@ -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]

Expand Down Expand Up @@ -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]])
Expand Down
36 changes: 12 additions & 24 deletions lshmm/vit_haploid.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand All @@ -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]

Expand All @@ -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:
Expand All @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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]

Expand All @@ -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
Expand Down Expand Up @@ -286,17 +276,15 @@ 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]
r_n = r / n

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
Expand Down

0 comments on commit 1cbfc64

Please sign in to comment.