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

tskit-dev · Aug 8, 2023 · 200b6ab · 200b6ab
1 parent 1679b1a
commit 200b6ab
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/python/tests/test_imputation.py b/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]