diff --git a/python/tests/test_imputation.py b/python/tests/test_imputation.py index acf0b43603..e239d936f9 100644 --- a/python/tests/test_imputation.py +++ b/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}")