Help with computing Hessian with JAX

[georgesle]

Hello,
I'm new to JAX and I'm working on certain optimization problem (fitting) with nonlinear constraints that often requires computing the Hessian matrix of the cost function which in my case is the sum of squared residuals divided by the error (chisq loss). Normally, I use scipy.optimize.minimize whenever I have to deal with explicit nonlinear constraints on model parameters, otherwise, scipy.optimize.curve_fit is just what the doctor ordered for a quick solution if there are no constraints.

For 1D case, I'd define a model function and wrap it with the loss function:

import numpy as np
from scipy.optimize import minimize
import matplotlib.pyplot as plt


# model function definition
def nonlinear(x, t):
    return np.sin(x[0]*t + x[1]) + x[2]

# wrapping model to loss
def chisq(x, func, t, y, sigma):
    term = (np.divide(func(x, t) - y, sigma))**2
    chisq = sum(term)
    return chisq

# nonlinear constraint
def cond1(x):
    return x[0]**0.5/x[1] - 2

cons = ({'type': 'eq', 'fun': cond1},
        )

# original model data
x0 = [3, 0.5, 2]
nt = 2000
t = np.linspace(0, 2*np.pi, nt)
noise = np.random.normal(0, 0.15, nt) + 0.01*t # adding some noise
y = nonlinear(x0, t)
y_obs = y + noise
noise[np.abs(noise) < 1e-4] = 1e-2


res = minimize(chisq, [2.5, 0.7, 2], args=(nonlinear, t, y_obs, noise), method='Nelder-Mead')
y_fit = nonlinear(res.x, t)
print(res.x)

bnds = ((2.5, 3.5), (0.5, 1.5), (1.5, 2.5))
res = minimize(chisq, [2.6, 0.7, 2], args=(nonlinear, t, y_obs, noise), method='SLSQP',
               constraints=cons, tol=1e-7)
y_fit_cons = nonlinear(res.x, t)
print(res.x, cond1(res.x))

plt.figure(figsize=(10, 8))
plt.plot(t, y_obs, label='y_obs')
plt.plot(t, y, label='y_true')
plt.plot(t, y_fit, color='orange', label='fit no constraint')
plt.plot(t, y_fit_cons, label='fit with constraint')
plt.legend()
plt.show()

When the fit converges, the solutions are:
[2.99998486 0.50061985 2.00045474] for unconstrained problem
[2.87661593 0.84802947 2.16570558] with constraints

I'm struggling from the JAX docs how to correctly implement the Hessian of the cost function. The idea is to test a few simple examples before building a much more complex model and many more constraints, but I'm stuck at the simplest case. How do I compute the Hessian matrix for optimal solution?
Thank you.

[mattjj]

Thanks for the question!

Can you just use jax.hessian?

[georgesle]

Yes, and I'm struggling hard how to make something like jax.hessian(chisq)(res.x, nonlinear, t, y_obs, noise) work or Hessian-vector products with grad-of-grad. Hopefully, a full model in the later stage could be developed, as soon as I manage to reproduce scipy.optimize.minimize with nonlinear constraints using JAX only.

Hessian is finally needed for computing the covariances and computing the error of the model parameters. This is implemented in many modules and libraries, however none of these have the option to include nonlinear constraints, so I have to rely on the minimize, and compute the missing parts myself.

[mattjj]

Hrm, res.x isn't a function though; it looks like it's the found solution value. jax.hessian needs to be applied to a function (i.e. a Python callable).

I haven't fully read and understood your code, but maybe you want something more like jax.hessian(lambda x: chisq(x, nonlinear, t, y_obs, noise))? (Similarly for Hessian-vector products, the function of which you want Hessian-vector products is something like lambda x: chisq(x, nonlinear, t, y_obs, noise), which you can use with one of the hvp recipes in the Autodiff Cookbook.)

Also notice that any JAX-transformed function needs to call into JAX libraries, like jax.numpy and jax.scipy, rather than numpy and scipy. Your imports may need to change.

[georgesle]

It works! I've obtained the matrix and now I can invert it. Any news on when the jax.scipy.optimize.minimize will support 'SLSQP' method? Would be good to have JAX supporting minimize method with constraints.

[mattjj]

That's great!

I don't know of anyone working on SLSQP. Though actually I realized the way you're calling scipy.optimize.minimize, with JAX stuff all on the inside (i.e. in its arguments, not in its caller), I think you can use the ordinary scipy.optimize.minimize rather than the jax.scipy.optimize.minimize wrapper, like in #936. So maybe just use ordinary scipy there after all, and then you'll have SLSQP?

[georgesle]

At the moment, I do ordinary scipy solver with JAX only for computing the hessian inverse after the good fit has been found.