feat: clean up walnuts a little bit

Author: aseyboldt

src/adapt_strategy.rs

@@ -339,11 +339,12 @@ where
         start: &State<M, P>,
         end: &State<M, P>,
         divergence_info: Option<&DivergenceInfo>,
+        num_substeps: u64,
     ) {
         self.collector1
-            .register_leapfrog(math, start, end, divergence_info);
+            .register_leapfrog(math, start, end, divergence_info, num_substeps);
         self.collector2
-            .register_leapfrog(math, start, end, divergence_info);
+            .register_leapfrog(math, start, end, divergence_info, num_substeps);
     }
 
     fn register_draw(&mut self, math: &mut M, state: &State<M, P>, info: &crate::nuts::SampleInfo) {

src/chain.rs

@@ -183,6 +183,7 @@ where
             &mut self.hamiltonian,
             &self.options,
             &mut self.collector,
+            self.draw_count < 70,
         )?;
         let mut position: Box<[f64]> = vec![0f64; math.dim()].into();
         state.write_position(math, &mut position);
@@ -237,6 +238,7 @@ pub struct NutsStatsBuilder<M: Math, A: AdaptStrategy<M>> {
     divergence_start_grad: Option<FixedSizeListBuilder<PrimitiveBuilder<Float64Type>>>,
     divergence_end: Option<FixedSizeListBuilder<PrimitiveBuilder<Float64Type>>>,
     divergence_momentum: Option<FixedSizeListBuilder<PrimitiveBuilder<Float64Type>>>,
+    non_reversible: Option<BooleanBuilder>,
     divergence_msg: Option<StringBuilder>,
 }
 
@@ -274,7 +276,9 @@ impl<M: Math, A: AdaptStrategy<M>> NutsStatsBuilder<M, A> {
             None
         };
 
-        let (div_start, div_start_grad, div_end, div_mom, div_msg) = if options.store_divergences {
+        let (div_start, div_start_grad, div_end, div_mom, non_rev, div_msg) = if options
+            .store_divergences
+        {
             let start_location_prim = PrimitiveBuilder::new();
             let start_location_list = FixedSizeListBuilder::new(start_location_prim, dim as i32);
 
@@ -288,17 +292,20 @@ impl<M: Math, A: AdaptStrategy<M>> NutsStatsBuilder<M, A> {
             let momentum_location_list =
                 FixedSizeListBuilder::new(momentum_location_prim, dim as i32);
 
+            let non_reversible = BooleanBuilder::new();
+
             let msg_list = StringBuilder::new();
 
             (
                 Some(start_location_list),
                 Some(start_grad_list),
                 Some(end_location_list),
                 Some(momentum_location_list),
+                Some(non_reversible),
                 Some(msg_list),
             )
         } else {
-            (None, None, None, None, None)
+            (None, None, None, None, None, None)
         };
 
         Self {
@@ -320,6 +327,7 @@ impl<M: Math, A: AdaptStrategy<M>> NutsStatsBuilder<M, A> {
             divergence_start_grad: div_start_grad,
             divergence_end: div_end,
             divergence_momentum: div_mom,
+            non_reversible: non_rev,
             divergence_msg: div_msg,
         }
     }
@@ -350,6 +358,7 @@ impl<M: Math, R: rand::Rng, A: AdaptStrategy<M>> StatTraceBuilder<M, NutsChain<M
             divergence_start_grad,
             divergence_end,
             divergence_momentum,
+            non_reversible,
             divergence_msg,
         } = self;
 
@@ -414,6 +423,14 @@ impl<M: Math, R: rand::Rng, A: AdaptStrategy<M>> StatTraceBuilder<M, NutsChain<M
             n_dim,
         );
 
+        if let Some(non_rev) = non_reversible.as_mut() {
+            if let Some(info) = div_info {
+                non_rev.append_value(info.non_reversible);
+            } else {
+                non_rev.append_null();
+            }
+        }
+
         if let Some(div_msg) = divergence_msg.as_mut() {
             if let Some(err) = div_info.and_then(|info| info.logp_function_error.as_ref()) {
                 div_msg.append_value(format!("{err}"));
@@ -447,6 +464,7 @@ impl<M: Math, R: rand::Rng, A: AdaptStrategy<M>> StatTraceBuilder<M, NutsChain<M
             divergence_start_grad,
             divergence_end,
             divergence_momentum,
+            non_reversible,
             divergence_msg,
         } = self;
 
@@ -541,6 +559,11 @@ impl<M: Math, R: rand::Rng, A: AdaptStrategy<M>> StatTraceBuilder<M, NutsChain<M
             &mut fields,
         );
 
+        if let Some(mut non_reversible) = non_reversible {
+            fields.push(Field::new("non_reversible", DataType::Boolean, true));
+            arrays.push(ArrayBuilder::finish(&mut non_reversible));
+        }
+
         let fields = Fields::from(fields);
         Some(StructArray::new(fields, arrays, None))
     }
