feat: add max_energy_error diagnostic

Author: aseyboldt

src/adapt_strategy.rs

@@ -476,7 +476,7 @@ mod test {
                 step_size: (),
                 mass_matrix: (),
             },
-            hamiltonian: (),
+            hamiltonian: -1i64,
             point: TransformedPointStatsOptions {
                 store_gradient: true,
                 store_unconstrained: true,

src/stepsize/adapt.rs

@@ -60,6 +60,8 @@ pub struct Strategy {
     pub last_sym_mean_tree_accept: f64,
     /// Last number of steps
     pub last_n_steps: u64,
+    /// Maximum absolute energy error observed in the last trajectory
+    pub last_max_energy_error: f64,
 }
 
 impl Strategy {
@@ -82,6 +84,7 @@ impl Strategy {
             last_n_steps: 0,
             last_sym_mean_tree_accept: 0.0,
             last_mean_tree_accept: 0.0,
+            last_max_energy_error: 0.0,
         }
     }
 
@@ -187,6 +190,7 @@ impl Strategy {
         self.last_mean_tree_accept = mean;
         self.last_sym_mean_tree_accept = mean_sym;
         self.last_n_steps = n_steps;
+        self.last_max_energy_error = collector.max_energy_error;
     }
 
     pub fn update_estimator_early(&mut self) {
@@ -258,6 +262,7 @@ pub struct Stats {
     pub mean_tree_accept: f64,
     pub mean_tree_accept_sym: f64,
     pub n_steps: u64,
+    pub max_energy_error: f64,
 }
 
 impl<M: Math> SamplerStats<M> for Strategy {
@@ -280,6 +285,7 @@ impl<M: Math> SamplerStats<M> for Strategy {
             mean_tree_accept: self.last_mean_tree_accept,
             mean_tree_accept_sym: self.last_sym_mean_tree_accept,
             n_steps: self.last_n_steps,
+            max_energy_error: self.last_max_energy_error,
         }
     }
 }

src/stepsize/dual_avg.rs

@@ -113,6 +113,7 @@ pub struct AcceptanceRateCollector {
     initial_energy: f64,
     pub(crate) mean: RunningMean,
     pub(crate) mean_sym: RunningMean,
+    pub(crate) max_energy_error: f64,
 }
 
 impl AcceptanceRateCollector {
@@ -121,6 +122,7 @@ impl AcceptanceRateCollector {
             initial_energy: 0.,
             mean: RunningMean::new(),
             mean_sym: RunningMean::new(),
+            max_energy_error: 0.,
         }
     }
 }
@@ -137,6 +139,7 @@ impl<M: Math, P: Point<M>> Collector<M, P> for AcceptanceRateCollector {
             Some(_) => {
                 self.mean.add(0.);
                 self.mean_sym.add(0.);
+                self.max_energy_error = f64::NEG_INFINITY;
             }
             None => {
                 let base_energy = self.initial_energy;
@@ -146,6 +149,10 @@ impl<M: Math, P: Point<M>> Collector<M, P> for AcceptanceRateCollector {
                 self.mean.add(diff.min(0.).exp());
                 self.mean_sym
                     .add(2. * diff.min(0.).exp() / (1. + diff.exp()));
+                let energy_error = diff;
+                if energy_error.abs() > self.max_energy_error.abs() {
+                    self.max_energy_error = energy_error;
+                }
             }
         };
     }
@@ -154,5 +161,6 @@ impl<M: Math, P: Point<M>> Collector<M, P> for AcceptanceRateCollector {
         self.initial_energy = state.energy();
         self.mean.reset();
         self.mean_sym.reset();
+        self.max_energy_error = 0.;
     }
 }