Lay out new structure for SumOfSquares package

6 files changed, 335 insertions, 1 deletions

diff --git a/sumofsquares/__init__.py b/sumofsquares/__init__.py
index da35128..d71d38f 100644
--- a/sumofsquares/__init__.py
+++ b/sumofsquares/__init__.py
@@ -1 +1,104 @@
-from sumofsquares.optvariable import from_names
\ No newline at end of file
+"""
+Sum of Squares Package
+======================
+
+This is a package to solve sum-of-squares optimization problems.
+
+Module Overview
+---------------
+
+Class diagram, arrows denote inheritance, whereas diamonds indicate composition.
+::
+                                                                                                                                          
+  ┏━━sumofsquares package━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓                                  
+  ┃ ┌─.abc──────────────────────────────────────────────────────────────────────────────────────────┐  ┃                                  
+  ┃ │ ┌────────────┐      ┌──────────────┐ ┌─────────┐    ┌───────┐   ┌──────────────┐              │  ┃                                  
+  ┃ │ │ SolverInfo │◇─┐   │ Solver(Enum) │ │ Problem │    │  Set  │   │  Constraint  │              │  ┃                                  
+  ┃ │ └────────────┘  │   └──────────────┘ └─────────┘    └───────┘   └──────────────┘              │  ┃                                  
+  ┃ │        ▲        │                         ▲             ▲               ▲                     │  ┃                                  
+  ┃ └────────┼────────┼─────────────────────────┼─────────────┼───────────────┼─────────────────────┘  ┃                                  
+  ┃          │        │            ┌────────────┘          ┌──┘               ├─────────────┐          ┃                                  
+  ┃          │        │            │                       │                  │             │          ┃                                  
+  ┃  ┌─.cvxop│t─────┐ │ ┌─.problems┼──────────┐ ┌─.constrai│nts───────────────┼─────────────┼───────┐  ┃                                  
+  ┃  │       │      │ │ │          │          │ │          │                  │             │       │  ┃                                  
+  ┃  │ ┌──────────┐ │ │ │ ┌─────────────────┐ │ │ ┌────────────────┐    ┌───────────┐ ┌───────────┐ │  ┃                                  
+  ┃  │ │CVXOptInfo│ │ │ │ │PutinarSOSProblem│ │ │ │ ArchimedeanSet │◇───│NonNegative│ │EqualToZero│ │  ┃                                  
+  ┃  │ └──────────┘ │ │ │ └─────────────────┘ │ │ └────────────────┘    └───────────┘ └───────────┘ │  ┃                                  
+  ┃  │              │ │ │ ┌────────┐          │ │                                                   │  ┃                                  
+  ┃  │              │ └─┼─│ Result │          │ │                                                   │  ┃                                  
+  ┃  │              │   │ └────────┘          │ │                                                   │  ┃                                  
+  ┃  └──────────────┘   └─────────────────────┘ └───────────────────────────────────────────────────┘  ┃                                  
+  ┃                                                                     ┌─.variable─────────────────┐  ┃  ┏━━━polymatrix package━━━━━━━━━┓
+  ┃                                                                     │                           │  ┃  ┃                              ┃
+  ┃                                                                     │    ┌───────────┐          │  ┃  ┃  ┌──────────┐   ┌──────────┐ ┃
+  ┃                                                                     │    │OptVariable│──────────┼──╋──╋─▶│ Variable │   │Expression│ ┃
+  ┃                                                                     │    └───────────┘          │  ┃  ┃  └──────────┘   └──────────┘ ┃
+  ┃                                                                     │          ▲                │  ┃  ┃                       │      ┃
+  ┃                                                                     │          │                │  ┃  ┃             ┌─────────┘      ┃
+  ┃                                                                     │          │                │  ┃  ┃             ◇                ┃
+  ┃                                                                     │ ┌────────────────┐        │  ┃  ┃  ┌─────────────────────┐     ┃
+  ┃                                                                     │ │OptVariableMixin│────────┼──╋──╋─▶│ ExpressionBaseMixin │     ┃
+  ┃                                                                     │ └────────────────┘        │  ┃  ┃  └─────────────────────┘     ┃
+  ┃                                                                     └───────────────────────────┘  ┃  ┃                              ┃
+  ┃                                                                                                    ┃  ┃                              ┃
+  ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛  ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
+
+"""
+from typing import Iterable
+
+import polymatrix as poly
+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 .problems import PutinarSOSProblem
+from .utils import partition
+from .variable import OptVariable, from_names
+
+
+def make_sos_constraint(expr: Expression, domain: Set | None = None) -> NonNegative:
+    return NonNegative(expr, domain)
+
+
+def make_problem(
+        cost: Expression,
+        constraints: Iterable[Constraint] = (),
+        solver: Solver = Solver.CVXOPT,
+        state: ExpressionState | None = None
+    ) -> Problem:
+    """
+    Create a sum-of-squares optimization problem.
+    """
+    if not state:
+        state = poly.make_state()
+
+        # Convert to polymatrix object to extract variables
+        state, _ = cost.apply(state)
+        for constr in constraints:
+            state, _ = constr.expression.apply(state)
+
+    # Collect decision variables from state object
+    def is_optvariable(v):
+        return isinstance(v, OptVariable)
+
+    poly_variables, opt_variables = partition(is_optvariable, state.indices.keys())
+    return PutinarSOSProblem(
+            cost=cost,
+            constraints=constraints,
+            variables=tuple(opt_variables),
+            polynomial_variables=tuple(poly_variables),
+            solver=solver,
+            state=state)
+
+
+def solve_problem(
+        cost: Expression,
+        constraints: Iterable[Constraint] = (),
+        solver: Solver = Solver.CVXOPT
+    ) -> tuple[Problem, Result]:
+    """
+    Solve a sum-of-squares optimization problem.
+    """
+    prob = make_problem(cost, constraints, solver)
+    return prob, prob.solve()
diff --git a/sumofsquares/abc.py b/sumofsquares/abc.py
new file mode 100644
index 0000000..482ba7f
--- /dev/null
+++ b/sumofsquares/abc.py
@@ -0,0 +1,86 @@
+""" Abstract base classes for Sum of Squares package """
+from __future__ import annotations
+from typing import Sequence, Any
+
+from abc import ABC, abstractmethod
+from enum import Enum, auto
+
+from polymatrix.expression.expression import Expression
+from polymatrix.expressionstate.abc import ExpressionState
+from polymatrix.variable.abc import Variable
+
+from sumofsquares.variable import OptVariable
+
+
+class Solver(Enum):
+    """ Enumeration for the supported solvers. """
+    CVXOPT = auto()
+    # MOSEK = auto()
+
+
+class SolverInfo(ABC):
+    """ Type that information returned by a specific solver. """
+
+
+class Result(ABC):
+    """ Result of an optimization problem. """
+    @abstractmethod
+    def value_of(self, var: OptVariable) -> float:
+        """ Retrieve value of variable. """
+
+
+class Set(ABC):
+    """ A set. In this context always a subset of :math:`\mathbf{R}^q` """
+
+
+class Constraint(ABC):
+    r"""
+    Optimization constraint. 
+    """
+    @property
+    @abstractmethod
+    def expression(self) -> Expression:
+        """ Expression under the constraint. """
+
+
+class Problem(ABC):
+    """ Optimization Problem. """
+    
+    # --- Properties ---
+
+    @property
+    @abstractmethod
+    def cost(self) -> Expression:
+        """ Expression of cost function. """
+
+    @property
+    @abstractmethod
+    def constraints(self) -> Sequence[Constraint]:
+        """ Optimization constraints. """
+
+    @property
+    @abstractmethod
+    def variables(self) -> Sequence[OptVariable]:
+        """ Optimization variables of the problem. """
+
+    @property
+    @abstractmethod
+    def polynomial_variables(self) -> Sequence[Variable]:
+        """ Polynomial variables, or varaibles that are not optimization varaibles. """
+
+    @property
+    @abstractmethod
+    def solver(self) -> Solver:
+        """ Which solver was / will be used to solve the problem. """
+
+
+    @property
+    @abstractmethod
+    def state(self) -> ExpressionState:
+        """ Expression state used to evaluate constraints and costs. """
+
+    # --- Methods ---
+
+    @abstractmethod
+    def solve(self) -> Result:
+        """ Solve the optimization problem. """
diff --git a/sumofsquares/canon.py b/sumofsquares/canon.py
new file mode 100644
index 0000000..4918758
--- /dev/null
+++ b/sumofsquares/canon.py
@@ -0,0 +1,11 @@
+""" Canonicalization of Problems """
+from .abc import Problem
+
+
+def log_canon(prob: Problem) -> Problem:
+    raise NotImplementedError
+
+
+def log_det_canon(prob: Problem) -> Problem:
+    raise NotImplementedError
+
diff --git a/sumofsquares/constraints.py b/sumofsquares/constraints.py
new file mode 100644
index 0000000..3634e61
--- /dev/null
+++ b/sumofsquares/constraints.py
@@ -0,0 +1,50 @@
+from dataclassabc import dataclassabc
+
+from polymatrix.expression.expression import Expression
+
+from .abc import Set, Constraint
+
+
+class ArchimedeanSet(Set, tuple[Expression]):
+    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
+    is a set of the form
+    ..math::
+
+        \mathbf{K} = \{
+            x : g_i(x) \geq 0,
+                g_i \in \mathbf{R}(x),
+                x \in \mathbf{R}^q,
+                i = 1, \ldots, N
+        \}
+        \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.
+    """
+
+
+@dataclassabc(frozen=True)
+class EqualToZero(Constraint):
+    """ Constrain an expression to be equal to zero. """
+    expression: Expression
+
+
+@dataclassabc(frozen=True)
+class NonNegative(Constraint):
+    """
+    Constrain an expression to be non-negative in a certain domain. If 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.
+    """
+    expression: Expression 
+    domain: ArchimedeanSet | None = None
diff --git a/sumofsquares/problems.py b/sumofsquares/problems.py
new file mode 100644
index 0000000..035e756
--- /dev/null
+++ b/sumofsquares/problems.py
@@ -0,0 +1,74 @@
+""" 
+Specific optimization problems that are related to Sum of Squares.
+"""
+from __future__ import annotations
+
+from itertools import filterfalse
+from dataclasses import replace
+from dataclassabc import dataclassabc
+from typing import Any, Sequence
+from typing_extensions import override
+
+import polymatrix as poly
+from polymatrix.expression.expression import Expression
+from polymatrix.expressionstate.abc import ExpressionState
+from polymatrix.variable.abc import Variable
+
+from .abc import Problem, Constraint, Solver, Result
+from .constraints import NonNegative
+from .cvxopt import solve_sos_cone as cvxopt_solve_sos_cone
+from .variable import OptVariable
+
+
+@dataclassabc(frozen=True)
+class SOSResult(Result):
+    values: dict[OptVariable, float]
+    solver_info: Any
+
+    @override
+    def value_of(self, var: OptVariable) -> float:
+        return self.values[var]
+
+
+@dataclassabc(frozen=True)
+class PutinarSOSProblem(Problem):
+    r"""
+    SOS problem with quadratic cost, linear constraints and optinally with
+    positivity constraints over an archimedean set :math:`\mathbf{K}`. To
+    constrain the positivity over :math:`\mathbf{K}` this problem uses
+    Putinar's positivstellensatz.
+    """
+    cost: Expression
+    constraints: Sequence[Constraint]
+    variables: Sequence[OptVariable]
+    polynomial_variables: Sequence[Variable]
+    solver: Solver
+    state: ExpressionState
+
+    def after_positivstellensatz(self) -> Problem:
+        """
+        Apply Putinar's positivstellensatz to convert the SOS constraints
+        into coefficient matching constraints 
+        """
+        def need_positivstellensatz(c):
+            return isinstance(c, NonNegative) and c.domain
+
+        # Add langrange-multiplier-like polynomials to SOS constraints, while
+        # leaving the rest untouched.
+        variables = list(self.variables)
+        constraints = list(filterfalse(need_positivstellensatz, self.constraints))
+        for c in filter(need_positivstellensatz, self.constraints):
+            # FIXME: implement
+            raise NotImplementedError
+
+        return replace(self, constraints=constraints, variables=variables)
+
+
+    @override
+    def solve(self) -> Result:
+        if not self.solver == Solver.CVXOPT:
+            raise NotImplementedError(f"Cannot solve problem {self} using solver {str(self.solver)}.")
+
+        prob = self.after_positivstellensatz()
+        result, info = cvxopt_solve_sos_cone(prob)
+        return SOSResult(result, info)
diff --git a/sumofsquares/utils.py b/sumofsquares/utils.py
new file mode 100644
index 0000000..c18bd08
--- /dev/null
+++ b/sumofsquares/utils.py
@@ -0,0 +1,10 @@
+from itertools import tee, filterfalse
+
+def partition(predicate, iterable):
+    """Partition entries into false entries and true entries.
+
+    If *predicate* is slow, consider wrapping it with functools.lru_cache().
+    """
+    # partition(is_odd, range(10)) → 0 2 4 6 8   and  1 3 5 7 9
+    t1, t2 = tee(iterable)
+    return filterfalse(predicate, t1), filter(predicate, t2)