Change structure of ConicProblem to allow interface to other solvers

5 files changed, 244 insertions, 139 deletions

diff --git a/sumofsquares/constraints.py b/sumofsquares/constraints.py
index b808bbb..844cc37 100644
--- a/sumofsquares/constraints.py
+++ b/sumofsquares/constraints.py
@@ -64,10 +64,10 @@ class NonNegative(Constraint[E]):
     the domain in set to `None` then it is interpreted as being non-negative
     everywhere.
 
-    **Note:** In this package, non-negativitiy constraint will always
-    eventually be replaced by the sufficient condition that the expression can
-    be written as a sum-of-squares, hence this constraint can also be
-    understood as `expression` being SOS.
+    **Note:** In this package, non-negativitiy of non-affine constraint will
+    always eventually be replaced by the sufficient condition that the
+    expression can be written as a sum-of-squares, hence this constraint can
+    also be understood as `expression` being SOS.
     """
 
     # This was part of the docstring but is not implemented yet
@@ -105,3 +105,30 @@ class ExponentialCone(Constraint[E]):
     a single point that must be inside of the exponential cone.
     """
     expression: E
+
+
+# TODO: Create missing classes below and add support in
+#  - SOSProblem.apply
+#  - InternalSOSProblem.to_conic_problem
+#  - solver.SOLVERNAME.solve_cone (for each solver)
+
+
+# @dataclassabc(frozen=True)
+# class SecondOrderCone(Constraint[E]):
+#     """ Constrain an expression to be in the second order cone. """
+#     expression: E
+#     t: E # FIXME: review this
+
+
+# @dataclassabc(frozen=True)
+# class DualExponentialCone(Constraint[E]):
+#     ...
+
+# @dataclassabc(frozen=True)
+# class PowerCone(Constraint[E]):
+#     ...
+
+
+# @dataclassabc(frozen=True)
+# class DualPowerCone(Constraint[E]):
+#     ...
diff --git a/sumofsquares/problems.py b/sumofsquares/problems.py
index e7b9490..02d5a1e 100644
--- a/sumofsquares/problems.py
+++ b/sumofsquares/problems.py
@@ -14,7 +14,7 @@ from numpy.typing import NDArray
 from typing import Any, Sequence
 from typing_extensions import override
 
-from polymatrix.expression.expression import Expression, VariableExpression, init_expression
+from polymatrix.expression.expression import Expression, VariableExpression
 from polymatrix.expressionstate import ExpressionState
 from polymatrix.polymatrix.mixins import PolyMatrixMixin
 from polymatrix.polymatrix.index import MonomialIndex
@@ -22,7 +22,8 @@ from polymatrix.variable import Variable
 
 from .abc import Problem, Constraint, Solver, Result
 from .constraints import NonNegative, EqualToZero, PositiveSemiDefinite, ExponentialCone
-from .solver.cvxopt import solve_sos_cone as cvxopt_solve_sos_cone
+from .solver.cvxopt import solve_cone as cvxopt_solve_cone
+from .solver.scs import solve_cone as scs_solve_cone
 from .utils import partition
 from .variable import OptVariable
 
@@ -43,47 +44,67 @@ class ConicProblem(Problem):
 
         minimize      .5 * x.T @ P @ x + q.T @ x
         
-        subject to    G @ x <= h
-                      A @ x = b
+        subject to    CONSTRAINTS
 
