# Numerical Methods II, Courant Institute, NYU, spring 2018
# http://www.math.nyu.edu/faculty/goodman/teaching/NumericalMethodsII2018/index.html
#  written by Jonathan Goodman (instructor)
#  see class notes Part 1 for more discussion

# Use a simple low order method to solve an ODE

import numpy as np
import matplotlib.pyplot as plt

#    Description of the mathematical problem

sig  = 10.
beta = 8./3.
rho  = 28.


T0 = 0.       # start time for the simulation
T  = 10.;     # ending time

d = 3         # this problem has x(t) with three components

x0 = np.ndarray(d)       # initial condition
x  = np.ndarray(d)       # the value of x(t)
f  = np.ndarray(d)       # hold the vector f(x)


x0[0] = 1.     # set the initial condition once for all runs
x0[1] = 0.
x0[2] = 2.

#    Numerical parameters for simulation

dt_list = [.002, .001, .0001, .00001]      #  Repeat the experiment for these values of \Delta t



#           setup for plots

dtp = .01                          # time increment for plotting, proposed
npt = int( (T - T0)/dtp) + 1      # number of recorded times, for plotting, including x0
dtp = (T - T0)/(npt -1)           # recompute the print time step so there ...
                                  # ... are an integer number of prints in total
tp  = np.ndarray(npt)             # a list of times at which to print the solution
for it in range(npt):
    tp[it] = T0 + it*dtp

xp = np.ndarray([d,npt])          # xp[k,it] is component k of x at time tp[it]
fig, ax = plt.subplots()          # see pyplot documentation


#    Do a simimulation for each value of Delta t

for dt in dt_list:

    x = x0                 # initialize the trajectory x(t0) = X0
    t = T0

    ipt = 0                # the index of the time to record the solution for printing
    tpn = T0               # the time of the next print recording
    
    while (t < T) :        # the main time step loop
    
        f[0] = sig*(x[1]-x[0])              # evaluate f(x)
        f[1] = x[0]*( rho - x[2]) - x[1]
        f[2] = x[0]*x[1] - beta*x[2]
        
        if (t + dt) > tpn:           # if the next time step takes you past the print time ...
        
            dtf = tpn - t            # time from the present to the print time, should be < dt
            xp[:,ipt]  = x + dtf*f   # record the solution at to the print time
              
            ipt        = ipt+1       # find the next print time
            if ipt >= npt:           # If there aren't any more, ...
                break                # ... stop computing, you're done.
            tpn = tp[ipt]            # Otherwise, get the next print time

        x    = x + dt*f     # take a forward Euler time step
        t    = t + dt       # increment time, so x is x(t)
    
    dt_last = T - t         # take a final time step that may be smaller than dt
    f[0] = sig*(x[1]-x[0])              # evaluate f(x)
    f[1] = x[0]*( rho - x[2]) - x[1]
    f[2] = x[0]*x[1] - beta*x[2]
    x    = x + dt*f         # should be x(T)
    
    dt_string = "dt = " + str(dt)
    ax.plot( tp, xp[2,:], label = dt_string)

    legend = ax.legend(loc='upper right', shadow=True, prop={'size':10})

    # Put a nicer background color on the legend.
    legend.get_frame().set_facecolor('#00FFCC')

title = "Lorenz system, Euler solver"
plt.title(title)
plt.xlabel("t")
plt.ylabel("z")
plt.grid()
fig.savefig("ODE_fE_Lorenz.pdf")