bugfixes in KKT conditions

7 files changed, 588 insertions, 300 deletions

diff --git a/polymatrix/init/initoptimization.py b/polymatrix/init/initoptimization.py
index c08c71f..653d248 100644
--- a/polymatrix/init/initoptimization.py
+++ b/polymatrix/init/initoptimization.py
@@ -1,13 +1,25 @@
 from polymatrix.impl.optimizationimpl import OptimizationImpl
+from polymatrix.init.initoptimizationstate import init_optimization_state
 from polymatrix.optimizationstate import OptimizationState
 
 
 def init_optimization(
-    state: OptimizationState,
-    equality_constraints: dict[int, dict[tuple[int, int], float]],
-    inequality_constraints: dict[int, dict[tuple[int, int], float]],
-    auxillary_equality: dict[int, dict[tuple[int, int], float]]
+    n_var: int,
+    equality_constraints: dict[int, dict[tuple[int, int], float]] = None,
+    inequality_constraints: dict[int, dict[tuple[int, int], float]] = None,
+    auxillary_equality: dict[int, dict[tuple[int, int], float]] = None,
 ):
+    state = init_optimization_state(n_var=n_var)
+
+    if equality_constraints is None:
+        equality_constraints = {}
+
+    if inequality_constraints is None:
+        inequality_constraints = {}
+
+    if auxillary_equality is None:
+        auxillary_equality = {}
+
     return OptimizationImpl(
         state=state,
         equality_constraints=equality_constraints,
diff --git a/polymatrix/init/initoptimizationstate.py b/polymatrix/init/initoptimizationstate.py
index e6e323d..372e225 100644
--- a/polymatrix/init/initoptimizationstate.py
+++ b/polymatrix/init/initoptimizationstate.py
@@ -4,9 +4,15 @@ from polymatrix.mixins.polymatrixmixin import PolyMatrixMixin
 
 def init_optimization_state(
     n_var: int,
-    n_param: int,
-    offset_dict: dict[tuple[PolyMatrixMixin, int], int],
+    n_param: int = None,
+    offset_dict: dict[tuple[PolyMatrixMixin, int], int] = None,
 ):
+    if n_param is None:
+        n_param = 0
+
+    if offset_dict is None:
+        offset_dict = {}
+
     return OptimizationStateImpl(
         n_var=n_var,
         n_param=n_param,
diff --git a/polymatrix/mixins/optimizationmixin.py b/polymatrix/mixins/optimizationmixin.py
index d9ea263..e44a36f 100644
--- a/polymatrix/mixins/optimizationmixin.py
+++ b/polymatrix/mixins/optimizationmixin.py
@@ -5,6 +5,7 @@ import functools
 import itertools
 
 import more_itertools
+from sympy import Equivalent
 from polymatrix.mixins.optimizationpipeopmixin import OptimizationPipeOpMixin
 from polymatrix.mixins.optimizationstatemixin import OptimizationStateMixin
 
@@ -21,7 +22,7 @@ class OptimizationMixin(abc.ABC):
 
     @property
     @abc.abstractmethod
-    def equality_constraints(self) -> dict[int, dict[tuple[int, int], float]]:
+    def equality_constraints(self) -> dict[int, dict[tuple[int, tuple[int, ...], tuple[int, ...]], float]]:
         ...
 
     @property
@@ -38,39 +39,43 @@ class OptimizationMixin(abc.ABC):
     def n_var(self) -> int:
         return self.state.n_var
 
-    @property
-    def n_equality_constraints(self) -> int:
-        if len(self.equality_constraints) == 0:
-            return 0
+    # @property
+    # def n_equality_constraints(self) -> int:
+    #     if len(self.equality_constraints) == 0:
+    #         return 0
 
-        return max(eq for eq, _ in self.equality_constraints.items()) + 1
+    #     return max(eq for eq, _ in self.equality_constraints.items()) + 1
 
-    @property
-    def n_inequality_constraints(self) -> int:
-        if len(self.inequality_constraints) == 0:
-            return 0
+    # @property
+    # def n_inequality_constraints(self) -> int:
+    #     if len(self.inequality_constraints) == 0:
+    #         return 0
 
-        return max(eq for eq, _ in self.inequality_constraints.items()) + 1
-        # return max(key[0] for _, d_data in self.inequality_constraints.items() for key, _ in d_data.items()) + 1
+    #     return max(eq for eq, _ in self.inequality_constraints.items()) + 1
+    #     # return max(key[0] for _, d_data in self.inequality_constraints.items() for key, _ in d_data.items()) + 1
 
-    @property
-    def n_auxillary_equality(self) -> int:
-        if len(self.auxillary_equality) == 0:
-            return 0
+    # @property
+    # def n_auxillary_equality(self) -> int:
+    #     if len(self.auxillary_equality) == 0:
+    #         return 0
 
-        return max(eq for eq, _ in self.auxillary_equality.items()) + 1
-        # return max(key[0] for _, d_data in self.auxillary_equality.items() for key, _ in d_data.items()) + 1
+    #     return max(eq for eq, _ in self.auxillary_equality.items()) + 1
+    #     # return max(key[0] for _, d_data in self.auxillary_equality.items() for key, _ in d_data.items()) + 1
 
     # def pipe(self, *operators: OptimizationPipeOpMixin):
     #     return functools.reduce(lambda obs, op: op(obs), operators, self)
 
     def add_equality_constraints(self, expr: tuple[tuple[PolyMatrixMixin, ...], ...]):
+        for term in expr:
+            for left, right in itertools.pairwise(term):
+                assert left.shape[1] == right.shape[0], f'{left} and {right} do not match'
+
         all_polymats = tuple(polymat for term in expr for polymat in term)
 
         # update offsets with unseen polymats
         state = self.state.update_offsets(all_polymats)
 
-        eq_constr_buffer = collections.defaultdict(lambda: collections.defaultdict(float))
+        equality_constraint = collections.defaultdict(float)
 
         for term in expr:
 
@@ -85,36 +90,36 @@ class OptimizationMixin(abc.ABC):
                 # n_var = 2, degree = 3
                 # cominations = [(0, 0, 0), (0, 0, 1), (0, 1, 1), (1, 1, 1)]
                 for combination in itertools.combinations_with_replacement(range(self.n_var), total_degree):
-                    for x_monom in more_itertools.distinct_permutations(combination):
+                    for x_monomial in more_itertools.distinct_permutations(combination):
 
                         def acc_func(acc, v):
                             last, _ = acc
                             new = last + v
 
-                            return new, x_monom[last:new]
+                            return new, x_monomial[last:new]
 
                         # monom=(1, 0, 1) -> monom1=(x2, x1), monom2=(x2)
                         # the monomial either refers to a free parameter, or to a value in the substitution table
-                        p_monoms = list(monom for _, monom in itertools.accumulate(degrees, acc_func, initial=(0, None)))[1:]
+                        rel_x_monomials = list(monom for _, monom in itertools.accumulate(degrees, acc_func, initial=(0, None)))[1:]
                         
                         def non_increasing(seq):
                             return all(y <= x for x, y in zip(seq, seq[1:]))
 
                         # (1,0) -> x2*x1 instead of (0,1)->x1*x2
-                        if all(non_increasing(monom) for monom in p_monoms):
+                        if all(non_increasing(rel_x_monomial) for rel_x_monomial in rel_x_monomials):
 
-                            col_defaults = [monomial_to_index(self.n_var, monom) for monom in p_monoms]
+                            col_defaults = [monomial_to_index(self.n_var, monom) for monom in rel_x_monomials]
 
                             n_rows_col = itertools.chain((m.shape[0] for m in term), (term[-1].shape[1],))
 
                             for curr_rows_col in itertools.product(*[range(e) for e in n_rows_col]):
                                 
-                                def gen_re_indexing():
-                                    for m, degree, (poly_row, poly_col), monom, col_default, subs, offset in zip(
-                                            term, degrees, itertools.pairwise(curr_rows_col), p_monoms, col_defaults, d_subs, d_offsets):
+                                def gen_re_indexing(): #term=term, degrees=degrees, curr_rows_col=curr_rows_col, p_monoms=p_monoms, col_defaults=col_defaults, d_subs=d_subs, d_offsets=d_offsets):
+                                    for polymat, degree, (poly_row, poly_col), rel_x_monomial, col_default, subs, offset in zip(
+                                            term, degrees, itertools.pairwise(curr_rows_col), rel_x_monomials, col_defaults, d_subs, d_offsets):
 
-                                        if m.re_index is not None:
-                                            re_index = m.re_index(degree, poly_row, poly_col, monom)
+                                        if polymat.re_index is not None:
+                                            re_index = polymat.re_index(degree, poly_row, poly_col, rel_x_monomial)
                                         else:
                                             re_index = None
 
@@ -138,11 +143,10 @@ class OptimizationMixin(abc.ABC):
 
                                         if subs_val is None:
                                             # the coefficient is selected after the reindexing
-
-                                            row = new_poly_row + new_poly_col * m.shape[0]
+                                            row = new_poly_row + new_poly_col * polymat.shape[0]
 
                                             # linearize parameter matrix
-                                            p_index = int(offset[0] + row + col * m.shape[0] * m.shape[1])
+                                            p_index = int(offset[0] + row + col * polymat.shape[0] * polymat.shape[1])
 
                                             assert p_index < offset[1], f'{p_index} is bigger than {offset[1]}'
 
@@ -164,25 +168,30 @@ class OptimizationMixin(abc.ABC):
                                     continue
 
                                 poly_row = data[0][0]
-                                p_monom = tuple(d[2] for d in data if d[4] is None)
-                                degree = len(p_monom)
-
-                                eq_constr_buffer[degree][poly_row, x_monom, p_monom] += value
+                                p_monomial = tuple(sorted(d[2] for d in data if d[4] is None))
+                                x_monomial_sorted = tuple(sorted(x_monomial))
 
-        # assign equations
-        rows_perm_set = set((eq_idx, perm) for eq_tuple_degree in eq_constr_buffer.values() for (eq_idx, perm, _) in eq_tuple_degree.keys())
-        eq_to_rows = {eq: idx for idx, eq in enumerate(rows_perm_set)}
+                                equality_constraint[poly_row, x_monomial_sorted, p_monomial] += value
 
-        eq_constr = collections.defaultdict(list)
+        # eq_constr = collections.defaultdict(list)
 
-        for degree, d_data in eq_constr_buffer.items():
-            for (eq_idx, perm, p_monoms), val in d_data.items():
-                row = eq_to_rows[(eq_idx, perm)]
-                eq_constr[row] += ((p_monoms, val),)
+        # for degree, d_data in eq_constr_buffer.items():
+        #     for (eq_idx, perm, p_monoms), val in d_data.items():
+        #         row = eq_to_rows[(eq_idx, perm)]
+        #         eq_constr[row] += ((p_monoms, val),)
 
         # eq_data = dict(gen_eq_data())
 
-        return dataclasses.replace(self, state=state, equality_constraints=eq_constr)
+        if len(self.equality_constraints) == 0:
+            constraint_index = 0
+        else:
+            constraint_index = max(self.equality_constraints.keys()) + 1
+
+        return dataclasses.replace(
+            self, 
+            state=state, 
+            equality_constraints=self.equality_constraints | {constraint_index: equality_constraint},
+        )
 
     def add_inequality_constraints(self, expr):
         # all_polymats = tuple(polymat for term in expr for polymat in term)
@@ -195,13 +204,11 @@ class OptimizationMixin(abc.ABC):
         aux_eq_buffer = collections.defaultdict(list)
         ineq_constr_buffer = collections.defaultdict(list)
 
-        ineq_idx = self.n_inequality_constraints
-        aux_eq_idx = self.n_auxillary_equality
+        ineq_idx = len(self.inequality_constraints)
+        aux_eq_idx = len(self.auxillary_equality)
 
-        for degree in polymat.degrees:
-            # print(f'{degree=}')
-            
-            assert degree != 0
+        for degree in polymat.degrees:            
+            # assert degree != 0
             assert degree % 2 == 0
 
             vec = tuple(itertools.combinations_with_replacement(range(state.n_var), int(degree/2)))
@@ -209,20 +216,22 @@ class OptimizationMixin(abc.ABC):
 
             # introduce auxillary parameters
             offset_x = state.n_param
-            # print(f'{offset_x=}')
 
             n_param_added = int(n_square_mat*(n_square_mat-1)/2)
             state = state.update_param(n_param_added)
-            # print(f'{n_param_added=}')
 
             offset_a, end_idx = state.offset_dict[(polymat, degree)]
 
             def get_param_index_of_poly_mat(row, col):
-                x_monom = vec[row] + vec[col]
-                reindex = polymat.re_index(degree, 1, 1, x_monom)
+                x_monomial = tuple(sorted(vec[row] + vec[col], reverse=True))
+
+                if polymat.re_index is None:
+                    reindex = None
+                else:
+                    reindex = polymat.re_index(degree, 1, 1, x_monomial)
 
                 if reindex is None:
-                    n_monom, factor = x_monom, 1.0
+                    n_monom, factor = x_monomial, 1.0
 
                 else:
                     _, _, n_monom, factor = polymat.re_index(degree, 1, 1, vec[row] + vec[col])
@@ -241,42 +250,46 @@ class OptimizationMixin(abc.ABC):
 
             for k in range(n_square_mat):
                 # f in f-v^T@x-r^2
-                monom, factor = get_param_index_of_poly_mat(k, k)
-                # ineq_constr_buffer[len(monom)][(ineq_idx, monom)] += factor
-                ineq_constr_buffer[ineq_idx] += ((monom, factor),)
-
-                # print(ineq_constr_buffer)
+                p_monomial, factor = get_param_index_of_poly_mat(k, k)
+                ineq_constr_buffer[ineq_idx] += ((p_monomial, factor),)
 
                 for row in range(k):
 
                     # v^T@x in f-v^T@x-r^2
-                    monom, factor = get_param_index_of_poly_mat(k, row)
-                    # ineq_constr_buffer[len(monom) + 1][(ineq_idx, monom + (offset_x + row,))] -= factor
-                    ineq_constr_buffer[ineq_idx] += ((monom + (offset_x + row,), -factor),)
-
-                    # print(f'{k}: {offset_x + row=}')
+                    p_monomial, factor = get_param_index_of_poly_mat(k, row)
+                    ineq_constr_buffer[ineq_idx] += ((p_monomial + (offset_x + row,), -factor),)
 
                     for col in range(k):
                         # P@x in P@x-v
-                        monom, factor = get_param_index_of_poly_mat(row, col)
-                        # aux_eq_buffer[len(monom) + 1][(aux_eq_idx+row, monom + (offset_x + col,))] += factor
-                        aux_eq_buffer[aux_eq_idx] += ((monom + (offset_x + col,), factor),)
+                        p_monomial, factor = get_param_index_of_poly_mat(row, col)
+                        aux_eq_buffer[aux_eq_idx] += ((p_monomial + (offset_x + col,), factor),)
 
                     # -v in P@x-v
-                    monom, factor = get_param_index_of_poly_mat(row, k)
-                    # aux_eq_buffer[len(monom)][(aux_eq_idx+row, monom)] -= factor
-                    aux_eq_buffer[aux_eq_idx] += ((monom, -factor),)
+                    p_monomial, factor = get_param_index_of_poly_mat(row, k)
+                    aux_eq_buffer[aux_eq_idx] += ((p_monomial, -factor),)
 
                     aux_eq_idx += 1
 
                 ineq_idx += 1
                 offset_x += k
 
-        # print(f'{ineq_idx - self.n_inequality_constraints + aux_eq_idx - self.n_auxillary_equality=}')
-        # print(f'{aux_eq_idx - self.n_auxillary_equality=}')
-
         return dataclasses.replace(self, inequality_constraints=ineq_constr_buffer, auxillary_equality=aux_eq_buffer, state=state)
 
+    def get_equality_constraints(self):
+        offset = 0
+        equality_constraints = collections.defaultdict(list)
+
+        for _, data in self.equality_constraints.items():
+            equation_set = set((poly_row, x_monomial) for (poly_row, x_monomial, _) in data.keys())
+            to_eq_index_dict = {eq: offset + idx for idx, eq in enumerate(equation_set)}
+            offset += len(equation_set)
+
+            for (poly_row, x_monomial, p_monomial), val in data.items():
+                eq_index = to_eq_index_dict[(poly_row, x_monomial)]
+                equality_constraints[eq_index] += ((p_monomial, val),)
+
+        return equality_constraints
+
     def minimize(self, cost_func=None):
         """
         - assume sum of squares cost function on variables x
@@ -293,17 +306,24 @@ class OptimizationMixin(abc.ABC):
 
         state = self.state
 
-        n_equality_constraints = self.n_equality_constraints
-        n_inequality_constraints = self.n_inequality_constraints
+        equality_constraints = self.get_equality_constraints()
+
+        n_equality_constraints = len(equality_constraints)
+
+        # if len(self.inequality_constraints) == 0:
+        #     n_inequality_constraints = 0
+        # else:
+        #     n_inequality_constraints = max(eq for eq, _ in self.inequality_constraints.items()) + 1
+        n_inequality_constraints = len(self.inequality_constraints)
 
         # remove unused parameters
         # ------------------------
 
-        monom_update_reverse = set(m for monom_vals in itertools.chain(self.equality_constraints.values(), self.inequality_constraints.values(), self.auxillary_equality.values()) for monom, _ in monom_vals for m in monom)
-        # print(tuple(p for p in param_indices if p not in monom_update_reverse))
-        monom_update = {monom: idx for idx, monom in enumerate(sorted(monom_update_reverse))}
+        # find used parameters
+        used_param_indices = set(param_idx for monomial_value_pairs in itertools.chain(equality_constraints.values(), self.inequality_constraints.values(), self.auxillary_equality.values()) for monomial, _ in monomial_value_pairs for param_idx in monomial)
+        param_update = {monom: idx for idx, monom in enumerate(sorted(used_param_indices))}
 
-        assert max(monom_update_reverse) == state.n_param - 1, f'{max(monom_update_reverse)=} is not {state.n_param - 1=}'
+        assert max(used_param_indices) == state.n_param - 1, f'{max(used_param_indices)=} is not {state.n_param - 1=}'
 
         # param_indices = tuple(start + idx for start, end in state.offset_dict.values() for idx in range(end - start))
         # print(tuple(m for m in monom_update_reverse if m not in param_indices))
@@ -311,20 +331,18 @@ class OptimizationMixin(abc.ABC):
 
         # the variables x are assumed to be the parameters of all registered polynomial matrices
         # todo: this would later come from a specific polynomial matrices
-        param_indices = tuple(monom_update[start + idx] for start, end in state.offset_dict.values() for idx in range(end - start) if start + idx in monom_update_reverse)
+        param_indices = tuple(param_update[start + idx] for start, end in state.offset_dict.values() for idx in range(end - start) if start + idx in used_param_indices)
 
-        equality_constraints = tuple((key, tuple((tuple(monom_update[m] for m in monom), val) for monom, val in monom_val)) for key, monom_val in self.equality_constraints.items())
-        inequality_constraints = tuple((key, tuple((tuple(monom_update[m] for m in monom), val) for monom, val in monom_val)) for key, monom_val in self.inequality_constraints.items())
-        auxillary_equality = tuple((key, tuple((tuple(monom_update[m] for m in monom), val) for monom, val in monom_val)) for key, monom_val in self.auxillary_equality.items())
+        equality_constraints = tuple((eq_index, tuple((tuple(param_update[m] for m in monom), val) for monom, val in monomial_value_pair)) for eq_index, monomial_value_pair in equality_constraints.items())
+        inequality_constraints = tuple((key, tuple((tuple(param_update[m] for m in monom), val) for monom, val in monom_val)) for key, monom_val in self.inequality_constraints.items())
+        auxillary_equality = tuple((key, tuple((tuple(param_update[m] for m in monom), val) for monom, val in monom_val)) for key, monom_val in self.auxillary_equality.items())
 
         # introduce variable nu for each equality/inequality
-        idx_nu = len(monom_update_reverse)
+        idx_nu = len(used_param_indices)
         idx_lambda = idx_nu + n_equality_constraints
         idx_r1 = idx_lambda + n_inequality_constraints
         idx_r2 = idx_r1 + n_inequality_constraints
-
-        # total_idx = idx_r2 + n_inequality_constraints
-        # print(f'{total_idx=}')
+        idx_dx = [idx_r2 + n_inequality_constraints]
 
         eq_buffer = collections.defaultdict(list)
 
@@ -337,77 +355,315 @@ class OptimizationMixin(abc.ABC):
         current_eq_offset += n_equality_constraints
         assert max(eq_buffer.keys()) == current_eq_offset - 1
 
-        # > inequality constraints - r1
+        # > inequality constraints - r1^2
         for eq, monom_val_list in inequality_constraints:
-            eq_buffer[current_eq_offset + eq] += monom_val_list + (((idx_r1 + eq,), -1),)
+            eq_buffer[current_eq_offset + eq] += monom_val_list + ((2 * (idx_r1 + eq,), -1),)
 
         current_eq_offset += n_inequality_constraints
         # assert max(eq_buffer.keys()) < current_eq_offset
         assert max(eq_buffer.keys()) == current_eq_offset - 1
 
-        # > x + nu * equality constraint + lambda * inequality constraint
-        # ---------------------------------------------------------------
+        # assert max(eq_buffer.keys()) < current_eq_offset + n_inequality_constraints, f'{max(eq_buffer.keys())=} is bigger than {current_eq_offset + n_inequality_constraints=}'
 
-        for index, param in enumerate(param_indices):
-            eq_buffer[current_eq_offset + index] += (((param,), 1),)
+        dx_dict = collections.defaultdict(dict)
 
-        # assert max(eq_buffer.keys()) < current_eq_offset + n_inequality_constraints, f'{max(eq_buffer.keys())=} is bigger than {current_eq_offset + n_inequality_constraints=}'
+        for _, aux_eq_data in auxillary_equality:
 
-        # differentiate equality constraints for each variable x
-        # def gen_derivative():
-        for ineq_constr_idx, eq_data in equality_constraints:
-            for p_monom, val in eq_data:
+            aux_variables = set(var for p_monomial, _ in aux_eq_data for var in p_monomial if var not in param_indices)
+            variables = set(var for p_monomial, _ in aux_eq_data for var in p_monomial if var in param_indices)
+
+            # def gen_equation():
+            # eq1 = (x, a d_a x)
+            # eq2 = (-1, a d_b x)
+            for var in variables:
+
+                # print(f'{var=}')
                 
-                p_monom_grp = collections.defaultdict(int)
-                for m in p_monom:
-                    p_monom_grp[m] += 1
+                # a x - b
+                # ((a, x), (-b,))
+                for p_monomial, val in aux_eq_data:
 
-                for m, counter in p_monom_grp.items():
-                    if m in param_indices:
+                    # print(f'{p_monomial=}')
+                    
+                    p_monomial_grp = dict(collections.Counter(p_monomial))
+
+                    # d_a, d_b
+                    # for var, counter in p_monomial_grp.items():
+                        # if var in param_indices:
+
+                    if var in p_monomial_grp:
                         def generate_monom():
-                            for i_m, i_counter in p_monom_grp.items():
-                                if m is i_m:
+                            # for each variable in the monomial
+                            for i_var, i_counter in p_monomial_grp.items():
+                                if var is i_var:
                                     sel_counter = i_counter - 1
                                 else:
                                     sel_counter = i_counter
 
+                                # print(f'{(i_var, sel_counter)=}')
+
                                 for _ in range(sel_counter):
-                                    yield i_m
+                                    yield i_var
+                        
+                        der_monomial = tuple(generate_monom())
+                        # if len(der_monomial):
+                        # print(f'{der_monomial=}')
+                        eq_buffer[current_eq_offset] += ((der_monomial, val * p_monomial_grp[var]),)
+
+                    for aux_var in aux_variables:
+                        if aux_var in p_monomial_grp:
+                            # create d_a x, d_b x, via the product rule
+                            def generate_monom_2():
+                                # for each variable in the monomial
+                                for i_var, i_counter in p_monomial_grp.items():
+                                    if i_var is aux_var:
+                                        sel_counter = i_counter - 1
+                                    else:
+                                        sel_counter = i_counter
+
+                                    for _ in range(sel_counter):
+                                        yield i_var
+
+                                if var not in dx_dict:
+                                    dx_dict[var][aux_var] = idx_dx[0]
+                                    idx_dx[0] += 1
+
+                                # product rule
+                                yield dx_dict[var][aux_var]
+                                
+                            der_monomial_2 = tuple(generate_monom_2())
+                            # print(f'{der_monomial_2=}')
+                            if len(der_monomial_2):
+                                eq_buffer[current_eq_offset] += ((der_monomial_2, val * p_monomial_grp[aux_var]),)
+
+                current_eq_offset += 1
+
+        # print(dx_dict)
 
-                        eq_idx = param_indices.index(m)
-                        der_monomial = tuple(generate_monom()) + (idx_nu + ineq_constr_idx,)
-                        eq_buffer[current_eq_offset + eq_idx] += ((der_monomial, val * counter),)
+        # @dataclasses.dataclass
+        # class VariableIndexingState:
+        #     index: int
 
-        for eq_constr_idx, eq_data in inequality_constraints:
-            for p_monom, val in eq_data:
+        #     def increment(self):
+        #         return dataclasses.replace(self, index=self.index + 1), self.index
+
+        def differentiate(data, var, dx_dict):
+            eq_buffer = tuple()
                 
-                p_monom_grp = collections.defaultdict(int)
-                for m in p_monom:
-                    p_monom_grp[m] += 1
+            # print(f'{var=}')
+
+            for p_monomial, val in data:
+
+                # count powers for each variable
+                p_monomial_grp = dict(collections.Counter(p_monomial))
+
+                # print(f'{p_monomial=}')
+
+                # for var, counter in p_monomial_grp.items():
+                if var in p_monomial_grp:
+                    def generate_monom():
+                        for i_var, i_counter in p_monomial_grp.items():
+                            if var is i_var:
+                                sel_counter = i_counter - 1
+                            else:
+                                sel_counter = i_counter
+
+                            for _ in range(sel_counter):
+                                yield i_var
+
+                    eq_buffer += ((
+                        tuple(generate_monom()), 
+                        -val * p_monomial_grp[var],
+                    ),)
+
+                if var in dx_dict:
+                    for aux_var in dx_dict[var]:
+                        if aux_var in p_monomial_grp:
+                            def generate_monom():
+                                for i_var, i_counter in p_monomial_grp.items():
+                                    if i_var is aux_var:
+                                        sel_counter = i_counter - 1
+                                    else:
+                                        sel_counter = i_counter
+
+                                    for _ in range(sel_counter):
+                                        yield i_var
+
+                                yield dx_dict[var][aux_var]
+
+                            # der_monomial = tuple(generate_monom())
+
+                            # if 1 < len(der_monomial):
+                            #     aux_eq_buffer += ((
+                            #         (der_monomial, 1), ((index_buffer[0],), -1)
+                            #     ),)
+                            #     variable_indices = (index_buffer[0],)
+                            # else:
+                            #     variable_indices = der_monomial
+                            
+                            eq_buffer += ((
+                                tuple(generate_monom()),
+                                -val * p_monomial_grp[aux_var],
+                            ),)
+
+                            # index_buffer[0] += 1
+
+            return eq_buffer
+
+        def reduce_to_quadratic_equation(eq, index):
+            eq_buffer = tuple()
+            aux_eq_buffer = tuple()
+            index_buffer = [index]
+
+            for monomial, val in eq:
+                n_aux = len(monomial) - 2
+
+                # print(f'{monomial=}')
+                # print(f'{n_aux=}')
+
+                if 0 < n_aux:
+                    eq_buffer += ((monomial[0:1] + (index_buffer[0],), val),)
+
+                    for i_aux in range(n_aux):
+                        if i_aux == n_aux - 1:
+                            aux_monom = monomial[i_aux+1:i_aux+3]
+                        else:
+                            aux_monom = monomial[i_aux+1:i_aux+2] + (index_buffer[0] + 1,)
+
+                        aux_eq_buffer += ((
+                            (aux_monom, 1.0), ((index_buffer[0],), -1.0)
+                        ),)
+                        index_buffer[0] += 1
+                else:
+                    eq_buffer += ((monomial, val),)
 
-                for m, counter in p_monom_grp.items():
-                    if m in param_indices:
-                        def generate_monom():
-                            for i_m, i_counter in p_monom_grp.items():
-                                if m is i_m:
-                                    sel_counter = i_counter - 1
-                                else:
-                                    sel_counter = i_counter
 
-                                for _ in range(sel_counter):
-                                    yield i_m
+            return eq_buffer, aux_eq_buffer, index_buffer[0]
+
+
+        # > x + nu * equality constraint + lambda * inequality constraint
+        # ---------------------------------------------------------------
+
+        for index, param in enumerate(param_indices):
+            eq_buffer[current_eq_offset + index] += (((param,), 1),)
+
+        # differentiate equality constraints for each variable x
+        for eq_constr_idx, eq_data in equality_constraints:
+            for var in param_indices:
+                equation = differentiate(eq_data, var, dx_dict)
+
+                equation = tuple((monomials + (idx_nu + eq_constr_idx,), val) for monomials, val in equation)
+
+                equation, aux_equations, idx_dx[0] = reduce_to_quadratic_equation(equation, idx_dx[0])
+
+                if len(equation):
+                    eq_idx = param_indices.index(var)
+                    # print(f'{equation=}')
+                    eq_buffer[current_eq_offset + eq_idx] += equation
+
+                # # print(f'{aux_equations=}')
+                # for aux_eq in aux_equations:
+                #     eq_buffer[next_current_eq_offset] = aux_eq
+                #     next_current_eq_offset += 1
+
+            # for p_monomial, val in ineq_data:
+                
+            #     p_monomial_grp = dict(collections.Counter(p_monomial))
+
+            #     for var, counter in p_monomial_grp.items():
+            #         if var in param_indices:
+            #             def generate_monom():
+            #                 for i_m, i_counter in p_monomial_grp.items():
+            #                     if var is i_m:
+            #                         sel_counter = i_counter - 1
+            #                     else:
+            #                         sel_counter = i_counter
+
+            #                     for _ in range(sel_counter):
+            #                         yield i_m
+
+            #             eq_idx = param_indices.index(var)
+            #             der_monomial = tuple(generate_monom()) + (idx_nu + eq_constr_idx,)
+            #             eq_buffer[current_eq_offset + eq_idx] += ((der_monomial, val * counter),)               
+
+        next_current_eq_offset = current_eq_offset + len(param_indices)
+
+        # differentiate inequality constraints for each variable x
+        for ineq_constr_idx, ineq_data in inequality_constraints:
+            # aux_variables = set(var for p_monomial, _ in ineq_data for var in p_monomial if var not in param_indices)
+            # variables = set(var for p_monomial, _ in ineq_data for var in p_monomial if var in param_indices)
+
+            for var in param_indices:
+                equation = differentiate(ineq_data, var, dx_dict)
+
+                equation = tuple((monomials + (idx_lambda + ineq_constr_idx,), val) for monomials, val in equation)
+
+                equation, aux_equations, idx_dx[0] = reduce_to_quadratic_equation(equation, idx_dx[0])
+                # idx_dx[0] = index
+
+                if len(equation):
+                    eq_idx = param_indices.index(var)
+                    # print(f'{equations=}')
+                    eq_buffer[current_eq_offset + eq_idx] += equation
+
+                # print(f'{aux_equations=}')
+                for aux_eq in aux_equations:
+                    eq_buffer[next_current_eq_offset] = aux_eq
+                    next_current_eq_offset += 1
+
+                # for p_monomial, val in ineq_data:
+
+                #     # count powers for each variable
+                #     p_monomial_grp = dict(collections.Counter(p_monomial))
+
+                #     # for var, counter in p_monomial_grp.items():
+                #     if var in p_monomial_grp:
+                #         def generate_monom():
+                #             for i_var, i_counter in p_monomial_grp.items():
+                #                 if var is i_var:
+                #                     sel_counter = i_counter - 1
+                #                 else:
+                #                     sel_counter = i_counter
+
+                #                 for _ in range(sel_counter):
+                #                     yield i_var
+
+                #         eq_idx = param_indices.index(var)
+                #         der_monomial = tuple(generate_monom()) + (idx_lambda + ineq_constr_idx,)
+                #         eq_buffer[current_eq_offset + eq_idx] += ((der_monomial, -val * counter),)
+
+                #     if var in dx_dict:
+                #         for aux_var in dx_dict[var]:
+                #             if aux_var in p_monomial_grp:
+                #                 def generate_monom():
+                #                     for i_var, i_counter in p_monomial_grp.items():
+                #                         if i_var is aux_var:
+                #                             sel_counter = i_counter - 1
+                #                         else:
+                #                             sel_counter = i_counter
+
+                #                         for _ in range(sel_counter):
+                #                             yield i_var
+
+                #                     yield dx_dict[var][aux_var]
+
+                #                 eq_idx = param_indices.index(var)
+                #                 der_monomial_2 = tuple(generate_monom())
+                #                 eq_buffer[next_current_eq_offset] = ((der_monomial_2, 1), ((idx_dx[0],), -1))
+                #                 next_current_eq_offset += 1
+
+                #                 eq_buffer[current_eq_offset + eq_idx] += (((idx_lambda + ineq_constr_idx, idx_dx[0]), -val * p_monomial_grp[aux_var]),)
+                #                 idx_dx[0] += 1
 
-                        eq_idx = param_indices.index(m)
-                        der_monomial = tuple(generate_monom()) + (idx_lambda + eq_constr_idx,)
-                        eq_buffer[current_eq_offset + eq_idx] += ((der_monomial, val * counter),)
+        # current_eq_offset += len(param_indices)
+        current_eq_offset = next_current_eq_offset
 
-        current_eq_offset += len(param_indices)
         # assert max(eq_buffer.keys()) < current_eq_offset, f'{max(eq_buffer.keys())=} is bigger than {current_eq_offset=}'
         assert max(eq_buffer.keys()) == current_eq_offset - 1
 
-        # > lambda - r2
+        # > lambda - r2^2
         for idx in range(n_inequality_constraints):
-            eq_buffer[current_eq_offset + idx] += (((idx_lambda + idx,), 1), ((idx_r2 + idx,), -1))
+            eq_buffer[current_eq_offset + idx] += (((idx_lambda + idx,), 1), (2 * (idx_r2 + idx,), -1))
 
         current_eq_offset += n_inequality_constraints
         # assert max(eq_buffer.keys()) < current_eq_offset
@@ -425,11 +681,11 @@ class OptimizationMixin(abc.ABC):
         for eq, data in auxillary_equality:
             eq_buffer[current_eq_offset + eq] += data
 
-        current_eq_offset += len(monom_update_reverse) - len(param_indices)
+        current_eq_offset += len(used_param_indices) - len(param_indices)
         # print(f'{current_eq_offset=}')
         # print(f'{len(monom_update_reverse) - len(param_indices)=}')
         assert max(eq_buffer.keys()) == current_eq_offset - 1, f'{max(eq_buffer.keys())} is not {current_eq_offset - 1}'
 
         # print(f'{current_eq_offset=}')
 
-        return eq_buffer, monom_update_reverse
+        return eq_buffer, param_update, current_eq_offset
diff --git a/polymatrix/mixins/optimizationstatemixin.py b/polymatrix/mixins/optimizationstatemixin.py
index ae688bb..731477c 100644
--- a/polymatrix/mixins/optimizationstatemixin.py
+++ b/polymatrix/mixins/optimizationstatemixin.py
@@ -48,9 +48,7 @@ class OptimizationStateMixin(abc.ABC):
 
     def update_offsets(self, polymats: tuple[PolyMatrixMixin]) -> 'OptimizationStateMixin':
         registered_polymats = set(polymat for polymat, _ in self.offset_dict.keys())
-        # print(registered_polymats)
         parametric_polymats = set(p for p in polymats if not p.is_constant and p not in registered_polymats)
-        # print(parametric_polymats)
 
         if len(parametric_polymats) == 0:
             return self
@@ -64,13 +62,15 @@ class OptimizationStateMixin(abc.ABC):
                         # number of terms is given by the maximum index + 1
                         number_of_terms = int(polymat.shape[0] * polymat.shape[1] * (monomial_to_index(self.n_var, degree*(self.n_var-1,)) + 1))
 
+                        # print(f'{polymat=}, {number_of_terms=}, {polymat.shape[0] * polymat.shape[1]=}, {monomial_to_index(self.n_var, degree*(self.n_var-1,)) + 1=}')
+
                         yield (polymat, degree), number_of_terms
 
             param_key, num_of_terms = tuple(zip(*gen_n_param_per_struct()))
             cum_sum = tuple(itertools.accumulate((self.n_param,) + num_of_terms))
             offset_dict = dict(zip(param_key, itertools.pairwise(cum_sum)))
 
-            return dataclasses.replace(self, offset_dict=offset_dict, n_param=cum_sum[-1])
+            return dataclasses.replace(self, offset_dict=self.offset_dict | offset_dict, n_param=cum_sum[-1])
 
     def update_param(self, n):
         return dataclasses.replace(self, n_param=self.n_param+n)
diff --git a/polymatrix/polysolver.py b/polymatrix/polysolver.py
index 24c7095..1f7bc8e 100644
--- a/polymatrix/polysolver.py
+++ b/polymatrix/polysolver.py
@@ -12,12 +12,11 @@ import time
 def substitude_x_add_a(data, a):    
 
     def kron_eye_and_a(n, data, a, indices):
-        permutations = more_itertools.distinct_permutations(indices)
-        # permutations = set(itertools.permutations(indices))
+        # permutations = more_itertools.distinct_permutations(indices)
+        permutations = set(itertools.permutations(indices))
         
         if not scipy.sparse.issparse(data):
-            a_mat = np.reshape(a, (-1, 1))
-            ops = (np.eye(n), a_mat)
+            ops = (np.eye(n), np.reshape(a, (-1, 1)))
 
             def acc_kron(perm):
                 *_, last = itertools.accumulate((ops[idx] for idx in perm), lambda acc, v: np.kron(acc, v))
@@ -103,6 +102,8 @@ def substitude_x_add_a(data, a):
                 yield key, d_data
     return dict(gen_subs_data())
 
+    # return subs_data
+
 
 def solve_poly_system(data, m):
     if isinstance(data[1], np.ndarray):
@@ -267,17 +268,24 @@ def outer_smart_solve(data, n_iter=10, a_max=1.0, irange=None, idegree=None):
     def gen_a_err():
         n = data[1].shape[0]
 
-        for idx_iter in range(n_iter):
+        for _ in range(n_iter):
             a_init = a_max * (2 * np.random.rand(n) - 1)
 
-            a, err = inner_smart_solve(data, irange=irange, idegree=idegree, a_init=a_init)
+            try:
+                a, err = inner_smart_solve(data, irange=irange, idegree=idegree, a_init=a_init)
 
-            subs_data = substitude_x_add_a(data, a[-1])
-            sol = solve_poly_system(subs_data, 6)
-            error_index = np.max(np.abs(eval_solution(subs_data, sol))) + max(a[-1])
-            # error_index = np.max(np.abs(eval_solution(subs_data, sol)))
+                subs_data = substitude_x_add_a(data, a[-1])
+                sol = solve_poly_system(subs_data, 6)
+                error_index = np.max(np.abs(eval_solution(subs_data, sol))) + max(a[-1])
+                # error_index = np.max(np.abs(eval_solution(subs_data, sol)))
+            
+            except:
+                print(f'nan error, continue')
+                # print(f'nan error for {a_init=}, continue')
+                # yield np.nan, a_init, np.nan
+                continue
 
-            # print(f'{error_index=}')
+            print(f'{error_index=}')
 
             yield error_index, a, err
 
diff --git a/polymatrix/sympyutils.py b/polymatrix/sympyutils.py
index 7413ba3..ac91b8f 100644
--- a/polymatrix/sympyutils.py
+++ b/polymatrix/sympyutils.py
@@ -55,23 +55,23 @@ def poly_to_data_coord(poly_list, x, degree = None):
 def poly_to_matrix(poly_list, x, power = None):
     data_coord_dict = poly_to_data_coord(poly_list, x, power)
 
-    n_free_symbols = len(x)
+    n_var = len(x)
     n_poly = len(poly_list)
     n_eq = sum((1 for row_eq in poly_list for e in row_eq))
 
     def gen_power_mat():
         if 0 not in data_coord_dict:
-            yield 0, np.zeros((n_free_symbols, 1))
+            yield 0, np.zeros((n_var, 1))
 
         # for all powers generate a matrix
         for degree, data_coord in data_coord_dict.items():
 
             # empty matrix
-            shape = (n_eq, n_free_symbols**degree)
+            shape = (n_eq, n_var**degree)
 
             # fill matrix
             if len(data_coord) == 0:
-                yield np.zeros((len(poly_list), n_free_symbols**degree))
+                yield np.zeros((len(poly_list), n_var**degree))
 
             else:
                 def gen_row_col_coord():
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,
         )