Fix bugs in PutinarPSatz, better names (variables -> symbols)

2 files changed, 92 insertions, 44 deletions

diff --git a/sumofsquares/canon.py b/sumofsquares/canon.py
index 8c4b541..9ba0639 100644
--- a/sumofsquares/canon.py
+++ b/sumofsquares/canon.py
@@ -9,11 +9,12 @@ from typing_extensions import override
 from polymatrix.expression.from_ import Expression
 from polymatrix.expression.init import init_variable_expr
 from polymatrix.expressionstate import ExpressionState
+from polymatrix.polymatrix.index import VariableIndex
 
 from .abc import Problem, Solver, Constraint, Result
 from .constraints import NonNegative, PositiveSemiDefinite, ExponentialCone
 from .problems import SOSProblem, InternalSOSProblem
-from .variable import from_name as opt_variable_from_name
+from .variable import OptSymbol, from_name as opt_variable_from_name
 from .utils import partition
 
 
@@ -26,7 +27,7 @@ class Canonicalization(Problem):
         s = f"Apply canonicalization procedure {self.__class__.__qualname__} to\n"
         return s + str(self.problem)
 
-    def __call__(self, problem : SOSProblem | None = None) -> SOSProblem:
+    def __call__(self, problem : SOSProblem) -> 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
@@ -35,8 +36,7 @@ class Canonicalization(Problem):
         # 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.
+        Apply the canonicalization procedure to the given problem.
 
         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
@@ -59,8 +59,7 @@ class Canonicalization(Problem):
     @property
     @override
     def solver(self) -> Solver:
-        # TODO: error message
-        raise AttributeError
+        return self.problem.solver
 
     @property
     @abstractmethod
@@ -92,6 +91,12 @@ class PutinarPSatz(Canonicalization):
 
     problem: SOSProblem
 
+    @property
+    @override
+    def cost(self) -> Expression:
+        # PSatz does not change cost function
+        return self.problem.cost
+
     @override
     def apply(self, state: ExpressionState) -> tuple[ExpressionState, InternalSOSProblem]:
 
@@ -102,18 +107,57 @@ class PutinarPSatz(Canonicalization):
         constraints = tuple(constraints)
         to_process = tuple(to_process)
 
-        # Evaluate expression to gather all variables
+        # Variables that are involved in the problem
+        variable_indices: set[VariableIndex] = set()
+
+        # Evaluate expressions of constraints tha need to be processed
+        # to gather all variables
         for c in to_process:
-            state, _ = c.expression.apply(state)
+            state, pm = c.expression.apply(state)
+            variable_indices |= pm.variables()
+
             for p in c.domain.polynomials:
-                state, _ = p.apply(state)
+                state, pm = p.apply(state)
+                variable_indices |= pm.variables()
+
+        # Get symbols of the variables
+        to_process_symbols = set(state.get_symbol_from_variable_index(v)
+                                 for v in variable_indices)
+
+        # Separate between symbols of optimization variables and those of polynomial variables
+        def is_opt(v):
+            return isinstance(v, OptSymbol)
 
+        polynomial_symbols, symbols = partition(is_opt, to_process_symbols)
+
+        # Get variables from rest of the problem
         state, prob = replace(self.problem, constraints=constraints).apply(state)
+        symbols = tuple(set(symbols) | set(prob.symbols))
+        polynomial_symbols = tuple(set(polynomial_symbols) | set(prob.polynomial_symbols))
+
+        # Vector of polynomial variables
+        x = poly.v_stack((1,) + tuple(init_variable_expr(v, state.get_shape(v))
+                                      for v in polynomial_symbols))
 
         # Create new constraints
         new_constraints: list[Constraint] = list(constraints)
         for i, constr in enumerate(to_process):
-            new_constr = constr.expression
+            # 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.linear_in(x).degree().max()
+
+            # Add first term in the quadratic module
+            d = constr_deg
+            if d.read(state).scalar().constant() % 2 != 0:
+                d += 1
+
+            monomials = x.combinations(poly.arange(d + 1)).half_newton_polytope(x)
+            coeffs = opt_variable_from_name(rf"\gamma_{i},c", shape=monomials.shape)
+
+            new_constr = constr.expression - coeffs.T @ monomials
+
+            # Add other terms of the quadratic module
             for j, domain_poly in enumerate(constr.domain.polynomials):
                 # Make a multiplier polynomial for a domain polynomial of a constraint.
                 # 
@@ -123,11 +167,7 @@ class PutinarPSatz(Canonicalization):
                 # :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()
+                domain_deg = domain_poly.linear_in(x).degree().max()
 
                 # Degree of multiplier needs to be an integer
                 d = constr_deg - domain_deg
@@ -140,19 +180,16 @@ class PutinarPSatz(Canonicalization):
                 if d.read(state).scalar().constant() % 2 != 0:
                     d += 1
 
-                x = poly.v_stack((1,) + tuple(init_variable_expr(v, state.get_shape(v))
-                                              for v in prob.polynomial_variables))
-
                 # FIXME: need to check that there is not a \gamma_i,j already!
                 # TODO: deterministically generate unique names
 
                 # TODO: better documentation
                 # Generate monomials of up to degree d and take only monomials
                 # that are in half of the newton polytope
