bugfixes in KKT conditions

1 files changed, 148 insertions, 142 deletions

diff --git a/test_polymatrix/test_polymatrix.py b/test_polymatrix/test_polymatrix.py
index f693df0..d8abeef 100644
--- a/test_polymatrix/test_polymatrix.py
+++ b/test_polymatrix/test_polymatrix.py
@@ -1,22 +1,55 @@
 import unittest
-
-from polymatrix.polystruct import init_equation, init_poly_matrix
-from polymatrix.utils import variable_to_index
+from polymatrix.init.initoptimization import init_optimization
+from polymatrix.init.initpolymatrix import init_poly_matrix
+from polymatrix.optimization import Optimization
+from polymatrix.polymatrix import PolyMatrix
 
 class TestPolyMatrix(unittest.TestCase):
+    # @staticmethod
+    # def assert_term_in_eq(result, degree, eq_idx, row_idx, value, monoms=None):
+    #     if monoms is None:
+    #         monoms = tuple()
+
+    #     assert degree in result.data, f'could not find {degree} in {result.data}'
+    #     degree_terms = result.data[degree]
+
+    #     key = eq_idx, monoms, row_idx
+    #     assert key in degree_terms, f'could not find {key} in {degree_terms}'
+    #     # eq_terms = degree_terms[key]
+    #     # assert row_idx in eq_terms, f'could not find {row_idx} in {eq_terms}'
+    #     assert degree_terms[key] == value, f'value {degree_terms[key]} and {value} do not match'
+
+    # @staticmethod
+    # def assert_term_in_eq(result, degree, row_idx, monoms, value):
+
+    #     assert degree in result.data, f'could not find {degree} in {result.data}'
+    #     degree_terms = result.data[degree]
+
+    #     key = row_idx, variable_to_index(result.n_param, monoms)
+    #     assert key in degree_terms, f'could not find {key} in {degree_terms}'
+
+    #     assert degree_terms[key] == value, f'value {degree_terms[key]} and {value} do not match'
+
     @staticmethod
-    def assert_term_in_eq(result, degree, eq_idx, row_idx, value, monoms=None):
-        if monoms is None:
-            monoms = tuple()
+    def assert_term_in_eq(
+        problem: Optimization,
+        poly_row: int, 
+        p_monomial: tuple[tuple[PolyMatrix, int, int], ...], 
+        value: float,
+        x_monomial: tuple[int, ...] = None,
+    ):
+        if x_monomial is None:
+            x_monomial = tuple()
+
+        offset_dict = problem.state.offset_dict
+        p_monomial = tuple(offset_dict[(polymat, degree)][0] + offset for polymat, degree, offset in p_monomial)
 
-        assert degree in result.terms, f'could not find {degree} in {result.terms}'
-        degree_terms = result.terms[degree]
+        equality_constraint = problem.equality_constraints[0]
 
