Add function to get forward and backward matrices from _tskit.lshmm WIP

Author: szhan

python/tests/test_imputation.py

@@ -6,6 +6,7 @@
 import numpy as np
 import pandas as pd
 
+import _tskit
 import tskit
 
 
@@ -259,3 +260,13 @@ def get_beagle_data(matrix_text, data_type):
         return df.sumSite.to_numpy().reshape((4, 4))[:, 0]
     else:
         raise ValueError(f"Unknown data type: {data_type}")
+
+
+def get_tskit_forward_backward_matrices(ts, h):
+    m = ts.num_sites
+    fm = _tskit.CompressedMatrix(ts)
+    bm = _tskit.CompressedMatrix(ts)
+    ls_hmm = _tskit.LsHmm(ts, np.zeros(m) + 0.1, np.zeros(m) + 0.1)
+    ls_hmm.forward_matrix(h, fm)
+    ls_hmm.backward_matrix(h, fm.normalisation_factor, bm)
+    return [fm, bm]