feat: add untested walnuts implementation

Author: aseyboldt

src/adapt_strategy.rs

@@ -503,6 +503,7 @@ mod test {
             store_unconstrained: true,
             check_turning: true,
             store_divergences: false,
+            walnuts_options: None,
         };
 
         let rng = {

src/euclidean_hamiltonian.rs

@@ -309,6 +309,7 @@ impl<M: Math, Mass: MassMatrix<M>> Hamiltonian<M> for EuclideanHamiltonian<M, Ma
         math: &mut M,
         start: &State<M, Self::Point>,
         dir: Direction,
+        step_size_factor: f64,
         collector: &mut C,
     ) -> LeapfrogResult<M, Self::Point> {
         let mut out = self.pool().new_state(math);
@@ -321,7 +322,7 @@ impl<M: Math, Mass: MassMatrix<M>> Hamiltonian<M> for EuclideanHamiltonian<M, Ma
             Direction::Backward => -1,
         };
 
-        let epsilon = (sign as f64) * self.step_size;
+        let epsilon = (sign as f64) * self.step_size * step_size_factor;
 
         start
             .point()

src/hamiltonian.rs

@@ -28,12 +28,36 @@ pub struct DivergenceInfo {
     pub logp_function_error: Option<Arc<dyn std::error::Error + Send + Sync>>,
 }
 
+impl DivergenceInfo {
+    pub fn new() -> Self {
+        DivergenceInfo {
+            start_momentum: None,
+            start_location: None,
+            start_gradient: None,
+            end_location: None,
+            energy_error: None,
+            end_idx_in_trajectory: None,
+            start_idx_in_trajectory: None,
+            logp_function_error: None,
+        }
+    }
+}
+
 #[derive(Debug, Copy, Clone)]
 pub enum Direction {
     Forward,
     Backward,
 }
 
