Create init function for PolyMatrixAsAffineExpression, fix methods

3 files changed, 93 insertions, 11 deletions

diff --git a/polymatrix/__init__.py b/polymatrix/__init__.py
index af1d5fa..a808d92 100644
--- a/polymatrix/__init__.py
+++ b/polymatrix/__init__.py
@@ -10,11 +10,14 @@ from polymatrix.expression import product as internal_product
 from polymatrix.expression.to import to_constant as internal_to_constant
 from polymatrix.expression.to import to_sympy as internal_to_sympy
 from polymatrix.denserepr.from_ import from_polymatrix
+from polymatrix.polymatrix.init import to_affine_expression
 
 Expression = internal_Expression
 ExpressionState = internal_ExpressionState
 
 init_expression_state = internal_init_expression_state
+make_state = init_expression_state
+
 from_ = internal_from
 v_stack = internal_v_stack
 h_stack = internal_h_stack
@@ -25,4 +28,5 @@ to_sympy_repr = internal_to_sympy
 to_sympy = internal_to_sympy
 to_matrix_repr = from_polymatrix
 to_dense = from_polymatrix
+to_affine = to_affine_expression
 from_names = internal_from_names
diff --git a/polymatrix/polymatrix/init.py b/polymatrix/polymatrix/init.py
index 9a85a0d..f9a3de1 100644
--- a/polymatrix/polymatrix/init.py
+++ b/polymatrix/polymatrix/init.py
@@ -5,8 +5,9 @@ import numpy as np
 
 from typing import TYPE_CHECKING
 
-from polymatrix.polymatrix.impl import BroadcastPolyMatrixImpl, PolyMatrixImpl
+from polymatrix.polymatrix.impl import BroadcastPolyMatrixImpl, PolyMatrixImpl, PolyMatrixAsAffineExpressionImpl
 from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict, MatrixIndex, MonomialIndex, VariableIndex
+from polymatrix.polymatrix.mixins import PolyMatrixAsAffineExpressionMixin
 
 if TYPE_CHECKING:
     from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin, PolyMatrixMixin
@@ -48,3 +49,29 @@ def init_broadcast_poly_matrix(
         data=data,
         shape=shape,
     )
+
+
+def to_affine_expression(p: PolyMatrixMixin) -> PolyMatrixAsAffineExpressionMixin:
+        
+    if isinstance(p, PolyMatrixAsAffineExpressionMixin):
+        return p
+
+    # get all monomials that come up in p, and sort them
+    monomials = sorted(set(m for _, entry in p.entries() for m
+                           in entry.monomials()))
+
+    nrows, ncols = p.shape
+    data = np.zeros((nrows, ncols * len(monomials)))
+    
+    # Generate slices to index sub-matrices
+    slices = {
+        m: tuple(range(i*ncols, (i+1)*ncols))
+        for i, m in enumerate(monomials)
+    }
+
+    for (row, col), poly in p.entries():
+        for monomial, coeff in poly.terms():
+            slice = slices[monomial]
+            data[row, slice[col]] += coeff
+
+    return PolyMatrixAsAffineExpressionImpl(data, p.shape, slices)
diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py
index 25b968c..b33379b 100644
--- a/polymatrix/polymatrix/mixins.py
+++ b/polymatrix/polymatrix/mixins.py
@@ -4,6 +4,7 @@ import abc
 import typing
 import itertools
 import functools
+import math
 import numpy as np
 
 from numpy.typing import NDArray
