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一个Function显示为所需的，但是随后在随后的计算中将无法正常工作（例如，派生将是抽象的，而不进行评估）。

任何帮助将非常感激！

Answer 1

您可以创建一个默认情况下不计算的自定义 Function 子类：

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 文档。