#!/usr/bin/env python
# coding: utf-8

# In[1]:


#numerical methods for solving initial value problems

#Euler method
def euler(f, a, b, ya, n):
    h = (b - a) / n
    from numpy import zeros, linspace
    x = linspace(a, b, n+1)
    y = zeros(n+1)
    y[0] = ya
    for i in range(n):
        y[i+1] = y[i] + h * f(x[i], y[i])  
    return x, y


# In[3]:


f = lambda t, y: t - y
a = 0
b = 2
ya = 1
n = 3
t, y = euler(f, a, b, ya, n)


# In[5]:


t20, y20 = euler(f, a, b, ya, 20)


# In[10]:


t200, y200 = euler(f, a, b, ya, 200)


# In[11]:


from matplotlib.pyplot import plot, figure, grid, title, xlabel, ylabel, legend, show
plot(t, y, t20, y20, t200, y200)
legend(['n=3', 'n=20', 'n=200'])


# In[12]:


#Heun method
def heun(f, a, b, ya, n):
    h = (b - a) / n
    from numpy import zeros, linspace
    x = linspace(a, b, n+1)
    y = zeros(n+1)
    y[0] = ya
    for i in range(n):
        K1 = h * f(x[i], y[i])
        K2 = h * f(x[i+1], y[i]+K1)
        y[i+1] = y[i] + (K1 + K2)/2  
    return x, y


# In[18]:


t_heun, y_heun = heun(f, a, b, ya, n)


# In[15]:


#RK4 method
def rk4(f, a, b, ya, n):
    h = (b - a) / n
    from numpy import zeros, linspace
    x = linspace(a, b, n+1)
    y = zeros(n+1)
    y[0] = ya
    for i in range(n):
        K1 = h * f(x[i], y[i])
        K2 = h * f(x[i] + h/2, y[i] + K1/2)
        K3 = h * f(x[i] + h/2, y[i] + K2/2)
        K4 = h * f(x[i+1], y[i] + K3)
        y[i+1] = y[i] + (K1 + 2*(K2 + K3) + K4)/6  
    return x, y


# In[17]:


t_rk4, y_rk4 = rk4(f, a, b, ya, n)


# In[19]:


plot(t, y, t20, y20, t200, y200, t_heun, y_heun, t_rk4, y_rk4)
legend(['Euler n=3', 'Euler n=20', 'Euler n=200', 'Heun n=3', 'RK4 n=3'])


# In[22]:


from scipy.integrate import solve_ivp
sol = solve_ivp(f, [a, b], [ya])


# In[23]:


sol


# In[26]:


plot(t, y, t_heun, y_heun, t_rk4, y_rk4, sol.t, sol.y[0,])
legend(['Euler n=3', 'Heun n=3', 'RK4 n=3', 'solve_ivp'])


# In[8]:


#solve_ivp for a system of first-order ODEs
f = lambda t, z: [t - z[1], z[0]]
z0 = [3, 2]
ts = [0, 2]
from scipy.integrate import solve_ivp
sol = solve_ivp(f, ts, z0, dense_output=True)


# In[4]:


sol


# In[7]:


from matplotlib.pyplot import plot, figure, grid, title, xlabel, ylabel, legend, show
plot(sol.t, sol.y[0, :], sol.t, sol.y[1, :])
legend(['z0(t)', 'z1(t)'])


# In[9]:


from numpy import linspace
t = linspace(0, 2, 100)
z = sol.sol(t)
plot(t, z[0, :], t, z[1, :])
legend(['z0(t)', 'z1(t)'])


# In[11]:


plot(z[0, :], z[1, :])
xlabel('z0(t)')
ylabel('z1(t)')


# In[16]:


#example from slides
#theta'' = -theta' - sin(theta) + sin(a*t)
#-->
#z0' = z1
#z1' = -z1 - sin(z0) + sin(a*t)

a = 0.1
from numpy import sin
f = lambda t, z: [z[1], -z[1] - sin(z[0]) + sin(a*t)]
z0 = [0, 0.1]
ts = [0, 200]
from scipy.integrate import solve_ivp
sol = solve_ivp(f, ts, z0, dense_output=True)



# In[17]:


t = linspace(0, 200, 100)
z = sol.sol(t)
plot(t, z[0, :])


# In[18]:


1 + 2j


# In[ ]: