add more info to report

Author: aloctavodia

pymc3/smc/sample_smc.py

@@ -196,12 +196,15 @@ def sample_smc(
         for i in range(chains):
             results.append((sample_smc_int(*params, random_seed[i], i, _log)))
 
-    traces, log_marginal_likelihoods = zip(*results)
+    traces, log_marginal_likelihoods, betas, accept_ratios, nsteps = zip(*results)
     trace = MultiTrace(traces)
     trace.report._n_draws = draws
     trace.report._n_tune = 0
     trace.report._t_sampling = time.time() - t1
     trace.report.log_marginal_likelihood = np.array(log_marginal_likelihoods)
+    trace.report.betas = betas
+    trace.report.accept_ratios = accept_ratios
+    trace.report.nsteps = nsteps
 
     return trace
 
@@ -239,6 +242,9 @@ def sample_smc_int(
         chain=chain,
     )
     stage = 0
+    betas = []
+    accept_ratios = []
+    nsteps = []
     smc.initialize_population()
     smc.setup_kernel()
     smc.initialize_logp()
@@ -252,5 +258,8 @@ def sample_smc_int(
         smc.mutate()
         smc.tune()
         stage += 1
+        betas.append(smc.beta)
+        accept_ratios.append(smc.acc_rate)
+        nsteps.append(smc.n_steps)
 
-    return smc.posterior_to_trace()
+    return smc.posterior_to_trace(), smc.log_marginal_likelihood, betas, accept_ratios, nsteps

pymc3/smc/smc.py

@@ -229,29 +229,6 @@ def mutate(self):
         self.acc_per_chain = np.mean(ac_, axis=0)
         self.acc_rate = np.mean(ac_)
 
-    def posterior_to_trace_bk(self):
-        """
-        Save results into a PyMC3 trace
-        """
-        lenght_pos = len(self.posterior)
-        varnames = [v.name for v in self.variables]
-        straces = []
-        with self.model:
-            chain_lenght = int(lenght_pos / 10)
-            for chain in range(10):
-                strace = NDArray(self.model)
-                strace.setup(chain_lenght, chain)
-                for i in range(chain_lenght):
-                    value = []
-                    size = 0
-                    for var in varnames:
-                        shape, new_size = self.var_info[var]
-                        value.append(self.posterior[i][size : size + new_size].reshape(shape))
-                        size += new_size
-                    strace.record({k: v for k, v in zip(varnames, value)})
-                straces.append(strace)
-        return MultiTrace(straces)
-
     def posterior_to_trace(self):
         """
         Save results into a PyMC3 trace
@@ -270,7 +247,7 @@ def posterior_to_trace(self):
                 value.append(self.posterior[i][size : size + new_size].reshape(shape))
                 size += new_size
             strace.record(point={k: v for k, v in zip(varnames, value)})
-        return strace, self.log_marginal_likelihood
+        return strace
 
 
 def logp_forw(out_vars, vars, shared):