如何用Sympy的EPath匹配表达式中的所有Sum节点？

问题场景 #

我了解到SymPy的EPath和XPath类似，但不确定具体用法。想要匹配表达式中所有的Sum节点，示例表达式如下：

import sympy
x = sympy.IndexedBase('x', real=True)
i = sympy.Symbol('i', integer=True, positive=True, finite=True)
n = sympy.Symbol('n', integer=True, positive=True, finite=True)
expr = sympy.sin(sympy.Sum(2*x[i], (i, 1, n))) ** 2 + sympy.Sum(x[i], (i, 1, n))

尝试用sympy.EPath(r"//Sum")时会报错ValueError: empty selector，只能通过指定层级（比如/*/Sum、/*/*/Sum）匹配部分节点，目前已经实现了一个遍历表达式树的替代方案。

正确的EPath解决方案 #

SymPy的EPath语法与XPath存在细微差异：要匹配任意深度的所有后代节点，需要使用**而非XPath中的//。因此匹配所有Sum节点的正确写法是**/Sum。

匹配所有Sum节点 #

from sympy import EPath

# 定义匹配任意层级Sum节点的EPath
path = EPath(r"**/Sum")
# 查找所有匹配的节点
all_sums = list(path.find(expr))
print(all_sums)
# 输出: [Sum(2*x[i], (i, 1, n)), Sum(x[i], (i, 1, n))]

批量修改匹配的Sum节点 #

如果需要对所有匹配的Sum节点执行修改操作，可以使用EPath.subs方法，无需手动遍历：

# 定义修改函数：将Sum节点乘以2
def modify_sum(sum_node):
    return 2 * sum_node

# 替换所有匹配的Sum节点
modified_expr = path.subs(modify_sum, expr)
print(modified_expr)
# 输出: sin(2*Sum(2*x[i], (i, 1, n)))**2 + 2*Sum(x[i], (i, 1, n))

关于自定义遍历函数 #

你实现的traverse_apply函数可以正常完成遍历修改的需求，但EPath是SymPy官方提供的更简洁的工具，适合快速定位和修改表达式树节点。如果需要更复杂的遍历逻辑（比如多节点类型匹配、自定义过滤规则），自定义遍历函数会更灵活。

小优化建议：将函数中的type(expr) == atype替换为isinstance(expr, atype)，这样可以支持子类节点的匹配（若存在Sum的子类），优化后的完整函数如下：

def traverse_apply(expr, atype, func):
    """
    Apply the function `func` to all the element in the expression `expr` of type `atype`.
    Return a new expression.
    """
    if isinstance(expr, atype):
        return func(expr)
    if isinstance(expr, sympy.Basic):
        if not expr.is_Atom:
            args, basic = expr.args, True
        else:
            return expr
    elif hasattr(expr, '__iter__'):
        args, basic = expr, False
    else:
        return expr
    
    args = list(args)
    indices = range(len(args))
    
    for i in indices:
        arg = args[i]
        args[i] = traverse_apply(arg, atype, func)

    if basic:
        return expr.func(*args)
    else:
        return expr.__class__(args)

内容的提问来源于stack exchange，提问作者Ruggero Turra