Refactor and njit

szhan · Sep 13, 2023 · 3830c64 · 3830c64
1 parent 2799ad3
commit 3830c64
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
@@ -9,6 +9,7 @@
 import _tskit
 
 
+@njit
 def convert_to_genetic_map_position(pos):
     """
     Convert genomic site positions (bp) to genetic map positions (cM).
@@ -50,7 +51,7 @@ def get_weights(genotyped_pos, imputed_pos):
     # Calculate weights for imputed markers.
     weights = np.zeros(x, dtype=np.float64)
     # Identify genotype markers k and k + 1 between imputed marker i
-    marker_interval_start = np.repeat(-2, x, dtype=np.int64)
+    marker_interval_start = np.zeros(x, dtype=np.int64) - 2
     idx_m = m - 1
     for idx_x in np.arange(x - 1, -1, -1):
         while imputed_cm[idx_x] < genotyped_cm[idx_m] and idx_m > -1: