Add a convenience function to get data from the toy data

Author: szhan

python/tests/test_imputation.py

@@ -219,3 +219,35 @@ def get_forward_backward_matrices():
     fwd_matrix_2 = convert_to_numpy(_fwd_matrix_text_2)
     bwd_matrix_2 = convert_to_numpy(_bwd_matrix_text_2)
     return [fwd_matrix_1, bwd_matrix_1, fwd_matrix_2, bwd_matrix_2]
+
+
+def get_test_data(matrix_text, par):
+    """Extracts data for checking forward or backward probability matrix calculations."""
+    x = convert_to_numpy(matrix_text)
+    if par == "switch":
+        # Switch probability, one per site
+        return x[:, 2].reshape((4, 4))[:, 2]
+    elif par == "mismatch":
+        # Mismatch probability
+        return x[:, 2].reshape((4, 4))[:, 4]
+    elif par == "ref_hap_allele":
+        # Allele in haplotype in reference panel
+        # 0 = ref allele, 1 = alt allele
+        return x[:, 2].reshape((4, 4))[:, 6]
+    elif par == "query_hap_allele":
+        # Allele in haplotype in query
+        # 0 = ref allele, 1 = alt allele
+        return x[:, 2].reshape((4, 4))[:, 7]
+    elif par == "shift":
+        # Shift factor
+        # TODO
+        pass
+    elif par == "scale":
+        # Scale factor
+        # TODO
+        pass
+    elif par == "sum":
+        # Sum of values over haplotypes
+        return x[:, 2].reshape((4, 4))[:, 10]
+    else:
+        raise ValueError(f"Unknown parameter: {par}")