Improve walnuts progress info and prepare release

CHANGELOG.md

@@ -2,7 +2,7 @@
 
 All notable changes to this project will be documented in this file.
 
-## [0.19.0] - 2026-10-02
+## [0.19.1] - 2026-10-15
 
 ### Bug Fixes
 
@@ -43,6 +43,14 @@ All notable changes to this project will be documented in this file.
 
 - Implement walnuts (Adrian Seyboldt)
 
+- Expose fisher distance and num_grads in progress (Adrian Seyboldt)
+
+- Add AnyTrace for generic sampler output (Adrian Seyboldt)
+
+- Update flow adaptation defaults (Adrian Seyboldt)
+
+- Limit number of walnuts steps (Adrian Seyboldt)
+
 
 ### Miscellaneous Tasks
 

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "nuts-rs"
-version = "0.19.0"
+version = "0.19.1"
 authors = [
   "Adrian Seyboldt <adrian.seyboldt@gmail.com>",
   "PyMC Developers <pymc.devs@gmail.com>",
@@ -20,7 +20,7 @@ codegen-units = 1
 [dependencies]
 rand = { version = "0.10.0", features = ["chacha"] }
 rand_distr = "0.6.0"
-itertools = "0.14.0"
+itertools = "0.15.0"
 thiserror = "2.0.3"
 anyhow = "1.0.72"
 faer = { version = "0.24.0", default-features = false, features = ["linalg"] }

src/adapt_strategy.rs

@@ -264,6 +264,10 @@ impl<M: Math, A: MassMatrixAdaptStrategy<M>> AdaptStrategy<M> for GlobalStrategy
     fn last_num_steps(&self) -> u64 {
         self.step_size.last_n_steps
     }
+
+    fn last_num_gradients(&self) -> u64 {
+        self.step_size.last_n_gradients
+    }
 }
 
 #[derive(Debug, Storable)]
@@ -419,6 +423,7 @@ mod test {
             max_energy_error: 1000.0,
             uturn_check_first_step: false,
             walnuts: None,
+            max_leapfrog_steps: None,
         };
 
         let rng = {

src/chain.rs

@@ -124,6 +124,7 @@ pub trait AdaptStrategy<M: Math>: SamplerStats<M> {
     fn new_collector(&self, math: &mut M) -> Self::Collector;
     fn is_tuning(&self) -> bool;
     fn last_num_steps(&self) -> u64;
+    fn last_num_gradients(&self) -> u64;
 }
 
 impl<M, R, A> Chain<M> for NutsChain<M, R, A>
@@ -185,6 +186,8 @@ where
             tuning: self.strategy.is_tuning(),
             step_size: self.hamiltonian.step_size(),
             num_steps: self.strategy.last_num_steps(),
+            num_gradients: self.strategy.last_num_gradients(),
+            fisher_distance: state.point().fisher_distance(math),
         };
 
         self.draw_count += 1;
