'''

Clément de la Salle - ENS de Lyon - M2 Physique Concepts et applications - 2020-2021

***

Rentrer une métrique dans la variable metrique, lancer le programme, afin de calculer les tenseurs de Christoffel, Riemann et Ricci.
Appeler la grandeur souhaitée parmis la liste ci-dessous (l'affiche est joli grâce à la fonction init_printing() de sympy).

Grandeurs calculées :

- Gamma[mu][alpha, beta] = Gamma^mu_{alpha, beta} :
    Contient les coefficients de Christoffel
    Il s'agit d'une liste de matrices (matrices de symboles sympy)
    Le premier indice est l'indice du haut (mu)
    Les deux suivants sont donnés dans le même crochet et désignent les indices du dessous (alpha, beta)

- Riemann[alpha][beta][mu, nu] = R^alpha_{beta, mu, nu} :
    Contient les coefficients du tenseur de Riemann

- Ricci[mu, nu] = R_{mu, nu} :
    Contient les coefficients du tenseur de Ricci

- R :
    Scalaire de Ricci

- Weyl[a][b][c, d] = Weyl_{a, b, c, d} :
    Contient les coefficients du tenseur de Weyl

- Equations_geodesiques :
    Liste contenant toutes les équations des géodésiques.
    Le temps propre est noté s.

Pour modifier la métrique :
- Renommer les variables commes bon vous semble en modifiant la lise variables canoniques
- Laisser les matrices Matrix de sympy qui permettent le calcul formel (pas de numpy !)
  Il y a un désavantage à cela : on ne peut utiliser que des matrices 2D, d'où l'utilisation de listes de matrices pour les gros tenseurs à 3 ou 4 indices
- Modifier la métrique (variable metrique) sous la forme d'une matrice carrée symétrique
- Ne pas utiliser la variable s qui est réservée pour le temps propre !

On peut utiliser pleinement le calcul formel offert par sympy :

- Pour introduire des paramètres dans la métrique, les ajouter avec une ligne de type
    a = Symbol('a')
  Puis utiliser ce paramètre comme un scalaire
  
  Ex: variables_canoniques = ['t', 'x']
      a = Symbol('a')
      metrique = Matrix([[-a**2 * x**2, 0],
                         [           0, 1]])
    
  Les calculs seront alors tous exprimés en fonction de a de façon analytique !

- On peut faire encore mieux en définissant des fonctions et laisser ainsi encore plus de liberté !
  Pour créer des fonctions quelconque, ajouter une ligne de type
    f, g = symbols('f g', cls = Function)
  Précisez de quelles variables ces fonctions dépendent, au moment de leur appel
    f(x,y)
    
  Ex: variables_canoniques = ['u', 'v', 'x', 'y']
      H = symbols('H', cls = Function)
        metrique = Matrix([[H(u, x, y), 1, 0, 0],
                           [         1, 0, 0, 0],
                           [         0, 0, 1, 0],
                           [         0, 0, 0, 1]])
                           
    Dans cet exemple, H ne dépend par exemple par de la variable v et cela est pris en compte : partial_v H = 0

  Les résultats feront alors intervenir les fonctions introduites et leurs dérivées.

'''

from sympy import *
init_printing()

variables_canoniques = ['t', 'x', 'y', 'z']
for var in variables_canoniques :
    exec(var + " = Symbol('" + var + "')")

# a = Symbol('a')                       # C'est le genre de ligne qui sert à introduire des paramètres quelconques dans la métrique (voir en-tête pour plus de détails)
# f, g = symbols('f g', cls = Function) # C'est le genre de ligne qui sert à introduire des fonctions quelconques dans la métrique (voir en-tête pour plus de détails)

# Définition de la métrique (exemple de la métrique de Minkowski)

metrique = Matrix([[-1, 0, 0, 0],
                   [ 0, 1, 0, 0],
                   [ 0, 0, 1, 0],
                   [ 0, 0, 0, 1]])
                   
