Skip to content

Commit

Permalink
Merge pull request #13 from astheeggeggs/fix-lint
Browse files Browse the repository at this point in the history 
Fixed lint errors, cleaned up the api and ensured all tests pass.
  • Loading branch information
@astheeggeggs
astheeggeggs authored Aug 30, 2022
2 parents b565924 + c21a1b8 commit 9ff2f8a
Show file tree
Hide file tree
Showing 16 changed files with 174 additions and 1,326 deletions.
384 changes: 123 additions & 261 deletions lshmm/api.py
View file

Large diffs are not rendered by default.

12 changes: 6 additions & 6 deletions lshmm/forward_backward/fb_diploid_samples_variants.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ def np_argmax(array, axis):
return np_apply_along_axis(np.argmax, axis, array)


@nb.jit
@nb.njit
def forwards_ls_dip(n, m, G, s, e, r, norm=True):
"""Matrix based diploid LS forward algorithm using numpy vectorisation."""
# Initialise the forward tensor
Expand Down Expand Up @@ -121,7 +121,7 @@ def forwards_ls_dip(n, m, G, s, e, r, norm=True):
return F, c, ll


@nb.jit
@nb.njit
def backwards_ls_dip(n, m, G, s, e, c, r):
"""Matrix based diploid LS backward algorithm using numpy vectorisation."""
# Initialise the backward tensor
Expand Down Expand Up @@ -161,7 +161,7 @@ def backwards_ls_dip(n, m, G, s, e, c, r):
return B


@nb.jit
@nb.njit
def forward_ls_dip_starting_point(n, m, G, s, e, r):
"""Naive implementation of LS diploid forwards algorithm."""
# Initialise the forward tensor
Expand Down Expand Up @@ -234,7 +234,7 @@ def forward_ls_dip_starting_point(n, m, G, s, e, r):
return F, ll


@nb.jit
@nb.njit
def backward_ls_dip_starting_point(n, m, G, s, e, r):
"""Naive implementation of LS diploid backwards algorithm."""
# Backwards
Expand Down Expand Up @@ -310,7 +310,7 @@ def backward_ls_dip_starting_point(n, m, G, s, e, r):
return B


@nb.jit
@nb.njit
def forward_ls_dip_loop(n, m, G, s, e, r, norm=True):
"""LS diploid forwards algoritm without vectorisation."""
# Initialise the forward tensor
Expand Down Expand Up @@ -422,7 +422,7 @@ def forward_ls_dip_loop(n, m, G, s, e, r, norm=True):
return F, c, ll


@nb.jit
@nb.njit
def backward_ls_dip_loop(n, m, G, s, e, c, r):
"""LS diploid backwards algoritm without vectorisation."""
# Initialise the backward tensor
Expand Down
273 changes: 0 additions & 273 deletions lshmm/forward_backward/fb_diploid_variants_samples.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -148,94 +148,13 @@ def backwards_ls_dip(n, m, G, s, e, c, r):
)

# One changes
# sum_j1 = np.tile(np.sum(B[l+1,:,:] * e[l+1, index], 0, keepdims=True), (n,1))
# sum_j1 = np_sum(B[l + 1, :, :], 0).repeat(n).reshape((-1, n)).T
# sum_j2 = np.tile(np.sum(B[l+1,:,:] * e[l+1, index], 1, keepdims=True), (1,n))
# sum_j2 = np_sum(B[l + 1, :, :], 1).repeat(n).reshape((-1, n))
sum_j = np_sum(B[l + 1, :, :] * e[l + 1, index], 0).repeat(n).reshape((-1, n))
B[l, :, :] += ((1 - r[l + 1]) * r_n[l + 1]) * (sum_j + sum_j.T)
B[l, :, :] *= 1 / c[l + 1]

return B


# def forward_ls_dip_starting_point(n, m, G, s, e, r):
# '''
# Unbelievably naive implementation of LS diploid forwards. Just to get something down
# that works.
# '''

