Fix bug in SOSProblem.apply concerning polynomial equality constraints

1 files changed, 20 insertions, 3 deletions

diff --git a/sumofsquares/problems.py b/sumofsquares/problems.py
index d884899..4256642 100644
--- a/sumofsquares/problems.py
+++ b/sumofsquares/problems.py
@@ -108,6 +108,7 @@ class ConicProblem(Problem):
 
 @dataclassabc(frozen=True)
 class ConicResult(Result):
+    """ Result of a Conic Problem """
     values: dict[OptVariable, float]
     solver_info: Any
 
@@ -194,11 +195,21 @@ class SOSProblem(Problem):
         polynomial_variables, variables = partition(is_optvariable, state.indices.keys())
         polynomial_variables = tuple(polynomial_variables) # because it is a generator
 
-        x = poly.v_stack(polynomial_variables)
+        x = poly.v_stack((1,) + polynomial_variables)
         for i, c in enumerate(self.constraints):
             if isinstance(c, EqualToZero):
-                state, pm = c.expression.apply(state)
-                constraints.append(replace(c, expression=pm))
+                state, deg = c.expression.degree().apply(state)
+
+                # Polynomial equality must be converted into coefficient
+                # matching condition
+                if deg.scalar().constant() > 1:
+                    state, pm = c.expression.linear_in(x).apply(state)
+                    constraints.append(replace(c, expression=pm))
+
+                # A normal (linear) equality
+                else:
+                    state, pm = c.expression.apply(state)
+                    constraints.append(replace(c, expression=pm))
 
             elif isinstance(c, NonNegative):
                 if c.domain:
@@ -257,6 +268,11 @@ class InternalSOSProblem(Problem):
     state: ExpressionState 
 
     def to_conic_problem(self) -> ConicProblem:
+        """
+        Conver the SOS problem into a Conic program. 
+
+        TODO: docstring
+        """
         cost = poly.to_affine(self.cost)
         if cost.degree > 2:
             raise ValueError("This package can solve at most quadratic conic programs, "
@@ -265,6 +281,7 @@ class InternalSOSProblem(Problem):
         is_qp = (cost.degree == 2)
 
         # Sizes
+        # See docstring of ConicProblem.constraints
         dims: dict[str, list[int]] = {
             "z": [], "l": [], "b": [], "q": [], "s": [],
             "ep": [], "ep*": [], "p": [], "p*": [],