Export energy for each sample in nuts

Author: aseyboldt

pymc3/step_methods/hmc/nuts.py

@@ -36,6 +36,7 @@ class NUTS(BaseHMC):
         'tree_size': np.float64,
         'diverging': np.bool,
         'energy_change': np.float64,
+        'energy': np.float64,
         'max_energy_change': np.float64,
     }]
 
@@ -262,6 +263,7 @@ def stats(self):
             'depth': self.depth,
             'mean_tree_accept': self.accept_sum / self.n_proposals,
             'energy_change': self.proposal.energy - self.start.energy,
+            'energy': self.proposal.energy,
             'tree_size': self.n_proposals,
             'max_energy_change': self.max_energy_change,
         }

pymc3/tests/sampler_fixtures.py

@@ -36,7 +36,6 @@ def test_kstest(self):
             elif samples.ndim == 2:
                 pvals = []
                 for samples_, cdf_ in zip(samples.T, cdf):
-                    print(samples_)
                     t, p = stats.kstest(samples_[::self.ks_thin], cdf=cdf_)
                     pvals.append(p)
                 t, p = stats.combine_pvalues(pvals)
@@ -115,7 +114,6 @@ def setUpClass(cls):
 
     def test_neff(self):
         if hasattr(self, 'min_n_eff'):
-            print(len(self.trace))
             n_eff = pm.effective_n(self.trace[self.burn:])
             for var in n_eff:
                 npt.assert_array_less(self.min_n_eff, n_eff[var])