add low rank modified mass matrix adaptation

.gitignore

@@ -11,3 +11,4 @@ notebooks/*.hpp
 perf.data*
 wheels
 .vscode/
+*~

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 All notable changes to this project will be documented in this file.
 
+## [0.13.0] - 2024-07-05
+
+### Features
+
+- Add low rank modified mass matrix adaptation (Adrian Seyboldt)
+
+
 ## [0.12.0] - 2024-06-29
 
 ### Features
@@ -45,11 +52,6 @@ All notable changes to this project will be documented in this file.
 - Add progress bar on terminal (Adrian Seyboldt)
 
 
-### Miscellaneous Tasks
-
-- Prepare release (Adrian Seyboldt)
-
-
 ## [0.11.0] - 2024-05-29
 
 ### Bug Fixes
@@ -78,8 +80,6 @@ All notable changes to this project will be documented in this file.
 
 - Update python dependencies (Adrian Seyboldt)
 
-- Prepare release (Adrian Seyboldt)
-
 
 ### Refactor
 

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "nutpie"
-version = "0.12.0"
+version = "0.13.0"
 authors = [
   "Adrian Seyboldt <[email protected]>",
   "PyMC Developers <[email protected]>",
@@ -22,7 +22,7 @@ name = "_lib"
 crate-type = ["cdylib"]
 
 [dependencies]
-nuts-rs = "0.11.0"
+nuts-rs = "0.12.0"
 numpy = "0.21.0"
 ndarray = "0.15.6"
 rand = "0.8.5"
@@ -33,7 +33,7 @@ rayon = "1.9.0"
 arrow = { version = "52.0.0", default-features = false, features = ["ffi"] }
 anyhow = "1.0.72"
 itertools = "0.13.0"
-bridgestan = "2.4.1"
+bridgestan = "2.5.0"
 rand_distr = "0.4.3"
 smallvec = "1.11.0"
 upon = { version = "0.8.1", default-features = false, features = [] }

cliff.toml

@@ -54,7 +54,7 @@ commit_parsers = [
   { message = "^refactor", group = "Refactor" },
   { message = "^style", group = "Styling" },
   { message = "^test", group = "Testing" },
-  { message = "^chore\\(release\\): prepare for", skip = true },
+  { message = "^chore: Prepare", skip = true },
   { message = "^chore", group = "Miscellaneous Tasks" },
   { body = ".*security", group = "Security" },
 ]

python/nutpie/sample.py

@@ -460,6 +460,7 @@ def sample(
     seed: Optional[int],
     save_warmup: bool,
     progress_bar: bool,
+    low_rank_modified_mass_matrix: bool = False,
     init_mean: Optional[np.ndarray],
     return_raw_trace: bool,
     blocking: Literal[True],
@@ -478,6 +479,7 @@ def sample(
     seed: Optional[int],
     save_warmup: bool,
     progress_bar: bool,
+    low_rank_modified_mass_matrix: bool = False,
     init_mean: Optional[np.ndarray],
     return_raw_trace: bool,
     blocking: Literal[False],
@@ -495,6 +497,7 @@ def sample(
     seed: Optional[int] = None,
     save_warmup: bool = True,
     progress_bar: bool = True,
+    low_rank_modified_mass_matrix: bool = False,
     init_mean: Optional[np.ndarray] = None,
     return_raw_trace: bool = False,
     blocking: bool = True,
@@ -569,6 +572,19 @@ def sample(
         for the progress bar (eg CSS).
     progress_rate: int, default=500
         Rate in ms at which the progress should be updated.
+    low_rank_modified_mass_matrix: bool, default=False
+        Allow adaptation to some posterior correlations using
+        a low-rank updated mass matrix. This is *experimental*
+        and details about this will probably change in the next
+        release.
+    mass_matrix_eigval_cutoff: float > 1, defaul=100
+        Ignore eigenvalues between cutoff and 1/cutoff in the
+        low-rank modified mass matrix estimate. Higher values
+        lead to worse correclation fitting, but increase
+        the performance of leapfrog steps.
+    mass_matrix_gamma: float > 0, default=1e-5
+        Regularisation parameter for the eigenvalues. Only
+        applicable with low_rank_modified_mass_matrix=True.
     **kwargs
         Pass additional arguments to nutpie._lib.PySamplerArgs
 
@@ -577,7 +593,11 @@ def sample(
     trace : arviz.InferenceData
         An ArviZ ``InferenceData`` object that contains the samples.
     """
-    settings = _lib.PyDiagGradNutsSettings(seed)
+
+    if low_rank_modified_mass_matrix:
+        settings = _lib.PyNutsSettings.LowRank(seed)
+    else:
+        settings = _lib.PyNutsSettings.Diag(seed)
     settings.num_tune = tune
     settings.num_draws = draws
     settings.num_chains = chains

src/pyfunc.rs

@@ -41,7 +41,7 @@ impl PyVariable {
                 let field = Arc::new(Field::new("item", DataType::Boolean, false));
                 DataType::FixedSizeList(field, tensor_type.size() as i32)
             }
-            ExpandDtype::ArrayFloat64 { tensor_type } => {
+            ExpandDtype::ArrayFloat64 { tensor_type: _ } => {
                 let field = Arc::new(Field::new("item", DataType::Float64, true));
                 DataType::List(field)
             }
@@ -303,11 +303,11 @@ impl ExpandDtype {
     #[getter]
     fn shape(&self) -> Option<Vec<usize>> {
         match self {
-            Self::BooleanArray {tensor_type} => { Some(tensor_type.shape.iter().cloned().collect()) },
-            Self::ArrayFloat64 {tensor_type} => { Some(tensor_type.shape.iter().cloned().collect()) },
-            Self::ArrayFloat32 {tensor_type} => { Some(tensor_type.shape.iter().cloned().collect()) },
-            Self::ArrayInt64 {tensor_type} => { Some(tensor_type.shape.iter().cloned().collect()) },
-            _ => { None },
+            Self::BooleanArray { tensor_type } => Some(tensor_type.shape.iter().cloned().collect()),
+            Self::ArrayFloat64 { tensor_type } => Some(tensor_type.shape.iter().cloned().collect()),
+            Self::ArrayFloat32 { tensor_type } => Some(tensor_type.shape.iter().cloned().collect()),
+            Self::ArrayInt64 { tensor_type } => Some(tensor_type.shape.iter().cloned().collect()),
+            _ => None,
         }
     }
 }