#!/usr/bin/env python
# coding: utf-8

# <script async src="https://www.googletagmanager.com/gtag/js?id=UA-59152712-8"></script>
# <script>
#   window.dataLayer = window.dataLayer || [];
#   function gtag(){dataLayer.push(arguments);}
#   gtag('js', new Date());
# 
#   gtag('config', 'UA-59152712-8');
# </script>
# 
# # Validating Runge Kutta Butcher tables using Truncated Taylor Series
# ## Authors: Zach Etienne & Brandon Clark
# 
# 
# ## This tutorial notebook is designed to validate the Butcher tables contained within the Butcher dictionary constructed in the [RK Butcher Table Dictionary](Tutorial-RK_Butcher_Table_Dictionary.ipynb) NRPy+ module.  
# 
# ### NRPy+ Source Code for this module: 
# * [MoLtimestepping/RK_Butcher_Table_Validation.py](../edit/MoLtimestepping/RK_Butcher_Table_Validation.py) Stores the `Validate` function for validating convergence orders for Runge Kutta methods
# * [MoLtimestepping/RK_Butcher_Table_Dictionary.py](../edit/MoLtimestepping/RK_Butcher_Table_Dictionary.py) [\[**tutorial**\]](Tutorial-RK_Butcher_Table_Dictionary.ipynb) Accesses the Butcher table dictionary `Butcher_dict` for known explicit Runge Kutta methods
# 
# ## Introduction:
# 
# Starting with the ODE (ordinary differential equation) initial value problem:
# $$
# y'(t) = f(y,t)\ \ \ y\left(t=0\right)=y_0,
# $$
# for various choices of $f(y,t)$, this module validates the Runge Kutta (RK) methods coded in [RK_Butcher_Table_Dictionary.py](../edit/MoLtimestepping/RK_Butcher_Table_Dictionary.py) [**tutorial notebook**](Tutorial-RK_Butcher_Table_Dictionary.ipynb) as follows:
# 
# Given $y_0$, and a smooth $f(y,t)$, all explicit RK methods provide an estimate for $y_1 = y\left(\Delta t\right)$, with an error term that is proportional to $\left(\Delta t\right)^m$, where $m$ is an integer typically greater than zero. This error term corresponds to the *local* truncation error. For RK4, for example, while the *total accumulated truncation error* (i.e., the accumulated error at a fixed final time $t_f$) is proportional to $\left(\Delta t\right)^4$, the *local* truncation error (i.e., the error after one arbitrarily chosen timestep $\Delta t$) is proportional to $\left(\Delta t\right)^5$.
# 
# If the exact solution $y(t)$ is known as a closed-form expression, then $y\left(\Delta t\right)$ can be *separately* written as a Taylor expansion about $y(t=0)$:
# 
# $$
# y\left(\Delta t\right) = \sum_{n=0}^\infty \frac{y^{(n)}(t=0)}{n!} \left(\Delta t\right)^n,
# $$
# where $y^{(n)}(t=0)$ is the $n$th derivative of $y(t)$ evaluated at $t=0$.
# 
# The above expression will be known exactly. Further if one chooses a numerical value for $y_0$ *and leaves $\Delta t$ unspecified*, any explicit RK method will provide an estimate for $y\left(\Delta t\right)$ of the form
# 
# $$
# y\left(\Delta t\right) = \sum_{n=0}^\infty a_n \left(\Delta t\right)^n,
# $$
# where $a_n$ *must* match the Taylor expansion of the *exact* solution at least up to and including terms proportional to $\left(\Delta t\right)^m$, where $m$ is the order of the local truncation error. If this is *not* the case, then the Butcher table was almost certainly *not* typed correctly.
# 
# Therefore, comparing the numerical result with unspecified $\Delta t$ against the exact Taylor series provides a convenient (though not perfectly robust) means to verify that the Butcher table for a given RK method was typed correctly. Multiple typos in the Butcher tables were found using this approach.
# 
# **Example from Z. Etienne's MATH 521 (Numerical Analysis) lecture notes:**
# 
# Consider the ODE
# $$
# y' = y - 2 t e^{-2t},\quad y(0)=y(t_0)=0.
# $$
# 
# * Solve this ODE exactly, then Taylor expand the solution about $t=0$ to
# approximate the solution at $y(t=\Delta t)$ to fifth order in $\Delta
# t$.
# * Next solve this ODE using Heun's method (second order in total accumulated truncation error, third order in local truncation error) {\it by hand} with a step size of
# $\Delta t$ to find $y(\Delta t)$. Confirm that the solution obtained
# when using Heun's method has an error term that is at worst
# $\mathcal{O}\left((\Delta t)^3\right)$. If the dominant error is
# proportional to a higher power of $\Delta t$, explain the discrepancy.
# 
# * Finally solve this ODE using the Ralston method {\it by hand}
#   with a step size of $\Delta t$ to find $y(\Delta t)$. Is the
#   coefficient on the dominant error term closer to the exact solution
#   than Heun's method?
# 
# We can solve this equation via the method of integrating factors,
# which states that ODEs of the form:
# $$
# y'(t) + p(t) y(t) = g(t)
# $$
# are solved via 
# $$
# y(t) = \frac{1}{\mu(t)} \left[ \int \mu(s) g(s) ds + c \right],
# $$
# where the integrating factor $\mu(t)$ is given by
# $$
# \mu(t) = \exp\left(\int p(t) dt\right)
# $$
# 
# Here, $p(t)=-1$ and $g(t) = - 2 t e^{-2t}$. Then
# 
# \begin{equation}
# \mu(t) = \exp\left(-\int dt\right) = e^{-t+c} = k e^{-t}
# \end{equation}
# and
# 
# \begin{align}
# y(t) &= e^t/k  \left[ \int k e^{-s} (- 2 s e^{-2s}) ds + c \right] = -2 e^t \left[ \int s e^{-3s} ds + c' \right] \\
# &= -2 e^t \left[ e^{-3 t} \left(-\frac{t}{3}-\frac{1}{9}\right) + c' \right] = -2 e^{-2t} \left(-\frac{t}{3}-\frac{1}{9}\right) -2 c' e^t \\
# &= e^{-2t} \left(2\frac{t}{3}+\frac{2}{9}\right) + c'' e^t \\
# \end{align}
# 
# If $y(0)=0$ then we can compute the integration constant $c''$, and
# $y(t)$ becomes
# $$
# y(t) = \frac{2}{9} e^{-2 t} \left(3 t + 1 - e^{3 t}\right).
# $$
# 
# The Taylor Series expansion of the exact solution about $t=0$
# evaluated at $y(\Delta t)$ yields
# $$
# y(\Delta t) = -(\Delta t)^2+(\Delta t)^3-\frac{3 (\Delta t)^4}{4}+\frac{23 (\Delta
#   t)^5}{60}-\frac{19 (\Delta t)^6}{120}+O\left((\Delta t)^7\right).
# $$
# 
# Next we evaluate $y(\Delta t)$ using Heun's method. We know $y(0)=y_0=0$ and
# $f(y,t)=y - 2 t e^{-2t}$, so
# \begin{align}
# k_1 &= \Delta t f(y(0),0) \\
#     &= \Delta t \times 0 \\
#     &= 0 \\
# k_2 &= \Delta t f(y(0)+k_1,0+\Delta t) \\
#    &= \Delta t f(y(0)+0,0+\Delta t) \\
#    &= \Delta t (-2 \Delta t e^{-2\Delta t}) \\
#    &= -2 (\Delta t)^2 e^{-2\Delta t} \\
# y(\Delta t) &= y_0 + \frac{1}{2} (k_1 + k_2) + \mathcal{O}\left((\Delta t)^3\right) \\
# &= 0 - (\Delta t)^2 e^{-2\Delta t} \\
# &= - (\Delta t)^2 ( 1 - 2 \Delta t + 2 (\Delta t)^2 + ...) \\
# &= - (\Delta t)^2 + 2 (\Delta t)^3 + \mathcal{O}\left((\Delta t)^4\right).
# \end{align}
# 
# Thus the coefficient on the $(\Delta t)^3$ term is wrong, but
# this is completely consistent with the fact that our stepping
# scheme is only third-order accurate in $\Delta t$.
# 
# In the below approach, the RK result is subtracted from the exact Taylor series result, as a check to determine whether the RK Butcher table was coded correctly; if it was not, then the odds are good that the RK results will not match to the expected local truncation error order. Multiple $f(y,t)$ are coded below to improve the robustness of this test.

