Minor refactoring

python/tests/test_imputation.py

@@ -132,7 +132,7 @@ def get_toy_data():
 # For each site, the sum of backward value over all haplotypes is calculated
 # before scaling and shifting.
 
-fwd_matrix_text_1 = """
+beagle_forward_matrix_text_1 = """
 m,h,probRec,probNoRec,noErrProb,errProb,refAl,queryAl,shiftFac,scaleFac,sumSite,val
 0,0,0.000000,1.000000,0.999900,0.000100,1,0,0.000000,1.000000,0.000100,0.000100
 0,1,0.000000,1.000000,0.999900,0.000100,0,0,0.000000,1.000000,1.000000,0.999900
@@ -152,7 +152,7 @@ def get_toy_data():
 3,3,1.000000,0.000000,0.999900,0.000100,1,0,0.250000,0.000000,0.250050,0.000025
 """
 
-bwd_matrix_text_1 = """
+beagle_backward_matrix_text_1 = """
 m,h,probRec,probNoRec,noErrProb,errProb,refAl,queryAl,shiftFac,scaleFac,sumSite,val
 3,0,-1,-1,-1,-1,-1,-1,-1,-1,-1,1.000000
 3,1,-1,-1,-1,-1,-1,-1,-1,-1,-1,1.000000
@@ -172,7 +172,7 @@ def get_toy_data():
 0,3,1.000000,0.000000,0.999900,0.000100,1,1,0.000000,0.250000,0.250050,0.250000
 """
 
-fwd_matrix_text_2 = """
+beagle_forward_matrix_text_2 = """
 m,h,probRec,probNoRec,noErrProb,errProb,refAl,queryAl,shiftFac,scaleFac,sumSite,val
 0,0,0.000000,1.000000,0.999900,0.000100,1,1,0.000000,1.000000,0.999900,0.999900
 0,1,0.000000,1.000000,0.999900,0.000100,0,1,0.000000,1.000000,1.000000,0.000100
@@ -192,7 +192,7 @@ def get_toy_data():
 3,3,1.000000,0.000000,0.999900,0.000100,1,1,0.250000,0.000000,0.749950,0.249975
 """
 
-bwd_matrix_text_2 = """
+beagle_backward_matrix_text_2 = """
 m,h,probRec,probNoRec,noErrProb,errProb,refAl,queryAl,shiftFac,scaleFac,sumSite,val
 3,0,-1,-1,-1,-1,-1,-1,-1,-1,-1,1.000000
 3,1,-1,-1,-1,-1,-1,-1,-1,-1,-1,1.000000
@@ -223,39 +223,39 @@ def convert_to_numpy(matrix_text):
     return df.val.to_numpy().reshape((4, 4))
 
 
-def get_forward_backward_matrices():
-    fwd_matrix_1 = convert_to_numpy(fwd_matrix_text_1)
-    bwd_matrix_1 = convert_to_numpy(bwd_matrix_text_1)
-    fwd_matrix_2 = convert_to_numpy(fwd_matrix_text_2)
-    bwd_matrix_2 = convert_to_numpy(bwd_matrix_text_2)
+def get_beagle_forward_backward_matrices():
+    fwd_matrix_1 = convert_to_numpy(beagle_forward_matrix_text_1)
+    bwd_matrix_1 = convert_to_numpy(beagle_backward_matrix_text_1)
+    fwd_matrix_2 = convert_to_numpy(beagle_forward_matrix_text_2)
+    bwd_matrix_2 = convert_to_numpy(beagle_backward_matrix_text_2)
     return [fwd_matrix_1, bwd_matrix_1, fwd_matrix_2, bwd_matrix_2]
 
 
-def get_test_data(matrix_text, par):
+def get_beagle_data(matrix_text, data_type):
     """Extracts data to check forward or backward probability matrix calculations."""
     df = pd.read_csv(io.StringIO(matrix_text))
-    if par == "switch":
+    if data_type == "switch":
         # Switch probability, one per site
         return df.probRec.to_numpy().reshape((4, 4))[:, 0]
-    elif par == "mismatch":
+    elif data_type == "mismatch":
         # Mismatch probability, one per site
         return df.errProb.to_numpy().reshape((4, 4))[:, 0]
-    elif par == "ref_hap_allele":
+    elif data_type == "ref_hap_allele":
         # Allele in haplotype in reference panel
         # 0 = ref allele, 1 = alt allele
         return df.refAl.to_numpy().reshape((4, 4))
-    elif par == "query_hap_allele":
+    elif data_type == "query_hap_allele":
         # Allele in haplotype in query
         # 0 = ref allele, 1 = alt allele
         return df.queryAl.to_numpy().reshape((4, 4))[:, 0]
-    elif par == "shift":
+    elif data_type == "shift":
         # Shift factor, one per site
         return df.shiftFac.to_numpy().reshape((4, 4))[:, 0]
-    elif par == "scale":
+    elif data_type == "scale":
         # Scale factor, one per site
         return df.scaleFac.to_numpy().reshape((4, 4))[:, 0]
-    elif par == "sum":
+    elif data_type == "sum":
         # Sum of values over haplotypes
         return df.sumSite.to_numpy().reshape((4, 4))[:, 0]
     else:
-        raise ValueError(f"Unknown parameter: {par}")
+        raise ValueError(f"Unknown data type: {data_type}")