Update InferDM for LowPass

xin-huang · Aug 7, 2024 · 56a439b · 56a439b
1 parent 1ee6a99
commit 56a439b
Showing 1 changed file with 34 additions and 0 deletions.
diff --git a/dadi_cli/InferDM.py b/dadi_cli/InferDM.py
@@ -13,6 +13,8 @@ def infer_demography(
     lower_bounds: list[float],
     fixed_params: list[float],
     misid: bool,
+    cov_args: list,
+    cov_inbreeding: list,
     cuda: bool,
     maxeval: int,
     maxtime: int,
@@ -82,6 +84,22 @@ def infer_demography(
         func = dadi.Numerics.make_anc_state_misid_func(func)
 
     func_ex = dadi.Numerics.make_extrap_func(func)
+
+    if cov_args != []:
+        try:
+            from dadi.LowPass.LowPass import make_low_pass_func_GATK_multisample as func_cov
+            import pickle
+        except ModuleNotFoundError:
+            raise ImportError("ERROR:\nCurrent dadi version does not support coverage model\n")
+        cov_dd = cov_args[0]
+        nseq = [int(ele) for ele in cov_args[1:]]
+        if cov_inbreeding == []:
+            Fx = None
+
+        cov_dd = pickle.load(open(cov_dd, 'rb'))
+        func_ex = func_cov(func_ex, cov_dd, fs.pop_ids, nseq, fs.sample_sizes, Fx=Fx)
+        print(func_cov)
+
     p0_len = len(p0)
     lower_bounds = convert_to_None(lower_bounds, p0_len)
     upper_bounds = convert_to_None(upper_bounds, p0_len)
@@ -131,6 +149,8 @@ def infer_global_opt(
     lower_bounds: list[float],
     fixed_params: list[float],
     misid: bool,
+    cov_args: list,
+    cov_inbreeding: list,
     cuda: bool,
     maxeval: int,
     maxtime: int,
@@ -191,6 +211,20 @@ def infer_global_opt(
 
     func_ex = dadi.Numerics.make_extrap_func(func)
 
+    if cov_args != []:
+        try:
+            from dadi.LowPass.LowPass import make_low_pass_func_GATK_multisample as func_cov
+            import pickle
+        except ModuleNotFoundError:
+            raise ImportError("ERROR:\nCurrent dadi version does not support coverage model\n")
+        cov_dd = cov_args[0]
+        nseq = [int(ele) for ele in cov_args[1:]]
+        if cov_inbreeding == []:
+            Fx = None
+
+        cov_dd = pickle.load(open(cov_dd, 'rb'))
+        func_ex = func_cov(func_ex, cov_dd, fs.pop_ids, nseq, fs.sample_sizes, Fx=Fx)
+
     p0_len = len(p0)
     lower_bounds = convert_to_None(lower_bounds, p0_len)
     upper_bounds = convert_to_None(upper_bounds, p0_len)