fix: incorrect momentum refresh in mclmc

Author: aseyboldt

src/dynamics/hamiltonian.rs

@@ -151,10 +151,9 @@ pub trait Hamiltonian<M: Math>: SamplerStats<M> + Sized {
     /// Perform one leapfrog step.
     ///
     /// `step_size_factor` scales the hamiltonian's base step size for this
-    /// step only: `actual_ε = hamiltonian.step_size() * step_size_factor`.
-    /// The actual step size used is stored on the output point as
-    /// `point.step_size`, so callers can compute importance weights as
-    /// `log(point.step_size) - energy_error`.
+    /// step only.
+    /// `energy_baseline` is the energy value against which the divergence
+    /// check (`|energy_error| >= max_energy_error`) is evaluated.
     ///
     /// Return either an unrecoverable error, a new state or a divergence.
     fn leapfrog<C: Collector<M, Self::Point>>(
@@ -163,6 +162,7 @@ pub trait Hamiltonian<M: Math>: SamplerStats<M> + Sized {
         start: &State<M, Self::Point>,
         dir: Direction,
         step_size_factor: f64,
+        energy_baseline: f64,
         collector: &mut C,
     ) -> LeapfrogResult<M, Self::Point>;
 

src/dynamics/transformed_hamiltonian.rs

@@ -558,6 +558,7 @@ impl<M: Math, T: Transformation<M>> Hamiltonian<M> for TransformedHamiltonian<M,
         start: &State<M, Self::Point>,
         dir: Direction,
         step_size_factor: f64,
+        energy_baseline: f64,
         collector: &mut C,
     ) -> LeapfrogResult<M, Self::Point> {
         let mut out = self.pool().new_state(math);
@@ -616,8 +617,14 @@ impl<M: Math, T: Transformation<M>> Hamiltonian<M> for TransformedHamiltonian<M,
 
         out_point.index_in_trajectory = start.index_in_trajectory() + sign;
 
-        let energy_error = out_point.energy_error();
-        if (energy_error > self.max_energy_error) | !energy_error.is_finite() {
+        let energy_error = out_point.energy() - energy_baseline;
+        let bad_energy = match self.kinetic_energy_kind {
+            KineticEnergyKind::Euclidean | KineticEnergyKind::ExactNormal => {
+                energy_error > self.max_energy_error
+            }
+            KineticEnergyKind::Microcanonical => energy_error.abs() >= self.max_energy_error,
+        };
+        if bad_energy | !energy_error.is_finite() {
             let divergence_info = DivergenceInfo {
                 logp_function_error: None,
                 start_location: Some(math.box_array(start.point().position())),
@@ -807,6 +814,7 @@ impl<M: Math, T: Transformation<M>> Hamiltonian<M> for TransformedHamiltonian<M,
 
         let half_step = self.step_size * factor / 2.0;
 
+        // TODO: Avoid array allocation
         let mut noise = math.new_array();
         math.array_gaussian(rng, &mut noise, &self.ones);
 

src/mclmc.rs

@@ -37,7 +37,7 @@ use crate::{
     Math, NutsError,
     chain::{AdaptStrategy, Chain, StatOptions},
     dynamics::{
-        DivergenceInfo, DivergenceStats, Hamiltonian, KineticEnergyKind, Point, State,
+        Direction, DivergenceInfo, DivergenceStats, Hamiltonian, KineticEnergyKind, Point, State,
         TransformedHamiltonian, TransformedPoint,
     },
     nuts::{Collector, NutsOptions},
@@ -168,6 +168,7 @@ where
     math: RefCell<M>,
     stats_options: StatOptions<M, A>,
     last_info: Option<MclmcInfo>,
+    tmp_velocity: M::Vector,
 }
 
 impl<M, R, A, T> MclmcChain<M, R, A, T>
@@ -191,6 +192,7 @@ where
     ) -> Self {
         let state = hamiltonian.pool().new_state(&mut math);
         let collector = adapt.new_collector(&mut math);
+        let tmp_velocity = math.new_array();
         Self {
             hamiltonian,
             collector,
@@ -207,6 +209,7 @@ where
             math: math.into(),
             stats_options,
             last_info: None,
+            tmp_velocity,
         }
     }
 
@@ -223,12 +226,8 @@ where
             .round()
             .max(1.0) as u64;
 
-        // ── First partial momentum refresh ───────────────────────────────────
-        self.hamiltonian
-            .partial_momentum_refresh(math, &mut self.state, &mut self.rng, 1.0)?;
-
-        // Reset the energy baseline to the post-refresh state so that
-        // energy_error() measures drift over this draw only.
+        // Reset the kinetic-energy accumulator so that energy() at the draw
+        // start is a clean baseline for MclmcInfo.energy_change.
         {
             let kind = self.hamiltonian.kinetic_energy_kind;
 
@@ -251,7 +250,12 @@ where
 
         use crate::dynamics::LeapfrogResult;
 
-        let mut current = self.state.clone();
+        let mut current = self.hamiltonian.copy_state(math, &self.state);
+
+        // Capture the draw-start energy once; used at the end to compute
+        // MclmcInfo.energy_change independently of the per-step baselines.
+        let draw_start_energy = current.point().energy();
+
         let mut divergence_info: Option<DivergenceInfo> = None;
         let mut steps_taken = 0u64;
 
@@ -264,14 +268,44 @@ where
 
         let mut remaining = num_steps;
         while remaining > 0 {
+            // debt == 2 indicates a previous unsuccessful leapfrog
+            if debt != 2 {
+                math.copy_into(&current.point().velocity, &mut self.tmp_velocity);
+            } else {
+                // Restore the old velocity and refresh momentum with smaller factor
+                // TODO: I think we should reuse the original gaussian noise?
+                math.copy_into(
+                    &self.tmp_velocity,
+                    &mut current.try_point_mut().unwrap().velocity,
+                );
+            }
+            self.hamiltonian
+                .partial_momentum_refresh(math, &mut current, &mut self.rng, factor)?;
+
+            // Use the post-refresh energy as the divergence baseline so that
+            // the leapfrog's energy_error measures only this single step's
+            // integration error (O(ε²)), not the cumulative drift from the
+            // draw start.  Without this, many small steps with a tight
+            // max_energy_error threshold can exhaust all halvings because the
+            // accumulated baseline already sits at the threshold.
+            let step_baseline = current.point().energy();
             let next = match self.hamiltonian.leapfrog(
                 math,
                 &current,
-                crate::dynamics::Direction::Forward,
+                Direction::Forward,
                 factor,
+                step_baseline,
                 &mut self.collector,
             ) {
-                LeapfrogResult::Ok(next) => next,
+                LeapfrogResult::Ok(mut next) => {
+                    self.hamiltonian.partial_momentum_refresh(
+                        math,
+                        &mut next,
+                        &mut self.rng,
+                        factor,
+                    )?;
+                    next
+                }
                 LeapfrogResult::Divergence(info) => {
                     if halvings >= max_halvings {
                         // Genuinely diverged — give up.
@@ -283,6 +317,7 @@ where
                     factor *= 0.5;
                     halvings += 1;
                     debt = 2;
+                    remaining += 1;
                     continue;
                 }
                 LeapfrogResult::Err(e) => {
@@ -318,7 +353,7 @@ where
             let mut next_state = self.hamiltonian.copy_state(math, &self.state);
             self.hamiltonian
                 .initialize_trajectory(math, &mut next_state, &mut self.rng)?;
-            let energy_change = current.point().energy_error();
+            let energy_change = current.point().energy() - draw_start_energy;
             let info = MclmcInfo {
                 energy_change,
                 diverging: true,
@@ -336,18 +371,14 @@ where
         };
         self.collector.register_draw(math, &current, &sample_info);
 
-        let energy_change = current.point().energy_error();
+        let energy_change = current.point().energy() - draw_start_energy;
 
-        // ── Second partial momentum refresh ──────────────────────────────────
-        let mut next_state = if current.try_point_mut().is_err() {
+        let next_state = if current.try_point_mut().is_err() {
             self.hamiltonian.copy_state(math, &current)
         } else {
             current
         };
 
-        self.hamiltonian
-            .partial_momentum_refresh(math, &mut next_state, &mut self.rng, 1.0)?;
-
         let info = MclmcInfo {
             energy_change,
             diverging: false,

src/nuts.rs

@@ -217,7 +217,14 @@ impl<M: Math, H: Hamiltonian<M>, C: Collector<M, H::Point>> NutsTree<M, H, C> {
             Direction::Forward => &self.right,
             Direction::Backward => &self.left,
         };
-        let end = match hamiltonian.leapfrog(math, start, direction, 1.0, collector) {
+        let end = match hamiltonian.leapfrog(
+            math,
+            start,
+            direction,
+            1.0,
+            start.point().initial_energy(),
+            collector,
+        ) {
             LeapfrogResult::Divergence(info) => return Ok(Err(info)),
             LeapfrogResult::Err(err) => return Err(NutsError::LogpFailure(err.into())),
             LeapfrogResult::Ok(end) => end,

src/stepsize/adapt.rs

@@ -109,8 +109,14 @@ impl Strategy {
 
         *hamiltonian.step_size_mut() = self.options.initial_step;
 
-        let state_next =
-            hamiltonian.leapfrog(math, &state, Direction::Forward, 1.0, &mut collector);
+        let state_next = hamiltonian.leapfrog(
+            math,
+            &state,
+            Direction::Forward,
+            1.0,
+            state.point().initial_energy(),
+            &mut collector,
+        );
 
         let LeapfrogResult::Ok(_) = state_next else {
             return Ok(());
@@ -126,7 +132,14 @@ impl Strategy {
         for _ in 0..100 {
             let mut collector = AcceptanceRateCollector::new();
             collector.register_init(math, &state, options);
-            let state_next = hamiltonian.leapfrog(math, &state, dir, 1.0, &mut collector);
+            let state_next = hamiltonian.leapfrog(
+                math,
+                &state,
+                dir,
+                1.0,
+                state.point().initial_energy(),
+                &mut collector,
+            );
             let LeapfrogResult::Ok(_) = state_next else {
                 *hamiltonian.step_size_mut() = self.options.initial_step;
                 return Ok(());