# <a id='toc'></a>
# 
# # Table of Contents
# $$\label{toc}$$
# 
# This notebook is organized as follows
# 
# 1. [Step 1](#initializenrpy): Initialize needed Python/NRPy+ modules
# 1. [Step 2](#table_validate) Validate Convergence Order of Butcher Tables
#     1. [Step 2.a](#rhs): Defining the right-hand side of the ODE
#     1. [Step 2.b](#validfunc): Defining a Validation Function
#     1. [Step 2.c](#rkvalid): Validating RK Methods against ODEs
# 1. [Step 3](#latex_pdf_output): Output this notebook to $\LaTeX$-formatted PDF file

# <a id='initializenrpy'></a>
# 
# # Step 1: Initialize needed Python/NRPy+ modules [Back to [top](#toc)\]
# $$\label{initializenrpy}$$
# 
# Let's start by importing all the needed modules from Python/NRPy+:

# In[1]:


import sympy as sp              # SymPy: The Python computer algebra package upon which NRPy+ depends
import numpy as np              # NumPy: A numerical methods module for Python
from MoLtimestepping.RK_Butcher_Table_Dictionary import Butcher_dict


# <a id='table_validate'></a>
# 
# # Step 2:  Validate Convergence Order of Butcher Tables [Back to [top](#toc)\]
# $$\label{table_validate}$$
# 
# 
# Each Butcher table/Runge Kutta method is tested by solving an ODE. Comparing the Taylor series expansions of the exact solution and the numerical solution as discussed in the **Introduction** above will confirm whether the method converges to the appropriate order. 

