sympy – Wanting feedback on Python code to calculate general relativity-related tensors/arrays using metric tensor as input

Please be gentle with me– I just began learning to code a few weeks ago as a hobby in order to support my other hobby of learning general relativity, so this is the first bit of code I’ve ever written. I would love to keep getting better at coding, though, so any feedback on ways to improve my code would be much appreciated.

What this code does is take as input the number of dimensions of a manifold, the coordinate labels being used, and the components of a metric, and outputs the non-zero components of the metric (exactly what was input, it just looks prettier), and also of the inverse metric, derivatives of the metric, the Christoffel symbols, the derivatives of the Christoffel symbols, the Riemann curvature tensor, the Ricci curvature tensor, the Ricci scalar, and the Einstein tensor (with 2 covariant indices, but also with 1 contravariant and 1 covariant).

For those of you who run the code, here are some useful tips on the user inputs (I will also include an example at the bottom): When inputting metric components, you can use ‘^’ instead of ‘**’ for exponents, and when multiplying a number by a symbol or something in parentheses, you don’t need to include a ‘*’. Feel free to include undefined functions- just make sure you include its arguments if you want it to be differentiated correctly (e.g. if you want a function that will have a non-zero derivative of x, then type ‘f(x)’ in your expression instead of just ‘f’). Also feel free to use greek letters (spelled out in English) when inputting coordinate labels and/or functions in your metric components.

Here is the code:

from sympy import *
from dataclasses import dataclass
from IPython.display import display as Idisplay
from IPython.display import Math

greek = ('alpha', 'beta', 'gamma', 'Gamma', 'delta', 'Delta', 'epsilon',
         'varepsilon', 'zeta', 'eta', 'theta', 'vartheta', 'Theta', 'iota',
         'kappa', 'lambda', 'Lambda', 'mu', 'nu', 'xi', 'Xi', 'pi', 'Pi',
         'rho', 'varrho', 'sigma', 'Sigma', 'tau', 'upsilon', 'Upsilon',
         'phi', 'varphi', 'Phi', 'chi', 'psi', 'Psi', 'omega', 'Omega')

n = int(input('Enter the number of dimensions:n'))
coords = ()
for i in range(n):
    coords.append(Symbol(str(input('Enter coordinate label %d:n' % i))))


@dataclass(frozen=False, order=True)
class Tensor:
    name: str
    symbol: str
    key: str
    components: list
    
    def rank(self):
        return self.key.count('*')
    
    def tensor_zeros(self, t=0):
        for i in range(self.rank()):
            t = (t,) * n
        return MutableDenseNDimArray(t)
    
    def coord_id(self, o):
        a = ()
        for i in range(self.rank()):
            c = int(o/(n**(self.rank() - i - 1)))
            a.append(str(coords(c)))
            if any(letter in a(i) for letter in greek) is True:
                a(i) = '\' + a(i) + ' '
            o -= c * (n**(self.rank() - i - 1))  
        x = self.key
        w = 0
        for i in x:
            if i == '*':
                x = x.replace('*', a(w), 1)
                w += 1
        return self.symbol + x
    
    def print_tensor(self):
        for o in range(len(self.components)):
            if self.components(o) != 0:
                Idisplay(Math(latex(Eq(Symbol(self.coord_id(o)),
                                       self.components(o)))))
        print('nn')


def assign(instance, thing):
    instance.components = thing.reshape(len(thing)).tolist()


def fix_input(expr):
    expr = expr.replace('^', '**')
    for i in range(len(expr)-1):
        if expr(i).isnumeric() and (expr(i+1).isalpha() or
                                    expr(i+1) == '('):
            expr = expr(:i+1) + '*' + expr(i+1:)
    return expr


metric = Tensor('metric tensor', 'g', '_**', ())
metric_inv = Tensor('inverse of metric tensor', 'g', '__**', ())
metric_d = Tensor('partial derivative of metric tensor', 'g', '_**,*', ())
Christoffel = Tensor('Christoffel symbol - 2nd kind', 'Gamma', '__*_**', ())
Christoffel_d = Tensor('partial derivative of Christoffel symbol',
                       'Gamma', '__*_**,*', ())
Riemann = Tensor('Riemann curvature tensor', 'R', '__*_***', ())
Ricci = Tensor('Ricci curvature tensor', 'R', '_**', ())
Einstein = Tensor('Einstein tensor', 'G', '_**', ())
Einstein_alt = Tensor('Einstein tensor', 'G', '__*_*', ())

