# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Tutorial 2

# # Imported modules

# %matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from scipy import optimize
from scipy import linalg


# # Format output

def output_banner():
    print(' Iter   Nfev     Step       Objective    Norm of g')



def output_iteration_info(k, nf, t, f, g):
    print('{0:5d} {1:6d} {2:10e} {3:10e} {4:10e}'.format(k, nf, t, f, linalg.norm(g, np.inf)))


def output_final_results(x, f, g, nf, ng, nh, k):   
    print('\n')
    print('          x:', x)
    print('        fun:', f)
    print('        jac:', g)
    print('norm of jac:', linalg.norm(g, np.inf))
    print('       nfev:', nf)
    print('       ngev:', ng)
    print('       nhev:', nh)
    print('        nit:', k)



# ## Example 1: The Rosenbrock function

# ### Objective

def objective(x):
    """
    Two-variable Rosenbrock function
    """
    return 100*(x[1]-x[0]**2)**2 + (1-x[0])**2


x = np.linspace(-5, 5, 1000)
y = np.linspace(-5, 5, 1000)
X, Y = np.meshgrid(x, y)
Z = objective(np.vstack([X.ravel(), Y.ravel()])).reshape((1000,1000))
plt.contour(X, Y, Z, np.arange(10)**5, cmap='RdGy')
plt.colorbar();
plt.text(1, 1, 'x', va='center', ha='center', color='red', fontsize=20);


# ### Gradient

def gradient(x):
    """
    Derivative of two-variable Rosenbrock function
    """
    return np.array([    
        400 * (x[0]**2 - x[1]) * x[0] + 2*(x[0]-1),
        200 * (x[1] - x[0]**2)
    ])


# ### Hessian

def hessian(x):
    """
    Hessian of two-variable Rosenbrock function
    """
    return np.array([
        [2 - 400 * (x[1] - 3 * x[0]**2), -400 * x[0]],
        [                   -400 * x[0],         200]
    ])


# +
## Suggested initial iterates
# -

x0 = np.array([-1.2, -1.2])
# x0 = np.array([-1, 0.8])
# x0 = np.array([-1.2, 1])
# x0 = np.array([0.4, 0.2]) # start in convex region near the solution (from <NAME>, Oxford University)
# x0 = np.array([-0.9, 1])  # start in nonconvex region (from <NAME>, Oxford University)
# x0 = np.array([-50, 40])  # start very far away (<NAME>, Oxford University)

# ## The steepest descent method (See Tutorial 1)

def steepest_descent_AllInOne(objective, gradient, x0):
    """Implementation of the steepest descent method with a backtracking-Armino linesearch.
    Adapted from steepdes.m, a matlab script which has been around on the internet for a while.
    Written by <NAME> (UCSD) and <NAME> (Stanford) for pegadogic use.
    """

    kmax = 100000
    jmax = 20

    dxmax = 1

    c1 = .0001
    c1 = 1/4
    beta = .5
    x = x0.astype(float)
    f = objective(x); nf = 1
    g = gradient(x); ng = 1
    
    k = 1
    
    output_banner()
    
    while ((linalg.norm(g, np.inf) > 1e-6) and (k <= kmax)):
        d = -g
        t = min(1, dxmax/linalg.norm(g, np.inf))
        xnew = x + t * d
        fnew = objective(xnew); nf += 1
        j = 1
        while ((fnew > f + t * c1 * np.inner(g,d)) and (j <= jmax)):
            t = t * beta
            xnew = x + t * d
            fnew = objective(xnew); nf += 1
            j  += 1
        if j > jmax:
            print('Armijo failed to make progress')
            break

        if (k%100 == 1): output_iteration_info(k, nf, t, f, g)

        x = xnew
        f = fnew
        g = gradient(x); ng += 1
        k += 1

    if k > kmax:
        print('Steepest descent failed to converge after maxiter iterations')

    output_final_results(x, f, g, nf, ng, 0, k);
    return x, f, g



steepest_descent_AllInOne(objective,gradient,x0);


# # Algorithms

# ## Linesearches

# ### Armijo

