Sympy - Rename part of an expression

Question

Say I have defined the following expression:

from sympy import *
N, D, i, j, d = symbols("N D i j d", integer=True)
beta, gamma = symbols(r'\beta \gamma')
X = IndexedBase("X", shape=(N, D))

# r(i, j) = euclidian distance between X[i] and X[j]
r = lambda i, j: sqrt(Sum((X[i, d] - X[j, d])**2, (d, 1, D)))
expr = r(i, j)**2 + r(i, j)

The expr variable now displays like this:

While this is fine for this minimal example, it gets quite messy in larger expressions. This really hinders my ability to see what happens later on when I compute sums over all r(i,j), derivatives etc.

My question: Is there a way to tell SymPy about r(i, j), such that it can be displayed as something like this:

while still behaving as before in subsequent expressions?

I know I could make r a Function, which would display as wanted, but then it would not work properly in a subsequent calculation (e.g. derivatives would be abstract and not evaluated).

Any help would be much appreciated!

Answer 1

You can make a custom Function subclass that doesn't evaluate by default:

class r(Function):
    @classmethod
    def eval(cls, i, j):
        return

    def doit(self, **kwargs):
        i, j = self.args
        return sqrt(Sum((X[i, d] - X[j, d])**2, (d, 1, D)))

eval tells it when to evaluate. Since it always returns None, it never evaluates. It also tells SymPy the function has two arguments. You can also have it return explicit values in some cases, if you like. For instance, you might want it to evaluate if i and j are explicit numbers.

@classmethod
def eval(cls, i, j):
    if i.is_Number and j.is_Number:
        return sqrt(Sum((X[i, d] - X[j, d])**2, (d, 1, D)))

With this you can use it as desired, and call expr.doit() when you want it to evaluate. You can also specifically define evaluation for certain functions to avoid doit. For example, derivatives:

def _eval_derivative(self, x):
    return self.doit()._eval_derivative(x)

This will make r(i, j).diff(i) evaluate immediately without having to call doit.

Other functions have similar methods you can define. See the SymPy documentation.

Answer 2

I don't really know if it may helps you but what about this:

from sympy import *
from sympy.utilities.lambdify import lambdify, implemented_function
N, D, i, j, d = symbols("N D i j d", integer=True)
beta, gamma = symbols(r'\beta \gamma')
X = IndexedBase("X", shape=(N, D))

r = implemented_function('r', lambda i, j: sqrt(Sum((X[i, d] - X[j, d])**2, (d, 1, D))));
expr = r(i, j)**2 + r(i, j)
print(expr)
r = lambdify((i,j), r(i,j))
print(diff(r(i,j), X[i,j]))

You can display your expression as you wish, then use lambdify() and makes it behave as it should. Just guessing, maybe it's useless for you as you probably prefer a way to maintain the same expression all along the code.