Minimizing the Rosenbrock function

[adtzlr]

Let us consider the problem of minimizing the Rosenbrock function (taken from [1]). With tensortrax, the gradient and the hessian are evaluated by forward-mode AD.

import tensortrax as tr
import tensortrax.math as tm

import numpy as np

def rosen(x):
    return tm.sum(100*(x[1:] - x[:-1]**2)**2 + (1 - x[:-1])**2.0)

x0 = np.array([1.3, 0.7, 0.8, 1.9, 1.2])

func = tr.function(rosen)
jac = tr.gradient(rosen)
hess = tr.hessian(rosen)

from scipy.optimize import minimize

res = minimize(func, x0, method="Newton-CG", jac=jac, hess=hess)

res.nit
>>> 21

res.x
>>> array([0.9999852 , 0.9999705 , 0.99994098, 0.99988173, 0.99976293])

References

[1] SciPy Documentation of scipy.optimize.minimize, https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html

[adtzlr]

Starting with tensortrax >= 0.4.0, it is possible to make use of the hessp-keyword (about two times faster runtime).

import tensortrax as tr
import tensortrax.math as tm

import numpy as np

def rosen(x):
    return tm.sum(100*(x[1:] - x[:-1]**2)**2 + (1 - x[:-1])**2)

x0 = np.array([1.3, 0.7, 0.8, 1.9, 1.2])

func = tr.function(rosen)
jac = tr.jacobian(rosen)
hessp = lambda x, v: tr.hessian_vector_product(rosen)(x, δx=v)

from scipy.optimize import minimize

res = minimize(func, x0, method="Newton-CG", jac=jac, hessp=hessp)