### Root Finding Methods ###

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import time 
import root_functions

# Example function: Expected number of clusters from Ewens distribution.
def enc(n,alpha):
  alpha = np.asarray([alpha]) if np.isscalar(alpha) else np.asarray(alpha)
  n_alpha = alpha.size
  result =  [0]*n_alpha
  for i in range(n_alpha):
    result[i] = np.sum(alpha[i]/(np.linspace(1,n,n)+alpha[i]-1))
  return np.array(result)

# Evaluate difference between expected number of clusters and target
n = 100
target = 6
def f(alpha):
  return enc(n,alpha) - target

# Method 0: Iteratively guess!
# Note: Not very much fun.
f(5)
f(3)
f(1)
f(2)
f(1.5)
f(1.2)
f(1.25)
f(1.23)
f(1.235)
f(1.236)
f(1.2355)
f(1.2352)
f(1.2351)

# Method 1: Grid method lays down a grid and evaluates the function at all points.
# Notes:
#   a. Accuracy depends on the fineness of grid.
#   b. Requires many evaluations of the function.
#   c. Conceptually scales beyond univariate functions,
#      but the curse of dimensionality is a practical
#      limitation.
x = np.linspace(.001,10,10000)
fx = f(x)
root = x[np.argmin(np.abs(fx))]
root
f(root)
plt.plot(x,fx)
plt.axvline(x=root,color='green')
plt.axhline(y=0,color='red')
plt.show()
plt.close()

def time1():
  tic = time.perf_counter()
  fx = f(x)
  root = x[np.argmin(np.abs(fx))]
  toc = time.perf_counter()
  return toc-tic

time1()


# Method 2: Bisection method sequentially reduces the width of window
# where the root is known to lie.
# Steps:
#   1. Select starting values x_0 < x_1 such that f(x_0)*f(x_1) < 0
#   2. Set x_2 = (x_0 + x_1)/2
#   3. If f(x_2)*f(x_0) < 0 set x_1 = x_2 else set x_0 = x_2
#   4. Repeat 2-3 until |x_1 - x_0| <= epsilon
#   5. Return (x_0 + x_1)/2
# Notes:
#   a. Is not easy to generalize beyond one-dimension.
#   b. Much more accurate and efficient than the grid method.


# You will provide your own implementation in your homework. :)
source("bisection.R")

# bisection function arguments:
#   1. a function taking one argument.
#   2. a value less than the root.
#   3. a value greater than the root.

# Let's make sure its working.
a = bisection(f,.01,10)
a

actual_enc = enc(100,a)
actual_enc-target

def time2():
  tic = time.perf_counter()
  bisection(f,.01,10)
  toc = time.perf_counter()
  return toc-tic

time2()

# Method 3: Newton-Raphson method
# Steps:
#   1. Let x_0 be an initial guess for the root.
#   2. Compute x_{n+1} = x_n - f(x_n) / f'(x_n), where f' is the derivative of f.
#   3. Repeat step 2 for n = 0, 1, ... until |x_n - x_{n-1}| <= epsilon
#   4. Return x_n
# Notes:
#   a. Much more computational efficient than the bisection method.
#   b. Just as accurate as the bisection method.
#   d. Requires the derivative of the function.
#   c. Generalizes well to high dimensions (if you can do the math!).

# Oops, it doesn't work for this example because we don't know the derivative of f.
# But, for the sake of demonstration, consider this example instead.
def g(x):
  return x**3-3*x+1

def gPrime(x):
  return 3*x**2-3

x = np.linspace(-2,2,1000)
plt.plot(x,g(x))
plt.axhline(y=0,color='red')
plt.show()

diff = float('inf')
x = .5
while (np.abs(diff) > .00000001):
  xnew = x-g(x)/gPrime(x)
  diff = xnew -x
  x = xnew

x
g(x)
plt.axvline(x,color='green')
plt.show()
plt.close()

# Method 4: Secant method.  Like the Newton-Raphson method, but numerically
# approximate the derivative.
# Steps:
#   1. Make two initial guesses x_0 and x_1 for the root.
#   2. Compute x_{n+1} = x_n - f(x_n) ( x_n - x_{n-1} ) / ( f(x_n) - f(x_{n-1}) )
#   3. Repeat step 2 for n = 1, 2, ... until |x_n - x_{n-1}| <= epsilon
#   4. Return x_n
# Notes:
#   a. As computational efficient as the Newton-Raphson method.
#   b. Just as accurate as the bisection and Newton-Raphson method.
#   d. Does **not** require the derivative of the function.
#   c. Generalizes well to high dimensions and doesn't require a lot of math.

# You will provide your own implementation in your homework. 

# secant function arguments:
#   1. a function taking one argument.
#   2. a guess as the value of the root.
#   3. another guess for the root.

# Back to the original example, since we don't need to derivative.
a = secant(f,2,3)
a

actual_enc = enc(100,a)
actual_enc-target

def time3():
  tic = time.perf_counter()
  secant(f,2,3)
  toc = time.perf_counter()
  return toc-tic

time3()