add polynomial operations for sos optimization

1 files changed, 86 insertions, 373 deletions

diff --git a/polymatrix/expression/mixins/derivativeexprmixin.py b/polymatrix/expression/mixins/derivativeexprmixin.py
index b403873..4fc48a5 100644
--- a/polymatrix/expression/mixins/derivativeexprmixin.py
+++ b/polymatrix/expression/mixins/derivativeexprmixin.py
@@ -10,6 +10,7 @@ from polymatrix.expression.init.initpolymatrix import init_poly_matrix
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
 from polymatrix.expression.polymatrix import PolyMatrix
 from polymatrix.expression.expressionstate import ExpressionState
+from polymatrix.expression.utils.getderivativemonomials import get_derivative_monomials
 from polymatrix.expression.utils.getvariableindices import get_variable_indices
 
 
@@ -36,112 +37,119 @@ class DerivativeExprMixin(ExpressionBaseMixin):
     ) -> tuple[ExpressionState, PolyMatrix]:
 
         state, underlying = self.underlying.apply(state=state)
+
+        assert underlying.shape[1] == 1, f'{underlying.shape=}'
         
         state, diff_wrt_variables = get_variable_indices(state, self.variables)
 
-        def get_derivative_terms(
-            monomial_terms, 
-            diff_wrt_variable: int, 
-            state: ExpressionState,
-            considered_variables: set,
-        ):
+        # def get_derivative_terms(
+        #     monomial_terms, 
+        #     diff_wrt_variable: int, 
+        #     state: ExpressionState,
+        #     considered_variables: set,
+        # ):
+
+        #     if self.introduce_derivatives:
 
-            if self.introduce_derivatives:
+        #         raise Exception('not implemented')
                 
-                def gen_new_variables():
-                    for monomial in monomial_terms.keys():
-                        for var in monomial:
-                            if var not in diff_wrt_variables and var not in considered_variables:
-                                yield var
+        #         def gen_new_variables():
+        #             for monomial in monomial_terms.keys():
+        #                 for var in monomial:
+        #                     if var not in diff_wrt_variables and var not in considered_variables:
+        #                         yield var
 
-                new_variables = set(gen_new_variables())
+        #         new_variables = set(gen_new_variables())
 
-                new_considered_variables = considered_variables | new_variables
+        #         new_considered_variables = considered_variables | new_variables
 
-                def acc_state_candidates(acc, new_variable):
-                    state, candidates = acc
+        #         def acc_state_candidates(acc, new_variable):
+        #             state, candidates = acc
 
-                    key = init_derivative_key(
-                        variable=new_variable,
-                        with_respect_to=diff_wrt_variable,
-                    )
-                    state = state.register(key=key, n_param=1)
+        #             key = init_derivative_key(
+        #                 variable=new_variable,
+        #                 with_respect_to=diff_wrt_variable,
+        #             )
+        #             state = state.register(key=key, n_param=1)
 
-                    state, auxillary_derivation_terms = get_derivative_terms(
-                        monomial_terms=state.auxillary_equations[new_variable],
-                        diff_wrt_variable=diff_wrt_variable,
-                        state=state,
-                        considered_variables=new_considered_variables,
-                    )
+        #             state, auxillary_derivation_terms = get_derivative_terms(
+        #                 monomial_terms=state.auxillary_equations[new_variable],
+        #                 diff_wrt_variable=diff_wrt_variable,
+        #                 state=state,
+        #                 considered_variables=new_considered_variables,
+        #             )
 
-                    if 1 < len(auxillary_derivation_terms):
-                        derivation_variable = state.offset_dict[key][0]
+        #             if 1 < len(auxillary_derivation_terms):
+        #                 derivation_variable = state.offset_dict[key][0]
 
