import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as ode

g_rask=9.814
r_keha=1.8
s_keha=np.pi*r_keha**2
c_keha=0.47
v_keha=4.*np.pi*r_keha**3/3
m_keha=15.
p_ohu=100000.
t_ohu=293.
mu_ohu=29.
r_gaas=8314.

ekin=[]
epot=[]

ro_ohu_0=p_ohu*mu_ohu/(r_gaas*t_ohu)
p1=g_rask*mu_ohu/(r_gaas*t_ohu)
p2=0.5*c_keha*s_keha/m_keha
p3=g_rask*v_keha/m_keha
def func(y,t):
   yd=np.zeros_like(y)
   yd[0]=y[1]
   ro=ro_ohu_0*np.exp(-p1*y[0])
   v=np.abs(y[1])
   yd[1]=-g_rask-p2*ro*v*y[1]+ro*p3
   return yd

y0=[0.,0.]
time=np.linspace(0,5000,num=10000)
res=ode.odeint(func,y0,time)
ekin=0.5*m_keha*res[:,1]
epot=m_keha*g_rask*res[:,0]

fig,[p1,p2,p3]=plt.subplots(1,3)
p1.plot(time,res[:,1])
p2.plot(time,res[:,0])
p3.plot(time,ekin,time,epot)
p1.grid()
p2.grid()
p3.grid()
plt.show()