feat: add mindepth option for nuts

Author: aseyboldt

src/adapt_strategy.rs

@@ -501,6 +501,7 @@ mod test {
             EuclideanHamiltonian::new(&mut math, mass_matrix, max_energy_error, step_size);
         let options = NutsOptions {
             maxdepth: 10u64,
+            mindepth: 0,
             store_gradient: true,
             store_unconstrained: true,
             check_turning: true,

src/nuts.rs

@@ -250,6 +250,7 @@ impl<M: Math, H: Hamiltonian<M>, C: Collector<M, H::Point>> NutsTree<M, H, C> {
 
 pub struct NutsOptions {
     pub maxdepth: u64,
+    pub mindepth: u64,
     pub store_gradient: bool,
     pub store_unconstrained: bool,
     pub check_turning: bool,
@@ -286,9 +287,13 @@ where
         tree = match tree.extend(math, rng, hamiltonian, direction, collector, options) {
             ExtendResult::Ok(tree) => tree,
             ExtendResult::Turning(tree) => {
-                let info = tree.info(false, None);
-                collector.register_draw(math, &tree.draw, &info);
-                return Ok((tree.draw, info));
+                if tree.depth < options.mindepth {
+                    tree
+                } else {
+                    let info = tree.info(false, None);
+                    collector.register_draw(math, &tree.draw, &info);
+                    return Ok((tree.draw, info));
+                }
             }
             ExtendResult::Diverging(tree, info) => {
                 let info = tree.info(false, Some(info));

src/sampler.rs

@@ -87,6 +87,9 @@ pub struct NutsSettings<A: Debug + Copy + Default> {
     /// The maximum tree depth during sampling. The number of leapfrog steps
     /// is smaller than 2 ^ maxdepth.
     pub maxdepth: u64,
+    /// The minimum tree depth during sampling. The number of leapfrog steps
+    /// is larger than 2 ^ mindepth.
+    pub mindepth: u64,
     /// Store the gradient in the SampleStats
     pub store_gradient: bool,
     /// Store each unconstrained parameter vector in the sampler stats
@@ -114,6 +117,7 @@ impl Default for DiagGradNutsSettings {
             num_tune: 400,
             num_draws: 1000,
             maxdepth: 10,
+            mindepth: 0,
             max_energy_error: 1000f64,
             store_gradient: false,
             store_unconstrained: false,
@@ -132,6 +136,7 @@ impl Default for LowRankNutsSettings {
             num_tune: 800,
             num_draws: 1000,
             maxdepth: 10,
+            mindepth: 0,
             max_energy_error: 1000f64,
             store_gradient: false,
             store_unconstrained: false,
@@ -152,6 +157,7 @@ impl Default for TransformedNutsSettings {
             num_tune: 1500,
             num_draws: 1000,
             maxdepth: 10,
+            mindepth: 0,
             max_energy_error: 20f64,
             store_gradient: false,
             store_unconstrained: false,
@@ -187,6 +193,7 @@ impl Settings for LowRankNutsSettings {
 
         let options = NutsOptions {
             maxdepth: self.maxdepth,
+            mindepth: self.mindepth,
             store_gradient: self.store_gradient,
             store_divergences: self.store_divergences,
             store_unconstrained: self.store_unconstrained,
@@ -246,6 +253,7 @@ impl Settings for DiagGradNutsSettings {
 
         let options = NutsOptions {
             maxdepth: self.maxdepth,
+            mindepth: self.mindepth,
             store_gradient: self.store_gradient,
             store_divergences: self.store_divergences,
             store_unconstrained: self.store_unconstrained,
@@ -302,6 +310,7 @@ impl Settings for TransformedNutsSettings {
 
         let options = NutsOptions {
             maxdepth: self.maxdepth,
+            mindepth: self.mindepth,
             store_gradient: self.store_gradient,
             store_divergences: self.store_divergences,
             store_unconstrained: self.store_unconstrained,