%pylab inline

from matplotlib import pyplot as plt
import numpy as np

from scipy.integrate import odeint

from colorline import colorline

# Display the help for the odeint function:
# The output can be shortened or suppressed by clicking or double clicking on the region to the left of the output

help(odeint)

a = 2

def f(y, t):
    return a*y

y0 = 1.0

t_output = np.arange(0, 6, 0.1)
print t_output

y_result = odeint(f, y0, t_output)
y_result = y_result[:, 0]  # convert the returned 2D array to a 1D array

#plt.plot(t_output, y_result)
colorline(t_output, y_result)
plt.plot(t_output, y0 * np.exp(a * t_output))


plt.plot(t_output, np.abs(y_result - y0 * np.exp(a * t_output)))

omega_squared = 4

def harmonic(x_vec, t):
    x, y = x_vec  # tuple unpacking
    
    return [y, -omega_squared*x]

y0 = [0.7, 0.5]

t_output = np.arange(0, 5, 0.1)

y_result = odeint(harmonic, y0, t_output)

#plt.figure(figsize = (5,5))

fig, (ax1, ax2, ax3) = plt.subplots(ncols = 3, figsize=(10,10))

xx, yy = y_result.T  # extract x and y cols

#ax1.plot(xx, yy)
plt.sca(ax1)
colorline(xx, yy, cmap='jet')
ax1.axis('scaled')

#ax2.plot(t_output, xx)
plt.sca(ax2)
colorline(t_output, xx, cmap='cool')
ax2.axis('scaled')

ax3.plot(t_output, yy)
ax3.axis('scaled')

plt.tight_layout()
plt.show()



x = np.linspace(-1.0, 1.0, 10)
XX, YY = np.meshgrid(x, x)

k = omega_squared
plt.quiver(XX, YY, YY, -k*XX, pivot='middle')
plt.axis('equal')

plt.streamplot(XX, YY, YY, -k*XX)
plt.quiver(XX, YY, YY, -k*XX, pivot = 'middle')
plt.axis('equal')

MM = np.array([[-1.5, 1.], [-1., -1.0]])

def matrix(x_vec, t):
    
    return np.dot(MM, x_vec)

y0 = [1, 1]

t_output = np.arange(0, 10, 0.1)
y_result = odeint(matrix, y0, t_output)

fig, (ax1, ax2, ax3) = plt.subplots(ncols = 3)

xx, yy = y_result.T  # extract x and y cols

ax1.plot(xx, yy)
ax1.axis('scaled')

ax2.plot(t_output, xx)
ax2.axis('scaled')

ax3.plot(t_output, yy)
ax3.axis('scaled')

plt.tight_layout()
plt.show()

from scipy.integrate import ode

help(ode)

omega_squared = 4

def harmonic_new(t, x_vec):
    x, y = x_vec  # tuple unpacking
    
    return [y, -omega_squared*x]

t_final = 10.0
dt = 0.1

y0 = [0.7, 0.5]
t0 = 0.0

y_result = []
t_output = []

# Initialisation:

backend = "dopri5"

solver = ode(harmonic_new)
solver.set_integrator(backend)  # nsteps=1
solver.set_initial_value(y0, t0)

y_result.append(y0)
t_output.append(t0)

while solver.successful() and solver.t < t_final:
    solver.integrate(solver.t + dt, step=1)
    
    y_result.append(solver.y)
    t_output.append(solver.t)
    
y_result = array(y_result)
t_output = array(t_output)

#print y_result

from scipy.integrate import ode

def my_odeint(f, y0, t):
    '''
    ODE integrator compatible with odeint, that uses ode underneath
    '''
    
    y0 = np.asarray(y0)
    
    backend = "dopri5"
    
    solver = ode(f)
    solver.set_integrator(backend)  # nsteps=1
    
    t0 = t[0]
    t_final = t[-1]
    
    solver.set_initial_value(y0, t0)
    
    y_result = [y0]
    
    i = 1
    current_t = t[i]
    
    while solver.successful() and solver.t < t_final:
        solver.integrate(current_t, step=1)
        i += 1
        if i < len(t):
            current_t = t[i]
    
        y_result.append(solver.y)
    
    return np.array(y_result)

t_output = np.arange(0, 10, 0.01)
y0 = [1, 1]

y_result = my_odeint(harmonic_new, y0, t_output)


fig, (ax1, ax2, ax3) = plt.subplots(ncols = 3, figsize=(10,10))

xx, yy = y_result.T  # extract x and y cols

#ax1.plot(xx, yy)
plt.sca(ax1)
colorline(xx, yy, cmap='jet')
ax1.axis('scaled')

#ax2.plot(t_output, xx)
plt.sca(ax2)
colorline(t_output, xx, cmap='cool')
ax2.axis('scaled')

plt.sca(ax3)
colorline(t_output, yy)
ax3.axis('scaled')

plt.tight_layout()
plt.show()

from numpy import zeros, linspace, array
from scipy.integrate import ode
from pylab import figure, show, xlabel, ylabel
from mpl_toolkits.mplot3d import Axes3D

def lorenz_sys(t, q):
    x = q[0]
    y = q[1]
    z = q[2]
    # sigma, rho and beta are global.
    f = [sigma * (y - x),
         rho*x - y - x*z,
         x*y - beta*z]
    return f


ic = [1.0, 2.0, 1.0]
t0 = 0.0
t1 = 100.0
dt = 0.01

sigma = 10.0
rho = 28.0
beta = 10.0/3

solver = ode(lorenz_sys)

t = []
sol = []
solver.set_initial_value(ic, t0)
#solver.set_integrator('dop853')
solver.set_integrator('dopri5')

while solver.successful() and solver.t < t1:
    solver.integrate(solver.t + dt)
    t.append(solver.t)
    sol.append(solver.y)

t = array(t)
sol = array(sol)

fig = figure()
ax = Axes3D(fig)
ax.plot(sol[:,0], sol[:,1], sol[:,2])
xlabel('x')
ylabel('y')
show()