Fix bug in .diff() and add DerivativeExpr.__str__

2 files changed, 6 insertions, 3 deletions

diff --git a/polymatrix/expression/expression.py b/polymatrix/expression/expression.py
index 49164d0..58a3883 100644
--- a/polymatrix/expression/expression.py
+++ b/polymatrix/expression/expression.py
@@ -173,8 +173,8 @@ class Expression(ExpressionBaseMixin, ABC):
     def diff(self, variables: Expression, introduce_derivatives: bool | None = None) -> Expression:
         return self.copy(
             underlying=diff(
-                expression=self,
-                variables=variables,
+                expression=self.underlying,
+                variables=variables.underlying,
                 introduce_derivatives=introduce_derivatives,
             ),
         )
diff --git a/polymatrix/expression/impl.py b/polymatrix/expression/impl.py
index 65cc2fe..aedc7da 100644
--- a/polymatrix/expression/impl.py
+++ b/polymatrix/expression/impl.py
@@ -133,7 +133,7 @@ class ConcatenateExprImpl(ConcatenateExprMixin):
 @dataclassabc.dataclassabc(frozen=True)
 class DerivativeExprImpl(DerivativeExprMixin):
     underlying: ExpressionBaseMixin
-    variables: tuple
+    variables: ExpressionBaseMixin
     introduce_derivatives: bool
     stack: tuple[FrameSummary]
 
@@ -141,6 +141,9 @@ class DerivativeExprImpl(DerivativeExprMixin):
     def __repr__(self):
         return f"{self.__class__.__name__}(variables={self.variables}, underlying={repr(self.underlying)})"
 
+    def __str__(self):
+        return f"∂_{self.variables} {self.underlying}"
+
 
 @dataclassabc.dataclassabc(frozen=True)
 class DiagExprImpl(DiagExprMixin):