+    See docstring of :py:attr:`ConicProblem.constraints` for supported types of
+    constraints.
     """
-    P: NDArray
-    q: NDArray
-    A: NDArray
-    b: NDArray
-    G: NDArray
-    h: NDArray
+    P: NDArray | None
+    """ Quadratic term of the cost function """
+
+    q: NDArray | None
+    """ Linear term of the cost function """
 
-    dims: dict
     is_qp: bool
+    dims: dict
+
+    # TODO: Maybe do a dict[type, list[...]] instead?
+    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
+    types of constraints
+    ::
+
+       z    zero cone                   (s = 0)
+       l    linear cone                 (s >= 0)
+       b    box cone                    (tl <= s <= tu)
+       q    second order cone           (|s|_2 <= t)
+       s    positive semidefinite cone  (s is PSD)
+       ep   exponential cone            ((x,y,z) in ExpCone)
+       ep*  dual exponential cone       ((u,v,w) in DualExpCone)
+       p    power cone
+       p*   dual power cone
+
+    Then the values are the coefficients of a LMI saved in a tuple `(linear
+    terms,  constant term)`. The linear terms is again a tuple, with length
+    equal the number of variables. The coefficients are left as matrices since
+    each solver may expect then to be vectorized in a different way.
+
+    See also solver.SOLVERNAME.solve_cone functions.
+    """
+
     solver: Solver
     variables: Sequence[OptVariable]
 
     @property
     @override
-    def cost(self) -> tuple[NDArray, NDArray]:
+    def cost(self) -> tuple[NDArray | None, NDArray | None]:
         return self.P, self.q
 
-    @property
-    @override
-    def constraints(self) -> dict[str, tuple[NDArray, NDArray]]:
-        return {
-            "eq": (self.A, self.b),
-            "ineq": (self.G, self.h)
-        }
-
     @override
     def solve(self, verbose: bool = False) -> ConicResult:
         match self.solver:
             case Solver.CVXOPT:
-                result, info = cvxopt_solve_sos_cone(self, verbose)
+                result, info = cvxopt_solve_cone(self, verbose)
+            case Solver.SCS:
+                result, info = scs_solve_cone(self, verbose)
             case _:
-                raise NotImplementedError
+                raise NotImplementedError(f"Solver {self.solver} is not supported (yet).")
 
         return ConicResult(values=result, solver_info=info)
 
 
-
 @dataclassabc(frozen=True)
 class ConicResult(Result):
     values: dict[OptVariable, float]
@@ -100,8 +121,8 @@ class ConicResult(Result):
             var = var.underlying
 
         if var not in self.values:
-            # TODO: error message
-            raise KeyError
+            raise KeyError(f"There is no result for the variable {var}, "
+                           "was the problem successfully solved?")
 
         return self.values[var]
 
@@ -129,8 +150,8 @@ class SOSProblem(Problem):
                 f"minimize  {self.cost}\n  " \
                 f"subject to ({n} constraint{'s' if n > 1 else ''})\n"
 
-        for c in self.constraints:
-            s += f"\t* {str(c)}\n"
+        for i, c in enumerate(self.constraints):
+            s += f"\tnr. {i}: {str(c)}\n"
 
         return s
 
@@ -165,14 +186,34 @@ class SOSProblem(Problem):
         polynomial_variables = tuple(polynomial_variables) # because it is a generator
 
         x = poly.v_stack(polynomial_variables)
-        for c in self.constraints:
-            # Scalar polynomial equality and non-negativity contraints need to
-            # be converted to (linear) coefficient matching constraints
-            if isinstance(c, NonNegative | EqualToZero):
-                state, pm = c.expression.quadratic_in(x).apply(state)
+        for i, c in enumerate(self.constraints):
+            if isinstance(c, EqualToZero):
+                state, pm = c.expression.apply(state)
+                constraints.append(replace(c, expression=pm))
+
+            elif isinstance(c, NonNegative):
+                if c.domain:
+                    raise ValueError(f"Cannot convert non-negativity constraint nr. {i} "
+                                     "to conic constraint. Domain restricted non-negativity "
+                                     "must be preprocessed by using a Positivstellensatz!")
 
+                state, deg = c.expression.degree().apply(state)
+
+                # Polynomial non-negativity contraints is converted to PSD
+                # constraint of SOS quadratic form
+                if deg.scalar().constant() > 1:
+                    state, pm = c.expression.quadratic_in(x).symmetric().apply(state)
+                    constraints.append(PositiveSemiDefinite(pm))
+
+                # A normal (linear) constraint
+                else:
+                    state, pm = c.expression.apply(state)
+                    constraints.append(replace(c, expression=pm))
+
+            # PSD constraint can be passed as-is
             elif isinstance(c, PositiveSemiDefinite):
                 state, pm = c.expression.apply(state)
