%matplotlib inline

import matplotlib
import numpy as np
import sympy as sy
from scipy import special, random, linalg, integrate
import matplotlib.pyplot as plt

def dX_dt(X, t, A):
    """ 
    Return the general action of a matrix on a vector $X$. 
    This allows for nice block multiplication.
    """
    return np.array([ A[0,0]*X[0] + A[0,1]*X[1],
                   A[1,0]*X[0] + A[1,1]*X[1]])
def gen_sol(A,x_0,t):
    """Compute general solution with the matrix exponential"""
    return np.matmul(linalg.expm(A*t),x_0)
def orbit_linear_2by2(A,x_0,T):
    """Compute linear orbit with initial data x_0 for t \in [start, stop]"""
    return np.apply_along_axis(lambda t:gen_sol(A,x_0,t),1,T)

### Define particular matrix
AA = np.array([[1,0],[0,-2]])

### Define a time interval
TT = np.linspace(-1.0,1.0,40).reshape(40,1)

### Plot some orbits by choosing different intial conditions (here we use matplotlib)
X_0 = np.array([1,1])
plt.plot(orbit_linear_2by2(AA,X_0,TT)[:,0],orbit_linear_2by2(AA,X_0,TT)[:,1])
X_0 = np.array([2,3])
plt.plot(orbit_linear_2by2(AA,X_0,TT)[:,0],orbit_linear_2by2(AA,X_0,TT)[:,1])
X_0 = np.array([-1,5])
plt.plot(orbit_linear_2by2(AA,X_0,TT)[:,0],orbit_linear_2by2(AA,X_0,TT)[:,1])
X_0 = np.array([-5,-1])
plt.plot(orbit_linear_2by2(AA,X_0,TT)[:,0],orbit_linear_2by2(AA,X_0,TT)[:,1])

###
x = np.linspace(-20, 10, 10)
y = np.linspace(-20, 10, 10)

X1 , Y1  = np.meshgrid(x, y)                       # create a grid
DX1, DY1 = dX_dt([X1, Y1],0,AA)
plt.quiver(X1, Y1, DX1, DY1)
plt.show()

### Nice option for rapid visualization s
from matplotlib.pyplot import cm
speed = np.sqrt(DX1**2 + DY1**2)
plt.streamplot(X1, Y1, DX1, DY1,          # data
               color=speed,         # array that determines the colour
               cmap=cm.cool,        # colour map
               linewidth=1,         # line thickness
               arrowstyle='->',     # arrow style
               arrowsize=1.5)       # arrow size
plt.show()

### Nonlinear equations

### We redefine the dX_dt function for each new ODE system
def dX_dt(X, t=0):
    """ 
    Logan, page 257, 1a
    """
    return np.array([ X[0] - X[1], 2.0*X[0]*(X[1] - 1.0)])

### Grid phase space
x = np.linspace(-3, 3, 20)
y = np.linspace(-3, 3, 20)
X1 , Y1  = np.meshgrid(x, y)

### Calculate vector field
DX1, DY1 = dX_dt([X1, Y1])

from matplotlib.pyplot import cm
speed = np.sqrt(DX1**2 + DY1**2)
plt.streamplot(X1, Y1, DX1, DY1,          # data
               color=speed,         # array that determines the colour
               cmap=cm.cool,        # colour map
               linewidth=1,         # line thickness
               arrowstyle='->',     # arrow style
               arrowsize=1)       # arrow size
plt.plot(x,x) # nullcline y = x
plt.plot(np.zeros_like(y),y) # nullcline x = 0
plt.plot(x,np.ones_like(x)) # nullcline y = 1

### long orbits
X_0 = np.array([0.1,1.1])
T = np.linspace(0,2,20)
X = integrate.odeint(dX_dt,X_0,T)
plt.plot(X[:,0],X[:,1],color='black')

### plot all the elements
plt.show()