@@ -220,7 +221,10 @@ class PolyMatrixAsAffineExpressionMixin(
 
     # -- Special methods for affine expression
 
-    def monomials_eval(self, x: tuple) -> dict[MonomialIndex, float]:
+    def monomials_eval(self, x: Sequence[float] | MatrixType, monomials: Iterable[MonomialIndex]) -> dict[MonomialIndex, float]:
+        raise NotImplementedError
+
+    def monomials_eval_all(self, x: Sequence[float] | MatrixType) -> dict[MonomialIndex, float]:
         r"""
         Evaluate all monomials up to degree :math:`d`.
 
@@ -228,6 +232,20 @@ class PolyMatrixAsAffineExpressionMixin(
         numerical values of :math:`\langle x \rangle_d = \{1, x_1, x_2, \ldots,
         x_q, x_1^2, x_1 x_2, \ldots, x_q^2, x_1^3, \ldots, x_q^d\}`.
         """
+
+        # Get involved variables
+        # TODO: this value should be cached
+        variables = sorted(set(v for m in self.slices.keys()
+                               for v in m))
+
+        # TODO: this value should be cached
+        q = len(variables)
+
+        if len(x) != q:
+            # TODO: improve error message, retrieve variable names from state?
+            raise ValueError(f"The value must be a {q} dimensional, however "
+                             f"{x} is {len(x)} dimensional.")
+
         def compute_with_triangles(d, x):
             # To compute all monomials we need to compute all integer points contained
             # in a d-dimensional triangle-like shape. For example given a q-dimensional
@@ -300,7 +318,7 @@ class PolyMatrixAsAffineExpressionMixin(
             return terms
 
         # Get the highest degree we need to compute from the monomial indices
-        # TODO: this value could be cached
+        # TODO: this value should be cached
         d = max(m.degree for m in self.slices.keys())
 
         # Add the constant term
@@ -309,7 +327,6 @@ class PolyMatrixAsAffineExpressionMixin(
         }
 
         # Compute all monomials and convert key to MonomialIndex
-        # TODO: modify compute_with_triangles to directly use MonomialIndex
         for key, val in compute_with_triangles(d, x).items():
             # See comment above for structure of key
             # Count occurrences to convert to MonomialIndex
@@ -320,7 +337,7 @@ class PolyMatrixAsAffineExpressionMixin(
                 else:
                     new_key[e] = 1
 
-            index = MonomialIndex(VariableIndex(k, v) for k, v in new_key.items())
+            index = MonomialIndex(VariableIndex(variables[k].index, v) for k, v in new_key.items())
             monomial_values[index] = val
         
         # Sort the values by monomial index, which have a total order
@@ -346,10 +363,9 @@ class PolyMatrixAsAffineExpressionMixin(
         r"""
         Iterate over the coefficients grouped by degree.
         """
-        monomials = self.slices.keys()
-        groups = itertools.groupby(monomials, lambda m: m.degree)
+        groups = itertools.groupby(self.slices.keys(), lambda m: m.degree)
         for degree, monomials in groups:
-            yield degree, np.vstack(list(self.data[m] for m in monomials))
+            yield degree, np.hstack(list(self.data[:, self.slices[m]] for m in monomials))
 
     def affine_eval(self, x: MatrixType) -> MatrixType:
         r"""
@@ -357,12 +373,47 @@ class PolyMatrixAsAffineExpressionMixin(
         :math:`p(x) = \sum_{\alpha \in \mathcal{E}\langle x \rangle_d} A_\alpha x^\alpha`
         at a given :math:`x`.
         """
-        # Compute all powers of x
-        monomials = np.array(tuple(self.monomials_eval(tuple(x)).items()))
 
+        # Get the number of variables involved
+        # TODO: this value should be cached
+        q = len(set(v.index
+                    for m in self.slices.keys()
+                    for v in m))
+
+        if len(x) != q:
+            # TODO: improve error message, retrieve variable names from state?
+            raise ValueError(f"The value must be a {q} dimensional, however "
+                             f"{x} is {len(x)} dimensional.")
+
+        # Get the highest degree we need to compute from the monomial indices
+        # TODO: this value should be cached
+        d = max(m.degree for m in self.slices.keys())
+
+        # Compute number of monomials that are generated by that
+        # TODO: this value should be cached
+        n = sum(math.comb(k + q - 1, k) for k in range(0, d+1))
+
+        # Heuristic on efficiency, computing all monomials exploits the
+        # structure to reduce the number of computations necessary. But the
+        # efficiency gain is useless if most of the computed terms are not
+        # used.
+        if len(self.slices) > n / 2:
+            # Compute all powers of x up to degree d
+            all_monomials = self.monomials_eval_all(x)
+            monomials = {
+                m: all_monomials[m]
+                for m in self.slices.keys()
+            }
+
+        else:
+            # Compute only the necessary terms
+            monomials = self.monomials_eval(x, self.slices.keys())
+
+        mvalues = np.array(tuple(monomials.values()))
+        
         # Evaluate the affine expression
         _, ncols = self.shape
-        return self.data @ (np.kron(monomials, np.eye(ncols)))
+        return self.data @ (np.kron(mvalues, np.eye(ncols)).T)
 
     def affine_eval_fn(self) -> Callable[[MatrixType], MatrixType]:
         r"""