import math
import numpy as np
import matplotlib.pyplot as plt
# Simulation parameters
#______________________________________________________________________________________
tmax = 2 # Simulation time
timeStep = 0.001 # Simulation time step
n = int(tmax/timeStep)
time = np.linspace(0, tmax, n, endpoint = False)
# Physical parameters
#______________________________________________________________________________________
ycmo = 10 # Initial height of center of mass
L = 1 # Length of spring
k = 5000# Spring constant
m = 1 # Mass m
M = 1 # Mass M
dyo = 1.5 # If 0 < ycmo < L ---> the spring is compressed, if ycmo ---> L , the spring is extended
g = 9.81 # Gravity
# Calculating body trajectories
#______________________________________________________________________________________
y1 = (ycmo + (M*L)/(M+m))*np.ones(n,float) - 0.5*g*pow(time,2) + \
( M*(dyo - L)/(M+m) )*np.sin( ( math.sqrt( k/( (M*m)/(M+m) ) ) )*time + ( (math.pi)/2 )*np.ones( n,float) )
y2 = (ycmo - (m*L)/(M+m))*np.ones(n,float) - 0.5*g*pow(time,2) - \
( m*(dyo - L)/(M+m) )*np.sin( ( math.sqrt( k/( (M*m)/(M+m) ) ) )*time + ( (math.pi)/2 )*np.ones( n,float) )
ycm = (m*y1 + M*y2)/(M+m)
# Draw trajectories
#______________________________________________________________________________________
plt.plot(time,y1)
plt.plot(time,y2)
plt.plot(time,ycm)
plt.title('y1(t),y2(t),ycm(t)')
plt.xlabel('t(s)')
plt.ylabel('y(m)')
plt.grid()
plt.legend(['y1(t)','y2(t)','ycm(t)'])
plt.show()
