# Written by:
# Kelly Roos
# Engineering Physics
# Bradley University
# rooster@bradley.edu | 309.677.2997
import numpy as np
import matplotlib.pyplot as plt
import math
# Input parameters for model
dt = 0.01 # time step (s)
vi = 0 # initial velocity (m/s)
yi = 0.1 # initial position (m), relative to eq position yeq (m)
m = 1.0 # mass (kg)
g = 9.8 # gravity (m/s^2)
k = 10 # spring constant (N/m)
t_steps = 2000 # number of time steps
# Defines the 1D arrays to be used in the computation and
# sets all values in the arrays to zero
time = np.zeros(t_steps)
y_Euler = np.zeros(t_steps)
y_EC = np.zeros(t_steps)
y_exact = np.zeros(t_steps)
v_Euler = np.zeros(t_steps)
v_EC = np.zeros(t_steps)
v_exact = np.zeros(t_steps)
energyEuler = np.zeros(t_steps);
energyEC = np.zeros(t_steps);
# Initial conditions
time[1] = 0
y_Euler[1] = yi
y_EC[1] = yi
y_exact[1] = yi
v_Euler[1] = vi
v_EC[1] = vi
v_exact[1] = vi
energyEuler[1]=0.5*m*v_Euler[1]**2+0.5*k*(y_Euler[1]+m*g/k)**2-m*g*(y_Euler[1]+m*g/k)
energyEC[1]=0.5*m*v_EC[1]**2+0.5*k*(y_EC[1]+m*g/k)**2-m*g*(y_EC[1]+m*g/k)
# Main loop: Euler algorithm, euler-Cromer, and evaluation of exact solutions for v and y
for i in range(2, t_steps):
time[i] = time[i-1] + dt
# Euler
v_Euler[i]=v_Euler[i-1]-k*y_Euler[i-1]*dt/m
y_Euler[i]=y_Euler[i-1]+v_Euler[i-1]*dt
# Euler-Cromer
v_EC[i]=v_EC[i-1]-k*y_EC[i-1]*dt/m
y_EC[i]=y_EC[i-1]+v_EC[i]*dt
# Exact (assuming vi=0)
v_exact[i]=-yi*math.sqrt(k/m)*math.sin(math.sqrt(k/m)*time[i])
y_exact[i]=yi*math.cos(math.sqrt(k/m)*time[i])
# Energies
energyEuler[i]=0.5*m*v_Euler[i]**2+0.5*k*(y_Euler[i]+m*g/k)**2-m*g*(y_Euler[i]+m*g/k)
energyEC[i]=0.5*m*v_EC[i]**2+0.5*k*(y_EC[i]+m*g/k)**2-m*g*(y_EC[i]+m*g/k)
# Plotting Results
plt.plot(time, v_Euler, time, v_EC)
plt.ylim((-0.3, 0.3))
plt.xlim((0, 20))
plt.xlabel('time (s)')
plt.ylabel('position (m)')
plt.title('position vs. time')
plt.show()