style: restructure packages

Author: aseyboldt

examples/zarr_async_trace.rs

@@ -24,7 +24,6 @@ use nuts_rs::{
 use nuts_storable::{HasDims, Value};
 use rand::Rng;
 use thiserror::Error;
-use zarrs::filesystem::FilesystemStore;
 use zarrs_object_store::AsyncObjectStore;
 
 /// A multivariate normal distribution model

src/adapt_strategy.rs

@@ -5,18 +5,18 @@ use nuts_storable::{HasDims, Storable};
 use rand::Rng;
 use serde::Serialize;
 
+use super::mass_matrix::MassMatrixAdaptStrategy;
+use super::stepsize::AcceptanceRateCollector;
+use super::stepsize::{StepSizeSettings, Strategy as StepSizeStrategy};
 use crate::{
     NutsError,
     chain::AdaptStrategy,
     euclidean_hamiltonian::EuclideanHamiltonian,
     hamiltonian::{DivergenceInfo, Hamiltonian, Point},
-    mass_matrix_adapt::MassMatrixAdaptStrategy,
     math_base::Math,
     nuts::{Collector, NutsOptions},
     sampler_stats::{SamplerStats, StatsDims},
     state::State,
-    stepsize_adapt::{StepSizeSettings, Strategy as StepSizeStrategy},
-    stepsize_dual_avg::AcceptanceRateCollector,
 };
 
 pub struct GlobalStrategy<M: Math, A: MassMatrixAdaptStrategy<M>> {
@@ -460,7 +460,7 @@ mod test {
 
     #[test]
     fn instanciate_adaptive_sampler() {
-        use crate::mass_matrix_adapt::Strategy;
+        use crate::mass_matrix::Strategy;
 
         let ndim = 10;
         let func = NormalLogp::new(ndim, 3.);

src/lib.rs

@@ -101,30 +101,20 @@
 mod adapt_strategy;
 mod chain;
 mod cpu_math;
-mod csv_storage;
 mod euclidean_hamiltonian;
 mod hamiltonian;
-mod hashmap_storage;
-mod low_rank_mass_matrix;
 mod mass_matrix;
-mod mass_matrix_adapt;
 mod math;
 mod math_base;
 mod model;
-#[cfg(feature = "ndarray")]
-mod ndarray_storage;
 mod nuts;
 mod sampler;
 mod sampler_stats;
 mod state;
-mod stepsize_adam;
-mod stepsize_adapt;
-mod stepsize_dual_avg;
+mod stepsize;
 mod storage;
 mod transform_adapt_strategy;
 mod transformed_hamiltonian;
-#[cfg(feature = "zarr")]
-mod zarr_storage;
 
 pub use nuts_derive::Storable;
 pub use nuts_storable::{HasDims, ItemType, Storable, Value};
@@ -143,16 +133,15 @@ pub use sampler::{
 };
 pub use sampler_stats::SamplerStats;
 
-pub use low_rank_mass_matrix::LowRankSettings;
-pub use mass_matrix_adapt::DiagAdaptExpSettings;
-pub use stepsize_adam::AdamOptions;
-pub use stepsize_adapt::{StepSizeAdaptMethod, StepSizeAdaptOptions, StepSizeSettings};
+pub use mass_matrix::DiagAdaptExpSettings;
+pub use mass_matrix::LowRankSettings;
+pub use stepsize::{AdamOptions, StepSizeAdaptMethod, StepSizeAdaptOptions, StepSizeSettings};
 pub use transform_adapt_strategy::TransformedSettings;
 
 #[cfg(feature = "zarr")]
-pub use zarr_storage::{ZarrAsyncConfig, ZarrAsyncTraceStorage, ZarrConfig, ZarrTraceStorage};
+pub use storage::{ZarrAsyncConfig, ZarrAsyncTraceStorage, ZarrConfig, ZarrTraceStorage};
 
-pub use csv_storage::{CsvConfig, CsvTraceStorage};
-pub use hashmap_storage::{HashMapConfig, HashMapValue};
+pub use storage::{CsvConfig, CsvTraceStorage};
+pub use storage::{HashMapConfig, HashMapValue};
 #[cfg(feature = "ndarray")]
-pub use ndarray_storage::{NdarrayConfig, NdarrayTrace, NdarrayValue};
+pub use storage::{NdarrayConfig, NdarrayTrace, NdarrayValue};

src/mass_matrix_adapt.rs src/mass_matrix/adapt.rssrc/mass_matrix_adapt.rs renamed to src/mass_matrix/adapt.rs

@@ -4,11 +4,11 @@ use nuts_derive::Storable;
 use rand::Rng;
 use serde::Serialize;
 
+use super::mass_matrix::{DiagMassMatrix, DrawGradCollector, MassMatrix, RunningVariance};
 use crate::{
     Math, NutsError,
     euclidean_hamiltonian::EuclideanPoint,
     hamiltonian::Point,
-    mass_matrix::{DiagMassMatrix, DrawGradCollector, MassMatrix, RunningVariance},
     nuts::{Collector, NutsOptions},
     sampler_stats::SamplerStats,
 };

src/low_rank_mass_matrix.rs src/mass_matrix/low_rank.rssrc/low_rank_mass_matrix.rs renamed to src/mass_matrix/low_rank.rs

@@ -5,12 +5,10 @@ use itertools::Itertools;
 use nuts_derive::Storable;
 use serde::Serialize;
 
+use super::adapt::MassMatrixAdaptStrategy;
+use super::mass_matrix::{DrawGradCollector, MassMatrix};
 use crate::{
-    Math, NutsError,
-    euclidean_hamiltonian::EuclideanPoint,
-    hamiltonian::Point,
-    mass_matrix::{DrawGradCollector, MassMatrix},
-    mass_matrix_adapt::MassMatrixAdaptStrategy,
+    Math, NutsError, euclidean_hamiltonian::EuclideanPoint, hamiltonian::Point,
     sampler_stats::SamplerStats,
 };
 
@@ -330,7 +328,8 @@ impl LowRankMassMatrixStrategy {
 
 fn rescale_points(draws: &mut Mat<f64>, grads: &mut Mat<f64>) -> Col<f64> {
     let (ndim, ndraws) = draws.shape();
-    let stds = Col::from_fn(ndim, |col| {
+    
+    Col::from_fn(ndim, |col| {
         let draw_mean = draws.row(col).sum() / (ndraws as f64);
         let grad_mean = grads.row(col).sum() / (ndraws as f64);
         let draw_std: f64 = draws
@@ -361,8 +360,7 @@ fn rescale_points(draws: &mut Mat<f64>, grads: &mut Mat<f64>) -> Col<f64> {
             .for_each(|val| *val = (*val - grad_mean) * grad_scale);
 
         std
-    });
-    stds
+    })
 }
 
 fn estimate_mass_matrix(
@@ -507,9 +505,7 @@ mod test {
     use rand::{Rng, SeedableRng, rngs::SmallRng};
     use rand_distr::StandardNormal;
 
-    use crate::low_rank_mass_matrix::mat_all_finite;
-
-    use super::{estimate_mass_matrix, spd_mean};
+    use super::{estimate_mass_matrix, mat_all_finite, spd_mean};
 
     #[test]
     fn test_spd_mean() {

src/mass_matrix.rs src/mass_matrix/mass_matrix.rssrc/mass_matrix.rs renamed to src/mass_matrix/mass_matrix.rs

@@ -0,0 +1,10 @@
+mod adapt;
+mod low_rank;
+mod mass_matrix;
+
+pub use adapt::DiagAdaptExpSettings;
+pub(crate) use adapt::MassMatrixAdaptStrategy;
+pub(crate) use adapt::Strategy;
+pub use low_rank::LowRankSettings;
+pub(crate) use low_rank::{LowRankMassMatrix, LowRankMassMatrixStrategy};
+pub(crate) use mass_matrix::{DiagMassMatrix, MassMatrix};

src/mass_matrix/mod.rs

@@ -24,9 +24,9 @@ use crate::{
     adapt_strategy::{EuclideanAdaptOptions, GlobalStrategy, GlobalStrategyStatsOptions},
     chain::{AdaptStrategy, Chain, NutsChain, StatOptions},
     euclidean_hamiltonian::EuclideanHamiltonian,
-    low_rank_mass_matrix::{LowRankMassMatrix, LowRankMassMatrixStrategy, LowRankSettings},
     mass_matrix::DiagMassMatrix,
-    mass_matrix_adapt::Strategy as DiagMassMatrixStrategy,
+    mass_matrix::Strategy as DiagMassMatrixStrategy,
+    mass_matrix::{LowRankMassMatrix, LowRankMassMatrixStrategy, LowRankSettings},
     math_base::Math,
     model::Model,
     nuts::NutsOptions,

src/sampler.rs

@@ -4,13 +4,13 @@ use rand::Rng;
 use rand_distr::Uniform;
 use serde::Serialize;
 
+use super::adam::{Adam, AdamOptions};
+use super::dual_avg::{AcceptanceRateCollector, DualAverage, DualAverageOptions};
 use crate::{
     Math, NutsError,
     hamiltonian::{Direction, Hamiltonian, LeapfrogResult, Point},
     nuts::{Collector, NutsOptions},
     sampler_stats::SamplerStats,
-    stepsize_adam::{Adam, AdamOptions},
-    stepsize_dual_avg::{AcceptanceRateCollector, DualAverage, DualAverageOptions},
 };
 use std::fmt::Debug;