Improve hessian matrix numerical precision implementing function using jax

Author: SimoneGasperini

requirements.txt

@@ -1,5 +1,6 @@
 numpy
 scipy
+jax
 matplotlib
 dwave-ocean-sdk
 scikit-learn

src/qunfold/utils.py

@@ -1,8 +1,8 @@
 import sys
+import jax
 import numpy as np
 import scipy as sp
 from tqdm import tqdm
-from scipy.optimize import approx_fprime
 from qunfold import QUnfolder
 
 try:
@@ -28,13 +28,17 @@ def compute_chi2(observed, expected):
     return chi2_red
 
 
-def approx_hessian(func, *point):
-    precision = np.sqrt(np.finfo(dtype=np.float32).eps)
-    epsilon = precision * np.array([max(1, x) for x in point])
-    function = lambda point: func(*point)
-    gradient = lambda point: approx_fprime(xk=point, f=function, epsilon=epsilon)
-    xk = np.array([x for x in point])
-    return approx_fprime(xk=xk, f=gradient, epsilon=epsilon)
+def approx_hessian(f, x):
+    x = jax.numpy.array(x, dtype=float)
+    n = len(x)
+    hessian = jax.numpy.zeros(shape=(n, n))
+    grad_f = jax.grad(f)
+    for i in range(n):
+        v = jax.numpy.zeros_like(x)
+        v = v.at[i].set(1.0)
+        _, hvp = jax.jvp(grad_f, primals=(x,), tangents=(v,))
+        hessian = hessian.at[:, i].set(hvp)
+    return hessian
 
 
 def lambda_optimizer(response, measured, truth, binning, num_reps=30, verbose=False, seed=None):

tests/test_utils.py

@@ -5,27 +5,26 @@
 from qunfold.utils import approx_hessian
 
 
-deg = st.integers(min_value=0, max_value=5)
+deg = st.integers(min_value=0, max_value=4)
 xi = st.integers(min_value=0, max_value=1e6)
 
 
 @settings(deadline=None)
 @given(degrees=st.tuples(deg, deg, deg), point=st.tuples(xi, xi, xi))
 def test_approx_hessian(degrees, point):
     varlist = [sympy.Symbol(f"x{i}") for i in range(len(degrees))]
-    poly_sympy = 0
+    poly_sympy = 0.0
     for deg, var in zip(degrees, varlist):
-        for exp in range(deg):
-            if np.random.rand() < 0.3:
-                coeff = np.random.rand() * 2 - 1
-                poly_sympy += coeff * var**exp
+        for exp in range(deg + 1):
+            coeff = np.random.rand() * 5 - 2
+            poly_sympy += coeff * var**exp
     n = len(varlist)
     hess_sympy = sympy.Matrix(np.zeros(shape=(n, n)))
     for i in range(n):
         for j in range(n):
             hess_sympy[i, j] = sympy.diff(poly_sympy, varlist[i], varlist[j])
     func = sympy.lambdify(args=varlist, expr=poly_sympy, modules="numpy")
     hess = sympy.lambdify(args=varlist, expr=hess_sympy, modules="numpy")
-    hessian1 = approx_hessian(func, *point)
+    hessian1 = approx_hessian(f=lambda x: func(*x), x=point)
     hessian2 = hess(*point)
-    assert np.allclose(hessian1, hessian2, rtol=0.01, atol=1)
+    assert np.allclose(hessian1, hessian2)