在 sympy 中扩展表达式树,将某些二项式保持在一起

Expanding expression trees in sympy, keeping certain binomials together

我在 sympy 中处理多元多项式,它们都表示为 (x_i + y_j) 形式的项的乘积之和,其中 i 和 j 是索引,我想要保持这种方式,即用一个 x 符号和一个 y 符号的总和来表达所有内容。

比如我要

(y_{1} + z_{2})*((0 + 1)*(y_{3} + z_{2}) + y_{1} + z_{1} + 0 + 0)

成为

(y_{1} + z_{2})*(y_{3} + z_{2}) + (y_{1} + z_{2})*(y_{1} + z_{1})

您可以做的第一件事是用虚拟替换符合模式的二项式并展开。问题是您将有一些悬而未决的术语组合在一起。如果你总是有两个,那就很容易了。如果你有更多,它将需要更多的工作和一个很好的定义,即哪个子索引 x 应该与哪个 y (或你想要配对的任何字母)一起使用。

所以让我们从你的未计算表达式开始,我们将其称为 u

获取免费符号(假设只有 xy 感兴趣):

>>> free = u.free_symbols

用唯一的虚拟变量替换现有的二项式

>>> reps = {}
>>> u.replace(lambda x:x.is_Add and len(x.args) == 2 and all(
...              i in free for i in x.args),
...          lambda x: reps.setdefault(x, Dummy()))
_Dummy_45*(_Dummy_47*(0 + 1) + y1 + z1)

现在展开

>>> expand(_)
_Dummy_45*_Dummy_47 + _Dummy_45*y1 + _Dummy_45*z1

并收集虚拟符号的产品

>>> _.replace(lambda x:x.is_Mul and len(x.args) == 2 and all(
...              i in reps.values() for i in x.args),
...    lambda x: reps.setdefault(x, Dummy())))
_Dummy_45*y1 + _Dummy_45*z1 + _Dummy_51

收集虚拟符号以使之前悬而未决的二项式出现

>>> collect(_, reps.values())
_Dummy_45*(y1 + z1) + _Dummy_51

现在用它们的值替换虚拟符号(这是 reps 中的键,所以我们必须反转字典):

>>> _.xreplace({v:k for k,v in reps.items()})
_Dummy_45*_Dummy_47 + (y1 + z1)*(y1 + z2)

再做一次

>>> _.xreplace({v:k for k,v in reps.items()})
(y1 + z1)*(y1 + z2) + (y1 + z2)*(y3 + z2)

以某种方式发布您希望看到的特定表达式 re-arranged 将有助于集中精力制定更可靠的解决方案,但这些技巧可以帮助您入门。这里也有一个函数,它在 Add 中将自由符号配对并用虚拟符号替换它们。

def collect_pairs(e, X, Y):
    free = e.free_symbols
    xvars, yvars = [[i for i in free if i.name.startswith(j)] for j in (X, Y)]
    reps = {}
    def do(e):
        if not e.is_Add: return e
        x, cy = sift(e.args, lambda x: x in xvars, binary=True)
        y, c = sift(cy, lambda x: x in yvars, binary=True)
        if x and len(x) != len(y): return e
        args = []
        for i,j in zip(ordered(x), ordered(y)):
            args.append(reps.setdefault(i+j, Dummy()))
        return Add(*(c + args))
    # hmmm...this destroys the reps and returns {}
    #return {v:k for k,v in reps.items()}, bottom_up(e, do)
    return reps, bottom_up(e, do)


>>> e1
(y1 + z2)*(y1 + y3 + z1 + z2)
>>> r, e = collect_pairs(e1,'y','z')
>>> expand(e).xreplace({v:k for k,v in r.items()})
(y1 + z1)*(y1 + z2) + (y1 + z2)*(y3 + z2)

这适用于完全展开的 e1 如果您首先考虑它:

>>> e2 = factor(expand(e1)); e2
(y1 + z2)*(y1 + y3 + z1 + z2)
>>> r, e = collect_pairs(e2, 'y', 'z')
>>> expand(e).xreplace({v:k for k,v in r.items()})
(y1 + z1)*(y1 + z2) + (y1 + z2)*(y3 + z2)

查看您最初发布的代码,我建议将二项式放在一起,只在最后替换它们,如下所示:

...
def single_variable_diff(perm_dict,k,m):
    ret_dict = {}
    for perm,val in perm_dict.items():
        if len(perm)<k:
            ret_dict[perm] = Add(ret_dict.get(perm,0), reps.setdefault(U(var2[k],var3[m]), Dummy())*val,evaluate=False)
        else:
            ret_dict[perm] = Add(ret_dict.get(perm,0), reps.setdefault(U(var2[perm[k-1]],var3[m]), Dummy())*val,evaluate=False)
...

reps = {}
U = lambda x,y: UnevaluatedExpr(Add(*ordered((x,y))))
ireps = lambda: {v:k for k,v in reps.items()}

perms=[]
curperm = []

...

coeff_perms.sort(key=lambda x: (inv(x),*x))

def touch(e):
    from sympy.core.traversal import bottom_up
    def do(e):
        return e if not e.args else e.func(*e.args)
    return bottom_up(e, do)

undo = ireps()
for perm in coeff_perms:
    val = touch(coeff_dict[perm]).expand().xreplace(undo))
    print(f"{str(perm):>{width}} {str(val)}")

(3, 4, 1, 2) 将以二项式乘积的形式给出,但有些元素不会 - 它们只是二项式的总和。为了将它们放在一起,您可以将它们创建为 UnevaluatedExpr,例如定义的 U lambda。我猜你不必使用 evaluated=False 并且不需要 touch 函数。