在SymPy中求解KKT方程

import sympy as sp
from sympy import symbols, Eq
from sympy.solvers import solve

n = 2  # 变量个数
m = 1  # 等式约束个数
p = 1  # 不等式约束个数

x = symbols('x1:%d' % (n+1))  # 定义变量x1, x2, ..., xn
g = symbols('g1:%d' % (m+1))  # 定义等式约束g1, g2, ..., gm
h = symbols('h1:%d' % (p+1))  # 定义不等式约束h1, h2, ..., hp

# 目标函数
f = x[0]**2 + x[1]**2

# 等式约束条件
eq_constraints = [x[0] + x[1] - g[0]]

# 不等式约束条件
ineq_constraints = [x[0] - h[0]]

# 松弛变量
s = symbols('s1:%d' % (p+1))

# KKT方程
kkt_eqs = []
kkt_eqs.append(Eq(sp.diff(f, x[0]) + sp.diff(eq_constraints[0]*s[0], x[0]), 0))
kkt_eqs.append(Eq(sp.diff(f, x[1]) + sp.diff(eq_constraints[0]*s[0], x[1]), 0))
kkt_eqs.append(Eq(eq_constraints[0], 0))
kkt_eqs.append(Eq(h[0] - x[0] + s[0], 0))
kkt_eqs.append(Eq(s[0] * h[0], 0))

solution = solve(kkt_eqs, x + g + h + s)