Prepare structure for logdet canon

3 files changed, 135 insertions, 45 deletions

diff --git a/sumofsquares/constraints.py b/sumofsquares/constraints.py
index 37486e1..61be087 100644
--- a/sumofsquares/constraints.py
+++ b/sumofsquares/constraints.py
@@ -69,3 +69,13 @@ class NonNegative(Constraint):
             s += f"\n\t\t over {self.domain}"
         return s
 
+
+@dataclassabc(frozen=True)
+class PositiveSemiDefinite(Constraint):
+    """ Constrain a matrix to be positive semidefinite. """
+    expression: Expression 
+
+
+@dataclassabc(frozen=True)
+class ExponentialCone(Constraint):
+    expression: Expression
diff --git a/sumofsquares/cvxopt.py b/sumofsquares/cvxopt.py
index 4cd803c..6a1339f 100644
--- a/sumofsquares/cvxopt.py
+++ b/sumofsquares/cvxopt.py
@@ -11,7 +11,7 @@ from polymatrix.expression.expression import Expression
 from polymatrix.polymatrix.index import MonomialIndex, VariableIndex
 
 from .abc import Problem, SolverInfo
-from .constraints import NonNegative, EqualToZero
+from .constraints import NonNegative, EqualToZero, PositiveSemiDefinite, ExponentialCone
 from .variable import OptVariable
 
 
@@ -25,7 +25,6 @@ class CVXOPTInfo(SolverInfo, UserDict):
         return self[key]
 
 
-
 def solve_sos_cone(prob: Problem, verbose: bool = False, *args, **kwargs) -> tuple[dict[OptVariable, float], CVXOPTInfo]:
     r"""
     Solve a conic problem in the cone of SOS polynomials :math:`\mathbf{\Sigma}_d(x)`.
@@ -85,47 +84,54 @@ def solve_sos_cone(prob: Problem, verbose: bool = False, *args, **kwargs) -> tup
 
     x = poly.v_stack(prob.polynomial_variables)
     for c in prob.constraints:
-        # Convert constraints to affine expression in decision variables
-        constr = poly.to_affine(c.expression.quadratic_in(x).read(prob.state))
-        if constr.degree > 1:
-            raise ValueError("CVXOpt can solve problem with constraints that are "
-                            f"at most linear, 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.values():
-            # Get the affine coefficient associated to this variable
-            linear = MonomialIndex((v,))
-            v_coeff = constr.affine_coefficient(linear)
-
-            # 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 CVXOPT format. "
-                                 "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)
+        if isinstance(c, PositiveSemiDefinite):
+            raise NotImplementedError
+
+        elif isinstance(c, ExponentialCone):
+            raise NotImplementedError
+
+        elif isinstance(c, EqualToZero | NonNegative):
+            # Convert constraints to affine expression in decision variables
+            constr = poly.to_affine(c.expression.quadratic_in(x).read(prob.state))
+            if constr.degree > 1:
+                raise ValueError("CVXOpt can solve problem with constraints that are "
+                                f"at most linear, 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.values():
+                # Get the affine coefficient associated to this variable
+                linear = MonomialIndex((v,))
+                v_coeff = constr.affine_coefficient(linear)
+
+                # 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 CVXOPT format. "
+                                     "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)
 
         else:
             raise NotImplementedError(f"Cannot convert constraint of type {type(c)} "
@@ -139,17 +145,28 @@ def solve_sos_cone(prob: Problem, verbose: bool = False, *args, **kwargs) -> tup
     b = cvxopt.matrix(np.hstack(b_rows)) if b_rows else None
     dims = {'l': l_dims, 'q': q_dims, 's': s_dims}
 
+    # print(q)
+    # print(G)
+    # print(h)
+    # print(A)
+    # print(b)
+    # print(dims)
+
     cvxopt.solvers.options['show_progress'] = verbose
 
     if not any((G, h, A, b)):
         raise ValueError("Optimization problem is unconstrained!")
 
+    # Add matrices for debugging
+    # info = {}
+    info = {"q": q, "G": G, "h": h, "A": A, "b": b, "dims": dims}
+
     if is_qp:
         P = cvxopt.matrix(P)
-        info = cvxopt.solvers.coneqp(P=P, q=q, G=G, h=h, A=A, b=b,
+        info |= cvxopt.solvers.coneqp(P=P, q=q, G=G, h=h, A=A, b=b,
                                      dims=dims, *args, **kwargs)
     else:
-        info = cvxopt.solvers.conelp(c=q, G=G, h=h, A=A, b=b,
+        info |= cvxopt.solvers.conelp(c=q, G=G, h=h, A=A, b=b,
                                      *args, **kwargs)
 
     results = {}
diff --git a/sumofsquares/expression.py b/sumofsquares/expression.py
new file mode 100644
index 0000000..08b4fbd
--- /dev/null
+++ b/sumofsquares/expression.py
@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Iterable
+from typing_extensions import override
+from dataclassabc import dataclassabc
+
+from polymatrix.expression.expression import ExpressionBaseMixin
+from polymatrix.expressionstate.abc import ExpressionState
+
+from .abc import Constraint
+from .constraints import NonNegative
+
+
+class SOSExpressionBaseMixin(ExpressionBaseMixin):
+    @override
+    def apply(self, state: ExpressionState) -> tuple[ExpressionState, ExpressionBaseMixin]:
+        # FIXME: think of name for this exception
+        raise RuntimeError
+
+
+    @abstractmethod
+    def recast(self) -> tuple[ExpressionBaseMixin, Iterable[Constraint]]:
+        """
+        Recast the expression into a form that can be directly used for
+        optimization, possibly by introducing new constraints.
+
+        The return values are the new expression to be minimized and the new
+        constraints that have to be added.
+        """
+
+
+class LogDetMixin(SOSExpressionBaseMixin):
+    """ Compute the sum of the logarithm of the eigenvalues. """
+
+    @property
+    @abstractmethod
+    def underlying(self) -> ExpressionBaseMixin:
+        """" Take the logdet of this expression """
+
+    @override
+    def recast(self) -> tuple[ExpressionBaseMixin, Iterable[Constraint]]:
+        # The problem
+        #
+        #   minimize logdet(A)
+        #
+        # is equivalent to solving
+        #
+        #   minimize      -t
+        #
+        #   subject to    [ A    Z       ]
+        #                 [ Z.T  diag(Z) ] >= 0
+        #
+        #                 Z lower triangular
+        #                 t <= sum_i log(Z[i, i])
+
+
+@dataclassabc(froze=True)
+class LogDetImpl(LogDetMixin):
+    underlying: ExpressionBaseMixin
+
+    def __str__(self):
+        return f"logdet({self.underlying})"