# Examples of use of the ODESolver package

from ODESolver import *
import numpy as np
import matplotlib.pyplot as plt

#### Example 1: introductory example
#### ODE: u'(t) = u(t) with initial condition u(0)=1

# First define right hand side of ODE
def f(u,t):
   return u

# Use Forward Euler method
FE = ForwardEuler(f)

# Set initial condition
FE.set_initial_condition(1)

# Compute solution on interval [0, 5]
time_points = np.linspace(0, 5, 500)
u,t = FE.solve(time_points)

# Plot solution
plt.plot(t, u)
plt.show()


#### Example 2: use a class to implement f(u,t)
#### ODE: u'(t) = a u(t) with initial condition u(0)=1
#### where the parameter a = 0.15, 0.30 or 0.45

# Since the right hand side f(u,t) of the ODE has a parameter a
# it is most convenient to implement f(u,t) as a class:
class Fnc:
   def __init__(self, a):
      self.a = a
      
   def __call__(self, u, t):
      return self.a * u

# We create three instances of Fnc corresponding to f(u,t)
# with parameter a = 0.15, a = 0.30 and a = 0.45:
f1 = Fnc(0.15)
f2 = Fnc(0.30)
f3 = Fnc(0.45)

# Use Forward Euler method
FE1 = ForwardEuler(f1)
FE2 = ForwardEuler(f2)
FE3 = ForwardEuler(f3)

# Set initial conditions
FE1.set_initial_condition(0.5)
FE2.set_initial_condition(0.5)
FE3.set_initial_condition(0.5)

# Compute solutions on interval [0, 5]
time_points = np.linspace(0, 5, 500)
u1,t1 = FE1.solve(time_points)
u2,t2 = FE2.solve(time_points)
u3,t3 = FE3.solve(time_points)

# Plot solutions
plt.subplot(2,2,1)
plt.plot(t1, u1)
plt.subplot(2,2,2)
plt.plot(t2, u2)
plt.subplot(2,2,3)
plt.plot(t3, u3)
plt.show()


#### Example 3: same as previous example, but more compact code
#### ODE: u'(t) = a u(t) with initial condition u(0)=1
#### where the parameter a = 0.15, 0.30 or 0.45

class Fnc:
   def __init__(self, a):
      self.a = a
      
   def __call__(self, u, t):
      return self.a * u

time_points = np.linspace(0, 5, 500)
a_values = [0.15, 0.30, 0.45]

for i in range(len(a_values)):
   f = Fnc(a_values[i])
   FE = ForwardEuler(f)
   FE.set_initial_condition(0.5)
   u,t = FE.solve(time_points)
   plt.subplot(2,2,i+1)
   plt.plot(t,u)
   
plt.show()


#### Example 4: logistic ODE and Runge-Kutta solution method
#### ODE: u'(t) = a u(t) (1-u(t)) with initial condition u(0)=0.01
#### where the parameter a = 2.2

class Fnc:
   def __init__(self, a):
      self.a = a
      
   def __call__(self, u, t):
      return self.a * u * (1-u)

f = Fnc(2.2)
FE = RungeKutta4(f)
FE.set_initial_condition(0.01)
u,t = FE.solve(np.linspace(0, 5, 500))
plt.plot(t,u)
plt.show()


#### Example 5: solving a system of two ODE's (= a vector ODE)
#### ODE: 
####     x'(t) = a y(t)
####     y'(t) = b x(t)
#### Initial conditions: 
####     x(0) = 1, y(0) = 0
#### We solve it below for a = 1 and b = -1

class Fnc:
   def __init__(self, a, b):
      self.a = a
      self.b = b
      
   def __call__(self, u, t):
      f1 = self.a * u[1]   # Corresponds to a*y(t)
      f2 = self.b * u[0]   # Corresponds to b*x(t)
      return np.array([f1,f2])

f = Fnc(a = 1.0, b = -1.0)
FE = ForwardEuler(f)
FE.set_initial_condition([1, 0])
u,t = FE.solve(np.linspace(0, 6, 500))
plt.plot(t, u[:,0], 'b-')  # Plot x(t)
plt.plot(t, u[:,1], 'r-')  # Plot y(t)
plt.show()