# user inputs metric:
g = eye(n)
while True:
    diag = str(input('Is metric diagonal? y for yes, n for non')).lower()
    if diag == 'y':
        for i in range(n):
            g(i, i) = sympify(fix_input(str(input(
                'What is g_(%s%s)?n' % (str(coords(i)), str(coords(i))
                )))))
    else:
        for i in range(n):
            for j in range(i, n):
                g(i, j) = sympify(fix_input(str(input(
                    'What is g_(%s%s)?n' % (str(coords(i)), str(coords(j))
                    )))))
                g(j, i) = g(i, j)
    if g.det() == 0:
        print('nMetric is singular, try again!n')
        continue
    else:
        break

# calculate everything:
# inverse metric:
g_inv = MutableDenseNDimArray(g.inv())
assign(metric_inv, g_inv)
g = MutableDenseNDimArray(g)
assign(metric, g)
# first derivatives of metric components:
g_d = metric_d.tensor_zeros()
for i in range(n):
    for j in range(i):
        for d in range(n):
            g_d(i, j, d) = g_d(j, i, d)
    for j in range(i, n):
        for d in range(n):
            g_d(i, j, d) = diff(g(i, j), coords(d))
assign(metric_d, g_d)
# Christoffel symbols for Levi-Civita connection (Gam^i_jk):
Gamma = Christoffel.tensor_zeros()
for i in range(n):
    for j in range(n):
        for k in range(j):
            Gamma(i, j, k) = Gamma(i, k, j)
        for k in range(j, n):
            for l in range(n):
                Gamma(i, j, k) += S(1)/2 * g_inv(i, l) * (
                    -g_d(j, k, l) + g_d(k, l, j) + g_d(l, j, k)
                    )
assign(Christoffel, Gamma)
# first derivatives of Christoffel symbols (Gam^i_jk,d):
Gamma_d = Christoffel_d.tensor_zeros()
for i in range(n):
    for j in range(n):
        for k in range(j):
            for d in range(n):
                Gamma_d(i, j, k, d) = Gamma_d(i, k, j, d)
        for k in range(j, n):
            for d in range(n):
                Gamma_d(i, j, k, d) = simplify(diff(Gamma(i, j, k),
                                                    coords(d)))
assign(Christoffel_d, Gamma_d)
# Riemann curvature tensor (R^i_jkl):
Rie = Riemann.tensor_zeros()
for i in range(n):
    for j in range(n):
        for k in range(n):
            for l in range(k):
                Rie(i, j, k, l) = -Rie(i, j, l, k)
            for l in range(k, n):
                Rie(i, j, k, l) = Gamma_d(i, j, l, k) - Gamma_d(i, j, k, l)
                for h in range(n):
                    Rie(i, j, k, l) += (Gamma(h, j, l) * Gamma(i, h, k)
                                    - Gamma(h, j, k) * Gamma(i, h, l))
                    Rie(i, j, k, l) = simplify(Rie(i, j, k, l))
assign(Riemann, Rie)
# Ricci curvature tensor (R_jl):
Ric = simplify(tensorcontraction(Rie, (0, 2)))
assign(Ricci, Ric)
# Ricci curvature scalar:
R = 0
for i in range(n):
    for j in range(n):
        R += g_inv(i, j) * Ric(i, j)
R = simplify(R)
# Einstein tensor (G_ij):
G = Einstein.tensor_zeros()
for i in range(n):
    for j in range(i):
        G(i, j) = G(j, i)
    for j in range(i, n):
        G(i, j) = simplify(Ric(i, j) - S(1)/2 * R * g(i, j))
assign(Einstein, G)
# G^i_j:
G_alt = Einstein_alt.tensor_zeros()
for i in range(n):
    for j in range(n):
        for k in range(n):
            G_alt(i, j) += g_inv(i, k) * G(k, j)
        G_alt(i, j) = simplify(G_alt(i, j))
assign(Einstein_alt, G_alt)

# print it all
print()
metric.print_tensor()
metric_inv.print_tensor()
metric_d.print_tensor()
Christoffel.print_tensor()
Christoffel_d.print_tensor()
Riemann.print_tensor()
Ricci.print_tensor()
if R != 0:
    Idisplay(Math(latex(Eq(Symbol('R'), R))))
    print('nn')
Einstein.print_tensor()
Einstein_alt.print_tensor()

Example user-inputs