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

如何解决Sympy：删除没有特定变量的术语

我正在尝试计算一些（很多）多元条件密度（即几个多元概率密度函数的乘法）。我能够正确设置和扩展矩阵，但现在想删除不包含 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',1)
Z_wg = sym.MatrixSymbol('Z_wg',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',1)
u    = sym.MatrixSymbol('u',1)     
Pi   = sym.MatrixSymbol('Pi',N)    
b    = sym.MatrixSymbol('b',1)

wg = sym.MatrixSymbol('w_g',1)     
wm = sym.MatrixSymbol('w_m',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() 添加到 expand 将防止一堆额外的循环（例如）：

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

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

谢谢！

解决方法

您可以提取表达式的 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',1)
Z_wg = sym.MatrixSymbol('Z_wg',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

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

,

您可以获得不包含 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,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