-                        state = dataclasses.replace(
-                            state,
-                            auxillary_equations=state.auxillary_equations | {derivation_variable: auxillary_derivation_terms},
-                        )
+        #                 state = dataclasses.replace(
+        #                     state,
+        #                     auxillary_equations=state.auxillary_equations | {derivation_variable: auxillary_derivation_terms},
+        #                 )
 
-                        return state, candidates + (new_variable,)
+        #                 return state, candidates + (new_variable,)
 
-                    else:
-                        return state, candidates
+        #             else:
+        #                 return state, candidates
 
-                *_, (state, confirmed_variables) = itertools.accumulate(
-                    new_variables,
-                    acc_state_candidates,
-                    initial=(state, tuple()),
-                )
+        #         *_, (state, confirmed_variables) = itertools.accumulate(
+        #             new_variables,
+        #             acc_state_candidates,
+        #             initial=(state, tuple()),
+        #         )
 
-            else:
-                confirmed_variables = tuple()
+        #     else:
+        #         confirmed_variables = tuple()
 
-            derivation_terms = collections.defaultdict(float)
+        #     derivation_terms = collections.defaultdict(float)
 
-            for monomial, value in monomial_terms.items():
+        #     for monomial, value in monomial_terms.items():
 
-                # count powers for each variable
-                monomial_cnt = dict(collections.Counter(monomial))
+        #         # # count powers for each variable
+        #         # monomial_cnt = dict(collections.Counter(monomial))
+        #         monomial_cnt = dict(monomial)
 
-                def differentiate_monomial(dependent_variable, derivation_variable=None):
-                    def gen_diff_monomial():
-                        for current_variable, current_count in monomial_cnt.items():
+        #         def differentiate_monomial(dependent_variable, derivation_variable=None):
+        #             def gen_diff_monomial():
+        #                 for current_variable, current_count in monomial:
 
-                            if current_variable is dependent_variable:
-                                sel_counter = current_count - 1
+        #                     if current_variable is dependent_variable:
+        #                         sel_counter = current_count - 1
 
-                            else:
-                                sel_counter = current_count
+        #                     else:
+        #                         sel_counter = current_count
 
-                            for _ in range(sel_counter):
-                                yield current_variable
+        #                     # for _ in range(sel_counter):
+        #                     #     yield current_variable
+        #                     yield current_variable, sel_counter
 
-                        if derivation_variable is not None:
-                            yield derivation_variable
+        #                 if derivation_variable is not None:
+        #                     yield derivation_variable
 
-                    diff_monomial = tuple(sorted(gen_diff_monomial()))
+        #             diff_monomial = tuple(sorted(gen_diff_monomial()))
 
-                    return diff_monomial, value * monomial_cnt[dependent_variable]
+        #             return diff_monomial, value * monomial_cnt[dependent_variable]
 
-                if diff_wrt_variable in monomial_cnt:
-                    diff_monomial, value = differentiate_monomial(diff_wrt_variable)
-                    derivation_terms[diff_monomial] += value
+        #         if diff_wrt_variable in monomial_cnt:
+        #             diff_monomial, value = differentiate_monomial(diff_wrt_variable)
+        #             derivation_terms[diff_monomial] += value
 
-                for candidate_variable in monomial_cnt.keys():
-                    if candidate_variable in considered_variables or candidate_variable in confirmed_variables:
-                        key = init_derivative_key(
-                            variable=candidate_variable,
-                            with_respect_to=diff_wrt_variable,
-                        )
-                        derivation_variable = state.offset_dict[key][0]
+        #         # only used if introduce_derivatives == True
+        #         for candidate_variable in monomial_cnt.keys():
+        #             if candidate_variable in considered_variables or candidate_variable in confirmed_variables:
+        #                 key = init_derivative_key(
+        #                     variable=candidate_variable,
+        #                     with_respect_to=diff_wrt_variable,
+        #                 )
+        #                 derivation_variable = state.offset_dict[key][0]
 
