"""
Example of how to use the ODESolver class to solve a
system of ODEs. Two key points:
- Let the right hand side function f return a list, tuple or array. Here
a list with 4 components.
- Let the initial condition U0 be a list, tuple or array, with the same
number of components as f.
"""

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


def f(u, t):
     x, vx, y, vy = u
     g = 9.81
     return [vx, 0, vy, -g]

# Initial condition, start at the origin:
x = 0
y = 0
# velocity magnitude and angle:
v0 = 5;
theta = 80*np.pi/180
vx = v0*np.cos(theta)
vy = v0*np.sin(theta)

U0 = [x, vx, y, vy]

solver= ForwardEuler(f)
solver.set_initial_condition(U0)
time_points = np.linspace(0, 1.0, 101)
u, t = solver.solve(time_points)
# u is an array of [x,vx,y,vy] arrays, plot y vs x:
x = u[:,0];  y = u[:,2]

plt.plot(x, y)
plt.show()