#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri May  3 14:59:30 2019

@author: kerkmann
"""

import numpy as np
import matplotlib.pyplot as plt

def Newton(f,df,x0,tol=1e-8,max_it=100):
    '''Newton algorithm to find root of function f'''
    # Input parameters:
    # f        function to find root of
    # df       derivative of f
    # x0       initial guess / starting value
    # tol      tolerance to finish iteration
    # max_it   maximum number of iterations
            
    # Output parameters:
    # x        solution
    # it       number of iterations
    # n        size of next step
    
    it = 0
    d = np.linalg.solve(df(x0),f(x0))
    x = x0
    
    while np.linalg.norm(d,np.inf) > tol and it < max_it:
        x = x - d
        d = np.linalg.solve(df(x),f(x))
        it += 1
    
    # Warning if maximimum number of iterations has been reached
    if it == max_it:
        print('Warning: Maximum number of iterations reached. Newton\'s method might not have converged.')
        
    return x, it, np.linalg.norm(d)

### MAIN PROGRAM ###
    
m = 100
h = 1/m*2*np.pi

x = np.linspace(0,2*np.pi,m+2)
alpha = 0.7

# initial guess
theta_0 = 0.7*np.ones(m)

# discretization
e = np.ones(m)
Gmat = np.diag(e[:-1],-1) - 2*np.diag(e) + np.diag(e[:-1],1)

G = lambda theta: 1/h**2*Gmat@theta + np.eye(m,1).flatten()*alpha/h**2 + np.eye(m,1,-m+1).flatten()*alpha/h**2 + np.sin(theta)

# Jacobian of G
J = lambda theta: 1/h**2*Gmat + np.diag(np.cos(theta))

# solve and print
y,k,n = Newton(G,J,theta_0)

print('Iterationen: {}, Aenderung zum naechsten Schritt: {}'.format(k,n))
plt.plot(x[1:-1],y,label='Anfangsdaten 0.7')

# different initial guess
theta_0 += np.sin(x[1:-1]*0.5)
y,k,n = Newton(G,J,theta_0)

print('Iterationen: {}, Aenderung zum naechsten Schritt: {}'.format(k,n))
plt.plot(x[1:-1],y,label='Anfangsdaten 0.7+sin(0.5t)')
plt.legend()