Sympy：删除没有特定变量的项

Question

我试图计算一些（很多）多元条件密度（即，几个多元概率密度函数的乘积）。我能够正确地设置和扩展矩阵，但现在我想删除一些项，例如在下面的等式（和代码）中，这些项不包含wg。在贴出的答案的帮助下，我能够开发出一个黑客解决方案；欢迎改进

更新：MWE

import sympy as sym
from IPython.display import display as disp
N  = 211
wg = sym.MatrixSymbol('w_g', N, 1)
wg_n = sym.MatrixSymbol('w_gn', N, 1)
Z_wg = sym.MatrixSymbol('Z_wg', N, N)

# pdf wg
pdf_wg = ((wg - wg_n).T * Z_wg.I * (wg - wg_n))
pdf_full = sym.expand(pdf_wg)
# pdf_full.collect(wg) # NotImplementedError: noncommutative scalars in MatMul are not supported.


# print (wg in pdf_full.atoms()) # False

# this gives what I want
terms = pdf_full.as_terms()[0]
for term in terms:
    if 'w_g,' in str(term[0].atoms()):
        disp (term[0])

更新2：更复杂的MWE

在这里，我试着抓住其中包含b的术语

import sympy as sym
from IPython.display import display as disp, Math

mu   = sym.symbols('mu')               # mean non GIA SSH trend
N    = 211
vec1 = sym.MatrixSymbol('1', N, 1)
u    = sym.MatrixSymbol('u', N, 1)     
Pi   = sym.MatrixSymbol('Pi', N, N)    
b    = sym.MatrixSymbol('b', N, 1)

wg = sym.MatrixSymbol('w_g', N, 1)     
wm = sym.MatrixSymbol('w_m', N, 1)     
bhat = mu*vec1 + wg + wm + u # convenience

pdf  = sym.expand((b - bhat).T * Pi.I * (b-bhat))
terms      = pdf.as_terms()[0]
good_terms = [] 
for term in terms:
    if b.args[0] in term[0].atoms():
        good_terms.append(term[0])

print ('Good terms:'); disp(sym.Add(*good_terms))

更新4：已解决 对于更复杂的表达式，将doit()添加到展开将防止出现大量额外的循环（例如）：

pdf  = sym.expand((b - bhat).T * Pi.I * (b-bhat)).doit()

更多信息可在对各种答案的评论中找到

谢谢

Answer 1

您可以获得不包含wg的术语，如：

In [53]: pdf_full.subs(wg, ZeroMatrix(N, 1)).doit()
Out[53]: 
    T     -1     
w_gn ⋅Z_wg  ⋅w_gn

然后可以从pdf_full中减去这些值：

In [54]: pdf_full - pdf_full.subs(wg, ZeroMatrix(N, 1)).doit()
Out[54]: 
   T     -1         T     -1           T     -1    
w_g ⋅Z_wg  ⋅w_g -w_g ⋅Z_wg  ⋅w_gn -w_gn ⋅Z_wg  ⋅w_g

Answer 2

您可以提取表达式的^{}并测试变量是否在其中：

from sympy import symbols

a, b, mug = symbols('a b mu_g')
expr1 = a * b + a * mug
expr2 = a * b
for expr in [expr1, expr2]:
    if mug in expr.atoms():
        print(expr, 'contains', mug)
    else:
        print(expr, 'does not contain', mug)

附：新问题的更新。对于MatrixSymbol符号存储为wg.args[0]（args[1]和args[2]是维度）：

import sympy as sym

N  = 211
wg = sym.MatrixSymbol('w_g', N, 1)
wg_n = sym.MatrixSymbol('w_gn', N, 1)
Z_wg = sym.MatrixSymbol('Z_wg', N, N)

pdf_wg = ((wg - wg_n).T * Z_wg.I * (wg - wg_n))
pdf_full = sym.expand(pdf_wg)

print (wg.args[0] in pdf_full.atoms()) # True

请注意，哈克的解决方案是，当w_g是最后一项或另一个名称以相同字符串结尾时，问题可能会出错