Sympy：删除没有特定变量的术语

Question

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

更新：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()

更多信息可以在各种答案的 cmets 中找到。

谢谢！

Answer 1

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

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)

PS：您的新问题的更新。对于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 是最后一项或另一个名称以相同的字符串结尾时，问题可能会出错。

Answer 2

你可以得到不包含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