# Verbose, but simple, Python3 program to calculate rainsplash erosion
# this is written in a very linear (non object oriented) form

import matplotlib.pyplot as plt     # using a MATLAB-like plotting package
import numpy as np           # this is optional, straight python deals with
                             # vectors as ordered lists, while numpy treats
                             # then as vectors if you do math, (doesn't matter here)

C     = 3e-4         # rainsplash coeff.in years
delx  = 1            # x positions or bins are 1 m wide
delt  = 10           # time steps are 10 years
total_time = 10000   # run for 10000 years

# make arrays of x and z values, include a place to put new (next time) z values
z = np.array([3,3,3,3,2.25,1.5,.75,0,0,0,0])
                     # a list of the current z elevations at x=0 to x=10 {11 nodes}
                     # converted to a numpy array for easier math
zn = np.empty_like(z)       # this is a cute way of making a new array of length z
zn[:] = z                   # to start our process we need to have zn the same as z
x = np.arange(0,len(z),delx)# make a list of x values

# make a figure to plot the results
fig1 = plt.figure()             # makes a separate figure on which to place a plot
ax   = fig1.add_subplot(1,1,1)  # makes a set of axes on the figure
plotn = 0                       # we don't want to plot every calc, this will let us count
ax.plot(x,z)                    # plot the original profile

for t in np.arange(0,total_time,delt):    # make time step from now+delt to total time
    z[:] = zn                             # update profile to new 'z's, and continue
    for n in np.arange(1,len(x)-1):       # calc the new z, for the interior nodes
        zn[n]=z[n] + (C*delt/(delx*delx))*(z[n+1]-2*z[n]+z[n-1])   #this is our equation from class
    # add boundary conditions
    zn[0]       = 3
    zn[len(z)-1]= 0  
    #plotting            
    plotn +=1
    if (plotn % 100) == 0:      # only plot every 100*delt years, % is the modulus operator
        ax.plot(x,zn)	

ax.set_title('Rainsplash on a fault scarp')
ax.set_xlabel('meters')
ax.set_ylabel('meters')
s = "evolution over " + str(total_time) + " years"
ax.text(4.5,2.7,s,fontsize=16)

plt.show()