#!/usr/bin/env python3

# integrates the simple model from Eq. (1) in Problem Set 2
# 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,N,params):
    r,a,b = params
    ndot = r*N*(-N**2+a*N-b)
    return ndot

# DEFINE PARAMETERS: this is what you might need to change
# constants r, a, and b
r = 1
a = 8 
b = 13

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

# timestep 
dt = 0.01

# initial condition 
N0 = 4 

# INTEGRATE AND PLOT
t = np.linspace(tstart,tend,int((tend-tstart)/dt))
N = np.zeros(len(t))
N[0] = N0 

model_instance = ode(model).set_integrator("dop853")
model_instance.set_initial_value(N[0]).set_f_params([r,a,b])

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

plt.plot(t,N)
plt.show()