@@ -565,6 +588,7 @@ impl<M: Math, R: rand::Rng, A: AdaptStrategy<M>> StatTraceBuilder<M, NutsChain<M
             divergence_start_grad,
             divergence_end,
             divergence_momentum,
+            non_reversible,
             divergence_msg,
         } = self;
 
@@ -659,6 +683,11 @@ impl<M: Math, R: rand::Rng, A: AdaptStrategy<M>> StatTraceBuilder<M, NutsChain<M
             &mut fields,
         );
 
+        if let Some(non_reversible) = non_reversible {
+            fields.push(Field::new("non_reversible", DataType::Boolean, true));
+            arrays.push(ArrayBuilder::finish_cloned(non_reversible));
+        }
+
         let fields = Fields::from(fields);
         Some(StructArray::new(fields, arrays, None))
     }

src/euclidean_hamiltonian.rs

@@ -309,7 +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,
+        step_size_splits: u64,
         collector: &mut C,
     ) -> LeapfrogResult<M, Self::Point> {
         let mut out = self.pool().new_state(math);
@@ -322,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 * step_size_factor;
+        let epsilon = (sign as f64) * self.step_size / (step_size_splits as f64);
 
         start
             .point()
@@ -334,17 +334,9 @@ impl<M: Math, Mass: MassMatrix<M>> Hamiltonian<M> for EuclideanHamiltonian<M, Ma
             if !logp_error.is_recoverable() {
                 return LeapfrogResult::Err(logp_error);
             }
-            let div_info = DivergenceInfo {
-                logp_function_error: Some(Arc::new(Box::new(logp_error))),
-                start_location: Some(math.box_array(start.point().position())),
-                start_gradient: Some(math.box_array(&start.point().gradient)),
-                start_momentum: Some(math.box_array(&start.point().momentum)),
-                end_location: None,
-                start_idx_in_trajectory: Some(start.point().index_in_trajectory()),
-                end_idx_in_trajectory: None,
-                energy_error: None,
-            };
-            collector.register_leapfrog(math, start, &out, Some(&div_info));
+            let error = Arc::new(Box::new(logp_error));
+            let div_info = DivergenceInfo::new_logp_function_error(math, start, error);
+            collector.register_leapfrog(math, start, &out, Some(&div_info), step_size_splits);
             return LeapfrogResult::Divergence(div_info);
         }
 
@@ -357,23 +349,21 @@ impl<M: Math, Mass: MassMatrix<M>> Hamiltonian<M> for EuclideanHamiltonian<M, Ma
 
         start.point().set_psum(math, out_point, dir);
 
+        // TODO: energy error measured relative to initial point or previous point?
         let energy_error = out_point.energy_error();
         if (energy_error > self.max_energy_error) | !energy_error.is_finite() {
-            let divergence_info = DivergenceInfo {
-                logp_function_error: None,
-                start_location: Some(math.box_array(start.point().position())),
-                start_gradient: Some(math.box_array(start.point().gradient())),
-                end_location: Some(math.box_array(&out_point.position)),
-                start_momentum: Some(math.box_array(&out_point.momentum)),
-                start_idx_in_trajectory: Some(start.index_in_trajectory()),
-                end_idx_in_trajectory: Some(out.index_in_trajectory()),
-                energy_error: Some(energy_error),
-            };
-            collector.register_leapfrog(math, start, &out, Some(&divergence_info));
+            let divergence_info = DivergenceInfo::new_energy_error_too_large(math, start, &out);
+            collector.register_leapfrog(
+                math,
+                start,
+                &out,
+                Some(&divergence_info),
+                step_size_splits,
+            );
             return LeapfrogResult::Divergence(divergence_info);
         }
 
-        collector.register_leapfrog(math, start, &out, None);
+        collector.register_leapfrog(math, start, &out, None, step_size_splits);
 
         LeapfrogResult::Ok(out)
     }
@@ -447,4 +437,8 @@ impl<M: Math, Mass: MassMatrix<M>> Hamiltonian<M> for EuclideanHamiltonian<M, Ma
     fn step_size_mut(&mut self) -> &mut f64 {
         &mut self.step_size
     }
+
+    fn max_energy_error(&self) -> f64 {
+        self.max_energy_error
+    }
 }

src/hamiltonian.rs

