path: root/polymatrix/expression/mixins/derivativeexprmixin.py

import abc
import collections
import dataclasses
import itertools
import typing
from polymatrix.expression.init.initderivativekey import init_derivative_key

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.getvariableindices import get_variable_indices


class DerivativeExprMixin(ExpressionBaseMixin):
    @property
    @abc.abstractmethod
    def underlying(self) -> ExpressionBaseMixin:
        ...

    @property
    @abc.abstractmethod
    def variables(self) -> typing.Union[tuple, ExpressionBaseMixin]:
        ...

    @property
    @abc.abstractmethod
    def introduce_derivatives(self) -> bool:
        ...

    # overwrites abstract method of `ExpressionBaseMixin`
    def _apply(
        self, 
        state: ExpressionState,
    ) -> tuple[ExpressionState, PolyMatrix]:

        state, underlying = self.underlying.apply(state=state)
        
        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,
        ):

            if self.introduce_derivatives:
                
                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_considered_variables = considered_variables | new_variables

                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)

                    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]

                        state = dataclasses.replace(
                            state,
                            auxillary_equations=state.auxillary_equations | {derivation_variable: auxillary_derivation_terms},
                        )

                        return state, candidates + (new_variable,)

                    else:
                        return state, candidates

                *_, (state, confirmed_variables) = itertools.accumulate(
                    new_variables,
                    acc_state_candidates,
                    initial=(state, tuple()),
                )

            else:
                confirmed_variables = tuple()

            derivation_terms = collections.defaultdict(float)

            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(sorted(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

                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

            return state, dict(derivation_terms)

        terms = {}

        for row in range(underlying.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):

                state, derivation_terms = get_derivative_terms(
                    monomial_terms=underlying_terms,
                    diff_wrt_variable=diff_wrt_variable,
                    state=state,
                    considered_variables=set(),
                )

                if 0 < len(derivation_terms):
                    terms[row, col] = derivation_terms

        poly_matrix = init_poly_matrix(
            terms=terms,
            shape=(underlying.shape[0], len(diff_wrt_variables)),
        )

        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