#-----------------------------------------------------------
# Lecture 14 -- NLEs with Scipy
#-----------------------------------------------------------

import numpy as np
import matplotlib.pyplot as plt

# Scipy is a module library that works with Numpy and Matplotlib to enable
# numerical calculations.
# * Numpy -- Arrays and matrices
# * Matplotlib -- Plotting
# * Scipy -- Numerical tools for Scientific/Engineering Calculations
#
# * "root" is the nonlinear solver in Scipy

import scipy.optimize as opt

# %%
#-----------------------------------------------------------
# A. Single Equations (1 equation, 1 unknown)
#-----------------------------------------------------------

# Steps to solving a single nonlinear equation using root:
# (1) import module: import scipy.optimize as opt
# (2) Put the equation in residual form, f(x) = 0
# (3) Define the python function, def <fn name> (<args>): return <f(x) function>
# (4) Set an initial guess for x
# (5) Call `root` function: x = root(<fn name>, <guess>).x
#       ** Note: root returns an object called a "dictionary" that is why you
#          put the .x after the function call. If not, you will get a 
#          whole bunch of extra info. Also, the answer (even if it just a 
#          single equation) will be stored in a numpy array.
#
# Additional information about root at the link:
#   https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.root.html#scipy.optimize.root

# ----------
# Example 
# ----------
# Solve x^2 +x - 5 = 0

def f(x):
    return x**2 + x - 5

x_plt = np.linspace(-4, 3)
plt.rc("font", size=14)
plt.figure()
plt.plot(x_plt, f(x_plt), 'k-')
plt.plot([-4, 3], [0, 0], 'r--')

xguess = 2.
x1 = opt.root(f,xguess).x
print('Example 1')
print("x =", x1)
print("residual = ", f(x1))

plt.plot(x1, f(x1), 'bo')

# ----------
# Practice 
# ----------
# Use a different initial guess to find the other root
# (Answer x2 = -2.79128785)

#%%
#-----------------------------------------------------------
# B. Systems of equations (n equations, n unknowns)
#-----------------------------------------------------------
#
# We can also use root to solve systems of nonlinear equations.
# Instead of f(x) = 0, we have:
#   f0 (x0, x1, ..., xn-1) = 0
#   f1 (x0, x1, ..., xn-1) = 0
#   ...
#   fn-1 (x0, x1, ..., xn-1) = 0
#
# We can write this more compactly as f_ (x_) = 0_, where f_, x_ 
# and 0_ are vectors.
#
# Steps to solving a system of nonlinear equations using root:
# (1) import module: import scipy.optimize as opt
# (2) Put the equation in residual form, f_(x_) = 0_
# (3) Define the python function, 
#           def <fn name> (<args>):
#               ...
#               return <f(x) function>
#     This is the different part. The function takes a vector argument and
#     returns a vector.
# (4) Set an initial guess for x (which is now a vector)
# (5) Call "root" function: 
#           x = opt.root(<fn name>, <guess>).x

# ----------
# Example 
# ----------
# Solve:
#   y + 2z = 0
#   sin(y)/z = 0

# Writing this in our standard notation:
#   x0 + 2*x1 = 0
#   sin(x0)/x1 = 0

def F(x):
  F0 = x[0] + 2*x[1]
  F1 = np.sin(x[0]) / x[1]
  return np.array([F0, F1])

x_guess = np.array([1, 2])
x = opt.root( F, x_guess ).x
print('Example 2')
print('(y, z) = ', x)
print("residual = ", F(x))

# ----------
# Practice 
# ----------
# Solve the following system of three equations in three unknowns
# 
# x^2 + y^2 = 4
# xy + yz = -1
# y^2 + z^2 = 2.
# 
# A reasonable guess for all variables is x = 1, y = 2, z = 1.
# (Answer: x = 1.60890092, y = -1.18803949, z = -0.76717806)