Remove state field from ConicProblem

3 files changed, 19 insertions, 23 deletions

diff --git a/examples/sostools/sosdemo2.py b/examples/sostools/sosdemo2.py
index d0cf81d..ece36b7 100644
--- a/examples/sostools/sosdemo2.py
+++ b/examples/sostools/sosdemo2.py
@@ -33,7 +33,7 @@ constraints = (
     sos.make_sos_constraint(- (dV @ f)),
 )
 
-prob, result = sos.solve_problem(cost=None, constraints=constraints, solver=Solver.CVXOPT)
+prob, state, result = sos.solve_problem(cost=None, constraints=constraints, solver=Solver.CVXOPT)
 print(prob)
 
 # TODO: create SolverStatus enumeration
diff --git a/sumofsquares/__init__.py b/sumofsquares/__init__.py
index bd5ec6a..58480fd 100644
--- a/sumofsquares/__init__.py
+++ b/sumofsquares/__init__.py
@@ -129,7 +129,7 @@ def maximize(cost, *args, **kwargs) -> SOSProblem:
     return make_problem(-cost, *args, **kwargs)
 
 
-def solve_problem(*args, verbose: bool = True, state: ExpressionState | None = None, **kwargs) -> tuple[Problem, Result]:
+def solve_problem(*args, state: ExpressionState | None = None, verbose: bool = True, **kwargs) -> tuple[Problem, ExpressionState, Result]:
     """
     Solve a sum-of-squares optimization problem.
     This function is just a shorthand for
@@ -139,4 +139,6 @@ def solve_problem(*args, verbose: bool = True, state: ExpressionState | None = N
         result = prob.solve(verbose)
     """
     prob = make_problem(*args, **kwargs)
-    return prob, prob.solve(verbose=verbose, state=state)
+    if state is None:
+        state = poly.make_state()
+    return (prob,) + prob.solve(verbose=verbose, state=state)
diff --git a/sumofsquares/problems.py b/sumofsquares/problems.py
index 7bcd629..14618fe 100644
--- a/sumofsquares/problems.py
+++ b/sumofsquares/problems.py
@@ -65,7 +65,7 @@ class ConicProblem(Problem):
     constraints: dict[str, list[tuple[tuple[NDArray, ...], NDArray]]]
     """ 
     Conic constraints are saved in a dictionary with the following form.
-    For the key we will use the following names to refer to various constraint
+    For the key we will use the following names to refer to various 
     types of constraints
     ::
 
@@ -162,12 +162,9 @@ class SOSProblem(Problem):
         return s
 
     @override
-    def solve(self, verbose: bool = False, state: ExpressionState | None = None) -> Result:
-        if state is None:
-            state = poly.make_state()
-
+    def solve(self, state: ExpressionState, verbose: bool = False) -> tuple[ExpressionState, Result]:
         state, internal_prob = self.apply(state)
-        return internal_prob.solve(verbose)
+        return internal_prob.solve(state, verbose)
 
     def apply(self, state: ExpressionState) -> tuple[ExpressionState, InternalSOSProblem]:
         """
@@ -207,6 +204,7 @@ class SOSProblem(Problem):
         x = poly.v_stack((1,) + tuple(
             init_variable_expr(v, state.get_shape(v))
             for v in polynomial_variables))
+
         for i, c in enumerate(self.constraints):
             if isinstance(c, EqualToZero):
                 state, deg = c.expression.degree().apply(state)
@@ -287,7 +285,7 @@ class SOSProblem(Problem):
 
         return state, InternalSOSProblem(cost, tuple(pm_constraints),
                                          tuple(variables), polynomial_variables,
-                                         self.solver, state)
+                                         self.solver)
 
 
 # FIXME: I hate this name, and the fact that this class needs to exist
@@ -308,12 +306,9 @@ class InternalSOSProblem(Problem):
     polynomial_variables: Sequence[Symbol]
     solver: Solver
 
-    # TODO: remove state field from this class, it is redundant
-    state: ExpressionState 
-
-    def to_conic_problem(self, verbose: bool = False) -> ConicProblem:
+    def to_conic_problem(self, state: ExpressionState, verbose: bool = False) -> ConicProblem:
         """
-        Conver the SOS problem into a Conic program. 
+        Convert the SOS problem into a Conic program. 
         """
         cost = poly.to_affine(self.cost)
         if cost.degree > 2:
@@ -336,7 +331,7 @@ class InternalSOSProblem(Problem):
         }
 
         # indices of variables in the optimization problem, sorted
-        variable_indices = sum(sorted(tuple(self.state.get_indices_as_variable_index(v))
+        variable_indices = sum(sorted(tuple(state.get_indices_as_variable_index(v))
                                       for v in self.variables), ())
 
         # cost linear term
@@ -357,9 +352,8 @@ class InternalSOSProblem(Problem):
             # see also polymatrix.index.MonomialIndex.__lt__
             P[np.triu_indices(n)] = cost_coeffs[2]
 
-            # Make symmetric. There is no 0.5 factor because there it is already
-            # there in the canonical form, see docstring of ConicProblem
-            P = (P + P.T)
+            # Make symmetric
+            P = .5 * (P + P.T)
 
         else:
             P = None
@@ -400,7 +394,7 @@ class InternalSOSProblem(Problem):
 
             else:
                 raise NotImplementedError(f"Cannot convert constraint of type {type(c)} "
-                                          "into a conic constriant (yet).")
+                                          "into a conic constraint (yet).")
 
         if all(len(cl) == 0 for cl in constraints.values()):
             raise ValueError("Optimization problem is unconstrained!")
@@ -414,11 +408,11 @@ class InternalSOSProblem(Problem):
         return ConicProblem(P=P, q=q, constraints=constraints,
                             dims=dims, is_qp=is_qp,
                             solver=self.solver,
-                            variables={v : self.state.get_shape(v)
+                            variables={v : state.get_shape(v)
                                 for v in self.variables
                             })
 
 
     @override
-    def solve(self, verbose: bool = False) -> Result:
-        return self.to_conic_problem(verbose).solve(verbose)
+    def solve(self, state: ExpressionState, verbose: bool = False) -> tuple[ExpressionState, Result]:
+        return state, self.to_conic_problem(state, verbose).solve(verbose)