import numpy as np
import matplotlib.pyplot as plt

"""
Wahadło matematyczne metodą Eulera-Cromera
"""

g=9.81       #Przyspieszenie ziemskie
l=50         #Długość wahadła
alfa=g/l  
r0=np.pi-0.1 #Początkowe położenie (kąt)
v0=0         #Początkowa prędkość (kątowa)
m=1          #Masa wahadła
dt=0.01      #Krok czasowy
N=20000

def F(x):
    return -alfa*np.sin(x)

def euler(r,v,dt,F):
    rnew=r+dt*v
    vnew=v+dt/m*F(rnew)
    #^W zwykłym Eulerze byłoby F(r)
    return rnew,vnew

def tor(r0,v0,N,dt,F):
    r=np.empty(N+1)
    v=np.empty(N+1)
    #^Można w zasadzie np.zeros, ale empty
    #marginalnie wydajniejsze - tylko alokuje
    #przestrzeń, nie zeruje
    r[0]=r0; v[0]=v0
    for i in range(1,N+1):
        r[i],v[i]=euler(r[i-1],v[i-1],dt,F)
    return r,v

def energia(r,v):
    T=m*l**2*v**2/2
    V=m*g*l*(1-np.cos(r))
    return T+V

t=np.linspace(0,dt*N,N+1)
r,v=tor(r0,v0,N,dt,F)
E=energia(r,v)
#Energia - pomijając małe oscylacje, jest stabilna.
#W zwykłym Eulerze niefizycznie systematycznie rośnie.

fig,(ax1,ax2,ax3) = plt.subplots(3,1)

ax1.plot(t,r)
ax1.set_title("Kąt")

ax2.plot(t,v)
ax2.set_title("Prędkość kątowa")

ax3.plot(t,E)
ax3.set_ylim(0,1.1*E.max())
ax3.set_title("Energia")
plt.tight_layout()
plt.show()