Add HessianExpr

5 files changed, 138 insertions, 4 deletions

diff --git a/polymatrix/expression/expression.py b/polymatrix/expression/expression.py
index fc4d3b6..2e4d544 100644
--- a/polymatrix/expression/expression.py
+++ b/polymatrix/expression/expression.py
@@ -184,6 +184,14 @@ class Expression(ExpressionBaseMixin, ABC):
             ),
         )
 
+    def hessian(self, variables: Expression) -> Expression:
+        return self.copy(
+            underlying=init.init_hessian_expr(
+                underlying=self.underlying,
+                variables=variables
+            )
+        )
+
     def eval(self, variable: tuple, value: tuple[float, ...] | None = None,) -> Expression:
         return self.copy(
             underlying=init.init_eval_expr(
diff --git a/polymatrix/expression/impl.py b/polymatrix/expression/impl.py
index a185802..f919dc7 100644
--- a/polymatrix/expression/impl.py
+++ b/polymatrix/expression/impl.py
@@ -29,6 +29,7 @@ from polymatrix.expression.mixins.fromstatemonad import FromStateMonadMixin
 from polymatrix.expression.mixins.fromsymmetricmatrixexprmixin import FromSymmetricMatrixExprMixin
 from polymatrix.expression.mixins.fromsympyexprmixin import FromSympyExprMixin
 from polymatrix.expression.mixins.halfnewtonpolytopeexprmixin import HalfNewtonPolytopeExprMixin
+from polymatrix.expression.mixins.hessianexprmixin import HessianExprMixin
 from polymatrix.expression.mixins.integrateexprmixin import IntegrateExprMixin
 from polymatrix.expression.mixins.legendreseriesmixin import LegendreSeriesMixin
 from polymatrix.expression.mixins.linearinexprmixin import LinearInExprMixin
@@ -255,6 +256,15 @@ class HalfNewtonPolytopeExprImpl(HalfNewtonPolytopeExprMixin):
 
 
 @dataclassabc.dataclassabc(frozen=True)
+class HessianExprImpl(HessianExprMixin):
+    underlying: ExpressionBaseMixin
+    variables: ExpressionBaseMixin
+
+    def __str__(self):
+        return f"(∂²_{self.variables} {self.underlying})"
+
+
+@dataclassabc.dataclassabc(frozen=True)
 class IntegrateExprImpl(IntegrateExprMixin):
     underlying: ExpressionBaseMixin
     variables: ExpressionBaseMixin
diff --git a/polymatrix/expression/init.py b/polymatrix/expression/init.py
index 6ed517e..15557f4 100644
--- a/polymatrix/expression/init.py
+++ b/polymatrix/expression/init.py
@@ -250,6 +250,14 @@ def init_half_newton_polytope_expr(
     )
 
 
+def init_hessian_expr(
+    underlying: ExpressionBaseMixin,
+    variables: ExpressionBaseMixin
+):
+    return polymatrix.expression.impl.HessianExprImpl(
+            underlying=underlying, variables=variables)
+
+
 def init_linear_matrix_in_expr(
     underlying: ExpressionBaseMixin,
     variable: int,
diff --git a/polymatrix/expression/mixins/hessianexprmixin.py b/polymatrix/expression/mixins/hessianexprmixin.py
new file mode 100644
index 0000000..0a89d5b
--- /dev/null
+++ b/polymatrix/expression/mixins/hessianexprmixin.py
@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing_extensions import override
+
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+
+from polymatrix.expressionstate import ExpressionState
+from polymatrix.polymatrix.index import PolyMatrixDict, PolyDict
+from polymatrix.polymatrix.init import init_poly_matrix
+from polymatrix.polymatrix.mixins import PolyMatrixMixin
+
+class HessianExprMixin(ExpressionBaseMixin):
+    """
+    Compute the hessian matrix (second derivative) of a scalar polynomial.
+    """
+
+    @property
+    @abstractmethod
+    def underlying(self) -> ExpressionBaseMixin:
+        """ Expression to take the derivative of """
+
+    @property
+    @abstractmethod
+    def variables(self) -> ExpressionBaseMixin:
+        """
+        Variables with respect to which the derivative is taken.
+        This must be a column vector.
+        """
+
+    @override
+    def apply(self, state: ExpressionState) -> tuple[ExpressionState, PolyMatrixMixin]:
+        state, underlying = self.underlying.apply(state)
+        state, variables = self.variables.apply(state)
+
+        # Check arguments
+        if underlying.shape != (1, 1):
+            raise ValueError("Cannot take Hessian of non-scalar expression "
+                             f"with shape {underlying.shape}")
+
+    
+        if variables.shape[1] != 1:
+            raise ValueError("Cannot take variable with repect to matrix "
+                             f"with shape {variables.shape}. It must be a column vector")
+
+
+        # number of variables
+        nvars, _ = variables.shape
+        variable_indices = []
+        for i in range(nvars):
+            monomials = tuple(variables.at(i, 0).monomials())
+            if len(monomials) > 1:
+                # FIXME improve error message
+                raise ValueError("Cannot take derivative with respect to polynomial")
+
+            monomial = monomials[0]
+            if monomial.degree != 1:
+                # FIXME improve error message
+                raise ValueError("Cannot take derivative with respect to non-linear term")
+
+            variable_indices.append(monomial[0].index)
+
+        # Actually compute derivative
+        result = PolyMatrixDict.empty()
+        for i in range(nvars):
+            # FIXME: make only upper triangular part and create symmetrix polymatrix
+            # FIXME: before that, take out SymmetrixPolyMatrix from SymmetrixExprMixin
+            # for j in range(i, nvars):
+            for j in range(nvars):
+                new_poly = PolyDict.differentiate(
+                    PolyDict.differentiate(
+                        underlying.scalar(),
+                        variable_indices[j]),
+                variable_indices[i])
+
+                if new_poly:
+                    result[i, j] = new_poly
+
+        return state, init_poly_matrix(result, shape=(nvars, nvars))
diff --git a/polymatrix/polymatrix/index.py b/polymatrix/polymatrix/index.py
index fd2a759..e4bd7d1 100644
--- a/polymatrix/polymatrix/index.py
+++ b/polymatrix/polymatrix/index.py
@@ -156,10 +156,27 @@ class MonomialIndex(tuple[VariableIndex]):
         return MonomialIndex.sort(result)
 
     @staticmethod
-    def differentiate(index: MonomialIndex, wrt: int) -> MonomialIndex | None:
-        # TODO: implement this and then PolyDict.differentiate to replace
-        # expression.utils.gederivativemonomials.differentiate_polynomial
-        raise NotImplementedError
+    def differentiate(index: MonomialIndex, wrt: int) -> tuple[int, MonomialIndex | None]:
+        """
+        Differentiate a monomial with respect to the variable with index `wrt`.
+        If `index` does not contain the variable `wrt` (the term disappears)
+        returns None. The function also returns the degree of the monomial
+        before it was differentiated.
+        """
+        old_pow, new_monom = None, []
+        for v in index:
+            if v.index == wrt:
+                old_pow = v.power
+                if old_pow > 1:
+                    new_monom.append(VariableIndex(index=v.index,
+                                                   power=old_pow - 1))
+            else:
+                new_monom.append(v)
+
+        if old_pow is None:
+            return 0, None
+        
+        return old_pow, MonomialIndex(tuple(new_monom))
         
 
 class PolyDict(UserDict[MonomialIndex, int | float]):
@@ -251,6 +268,18 @@ class PolyDict(UserDict[MonomialIndex, int | float]):
 
         return p
 
+    @staticmethod
+    def differentiate(poly: PolyDict, wrt: int) -> PolyDict:
+        """ Differentiate a varaible with respect to a variable """
+        p = PolyDict.empty()
+        for monomial, coeff in poly.terms():
+            old_degree, m = MonomialIndex.differentiate(monomial, wrt)
+            # term has disappeared
+            if m is not None:
+                p[m] = coeff * old_degree
+
+        return p
+            
 
 class MatrixIndex(NamedTuple):
     """ Index to represent an entry of a matrix. """