From 04a3c4e6dbcf76a26a5283111d84c64143f98fc3 Mon Sep 17 00:00:00 2001
From: aphc14 <177544929+aphc14@users.noreply.github.com>
Date: Thu, 17 Apr 2025 22:59:13 +1000
Subject: [PATCH] Add regularisation term before Cholesky decomposition

---
 pymc_extras/inference/pathfinder/pathfinder.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pymc_extras/inference/pathfinder/pathfinder.py b/pymc_extras/inference/pathfinder/pathfinder.py
index 846c00fa9..e2eed76e9 100644
--- a/pymc_extras/inference/pathfinder/pathfinder.py
+++ b/pymc_extras/inference/pathfinder/pathfinder.py
@@ -489,6 +489,7 @@ def bfgs_sample_dense(
 
     N = x.shape[-1]
     IdN = pt.eye(N)[None, ...]
+    IdN += IdN * REGULARISATION_TERM
 
     # inverse Hessian
     H_inv = (
@@ -568,9 +569,13 @@ def bfgs_sample_sparse(
     # qr_input: (L, N, 2J)
     qr_input = inv_sqrt_alpha_diag @ beta
     (Q, R), _ = pytensor.scan(fn=pt.nlinalg.qr, sequences=[qr_input], allow_gc=False)
+
     IdN = pt.eye(R.shape[1])[None, ...]
+    IdN += IdN * REGULARISATION_TERM
+
     Lchol_input = IdN + R @ gamma @ pt.matrix_transpose(R)
 
+    # TODO: make robust Lchol calcs more robust, ie. try exceptions, increase REGULARISATION_TERM if non-finite exists
     Lchol = pt.linalg.cholesky(Lchol_input, lower=False, check_finite=False, on_error="nan")
 
     logdet = 2.0 * pt.sum(pt.log(pt.abs(pt.diagonal(Lchol, axis1=-2, axis2=-1))), axis=-1)