From 56a439b4be5f90d4cc819b8a67ec1d22e537e8d6 Mon Sep 17 00:00:00 2001
From: tjstruck <tjstruck@email.arizona.edu>
Date: Tue, 6 Aug 2024 18:30:51 -0700
Subject: [PATCH] Update InferDM for LowPass

---
 dadi_cli/InferDM.py | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/dadi_cli/InferDM.py b/dadi_cli/InferDM.py
index 21c26ffd..58740844 100644
--- 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)