WIP Add walnuts implementation

src/adapt_strategy.rs

@@ -143,6 +143,8 @@ impl<M: Math, A: MassMatrixAdaptStrategy<M>> AdaptStrategy<M> for GlobalStrategy
         self.step_size.update(&collector.collector1);
 
         if draw >= self.num_tune {
+            // Needed for step size jitter
+            self.step_size.update_stepsize(rng, hamiltonian, true);
             self.tuning = false;
             return Ok(());
         }
@@ -194,14 +196,14 @@ impl<M: Math, A: MassMatrixAdaptStrategy<M>> AdaptStrategy<M> for GlobalStrategy
                 self.step_size
                     .init(math, options, hamiltonian, &position, rng)?;
             } else {
-                self.step_size.update_stepsize(hamiltonian, false)
+                self.step_size.update_stepsize(rng, hamiltonian, false)
             }
             return Ok(());
         }
 
         self.step_size.update_estimator_late();
         let is_last = draw == self.num_tune - 1;
-        self.step_size.update_stepsize(hamiltonian, is_last);
+        self.step_size.update_stepsize(rng, hamiltonian, is_last);
         Ok(())
     }
 
@@ -339,11 +341,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) {
@@ -503,6 +506,7 @@ mod test {
             store_unconstrained: true,
             check_turning: true,
             store_divergences: false,
+            walnuts_options: None,
         };
 
         let rng = {

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,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_splits: u64,
         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_splits as f64);
 
         start
             .point()
@@ -333,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);
         }
 
@@ -356,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)
     }
@@ -446,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
+    }
 }