-                monomials = x.combinations(poly.arange(d + 1)) # .half_newton_polytope(x)
-                gamma = opt_variable_from_name(rf"\gamma_{i},{j}", shape=monomials.shape)
+                monomials = x.combinations(poly.arange(d + 1)).half_newton_polytope(x)
+                coeffs = opt_variable_from_name(rf"\gamma_{i},{j}", shape=monomials.shape)
 
-                multiplier = gamma.T @ monomials
+                multiplier = coeffs.T @ monomials
 
                 # Subtract multiplier from expression and impose that it is also SOS
                 new_constr -= multiplier * domain_poly
@@ -211,11 +248,18 @@ class LogDet(Canonicalization):
     """
     problem: SOSProblem
 
+    @property
+    @override
+    def cost(self) -> Expression:
+        return self(self.problem).cost
+
+    @property
     @override
-    def __call__(self, problem : SOSProblem | None = None) -> SOSProblem:
-        if problem is None:
-            problem = self.problem
+    def constraints(self) -> Expression:
+        return self(self.problem).constraints
 
+    @override
+    def __call__(self, problem : SOSProblem) -> SOSProblem:
         A = problem.cost
 
         # FIXME: should check for name clashes in state object
@@ -236,7 +280,6 @@ class LogDet(Canonicalization):
 
         return replace(problem, cost=new_cost, constraints=new_constraints)
 
-
     @override
     def apply(self, state: ExpressionState) -> tuple[ExpressionState, InternalSOSProblem]:
-        return self().apply(state)
+        return self(self.problem).apply(state)
diff --git a/sumofsquares/problems.py b/sumofsquares/problems.py
index a646d0e..9990264 100644
--- a/sumofsquares/problems.py
+++ b/sumofsquares/problems.py
@@ -179,31 +179,31 @@ class SOSProblem(Problem):
         variable_indices: set[VariableIndex] = set()
         constraints: list[Constraint] = []
 
+        # Collect all variables that are involved in the problem
         state, cost = self.cost.apply(state)
-        variable_indices.update(cost.variables())
-
-        # Compute the polymatrix of each constraint expression. Even though the
-        # result may not be used, this is necessary to account for the case of
-        # variables that are only present in the constraints, which the state
-        # object may not contain yet.
+        variable_indices |= cost.variables()
 
         for c in self.constraints:
             state, pm  = c.expression.apply(state)
-            variable_indices.update(pm.variables())
+            variable_indices |= pm.variables()
 
-        variables = set(state.get_symbol_from_variable_index(v)
-                        for v in variable_indices)
+        # Get symbols of the variables
+        all_symbols = set(state.get_symbol_from_variable_index(v)
+                      for v in variable_indices)
 
-        # Collect optimization variables
+        # Separate between symbols of optimization variables and those of polynomial variables
         def is_opt(v):
             return isinstance(v, OptSymbol)
 
-        polynomial_variables, variables = partition(is_opt, variables)
-        polynomial_variables = tuple(polynomial_variables) # because it is a generator
+        polynomial_symbols, symbols = partition(is_opt, all_symbols)
+
+        # evaluate generators
+        symbols = tuple(symbols)
+        polynomial_symbols = tuple(polynomial_symbols)
 
         x = poly.v_stack((1,) + tuple(
             init_variable_expr(v, state.get_shape(v))
-            for v in polynomial_variables))
+            for v in polynomial_symbols))
 
         for i, c in enumerate(self.constraints):
             if isinstance(c, EqualToZero):
@@ -287,7 +287,7 @@ class SOSProblem(Problem):
         #         raise NotImplementedError(f"Cannot process constraint of type {ctype} (yet).")
 
         return state, InternalSOSProblem(cost, tuple(pm_constraints),
-                                         tuple(variables), polynomial_variables,
+                                         symbols, polynomial_symbols,
                                          self.solver)
 
 
@@ -305,8 +305,8 @@ class InternalSOSProblem(Problem):
     """
     cost: PolyMatrixMixin
     constraints: Sequence[Constraint[PolyMatrixMixin]]
-    variables: Sequence[OptSymbol]
-    polynomial_variables: Sequence[Symbol]
+    symbols: Sequence[OptSymbol]
+    polynomial_symbols: Sequence[Symbol]
     solver: Solver
 
     def to_conic_problem(self, state: ExpressionState, verbose: bool = False) -> ConicProblem:
@@ -335,7 +335,7 @@ class InternalSOSProblem(Problem):
 
         # indices of variables in the optimization problem, sorted
         variable_indices = sum(sorted(tuple(state.get_indices_as_variable_index(v))
-                                      for v in self.variables), ())
+                                      for v in self.symbols), ())
 
         # cost linear term
         cost_coeffs = dict(cost.affine_coefficients_by_degrees(variable_indices))
@@ -374,6 +374,11 @@ class InternalSOSProblem(Problem):
             # Get constant and linear terms
             constant = constr.affine_coefficient(MonomialIndex.constant())
             linear = tuple(constr.affine_coefficient(v) for v in variable_indices)
+
+            # constraint has disappeared
+            if not (np.any(constant) or np.any(linear)):
+                # print("Constraint disappeared!")
+                continue
             
             if isinstance(c, EqualToZero):
                 dims["z"].append(nrows)
@@ -420,7 +425,7 @@ class InternalSOSProblem(Problem):
                             dims=dims, is_qp=is_qp,
                             solver=self.solver,
                             variables={v : state.get_shape(v)
-                                for v in self.variables
+                                for v in self.symbols
                             })