-                        diff_monomial, value = differentiate_monomial(
-                            dependent_variable=candidate_variable, 
-                            derivation_variable=derivation_variable,
-                        )
-                        derivation_terms[diff_monomial] += value
+        #                 diff_monomial, value = differentiate_monomial(
+        #                     dependent_variable=candidate_variable, 
+        #                     derivation_variable=derivation_variable,
+        #                 )
+        #                 derivation_terms[diff_monomial] += value
 
-            return state, dict(derivation_terms)
+        #     return state, dict(derivation_terms)
 
         terms = {}
 
@@ -155,11 +163,12 @@ class DerivativeExprMixin(ExpressionBaseMixin):
             # derivate each variable and map result to the corresponding column
             for col, diff_wrt_variable in enumerate(diff_wrt_variables):
 
-                state, derivation_terms = get_derivative_terms(
+                state, derivation_terms = get_derivative_monomials(
                     monomial_terms=underlying_terms,
                     diff_wrt_variable=diff_wrt_variable,
                     state=state,
                     considered_variables=set(),
+                    introduce_derivatives=self.introduce_derivatives,
                 )
 
                 if 0 < len(derivation_terms):
@@ -171,299 +180,3 @@ class DerivativeExprMixin(ExpressionBaseMixin):
         )
 
         return state, poly_matrix   
