move init expr functions into a single file

1 files changed, 21 insertions, 2 deletions

diff --git a/polymatrix/expression/mixins/derivativeexprmixin.py b/polymatrix/expression/mixins/derivativeexprmixin.py
index 183608b..15873b6 100644
--- a/polymatrix/expression/mixins/derivativeexprmixin.py
+++ b/polymatrix/expression/mixins/derivativeexprmixin.py
@@ -8,9 +8,19 @@ from polymatrix.polymatrix.polymatrix import PolyMatrix
 from polymatrix.expressionstate.expressionstate import ExpressionState
 from polymatrix.expression.utils.getderivativemonomials import get_derivative_monomials
 from polymatrix.expression.utils.getvariableindices import get_variable_indices_from_variable
+from polymatrix.utils.getstacklines import FrameSummary
+from polymatrix.utils.tooperatorexception import to_operator_exception
 
 
 class DerivativeExprMixin(ExpressionBaseMixin):
+    """
+    differentiate w.r.t. x:
+
+    [[x**2]]  ->  [[2*x]]
+
+    introduce_derivatives: not used at the moment
+    """
+
     @property
     @abc.abstractmethod
     def underlying(self) -> ExpressionBaseMixin:
@@ -26,6 +36,11 @@ class DerivativeExprMixin(ExpressionBaseMixin):
     def introduce_derivatives(self) -> bool:
         ...
 
+    @property
+    @abc.abstractmethod
+    def stack(self) -> tuple[FrameSummary]:
+        ...
+
     # overwrites abstract method of `ExpressionBaseMixin`
     def apply(
         self, 
@@ -35,7 +50,11 @@ class DerivativeExprMixin(ExpressionBaseMixin):
         state, underlying = self.underlying.apply(state=state)        
         state, diff_wrt_variables = get_variable_indices_from_variable(state, self.variables)
 
-        assert underlying.shape[1] == 1, f'{underlying.shape=}'
+        if not (underlying.shape[1] == 1):
+            raise AssertionError(to_operator_exception(
+                message=f'{underlying.shape[1]=} is not 1',
+                stack=self.stack,
+            ))
 
         terms = {}
 
@@ -52,7 +71,7 @@ class DerivativeExprMixin(ExpressionBaseMixin):
                     monomial_terms=underlying_terms,
                     diff_wrt_variable=diff_wrt_variable,
                     state=state,
-                    considered_variables=set(),
+                    considered_variables=set(diff_wrt_variables),
                     introduce_derivatives=self.introduce_derivatives,
                 )