从 sympy 表达式中获取线性和二次项

Question

我有一个通用表达式E ，它依赖于X(t) （变量数组）和dX(t) 。 有没有一种简单的方法可以从E获得linear和quadratic项？ 我的实际代码仅在E简单时才有效。

我的目标是得到矩阵M ， V ， K ， A ， B和C使得E = dX @ M @ dX + dX @ V @ X + X @ K @ X + A @ dX + B @ X + C

举个例子

import sympy as sp

t = sp.symbols("t")
x = sp.Function("x")(t)
y = sp.Function("y")(t)
dx = sp.diff(x, t)
dy = sp.diff(y, t)
X = sp.Array([x, y])
dX = sp.Array([dx, dy])

E = 7*(dx**2) + 2*dx*dy + 6*(dy**2)  # M = [[7, 1],
                                     #      [1, 6]]
E += 4*dx*x + 6*dx*y + 3*dy*x + 5*dy*y  # V = [[4, 6],
                                        #      [3, 5]]
E += (-3)*x**2 + (-4)*x*y + (-5)*y**2  # K = [[-3, -2],
                                       #      [-2, -5]]
E += 1*dx + (-7)*dy  # A = [1, -7]
E += (-5)*x + 9*y  # B = [-5, 9]
E += 80  # C = 80

到目前为止，我所做的只是取导数并手动减去这些项。 当我有常量值时，以下代码可以正常工作：

diff = sp.derive_by_array
def dot(A, B):
    # Matrix multiplication of A and B
    ndimA = len(sp.shape(A))
    C = sp.tensorproduct(A, B)
    D = sp.tensorcontraction(C, (ndimA-1, ndimA))
    return D

M = diff( diff(E, dX), dX)/2
E -= dot(dX, dot(M, dX))
V = diff( diff(E, X), dX)
E -= dot(dX, dot(V, X))
K = diff( diff(E, X), X)/2
E -= dot(X, dot(K, X))
A = diff(E, dX)
E -= dot(A, dX)
B = diff(E, X)
E -= dot(B, X)
C = sp.expand(E)

但是当我输入非常数项时，结果是错误的。 例如E = cos(x) ：

# expected
M = [[0, 0],
     [0, 0]]
V = [[0, 0],
     [0, 0]]
K = [[0, 0],
     [0, 0]]
A = [0, 0]
B = [0, 0]
C = cos(x)

# gotten
M = [[0, 0],
     [0, 0]]
V = [[0, 0],
     [0, 0]]
K = [[-cos(x)/2, 0],
     [0, 0]]
A = [0, 0]
B = [-x(t)**2*sin(x(t))/2 + x(t)*cos(x(t)) - sin(x(t)), 0]
C = x(t)**3*sin(x(t))/2 - x(t)**2*cos(x(t))/2 + x(t)*sin(x(t)) + cos(x(t))

我期待使用高级表达式操作和 function srepr来分解E的解决方案，但我认为它应该存在一种更简单的方法来做到这一点。

Answer 1

我找到了一个使用符号矩阵的解决方案，并使用sympy.solve来找到解决方案。 我创建symbolic function arrays ，然后计算标量Esup作为所需的表达式，并求解以找到每个矩阵内的值。

import numpy as np

all_variables = []
def MatrixFunction(name, shape):
    ndim = len(shape)
    M = np.zeros(shape, dtype="object")
    if ndim == 1:
        for i in range(shape[0]):
            M[i] = sp.Function(f"{name}{i}")(x, y)
            all_variables.append(M[i])
    elif ndim == 2:
        for i in range(shape[0]):
            for j in range(shape[1]):
                M[i, j] = sp.Function(f"{name}{i}{j}")(x, y)
                all_variables.append(M[i, j])
    return sp.MutableDenseNDimArray(M)

M = MatrixFunction("M", [len(X), len(X)])
V = MatrixFunction("V", [len(X), len(X)])
K = MatrixFunction("K", [len(X), len(X)])
A = MatrixFunction("A", [len(X)])
B = MatrixFunction("B", [len(X)])
C = sp.Function("C")(x, y)
all_variables.append(C)
equations = []
for i in range(len(X)):
    for j in range(i+1, len(X)):
        equations.append(M[i, j] - M[j, i])
        equations.append(K[i, j] - K[j, i])

Esup = dot(dX, dot(M, dX))
Esup += dot(dX, dot(V, X))
Esup += dot(X, dot(K, X))
Esup += dot(A, dX)
Esup += dot(B, X)
Esup += C

dife = E - Esup
poly = sp.Poly(dife, (x, y, dx, dy))
equations += poly.coeffs()

solution = sp.solve(equations, all_variables)
for var, val in solution.items():
    C = C.subs(var, val)
    for i in range(len(X)):
        A[(i,)] = A[(i,)].subs(var, val)
        B[(i,)] = B[(i,)].subs(var, val)
        for j in range(len(X)):
            M[i, j] = M[i, j].subs(var, val)
            V[i, j] = V[i, j].subs(var, val)
            K[i, j] = K[i, j].subs(var, val)