update to print Ne, rec

Author: stsmall

Diff for: project/sim_modules/readconfig.py

@@ -195,10 +195,8 @@ def read_config_stats(configFile):
         assert max(size_list) <= n_haps
     sim_dt["spatial_sfs"] = config.getboolean(pop1, "compute_spatial_SFS")
     spatfold = config.getboolean(pop1, "fold_spatial_SFS")
-    spatagg = config.getboolean(pop1, "aggregate_spatial_SFS")
     sim_dt["sfs"] = config.getboolean(pop1, "compute_SFS")
     sfsfold = config.getboolean(pop1, "fold_SFS")
-    sfsagg = config.getboolean(pop1, "aggregate_SFS")
     sim_dt["ld"] = config.getboolean(pop1, "compute_LD")
     if sim_dt["ld"]:
         nb_times = config.getint(pop1, "pop_time_changes")
@@ -228,8 +226,8 @@ def read_config_stats(configFile):
                 "win_size2": int(win_size_2),
                 "calc_stats": stat_list,
                 "afibs_fold": afibsfold,
-                "spat_fold_agg": (spatfold, spatagg),
-                "sfs_fold_agg": (sfsfold, sfsagg),
+                "spat_fold_agg": (spatfold),
+                "sfs_fold_agg": (sfsfold),
                 "pw_quants": quant_list
                 }
     if "ibs" in stat_list:

Diff for: project/sim_modules/sim_msprime.py