# <a id='rhs'></a>
# 
# ## Step 2.a:  Defining the right-hand side of the ODE [Back to [top](#toc)\]
# $$\label{rhs}$$
# 
# Consider the form of ODE $y'=f(y,t)$. The following begins to construct a dictionary `rhs_dict` of right-hand side functions for us to validate explicit Runge Kutta methods. The most up-to-date catalog of functions stored in `rhs_dict` can be found in the [RK_Butcher_Table_Validation.py](../edit/MoLtimestepping/RK_Butcher_Table_Validation.py) module. 

# In[2]:


def fypt(y,t): # Yields expected convergence order for all cases
    #            except DP6 which converge to higher order (7, respectively)
    return y+t

def fy(y,t): # Yields expected convergence order for all cases
    return y

def feypt(y,t): # Yields expected convergence order for all cases
    return sp.exp(1.0*(y+t))

def ftpoly6(y,t): # Yields expected convergence order for all cases, L6 has 0 error
    return 2*t**6-389*t**5+15*t**4-22*t**3+81*t**2-t+42
rhs_dict = {'ypt':fypt, 'y':fy, 'eypt':feypt, 'tpoly6':ftpoly6}


# <a id='validfunc'></a>
# 
# ## Step 2.b: Defining a Validation Function  [Back to [top](#toc)\]
# $$\label{validfunc}$$
# 
# To validate each Butcher table we compare the exact solutions to ODEs with the numerical solutions using the Runge Kutta scheme built into each Butcher table. The following is a function that
# 
# 1. Solves the ODE exactly,
# 2. Solves the ODE numerically for a given Butcher table, and
# 3. Compares the two solutions and checks for the order of convergence by returning their difference.
# 
# The `Validate()` function inputs a specified `Butcher_key`, the starting guess solution and time `y_n`, `t_n` and the right-hand side of the ODE corresponding to a specified initial value problem, `rhs_key`.
# 
# 
# 

# In[3]:


from MoLtimestepping.RK_Butcher_Table_Dictionary import Butcher_dict
def Validate(Butcher_key, yn, tn, rhs_key):
    # 1. First we solve the ODE exactly
    y = sp.Function('y')
    sol = sp.dsolve(sp.Eq(y(t).diff(t), rhs_dict[rhs_key](y(t), t)), y(t)).rhs
    constants = sp.solve([sol.subs(t,tn)-yn])
    exact = sol.subs(constants)

    # 2. Now we solve the ODE numerically using specified Butcher table

    # Access the requested Butcher table
    Butcher = Butcher_dict[Butcher_key][0]
    # Determine number of predictor-corrector steps
    L = len(Butcher)-1
    # Set a temporary array for update values
    k = np.zeros(L, dtype=object)
    # Initialize intermediate variable
    yhat = 0
    # Initialize the updated solution
    ynp1 = 0
    for i in range(L):
        #Initialize and approximate update for solution
        yhat = yn
        for j in range(i):
            # Update yhat for solution using a_ij Butcher table coefficients
            yhat += Butcher[i][j+1]*k[j]
            if Butcher_key == "DP8" or Butcher_key == "L6":
                yhat = 1.0*sp.N(yhat,20) # Otherwise the adding of fractions kills performance.
        # Determine the next corrector variable k_i using c_i Butcher table coefficients
        k[i] = dt*rhs_dict[rhs_key](yhat, tn + Butcher[i][0]*dt)
        # Update the solution at the next iteration ynp1 using Butcher table coefficients
        ynp1 += Butcher[L][i+1]*k[i]
    # Finish determining the solution for the next iteration
    ynp1 += yn

    # Determine the order of the RK method
    order = Butcher_dict[Butcher_key][1]+2
    # Produces Taylor series of exact solution at t=tn about t = 0 with the specified order
    exact_series = sp.series(exact.subs(t, dt),dt, 0, order)
    num_series = sp.series(ynp1, dt, 0, order)
    diff = exact_series-num_series
    return diff


# <a id='rkvalid'></a>
# 
# ## Step 2.c: Validating RK Methods against ODEs  [Back to [top](#toc)\]
# $$\label{rkvalid}$$
# 
# The following makes use of the `Validate()` function above to demonstrate that each method within the Butcher table dictionary converges to the expected order for the given right-hand side expression. 

# In[4]:


t, dt = sp.symbols('t dt')
# Set initial conditions
t0 = 0
y0 = 1
# Set RHS of ODE
function = 'ypt'# This can be changed, just be careful that the initial conditions are satisfied
for key,value in Butcher_dict.items():
    print("RK method: \""+str(key)+"\".")
    y = sp.Function('y')
    print(" When solving y'(t) = "+str(rhs_dict[function](y(t),t))+", y("+str(t0)+")="+str(y0)+",")
    local_truncation_order = list(value)[1]+1
    print(" the first nonzero term should have local truncation error proportional to O(dt^"+str(local_truncation_order)+") or a higher power of dt.")
    print("Subtracting the numerical result from the exact Taylor expansion, we find a local truncation error of:")
    sp.pretty_print(Validate(key, y0, t0, function))
#     print("\n")
    print(" (Coefficients of order 1e-15 or less may generally be ignored, as these are at roundoff error.)\n")


# <a id='latex_pdf_output'></a>
# 
# # Step 3: Output this notebook to $\LaTeX$-formatted PDF file \[Back to [top](#toc)\]
# $$\label{latex_pdf_output}$$
# 
# The following code cell converts this Jupyter notebook into a proper, clickable $\LaTeX$-formatted PDF file. After the cell is successfully run, the generated PDF may be found in the root NRPy+ tutorial directory, with filename
# [Tutorial-RK_Butcher_Table_Validation.pdf](Tutorial-RK_Butcher_Table_Validation.pdf) (Note that clicking on this link may not work; you may need to open the PDF file through another means.)

# In[5]:


import cmdline_helper as cmd    # NRPy+: Multi-platform Python command-line interface
cmd.output_Jupyter_notebook_to_LaTeXed_PDF("Tutorial-RK_Butcher_Table_Validation")