SymPy中如何定义符号大小的符号数组并完成符号求导？

需求与概念代码 #

要对「数据集每个元素应用函数后的求和式」做符号求导，概念代码如下：

from sympy import Symbol, sum, diff

a = Symbol('a')
n = Symbol('n')
k = Symbol('k')

# 此处data的定义是核心问题
cost = sum(a*data[k] + k, (k, 0, n-1))
diff(cost, a)

期望得到的符号化结果：

sum(data[k], (k, 0, n-1))

核心问题 #

无法正确指定data为由符号n定义大小的数组：

• 尝试用函数实现data时，求导会触发链式法则，不符合需求；
• 若指定固定大小的数组，求和式会被完全展开（当n=1000时结果毫无意义）；
• 实际场景中涉及的函数无法简单分解，必须保留符号化的求和形式。

解决方案 #

使用SymPy的IndexedBase定义符号化数组，它支持用符号下标和符号长度的求和范围，不会展开求和项，求导后能得到期望的符号化结果。

修正后的代码：

from sympy import Symbol, sum, diff, IndexedBase

a = Symbol('a')
n = Symbol('n')
k = Symbol('k')
# 用IndexedBase定义符号数组data
data = IndexedBase('data')

cost = sum(a*data[k] + k, (k, 0, n-1))
result = diff(cost, a)
print(result)

运行后输出：

Sum(data[k], (k, 0, n - 1))

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