@@ -223,6 +226,9 @@ where
 pub struct NutsStats<P: HasDims, H: Storable<P>, A: Storable<P>, D: Storable<P>> {
     pub depth: u64,
     pub maxdepth_reached: bool,
+    /// Whether the trajectory was stopped because extending it would have
+    /// needed more than `max_leapfrog_steps` leapfrog steps.
+    pub max_leapfrog_steps_reached: bool,
     /// Whether the trajectory was stopped because a WALNUTS macro step was
     /// not reversible. `None` without WALNUTS.
     pub irreversible: Option<bool>,
@@ -287,6 +293,7 @@ impl<M: Math, R: rand::Rng, A: AdaptStrategy<M>> SamplerStats<M> for NutsChain<M
         NutsStats {
             depth: info.depth,
             maxdepth_reached: info.reached_maxdepth,
+            max_leapfrog_steps_reached: info.reached_max_leapfrog_steps,
             irreversible: self.options.walnuts.map(|_| info.irreversible),
             min_micro_steps: self.options.walnuts.map(|w| w.min_micro_steps),
             chain: self.chain,

src/dynamics/hamiltonian.rs

@@ -174,6 +174,10 @@ pub trait Point<M: Math>: Sized + SamplerStats<M> + Debug {
 
     fn initial_energy(&self) -> f64;
 
+    /// `‖z + ∇z‖²` in the transformed space, a single-draw estimate of the
+    /// Fisher divergence between the target and the standard normal.
+    fn fisher_distance(&self, math: &mut M) -> f64;
+
     fn new(math: &mut M) -> Self;
     fn copy_into(&self, math: &mut M, other: &mut Self);
 }

src/dynamics/transformed_hamiltonian.rs

@@ -141,8 +141,7 @@ impl<M: Math> SamplerStats<M> for TransformedPoint<M> {
             transformed_position = Some(math.box_array(&self.transformed_position));
             transformed_gradient = Some(math.box_array(&self.transformed_gradient));
         }
-        let fisher_distance =
-            math.sq_norm_sum(&self.transformed_position, &self.transformed_gradient);
+        let fisher_distance = self.fisher_distance(math);
         PointStats {
             index_in_trajectory: self.index_in_trajectory,
             logp: self.logp,
@@ -371,6 +370,10 @@ impl<M: Math> Point<M> for TransformedPoint<M> {
         self.logp
     }
 
+    fn fisher_distance(&self, math: &mut M) -> f64 {
+        math.sq_norm_sum(&self.transformed_position, &self.transformed_gradient)
+    }
+
     fn new(math: &mut M) -> Self {
         Self {
             untransformed_position: math.new_array(),

src/external_adapt_strategy.rs

@@ -271,4 +271,8 @@ impl<M: Math> AdaptStrategy<M> for ExternalTransformAdaptation {
     fn last_num_steps(&self) -> u64 {
         self.step_size.last_n_steps
     }
+
+    fn last_num_gradients(&self) -> u64 {
+        self.step_size.last_n_gradients
+    }
 }

src/lib.rs

@@ -144,6 +144,7 @@ pub use transform::LowRankSettings;
 #[cfg(feature = "zarr")]
 pub use storage::{ZarrAsyncConfig, ZarrAsyncTraceStorage, ZarrConfig, ZarrTraceStorage};
 
+pub use storage::{AnyStorageConfig, AnyTrace};
 pub use storage::{CsvConfig, CsvTraceStorage};
 pub use storage::{HashMapConfig, HashMapResult, HashMapValue};
 #[cfg(feature = "ndarray")]

src/mclmc.rs

@@ -82,6 +82,9 @@ pub struct MclmcInfo {
     pub divergence_info: Option<DivergenceInfo>,
     /// Number of leapfrog steps actually taken.
     pub num_steps: u64,
+    /// Number of gradient evaluations, including diverging steps that were
+    /// retried with a smaller step size.
+    pub num_gradients: u64,
     pub average_step_size: f64,
 }
 
@@ -267,6 +270,7 @@ where
 
         let mut divergence_info: Option<DivergenceInfo> = None;
         let mut steps_taken = 0u64;
+        let mut num_gradients = 0u64;
 
         // `factor` is the current step size multiplier (always a power of 2).
         let mut factor = 1.0_f64;
@@ -296,6 +300,7 @@ where
             // max_energy_error threshold can exhaust all halvings because the
             // accumulated baseline already sits near the threshold.
             let step_baseline = current.point().energy();
+            num_gradients += 1;
             // We normalize the max energy error by the length of the step
             match self.hamiltonian.leapfrog(
                 math,
@@ -372,13 +377,15 @@ where
                 // TODO: This clone is annoing, get rid of it.
                 divergence_info: divergence_info.clone(),
                 num_steps: steps_taken,
+                num_gradients,
                 average_step_size: time / steps_taken.to_f64().unwrap(),
             };
             let sample_info = crate::nuts::SampleInfo {
                 depth: steps_taken,
                 divergence_info,
                 reached_maxdepth: false,
                 irreversible: false,
+                reached_max_leapfrog_steps: false,
             };
             self.collector.register_draw(math, &current, &sample_info);
             return Ok((next_state, info));
@@ -392,6 +399,7 @@ where
             divergence_info: None,
             reached_maxdepth: false,
             irreversible: false,
+            reached_max_leapfrog_steps: false,
         };
         self.collector.register_draw(math, &current, &sample_info);
 
@@ -404,6 +412,7 @@ where
             diverging: false,
             divergence_info: None,
             num_steps: steps_taken,
+            num_gradients,
             average_step_size: time / steps_taken.to_f64().unwrap(),
         };
 
@@ -513,12 +522,12 @@ where
 
         let (state, info) = self.mclmc_kernel(resample_velocity)?;
 
-        let position: Box<[f64]> = {
+        let (position, fisher_distance): (Box<[f64]>, f64) = {
             let mut math_ = self.math.borrow_mut();
             let math = math_.deref_mut();
             let mut pos = vec![0f64; math.dim()];
             state.write_position(math, &mut pos);
-            pos.into()
+            (pos.into(), state.point().fisher_distance(math))
         };
 
         let mut progress = Progress {
@@ -528,6 +537,8 @@ where
             tuning: true,
             step_size: self.hamiltonian.step_size(),
             num_steps: info.num_steps,
+            num_gradients: info.num_gradients,
+            fisher_distance,
         };
 
         // The collector was already fed during mclmc_kernel via register_leapfrog