+                constraints.append(replace(c, expression=pm))
 
             elif isinstance(c, ExponentialCone):
                 raise NotImplementedError
@@ -180,8 +221,6 @@ class SOSProblem(Problem):
             else:
                 raise NotImplementedError
 
-            constraints.append(replace(c, expression=pm))
-
         return state, InternalSOSProblem(cost, tuple(constraints),
                                          tuple(variables), polynomial_variables,
                                          self.solver, state)
@@ -204,6 +243,7 @@ class InternalSOSProblem(Problem):
     variables: Sequence[OptVariable]
     polynomial_variables: Sequence[Variable]
     solver: Solver
+
     # TODO: remove state field from this class, it is redundant
     state: ExpressionState 
 
@@ -215,19 +255,27 @@ class InternalSOSProblem(Problem):
 
         is_qp = (cost.degree == 2)
 
-        A_rows, b_rows = [], []
-        G_rows, h_rows= [], []
-        l_dims, q_dims, s_dims = [], [], []
+        # Sizes
+        dims: dict[str, list[int]] = {
+            "z": [], "l": [], "b": [], "q": [], "s": [],
+            "ep": [], "ep*": [], "p": [], "p*": [],
+        }
+        
+        # See docstring of ConicProblem.constraints
+        constraints: dict[str, list[tuple[tuple[NDArray, ...], NDArray]]] = {
+            "z": [], "l": [], "b": [], "q": [], "s": [],
+            "ep": [], "ep*": [], "p": [], "p*": [],
+        }
 
-        # indices of variables in the optimization problem
+        # indices of variables in the optimization problem, sorted
         variable_indices = sum(sorted(tuple(self.state.get_indices_as_variable_index(v))
                                       for v in self.variables), ())
 
         # cost linear term
         cost_coeffs = dict(cost.affine_coefficients_by_degrees(variable_indices))
-        q = cost_coeffs.get(1) 
-        if q is None and is_qp:
-            q = np.zeros((len(variable_indices), 1))
+        q = cost_coeffs.get(1)
+        if q is not None:
+            q = q.reshape((-1, 1))
 
         # cost quadratic term
         if is_qp:
@@ -238,98 +286,55 @@ class InternalSOSProblem(Problem):
             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 of CVXOPT
+            # there in the canonical form, see docstring of ConicProblem
             P = (P + P.T)
 
         else:
             P = None
 
         for c in self.constraints:
-            if isinstance(c, PositiveSemiDefinite):
-                constr = poly.to_affine(c.expression)
-
-                nrows, ncols = constr.shape
-                if nrows != ncols:
-                    raise ValueError(f"PSD constraint cannot contain non-square matrix of shape ({nrows, ncols})!")
+            constr = poly.to_affine(c.expression)
+            nrows, ncols = constr.shape
 
-                s_dims.append(nrows)
+            if constr.degree > 1:
+                raise ValueError("To convert to conic constraints must be linear or affine "
+                                f"but {str(c.expression)} has degree {constr.degree}.")
 
-                # constant term
-                c_coeff = constr.affine_coefficient(MonomialIndex.constant())
-                c_stacked = c_coeff.T.reshape((nrows * ncols,))
+            # Get constant and linear terms
+            constant = constr.affine_coefficient(MonomialIndex.constant())
+            linear = tuple(constr.affine_coefficient(v) for v in variable_indices)
+            
+            if isinstance(c, EqualToZero):
+                dims["z"].append(nrows)
+                constraints["z"].append((linear, constant))
 