-
-        # def get_derivative_terms(
-        #     monomial_terms, 
-        #     diff_wrt_variable: int, 
-        #     state: PolyMatrixExprState,
-        #     considered_variables: set,
-        #     # implement_derivation: bool,
-        # ):
-        #     derivation_terms = collections.defaultdict(float)
-
-        #     other_independent_variables = tuple(var for var in diff_wrt_variables if var is not diff_wrt_variable)
-
-        #     # print(other_independent_variables)
-        #     # print(tuple(variable for monomial in monomial_terms.keys() for variable in monomial))
-
-        #     if sum(variable not in other_independent_variables for monomial in monomial_terms.keys() for variable in monomial) < 2:
-        #         return {}, state
-
-        #     # if not implement_derivation:
-        #     #     implement_derivation = any(diff_wrt_variable in monomial for monomial in monomial_terms.keys())
-
-        #     for monomial, value in monomial_terms.items():
-
-        #         # count powers for each variable
-        #         monomial_cnt = dict(collections.Counter(monomial))
-
-        #         def differentiate_monomial(dependent_variable, derivation_variable=None):
-        #             def gen_diff_monomial():
-        #                 for current_variable, current_count in monomial_cnt.items():
-
-        #                     if current_variable is dependent_variable:
-        #                         sel_counter = current_count - 1
-
-        #                     else:
-        #                         sel_counter = current_count
-
-        #                     for _ in range(sel_counter):
-        #                         yield current_variable
-
-        #                 if derivation_variable is not None:
-        #                     yield derivation_variable
-
-        #             diff_monomial = tuple(gen_diff_monomial())
-
-        #             return diff_monomial, value * monomial_cnt[dependent_variable]
-
-        #         if diff_wrt_variable in monomial_cnt:
-        #             diff_monomial, value = differentiate_monomial(diff_wrt_variable)
-        #             derivation_terms[diff_monomial] += value
-
-        #         if self.introduce_derivatives:
-
-        #             def gen_derivation_keys():
-        #                 for variable in monomial_cnt.keys():
-        #                     if variable not in diff_wrt_variables:
-        #                         yield variable
-
-        #             candidate_variables = tuple(gen_derivation_keys())
-
-        #             new_considered_derivations = considered_variables | set(candidate_variables)
-
-        #             for candidate_variable in candidate_variables:
-
-        #                 # introduce new auxillary equation
-        #                 if candidate_variable not in considered_variables:
-        #                     auxillary_derivation_terms, state = get_derivative_terms(
-        #                         monomial_terms=state.auxillary_equations[candidate_variable],
-        #                         diff_wrt_variable=diff_wrt_variable,
-        #                         state=state,
-        #                         considered_variables=new_considered_derivations,
-        #                         # implement_derivation=implement_derivation,
-        #                     )
-
-        #                     if 0 < len(auxillary_derivation_terms):
-        #                         key = init_derivative_key(
-        #                             variable=candidate_variable,
-        #                             with_respect_to=diff_wrt_variable,
-        #                         )
-        #                         state = state.register(key=key, n_param=1)
-        #                         derivation_variable = state.offset_dict[key][0]
-
-        #                         state = dataclasses.replace(
-        #                             state,
-        #                             auxillary_equations=state.auxillary_equations | {derivation_variable: auxillary_derivation_terms},
-        #                         )
-
-        #                 else:
-
-        #                     key = init_derivative_key(
-        #                         variable=candidate_variable,
-        #                         with_respect_to=diff_wrt_variable,
-        #                     )
-        #                     state = state.register(key=key, n_param=1)
-        #                     derivation_variable = state.offset_dict[key][0]
-
-        #                     diff_monomial, value = differentiate_monomial(
-        #                         dependent_variable=candidate_variable, 
-        #                         derivation_variable=derivation_variable,
-        #                     )
-        #                     derivation_terms[diff_monomial] += value
-
-        #     return dict(derivation_terms), state
-
-        # terms = {}
-
-        # for row in range(self.shape[0]):
-
-        #     try:
-        #         underlying_terms = underlying.get_poly(row, 0)
-        #     except KeyError:
-        #         continue
-
-        #     # derivate each variable and map result to the corresponding column
-        #     for col, diff_wrt_variable in enumerate(diff_wrt_variables):
-
-        #         derivation_terms, state = get_derivative_terms(
-        #             monomial_terms=underlying_terms,
-        #             diff_wrt_variable=diff_wrt_variable,
-        #             state=state,
-        #             considered_variables=set(),
-        #             # implement_derivation=False,
-        #         )
-
-        #         if 0 < len(derivation_terms):
-        #             terms[row, col] = derivation_terms
-
-        # poly_matrix = init_poly_matrix(
-        #     terms=terms,
-        #     shape=self.shape,
-        # )
-
-        # return state, poly_matrix   
-
-
-
-        # state = [state]
-
-        # state, underlying = self.underlying.apply(state=state)
-
-        # match self.variables:
-        #     case ExpressionBaseMixin():
-        #         assert self.variables.shape[1] == 1
-
-        #         state, dependent_variables = self.variables.apply(state)
-
-        #         def gen_indices():
-        #             for row in range(dependent_variables.shape[0]):
-        #                 for monomial in dependent_variables.get_poly(row, 0).keys():
-        #                     yield monomial[0]
-
-        #         variable_indices = tuple(sorted(gen_indices()))
-
-        #     case _:
-        #         def gen_indices():
-        #             for variable in self.variables:
-        #                 if variable in state.offset_dict:
-        #                     yield state.offset_dict[variable][0]
-
-        #         variable_indices = tuple(sorted(gen_indices()))
-
-        # terms = {}
-        # derivations_keys = set()
-
-        # # derivate each variable and map result to the corresponding column
-        # for col, derivation_variable in enumerate(variable_indices):
-            
-        #     def get_derivative_terms(monomial_terms):
-
-        #         terms_row_col = collections.defaultdict(float)
-
-        #         for monomial, value in monomial_terms.items():
-
-        #             # count powers for each variable
-        #             monomial_cnt = dict(collections.Counter(monomial))
-
-        #             variable_candidates = tuple()
-
-        #             if derivation_variable in monomial_cnt:
-        #                 variable_candidates += ((derivation_variable, None),)
-
-        #             if self.introduce_derivatives:
-        #                 def gen_dependent_variables():
-        #                     for dependent_variable in monomial_cnt.keys():
-        #                         if dependent_variable not in variable_indices:
-        #                             derivation_key = init_derivative_key(
-        #                                 variable=dependent_variable,
-        #                                 with_respect_to=derivation_variable,
-        #                             )
-        #                             derivations_keys.add(derivation_key)
-        #                             state = state.register(key=derivation_key, n_param=1)
-        #                             yield dependent_variable, derivation_key
-
-        #                 variable_candidates += tuple(gen_dependent_variables())
-
-        #             for variable_candidate, derivation_key in variable_candidates:
-
-        #                 def generate_monomial():
-        #                     for current_variable, current_count in monomial_cnt.items():
-
-        #                         if current_variable is variable_candidate:
-        #                             sel_counter = current_count - 1
-
-        #                         else:
-        #                             sel_counter = current_count
-
-        #                         for _ in range(sel_counter):
-        #                             yield current_variable
-
-        #                     if derivation_key is not None:
-        #                         yield state.offset_dict[derivation_key][0]
-
-        #                 col_monomial = tuple(generate_monomial())
-
-        #                 terms_row_col[col_monomial] += value * monomial_cnt[variable_candidate]
-
-        #         return dict(terms_row_col)
-
-        #     for row in range(self.shape[0]):
-
-        #         try:
-        #             underlying_terms = underlying.get_poly(row, 0)
-        #         except KeyError:
-        #             continue
-
-        #         derivative_terms = get_derivative_terms(underlying_terms)
-
-        #         if 0 < len(derivative_terms):
-        #             terms[row, col] = derivative_terms
-
-        # derivation_variables = collections.defaultdict(list)
-        # for derivation_key in derivations_keys:
-        #     derivation_variables[derivation_key.with_respect_to].append(derivation_key)
-
-        # aux_der_terms = []
-
-        # for derivation_variable, derivation_keys in derivation_variables.items():
-
-        #     dependent_variables = tuple(derivation_key.variable for derivation_key in derivation_keys)
-
-            
-        #     for aux_terms in state.auxillary_equations:
-
-        #         # only intoduce a new auxillary equation if there is a monomial containing at least one dependent variable 
-        #         if any(variable in dependent_variables for monomial in aux_terms.keys() for variable in monomial):
-
-        #             terms_row_col = collections.defaultdict(float)
-
-        #             # for each monomial
-        #             for aux_monomial, value in aux_terms.items():
-
-        #                 # count powers for each variable
-        #                 monomial_cnt = dict(collections.Counter(aux_monomial))
-
-        #                 variable_candidates = tuple()
-
-        #                 if derivation_variable in monomial_cnt:
-        #                     variable_candidates += ((derivation_variable, None),)
-
-        #                 # add dependent variables
-        #                 variable_candidates += tuple((derivation_key.variable, derivation_key) for derivation_key in derivation_keys if derivation_key.variable in monomial_cnt)
-
-        #                 for variable_candidate, derivative_key in variable_candidates:
-
-        #                     def generate_monomial():
-        #                         for current_variable, current_count in monomial_cnt.items():
-
-        #                             if current_variable is variable_candidate:
-        #                                 sel_counter = current_count - 1
-
-        #                             else:
-        #                                 sel_counter = current_count
-
-        #                             for _ in range(sel_counter):
-        #                                 yield current_variable
-
-        #                         if derivative_key is not None:
-        #                             yield state.offset_dict[derivative_key][0]
-
-        #                     col_monomial = tuple(generate_monomial())
-
-        #                     terms_row_col[col_monomial] += value * monomial_cnt[variable_candidate]
-
-        #             if 0 < len(terms_row_col):
-        #                 aux_der_terms.append(dict(terms_row_col))
-        
-        # state = dataclasses.replace(
-        #     state,
-        #     auxillary_equations=state.auxillary_equations + tuple(aux_der_terms),
-        # )
-
-        # poly_matrix = init_poly_matrix(
-        #     terms=terms,
-        #     shape=self.shape,
-        # )
-
-        # return state, poly_matrix