#! /usr/bin/env
#
def lesson02 ( ):

#*****************************************************************************80
#
## lesson02 presents the nonlinear convection equation.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license.
#
#  Modified:
#
#    14 November 2015
#
#  Author:
#
#    Initial version by Lorena Barba
#    This version by John Burkardt
#
#-------------------------------------------------------------------------------
Step 2: Non-linear Convection
Now we're going to implement non-linear convection using the same methods as in step 1. The 1D convection equation is:
∂u∂t+u∂u∂x=0
Instead of a constant factor c multiplying the second term, now we have the solution u multiplying it. Thus, the second term of the equation is now non-linear We're going to use the same discretization as in Step 1 — forward difference in time and backward difference in space. Here is the discretized equation.
un+1i−uniΔt+uniuni−uni−1Δx=0
Solving for the only unknown term, un+1i, yields:
un+1i=uni−uniΔtΔx(uni−uni−1)
As before, the Python code starts by loading the necessary libraries. Then, we declare some variables that determine the discretization in space and time (you should experiment by changing these parameters to see what happens). Then, we create the initial condition u0 by initializing the array for the solution using u=2 @ 0.5≤x≤1 and u=1 everywhere else in (0,2) (i.e., a hat function).
#-------------------------------------------------------------------------------
import numpy                 #we're importing numpy and calling it np locally
from matplotlib import pyplot    #and our 2D plotting library, calling it plt
%matplotlib inline


nx = 41
dx = 2./(nx-1)
nt = 20    #nt is the number of timesteps we want to calculate
dt = .025  #dt is the amount of time each timestep covers (delta t)

u = numpy.ones(nx)      #as before, we initialize u with every value equal to 1.
u[.5/dx : 1/dx+1]=2  #then set u = 2 between 0.5 and 1 as per our I.C.s

un = numpy.ones(nx) #initialize our placeholder array un, to hold the time-stepped solution
#-------------------------------------------------------------------------------
The code snippet below is unfinished. 
We have copied over the line from Step 1 that executes the time-stepping update. 
Can you edit this code to execute the non-linear convection instead?
#-------------------------------------------------------------------------------
for n in range(nt):  #iterate through time
    un = u.copy() ##copy the existing values of u into un
    for i in range(1,nx):  ##now we'll iterate through the u array
    
     ###This is the line from Step 1, copied exactly.  Edit it for our new equation.
     ###then uncomment it and run the cell to evaluate Step 2   
      
           ###u[i] = un[i]-c*dt/dx*(un[i]-un[i-1]) 

        
pyplot.plot(numpy.linspace(0,2,nx),u) ##Plot the results
#-------------------------------------------------------------------------------
What do you observe about the evolution of the hat function under the non-linear convection equation? What happens when you change the numerical parameters and run again?
#-------------------------------------------------------------------------------
For a careful walk-through of the discretization of the convection equation 
with finite differences (and all steps from 1 to 4), watch Video Lesson 4 by 
Prof. Barba on YouTube.
#-------------------------------------------------------------------------------