# # Initialise the forward tensor
# F = np.zeros((m,n,n))
# F[0,:,:] = 1/(n**2)
# index = (
# 4*np.equal(G[0,:,:], s[0,0]).astype(np.int64) +
# 2*(G[0,:,:] == 1).astype(np.int64) +
# np.int64(s[0,0] == 1)
# )
# F[0,:,:] *= e[0, index]
# r_n = r/n

# for l in range(1,m):

# # Determine the various components
# F_no_change = np.zeros((n,n))
# F_j1_change = np.zeros(n)
# F_j2_change = np.zeros(n)
# F_both_change = 0

# for j1 in range(n):
# for j2 in range(n):
# F_no_change[j1, j2] = (1-r[l])**2 * F[l-1, j1, j2]

# for j1 in range(n):
# for j2 in range(n):
# F_both_change += r_n[l]**2 * F[l-1, j1, j2]

# for j1 in range(n):
# for j2 in range(n): # This is the variable to sum over - it changes
# F_j2_change[j1] += (1 - r[l]) * r_n[l] * F[l-1, j1, j2]

# for j2 in range(n):
# for j1 in range(n): # This is the variable to sum over - it changes
# F_j1_change[j2] += (1 - r[l]) * r_n[l] * F[l-1, j1, j2]

# F[l,:,:] = F_both_change

# for j1 in range(n):
# F[l, j1, :] += F_j2_change

# for j2 in range(n):
# F[l, :, j2] += F_j1_change

# for j1 in range(n):
# for j2 in range(n):
# F[l, j1, j2] += F_no_change[j1, j2]

# for j1 in range(n):
# for j2 in range(n):
# # What is the emission?
# if s[0,l] == 1:
# # OBS is het
# if G[l, j1, j2] == 1: # REF is het
# F[l, j1, j2] *= e[l,BOTH_HET]
# else: # REF is hom
# F[l, j1, j2] *= e[l,REF_HOM_OBS_HET]
# else:
# # OBS is hom
# if G[l, j1, j2] == 1: # REF is het
# F[l, j1, j2] *= e[l,REF_HET_OBS_HOM]
# else: # REF is hom
# if G[l, j1, j2] == s[0,l]: # Equal
# F[l, j1, j2] *= e[l,EQUAL_BOTH_HOM]
# else: # Unequal
# F[l, j1, j2] *= e[l,UNEQUAL_BOTH_HOM]

# ll = np.log10(np.sum(F[l,:,:]))

# return F, ll


@nb.njit
def forward_ls_dip_starting_point(n, m, G, s, e, r):
"""Naive implementation of LS diploid forwards algorithm."""
Expand Down Expand Up @@ -312,79 +231,6 @@ def forward_ls_dip_starting_point(n, m, G, s, e, r):
return F, ll


# def backward_ls_dip_starting_point(n, m, G, s, e, r):
# '''
# Unbelievably naive implementation of LS diploid backwards. Just to get something down
# that works.
# '''

# # Backwards
# B = np.zeros((m,n,n))

# # Initialise
# B[m-1, :, :] = 1
# r_n = r/n

# for l in range(m-2, -1, -1):

# # Determine the various components
# B_no_change = np.zeros((n,n))
# B_j1_change = np.zeros(n)
# B_j2_change = np.zeros(n)
# B_both_change = 0

# # Evaluate the emission matrix at this site, for all pairs
# e_tmp = np.zeros((n,n))
# for j1 in range(n):
# for j2 in range(n):
# # What is the emission?
# if s[0,l+1] == 1:
# # OBS is het
# if G[l+1, j1, j2] == 1: # REF is het
# e_tmp[j1, j2] = e[l+1, BOTH_HET]
# else: # REF is hom
# e_tmp[j1, j2] = e[l+1, REF_HOM_OBS_HET]
# else:
# # OBS is hom
# if G[l+1, j1, j2] == 1: # REF is het
# e_tmp[j1, j2] = e[l+1,REF_HET_OBS_HOM]
# else: # REF is hom
# if G[l+1, j1, j2] == s[0,l+1]: # Equal
# e_tmp[j1, j2] = e[l+1,EQUAL_BOTH_HOM]
# else: # Unequal
# e_tmp[j1, j2] = e[l+1,UNEQUAL_BOTH_HOM]

