Lowpass update - automate nseq (#133)

* Update dadi-cli LowPass integration to automatically generate a dictionary with population IDs and the maximum potential sequenced samples using dadi.make_data_dict_vcf's extract_ploidy argument

Author: tjstruck

dadi_cli/GenerateFs.py

@@ -90,19 +90,25 @@ def generate_fs(
         # multiply number of individuals subsamples by the ploidy to get sample size
         projections = [individuals*ploidy for individuals in subsample]
         print(projections, ploidy, subsample)
+    elif calc_coverage:
+        dd, ploidy = dadi.Misc.make_data_dict_vcf(vcf_filename=vcf, popinfo_filename=pop_info, calc_coverage=calc_coverage, extract_ploidy=True)
+        import collections
+        nseq = collections.defaultdict(int)
+        for line in open(pop_info).readlines():
+            nseq[line.strip().split()[-1]] += 1 * ploidy
     else:
         dd = dadi.Misc.make_data_dict_vcf(vcf_filename=vcf, popinfo_filename=pop_info, calc_coverage=calc_coverage)
 
     # Moved this lower, since using subsamples make projections not required
     if len(pop_ids) != len(projections):
         raise ValueError("The lengths of `pop_ids` and `projections` must match.")
 
-    if calc_coverage:    
+    if calc_coverage:
         import pickle
         import dadi.LowPass.LowPass as lp
         cov_dist = lp.compute_cov_dist(dd, pop_ids)
         print(f"\nSaving coverage distribution data in pickle named:\n{output}.coverage.pickle\n")
-        pickle.dump([cov_dist, calc_coverage], open(f"{output}.coverage.pickle","wb"))
+        pickle.dump([cov_dist, nseq], open(f"{output}.coverage.pickle","wb"))
 
     if bootstrap is None:
         fs = dadi.Spectrum.from_data_dict(

dadi_cli/InferDFE.py

@@ -117,13 +117,15 @@ def infer_dfe(
             from dadi.LowPass.LowPass import make_low_pass_func_GATK_multisample as func_cov
         except ModuleNotFoundError:
             raise ImportError("ERROR:\nCurrent dadi version does not support coverage model\n")
-        # nseq = [int(ele) for ele in cov_args[1:]]
+        nseq = []
+        for pop in fs.pop_ids:
+            nseq.append(cov_args[1][pop])
         if cov_inbreeding == []:
             Fx = None
         else:
             Fx = cov_inbreeding
 
-        func = func_cov(func, cov_args[0], fs.pop_ids, cov_args[1], fs.sample_sizes, Fx=Fx)
+        func = func_cov(func, cov_args[0], fs.pop_ids, nseq, fs.sample_sizes, Fx=Fx)
 
     p0_len = len(p0)
     # lower_bounds = convert_to_None(lower_bounds, p0_len)

dadi_cli/InferDM.py

@@ -90,14 +90,15 @@ def infer_demography(
             from dadi.LowPass.LowPass import make_low_pass_func_GATK_multisample as func_cov
         except ModuleNotFoundError:
             raise ImportError("ERROR:\nCurrent dadi version does not support coverage model\n")
-        # nseq = [int(ele) for ele in cov_args[1]]
-        # cov_args[1] should be a list with ints, so no need to convert to int
+        nseq = []
+        for pop in fs.pop_ids:
+            nseq.append(cov_args[1][pop])
         if cov_inbreeding == []:
             Fx = None
         else:
             Fx = cov_inbreeding
 
-        func_ex = func_cov(func_ex, cov_args[0], fs.pop_ids, cov_args[1], fs.sample_sizes, Fx=Fx)
+        func_ex = func_cov(func_ex, cov_args[0], fs.pop_ids, nseq, fs.sample_sizes, Fx=Fx)
 
     p0_len = len(p0)
     if fixed_params == -1:
@@ -215,13 +216,15 @@ def infer_global_opt(
             import pickle
         except ModuleNotFoundError:
             raise ImportError("ERROR:\nCurrent dadi version does not support coverage model\n")
-        # nseq = [int(ele) for ele in cov_args[1:]]
+        nseq = []
+        for pop in fs.pop_ids:
+            nseq.append(cov_args[1][pop])
         if cov_inbreeding == []:
             Fx = None
         else:
             Fx = cov_inbreeding
 
-        func_ex = func_cov(func_ex, cov_args[0], fs.pop_ids, cov_args[1], fs.sample_sizes, Fx=Fx)
+        func_ex = func_cov(func_ex, cov_args[0], fs.pop_ids, nseq, fs.sample_sizes, Fx=Fx)
 
     p0_len = len(p0)
     if fixed_params == -1:

dadi_cli/Plot.py

@@ -183,13 +183,15 @@ def plot_fitted_demography(
             from dadi.LowPass.LowPass import make_low_pass_func_GATK_multisample as func_cov
         except ModuleNotFoundError:
             raise ImportError("ERROR:\nCurrent dadi version does not support coverage model\n")
-        # nseq = [int(ele) for ele in cov_args[1:]]
+        nseq = []
+        for pop in fs.pop_ids:
+            nseq.append(cov_args[1][pop])
         if cov_inbreeding == []:
             Fx = None
         else:
             Fx = cov_inbreeding
 
-        func_ex = func_cov(func_ex, cov_args[0], fs.pop_ids, cov_args[1], fs.sample_sizes, Fx=Fx)
+        func_ex = func_cov(func_ex, cov_args[0], fs.pop_ids, nseq, fs.sample_sizes, Fx=Fx)
 
     model = func_ex(popt, ns, pts_l)
 
@@ -304,13 +306,15 @@ def plot_fitted_dfe(
             from dadi.LowPass.LowPass import make_low_pass_func_GATK_multisample as func_cov
         except ModuleNotFoundError:
             raise ImportError("ERROR:\nCurrent dadi version does not support coverage model\n")
-        # nseq = [int(ele) for ele in cov_args[1:]]
+        nseq = []
+        for pop in fs.pop_ids:
+            nseq.append(cov_args[1][pop])
         if cov_inbreeding == []:
             Fx = None
         else:
             Fx = cov_inbreeding
 
-        func = func_cov(func, cov_args[0], fs.pop_ids, cov_args[1], fs.sample_sizes, Fx=Fx)
+        func = func_cov(func, cov_args[0], fs.pop_ids, nseq, fs.sample_sizes, Fx=Fx)
     # Get expected SFS for MLE
     if (cache1d != None) and (cache2d != None):
         model = func(sele_popt, None, spectra1d, spectra2d, pdf, pdf2, theta, None)

dadi_cli/parsers/generate_fs_parsers.py

@@ -173,13 +173,10 @@ def add_generate_fs_parsers(subparsers: argparse.ArgumentParser) -> None:
 
     parser.add_argument(
         "--calc-coverage",
-        nargs="+",
         default=False,
-        type=list,
+        action="store_true",
         dest="calc_coverage",
-        help="Pass in the total number of sequenced haploids (nseq) for each population:\n"+
-        "ex. --calc-coverage 22 12, if the total number a sequenced haploids was 22 for population 1 and 12 for population 2.\n" +
-        "And store coverage information of sites and nseq in <output>.coverage.pickle object. Default: False.",
+        help="Store coverage information of sites and total haploids sequenced per population in <output>.coverage.pickle object. Default: False.",
     )
     add_output_argument(parser)
     add_seed_argument(parser)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dadi-cli"
-version = "0.9.10"
+version = "0.9.11"
 requires-python = ">=3.9"
 description = "A command line interface for dadi"
 readme = "README.md"

tests/example_data/LowPass-files/cov.fs.coverage.pickle

@@ -424,13 +424,13 @@ def test_generate_fs_lowpass_dd(data):
         bootstrap=None,
         chunk_size=None,
         masking="",
-        calc_coverage=[20, 20],
+        calc_coverage=True,
         seed=None,
     )
     import pickle
     cov_dist, nseq = pickle.load(open("tests/test_results/cov-test.fs.coverage.pickle", "rb"))
 
-    assert nseq == [20, 20]
+    assert nseq == {'pop1': 20, 'pop2': 20}
 
 def test_generate_fs_lowpass_fs(data):
     generate_fs(
@@ -445,7 +445,7 @@ def test_generate_fs_lowpass_fs(data):
         bootstrap=None,
         chunk_size=None,
         masking="",
-        calc_coverage=[20, 20],
+        calc_coverage=True,
         seed=None,
     )
     dadi_cli_fs = dadi.Spectrum.from_file("tests/test_results/cov-test.fs")

tests/test_GenerateFs.py

@@ -60,7 +60,7 @@ def gen_fs_args():
     gen_fs_args.pop_info = "tests/example_data/LowPass-files/cov_popfile_2D.txt"
     gen_fs_args.pop_ids = ["pop1", "pop2"]
     gen_fs_args.polarized = False
-    gen_fs_args.calc_coverage = [20, 20]
+    gen_fs_args.calc_coverage = True
     _run_generate_fs(gen_fs_args)
     os.remove(gen_fs_args.output)
     os.remove(gen_fs_args.output+".coverage.pickle")