+impl Direction {
+    pub fn reverse(&self) -> Self {
+        match self {
+            Direction::Forward => Direction::Backward,
+            Direction::Backward => Direction::Forward,
+        }
+    }
+}
+
 impl Distribution<Direction> for StandardUniform {
     fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Direction {
         if rng.random::<bool>() {
@@ -82,6 +106,7 @@ pub trait Hamiltonian<M: Math>: SamplerStats<M> + Sized {
         math: &mut M,
         start: &State<M, Self::Point>,
         dir: Direction,
+        step_size_factor: f64,
         collector: &mut C,
     ) -> LeapfrogResult<M, Self::Point>;
 

src/nuts.rs

@@ -120,7 +120,7 @@ impl<M: Math, H: Hamiltonian<M>, C: Collector<M, H::Point>> NutsTree<M, H, C> {
         H: Hamiltonian<M>,
         R: rand::Rng + ?Sized,
     {
-        let mut other = match self.single_step(math, hamiltonian, direction, collector) {
+        let mut other = match self.single_step(math, hamiltonian, direction, options, collector) {
             Ok(Ok(tree)) => tree,
             Ok(Err(info)) => return ExtendResult::Diverging(self, info),
             Err(err) => return ExtendResult::Err(err),
@@ -213,19 +213,141 @@ impl<M: Math, H: Hamiltonian<M>, C: Collector<M, H::Point>> NutsTree<M, H, C> {
         math: &mut M,
         hamiltonian: &mut H,
         direction: Direction,
+        options: &NutsOptions,
         collector: &mut C,
     ) -> Result<std::result::Result<NutsTree<M, H, C>, DivergenceInfo>> {
         let start = match direction {
             Direction::Forward => &self.right,
             Direction::Backward => &self.left,
         };
-        let end = match hamiltonian.leapfrog(math, start, direction, collector) {
-            LeapfrogResult::Divergence(info) => return Ok(Err(info)),
-            LeapfrogResult::Err(err) => return Err(NutsError::LogpFailure(err.into())),
-            LeapfrogResult::Ok(end) => end,
+
+        let (log_size, end) = match options.walnuts_options {
+            Some(ref options) => {
+                // Walnuts implementation
+                // TODO: Shouldn't all be in this one big function...
+                let mut step_size_factor = 1.0;
+                let mut num_steps = 1;
+                let mut current = start.clone();
+
+                let mut success = false;
+
+                'step_size_search: for _ in 0..options.max_step_size_halvings {
+                    current = start.clone();
+                    let mut min_energy = current.energy();
+                    let mut max_energy = min_energy;
+
+                    for _ in 0..num_steps {
+                        current = match hamiltonian.leapfrog(
+                            math,
+                            &current,
+                            direction,
+                            step_size_factor,
+                            collector,
+                        ) {
+                            LeapfrogResult::Ok(state) => state,
+                            LeapfrogResult::Divergence(_) => {
+                                num_steps *= 2;
+                                step_size_factor *= 0.5;
+                                continue 'step_size_search;
+                            }
+                            LeapfrogResult::Err(err) => {
+                                return Err(NutsError::LogpFailure(err.into()));
+                            }
+                        };
+
+                        // Update min/max energies
+                        let current_energy = current.energy();
+                        min_energy = min_energy.min(current_energy);
+                        max_energy = max_energy.max(current_energy);
+                    }
+
+                    if max_energy - min_energy > options.max_energy_error {
+                        num_steps *= 2;
+                        step_size_factor *= 0.5;
+                        continue 'step_size_search;
+                    }
+
+                    success = true;
+                    break 'step_size_search;
+                }
+
+                if !success {
+                    // TODO: More info
+                    return Ok(Err(DivergenceInfo::new()));
+                }
+
+                // TODO
+                let back = direction.reverse();
+                let mut current_backward;
+
+                let mut reversible = true;
+
+                'rev_step_size: while num_steps >= 2 {
+                    num_steps /= 2;
+                    step_size_factor *= 0.5;
+
+                    // TODO: Can we share code for the micro steps in the two directions?
+                    current_backward = current.clone();
+
+                    let mut min_energy = current_backward.energy();
+                    let mut max_energy = min_energy;
+
+                    for _ in 0..num_steps {
+                        current_backward = match hamiltonian.leapfrog(
+                            math,
+                            &current_backward,
+                            back,
+                            step_size_factor,
+                            collector,
+                        ) {
+                            LeapfrogResult::Ok(state) => state,
+                            LeapfrogResult::Divergence(_) => {
+                                // We also reject in the backward direction, all is good so far...
+                                continue 'rev_step_size;
+                            }
+                            LeapfrogResult::Err(err) => {
+                                return Err(NutsError::LogpFailure(err.into()));
+                            }
+                        };
+
+                        // Update min/max energies
+                        let current_energy = current_backward.energy();
+                        min_energy = min_energy.min(current_energy);
+                        max_energy = max_energy.max(current_energy);
+                        if max_energy - min_energy > options.max_energy_error {
+                            // We reject also in the backward direction, all good so far...
+                            continue 'rev_step_size;
+                        }
+                    }
+
+                    // We did not reject in the backward direction, so we are not reversible
+                    reversible = false;
+                    break;
+                }
+
+                if reversible {
+                    let log_size = -current.point().energy_error();
+                    (log_size, current)
+                } else {
+                    // TODO: More info
+                    return Ok(Err(DivergenceInfo::new()));
+                }
+            }
+            None => {
+                // Classical NUTS
+                //
+                let end = match hamiltonian.leapfrog(math, start, direction, 1.0, collector) {
+                    LeapfrogResult::Divergence(info) => return Ok(Err(info)),
+                    LeapfrogResult::Err(err) => return Err(NutsError::LogpFailure(err.into())),
+                    LeapfrogResult::Ok(end) => end,
+                };
+
+                let log_size = -end.point().energy_error();
+
+                (log_size, end)
+            }
         };
 
-        let log_size = -end.point().energy_error();
         Ok(Ok(NutsTree {
             right: end.clone(),
             left: end.clone(),
@@ -248,12 +370,21 @@ impl<M: Math, H: Hamiltonian<M>, C: Collector<M, H::Point>> NutsTree<M, H, C> {
     }
 }
 
+#[derive(Debug, Clone, Copy)]
+pub struct WalnutsOptions {
+    pub max_energy_error: f64,
+    pub max_step_size_halvings: u64,
+}
+
+#[derive(Debug, Clone, Copy)]
 pub struct NutsOptions {
     pub maxdepth: u64,
     pub store_gradient: bool,
     pub store_unconstrained: bool,
     pub check_turning: bool,
     pub store_divergences: bool,
+
+    pub walnuts_options: Option<WalnutsOptions>,
 }
 
 pub(crate) fn draw<M, H, R, C>(

src/sampler.rs

@@ -24,7 +24,7 @@ use crate::{
     mass_matrix::DiagMassMatrix,
     mass_matrix_adapt::Strategy as DiagMassMatrixStrategy,
     math_base::Math,
-    nuts::NutsOptions,
+    nuts::{NutsOptions, WalnutsOptions},
     sampler_stats::{SamplerStats, StatTraceBuilder},
     transform_adapt_strategy::{TransformAdaptation, TransformedSettings},
     transformed_hamiltonian::{TransformedHamiltonian, TransformedPointStatsOptions},
@@ -102,6 +102,7 @@ pub struct NutsSettings<A: Debug + Copy + Default> {
 
     pub num_chains: usize,
     pub seed: u64,
+    pub walnuts_options: Option<WalnutsOptions>,
 }
 
 pub type DiagGradNutsSettings = NutsSettings<EuclideanAdaptOptions<DiagAdaptExpSettings>>;
@@ -122,6 +123,7 @@ impl Default for DiagGradNutsSettings {
             check_turning: true,
             seed: 0,
             num_chains: 6,
+            walnuts_options: None,
         }
     }
 }
@@ -140,6 +142,7 @@ impl Default for LowRankNutsSettings {
             check_turning: true,
             seed: 0,
             num_chains: 6,
+            walnuts_options: None,
         };
         vals.adapt_options.mass_matrix_update_freq = 10;
         vals
@@ -160,6 +163,7 @@ impl Default for TransformedNutsSettings {
             check_turning: true,
             seed: 0,
             num_chains: 1,
+            walnuts_options: None,
         }
     }
 }
@@ -191,6 +195,7 @@ impl Settings for LowRankNutsSettings {
             store_divergences: self.store_divergences,
             store_unconstrained: self.store_unconstrained,
             check_turning: self.check_turning,
+            walnuts_options: self.walnuts_options,
         };
 
         let rng = rand::rngs::SmallRng::try_from_rng(&mut rng).expect("Could not seed rng");
@@ -250,6 +255,7 @@ impl Settings for DiagGradNutsSettings {
             store_divergences: self.store_divergences,
             store_unconstrained: self.store_unconstrained,
             check_turning: self.check_turning,
+            walnuts_options: self.walnuts_options,
         };
 
         let rng = rand::rngs::SmallRng::try_from_rng(&mut rng).expect("Could not seed rng");
@@ -306,6 +312,7 @@ impl Settings for TransformedNutsSettings {
             store_divergences: self.store_divergences,
             store_unconstrained: self.store_unconstrained,
             check_turning: self.check_turning,
+            walnuts_options: self.walnuts_options,
         };
 
         let rng = rand::rngs::SmallRng::try_from_rng(&mut rng).expect("Could not seed rng");

src/stepsize_adapt.rs

@@ -52,7 +52,8 @@ impl Strategy {
 
         *hamiltonian.step_size_mut() = self.options.initial_step;
 
-        let state_next = hamiltonian.leapfrog(math, &state, Direction::Forward, &mut collector);
+        let state_next =
+            hamiltonian.leapfrog(math, &state, Direction::Forward, 1.0, &mut collector);
 
         let LeapfrogResult::Ok(_) = state_next else {
             return Ok(());
@@ -68,7 +69,7 @@ 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, &mut collector);
+            let state_next = hamiltonian.leapfrog(math, &state, dir, 1.0, &mut collector);
             let LeapfrogResult::Ok(_) = state_next else {
                 *hamiltonian.step_size_mut() = self.options.initial_step;
                 return Ok(());

src/transformed_hamiltonian.rs

@@ -456,6 +456,7 @@ impl<M: Math> Hamiltonian<M> for TransformedHamiltonian<M> {
         math: &mut M,
         start: &State<M, Self::Point>,
         dir: Direction,
+        step_size_factor: f64,
         collector: &mut C,
     ) -> LeapfrogResult<M, Self::Point> {
         let mut out = self.pool().new_state(math);
@@ -469,7 +470,7 @@ impl<M: Math> Hamiltonian<M> for TransformedHamiltonian<M> {
             Direction::Backward => -1,
         };
 
-        let epsilon = (sign as f64) * self.step_size;
+        let epsilon = (sign as f64) * self.step_size * step_size_factor;
 
         start
             .point()