# for j1 in range(n):
# for j2 in range(n):
# B_no_change[j1, j2] = (1-r[l+1])**2 * B[l+1,j1,j2] * e_tmp[j1, j2]

# for j1 in range(n):
# for j2 in range(n):
# B_both_change += r_n[l+1]**2 * e_tmp[j1, j2] * B[l+1,j1,j2]

# for j1 in range(n):
# for j2 in range(n): # This is the variable to sum over - it changes
# B_j2_change[j1] += (1 - r[l+1]) * r_n[l+1] * B[l+1,j1,j2] * e_tmp[j1, j2]

# for j2 in range(n):
# for j1 in range(n): # This is the variable to sum over - it changes
# B_j1_change[j2] += (1 - r[l+1]) * r_n[l+1] * B[l+1,j1,j2] * e_tmp[j1, j2]

# B[l,:,:] = B_both_change

# for j1 in range(n):
# B[l, j1, :] += B_j2_change

# for j2 in range(n):
# B[l, :, j2] += B_j1_change

# for j1 in range(n):
# for j2 in range(n):
# B[l, j1, j2] += B_no_change[j1, j2]

# return B


@nb.njit
def backward_ls_dip_starting_point(n, m, G, s, e, r):
"""Naive implementation of LS diploid backwards algorithm."""
Expand Down Expand Up @@ -461,68 +307,6 @@ def backward_ls_dip_starting_point(n, m, G, s, e, r):
return B


# def forward_ls_dip_loop(n, m, G, s, e, r):
# '''
# LS diploid forwards with lots of loops.
# '''

# # Initialise the forward tensor
# F = np.zeros((m,n,n))
# F[0,:,:] = 1/(n**2)
# index = (
# 4*np.equal(G[0,:,:], s[0,0]).astype(np.int64) +
# 2*(G[0,:,:] == 1).astype(np.int64) +
# np.int64(s[0,0] == 1)
# )
# F[0,:,:] *= e[0, index]
# r_n = r/n

# for l in range(1,m):

# # Determine the various components
# F_no_change = np.zeros((n,n))
# F_j1_change = np.zeros(n)
# F_j2_change = np.zeros(n)
# F_both_change = 0

# for j1 in range(n):
# for j2 in range(n):
# F_no_change[j1, j2] = (1-r[l])**2 * F[l-1,j1,j2]
# F_j1_change[j1] += (1 - r[l]) * r_n[l] * F[l-1,j2,j1]
# F_j2_change[j1] += (1 - r[l]) * r_n[l] * F[l-1,j1,j2]
# F_both_change += r_n[l]**2 * F[l-1,j1,j2]

# F[l,:,:] = F_both_change

# for j1 in range(n):
# F[l, j1, :] += F_j2_change
# F[l, :, j1] += F_j1_change
# for j2 in range(n):
# F[l, j1, j2] += F_no_change[j1, j2]

# for j1 in range(n):
# for j2 in range(n):
# # What is the emission?
# if s[0,l] == 1:
# # OBS is het
# if G[l, j1, j2] == 1: # REF is het
# F[l, j1, j2] *= e[l, BOTH_HET]
# else: # REF is hom
# F[l, j1, j2] *= e[l, REF_HOM_OBS_HET]
# else:
# # OBS is hom
# if G[l, j1, j2] == 1: # REF is het
# F[l, j1, j2] *= e[l, REF_HET_OBS_HOM]
# else: # REF is hom
# if G[l, j1, j2] == s[0,l]: # Equal
# F[l, j1, j2] *= e[l, EQUAL_BOTH_HOM]
# else: # Unequal
# F[l, j1, j2] *= e[l, UNEQUAL_BOTH_HOM]

