Modify how to classify transitions

lshmm/core.py

@@ -335,9 +335,6 @@ def get_emission_matrix_haploid_tstv(mu, kappa=None):
         np.zeros((num_sites, num_alleles, num_alleles), dtype=np.float64) - 1
     )
 
-    # Define transitions: A <-> G and C <-> T.
-    transitions = [(0, 2), (2, 0), (1, 3), (3, 1)]
-
     for i in range(num_sites):
         for j in range(num_alleles):
             for k in range(num_alleles):
@@ -346,7 +343,10 @@ def get_emission_matrix_haploid_tstv(mu, kappa=None):
                 else:
                     mu_over_two_plus_kappa = mu[i] / (2.0 + kappa)
                     emission_matrix[i, j, k] = mu_over_two_plus_kappa
-                    if (j, k) in transitions:
+                    # Transitions: A <-> G and C <-> T.
+                    is_transition_AG = j in [0, 2] and k in [0, 2]
+                    is_transition_CT = j in [1, 3] and k in [1, 3]
+                    if is_transition_AG or is_transition_CT:
                         emission_matrix[i, j, k] *= kappa
 
         row_sum = np.sum(emission_matrix[i, j, :])