Review of Lab 9 for use in Second Project

Application to Lorenz system

from resources306 import *
%config InlineBackend.figure_format = 'retina'

def euler_step(F,t,X,h):
    newX = X + h*F(t,X)  # this vector arithmetic is done elementwise
    newt = t + h
    return newt,newX

def rk_step(F,t,X,h):  # this is a plug-in replacement for euler_step, that is more accurate
    k1 = F(t,X)  
    k2 = F(t+h/2,X+k1*h/2)
    k3 = F(t+h/2,X+k2*h/2)
    k4 = F(t+h  ,X+k3*h  )
    newX = X + h*(k1+2*k2+2*k3+k4)/6
    newt = t + h
    return newt,newX

# The above is for any system of first-order differential equations.

# Below we specify a particular system.

def F(t,X): 
    # This function computes and returns the rates of change of the dependent variables according to our DEs.
    # X is the vector of the current values of the dependent variables.
    # The independent variable t may or may not be needed.
    
    # unpack elements of the dependent variable vector X
    x,y,z = X
    
    # compute and return the rates of change of the dependent variables
    # Lorenz system
    xprime = 10*y - 10*x
    yprime = 28*x - y - x*z
    zprime = -8/5*z + x*y
    return np.array([xprime,yprime,zprime])

d = 3   # dimension of system

t0     = 0
tfinal = 10

nsteps = 2000
h = (tfinal-t0)/nsteps

t = t0
X = -4,-8,15   # Here is where the initial values of the dependent variables are specified

data = np.zeros((nsteps+1,1+d))  # we will store t and X in the rows of this array called "data"
data[0] = t,*X

for i in range(nsteps):
    t,X = euler_step(F,t,X,h)
    data[i+1] = t,*X

#print(data)

t,x,y,z = data.T
plt.figure()
plt.plot(t,x,label='x')
plt.plot(t,y,label='y')
plt.plot(t,z,label='z')
plt.xlabel('t')
plt.legend();

plt.figure()
plt.plot(y,z,'m',alpha=0.6)
plt.plot(y[ 0],z[ 0],'go',alpha=.5,label='initial')
plt.plot(y[-1],z[-1],'ro',alpha=.5,label='final')
plt.xlabel('y'); plt.ylabel('z',rotation=0);
plt.legend();