# ll = np.log10(np.sum(F[l,:,:]))
# return F, ll


@nb.njit
def forward_ls_dip_loop(n, m, G, s, e, r, norm=True):
"""LS diploid forwards algoritm without vectorisation."""
Expand Down Expand Up @@ -637,63 +421,6 @@ def forward_ls_dip_loop(n, m, G, s, e, r, norm=True):
return F, c, ll


# def backward_ls_dip_loop(n, m, G, s, e, r):
# '''
# LS diploid backwards with lots of loops.
# '''

# # Initialise the backward tensor
# B = np.zeros((m,n,n))
# B[m-1, :, :] = 1
# r_n = r/n

# for l in range(m-2, -1, -1):

# # Determine the various components
# B_no_change = np.zeros((n,n))
# B_j1_change = np.zeros(n)
# B_j2_change = np.zeros(n)
# B_both_change = 0

# # Evaluate the emission matrix at this site, for all pairs
# e_tmp = np.zeros((n,n))
# for j1 in range(n):
# for j2 in range(n):
# # What is the emission?
# if s[0,l+1] == 1:
# # OBS is het
# if G[l+1, j1, j2] == 1: # REF is het
# e_tmp[j1, j2] = e[l+1, BOTH_HET]
# else: # REF is hom
# e_tmp[j1, j2] = e[l+1, REF_HOM_OBS_HET]
# else:
# # OBS is hom
# if G[l+1, j1, j2] == 1: # REF is het
# e_tmp[j1, j2] = e[l+1, REF_HET_OBS_HOM]
# else: # REF is hom
# if G[l+1, j1, j2] == s[0,l+1]: # Equal
# e_tmp[j1, j2] = e[l+1, EQUAL_BOTH_HOM]
# else: # Unequal
# e_tmp[j1, j2] = e[l+1, UNEQUAL_BOTH_HOM]

# for j1 in range(n):
# for j2 in range(n):
# B_no_change[j1, j2] = (1-r[l+1])**2 * B[l+1,j1,j2] * e_tmp[j1, j2]
# B_j2_change[j1] += (1 - r[l+1]) * r_n[l+1] * B[l+1,j1,j2] * e_tmp[j1, j2]
# B_j1_change[j1] += (1 - r[l+1]) * r_n[l+1] * B[l+1,j2,j1] * e_tmp[j2, j1]
# B_both_change += r_n[l+1]**2 * e_tmp[j1, j2] * B[l+1,j1,j2]

# B[l,:,:] = B_both_change

# for j1 in range(n):
# B[l, j1, :] += B_j2_change
# B[l, :, j1] += B_j1_change
# for j2 in range(n):
# B[l, j1, j2] += B_no_change[j1, j2]

# return B


@nb.njit
def backward_ls_dip_loop(n, m, G, s, e, c, r):
"""LS diploid backwards algoritm without vectorisation."""
Expand Down
4 changes: 2 additions & 2 deletions lshmm/forward_backward/fb_haploid_samples_variants.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import numpy as np


@nb.jit
@nb.njit
def forwards_ls_hap(n, m, H, s, e, r, norm=True):
"""Matrix based haploid LS forward algorithm using numpy vectorisation."""
# Initialise
Expand Down Expand Up @@ -39,7 +39,7 @@ def forwards_ls_hap(n, m, H, s, e, r, norm=True):
return F, c, ll


@nb.jit
@nb.njit
def backwards_ls_hap(n, m, H, s, e, c, r):
"""Matrix based haploid LS backward algorithm using numpy vectorisation."""
# Initialise
Expand Down
Loading

0 comments on commit 9ff2f8a

Please sign in to comment.