Fix tests and add doc for new statistics

Author: aseyboldt

pymc/step_methods/hmc/hmc.py

@@ -51,6 +51,8 @@ class HamiltonianMC(BaseHMC):
             "process_time_diff": np.float64,
             "perf_counter_diff": np.float64,
             "perf_counter_start": np.float64,
+            "largest_eigval": np.float64,
+            "smallest_eigval": np.float64,
         }
     ]
 
@@ -162,7 +164,16 @@ def _hamiltonian_step(self, start, p0, step_size):
             "model_logp": state.model_logp,
         }
         # Retrieve State q and p data from respective RaveledVars
-        end = State(end.q.data, end.p.data, end.v, end.q_grad, end.energy, end.model_logp)
+        end = State(
+            end.q.data,
+            end.p.data,
+            end.v,
+            end.q_grad,
+            end.energy,
+            end.model_logp,
+            end.index_in_trajectory,
+        )
+        stats.update(self.potential.stats())
         return HMCStepData(end, accept_stat, div_info, stats)
 
     @staticmethod

pymc/step_methods/hmc/nuts.py

@@ -78,6 +78,12 @@ class NUTS(BaseHMC):
       by the python standard library `time.perf_counter` (wall time).
     - `perf_counter_start`: The value of `time.perf_counter` at the beginning
       of the computation of the draw.
+    - `index_in_trajectory`: This is usually only interesting for debugging
+      purposes. This indicates the position of the posterior draw in the
+      trajectory. Eg a -4 would indicate that the draw was the result of the
+      fourth leapfrog step in negative direction.
+    - `largest_eigval` and `smallest_eigval`: Experimental statistics for
+      some mass matrix adaptation algorithms. This is nan if it is not used.
 
     References
     ----------

pymc/tests/test_step.py

@@ -595,6 +595,9 @@ def test_sampler_stats(self):
             "perf_counter_diff",
             "perf_counter_start",
             "process_time_diff",
+            "index_in_trajectory",
+            "largest_eigval",
+            "smallest_eigval",
         }
         assert trace.stat_names == expected_stat_names
         for varname in trace.stat_names: