fix: correctly specify dims of some sample stats

Author: aseyboldt

src/adapt_strategy.rs

@@ -210,6 +210,7 @@ pub struct GlobalStrategyStats<P: HasDims, S: Storable<P>, M: Storable<P>> {
     pub step_size: S,
     #[storable(flatten)]
     pub mass_matrix: M,
+    pub tuning: bool,
     #[storable(ignore)]
     _phantom: std::marker::PhantomData<fn() -> P>,
 }
@@ -243,6 +244,7 @@ where
                 self.step_size.extract_stats(math, ())
             },
             mass_matrix: self.mass_matrix.extract_stats(math, opt.mass_matrix),
+            tuning: self.tuning,
             _phantom: PhantomData,
         }
     }

src/chain.rs

@@ -217,7 +217,7 @@ pub struct NutsStats<P: HasDims, H: Storable<P>, A: Storable<P>, D: Storable<P>>
     pub draw: u64,
     pub energy_error: f64,
     #[storable(dims("unconstrained_parameter"))]
-    pub unconstrained: Option<Vec<f64>>,
+    pub unconstrained_draw: Option<Vec<f64>>,
     #[storable(dims("unconstrained_parameter"))]
     pub gradient: Option<Vec<f64>>,
     #[storable(flatten)]
@@ -289,7 +289,7 @@ impl<M: Math, R: rand::Rng, A: AdaptStrategy<M>> SamplerStats<M> for NutsChain<M
             chain: self.chain,
             draw: self.draw_count,
             energy_error: point.energy_error(),
-            unconstrained: Some(math.box_array(point.position()).into_vec()),
+            unconstrained_draw: Some(math.box_array(point.position()).into_vec()),
             gradient: Some(math.box_array(point.gradient()).into_vec()),
             hamiltonian: hamiltonian_stats,
             adapt: adapt_stats,

src/mass_matrix/low_rank.rs

@@ -1,4 +1,5 @@
 use std::collections::VecDeque;
+use std::iter::repeat;
 
 use faer::{Col, ColRef, Mat, MatRef, Scale};
 use itertools::Itertools;
@@ -127,15 +128,17 @@ impl Default for LowRankSettings {
         Self {
             store_mass_matrix: false,
             gamma: 1e-5,
-            eigval_cutoff: 100f64,
+            eigval_cutoff: 2f64,
         }
     }
 }
 
 #[derive(Debug, Storable)]
 pub struct MatrixStats {
-    pub eigvals: Option<Vec<f64>>,
-    pub stds: Option<Vec<f64>>,
+    #[storable(dims("unconstrained_parameter"))]
+    pub mass_matrix_eigvals: Option<Vec<f64>>,
+    #[storable(dims("unconstrained_parameter"))]
+    pub mass_matrix_stds: Option<Vec<f64>>,
     pub num_eigenvalues: u64,
 }
 
@@ -145,14 +148,18 @@ impl<M: Math> SamplerStats<M> for LowRankMassMatrix<M> {
 
     fn extract_stats(&self, math: &mut M, _opt: Self::StatsOptions) -> Self::Stats {
         if self.settings.store_mass_matrix {
+            let stds = Some(math.box_array(&self.stds));
             let eigvals = self
                 .inner
                 .as_ref()
                 .map(|inner| math.eigs_as_array(&inner.vals));
-            let stds = Some(math.box_array(&self.stds));
+            let mut eigvals = eigvals.map(|x| x.into_vec());
+            if let Some(ref mut eigvals) = eigvals {
+                eigvals.extend(repeat(f64::NAN).take(stds.as_ref().unwrap().len() - eigvals.len()));
+            }
             MatrixStats {
-                eigvals: eigvals.map(|x| x.into_vec()),
-                stds: stds.map(|x| x.into_vec()),
+                mass_matrix_eigvals: eigvals,
+                mass_matrix_stds: stds.map(|x| x.into_vec()),
                 num_eigenvalues: self
                     .inner
                     .as_ref()
@@ -161,9 +168,13 @@ impl<M: Math> SamplerStats<M> for LowRankMassMatrix<M> {
             }
         } else {
             MatrixStats {
-                eigvals: None,
-                stds: None,
-                num_eigenvalues: 0,
+                mass_matrix_eigvals: None,
+                mass_matrix_stds: None,
+                num_eigenvalues: self
+                    .inner
+                    .as_ref()
+                    .map(|inner| inner.num_eigenvalues)
+                    .unwrap_or(0),
             }
         }
     }

src/transform_adapt_strategy.rs

@@ -43,14 +43,18 @@ pub struct TransformAdaptation {
 }
 
 #[derive(Debug, Storable)]
-pub struct Stats {}
+pub struct Stats {
+    tuning: bool,
+}
 
 impl<M: Math> SamplerStats<M> for TransformAdaptation {
     type Stats = Stats;
     type StatsOptions = ();
 
     fn extract_stats(&self, _math: &mut M, _opt: Self::StatsOptions) -> Self::Stats {
-        Stats {}
+        Stats {
+            tuning: self.tuning,
+        }
     }
 }
 

src/transformed_hamiltonian.rs

@@ -26,7 +26,9 @@ pub struct TransformedPoint<M: Math> {
 #[derive(Debug, Storable)]
 pub struct PointStats {
     pub fisher_distance: f64,
+    #[storable(dims("unconstrained_parameter"))]
     pub transformed_position: Option<Vec<f64>>,
+    #[storable(dims("unconstrained_parameter"))]
     pub transformed_gradient: Option<Vec<f64>>,
     pub transformation_index: i64,
 }