@@ -16,6 +16,7 @@ use crate::{
 ///   a cutoff value or nan.
 /// - The logp function caused a recoverable error (eg if an ODE solver
 ///   failed)
+#[non_exhaustive]
 #[derive(Debug, Clone)]
 pub struct DivergenceInfo {
     pub start_momentum: Option<Box<[f64]>>,
@@ -26,6 +27,7 @@ pub struct DivergenceInfo {
     pub end_idx_in_trajectory: Option<i64>,
     pub start_idx_in_trajectory: Option<i64>,
     pub logp_function_error: Option<Arc<dyn std::error::Error + Send + Sync>>,
+    pub non_reversible: bool,
 }
 
 impl DivergenceInfo {
@@ -39,8 +41,67 @@ impl DivergenceInfo {
             end_idx_in_trajectory: None,
             start_idx_in_trajectory: None,
             logp_function_error: None,
+            non_reversible: false,
         }
     }
+
+    pub fn new_energy_error_too_large<M: Math>(
+        math: &mut M,
+        start: &State<M, impl Point<M>>,
+        stop: &State<M, impl Point<M>>,
+    ) -> Self {
+        DivergenceInfo {
+            logp_function_error: None,
+            start_location: Some(math.box_array(start.point().position())),
+            start_gradient: Some(math.box_array(start.point().gradient())),
+            // TODO
+            start_momentum: None,
+            start_idx_in_trajectory: Some(start.index_in_trajectory()),
+            end_location: Some(math.box_array(&stop.point().position())),
+            end_idx_in_trajectory: Some(stop.index_in_trajectory()),
+            // TODO
+            energy_error: None,
+            non_reversible: false,
+        }
+    }
+
+    pub fn new_logp_function_error<M: Math>(
+        math: &mut M,
+        start: &State<M, impl Point<M>>,
+        logp_function_error: Arc<dyn std::error::Error + Send + Sync>,
+    ) -> Self {
+        DivergenceInfo {
+            logp_function_error: Some(logp_function_error),
+            start_location: Some(math.box_array(start.point().position())),
+            start_gradient: Some(math.box_array(start.point().gradient())),
+            // TODO
+            start_momentum: None,
+            start_idx_in_trajectory: Some(start.index_in_trajectory()),
+            end_location: None,
+            end_idx_in_trajectory: None,
+            energy_error: None,
+            non_reversible: false,
+        }
+    }
+
+    pub fn new_not_reversible<M: Math>(math: &mut M, start: &State<M, impl Point<M>>) -> Self {
+        // TODO add info about what went wrong
+        DivergenceInfo {
+            logp_function_error: None,
+            start_location: Some(math.box_array(start.point().position())),
+            start_gradient: Some(math.box_array(start.point().gradient())),
+            // TODO
+            start_momentum: None,
+            start_idx_in_trajectory: Some(start.index_in_trajectory()),
+            end_location: None,
+            end_idx_in_trajectory: None,
+            energy_error: None,
+            non_reversible: true,
+        }
+    }
+    pub fn new_max_step_size_halvings<M: Math>(math: &mut M, num_steps: u64, info: Self) -> Self {
+        info // TODO
+    }
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -106,10 +167,44 @@ pub trait Hamiltonian<M: Math>: SamplerStats<M> + Sized {
         math: &mut M,
         start: &State<M, Self::Point>,
         dir: Direction,
-        step_size_factor: f64,
+        step_size_splits: u64,
         collector: &mut C,
     ) -> LeapfrogResult<M, Self::Point>;
 
+    fn split_leapfrog<C: Collector<M, Self::Point>>(
+        &mut self,
+        math: &mut M,
+        start: &State<M, Self::Point>,
+        dir: Direction,
+        num_steps: u64,
+        collector: &mut C,
+        max_error: f64,
+    ) -> LeapfrogResult<M, Self::Point> {
+        let mut state = start.clone();
+
+        let mut min_energy = start.energy();
+        let mut max_energy = min_energy;
+
+        for _ in 0..num_steps {
+            state = match self.leapfrog(math, &state, dir, num_steps, collector) {
+                LeapfrogResult::Ok(state) => state,
+                LeapfrogResult::Divergence(info) => return LeapfrogResult::Divergence(info),
+                LeapfrogResult::Err(err) => return LeapfrogResult::Err(err),
+            };
+            let energy = state.energy();
+            min_energy = min_energy.min(energy);
+            max_energy = max_energy.max(energy);
+
+            // TODO: walnuts papers says to use abs, but c++ code doesn't?
+            if max_energy - min_energy > max_error {
+                let info = DivergenceInfo::new_energy_error_too_large(math, start, &state);
+                return LeapfrogResult::Divergence(info);
+            }
+        }
+
+        LeapfrogResult::Ok(state)
+    }
+
     fn is_turning(
         &self,
         math: &mut M,
@@ -141,4 +236,6 @@ pub trait Hamiltonian<M: Math>: SamplerStats<M> + Sized {
 
     fn step_size(&self) -> f64;
     fn step_size_mut(&mut self) -> &mut f64;
+
+    fn max_energy_error(&self) -> f64;
 }

src/lib.rs

@@ -108,7 +108,7 @@ pub use chain::Chain;
 pub use cpu_math::{CpuLogpFunc, CpuMath};
 pub use hamiltonian::DivergenceInfo;
 pub use math_base::{LogpError, Math};
-pub use nuts::NutsError;
+pub use nuts::{NutsError, WalnutsOptions};
 pub use sampler::{
     sample_sequentially, ChainOutput, ChainProgress, DiagGradNutsSettings, DrawStorage,
     LowRankNutsSettings, Model, NutsSettings, Progress, ProgressCallback, Sampler,