-        key = eq_idx, monoms, row_idx
-        assert key in degree_terms, f'could not find {key} in {degree_terms}'
-        # eq_terms = degree_terms[key]
-        # assert row_idx in eq_terms, f'could not find {row_idx} in {eq_terms}'
-        assert degree_terms[key] == value, f'value {degree_terms[key]} and {value} do not match'
+        key = (poly_row, x_monomial, p_monomial)
+        assert key in equality_constraint, f'could not find {key} in {equality_constraint}'
+
+        assert equality_constraint[key] == value, f'value {equality_constraint[key]} and {value} do not match'
 
     def test_param_matrix_param_d0_vector_degree_d0(self):
         """
@@ -25,50 +58,44 @@ class TestPolyMatrix(unittest.TestCase):
 
         n_var = 2
 
-        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
-        vec = init_poly_matrix(degrees=(0,), shape=(n_var, 1))
+        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(name='vec', degrees=(0,), shape=(n_var, 1))
 
-        eq = init_equation(
-            terms = [(mat, vec)],
-            n_var = n_var,
+        problem = init_optimization(
+            n_var=n_var,
+        ).add_equality_constraints(
+            expr=[(mat, vec)],
         )
 
-        result = eq.create()
-        offset_dict = eq.offset_dict
-
         # a11 v11
         self.assert_term_in_eq(
-            result = result,
-            degree = 2,
-            eq_idx = 0,
-            row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 0)]),
+            problem = problem,
+            poly_row = 0,
+            p_monomial = ((mat, 0, 0), (vec, 0, 0)),
             value = 1,
         )
 
         # a12 v21
         self.assert_term_in_eq(
-            result = result,
-            degree = 2,
-            eq_idx = 0,
-            row_idx = (offset_dict[(mat, 0)]+2, offset_dict[(vec, 0)]+1),
+            problem = problem,
+            poly_row = 0,
+            p_monomial = ((mat, 0, 2), (vec, 0, 1)),
             value = 1,
         )
 
         # a21 v11
         self.assert_term_in_eq(
-            result = result,
-            degree = 2,
-            eq_idx = 1,
-            row_idx = (offset_dict[(mat, 0)]+1, offset_dict[(vec, 0)]),
+            problem = problem,
+            poly_row = 1,
+            p_monomial = ((mat, 0, 1), (vec, 0, 0)),
             value = 1,
         )
 
         # a22 v21
         self.assert_term_in_eq(
-            result = result,
-            degree = 2,
-            eq_idx = 1,
-            row_idx = (offset_dict[(mat, 0)]+3, offset_dict[(vec, 0)]+1),
+            problem = problem,
+            poly_row = 1,
+            p_monomial = ((mat, 0, 3), (vec, 0, 1)),
             value = 1,
         )
 
@@ -79,33 +106,29 @@ class TestPolyMatrix(unittest.TestCase):
 
         n_var = 2
 
-        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
-        vec = init_poly_matrix(degrees=(0,1), shape=(n_var, 1))
+        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(name='vec', degrees=(0, 1), shape=(n_var, 1))
 
-        eq = init_equation(
-            terms = [(mat, vec)],
-            n_var = n_var,
+        problem = init_optimization(
+            n_var=n_var,
+        ).add_equality_constraints(
+            expr=[(mat, vec)],
         )
 
-        result = eq.create()
-        offset_dict = eq.offset_dict
-
         # a11 v011
         self.assert_term_in_eq(
-            result = result,
-            degree = 2,
-            eq_idx = 0,
-            row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 0)]),
+            problem = problem,
+            poly_row = 0,
+            p_monomial = ((mat, 0, 0), (vec, 0, 0)),
             value = 1,
         )
 
         # a11 v011
         self.assert_term_in_eq(
-            result = result,
-            degree = 2,
-            eq_idx = 0,
-            monoms=(0,),
-            row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 1)]),
+            problem = problem,
+            poly_row = 0,
+            x_monomial=(0,),
+            p_monomial = ((mat, 0, 0), (vec, 1, 0)),
             value = 1,
         )
 
@@ -116,32 +139,29 @@ class TestPolyMatrix(unittest.TestCase):
 
         n_var = 2
 
-        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
-        vec = init_poly_matrix(subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
+        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(name='vec', subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
 
-        eq = init_equation(
-            terms = [(mat, vec)],
-            n_var = n_var,
+        problem = init_optimization(
+            n_var=n_var,
+        ).add_equality_constraints(
+            expr=[(mat, vec)],
         )
 
-        result = eq.create()
-        offset_dict = eq.offset_dict
 
         # a11
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            eq_idx = 0,
-            row_idx = (offset_dict[(mat, 0)],),
+            problem = problem,
+            poly_row = 0,
+            p_monomial = ((mat, 0, 0),),
             value = 1,
         )
 
         # a21
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            eq_idx = 1,
-            row_idx = (offset_dict[(mat, 0)]+1,),
+            problem = problem,
+            poly_row = 1,
+            p_monomial = ((mat, 0, 1),),
             value = 1,
         )
 
@@ -151,45 +171,40 @@ class TestPolyMatrix(unittest.TestCase):
         """
 
         n_var = 2
-
-        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
-        vec = init_poly_matrix(subs={1: {(0, 0, 0): 1, (0, 0, 1): 0, (1, 0, 0): 0, (1, 0, 1): 1}}, shape=(n_var, 1))
-
-        eq = init_equation(
-            terms = [(mat, vec)],
-            n_var = n_var,
+        
+        mat = init_poly_matrix(name='mat', degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(name='vec', subs={1: {(0, 0, 0): 1, (0, 0, 1): 0, (1, 0, 0): 0, (1, 0, 1): 1}}, shape=(n_var, 1))
+
+        problem = init_optimization(
+            n_var=n_var,
+        ).add_equality_constraints(
+            expr=[(mat, vec)],
         )
 
-        result = eq.create()
-        offset_dict = eq.offset_dict
-
         # a11
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            eq_idx = 0,
-            monoms=(0,),
-            row_idx = (offset_dict[(mat, 0)],),
+            problem = problem,
+            poly_row = 0,
+            x_monomial=(0,),
+            p_monomial = ((mat, 0, 0),),
             value = 1,
         )
 
         # a12
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            eq_idx = 0,
-            monoms=(1,),
-            row_idx = (offset_dict[(mat, 0)]+2,),
+            problem = problem,
+            poly_row = 0,
+            x_monomial=(1,),
+            p_monomial = ((mat, 0, 2),),
             value = 1,
         )
 
         # a21
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            eq_idx = 1,
-            monoms=(0,),
-            row_idx = (offset_dict[(mat, 0)]+1,),
+            problem = problem,
+            poly_row = 1,
+            x_monomial=(0,),
+            p_monomial = ((mat, 0, 1),),
             value = 1,
         )
 
