SymPy中替换IndexedBase对象：用联合分布q替换熵表达式中的p

问题描述 #

定义了基于离散分布p（IndexedBase类型）的熵表达式：

import sympy as sym

x = sym.Idx("x")
p = sym.IndexedBase("p")
N_X = sym.Symbol("N_X", positive=True)

h = -sym.summation(p[x] * sym.log(p[x]), (x, 1, N_X))

尝试将p[x]替换为联合分布q[x,y]对y的求和时，执行以下代码触发TypeError：

y = sym.Idx("y")
N_Y = sym.Symbol("N_Y", positive=True)
q = sym.IndexedBase("q")
h.subs(p, sym.summation(q[x,y], (y, 1, N_Y)))

错误信息：

TypeError: 
The base can only be replaced with a string, Symbol, IndexedBase or an
object with a method for getting items (i.e. an object with a
`__getitem__` method).

解决方案 #

不能直接替换整个IndexedBase对象p，而是要替换具体的索引项p[x]。因为SymPy要求替换IndexedBase时，替换对象必须是同类类型（如IndexedBase、Symbol等），而我们实际需要替换的是p[x]这个索引后的表达式，直接针对p[x]做替换即可：

import sympy as sym

# 定义原熵表达式
x = sym.Idx("x")
p = sym.IndexedBase("p")
N_X = sym.Symbol("N_X", positive=True)
h = -sym.summation(p[x] * sym.log(p[x]), (x, 1, N_X))

# 定义联合分布相关符号
y = sym.Idx("y")
N_Y = sym.Symbol("N_Y", positive=True)
q = sym.IndexedBase("q")

# 正确替换：针对p[x]进行替换
h_substituted = h.subs(p[x], sym.summation(q[x, y], (y, 1, N_Y)))
print(h_substituted)

原理说明 #

• 原错误的原因是试图用一个求和表达式（非IndexedBase类型）替换IndexedBase对象p，不符合SymPy的替换规则。
• 直接替换p[x]时，SymPy会识别表达式中所有p[x]的实例，将其替换为对应的求和式，同时保留外层对x的求和逻辑，不会出现索引冲突。

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