I am trying to solve KKT equations using sympy. All of the equations are symbolic and contain constants that are not given as numbers but as symbols. Alongside with the equations, there are also inequality constraints. Is it possible to do this in sympy? If not, are there any alternatives?
An example would be:
Doing your question by hand, we get the first Lagrangian to be L = x**2 - bx + 1 - lambda(x - a).
Differentiating with respect to x and lambda and setting to 0, we solve to get x = a, lambda = 2a - b.
So if lambda < 0, we repeat with the new Lagrangian, L = x**2 - bx + 1.
This gives us 2 cases: If 2a - b < 0, then x = b/2, otherwise, x = a.
The following code reproduces this logic. Copy paste this into a file called lagrangian.py and see example 4 for your specific query.
"""
KKT Conditions:
Goal:
To minimise/maximise f(x_) subject to gi(x_) >= 0 for all i and hi(x_) == 0 for all i
where x_ refers to a vector x.
Variables with a `*` after them are optimal quantities.
1. gi(x*) and hi(x*) is feasible (that is, they are satisfied)
2. (df/dxj)(x*) - sum(lambdai* * (dgi/dxj)(x*)) - sum(mui* * (dhi/dxj)(x*)) = 0 for all j
3. lambdai* * gi(x*) = 0 for all i
4. lambdai >= 0 for all i
"""
from sympy import *
from typing import Iterable
def kkt(f, g=None, h=None, x=None):
"""
Finds the optimal values of `x` for `f` given the equalities `g[i] >= 0` for all i
and `h[i] == 0` for all i.
TODO: Remove the private variables _lambda and _mu from the output.
TODO: Make the output more user friendly.
Examples:
>>> from sympy import *
>>> from lagrangian import kkt
>>> x = list(symbols("x:3"))
Example 0 the most basic
>>> kkt(x[0]**2)
([x0], FiniteSet((0,)))
Example 1 from References
>>> kkt(2 * x[0] ** 2 + x[1] ** 2 + 4 * x[2]**2,
... h=[x[0] + 2*x[1] - x[2] - 6, 2 * x[0] - 2 * x[1] + 3 * x[2] - 12])
([_mu0, _mu1, x0, x1, x2], FiniteSet((504/67, 424/67, 338/67, 80/67, 96/67)))
Example 2 from References
>>> kkt(x[0] ** 2 + 2 * x[1] ** 2 + 3 * x[2] ** 2,
... [5 * x[0] - x[1] - 3 * x[2] - 3,
... 2 * x[0] + x[1] + 2 * x[2] - 6])
([_lambda0, x0, x1, x2], FiniteSet((72/35, 72/35, 18/35, 24/35)))
Example 3 from References
>>> kkt(4 * x[0] ** 2 + 2 * x[1] ** 2,
... [-2 * x[0] - 4 * x[1] + 15],
... [3 * x[0] + x[1] - 8])
([_mu0, x0, x1], FiniteSet((64/11, 24/11, 16/11)))
Example 4 for general KKT
>>> t, a, b = symbols("x a b")
>>> kkt(t ** 2 - b * t + 1, t - a, x=t)
Piecewise((([x], FiniteSet((b/2,))), 2*a - b < 0), (([_lambda0, x], FiniteSet((2*a - b, a))), True))
Warnings:
This function uses recursion and if queries are such as example 4 with many inequality
conditions, one will experience heavy performance issues.
References:
.. [1] http://apmonitor.com/me575/index.php/Main/KuhnTucker
Disadvantages:
- Does not allow for arbitrary number of inequalities and equalities.
- Does not work for functions of f that have an infinite
number of turning points and no equality constraints (I think)
"""
# begin sanity checks
if not g:
g = []
if not h:
h = []
if not x:
x = list(f.free_symbols)
if not isinstance(g, Iterable):
g = [g]
if not isinstance(h, Iterable):
h = [h]
if not isinstance(x, Iterable):
x = [x]
# end sanity checks
def grad(func):
"""Returns the grad of an expression or function `func`"""
grad_f = Matrix(len(x), 1, derive_by_array(func, x))
return grad_f
# define our dummy variables for the inequalities
if g:
_lambdas = list(symbols(f"_lambda:{len(g)}", real=True))
sum_g = sum([_lambdas[i] * g[i] for i in range(len(g))])
else:
_lambdas = []
sum_g = 0
# define our dummy variables for the equalities
if h:
_mus = list(symbols(f"_mu:{len(h)}", real=True))
sum_h = sum([_mus[i] * h[i] for i in range(len(h))])
else:
_mus = []
sum_h = 0
# define the lagrangian
lagrangian = f - sum_g - sum_h
# find grad of lagrangian
grad_l = grad(lagrangian)
# 1. feasibility conditions
feasible = Matrix(len(g) + len(h), 1, g + h)
# 2. combine everything into a vector equation
eq = grad_l.row_insert(0, feasible)
solution = solve(eq, x + _lambdas + _mus, set=True)
# 4. remove all non-binding inequality constraints
# for each _lambda solution, add a new solution with that inequality removed
# in the case that it is false.
pieces = []
for i, lamb in enumerate(_lambdas):
new_g = g[:i] + g[i + 1:]
lamb_sol_index = [i for i, s in enumerate(solution[0]) if s == lamb][0]
for solution_piece in solution[1]:
lamb_sol = solution_piece[lamb_sol_index]
try:
if lamb_sol >= 0: # we dont need to check the next kkt if the value is known.
continue
except TypeError: # error when inequality cannot be identified
pass
pieces.append((kkt(f, new_g, h, x), lamb_sol < 0))
pieces.append((solution, True))
return Piecewise(*pieces)
