feat: Add transforming adaptation

Author: aseyboldt

src/adapt_strategy.rs

@@ -90,7 +90,7 @@ impl<M: Math, A: MassMatrixAdaptStrategy<M>> AdaptStrategy<M> for GlobalStrategy
     >;
     type Options = EuclideanAdaptOptions<A::Options>;
 
-    fn new(math: &mut M, options: Self::Options, num_tune: u64) -> Self {
+    fn new(math: &mut M, options: Self::Options, num_tune: u64, chain: u64) -> Self {
         let num_tune_f = num_tune as f64;
         let step_size_window = (options.step_size_window * num_tune_f) as u64;
         let early_end = (options.early_window * num_tune_f) as u64;
@@ -100,7 +100,7 @@ impl<M: Math, A: MassMatrixAdaptStrategy<M>> AdaptStrategy<M> for GlobalStrategy
 
         Self {
             step_size: StepSizeStrategy::new(options.dual_average_options),
-            mass_matrix: A::new(math, options.mass_matrix_options, num_tune),
+            mass_matrix: A::new(math, options.mass_matrix_options, num_tune, chain),
             options,
             num_tune,
             early_end,
@@ -116,13 +116,13 @@ impl<M: Math, A: MassMatrixAdaptStrategy<M>> AdaptStrategy<M> for GlobalStrategy
         math: &mut M,
         options: &mut NutsOptions,
         hamiltonian: &mut Self::Hamiltonian,
-        state: &State<M, <Self::Hamiltonian as Hamiltonian<M>>::Point>,
+        position: &[f64],
         rng: &mut R,
     ) -> Result<(), NutsError> {
         self.mass_matrix
-            .init(math, options, hamiltonian, state, rng)?;
+            .init(math, options, hamiltonian, position, rng)?;
         self.step_size
-            .init(math, options, hamiltonian, state, rng)?;
+            .init(math, options, hamiltonian, position, rng)?;
         Ok(())
     }
 
@@ -186,8 +186,9 @@ impl<M: Math, A: MassMatrixAdaptStrategy<M>> AdaptStrategy<M> for GlobalStrategy
             // First time we change the mass matrix
             if did_change & self.has_initial_mass_matrix {
                 self.has_initial_mass_matrix = false;
+                let position = math.box_array(state.point().position());
                 self.step_size
-                    .init(math, options, hamiltonian, state, rng)?;
+                    .init(math, options, hamiltonian, &position, rng)?;
             } else {
                 self.step_size.update_stepsize(hamiltonian, false)
             }
@@ -403,7 +404,18 @@ pub mod test_logps {
             unimplemented!()
         }
 
-        fn transformed_logp(
+        fn init_from_transformed_position(
+            &mut self,
+            _params: &Self::TransformParams,
+            _untransformed_position: &mut [f64],
+            _untransformed_gradient: &mut [f64],
+            _transformed_position: &[f64],
+            _transformed_gradient: &mut [f64],
+        ) -> Result<(f64, f64), Self::LogpError> {
+            unimplemented!()
+        }
+
+        fn init_from_untransformed_position(
             &mut self,
             _params: &Self::TransformParams,
             _untransformed_position: &[f64],
@@ -424,15 +436,19 @@ pub mod test_logps {
             unimplemented!()
         }
 
-        fn new_transformation(
+        fn new_transformation<R: rand::Rng + ?Sized>(
             &mut self,
+            _rng: &mut R,
             _untransformed_position: &[f64],
             _untransfogmed_gradient: &[f64],
         ) -> Result<Self::TransformParams, Self::LogpError> {
             unimplemented!()
         }
 
-        fn transformation_id(&self, _params: &Self::TransformParams) -> i64 {
+        fn transformation_id(
+            &self,
+            _params: &Self::TransformParams,
+        ) -> Result<i64, Self::LogpError> {
             unimplemented!()
         }
     }
@@ -462,7 +478,8 @@ mod test {
         let max_energy_error = 1000f64;
         let step_size = 0.1f64;
 
-        let hamiltonian = EuclideanHamiltonian::new(mass_matrix, max_energy_error, step_size);
+        let hamiltonian =
+            EuclideanHamiltonian::new(&mut math, mass_matrix, max_energy_error, step_size);
         let options = NutsOptions {
             maxdepth: 10u64,
             store_gradient: true,

src/chain.rs

@@ -6,7 +6,7 @@ use crate::{
     hamiltonian::{Hamiltonian, Point},
     nuts::{draw, Collector, NutsOptions, NutsSampleStats, NutsStatsBuilder},
     sampler_stats::SamplerStats,
-    state::{State, StatePool},
+    state::State,
     Math, NutsError, Settings,
 };
 
@@ -35,7 +35,6 @@ where
     R: rand::Rng,
     A: AdaptStrategy<M>,
 {
-    pool: StatePool<M, <A::Hamiltonian as Hamiltonian<M>>::Point>,
     hamiltonian: A::Hamiltonian,
     collector: A::Collector,
     options: NutsOptions,
@@ -56,18 +55,15 @@ where
 {
     pub fn new(
         mut math: M,
-        hamiltonian: A::Hamiltonian,
+        mut hamiltonian: A::Hamiltonian,
         strategy: A,
         options: NutsOptions,
         rng: R,
         chain: u64,
     ) -> Self {
-        let pool_size: usize = options.maxdepth.checked_mul(2).unwrap().try_into().unwrap();
-        let pool = hamiltonian.new_pool(&mut math, pool_size);
-        let init = pool.new_state(&mut math);
+        let init = hamiltonian.pool().new_state(&mut math);
         let collector = strategy.new_collector(&mut math);
         NutsChain {
-            pool,
             hamiltonian,
             collector,
             options,
@@ -87,14 +83,14 @@ pub trait AdaptStrategy<M: Math>: SamplerStats<M> {
     type Collector: Collector<M, <Self::Hamiltonian as Hamiltonian<M>>::Point>;
     type Options: Copy + Send + Debug + Default;
 
-    fn new(math: &mut M, options: Self::Options, num_tune: u64) -> Self;
+    fn new(math: &mut M, options: Self::Options, num_tune: u64, chain: u64) -> Self;
 
     fn init<R: Rng + ?Sized>(
         &mut self,
         math: &mut M,
         options: &mut NutsOptions,
         hamiltonian: &mut Self::Hamiltonian,
-        state: &State<M, <Self::Hamiltonian as Hamiltonian<M>>::Point>,
+        position: &[f64],
         rng: &mut R,
     ) -> Result<(), NutsError>;
 
@@ -151,24 +147,20 @@ where
     type AdaptStrategy = A;
 
     fn set_position(&mut self, position: &[f64]) -> Result<()> {
-        let state = self
-            .hamiltonian
-            .init_state(&mut self.math, &mut self.pool, position)?;
-        self.init = state;
         self.strategy.init(
             &mut self.math,
             &mut self.options,
             &mut self.hamiltonian,
-            &self.init,
+            position,
             &mut self.rng,
         )?;
+        self.init = self.hamiltonian.init_state(&mut self.math, position)?;
         Ok(())
     }
 
     fn draw(&mut self) -> Result<(Box<[f64]>, Self::Stats)> {
         let (state, info) = draw(
             &mut self.math,
-            &mut self.pool,
             &mut self.init,
             &mut self.rng,
             &mut self.hamiltonian,

src/cpu_math.rs

@@ -360,15 +360,15 @@ impl<F: CpuLogpFunc> Math for CpuMath<F> {
         )
     }
 
-    fn transformed_logp(
+    fn init_from_untransformed_position(
         &mut self,
         params: &Self::TransformParams,
         untransformed_position: &Self::Vector,
         untransformed_gradient: &mut Self::Vector,
         transformed_position: &mut Self::Vector,
         transformed_gradient: &mut Self::Vector,
     ) -> Result<(f64, f64), Self::LogpErr> {
-        self.logp_func.transformed_logp(
+        self.logp_func.init_from_untransformed_position(
             params,
             untransformed_position.as_slice(),
             untransformed_gradient.as_slice_mut(),
@@ -377,11 +377,28 @@ impl<F: CpuLogpFunc> Math for CpuMath<F> {
         )
     }
 
+    fn init_from_transformed_position(
+        &mut self,
+        params: &Self::TransformParams,
+        untransformed_position: &mut Self::Vector,
+        untransformed_gradient: &mut Self::Vector,
+        transformed_position: &Self::Vector,
+        transformed_gradient: &mut Self::Vector,
+    ) -> Result<(f64, f64), Self::LogpErr> {
+        self.logp_func.init_from_transformed_position(
+            params,
+            untransformed_position.as_slice_mut(),
+            untransformed_gradient.as_slice_mut(),
+            transformed_position.as_slice(),
+            transformed_gradient.as_slice_mut(),
+        )
+    }
+
     fn update_transformation<'a, R: rand::Rng + ?Sized>(
         &'a mut self,
         rng: &mut R,
-        untransformed_positions: impl Iterator<Item = &'a Self::Vector>,
-        untransformed_gradients: impl Iterator<Item = &'a Self::Vector>,
+        untransformed_positions: impl ExactSizeIterator<Item = &'a Self::Vector>,
+        untransformed_gradients: impl ExactSizeIterator<Item = &'a Self::Vector>,
         params: &'a mut Self::TransformParams,
     ) -> Result<(), Self::LogpErr> {
         self.logp_func.update_transformation(
@@ -392,18 +409,22 @@ impl<F: CpuLogpFunc> Math for CpuMath<F> {
         )
     }
 
-    fn new_transformation(
+    fn new_transformation<R: rand::Rng + ?Sized>(
         &mut self,
+        rng: &mut R,
         untransformed_position: &Self::Vector,
         untransfogmed_gradient: &Self::Vector,
+        chain: u64,
     ) -> Result<Self::TransformParams, Self::LogpErr> {
         self.logp_func.new_transformation(
+            rng,
             untransformed_position.as_slice(),
             untransfogmed_gradient.as_slice(),
+            chain,
         )
     }
 
-    fn transformation_id(&self, params: &Self::TransformParams) -> i64 {
+    fn transformation_id(&self, params: &Self::TransformParams) -> Result<i64, Self::LogpErr> {
         self.logp_func.transformation_id(params)
     }
 }
@@ -417,35 +438,58 @@ pub trait CpuLogpFunc {
 
     fn inv_transform_normalize(
         &mut self,
-        params: &Self::TransformParams,
-        untransformed_position: &[f64],
-        untransofrmed_gradient: &[f64],
-        transformed_position: &mut [f64],
-        transformed_gradient: &mut [f64],
-    ) -> Result<f64, Self::LogpError>;
+        _params: &Self::TransformParams,
+        _untransformed_position: &[f64],
+        _untransformed_gradient: &[f64],
+        _transformed_position: &mut [f64],
+        _transformed_gradient: &mut [f64],
+    ) -> Result<f64, Self::LogpError> {
+        unimplemented!()
+    }
 
-    fn transformed_logp(
+    fn init_from_untransformed_position(
         &mut self,
-        params: &Self::TransformParams,
-        untransformed_position: &[f64],
-        untransformed_gradient: &mut [f64],
-        transformed_position: &mut [f64],
-        transformed_gradient: &mut [f64],
-    ) -> Result<(f64, f64), Self::LogpError>;
+        _params: &Self::TransformParams,
+        _untransformed_position: &[f64],
+        _untransformed_gradient: &mut [f64],
+        _transformed_position: &mut [f64],
+        _transformed_gradient: &mut [f64],
+    ) -> Result<(f64, f64), Self::LogpError> {
+        unimplemented!()
+    }
+
+    fn init_from_transformed_position(
+        &mut self,
+        _params: &Self::TransformParams,
+        _untransformed_position: &mut [f64],
+        _untransformed_gradient: &mut [f64],
+        _transformed_position: &[f64],
+        _transformed_gradient: &mut [f64],
+    ) -> Result<(f64, f64), Self::LogpError> {
+        unimplemented!()
+    }
 
     fn update_transformation<'a, R: rand::Rng + ?Sized>(
         &'a mut self,
-        rng: &mut R,
-        untransformed_positions: impl Iterator<Item = &'a [f64]>,
-        untransformed_gradients: impl Iterator<Item = &'a [f64]>,
-        params: &'a mut Self::TransformParams,
-    ) -> Result<(), Self::LogpError>;
+        _rng: &mut R,
+        _untransformed_positions: impl ExactSizeIterator<Item = &'a [f64]>,
+        _untransformed_gradients: impl ExactSizeIterator<Item = &'a [f64]>,
+        _params: &'a mut Self::TransformParams,
+    ) -> Result<(), Self::LogpError> {
+        unimplemented!()
+    }
 
-    fn new_transformation(
+    fn new_transformation<R: rand::Rng + ?Sized>(
         &mut self,
-        untransformed_position: &[f64],
-        untransfogmed_gradient: &[f64],
-    ) -> Result<Self::TransformParams, Self::LogpError>;
+        _rng: &mut R,
+        _untransformed_position: &[f64],
+        _untransformed_gradient: &[f64],
+        _chain: u64,
+    ) -> Result<Self::TransformParams, Self::LogpError> {
+        unimplemented!()
+    }
 
-    fn transformation_id(&self, params: &Self::TransformParams) -> i64;
+    fn transformation_id(&self, _params: &Self::TransformParams) -> Result<i64, Self::LogpError> {
+        unimplemented!()
+    }
 }