Work on early tuning schedule

Author: aseyboldt

src/adapt_strategy.rs

@@ -181,16 +181,16 @@ impl<F: CpuLogpFunc> AdaptStrategy for ExpWindowDiagAdapt<F> {
         potential: &mut Self::Potential,
         state: &<Self::Potential as Hamiltonian>::State,
     ) {
-        self.exp_variance_draw.set_variance(iter::repeat(1f64));
+        self.exp_variance_draw.set_variance(iter::repeat(0f64));
         self.exp_variance_draw.set_mean(state.q.iter().copied());
         self.exp_variance_grad
             .set_variance(state.grad.iter().map(|&val| {
                 let diag = if !self.settings.grad_init {
                     1f64
                 } else {
                     let out = val * val;
-                    let out = out.clamp(LOWER_LIMIT * LOWER_LIMIT, UPPER_LIMIT * UPPER_LIMIT);
-                    if (out == 0f64) | (!out.is_finite()) {
+                    let out = out.clamp(LOWER_LIMIT, UPPER_LIMIT);
+                    if !out.is_finite() {
                         1f64
                     } else {
                         out
@@ -205,12 +205,13 @@ impl<F: CpuLogpFunc> AdaptStrategy for ExpWindowDiagAdapt<F> {
                 self.exp_variance_draw.current(),
                 self.exp_variance_grad.current(),
             )
-            .map(|(draw, grad)| {
-                let val = (draw / grad).sqrt().clamp(LOWER_LIMIT, UPPER_LIMIT);
+            .map(|(_draw, grad)| {
+                //let val = (1f64 / grad).clamp(LOWER_LIMIT, UPPER_LIMIT);
+                let val = (1f64 / grad).sqrt().clamp(LOWER_LIMIT, UPPER_LIMIT);
                 if val.is_finite() {
-                    val
+                    Some(val)
                 } else {
-                    1f64
+                    Some(1f64)
                 }
             }),
         );
@@ -227,25 +228,44 @@ impl<F: CpuLogpFunc> AdaptStrategy for ExpWindowDiagAdapt<F> {
             return;
         }
 
+        let is_early = (draw as f64) < self.settings.early_ratio * (self.num_tune as f64);
+
+
         let count = self.exp_variance_draw_bg.count();
 
-        let early_switch = (count == self.settings.early_window_switch_freq)
-            & (draw < self.settings.window_switch_freq);
+        let switch_freq = if is_early {
+            self.settings.early_window_switch_freq
+        } else {
+            self.settings.window_switch_freq
+        };
+
+        let variance_decay = if is_early {
+            self.settings.early_variance_decay
+        } else {
+            self.settings.variance_decay
+        };
 
-        if early_switch | ((draw % self.settings.window_switch_freq == 0) & (count > 5)) {
+        let switch = count >= switch_freq;
+
+        if switch {
+            assert!(count == switch_freq);
             self.exp_variance_draw = std::mem::replace(
                 &mut self.exp_variance_draw_bg,
-                ExpWeightedVariance::new(self.dim, self.settings.variance_decay, true),
+                ExpWeightedVariance::new(self.dim, variance_decay, true),
             );
             self.exp_variance_grad = std::mem::replace(
                 &mut self.exp_variance_grad_bg,
-                ExpWeightedVariance::new(self.dim, self.settings.variance_decay, true),
+                ExpWeightedVariance::new(self.dim, variance_decay, true),
             );
 
             self.exp_variance_draw_bg
                 .set_mean(collector.draw.iter().copied());
+            self.exp_variance_draw_bg
+                .set_variance(iter::repeat(0f64));
             self.exp_variance_grad_bg
-                .set_mean(collector.grad.iter().copied());
+                .set_mean(iter::repeat(0f64));
+            self.exp_variance_grad_bg
+                .set_variance(collector.grad.iter().map(|&x| x * x));
         } else if collector.is_good {
             self.exp_variance_draw
                 .add_sample(collector.draw.iter().copied());
@@ -257,23 +277,23 @@ impl<F: CpuLogpFunc> AdaptStrategy for ExpWindowDiagAdapt<F> {
                 .add_sample(collector.grad.iter().copied());
         }
 
+        //if (is_early & (self.exp_variance_draw.count() > 2)) | (!is_early & switch) {
         if self.exp_variance_draw.count() > 2 {
             assert!(self.exp_variance_draw.count() == self.exp_variance_grad.count());
-            if (self.settings.grad_init) | (draw > self.settings.window_switch_freq) {
-                potential.mass_matrix.update_diag(
-                    izip!(
-                        self.exp_variance_draw.current(),
-                        self.exp_variance_grad.current(),
-                    )
-                    .map(|(draw, grad)| {
-                        let mut val = (draw / grad).sqrt().clamp(LOWER_LIMIT, UPPER_LIMIT);
-                        if !val.is_finite() {
-                            val = 1f64;
-                        }
-                        val
-                    }),
-                );
-            }
+            potential.mass_matrix.update_diag(
+                izip!(
+                    self.exp_variance_draw.current(),
+                    self.exp_variance_grad.current(),
+                )
+                .map(|(draw, grad)| {
+                    let val = (draw / grad).sqrt().clamp(LOWER_LIMIT, UPPER_LIMIT);
+                    if !val.is_finite() {
+                        None
+                    } else {
+                        Some(val)
+                    }
+                }),
+            );
         }
     }
 

src/mass_matrix.rs

@@ -33,7 +33,7 @@ impl DiagMassMatrix {
         }
     }
 
-    pub(crate) fn update_diag(&mut self, new_variance: impl Iterator<Item = f64>) {
+    pub(crate) fn update_diag(&mut self, new_variance: impl Iterator<Item = Option<f64>>) {
         update_diag(&mut self.variance, &mut self.inv_stds, new_variance);
     }
 }
@@ -44,13 +44,16 @@ impl DiagMassMatrix {
 fn update_diag(
     variance_out: &mut [f64],
     inv_std_out: &mut [f64],
-    new_variance: impl Iterator<Item = f64>,
+    new_variance: impl Iterator<Item = Option<f64>>,
 ) {
-    izip!(variance_out, inv_std_out, new_variance,).for_each(|(var, inv_std, x)| {
-        assert!(x.is_finite(), "Illegal value on mass matrix: {}", x);
-        assert!(x > 0f64, "Illegal value on mass matrix: {}", x);
-        *var = x;
-        *inv_std = (1. / x).sqrt();
+    izip!(variance_out, inv_std_out, new_variance).for_each(|(var, inv_std, x)| {
+        if let Some(x) = x {
+            assert!(x.is_finite(), "Illegal value on mass matrix: {}", x);
+            assert!(x > 0f64, "Illegal value on mass matrix: {}", x);
+            //assert!(*var != x, "No change in mass matrix from {} to {}", *var, x);
+            *var = x;
+            *inv_std = (1. / x).sqrt();
+        };
     });
 }
 
@@ -135,7 +138,6 @@ fn add_sample(self_: &mut ExpWeightedVariance, value: impl Iterator<Item = f64>)
                 //    assert!(x - *mean != 0f64, "var = {}, mean = {}, x = {}, delta = {}, count = {}", var, mean, x, x - *mean, self_.count);
                 //}
                 let delta = x - *mean;
-                //*mean += self_.alpha * delta;
                 *mean = self_.alpha.mul_add(delta, *mean);
                 *var = (1f64 - self_.alpha) * (*var + self_.alpha * delta * delta);
             },
@@ -165,6 +167,8 @@ pub struct DiagAdaptExpSettings {
     /// Switch to a new variance estimator every `window_switch_freq` draws.
     pub window_switch_freq: u64,
     pub early_window_switch_freq: u64,
+    /// The ratio of the adaptation steps that is considered "early"
+    pub early_ratio: f64,
     pub grad_init: bool,
 }
 
@@ -174,9 +178,10 @@ impl Default for DiagAdaptExpSettings {
             variance_decay: 0.02,
             final_window: 50,
             store_mass_matrix: false,
-            window_switch_freq: 50,
-            early_window_switch_freq: 10,
+            window_switch_freq: 200,
+            early_window_switch_freq: 20,
             early_variance_decay: 0.1,
+            early_ratio: 0.4,
             grad_init: true,
         }
     }
@@ -201,14 +206,10 @@ impl DrawGradCollector {
 impl Collector for DrawGradCollector {
     type State = State;
 
-    fn register_draw(&mut self, state: &Self::State, info: &crate::nuts::SampleInfo) {
+    fn register_draw(&mut self, state: &Self::State, _info: &crate::nuts::SampleInfo) {
         self.draw.copy_from_slice(&state.q);
         self.grad.copy_from_slice(&state.grad);
         let idx = state.index_in_trajectory();
-        if let Some(_) = info.divergence_info {
-            self.is_good = (idx <= -4) | (idx >= 4);
-        } else {
-            self.is_good = idx != 0;
-        }
+        self.is_good = idx != 0;
     }
 }