# Exemple de la métrique pour la 2-sphère (décommenter ces lignes pour faire le test) :

# variables_canoniques = ['theta', 'phi']
# for var in variables_canoniques :
#     exec(var + " = Symbol('" + var + "')")
# metrique = Matrix([[1,     0        ],
#                    [0, sin(theta)**2]])
                     
# Calcul de la métrique inverse
metrique_inverse = metrique**(-1)

N = metrique.shape[0]                               # Dimension de l'espace-temps

metrique_derivee = [zeros(N, N) for k in range(N)]  # Calcul des dérivées de la métrique : metrique_derivee[rho][mu, nu] = partial_rho g_{mu, nu}
for rho in range(N) :
    for mu in range(N) :
        for nu in range(N) :
            metrique_derivee[rho][mu, nu] = metrique[mu, nu].diff(variables_canoniques[rho])

Gamma = [zeros(N, N) for k in range(N)]             # Calcul des coefficients de Christoffel (voir en-tête pour les conventions sur les indices)
for mu in range(N) :
    for alpha in range(N) :
        for beta in range(N) :
            for nu in range(N) :
                Gamma[mu][alpha, beta] += 1/2 * metrique_inverse[mu, nu] * (metrique_derivee[alpha][beta, nu] + metrique_derivee[beta][alpha, nu] - metrique_derivee[nu][alpha, beta])

Riemann = [[zeros(N, N) for k in range(N)] for i in range(N)]   # Calcul du tenseur de Riemann (voir en-tête pour les conventions sur les indices)
for alpha in range(N) :
    for beta in range(N) :
        for mu in range(N) :
            for nu in range(N) :
                Riemann[alpha][beta][mu, nu] = Gamma[alpha][beta, nu].diff(variables_canoniques[mu]) - Gamma[alpha][beta, mu].diff(variables_canoniques[nu])
                for kappa in range(N) :
                    Riemann[alpha][beta][mu, nu] += Gamma[alpha][kappa, mu] * Gamma[kappa][beta, nu] - Gamma[alpha][kappa, nu] * Gamma[kappa][beta, mu]

Ricci = zeros(N, N)             # Calcul du tenseur de Ricci (voir en-tête pour les conventions sur les indices)
for mu in range(N) :
    for nu in range(N) :
        for rho in range(N) :
            Ricci[mu, nu] += Riemann[rho][mu][rho, nu]

R = 0                           # Calcul du scalaire de Ricci
for mu in range(N) :
    for nu in range(N) :
        R += Ricci[mu, nu] * metrique_inverse[mu, nu]

Weyl = [[zeros(N, N) for k in range(N)] for i in range(N)]
for a in range(N) :
    for b in range(N) :
        for c in range(N) :
            for d in range(N) :
                somme = 0
                for e in range(N) :
                    somme += Riemann[e][b][c, d] * metrique_inverse[a, e]
                Weyl[a][b][c, d] = somme - 1/2 * (Ricci[a, c] * metrique[b, d] - Ricci[a, d] * metrique[b, c] + Ricci[b, d] * metrique[a, c] - Ricci[b, c] * metrique[a, d]) + 1/6 * (metrique[a, c] * metrique[b, d] - metrique[a, d] * metrique[b, c])

Equations_geodesiques = []
s = Symbol('s')             # C'est le temps propre
exec(','.join(variables_canoniques) + '= symbols("' + ' '.join(variables_canoniques) + '", cls = Function)')    # Les variables sont converties en fonction de s
for mu in range(N) :
    somme = 0
    for rho in range(N) :
        for nu in range(N) :
            exec('somme += Gamma[' + str(mu) + '][' + str(rho) + ',' + str(nu) + ']*' + variables_canoniques[rho] + '(s).diff(s)*' + variables_canoniques[nu] + '(s).diff(s)')
    exec('somme +=' + variables_canoniques[mu] + '(s).diff(s).diff(s)')
    Equations_geodesiques.append(Eq(somme, 0))