@@ -264,9 +264,13 @@ def simulate_msprime(model_dict, demo_dataframe, param_df, sim_number: int,
             low, high = mut_rate
             mu = np.random.uniform(low, high, sim_number)
         else:
-            mu = mut_rate
-            if order:
+            if len(mut_rate) < sim_number:
+                mu = np.random.choice(mut_rate, sim_number)
+            elif order:
                 l_mu = len(mu)
+            else:
+                mu = mut_rate
+
     else:
         mu = [mut_rate] * sim_number
 
@@ -280,9 +284,12 @@ def simulate_msprime(model_dict, demo_dataframe, param_df, sim_number: int,
             rec = np.random.uniform(low, high, sim_number)
             # rec = np.random.exponential(rec_rate, sim_number)
         else:
-            rec = rec_rate
-            if order:
+            if len(rec_rate) < sim_number:
+                rec = np.random.choice(rec_rate, sim_number)
+            elif order:
                 l_rec = len(rec)
+            else:
+                rec = rec_rate
     else:
         rec = [rec_rate] * sim_number
 
@@ -299,9 +306,12 @@ def simulate_msprime(model_dict, demo_dataframe, param_df, sim_number: int,
             low, high = effective_size
             scaled_Ne = np.random.randint(low, high, sim_number) * ploidy
         else:
-            scaled_Ne = effective_size
-            if order:
+            if len(effective_size) < sim_number:
+                scaled_Ne = np.random.choice(effective_size, sim_number)
+            elif order:
                 l_ne = len(scaled_Ne)
+            else:
+                scaled_Ne = effective_size
     else:
         scaled_Ne = [effective_size * ploidy] * sim_number
     # =========================================================================

Diff for: project/stat_modules/afs.py

@@ -59,8 +59,8 @@ def spatial_histo_fast(pos, count, M, dmax=np.inf):
     return dist
 
 
-def asfs_stats(gt, pos, fold, agg):
-    """Calculate the aggregate allele frequence spectrum.
+def asfs_stats(gt, pos, fold):
+    """Calculate the allele frequence spectrum.
 
     With many individuals the SFS becomes unwieldy, here I collapse the
     intermediate frequencies into 1 bin. This differs from the one above by
@@ -75,8 +75,6 @@ def asfs_stats(gt, pos, fold, agg):
         DESCRIPTION.
     pos : TYPE
         DESCRIPTION.
-    agg : bool
-        if True, return an aggregate SFS, 0,1,2, 3:-2, -2, -1
     fold : bool
         if True, return folded SFS
 
@@ -94,12 +92,6 @@ def asfs_stats(gt, pos, fold, agg):
         sfsp = (allel.sfs(gtseg.count_alleles()[:, 1], gtseg.shape[1]))[1:-1]
     tots = np.sum(sfsp)
     sfs = sfsp / tots
-    if agg:
-        # TODO: test best binning.
-        sfs[3] = sum(sfs[3:-2])
-        sfs[4] = sfs[-2]
-        sfs[5] = sfs[-1]
-        sfs = sfs[:6]
 
     return sfs
 

Diff for: project/stat_modules/pwpopstats.py

@@ -16,7 +16,7 @@
 from project.stat_modules.sequtils import get_seg, get_ac_seg, h2gt, pop2seg
 
 
-def d_tajD(p1, pos, gt, win_size, length_bp):
+def d_tajD(p1, pos, gt, win_size, length_bp, quants):
     """Compute the difference in Tajima’s D between two populations in moving windows.
 
     Parameters
@@ -49,6 +49,11 @@ def d_tajD(p1, pos, gt, win_size, length_bp):
     if size < 5:
         size = 5
     dtajd_ = allel.moving_delta_tajima_d(ac1_seg, ac2_seg, size=size)
+
+    if quants[0] < 0:
+        dtajdq = [np.nanmean(dtajd_)]
+    else:
+        dtajdq = np.nanquantiles(dtajd_, quants)
     return dtajd_
 
 
@@ -86,7 +91,10 @@ def fst(p1, pos, gt, quants):
     ac2_seg = ac2.compress(loc_asc, axis=0)
     num, den = allel.hudson_fst(ac1_seg, ac2_seg)
     fst_snp = num / den
-    fst_ = np.nanquantile(fst_snp, quants)
+    if quants[0] < 0:
+        fst_ = [np.nanmean(fst_snp)]
+    else:
+        fst_ = np.nanquantile(fst_snp, quants)
     return fst_
 
 
@@ -141,15 +149,21 @@ def dmin(p1, pos, gt, win_size, length_bp):
     return dmin_
 
 
-def gmin(p1, pos, gt, win_size, length_bp):
+def gmin(p1, pos, gt, win_size, length_bp, quants):
     """Calculate minimum pairwise difference between two populations.
 
     dmin / dxy
     """
     dmin_ = dmin(p1, pos, gt, win_size, length_bp)
     dxy_ = dxy(p1, pos, gt, win_size, length_bp)
     gmin_win = dmin_/dxy_
-    return gmin_win
+
+    if quants[0] < 0:
+        gmin_ = [np.nanmean(gmin_win)]
+    else:
+        gmin_ = np.nanquantile(gmin_win, quants)
+
+    return gmin_
 
 
 def pi_fx(pos, gt, win_size, length_bp):
@@ -174,7 +188,10 @@ def dd1_2(p1, pos, gt, win_size, length_bp, quants):
     for p in [p1_, p2_]:
         gtpop = gt.take(p, axis=1)
         pi_ = pi_fx(pos, gtpop, win_size, length_bp)
-        dd12_ls.extend(np.nanquantile((dmin_/pi_), quants))
+        if quants[0] < 0:
+            dd12_ls.append(np.nanmean(dmin_/pi_))
+        else:
+            dd12_ls.extend(np.nanquantile((dmin_/pi_), quants))
     return dd12_ls
 
 
@@ -214,7 +231,10 @@ def ddRank1_2(p1, pos, gt, win_size, length_bp, quants):
         gtpop = gt.take(p, axis=1)
         pw_win = pw_within(pos, gtpop, win_size, length_bp)
         ddR12 = [stats.percentileofscore(pw, dmin_[i]) for i, pw in enumerate(pw_win)]
-        ddR12_ls.extend(np.nanquantile(ddR12, quants))
+        if quants[0] < 0:
+            ddR12_ls.append(np.nanmean(ddR12))
+        else:
+            ddR12_ls.extend(np.nanquantile(ddR12, quants))
     return ddR12_ls
 
 
@@ -233,7 +253,7 @@ def ld_window(pos, hap, win_size, length_bp):
     return ld_wins
 
 
-def zx(p1, pos, hap, win_size, length_bp, randn=500):
+def zx(p1, pos, hap, win_size, length_bp, quants, randn=500):
     """FILET statistic.
 
     Zx = (Zn_s1 + Zn_s2)/(2*Zn_sg)
@@ -259,7 +279,12 @@ def zx(p1, pos, hap, win_size, length_bp, randn=500):
     zn2 = np.array(ld_window(pos, h2, win_size, length_bp))
     zx_ = (zn1 + zn2) / (2 * zn_sg)
 
-    return zx_
+    if quants[0] < 0:
+        zxq = [np.nanmean(zx_)]
+    else:
+        zxq = np.nanquantile(zx_, quants)
+
+    return zxq
 
 
 def ibs_max(pos, hap, dmax):

Diff for: project/stat_modules/sumstats.py

@@ -86,7 +86,7 @@ def hap_het(self):
 
     def spatial_sfs(self):
         stats_ls = []
-        fold, agg = self.stats["spat_fold_agg"]
+        fold = self.stats["spat_fold"]
         for p in self.stats["pop_config"]:
             hap_p, pos_p, counts_p = self.split_pop(p)
             n_haplo = len(p)
@@ -97,12 +97,6 @@ def spatial_sfs(self):
                 haps = int(len(p)/2)
                 spat_fold = spat_fold[0:haps]
                 spat_ = spat_fold
-            if agg:
-                spat_ag = spat_[0:2]
-                spat_ag.append(np.mean(spat_[3:-2]))
-                spat_ag.append(spat_[-2])
-                spat_ag.append(spat_[-1])
-                spat_ = spat_ag
             stats_ls.extend(spat_)
 
         return stats_ls
@@ -159,14 +153,14 @@ def ld(self):
 
         return stats_ls
 
-    def sfs(self, agg=False, fold=False):
-        fold, agg = self.stats["sfs_fold_agg"]
+    def sfs(self, fold=False):
+        fold = self.stats["sfs_fold"]
         gt = allel.HaplotypeArray(self.haparr.T)
         pos = allel.SortedIndex(self.pos)
         stats_ls = []
         for pop in self.stats["pop_config"]:
             gtpop = gt.take(pop, axis=1)
-            sfs = afs.asfs_stats(gtpop, pos, fold, agg)
+            sfs = afs.asfs_stats(gtpop, pos, fold)
             stats_ls.extend(sfs)
 
         return stats_ls
@@ -191,8 +185,8 @@ def delta_tajD(self):
         stats_ls = []
         for p1, p2 in combinations(self.stats["pop_config"], 2):
             gtpops = gt.take(p1+p2, axis=1)
-            flt = pwpopstats.d_tajD(len(p1), pos, gtpops, win_size, length_bp)
-            stats_ls.extend(np.nanquantiles(flt, quants))
+            flt = pwpopstats.d_tajD(len(p1), pos, gtpops, win_size, length_bp, quants)
+            stats_ls.extend(flt)
         return stats_ls
 
     def ld_2pop(self, D=0):
@@ -244,7 +238,11 @@ def dXY(self):
         for p1, p2 in combinations(self.stats["pop_config"], 2):
             gtpops = gt.take(p1+p2, axis=1)
             flt = pwpopstats.dxy(len(p1), pos, gtpops, win_size, length_bp)
-            stats_ls.extend(np.nanquantile(flt, quants))
+            if quants[0] < 0:
+                dxy_ = [np.nanmean(flt)]
+            else:
+                dxy_ = np.nanquantile(flt, quants)
+            stats_ls.extend(dxy_)
         return stats_ls
 
     def dmin(self):
@@ -256,8 +254,12 @@ def dmin(self):
         stats_ls = []
         for p1, p2 in combinations(self.stats["pop_config"], 2):
             gtpops = gt.take(p1+p2, axis=1)
-            flt = pwpopstats.dmin(len(p1), pos, gtpops, win_size, length_bp)
-            stats_ls.extend(np.nanquantile(flt, quants))
+            flt = pwpopstats.dmin(len(p1), pos, gtpops, win_size, length_bp, quants)
+            if quants[0] < 0:
+                dminq = [np.nanmean(flt)]
+            else:
+                dminq = np.nanquantile(flt, quants)
+            stats_ls.extend(dminq)
         return stats_ls
 
     def gmin(self):
@@ -269,8 +271,8 @@ def gmin(self):
         stats_ls = []
         for p1, p2 in combinations(self.stats["pop_config"], 2):
             gtpops = gt.take(p1+p2, axis=1)
-            flt = pwpopstats.gmin(len(p1), pos, gtpops, win_size, length_bp)
-            stats_ls.extend(np.nanquantile(flt, quants))
+            flt = pwpopstats.gmin(len(p1), pos, gtpops, win_size, length_bp, quants)
+            stats_ls.extend(flt)
         return stats_ls
 
     def dd12(self):
@@ -307,8 +309,8 @@ def Zx(self):
         stats_ls = []
         for p1, p2 in combinations(self.stats["pop_config"], 2):
             hap_p, pos_p, counts_p = self.split_pop(p1+p2)
-            flt = pwpopstats.zx(len(p1), pos_p, hap_p, win_size, length_bp)
-            stats_ls.extend(np.nanquantile(flt, quants))
+            flt = pwpopstats.zx(len(p1), pos_p, hap_p, win_size, length_bp, quants)
+            stats_ls.extend(flt)
         return stats_ls
 
     def IBS_maxXY(self):

Diff for: project/stat_modules/write_stats.py

@@ -37,19 +37,15 @@ def headers(out_file, stats_dt, obs=False):
     sub_pops = [len(i) for i in stats_dt["pop_config"]]
     n_pops = len(sub_pops)
     quants = len(stats_dt["pw_quants"])
-    sfsfold, sfsagg = stats_dt["sfs_fold_agg"]
-    spatfold, spatagg = stats_dt["spat_fold_agg"]
+    sfsfold = stats_dt["sfs_fold"]
+    spatfold = stats_dt["spat_fold"]
     afibsfold = stats_dt["afibs_fold"]
     sfs_n = [sub-1 for sub in sub_pops]
     if sfsfold:
         sfs_n = [floor((sub)/2) for sub in sub_pops]
-    if sfsagg:
-        sfs_n = [6] * n_pops
     spat_n = [sub-1 for sub in sub_pops]
     if spatfold:
         spat_n = [floor((sub)/2) for sub in sub_pops]
-    if spatagg:
-        spat_n = [6] * n_pops
     afibs_n = [sub-2 for sub in sub_pops]
     if afibsfold:
         afibs_n = [floor((sub-1)/2) for sub in sub_pops]

Diff for: stat_test.py

@@ -106,18 +106,15 @@ def test_afibs():
 
 def test_sfs():
     # unfold
-    sfs = afs.asfs_stats(gtpop, pos_s, False, False)
+    sfs = afs.asfs_stats(gtpop, pos_s, False)
     ss = np.array([0.43859649, 0.16666667, 0.14912281, 0.01754386, 0.05263158,
                    0.04385965, 0.0877193, 0.01754386, 0.02631579])
     assert(all(np.isclose(ss, sfs)) is True)
     # fold
-    sfs = afs.asfs_stats(gtpop, pos_s, True, False)
+    sfs = afs.asfs_stats(gtpop, pos_s, True)
     ss = np.array([0.46491228, 0.18421053, 0.23684211, 0.06140351, 0.05263158])
     assert(all(np.isclose(ss, sfs)) is True)
-    # agg
-    sfs = afs.asfs_stats(gtpop, pos_s, False, True)
-    ss = np.array([0.43859649, 0.16666667, 0.14912281, 0.20175439, 0.01754386,
-                   0.02631579])
+
     assert(all(np.isclose(ss, sfs)) is True)
 
 
@@ -170,14 +167,6 @@ def test_spatial_sfs():
     spat_f = spat_fold
     ss = np.array([13881.,  7868., 15864., 32643., 24952.])
     assert(all(np.isclose(ss, spat_f)) is True)
-    # agg
-    spat_ag = list(spat_[0:3])
-    spat_ag.append(np.mean(spat_[3:-2]))
-    spat_ag.append(spat_[-2])
-    spat_ag.append(spat_[-1])
-    spat_a = np.array(spat_ag)
-    ss = np.array([1419.0, 7868.0, 5648.0, 13833.75, 0.0, 12462.0])
-    assert(all(np.isclose(ss, spat_a)) is True)
 
 
 def test_ld():
@@ -231,7 +220,7 @@ def test_dmin():
 
 
 def test_gmin():
-    flt = pwpopstats.gmin(len(pop1), pos_s, gt, win_size, length_bp)
+    flt = pwpopstats.gmin(len(pop1), pos_s, gt, win_size, length_bp, quants)
     stats_ls = np.quantile(flt, quants)
     ss = np.array([0.09416196, 0.35460993, 0.40257649, 0.48543689, 0.49382716])
     assert(all(np.isclose(ss, stats_ls)) is True)
@@ -251,7 +240,7 @@ def test_ddRank12():
 
 
 def test_Zx():
-    flt = pwpopstats.zx(len(pop1), pos, haps, win_size, length_bp)
+    flt = pwpopstats.zx(len(pop1), pos, haps, win_size, length_bp, quants)
     stats_ls = np.nanquantile(flt, quants)
     ss = np.array([0. , 0.21724138, 0.65702271, 1.21374723, 1.78181818])
     assert(all(np.isclose(ss, stats_ls)) is True)