import numpy as np
import matplotlib.pyplot as plt

M=5.97*pow(10,24)
R=6400000.0

def sign(a):
    xx=a
    if abs(xx)>0.0:
        return xx/abs(xx)
    else:
        return 1.0

def gravity(M,R):
    G=6.67*pow(10.0,-11.0)
    g=(G*M/R**2)
    return g

def posp(v,S,c,m,g):
    fi=np.arctan(v[1]/v[0])
    fi=abs(fi)
    abs_v=np.sqrt(v.dot(v))
    Fu=0.5*1.3*S*c*abs_v
    ax=-sign(v[0])*Fu*np.cos(fi)/m
    ay=-g-sign(v[1])*Fu*np.sin(fi)/m
    az=0.0
    return np.array([ax,ay,az])

def narisi_graf():
    data=np.loadtxt("rezultati_simulacije_upor.dat")
    fig = plt.figure(figsize=(24.0,6.5)) 
    ax1 = fig.add_subplot(131)
    ax1.plot(data[:,1],data[:,2],color='red',linewidth=2.5)
    ax1.grid(True)
    ax1.set_xlim([min(data[:,1])-0.1,max(data[:,1]+0.1)])
    ax1.set_ylim([min(data[:,2])-0.1,max(data[:,2]+0.1)])
    ax1.set_ylabel('$y$ [m]')
    ax1.set_xlabel('$x$ [m]')
    
    ax2 = fig.add_subplot(132)
    ax2.plot(data[:,0],data[:,3],color='blue',linewidth=2.5)
    ax2.grid(True)
    ax2.set_xlim([min(data[:,0])-0.1,max(data[:,0]+0.1)])
    ax2.set_ylim([min(data[:,3])-0.1,max(data[:,3]+0.1)])
    ax2.set_xlabel('$t$ [s]')
    ax2.set_ylabel('$v [\\frac{m}{s}]$')

    ax3 = fig.add_subplot(133)
    ax3.plot(data[:,0],data[:,4],color='olive',linewidth=2.5)
    ax3.grid(True)
    ax3.set_xlim([min(data[:,0])-0.1,max(data[:,0]+0.1)])
    ax3.set_ylim([min(data[:,4])-0.1,max(data[:,4])+0.1])
    ax3.set_ylabel('$a [\\frac{m}{s^2}]$')
    ax3.set_xlabel('$t$ [s]')

    plt.savefig("posevni_med_upor.png",bbox_inches='tight',dpi=80)
    plt.show()
    fig.clear()

dt=0.001
pi=np.pi
fi=45.0*(pi/180.0)
v0=90.0
g=gravity(M,R)

#LASTNOSTI IZSTRELKA#
C=0.47
rho=float(input("Vnesite gostoto telessa [kg/m^3]:"))
m=float(input("Vnesite maso telesa [kg]:"))
r=pow(3.0*m/(4.0*pi*rho),1.0/3.0)
S=4.0*pi*r**2


pos=np.array([0.0,0.0,0.0])
v=np.array([v0*np.cos(fi),v0*np.sin(fi),0.0])
a=posp(v,S,C,m,g)

# PREDVIDENA NAJVECJA VISINA, DOMET IN CAS TRAJANJA SIMULACIJE
t_fin=2.0*v0*np.sin(fi)/g
D=v0*np.cos(fi)*t_fin
h_max=v0*np.sin(fi)*t_fin*0.5-0.5*g*(t_fin*0.5)**2

print("cas leta = %.2f s\t"%(t_fin),"domet je %.2f m\t"%(D),"najvisja tocka leta = %.2f m"%(h_max))

data=open("rezultati_simulacije_upor.dat","w+")

t=0.0
i=0
ii=0
while pos[1]>-0.0001:
    pos+=v*dt
    v+=a*dt
    a=posp(v,S,C,m,g)
    t+=dt
    data.write("%f\t%f\t%f\t%f\t%f\n"%(t,pos[0],pos[1],np.sqrt(v.dot(v)),np.sqrt(a.dot(a))))
    ii+=1
print ("SIMULACIJA JE KONCANA")
data.close()
narisi_graf()

Vnesite gostoto telessa:500
Vnesite maso telesa0.1
cas leta = 13.09 s	 domet je 833.19 m	 najvisja tocka leta = 208.30 m
SIMULACIJA JE KONCANA