From 46551e196c4d1ed4792ac5120be61473bcadb278 Mon Sep 17 00:00:00 2001
From: szhan <shing.zhan@gmail.com>
Date: Tue, 2 Jul 2024 17:21:00 +0100
Subject: [PATCH] Update

---
 lshmm/core.py                          | 10 +++++-----
 tests/test_nontree_vit_haploid_tstv.py | 12 ++++++------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/lshmm/core.py b/lshmm/core.py
index 1f54ff8..7978a10 100644
--- a/lshmm/core.py
+++ b/lshmm/core.py
@@ -303,19 +303,19 @@ def get_emission_probability_haploid(ref_allele, query_allele, site, emission_ma
 
 
 @jit.numba_njit
-def get_emission_matrix_hkylike(mu, kappa=None):
+def get_emission_matrix_haploid_hkylike(mu, kappa=None):
     """
-    Return an emission probability matrix that allows for
-    differences between transition and transversion rates.
+    Return an emission probability matrix that allows for mutational bias
+    towards transitions or transversions.
 
-    When `kappa` is set to None, it is equivalent to setting it to 1.
+    When `kappa` is set to None, it defaults to 1.
 
     :param float mu: Probability of mutation to any allele.
     :param float kappa: Transition-to-transversion rate ratio.
     """
     if kappa is not None:
         if kappa <= 0:
-            raise ValueError("Ts/tv ratio must be positive.")
+            raise ValueError("Transition-to-transversion rate ratio must be positive.")
     num_sites = len(mu)
     num_alleles = 4  # Assume that ACGT are encoded as 0 to 3.
     emission_matrix = (
diff --git a/tests/test_nontree_vit_haploid_tstv.py b/tests/test_nontree_vit_haploid_tstv.py
index ea655b8..ad0d56c 100644
--- a/tests/test_nontree_vit_haploid_tstv.py
+++ b/tests/test_nontree_vit_haploid_tstv.py
@@ -16,11 +16,11 @@ def verify(self, ts, include_ancestors):
         n = H.shape[1]
 
         r_s = [
-            np.append([0], np.zeros(m - 1) + 0.01),
-            np.append([0], np.random.rand(m - 1)),
-            np.append([0], 1e-5 * (np.random.rand(m - 1) + 0.5) / 2),
-            np.append([0], np.zeros(m - 1) + 0.2),
-            np.append([0], np.zeros(m - 1) + 1e-6),
+            np.zeros(m) + 0.01,
+            np.random.rand(m),
+            1e-5 * (np.random.rand(m - 1) + 0.5) / 2,
+            np.zeros(m - 1) + 0.2,
+            np.zeros(m - 1) + 1e-6,
         ]
         mu_s = [
             np.zeros(m) + 0.01,
@@ -32,7 +32,7 @@ def verify(self, ts, include_ancestors):
         kappa_s = [1 / 2, 1 / 4, 1, 3 / 2, 2]
 
         for s, r, mu, kappa in itertools.product(queries, r_s, mu_s, kappa_s):
-            e = core.get_emission_matrix_hkylike(mu, kappa)
+            e = core.get_emission_matrix_haploid_hkylike(mu, kappa)
 
             V_vs, P_vs, ll_vs = vh.forwards_viterbi_hap_naive(
                 n=n,