def armijo(obj, grad, x0, f0, g0, t0, d, nf, ng):
    
    """ Backtracking-Armijo linesearch """

    c1 = 1e-4 
    
    iterMax = 20
    
    gtd0 = np.inner(g0,d)
    
    if (gtd0 >= 0):
        raise SystemExit("Armijo: Direction provided is not a descent direction.")
    
    t = t0

    for k in range(iterMax):
        x = x0 + t*d
        f = obj(x)
        if (f < f0 + c1*t*gtd0):
            g = grad(x)
            return x, f, g, t, nf + k + 1, ng + 1
        else:
            t = t/2
    
    raise SystemExit("Armijo: Maximum Iterations exceeded.")


# ### Wolfe

def wolfe(obj, grad, x0, f0, g0, t0, d, nf, ng):
    
    """
    """
    
    c1 = 1e-4
    c1 = 1/4
    c2 = 0.90
    
    iterMax = 100
    
    a  = 0
    b  = np.inf
    gtd0 = np.inner(g0,d)

    if (gtd0 >= 0):
        raise SystemExit("Wolfe: Direction not a descent direction.")

    t  = t0;

    for k in range(iterMax):
        x = x0 + t*d
        f = obj(x); nf += 1
        if (f > f0 + c1*t*gtd0):
            b = t
            t = (a+b)/2
        else:
            g = grad(x); ng += 1
            if (np.inner(g,d) < c2*gtd0):
                a = t
                if (b == np.inf):
                    t = 2*t
                else:
                    t = (a+b)/2
            else:
                return x, f, g, t, nf, ng
    

    raise SystemExit("WOLFE: Maximum Iterations exceeded.")


# ## Update Quasi-Newton Matrices 

def HUpdate_H(H, s, y):
    
    """ Updates matrix H 
    """



# +
# Descent directions
# -

# ## Steepest descent

def steepest_descent(objective, gradient, linesearch, x0):
    
    """Steepest gradient descent."""

    maxiter = 20000

    x = x0.astype(float)
    f = objective(x); nf = 1
    g = gradient(x); ng = 1
    
    dxmax = 1
    
    k = 0
    
    output_banner()
    
    while ((linalg.norm(g, np.inf) > 1e-6) and (k < maxiter)):
        t = min(1, dxmax/linalg.norm(g, np.inf))
        d = -g
        x, f, g, t, nf, ng = linesearch(objective, gradient, x, f, g, 1, d, nf, ng)
        k += 1
        if (k%100 == 1): output_iteration_info(k, nf, t, f, g)

    output_final_results(x, f, g, nf, ng, 0, k);
    return x, f, g, nf, ng, k;


steepest_descent(objective,gradient, wolfe, x0);


# ## Newton's method

def newton(objective, gradient, hessian, x0):
    """Implementation of the Newton method."""

    maxiter = 500
    sigma = 1e-4
    beta = .5

    x = x0.astype(float)
    f = objective(x); nf = 1
    g = gradient(x); ng = 1
    h = hessian(x); nh = 1
    
    k = 0
    
    output_banner()
    
    while ((linalg.norm(g, np.inf) > 1e-10) and (k <= maxiter)):
        d = - linalg.solve(h,g)
        t = 1
        xnew = x + t * d
        fnew = objective(xnew); nf += 1
        j = 1
        while ((fnew > f + t * sigma * np.inner(g,d)) and (j <= 15)):
            t = t * beta
            xnew = x + t * d
            fnew = objective(xnew); nf += 1
            j += 1
        if j > 15:
            print('Armijo failed to make progress')
            break
        x = xnew
        f = fnew
        g = gradient(x); ng += 1 
        h = hessian(x); nh +=1
        k += 1
        output_iteration_info(k, nf, t, f, g)


    if k > maxiter:
        print('Newton method failed to converge after maxiter iterations')
    output_final_results(x, f, g, nf, ng, nh, k)
    return x, f, g, nf, ng, nh, k


newton(objective, gradient, hessian, x0);

optimize.minimize(objective, x0, method="Newton-CG", jac=gradient, hess=hessian)    


# ## Modified Newton's method

