Clean up canonicalization classes

1 files changed, 80 insertions, 59 deletions

diff --git a/sumofsquares/canon.py b/sumofsquares/canon.py
index 0874e2c..17ad187 100644
--- a/sumofsquares/canon.py
+++ b/sumofsquares/canon.py
@@ -31,6 +31,24 @@ class Canonicalization:
         s = f"Apply canonicalization procedure {self.__class__.__qualname__} to\n"
         return s + str(self.problem)
 
+    @abstractmethod
+    def __call__(self, problem : SOSProblem | None = None) -> SOSProblem:
+        # TODO: is this a good idea? probably not but I don't want to introduce
+        # a new type of class for "horizontal" function in the commutative
+        # diagram. The idea is to make canon objects behave like functions
+        # so that given a problem P, you can do Canon2(Canon1(P)), but this
+        # needs a default implementation if the transformation is a stateful
+        # computation, which would need a statemonad, defeating the purpose of
+        # creating this class in the first place.
+        """
+        Apply the canonicalization procedure to the given problem, or to
+        `self.problem` if `problem` is None.
+
+        Override this method to write transformation if it can be written
+        without having to apply the state. I.e. it is not stateful in the
+        problem.
+        """
+
     @property
     @abstractmethod
     def problem(self) -> SOSProblem:
@@ -45,7 +63,12 @@ class Canonicalization:
 
     @abstractmethod
     def apply(self, state: ExpressionState) -> tuple[ExpressionState, InternalSOSProblem]:
-        """ Apply the canonicalization procedure """
+        """
+        Apply the canonicalization procedure.
+
+        Override this method to write a transformation if it is stateful in the
+        problem, and hence "needs to be performed later".
+        """
 
 
 @dataclassabc(frozen=True)
@@ -58,74 +81,59 @@ 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_state_monad(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
+    @override
     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
+        constraints = tuple(constraints)
+        to_process = tuple(to_process)
+
+        # Evaluate expression to gather all variables
+        for c in to_process:
+            state, _ = c.expression.apply(state)
+            for p in c.domain.polynomials:
+                state, _ = p.apply(state)
 
-        state, prob = self.problem.apply(state)
+        state, prob = replace(self.problem, constraints=constraints).apply(state)
 
+        # Create new constraints
         new_constraints: list[Constraint] = list(constraints)
         for i, constr in enumerate(to_process):
             new_constr = constr.expression
             for domain_poly in constr.domain.polynomials:
-                # FIXME: need to check that there is not a \gamma_{i}!
-                m = PutinarPSatz.make_multiplier(prob.variables,
-                                                 prob.polynomial_variables,
-                                                 constr, domain_poly, rf"\gamma_{i}")
+                # 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.
+
+                # FIXME: create an extension poly_degree and opt_degree to get
+                # the maximul degree of the polynomial and the degree of the
+                # optimization variables respectively.
+                constr_deg = constr.expression.eval({v: 1. for v in prob.variables}).degree()
+                domain_deg = domain_poly.eval({v: 1. for v in prob.variables}).degree()
+
+                d = constr_deg - domain_deg
+                x = poly.v_stack((1,) + prob.polynomial_variables)
+
+                # FIXME: proper error here
+                assert d.read(state).scalar().constant() >= 0, \
+                        "Degree of domain polynomial is higher than constraint polynomial!"
+
+                # FIXME: need to check that there is not a \gamma_{i} already!
+                monomials = x.combinations(poly.arange(d + 1))
+                c = opt_variable_from_name(rf"\gamma_{i}", shape=monomials.shape)
+
+                multiplier = c.T @ monomials
 
                 # Subtract multiplier from expression and impose that it is also SOS
-                new_constr -= m * domain_poly
-                new_constraints.append(NonNegative(m))
+                new_constr -= multiplier * domain_poly
+                new_constraints.append(NonNegative(multiplier))
 
             new_constraints.append(NonNegative(new_constr))
         return replace(self.problem, constraints=new_constraints).apply(state)
@@ -133,7 +141,13 @@ class PutinarPSatz(Canonicalization):
 
 @dataclassabc(frozen=True)
 class LogDet(Canonicalization):
+    # TODO: reconsider extension of polymatrix Expression objects to introduce
+    # logdet expression. The current solution works but is not ideal, as it may
+    # be a bit counterintuitive.
     """
+    Create a new problem, whose cost function is the logdet of the given
+    problem's cost function.
+
     The problem
     
       maximize logdet(A)
@@ -171,8 +185,11 @@ class LogDet(Canonicalization):
     problem: SOSProblem
 
     @override
-    def apply(self, state: ExpressionState) -> tuple[ExpressionState, InternalSOSProblem]:
-        A = self.problem.cost
+    def __call__(self, problem : SOSProblem | None = None) -> SOSProblem:
+        if problem is None:
+            problem = self.problem
+
+        A = problem.cost
 
         # FIXME: should check for name clashes in state object
         # i.e. what if there is already a variable named u_logdet / Z_logdet?
@@ -188,7 +205,11 @@ class LogDet(Canonicalization):
         E = poly.h_stack((Z.diag(), poly.ones((n, 1)), u))
 
         new_cost = - u.T.sum()
-        new_constraints = [PositiveSemiDefinite(Q), ExponentialCone(E)]
+        new_constraints = tuple(problem.constraints) + (PositiveSemiDefinite(Q), ExponentialCone(E))
+
+        return replace(problem, cost=new_cost, constraints=new_constraints)
 
-        raise NotImplementedError
 
+    @override
+    def apply(self, state: ExpressionState) -> tuple[ExpressionState, InternalSOSProblem]:
+        return self().apply(state)