@@ -200,22 +215,19 @@ class TestPolyMatrix(unittest.TestCase):
 
         n_var = 2
 
-        mat = init_poly_matrix(subs={0: {(0, 0, 0): 1, (0, 1, 0): 1, (1, 0, 0): 1, (1, 1, 0): 1}}, shape=(n_var, n_var))
-        vec = init_poly_matrix(subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
+        mat = init_poly_matrix(name='mat', subs={0: {(0, 0, 0): 1, (0, 1, 0): 1, (1, 0, 0): 1, (1, 1, 0): 1}}, shape=(n_var, n_var))
+        vec = init_poly_matrix(name='vec', subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
 
-        eq = init_equation(
-            terms = [(mat, vec)],
-            n_var = n_var,
+        problem = init_optimization(
+            n_var=n_var,
+        ).add_equality_constraints(
+            expr=[(mat, vec)],
         )
 
-        result = eq.create()
-        offset_dict = eq.offset_dict
-
         self.assert_term_in_eq(
-            result = result,
-            degree = 0,
-            eq_idx = 0,
-            row_idx = tuple(),
+            problem = problem,
+            poly_row = 0,
+            p_monomial = tuple(),
             value = 2,
         )
     
@@ -233,33 +245,30 @@ class TestPolyMatrix(unittest.TestCase):
         n_var = 2
 
         mat = init_poly_matrix(
+            name='mat', 
             degrees=(0,), 
             re_index=skew_symmetric,
             shape=(n_var, n_var),
         )
-        vec = init_poly_matrix(subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
+        vec = init_poly_matrix(name='vec', subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
 
-        eq = init_equation(
-            terms = [(mat, vec)],
-            n_var = n_var,
+        problem = init_optimization(
+            n_var=n_var,
+        ).add_equality_constraints(
+            expr=[(mat, vec)],
         )
 
-        result = eq.create()
-        offset_dict = eq.offset_dict
-
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            eq_idx = 0,
-            row_idx = (offset_dict[(mat, 0)] + 2,),
+            problem = problem,
+            poly_row = 0,
+            p_monomial = ((mat, 0, 2),),
             value = 1,
         )
 
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            eq_idx = 1,
-            row_idx = (offset_dict[(mat, 0)] + 2,),
+            problem = problem,
+            poly_row = 1,
+            p_monomial = ((mat, 0, 2),),
             value = -1,
         )
 
@@ -285,46 +294,43 @@ class TestPolyMatrix(unittest.TestCase):
         n_var = 2
 
         mat = init_poly_matrix(
+            name='mat', 
             subs={0: {(0, 0, 0): 1, (0, 1, 0): 1, (1, 0, 0): 1, (1, 1, 0): 1}}, 
             shape=(n_var, n_var),
         )
         vec = init_poly_matrix(
+            name='vec', 
             degrees=(1,),
             re_index=gradient,
             shape=(n_var, 1),
         )
 
-        eq = init_equation(
-            terms = [(mat, vec)],
-            n_var = n_var,
+        problem = init_optimization(
+            n_var=n_var,
+        ).add_equality_constraints(
+            expr=[(mat, vec)],
         )
 
-        result = eq.create()
-        offset_dict = eq.offset_dict
-
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            monoms=(0,),
-            eq_idx = 0,
-            row_idx = (offset_dict[(vec, 1)],),
+            problem = problem,
+            x_monomial=(0,),
+            poly_row = 0,
+            p_monomial = ((vec, 1, 0),),
             value = 2,
         )
 
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            monoms=(1,),
-            eq_idx = 0,
-            row_idx = (offset_dict[(vec, 1)]+2,),
+            problem = problem,
+            x_monomial=(1,),
+            poly_row = 0,
+            p_monomial = ((vec, 1, 2),),
             value = 1,
         )
 
         self.assert_term_in_eq(
-            result = result,
-            degree = 1,
-            monoms=(0,),
-            eq_idx = 1,
-            row_idx = (offset_dict[(vec, 1)]+2,),
+            problem = problem,
+            x_monomial=(0,),
+            poly_row = 1,
+            p_monomial = ((vec, 1, 2),),
             value = 1,
         )