Expand index notation equation using sympy

Question

Below I have an equation written using index notation. This equation can be expressed with the six equations on the figure.

The first equation is expanded using index notation (einstein notation: https://en.wikipedia.org/wiki/Einstein_notation). In U_k,k the comma is a convention for derivative. Since we have repeated indices (k,k) we apply the summation convention and get (du_1/dx_1 + du_2/dx_2 + du_3/dx_3). On the figure the terms u_1, u_2 and u_3 are written as u, v and w and they are differentiated by x_1, x_2 and x_3 (x, y, and z).

I am new to python and Sympy and I saw that there is a tensor module but I could not see if there is something already implemented in Sympy where I can write the first equation (with index) and from that obtain the other six relations.

Answer 1

Currently, there is no straightforward way to do what you ask.

I shall suggest a trick that should work, by using IndexedBase (that is, symbols indexed by other symbols), followed by a replacement.

Declare your symbols:

U = IndexedBase("U")
l = symbols("lambda")
var("mu u v w x y z i j k")

Declare your expressions (no Einstein summation, so explicitly put a Sum):

sij = l*Sum(U[k, k], (k, 0, 2)) * KroneckerDelta(i, j) + mu*(U[i, j] + U[j, i])

Define a replacement function: will match U[0, 0] to Derivative(u, x) and so on according to the index:

def replace_U(uij):
    i1, i2 = uij.indices
    return Derivative([u, v, w][i1], [x, y, z][i2])

Now, run over all i, j indices replacing U[i, j] with integer values first, then replacing expressions like U[0, 0], ...

for ii in range(3):
    for ji in range(3):
        if ii < ji:
            continue
        pprint(sij.doit().xreplace({i: ii, j: ji}).replace(lambda v: v.base == U, replace_U))

Remember: Sum.doit() expands the summation. Condition ii >= ji is used to avoid printing the same expressions twice (your equation is symmetric in i, j).

Note that U[i, j] in the code above is just a symbol, i and j have no special meaning except as indices to U. The replacement function assigns derivatives to them.

The output I get is:

  ⎛d       d       d    ⎞       d    
λ⋅⎜──(u) + ──(v) + ──(w)⎟ + 2⋅μ⋅──(u)
  ⎝dx      dy      dz   ⎠       dx   
  ⎛d       d    ⎞
μ⋅⎜──(u) + ──(v)⎟
  ⎝dy      dx   ⎠
  ⎛d       d       d    ⎞       d    
λ⋅⎜──(u) + ──(v) + ──(w)⎟ + 2⋅μ⋅──(v)
  ⎝dx      dy      dz   ⎠       dy   
  ⎛d       d    ⎞
μ⋅⎜──(u) + ──(w)⎟
  ⎝dz      dx   ⎠
  ⎛d       d    ⎞
μ⋅⎜──(v) + ──(w)⎟
  ⎝dz      dy   ⎠
  ⎛d       d       d    ⎞       d    
λ⋅⎜──(u) + ──(v) + ──(w)⎟ + 2⋅μ⋅──(w)
  ⎝dx      dy      dz   ⎠       dz   

I am new to python and Sympy and I saw that there is a tensor module but I could not see if there is something already implemented in Sympy where I can write the first equation (with index) and from that obtain the other six relations.

There is sympy.tensor.tensor, which supports abstract index notation (a restricted kind of Einstein summation). Unfortunately, it does not support derivatives.

If you want to work on components, there's a recent addition: sympy.tensor.array. It provides multidimensional arrays, tensor products and contractions, and derivatives by arrays.

Answer 2

As Francesco Bonazzi stated, there is another way to implement equations using sympy.tensor.array. This method has the advantage of being easy to implement derivatives.

The sample code (for jupyter) is as follows:

from sympy import *
lam, mu, x, y, z = symbols("lambda mu x y z")
u, v, w = symbols("u v w", cls=Function)
du = derive_by_array([u(x, y, z), v(x, y, z), w(x, y, z)], [x, y, z])
sig = lam * tensorcontraction(du, (0, 1)) * Array(eye(3)) + mu * (du + transpose(du))
for i in range(3):
    for j in range(i, 3):
        display(Eq(Symbol("sigma_{}{}".format(i+1, j+1)), sig[i, j]))

Output: