# 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()