-                # linear terms
-                l_stacked_rows = []
-                for v in variable_indices:
-                    v_coeff = constr.affine_coefficient(v)
-                    v_stacked = v_coeff.T.reshape((nrows * ncols,))
-                    l_stacked_rows.append(v_stacked)
+            elif isinstance(c, NonNegative):
+                dims["l"].append(nrows)
+                constraints["l"].append((linear, constant))
 
-                l_stacked = np.vstack(l_stacked_rows)
+            elif isinstance(c, PositiveSemiDefinite):
+                if nrows != ncols:
+                    raise ValueError(f"PSD constraint cannot contain non-square matrix of shape ({nrows, ncols})!")
 
-                h_rows.append(c_stacked)
-                G_rows.append(-1 * l_stacked)
+                dims["s"].append(nrows)
+                constraints["s"].append((linear, constant))
 
             elif isinstance(c, ExponentialCone):
-                # TODO: implement
-                raise NotImplementedError
-
-            elif isinstance(c, EqualToZero | NonNegative):
-                constr = poly.to_affine(c.expression)
-                if constr.degree > 1:
-                    raise ValueError("To convert to conic constraints must be linear or affine "
-                                    f"but {str(c.expression)} has degree {constr.degree}.")
-
-                # For each constraint we add PSD matrix
-                nrows, ncols = constr.shape
-                s_dims.append(nrows)
-
-                # Constant term
-                c_coeff = constr.affine_coefficient(MonomialIndex.constant())
-                c_stacked = c_coeff.T.reshape((nrows * ncols,))
-
-                # Linear term
-                l_stacked_rows = []
-                for v in variable_indices:
-                    # Get the affine coefficient associated to this variable
-                    v_coeff = constr.affine_coefficient(v)
-
-                    # stack the columns of coeff matrix into a fat row vector (column major order).
-                    v_stacked = v_coeff.T.reshape((nrows * ncols,))
-                    l_stacked_rows.append(v_stacked)
-
-                l_stacked = np.vstack(l_stacked_rows)
-
-                if isinstance(c, EqualToZero):
-                    A_rows.append(l_stacked)
-                    # Factor of -1 because it is on the right hand side
-                    b_rows.append(-1 * c_stacked)
-
-                elif isinstance(c, NonNegative):
-                    if c.domain:
-                        raise ValueError("Cannot convert non-negativity constraint to conic constraint. "
-                                         "Domain restricted non-negativity must be "
-                                         "preprocessed by using a Positivstellensatz!")
-
-                    # Factor of -1 because it is on the right hand side
-                    G_rows.append(-1 * l_stacked)
-                    h_rows.append(c_stacked)
+                # Interpret row-wise
+                dims["ep"].extend(1 for _ in range(nrows))
+                for i in range(nrows):
+                    constraints["ep"].append((
+                        tuple(m[i, :] for m in linear),
+                        constant[i, :]))
 
             else:
                 raise NotImplementedError(f"Cannot convert constraint of type {type(c)} "
-                                          "into a conic constriant.")
-
-        G = np.hstack(G_rows).T if G_rows else None
-        h = np.hstack(h_rows) if h_rows else None
-        A = np.hstack(A_rows).T if A_rows else None
-        b = np.hstack(b_rows) if b_rows else None
-
-        dims = {'l': l_dims, 'q': q_dims, 's': s_dims}
+                                          "into a conic constriant (yet).")
 
-        if all((G is None, h is None, A is None, b is None)):
+        if all(len(cl) == 0 for cl in constraints.values()):
             raise ValueError("Optimization problem is unconstrained!")
 
