[core] minor bug fixes

Author: terhorst

smcpp/analysis.py

@@ -52,6 +52,11 @@ def __init__(self, files, args):
         self._recode_nonseg(args.nonseg_cutoff)
 
 
+    @property
+    def ns(self):
+        return np.array([sum(c.n) for c in self._contigs])
+
+
     def _init_optimizer(self, outdir, algorithm, xtol, ftol, single):
         self._optimizer = self._OPTIMIZER_CLS(self, algorithm, xtol, ftol, single)
         if outdir:
@@ -88,7 +93,6 @@ def _load_data(self, files):
             assert c.data.shape[1] == 1 + 3 * len(c.n)
             logger.debug(c)
         logger.info("%d population%s", self.npop, "" if self.npop == 1 else "s")
-        ns = self._ns = np.array([sum(c.n) for c in self._contigs])
 
     def _validate_data(self):
         for c in self._contigs:
@@ -147,7 +151,7 @@ def _perform_thinning(self, thinning):
         if isinstance(thinning, int):
             thinning = np.array([thinning] * len(self._contigs))
         if thinning is None:
-            thinning = (1000 * np.log(2 + self._ns)).astype("int")   # 500  * ns
+            thinning = (1000 * np.log(2 + self.ns)).astype("int")   # 500  * ns
         thinning[npop > 1] = 0
         if np.any(thinning > 1):
             logger.info("Thinning...")
@@ -237,7 +241,7 @@ def _calculate_t1_tK(self, args):
         Ne = self._watterson / self._theta
         logger.debug("Ne: %f", Ne)
         Ne *= 2 * self._N0
-        n = 2 + max(self._ns)
+        n = 2 + max(self.ns)
         t1 = args.t1 or 200 + np.log(.9) / (-n * (n - 1) / 2) * Ne
         if t1 <= 0:
             logger.error("--t1 should be >0")
@@ -287,7 +291,7 @@ def _init_inference_manager(self, polarization_error):
             im.rho = self._rho
             im.alpha = self._alpha = .01
             self._ims[pid] = im
-        self._max_n = np.max(list(max_n.values()), axis=0)
+        self._max_n = np.max(list(map(sum, max_n.values())), axis=0)
         self._model.randomize()
 
     @property
@@ -366,9 +370,9 @@ def npop(self):
 
     def dump(self, filename):
         'Dump result of this analysis to :filename:.'
-        d = {'theta': self._theta, 'rho': self._rho}
+        d = {'theta': self._theta, 'rho': self._rho, 'alpha': self._alpha}
         d['model'] = self.model.to_dict()
-        d['hidden_states'] = {k: v.tolist() for k, v in self._hidden_states.items()}
+        d['hidden_states'] = {k: list(v) for k, v in self._hidden_states.items()}
         json.dump(d, open(filename + ".json", "wt"), sort_keys=True, indent=4)
 
 
@@ -456,18 +460,11 @@ def _init_model(self, N0, spline_class):
         self._y0 = y0 = self._het / (2. * self._theta)
         logger.debug("Long term avg. effective population size: %f", y0)
         y0 = np.log(y0)
-        if self.npop == 1:
-            self._model = SMCModel(
-                self._knots, self._N0,
-                spline_class, self._populations[0])
-            self._model[:] = y0
-        else:
-            split = self.rescale(.5 * tK)  # just pick the midpoint as a starting value.
-            mods = []
-            for pid in self._populations:
-                mods.append(SMCModel(self._knots, self._N0, spline_class, pid))
-                mods[-1][:] = y0
-            self._model = SMCTwoPopulationModel(mods[0], mods[1], split)
+        assert self.npop == 1
+        self._model = SMCModel(
+            self._knots, self._N0,
+            spline_class, self._populations[0])
+        self._model[:] = y0
 
     def _preinitialize(self):
         args = self._args
@@ -530,7 +527,7 @@ def __init__(self, files, args):
         self._perform_thinning(args.thinning)
         # Further initialization
         self._init_inference_manager(args.polarization_error)
-        self._init_optimizer(args.outdir, args.algorithm, args.xtol, args.ftol)
+        self._init_optimizer(args.outdir, args.algorithm, args.xtol, args.ftol, True)
         self._niter = 1
 
     def _validate_data(self):
@@ -556,6 +553,7 @@ def _init_model(self, pop1, pop2):
         m1 = _model_cls_d[d['model']['class']].from_dict(d['model'])
         d = json.load(open(pop2, "rt"))
         m2 = _model_cls_d[d['model']['class']].from_dict(d['model'])
+        self._hidden_states.update(d['hidden_states'])
         assert d['theta'] == self._theta
         self._max_split = m2._knots[-(len(smcpp.defaults.additional_knots) + 1)]
         self._model = SMCTwoPopulationModel(m1, m2, self._max_split * 0.5)

smcpp/commands/posterior.py

@@ -89,13 +89,14 @@ def main(self, args):
             all_obs = estimation_tools.thin_dataset(all_obs, [args.thinning] * len(all_obs))
         if npop == 1:
             im = _smcpp.PyOnePopInferenceManager(
-                n[0], all_obs, hidden_states, contig.key, args.polarization_error)
+                n[0], all_obs, hidden_states, contig.key[0], args.polarization_error)
         else:
             assert npop == 2
             im = _smcpp.PyTwoPopInferenceManager(
-                *n, *a, all_obs, hidden_states, contig.key, args.polarization_error)
+                *n, *a, all_obs, hidden_states, contig.key[0], args.polarization_error)
         im.theta = j['theta']
         im.rho = j['rho']
+        im.alpha = j['alpha']
         im.save_gamma = True
         im.model = m
         im.E_step()

src/transition.cpp

@@ -248,6 +248,7 @@ HJTransition<T>::HJTransition(const PiecewiseConstantRateFunction<T> &eta, const
     T small = eta.zero() + 1e-20;
     this->Phi = this->Phi.unaryExpr([small] (const T &x) { if (x < 1e-20) return small; return x; });
     CHECK_NAN(this->Phi);
+    return;
     const double beta = 1e-5;
     T p2 = eta.zero() + beta / this->M;
     Matrix<T> Phi2(this->M, this->M);

test/integration/test.sh

@@ -4,8 +4,8 @@ set -e
 $SMC vcf2smc -v example/example.vcf.gz /tmp/example.1.smc.gz 1 msp1:msp_0,msp_1,msp_2
 $SMC vcf2smc -d msp_0 msp_0 example/example.vcf.gz /tmp/example.2.smc.gz 1 msp2:msp_0,msp_3,msp_4
 $SMC vcf2smc -d msp_1 msp_1 example/example.vcf.gz /tmp/example.12.smc.gz 1 msp1:msp_1,msp_2 msp2:msp_3,msp_4,msp_0
-$SMC estimate -o /tmp/out/1 --unfold --em-iterations 1 1.25e-8 /tmp/example.1.smc.gz
-$SMC estimate -p 0.01 -r 1e-8 -o /tmp/out/2 --em-iterations 1 1.25e-8 /tmp/example.2.smc.gz
+$SMC estimate -o /tmp/out/1 --unfold --knots 5 --em-iterations 1 1.25e-8 /tmp/example.1.smc.gz
+$SMC estimate -p 0.01 -r 1e-8 -o /tmp/out/2 --knots 5 --em-iterations 1 1.25e-8 /tmp/example.2.smc.gz
 $SMC split -o /tmp/out/split --em-iterations 1 \
     /tmp/out/1/model.final.json \
     /tmp/out/2/model.final.json \