Add docstrings to indicate functions exposed via API

lshmm/fb_diploid.py

@@ -234,7 +234,11 @@ def backward_ls_dip_starting_point(n, m, G, s, e, r):
 
 @jit.numba_njit
 def forward_ls_dip_loop(n, m, G, s, e, r, norm=True):
-    """An implementation without vectorisation."""
+    """
+    An implementation without vectorisation.
+
+    This is exposed via the API.
+    """
     # Initialise
     F = np.zeros((m, n, n))
     for j1 in range(n):
@@ -316,7 +320,11 @@ def forward_ls_dip_loop(n, m, G, s, e, r, norm=True):
 
 @jit.numba_njit
 def backward_ls_dip_loop(n, m, G, s, e, c, r):
-    """An implementation without vectorisation."""
+    """
+    An implementation without vectorisation.
+
+    This is exposed via the API.
+    """
     # Initialise
     B = np.zeros((m, n, n))
     B[m - 1, :, :] = 1

lshmm/fb_haploid.py

@@ -11,7 +11,11 @@
 
 @jit.numba_njit
 def forwards_ls_hap(n, m, H, s, e, r, norm=True):
-    """A matrix-based implementation using Numpy."""
+    """
+    A matrix-based implementation using Numpy vectorisation.
+
+    This is exposed via the API.
+    """
     F = np.zeros((m, n))
     r_n = r / n
 
@@ -66,7 +70,11 @@ def forwards_ls_hap(n, m, H, s, e, r, norm=True):
 
 @jit.numba_njit
 def backwards_ls_hap(n, m, H, s, e, c, r):
-    """A matrix-based implementation using Numpy."""
+    """
+    A matrix-based implementation using Numpy vectorisation.
+
+    This is exposed via the API.
+    """
     B = np.zeros((m, n))
     for i in range(n):
         B[m - 1, i] = 1

lshmm/vit_diploid.py

@@ -109,7 +109,11 @@ def forwards_viterbi_dip_naive_low_mem(n, m, G, s, e, r):
 
 @jit.numba_njit
 def forwards_viterbi_dip_low_mem(n, m, G, s, e, r):
-    """An implementation with reduced memory."""
+    """
+    An implementation with reduced memory.
+
+    This is exposed via the API.
+    """
     # Initialise
     V = np.zeros((n, n))
     V_prev = np.zeros((n, n))
@@ -357,7 +361,11 @@ def forwards_viterbi_dip_naive_full_vec(n, m, G, s, e, r):
 
 @jit.numba_njit
 def backwards_viterbi_dip(m, V_last, P):
-    """Run a backwards pass to determine the most likely path."""
+    """
+    Run a backwards pass to determine the most likely path.
+
+    This is exposed via the API.
+    """
     assert V_last.ndim == 2
     assert V_last.shape[0] == V_last.shape[1]
 
@@ -418,12 +426,17 @@ def backwards_viterbi_dip_no_pointer(
 
 
 def get_phased_path(n, path):
+    """This is exposed via the API."""
     return np.unravel_index(path, (n, n))
 
 
 @jit.numba_njit
 def path_ll_dip(n, m, G, phased_path, s, e, r):
-    """Evaluate log-likelihood path through a reference panel which results in sequence."""
+    """
+    Evaluate log-likelihood path through a reference panel which results in sequence.
+
+    This is exposed via the API.
+    """
     emission_index = core.get_index_in_emission_matrix_diploid(
         ref_allele=G[0, phased_path[0][0], phased_path[1][0]], query_allele=s[0, 0]
     )

lshmm/vit_haploid.py

@@ -174,7 +174,11 @@ def forwards_viterbi_hap_low_mem_rescaling(n, m, H, s, e, r):
 
 @jit.numba_njit
 def forwards_viterbi_hap_lower_mem_rescaling(n, m, H, s, e, r):
-    """An implementation with even smaller memory footprint that exploits the Markov structure."""
+    """
+    An implementation with even smaller memory footprint that exploits the Markov structure.
+
+    This is exposed via the API.
+    """
     V = np.zeros(n)
     for i in range(n):
         emission_idx = core.get_index_in_emission_matrix_haploid(
@@ -251,7 +255,11 @@ def forwards_viterbi_hap_lower_mem_rescaling_no_pointer(n, m, H, s, e, r):
 # Speedier version, variants x samples
 @jit.numba_njit
 def backwards_viterbi_hap(m, V_last, P):
-    """Run a backwards pass to determine the most likely path."""
+    """
+    Run a backwards pass to determine the most likely path.
+
+    This is exposed via API.
+    """
     assert len(V_last.shape) == 1
     path = np.zeros(m, dtype=np.int64)
     path[m - 1] = np.argmax(V_last)
@@ -279,7 +287,11 @@ def backwards_viterbi_hap_no_pointer(m, V_argmaxes, recombs):
 
 @jit.numba_njit
 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."""
+    """
+    Evaluate the log-likelihood of a path through a reference panel resulting in a sequence.
+
+    This is exposed via the API.
+    """
     emission_idx = core.get_index_in_emission_matrix_haploid(
         ref_allele=H[0, path[0]], query_allele=s[0, 0]
     )