Sympy - переименовать часть выражения

Question

Скажите, что я определил следующее выражение:

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)

Переменная expr теперь отображается так:

Хотя это хорошо для этого минимального примера, в больших выражениях он выглядит довольно грязно.Это действительно мешает мне видеть, что происходит позже, когда я вычисляю суммы по всем r(i,j), производным и т. Д.

Мой вопрос : есть ли способ сообщить SymPy о r(i, j)такой, что его можно отобразить примерно так:

, продолжая вести себя так же, как и раньше, в последующих выражениях?

Я знаю, что мог бы сделать r a Function, который выглядел бы как хотел, но тогда он не работал бы должным образом в последующемрасчет (например, производные будут абстрактными и не оценены).

Любая помощь будет высоко ценится!

Answer 1

Вы можете создать пользовательский подкласс функций, который не оценивается по умолчанию:

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 указывает, когда выполнять оценку.Поскольку он всегда возвращает None, он никогда не оценивает.Он также сообщает SymPy, что функция имеет два аргумента.Вы также можете получить явные значения в некоторых случаях, если хотите.Например, вы можете захотеть, чтобы он вычислял, являются ли i и j явными числами.

@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)))

При этом вы можете использовать его по своему усмотрению и вызывать expr.doit(), когда хотите, чтобы он вычислялся,Вы также можете специально определить оценку для определенных функций, чтобы избежать doit.Например, производные:

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

Это приведет к немедленной оценке r(i, j).diff(i) без необходимости вызова doit.

Другие функции имеют аналогичные методы, которые вы можете определить.Смотрите документацию SymPy.

Answer 2

Я действительно не знаю, может ли это помочь вам, но как на счет этого:

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]))

Вы можете отобразить свое выражение по своему желанию, затем использовать lambdify() иведет себя как следует.Просто догадка, может быть, это бесполезно для вас, поскольку вы, вероятно, предпочитаете способ поддерживать одно и то же выражение во всем коде.