feat: Add option to use draw base mass matrix estimate

Author: aseyboldt

python/nutpie/sample.py

@@ -428,6 +428,10 @@ def sample(
     return_raw_trace: bool, default=False
         Return the raw trace object (an apache arrow structure)
         instead of converting to arviz.
+    use_grad_based_mass_matrix: bool, default=True
+        Use a mass matrix estimate that is based on draw and gradient
+        variance. Set to `False` to get mass matrix adaptation more
+        similar to PyMC and Stan.
     **kwargs
         Pass additional arguments to nutpie._lib.PySamplerArgs
 

src/wrapper.rs

@@ -222,6 +222,16 @@ impl PyDiagGradNutsSettings {
             .mass_matrix_options
             .store_mass_matrix = val;
     }
+
+    #[getter]
+    fn use_grad_based_mass_matrix(&self) -> bool {
+        self.0.mass_matrix_adapt.mass_matrix_options.use_grad_based_estimate
+    }
+
+    #[setter(use_grad_based_mass_matrix)]
+    fn set_use_grad_based_mass_matrix(&mut self, val: bool) {
+        self.0.mass_matrix_adapt.mass_matrix_options.use_grad_based_estimate = val
+    }
 }
 
 pub(crate) enum SamplerState {