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