def modified_newton(objective, gradient, hessian, x0):
    """Implementation of the Modified-Newton method with positive-definitess check.
    """

    maxiter = 3000
    sigma = 1e-4
    armax = 50
    
    beta = 1e-3

    x = x0.astype(float)
    n = len(x)
    f = objective(x); nf = 1
    g = gradient(x); ng = 1
    h = hessian(x); nh = 1
    
    k = 1
    
    output_banner()
    
    while ((np.linalg.norm(g, np.inf) > 1e-10) and (k <= maxiter)):
        
        hmin = min(np.diagonal(h)) 

        if hmin > 0:
            tau = 0
        else:
            tau = - hmin + beta

        count = 1
        while count <= 10:
            try:
                c, low = linalg.cho_factor(h + tau * np.eye(n))
                break
            except:
                tau = max(2 * tau, beta)
            count +=1
            
        if count > 10:
            print('Failed to obtain a positive definite modified Hessian.')
            break
                
        d = - linalg.cho_solve((c, low), g)

        t = 1
        xnew = x + t * d
        fnew = objective(xnew); nf += 1
        j = 1
        while ((fnew > f + t * sigma * np.inner(g,d)) and (j <= armax)):
            t = t/2
            xnew = x + t * d
            fnew = objective(xnew); nf += 1
            j += 1
        if j > armax:
            print('Armijo failed to make progress')
            break

        x = xnew
        f = fnew
        g = gradient(x); ng += 1 
        h = hessian(x); nh +=1
        output_iteration_info(k, nf, t, f, g)
        k += 1

    if k > maxiter:
        print('Newton method failed to converge after maxiter iterations')
    output_final_results(x, f, g, nf, ng, nh, k)
    return x, f, g, nf, ng, nh, k



modified_newton(objective,gradient,hessian,x0);

optimize.minimize(objective, x0, method="Newton-CG", jac=gradient, hess=hessian)    


# ## BFGS

def BFGSWolfe(objective, gradient, x0):

    maxIter = 5000;

    eps = 1e-6;
    
    x = x0;   

    f = objective(x); nf = 1
    g = gradient(x); ng = 1

    I = np.eye(len(x))
    H = I

    output_banner()
    
    k = 1

    while ((linalg.norm(g, np.inf) > eps) and (k <= maxIter)):
        d = - np.dot(H, g)
        xnew, fnew, gnew, t, nf, ng = wolfe(objective,gradient,x,f,g,1,d,nf,ng);
        s  = xnew - x
        y  = gnew - g
        r  = 1/np.dot(y,s)
        H  = np.dot((I - r * np.outer(s,y)), np.dot(H, (I - r * np.outer(y,s)))) + r * np.outer(s,s)
        x  = xnew
        f  = fnew
        g  = gnew
        output_iteration_info(k, nf, t, f, g)
        k += 1
    
    if k > maxIter:
        print('BFGS method failed to converge after maxiter iterations')
    output_final_results(x, f, g, nf, ng, 0, k)
    return x, f, g, nf, ng, 0, k


BFGSWolfe(objective, gradient, x0);

optimize.minimize(objective, x0, method="BFGS", jac=gradient) 

# ## Conjugate gradient

optimize.minimize(objective, x0, method="CG", jac=gradient)    


# # Generalized Rosenbrock function

# +
def generalized_rosen(x):
    """The Rosenbrock function as per scipy documentation"""
    return sum(100.0*(x[1:]-x[:-1]**2.0)**2.0 + (1-x[:-1])**2.0)

def generalized_rosen_der(x):
    xm = x[1:-1]
    xm_m1 = x[:-2]
    xm_p1 = x[2:]
    der = np.zeros_like(x)
    der[1:-1] = 200*(xm-xm_m1**2) - 400*(xm_p1 - xm**2)*xm - 2*(1-xm)
    der[0] = -400*x[0]*(x[1]-x[0]**2) - 2*(1-x[0])
    der[-1] = 200*(x[-1]-x[-2]**2)
    return der

def generalized_rosen_hess(x):
    x = np.asarray(x)
    H = np.diag(-400*x[:-1],1) - np.diag(400*x[:-1],-1)
    diagonal = np.zeros_like(x)
    diagonal[0] = 1200*x[0]**2-400*x[1]+2
    diagonal[-1] = 200
    diagonal[1:-1] = 202 + 1200*x[1:-1]**2 - 400*x[2:]
    H = H + np.diag(diagonal)
    return H


# -

x0 = np.array([1.3, 0.7, 0.8, 1.9, 1.2])
x0 = 100*np.random.rand(100)

modified_newton(generalized_rosen, generalized_rosen_der, generalized_rosen_hess, x0);

BFGSWolfe(generalized_rosen, generalized_rosen_der, x0);

res = optimize.minimize(generalized_rosen, x0, method='BFGS', jac=generalized_rosen_der, options={'gtol': 1e-8, 'disp': True});
res

StarcoderdataJupyterScriptsDedupFiltered