import scipy.

stats as st
from math import *
import numpy as np
import matplotlib.pyplot as plt

T=0.5
U_0=[1,1] #condition initiale
R2=(U_0[0]**2 + U_0[1]**2) #|u(t)|**2 = cste

## 1-ère partie : La solution exacte

def U_exacte(t): #solution du problème de Cauchy
return [U_0[0]*cos(R2*t)+U_0[1]*sin(R2*t) , U_0[1]*cos(R2*t)-U_0[0]*sin(R2*t)]

def sol_exacte(T,N): #liste des valeurs prisent par la solution aux temps tn
h=T/N
U=[U_0]
for i in range(1,N+1):
U.append(U_exacte(i*h))
return U

## 2-ème partie : la solution approchée

def f(U): #prend en argument : le temps, la partie réelle et la partie imaginaire
return np.array([R2*U[1] , -R2*U[0]]) #renvoie un vecteur de taille 2

def RK2(T,N): #solution approchée

h=T/N
U=[U_0]
for i in range(N):
U.append( U[i]+ h*f( U[i]+h/2*f(U[i]) ) )
return U

## 3-ième partie : le calcul de l'ordre du schéma

def erreur_max(T,N): #erreur maximale pour la partie réelle
M=[]
for i in range(1,N+1):
M.append(np.linalg.norm(np.array(sol_exacte(T,N)[i])-np.array(RK2(T,N)[i])))
return max(M)

def ordre_schema(T): #ordre du shéma pour la partie réelle

M=[]
L=[]
for i in range(20,200,20):
 M.append(erreur_max(T,i))
L.append(T/i)
return M,L

plt.plot(np.log(ordre_schema(T)[1]),np.log(ordre_schema(T)[0]))
plt.show()
print(st.linregress(np.log(ordre_schema(T)[1]),np.log(ordre_schema(T)[0]))[0])