Solve a constrained optimization problem with a symbolic number of variables using SymPy

Question

I'm trying to solve a constrained minimization problem using SymPy. For a fixed number of variables, say w1, w2 , I'm able to do this in the following way:

from sympy import *

w1, w2 = var('w1, w2', real = True)
n1, n2 = symbols('n1, n2', integer = True)
p1, p2 = symbols('p1, p2', real = True)
f = w1**2 / (n1 * p1) + w2**2 / (n2 * p2)
g = w1 + w2 - 1

lam = symbols('lambda', real = True)
L = f - lam * g

gradL = [diff(L, c) for c in [w1, w2]]
KKT_eqs = gradL + [g]

stationary_points = solve(KKT_eqs, [w1, w2, lam], dict = True)

Are we able to solve this problem for a variable number, say k , of variables? I've tried the following:

from sympy import *

i = symbols('i', cls = Idx)
k = symbols('k', integer = True)

w = IndexedBase('w', real = True)

n = IndexedBase('n', integer = True)
p = IndexedBase('p', real = True)

f = summation(w[i]**2 / (n[i] * p[i]), (i, 1, k))
g = summation(w[i], (i, 1, k)) - 1

lam = symbols('lambda', real = True)
L = f - lam * g

However, I wasn't able to figure out how I need to adopt the remainder of the code.

(I'm new to python so please bear with me.)

Answer 1

One thing you can do is look for a pattern within the concrete cases:

>>> from sympy import *
... from sympy.abc import i
... w = IndexedBase('w')
... np = IndexedBase('np')
... lam = symbols('lambda', real = True)
... def go(n):
...  ww = [w[i] for i in range(n)]
...  f = Add(*[wi**2/np[i] for i,wi in enumerate(ww)])
...  g = Add(*ww) - 1
...  L = f - lam * g
...  gradL = [diff(L, c) for c in ww]
...  KKT_eqs = gradL + [g]
...  return solve(KKT_eqs, ww + [lam], dict = True)
>>> go(2)
[{lambda: 2/(np[0] + np[1]), w[0]: np[0]/(np[0] + np[1]), w[1]: np[1]/(np[0] + np[1])}]
>>> go(3)
[{lambda: 2/(np[0] + np[1] + np[2]), w[0]: np[0]/(np[0] + np[1] + np[2]), w[1]: np[1]/(np[0] + np[1] + np[2]), w[2]: np[2]/(np[0] + np[1] + np[2])}]

Note: since n[i]*p[i] always appear together, those two variables have been combined into one. Do you see the pattern for the solutions? Try go(4) if you don't.

To generalize without concrete insight you might be able to do something with MatrixExpr .