-        return ConicProblem(P=P, q=q, A=A, b=b, G=G, h=h,
+        return ConicProblem(P=P, q=q, constraints=constraints,
                             dims=dims, is_qp=is_qp,
                             solver=self.solver,
                             variables=self.variables)
diff --git a/sumofsquares/solver/cvxopt.py b/sumofsquares/solver/cvxopt.py
index 87e0b0f..e41633f 100644
--- a/sumofsquares/solver/cvxopt.py
+++ b/sumofsquares/solver/cvxopt.py
@@ -7,11 +7,13 @@ import cvxopt
 import numpy as np
 import math
 
+from numpy.typing import NDArray
 from collections import UserDict
 from typing import TYPE_CHECKING
+from pprint import pprint
 
 from ..abc import SolverInfo
-from ..error import SolverError
+from ..error import SolverError, NotSupportedBySolver
 from ..variable import OptVariable
 
 if TYPE_CHECKING:
@@ -29,8 +31,18 @@ class CVXOPTInfo(SolverInfo, UserDict):
         return self[key]
 
 
-def solve_sos_cone(prob: ConicProblem, verbose: bool = False,
-                   *args, **kwargs) -> tuple[dict[OptVariable, float], CVXOPTInfo]:
+def vectorize_matrix(m: NDArray) -> NDArray:
+    r"""
+    Vectorize a symmetric matrix for CVXOPT by stacking its columns. So,
+    a :math:`n\times n` matrix becomes an :math:`n^2 \times 1` column vector.
+    """
+    # if m.shape[0] != m.shape[1]:
+    #     raise ValueError("Matrix must be symmetric, and hence square!")
+    return np.hstack(m.T.tolist()).reshape((-1, 1))
+
+
+def solve_cone(prob: ConicProblem, verbose: bool = False,
+               *args, **kwargs) -> tuple[dict[OptVariable, NDArray | float], CVXOPTInfo]:
     r"""
     Any `*args` and `**kwargs` other than `prob` are passed to the CVXOPT
     solver.
@@ -43,32 +55,84 @@ def solve_sos_cone(prob: ConicProblem, verbose: bool = False,
     #                  A @ x = b
     #                  s >= 0
     #
-    # This matches the structure of ConicProblem, so the values can be used
-    # directly.
-    try:
-        # Some are transposed because CVXOPT uses column major arrays
-        q = cvxopt.matrix(prob.q.T)
-        A = cvxopt.matrix(prob.A) if prob.A is not None else None
-        b = cvxopt.matrix(prob.b.T) if prob.b is not None else None
-        G = cvxopt.matrix(prob.G) if prob.G is not None else None
-        h = cvxopt.matrix(prob.h.T) if prob.h is not None else None
+    # Importantly, CVXOPT can only solve problems in the l, q and s cones.
 
-        if prob.is_qp:
-            P = cvxopt.matrix(prob.P)
+    if prob.constraints["b"]:
+        raise NotImplementedError("Automatic conversion of box constraints "
+                                  "into linear constraints is not supported yet.")
+
+    if prob.constraints["ep"] or prob.constraints["ep*"]:
+        raise NotSupportedBySolver("CVXOPT cannot solve problems with the exponential cone.")
+
+    if prob.constraints["p"] or prob.constraints["p*"]:
+        raise NotSupportedBySolver("CVXOPT cannot solve problems with the power cone.")
+
+    # Cost function
+    q = cvxopt.matrix(prob.q) if prob.q is not None else None
+    P = cvxopt.matrix(prob.P) if prob.P is not None else None
+
+    # Constraint matrices
+    A, b, G, h = None, None, None, None
+
+    # Equality constraints
+    A_rows, b_rows = [], []
+    for (linear, constant) in prob.constraints["z"]:
+        b_rows.append(constant)
+        A_rows.append(np.hstack(tuple(vectorize_matrix(m) for m in linear)).T)
+        
+    if A_rows:
+        # b is multiplied by -1 because it is on the RHS
+        b = cvxopt.matrix(np.vstack(b_rows)) * -1
+        A = cvxopt.matrix(np.vstack(A_rows))
 
+    # Inequality, second order and PSD constraints
+    G_rows, h_rows = [], []
+
+    for (linear, constant) in prob.constraints["l"]:
+        h_rows.append(constant)
+        G_rows.append(np.hstack(tuple(m for m in linear)).T)
+
+    if prob.constraints["q"]:
+        raise NotImplementedError("SOC constraints are not implemented yet.")
+
+    for (linear, constant) in prob.constraints["s"]:
+        h_rows.append(vectorize_matrix(constant))
+        G_rows.append(np.hstack(tuple(vectorize_matrix(m) for m in linear)))
+
+    # pprint(f"{G_rows = }")
+    # pprint(f"{h_rows = }")
+        
+    if G_rows:
+        # h is multiplied by -1 because it is on the RHS
+        h = cvxopt.matrix(np.vstack(h_rows).reshape((-1, 1))) * -1
+        G = cvxopt.matrix(np.vstack(G_rows))
+
+
+    # Solve the problem
+    # pprint({"A": A, "b": b, "G": G, "h": h})
+
+    try:
+        if prob.is_qp:
+            assert P is not None, "Solving LP with QP solver! Something has gone wrong"
             info = cvxopt.solvers.coneqp(P=P, q=q, G=G, h=h, A=A, b=b,
                                          dims=prob.dims, *args, **kwargs)
         else:
+            assert q is not None, "LP has no objective! Something is very wrong"
             info = cvxopt.solvers.conelp(c=q, G=G, h=h, A=A, b=b,
                                          *args, **kwargs)
 
+    except AssertionError as e:
+        raise e
+
     except Exception as e:
         raise SolverError("CVXOpt can't solve this problem, "
-                          "see previous exception why.") from e
+                          "see previous exception for details on why.") from e
+
 
-    results = {}
+    if info['x'] is None:
+        return {}, CVXOPTInfo(info)
 
-    i = 0
+    results, i = {}, 0
     for variable in prob.variables:
         num_indices = math.prod(variable.shape)
         values = np.array(info['x'][i:i+num_indices]).reshape(variable.shape)
diff --git a/sumofsquares/solver/mosek.py b/sumofsquares/solver/mosek.py
index e0b73fc..763555b 100644
--- a/sumofsquares/solver/mosek.py
+++ b/sumofsquares/solver/mosek.py
@@ -39,8 +39,8 @@ def setup(license_file: Path | str | None = None):
     MOSEK_ENV = mosek.Env(licensefile=str(license_file))
 
 
-def solve_sos_cone(prob: Problem, verbose: bool = False,
-                   *args, **kwargs) -> tuple[dict[OptVariable, float], MOSEKInfo]:
+def solve_cone(prob: Problem, verbose: bool = False,
+               *args, **kwargs) -> tuple[dict[OptVariable, float], MOSEKInfo]:
     r"""
     Solve a conic problem in the cone of SOS polynomials
     :math:`\mathbf{\Sigma}_d(x)` using MOSEK.
diff --git a/sumofsquares/solver/scs.py b/sumofsquares/solver/scs.py
index 3106df2..4f17a24 100644
--- a/sumofsquares/solver/scs.py
+++ b/sumofsquares/solver/scs.py
@@ -6,6 +6,7 @@ from __future__ import annotations
 import scs
 
 from collections import UserDict
+from numpy.typing import NDArray
 from typing import TYPE_CHECKING
 
 from ..abc import Problem, SolverInfo
@@ -25,8 +26,16 @@ class SCSInfo(SolverInfo, UserDict):
         return self[key]
 
 
-def solve_sos_cone(prob: ConicProblem, verbose: bool = False,
-                   *args, **kwargs) -> tuple[dict[OptVariable, float], SCSInfo]:
+def vectorize_matrix(m: NDArray) -> NDArray:
+    r"""
+    Vectorize a symmetric matrix for SCS by storing only a half of the entries
+    and multiplying the off-diaognal elements by :math:`\sqrt{2}`.
+    """
+    raise NotImplementedError
+
+
+def solve_cone(prob: ConicProblem, verbose: bool = False,
+               *args, **kwargs) -> tuple[dict[OptVariable, float], SCSInfo]:
     r"""
     Solve a conic problem in the cone of SOS polynomials
     :math:`\mathbf{\Sigma}_d(x)` using SCS.