Implement transformation using Positivstellensatz for constrained optimization

2 files changed, 78 insertions, 28 deletions

diff --git a/sumofsquares/problems.py b/sumofsquares/problems.py
index 8b3eaae..309b13c 100644
--- a/sumofsquares/problems.py
+++ b/sumofsquares/problems.py
@@ -3,21 +3,25 @@ Specific optimization problems that are related to Sum of Squares.
 """
 from __future__ import annotations
 
-from itertools import filterfalse
-from dataclasses import replace
+import math
+import polymatrix as poly
+
 from dataclassabc import dataclassabc
+from dataclasses import replace
+from itertools import chain
 from typing import Any, Sequence
 from typing_extensions import override
 
-import polymatrix as poly
-from polymatrix.expression.expression import Expression, VariableExpression
+from polymatrix.expression.expression import Expression, VariableExpression, init_expression
 from polymatrix.expressionstate.abc import ExpressionState
+from polymatrix.statemonad.init import init_state_monad
 from polymatrix.variable.abc import Variable
 
 from .abc import Problem, Constraint, Solver, Result
 from .constraints import NonNegative
 from .cvxopt import solve_sos_cone as cvxopt_solve_sos_cone
-from .variable import OptVariable
+from .utils import partition
+from .variable import OptVariable, from_name as opt_variable_from_name
 
 
 @dataclassabc(frozen=True)
@@ -62,24 +66,71 @@ class PutinarSOSProblem(Problem):
         Apply Putinar's positivstellensatz to convert the SOS constraints
         into coefficient matching constraints 
         """
-        def need_positivstellensatz(c):
-            return isinstance(c, NonNegative) and c.domain
+        def need_positivstellensatz(c: Constraint) -> bool:
+            return isinstance(c, NonNegative) and (c.domain is not None)
 
         # Add langrange-multiplier-like polynomials to SOS constraints, while
         # leaving the rest untouched.
-        variables = list(self.variables)
-        constraints = filterfalse(need_positivstellensatz, self.constraints)
-        multipliers = []
-        for i, c in enumerate(filter(need_positivstellensatz, self.constraints)):
-            d = c.expression.degree()
-            x = poly.v_stack(self.polynomial_variables)
-            # FIXME: check that there are no variables with this names
-            # m = x.combinations(d).parametrize(fr"\gamma_{i}")
-            # multipliers.append(NonNegative(m))
-            raise NotImplementedError
-
-        return replace(self, variables=variables,
-                       constraints=tuple(constraints) + tuple(multipliers))
+        variables = tuple(self.variables)
+        constraints, to_process = partition(need_positivstellensatz, self.constraints)
+
+        # New variables and constraints that will be introduced by the
+        # positivstellensatz
+        # new_variables = []
+        new_constraints = []
+
+        # Number of variables in the original problem
+        q = len(self.variables)
+        x = poly.v_stack((1,) + tuple(self.polynomial_variables))
+
+        for i, constr in enumerate(to_process):
+            # For each polynomial that defines the archimedean set
+            # we create a multiplier polynomial of the same degree
+
+            # FIXME: this is not optimal, look into how to determine the
+            # smallest degree needed to make the positivstellensatz work.
+            new_constr = constr.expression
+            for domain_poly in constr.domain.polynomials:
+                # To know the degree of the polynomial we need to evaluate the
+                # expression, hence we use a monad here
+                def make_multiplier(state):
+                    state, p = domain_poly.degree().apply(state)
+
+                    d = p.at(0, 0).constant()
+                    n = sum(math.comb(k + q - 1, k) for k in range(0, d+1))
+
+                    # TODO: check that there is not a variable with this name already
+                    # Coefficients of the multiplier polynomial 
+                    c = opt_variable_from_name(rf"\gamma_{i}", shape=(1,n))
+                    multiplier = c @ x.combinations(tuple(range(d +1)))
+
+                    return multiplier.apply(state)
+
+                m = poly.from_statemonad(init_state_monad(make_multiplier))
+
+                # Subtract multiplier from expression and impose that it is also SOS
+                new_constr -= m * domain_poly
+                new_constraints.append(NonNegative(m))
+
+            new_constraints.append(NonNegative(new_constr))
+
+        # FIXME: this is a dirty trick to make it work for now,
+        # it is not efficient, replace with correct solution
+
+        # Convert to polymatrix object to extract variables
+        new_state = self.state
+        for constr in new_constraints:
+            new_state, _ = constr.expression.apply(new_state)
+
+        def is_optvariable(v):
+            return isinstance(v, OptVariable)
+
+        new_variables = filter(is_optvariable, new_state.indices.keys())
+
+        return replace(self,
+                       state=new_state,
+                       variables=tuple(new_variables),
+                       constraints=tuple(constraints) + tuple(new_constraints))
 
 
     @override
diff --git a/sumofsquares/variable.py b/sumofsquares/variable.py
index 12df4df..6dec3fb 100644
--- a/sumofsquares/variable.py
+++ b/sumofsquares/variable.py
@@ -60,13 +60,12 @@ def init_opt_variable_expr(name, shape):
     return OptVariableImpl(name, shape)
 
 
-def from_names(names: str, shape: tuple[int, int] = (1, 1)) -> tuple[VariableExpression] | VariableExpression:
-    """ Construct one or multiple variables from comma separated a list of names. """
-    variables = tuple(init_variable_expression(
-        underlying=init_opt_variable_expr(name.strip(), shape))
-                      for name in names.split(","))
+def from_name(name: str, shape: tuple[int, int] = (1, 1)) -> VariableExpression:
+    """ Construct an optimization variable. """
+    return init_variable_expression(underlying=init_opt_variable_expr(name, shape))
 
-    if len(variables) == 1:
-        return variables[0]
 
-    return variables
+def from_names(names: str, shape: tuple[int, int] = (1, 1)) -> Iterable[VariableExpression]:
+    """ Construct one or multiple variables from comma separated a list of names. """
+    for name in names.split(","):
+        yield from_name(name.strip(), shape)