#!/usr/bin/env python3

# integrates the Lorenz model
# we use the Scipy ode package: "dop853" refers to an eighth-order Runge-Kutta method (basically a fancy Euler method)

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import ode
plt.rcParams.update({'font.size':14})

# DEFINE FUNCTIONS
# the model itself
def model(t,x,params):
    Pr,b,r = params
    xdot = np.zeros(3)
    xdot[0] = Pr*x[1] - Pr*x[0]
    xdot[1] = -x[0]*x[2] + r*x[0] - x[1]
    xdot[2] = x[0]*x[1]-b*x[2]

    return xdot

# DEFINE PARAMETERS
Pr = 10  
b = 8/3 
r = 1 

# set up start and end times of the integration
tstart = 0
tend = 200

# timestep 
dt = 0.002

# initial condition 
x0 = [0.01, 0.01, 0.01]

# INTEGRATE AND PLOT
t = np.linspace(tstart,tend,int((tend-tstart)/dt))
x = np.zeros((len(t),3))

x[0,:] = x0 

model_instance = ode(model).set_integrator("dop853")
model_instance.set_initial_value(x[0,:]).set_f_params([Pr,b,r])

for i in range(1,len(t)):
    model_instance.integrate(model_instance.t+dt)
    x[i,:] = model_instance.y

plt.figure()
plt.plot(t,x[:,0])
plt.xlabel('Time')
plt.ylabel('x')

plt.show()