Separate out Putinar positivstellensatz as function

3 files changed, 109 insertions, 64 deletions

diff --git a/sumofsquares/__init__.py b/sumofsquares/__init__.py
index 71aafde..99438d6 100644
--- a/sumofsquares/__init__.py
+++ b/sumofsquares/__init__.py
@@ -51,7 +51,7 @@ from polymatrix.expression.expression import Expression
 from polymatrix.expressionstate.abc import ExpressionState
 
 from .abc import Problem, Result, Set, Constraint, Solver
-from .constraints import ArchimedeanSet, NonNegative
+from .constraints import BasicSemialgebraicSet, NonNegative
 from .problems import PutinarSOSProblem
 from .utils import partition
 from .variable import OptVariable, from_names
@@ -62,7 +62,7 @@ def make_sos_constraint(expr: Expression, domain: Set | None = None) -> NonNegat
 
 
 def make_set(*polynomials: Iterable[Expression]):
-    return ArchimedeanSet(polynomials)
+    return BasicSemialgebraicSet(polynomials)
 
 
 def make_problem(
diff --git a/sumofsquares/constraints.py b/sumofsquares/constraints.py
index 61be087..4209ffe 100644
--- a/sumofsquares/constraints.py
+++ b/sumofsquares/constraints.py
@@ -1,15 +1,30 @@
+import math
+
 from dataclasses import dataclass
 from dataclassabc import dataclassabc
 from typing_extensions import override
-from typing import NamedTuple
+from typing import Iterable
+
+import polymatrix as poly
 
 from polymatrix.expression.expression import Expression
+from polymatrix.expressionstate.abc import ExpressionState
+from polymatrix.polymatrix.abc import PolyMatrix
+from polymatrix.statemonad.init import init_state_monad
+from polymatrix.variable.abc import Variable
 
 from .abc import Set, Constraint
+from .error import AlgebraicError
+from .variable import OptVariable, from_name as opt_variable_from_name
+
+
+# ┏━┓┏━╸╺┳╸┏━┓
+# ┗━┓┣╸  ┃ ┗━┓
+# ┗━┛┗━╸ ╹ ┗━┛
 
 
 @dataclass(frozen=True)
-class ArchimedeanSet(Set):
+class BasicSemialgebraicSet(Set):
     r"""
     A set that is described by the intersection of the positive loci of a
     finite number of polynomials (under some conditions). Roughly speaking this
@@ -25,18 +40,19 @@ class ArchimedeanSet(Set):
         \subseteq \mathbf{R}^q.
 
     This tuple stores the :math:`g_i(x)`.
-
-    **Important note:** this package does not check whether the set is actually
-    Archimedean! It is assumed that the given :math:`g_i` make
-    :math:`\mathbf{K}` Archimedean.
     """
-    polynomials: tuple[Expression]
+    polynomials: tuple[Expression, ...]
 
     def __str__(self):
         polynomials = ', '.join(map(str, self.polynomials))
         return f"{self.__class__.__qualname__}({polynomials})"
 
 
+# ┏━╸┏━┓┏┓╻┏━┓╺┳╸┏━┓┏━┓╻┏┓╻╺┳╸┏━┓
+# ┃  ┃ ┃┃┗┫┗━┓ ┃ ┣┳┛┣━┫┃┃┗┫ ┃ ┗━┓
+# ┗━╸┗━┛╹ ╹┗━┛ ╹ ╹┗╸╹ ╹╹╹ ╹ ╹ ┗━┛
+
+
 @dataclassabc(frozen=True)
 class EqualToZero(Constraint):
     """ Constrain an expression to be equal to zero. """
@@ -60,7 +76,7 @@ class NonNegative(Constraint):
     understood as `expression` being SOS.
     """
     expression: Expression 
-    domain: ArchimedeanSet | None = None
+    domain: BasicSemialgebraicSet | None = None
 
     @override
     def __str__(self):
@@ -79,3 +95,63 @@ class PositiveSemiDefinite(Constraint):
 @dataclassabc(frozen=True)
 class ExponentialCone(Constraint):
     expression: Expression
+
+
+# ┏━┓┏━┓┏━┓╻╺┳╸╻╻ ╻┏━┓╺┳╸┏━╸╻  ╻  ┏━╸┏┓╻┏━┓┏━┓┏━╸╺┳╸╺━┓┏━╸
+# ┣━┛┃ ┃┗━┓┃ ┃ ┃┃┏┛┗━┓ ┃ ┣╸ ┃  ┃  ┣╸ ┃┗┫┗━┓┣━┫┣╸  ┃ ┏━┛┣╸ 
+# ╹  ┗━┛┗━┛╹ ╹ ╹┗┛ ┗━┛ ╹ ┗━╸┗━╸┗━╸┗━╸╹ ╹┗━┛╹ ╹┗━╸ ╹ ┗━╸┗━╸
+
+def putinar_psatz(
+        constr: NonNegative,
+        variables: Iterable[OptVariable],
+        polynomial_variables: Iterable[Variable],
+        multiplier_name: str
+    ) -> Iterable[Constraint]:
+    """
+    Apply Putinar's positivstellensatz to a non-negativity constraint over an
+    archimedean domain.
+    """
+    if not constr.domain:
+        # TODO: maybe raise an exception to warn the user?
+        return (constr,)
+
+    # Number of variables in the original problem
+    q = sum(math.prod(v.shape) for v in variables)
+    x = poly.v_stack((1,) + tuple(polynomial_variables))
+
+    new_constraints = []
+    new_constr = constr.expression
+
+    for domain_poly in constr.domain.polynomials:
+        # To know the degree of the domain polynomial we need to evaluate the
+        # expression, hence we use a monad here
+        def make_multiplier( state: ExpressionState) -> tuple[ExpressionState, PolyMatrix]:
+            state, constr_pm = constr.expression.degree().apply(state)
+            state, domain_poly_pm = domain_poly.degree().apply(state)
+            
+            constr_deg = constr_pm.at(0, 0).constant()
+            domain_poly_deg = domain_poly_pm.at(0, 0).constant()
+
+            # degree of multiplier
+            d = constr_deg - domain_poly_deg
+            if d < 0:
+                raise AlgebraicError("Cannot create Positivstellensatz multiplier because "
+                                     f"constraint polynomial has degree {constr_deg} and domain "
+                                     f"polynomial {domain_poly} has degree {domain_poly_deg}")
+            
+            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(multiplier_name, 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))
+
+    return new_constraints
diff --git a/sumofsquares/problems.py b/sumofsquares/problems.py
index 3ec3f54..3105eae 100644
--- a/sumofsquares/problems.py
+++ b/sumofsquares/problems.py
@@ -14,11 +14,10 @@ from typing_extensions import override
 
 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 .constraints import NonNegative, putinar_psatz
 from .error import AlgebraicError
 from .solver.cvxopt import solve_sos_cone as cvxopt_solve_sos_cone
 from .utils import partition
@@ -48,6 +47,22 @@ class SOSResult(Result):
 
 
 @dataclassabc(frozen=True)
+class SchmuedgenSOSProblem(Problem):
+    r"""
+    SOS problem with quadratic cost, linear constraints and optinally with
+    positivity constraints over a basic closed semialgebraic set
+    :math:`\mathbf{K}`. To constrain the positivity over :math:`\mathbf{K}`
+    this problem uses Schmüdgen's Positivstellensatz.
+    """
+    cost: Expression
+    constraints: Sequence[Constraint]
+    variables: Sequence[OptVariable]
+    polynomial_variables: Sequence[Variable]
+    solver: Solver
+    state: ExpressionState
+
+
+@dataclassabc(frozen=True)
 class PutinarSOSProblem(Problem):
     r"""
     SOS problem with quadratic cost, linear constraints and optinally with
@@ -72,60 +87,14 @@ class PutinarSOSProblem(Problem):
 
         # Add langrange-multiplier-like polynomials to SOS constraints, while
         # leaving the rest untouched.
-        variables = tuple(self.variables)
         constraints, to_process = partition(need_positivstellensatz, self.constraints)
+        new_constraints: list[Constraint] = []
+        for i, c in enumerate(to_process):
+            new_constraints.extend(putinar_psatz(c, self.variables, self.polynomial_variables, r"\gamma_{i}"))
 
-        # 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):
-            # 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, constr_pm = constr.expression.degree().apply(state)
-                    state, domain_poly_pm = domain_poly.degree().apply(state)
-                    
-                    constr_deg = constr_pm.at(0, 0).constant()
-                    domain_poly_deg = domain_poly_pm.at(0, 0).constant()
-
-                    # degree of multiplier
-                    d = constr_deg - domain_poly_deg
-                    if d < 0:
-                        raise AlgebraicError("Cannot create Positivstellensatz multiplier because "
-                                             f"constraint polynomial has degree {constr_deg} and domain "
-                                             f"polynomial {domain_poly} has degree {domain_poly_deg}")
-                    
-                    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
+        # Convert to polymatrix object to extract variables created by p-satz
+        # FIXME: this is a dirty trick to make it work for now, it is not
+        # efficient, is there a better way?
         new_state = self.state
         for c in new_constraints:
             new_state, _ = c.expression.apply(new_state)