1 files changed, 125 insertions, 5 deletions
diff --git a/sumofsquares/canon.py b/sumofsquares/canon.py
index 4918758..e5960e7 100644
--- a/sumofsquares/canon.py
+++ b/sumofsquares/canon.py
@@ -1,11 +1,131 @@
 """ Canonicalization of Problems """
-from .abc import Problem
+import math
+import polymatrix as poly
 
+from abc import abstractmethod
+from dataclassabc import dataclassabc
+from dataclasses import replace
+from typing import Sequence
+from typing_extensions import override
 
-def log_canon(prob: Problem) -> Problem:
-    raise NotImplementedError
+from polymatrix.expression.from_ import Expression, from_statemonad
+from polymatrix.expressionstate import ExpressionState
+from polymatrix.statemonad import init_state_monad
+from polymatrix.polymatrix.mixins import PolyMatrixMixin
+from polymatrix.variable import Variable
 
+from polymatrix.expression.init import (
+        init_concatenate_expr,
+        init_diag_expr,
+        init_lower_triangular_expr,
+        init_transpose_expr)
+
+from .abc import Constraint, Result
+from .constraints import NonNegative, PositiveSemiDefinite, ExponentialCone
+from .error import AlgebraicError
+from .problems import SOSProblem, InternalSOSProblem
+from .variable import OptVariable, init_opt_variable_expr, from_name as opt_variable_from_name
+from .utils import partition
+
+
+class Canonicalization:
+    """ Transform a problem into another problem. """
+
+    def __str__(self):
+        s = f"Apply canonicalization procedure {self.__class__.__qualname__} to\n"
+        return s + str(self.problem)
+
+    @property
+    @abstractmethod
+    def problem(self) -> SOSProblem:
+        """ The problem that will be canonicalized """
+
+    def solve(self, verbose: bool = False, state: ExpressionState | None = None) -> Result:
+        if state is None:
+            state = poly.make_state()
+
+        state, internal_prob = self.apply(state)
+        return internal_prob.solve(verbose)
+
+    @abstractmethod
+    def apply(self, state: ExpressionState) -> tuple[ExpressionState, InternalSOSProblem]:
+        """ Apply the canonicalization procedure """
+
+
+@dataclassabc(frozen=True)
+class PutinarPSatz(Canonicalization):
+    problem: SOSProblem
+
+    @staticmethod
+    def make_multiplier(variables: Sequence[OptVariable],
+                        polynomial_variables: Sequence[Variable],
+                        constr: Expression,
+                        domain_poly: Expression,
+                        multiplier_name: str) -> Expression:
+        r"""
+        Make a multiplier polynomial for a domain polynomial of a constraint.
+
+        Since Putinar's PSatz does not explicity tell us the order of the
+        multiplier this function will make a multiplier, let's call it
+        :math:`\gamma`, such that :math:`\deg(\gamma * p) = \deg(c)`, wherein
+        :math:`p` is the domain polynomial and :math:`c` is the contraint
+        polynomial.
+        """
+        # Number of variables in the original problem
+        q = sum(math.prod(v.shape) for v in variables)
+        x = poly.v_stack((1,) + tuple(polynomial_variables))
+
+        def make_multiplier_later(state: ExpressionState) -> tuple[ExpressionState, PolyMatrixMixin]:
+            state, constr_pm = constr.expression.degree().apply(state)
+            state, domain_poly_pm = domain_poly.degree().apply(state)
+            
+            constr_deg = constr_pm.at(0, 0).constant()
+            domain_poly_deg = domain_poly_pm.at(0, 0).constant()
+
+            # degree of multiplier
+            d = constr_deg - domain_poly_deg
+            if d < 0:
+                raise AlgebraicError("Cannot create Positivstellensatz multiplier because "
+                                     f"constraint polynomial has degree {constr_deg} and domain "
+                                     f"polynomial {domain_poly} has degree {domain_poly_deg}")
+            
+            n = sum(math.comb(k + q - 1, k) for k in range(0, d+1))
+
+            # TODO: check that there is not a variable with this name already
+            # Coefficients of the multiplier polynomial 
+            c = opt_variable_from_name(multiplier_name, shape=(1,n))
+            multiplier = c @ x.combinations(tuple(range(d +1)))
+
+            return multiplier.apply(state)
+        return poly.from_statemonad(init_state_monad(make_multiplier_later))
+
+    # TODO: rewrite this function, this is adapted from how it was done
+    # previously, and it works but it's not efficient.
+    @abstractmethod
+    def apply(self, state: ExpressionState) -> tuple[ExpressionState, InternalSOSProblem]:
+
+        def need_positivstellensatz(c: Constraint) -> bool:
+            return isinstance(c, NonNegative) and (c.domain is not None)
+
+        constraints, to_process = partition(need_positivstellensatz, self.problem.constraints)
+        constraints = tuple(constraints) # because it is a generator
+
+        state, prob = self.problem.apply(state)
+
+        new_constraints: list[Constraint] = list(constraints)
+        for i, constr in enumerate(to_process):
+            new_constr = constr.expression
+            for domain_poly in constr.domain.polynomials:
+                m = PutinarPSatz.make_multiplier(prob.variables,
+                                                 prob.polynomial_variables,
+                                                 constr, domain_poly, rf"\gamma_{i}")
+
+                # Subtract multiplier from expression and impose that it is also SOS
+                new_constr -= m * domain_poly
+                new_constraints.append(NonNegative(m))
+
+            new_constraints.append(NonNegative(new_constr))
+
+        return replace(self.problem, constraints=new_constraints).apply(state)
 
-def log_det_canon(prob: Problem) -> Problem:
-    raise NotImplementedError