Split FromTuple into FromNumbers, FromNumpy and FromSympy

7 files changed, 232 insertions, 148 deletions

diff --git a/polymatrix/expression/impl.py b/polymatrix/expression/impl.py
index c921caa..d4f5429 100644
--- a/polymatrix/expression/impl.py
+++ b/polymatrix/expression/impl.py
@@ -1,4 +1,6 @@
-import typing
+import numpy.typing
+import sympy
+
 import dataclassabc
 from polymatrix.expression.mixins.integrateexprmixin import IntegrateExprMixin
 
@@ -22,7 +24,9 @@ from polymatrix.expression.mixins.filterexprmixin import FilterExprMixin
 from polymatrix.expression.mixins.fromsymmetricmatrixexprmixin import (
     FromSymmetricMatrixExprMixin,
 )
-from polymatrix.expression.mixins.fromtupleexprmixin import FromTupleExprMixin
+from polymatrix.expression.mixins.fromnumbersexprmixin import FromNumbersExprMixin
+from polymatrix.expression.mixins.fromnumpyexprmixin import FromNumpyExprMixin
+from polymatrix.expression.mixins.fromsympyexprmixin import FromSympyExprMixin
 from polymatrix.expression.mixins.fromtermsexprmixin import (
     FromPolynomialDataExprMixin,
     PolynomialMatrixTupledData,
@@ -162,18 +166,23 @@ class FilterExprImpl(FilterExprMixin):
 
 
 @dataclassabc.dataclassabc(frozen=True)
-class FromSymmetricMatrixExprImpl(FromSymmetricMatrixExprMixin):
-    underlying: ExpressionBaseMixin
+class FromNumbersExprImpl(FromNumbersExprMixin):
+    data: tuple[tuple[int | float]]
 
 
 @dataclassabc.dataclassabc(frozen=True)
-class FromTupleExprImpl(FromTupleExprMixin):
-    data: tuple[tuple[float]]
-    stack: tuple[FrameSummary]
+class FromNumpyExprImpl(FromNumpyExprMixin):
+    data: numpy.typing.NDArray
 
-    # implement custom __repr__ method that returns a representation without the stack
-    def __repr__(self):
-        return f"{self.__class__.__name__}(data={self.data})"
+
+@dataclassabc.dataclassabc(frozen=True)
+class FromSympyExprImpl(FromSympyExprMixin):
+    data: sympy.Expr | sympy.Matrix | tuple[tuple[sympy.Expr]]
+
+
+@dataclassabc.dataclassabc(frozen=True)
+class FromSymmetricMatrixExprImpl(FromSymmetricMatrixExprMixin):
+    underlying: ExpressionBaseMixin
 
 
 @dataclassabc.dataclassabc(frozen=True)
diff --git a/polymatrix/expression/init.py b/polymatrix/expression/init.py
index 7d708f6..5acb40d 100644
--- a/polymatrix/expression/init.py
+++ b/polymatrix/expression/init.py
@@ -1,5 +1,6 @@
 import typing
 import numpy as np
+import numpy.typing as npt
 import sympy
 
 from polymatrix.expression.typing import FromDataTypes
@@ -13,7 +14,7 @@ from polymatrix.utils.getstacklines import FrameSummary
 from polymatrix.utils.getstacklines import get_stack_lines
 from polymatrix.expression.utils.formatsubstitutions import format_substitutions
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
-from polymatrix.expression.impl import FromTupleExprImpl, AdditionExprImpl
+from polymatrix.expression.impl import AdditionExprImpl
 
 
 def init_addition_expr(
@@ -156,6 +157,17 @@ def init_from_symmetric_matrix_expr(
     )
 
 
+def init_from_numbers_expr(data: tuple[tuple[int | float]]):
+    return polymatrix.expression.impl.FromNumbersExprImpl(data=data)
+
+
+def init_from_numpy_expr(data: npt.NDArray):
+    return polymatrix.expression.impl.FromNumpyExprImpl(data=data)
+
+
+def init_from_sympy_expr(data: sympy.Expr | sympy.Matrix | tuple[tuple[sympy.Expr]]):
+    return polymatrix.expression.impl.FromSympyExprImpl(data=data)
+
 
 # NP: this function should be split up into smaller functions, one for each "from" type
 # NP: and each "from" should be documented, explaining how it is interpreted.
@@ -173,44 +185,32 @@ def init_from_expr_or_none(
         return data.flat_map(lambda inner_data: init_from_expr_or_none(inner_data))
 
     elif isinstance(data, np.ndarray):
-        assert len(data.shape) <= 2
+        return init_from_numpy_expr(data)
 
-        def gen_elements():
-            for row in data:
-                if isinstance(row, np.ndarray):
-                    yield tuple(row)
-                else:
-                    yield (row,)
+    elif isinstance(data, sympy.Expr | sympy.Matrix):
+        return init_from_sympy_expr(data)
 
-        data = tuple(gen_elements())
+    if isinstance(data, tuple):
+        if len(data) < 1:
+            return None
 
-    elif isinstance(data, sympy.Matrix):
-        data = tuple(tuple(i for i in data.row(row)) for row in range(data.rows))
+        if isinstance(data[0], tuple):
+            if len(data[0]) < 1:
+                return None
 
-    elif isinstance(data, sympy.Expr):
-        data = ((sympy.expand(data),),)
+            if isinstance(data[0][0], sympy.Expr):
+                return init_from_sympy_expr(data)
 
-    elif isinstance(data, tuple):
-        if isinstance(data[0], tuple):
-            n_col = len(data[0])
-            assert all(len(col) == n_col for col in data)
+            elif isinstance(data[0][0], int | float):
+                return init_from_numbers_expr(data)
 
-        else:
-            data = tuple((e,) for e in data)
+        elif isinstance(data[0], sympy.Expr):
+            return init_from_sympy_expr((data,))
 
     elif isinstance(data, ExpressionBaseMixin):
         return data
 
-    elif isinstance(data, (float, int, np.number)):
-        data = ((data,),)
-
-    else:
-        return None
-
-    return FromTupleExprImpl(
-        data=data,
-        stack=get_stack_lines(),
-    )
+    return None
 
 
 def init_from_expr(data: FromDataTypes):
diff --git a/polymatrix/expression/mixins/fromnumbersexprmixin.py b/polymatrix/expression/mixins/fromnumbersexprmixin.py
new file mode 100644
index 0000000..c35f955
--- /dev/null
+++ b/polymatrix/expression/mixins/fromnumbersexprmixin.py
@@ -0,0 +1,39 @@
+from abc import abstractmethod
+from typing_extensions import override
+
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.expressionstate.mixins import ExpressionStateMixin
+
+from polymatrix.polymatrix.mixins import PolyMatrixMixin
+from polymatrix.polymatrix.init import init_poly_matrix
+from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict, MonomialIndex
+
+
+class FromNumbersExprMixin(ExpressionBaseMixin):
+    """
+    Make an expression from a tuple of tuples of numbers (constant). The tuple
+    of tuples is interpreted as a matrix stored with row major ordering.
+
+    ..code:: py
+        m = polymatrix.from_numbers(((0, 1), (1, 0))
+    """
+
+    @property
+    @abstractmethod
+    def data(self) -> tuple[tuple[int | float]]:
+        """ The matrix of numbers in row major order. """
+        # TODO: allow numbers tuple, interpreted as row vector and just
+        # numbers, interpreted as scalars
+
+    @override
+    def apply(self, state: ExpressionStateMixin) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
+        p = PolyMatrixDict.empty()
+
+        for r, row in enumerate(self.data):
+            for c, entry in enumerate(row):
+                p[r, c] = PolyDict({
+                    MonomialIndex.constant(): entry
+                })
+
+        shape = (r + 1, c + 1)
+        return state, init_poly_matrix(p, shape)
diff --git a/polymatrix/expression/mixins/fromnumpyexprmixin.py b/polymatrix/expression/mixins/fromnumpyexprmixin.py
new file mode 100644
index 0000000..92132f8
--- /dev/null
+++ b/polymatrix/expression/mixins/fromnumpyexprmixin.py
@@ -0,0 +1,49 @@
+import math
+
+from abc import abstractmethod
+from typing_extensions import override, cast
+from numpy.typing import NDArray
+
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.expressionstate.mixins import ExpressionStateMixin
+
+from polymatrix.polymatrix.mixins import PolyMatrixMixin
+from polymatrix.polymatrix.init import init_poly_matrix 
+from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict, MonomialIndex
+
+
+class FromNumpyExprMixin(ExpressionBaseMixin):
+    """
+    Make a (constant) expression from a numpy array.
+
+    ..code:: py
+        identity = polymatrix.from_numpy(np.eye(3))
+    """
+
+    @property
+    @abstractmethod
+    def data(self) -> NDArray:
+        """ The Numpy array. """
+
+    @override
+    def apply(self, state: ExpressionStateMixin) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
+        p = PolyMatrixDict.empty()
+
+        if len(self.data.shape) > 2:
+            raise ValueError("Cannot construct expression from numpy array with "
+                             f"shape {self.data.shape}, only vectors and matrices are allowed")
+
+        # Case when it is a (n,) array
+        data = self.data
+        if len(data.shape) == 1:
+            data = data.reshape(self.data.shape[0], 1)
+
+        nrows, ncols = data.shape
+        for r in range(nrows):
+            for c in range(ncols):
+                if not math.isclose(self.data[r, c], 0):
+                    p[r, c] = PolyDict({
+                        MonomialIndex.constant(): data[r, c]
+                    })
+
+        return state, init_poly_matrix(p, (nrows, ncols))
diff --git a/polymatrix/expression/mixins/fromsympyexprmixin.py b/polymatrix/expression/mixins/fromsympyexprmixin.py
new file mode 100644
index 0000000..a0003fb
--- /dev/null
+++ b/polymatrix/expression/mixins/fromsympyexprmixin.py
@@ -0,0 +1,93 @@
+import sympy
+import math
+
+from abc import abstractmethod
+from typing_extensions import override, cast
+
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.expressionstate.mixins import ExpressionStateMixin
+
+from polymatrix.polymatrix.abc import PolyMatrix
+from polymatrix.polymatrix.init import init_poly_matrix 
+from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict, MonomialIndex, VariableIndex
+
+
+class FromSympyExprMixin(ExpressionBaseMixin):
+    """
+    Make an expression from a sympy expression. The sympy expression can be
+    scalar or matrix (`sympy.Matrix`).
+
+    ..code:: py
+        x, y = sympy.symbols('x, y')
+
+        # scalar bivariate polynomial
+        p_sym = x**2 + y + 2
+        p = polymatrix.from_sympy(p_sym)
+
+        # 2D vector field
+        q_sym = sympy.Matrix(((x + y), (x**2 + 1)))
+        q = polymatrix.from_sympy(q_sym)
+    """
+
+    @property
+    @abstractmethod
+    def data(self) -> sympy.Expr | tuple[tuple[sympy.Expr]] | sympy.Matrix:
+        """ The sympy objects. """
+
+    @override
+    def apply(self, state: ExpressionStateMixin) -> tuple[ExpressionStateMixin, PolyMatrix]:
+        # Unpack if it is a sympy matrix
+        data = sympy.expand(self.data)
+        if isinstance(data, sympy.Matrix):
+            data = ((entry for entry in data.row(r))
+                    for r in range(data.rows))
+
+        # Pack if is raw (scalar) expression
+        if isinstance(data, sympy.Expr):
+            data = ((data,),)
+
+        # Convert to polymatrix
+        polymatrix = PolyMatrixDict.empty()
+        for r, row in enumerate(data):
+            for c, entry in enumerate(row):
+                # Special case of constant polynomial
+                if isinstance(entry, int | float):
+                    if not math.isclose(entry, 0):
+                        polymatrix[r, c] = PolyDict({
+                            MonomialIndex.constant(): entry,
+                        })
+                    continue
+
+                try:
+                    sympy_poly = sympy.poly(entry)
+
+                except sympy.polys.polyerrors.BasePolynomialError as e:
+                    raise ValueError(f"Cannot convert sympy expression {entry} "
+                    "into a polynomial, are you sure it is a polynomial?") from e
+
+                # Convert sympy variables to our variables
+                from polymatrix.expression.init import init_variable_expr
+                sympy_to_var = {
+                    sympy_idx: init_variable_expr(var.name)
+                    for sympy_idx, var in enumerate(sympy_poly.gens)
+                }
+
+                poly = PolyDict.empty()
+                for coeff, monom in zip(sympy_poly.coeffs(), sympy_poly.monoms()):
+                    # sympy monomial is stored as multi-index, eg. for a
+                    # multivariate polynomial with three variables (x, y, z)
+                    # the index is x*y**2 = (1, 2, 0)
+                    m: list[VariableIndex] = []
+                    for i, exponent in enumerate(monom):
+                        var = sympy_to_var[i]
+                        state, idx = state.index(var)
+                        # idx.start because var is a scalar
+                        m.append(VariableIndex(idx.start, exponent))
+
+                    poly[m] = coeff
+                polymatrix[r, c] = poly
+
+        shape = (r + 1, c + 1)
+        return state, init_poly_matrix(polymatrix, shape)
+
+
diff --git a/polymatrix/expression/mixins/fromtupleexprmixin.py b/polymatrix/expression/mixins/fromtupleexprmixin.py
deleted file mode 100644
index 35c6e5d..0000000
--- a/polymatrix/expression/mixins/fromtupleexprmixin.py
+++ /dev/null
@@ -1,106 +0,0 @@
-import abc
-import math
-import sympy
-import numpy as np
-
-from polymatrix.expression.typing import FromDataTypes
-
-from polymatrix.utils.getstacklines import FrameSummary
-from polymatrix.utils.tooperatorexception import to_operator_exception
-from polymatrix.polymatrix.init import init_poly_matrix
-from polymatrix.polymatrix.mixins import PolyMatrixMixin
-from polymatrix.expressionstate.mixins import ExpressionStateMixin
-from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
-
-
-class FromTupleExprMixin(ExpressionBaseMixin):
-    @property
-    @abc.abstractmethod
-    def data(self) -> tuple[tuple[FromDataTypes]]: ...
-
-    @property
-    @abc.abstractmethod
-    def stack(self) -> tuple[FrameSummary]: ...
-
-    # overwrites the abstract method of `ExpressionBaseMixin`
-    # FIXME: break up this mess
-    def apply(
-        self,
-        state: ExpressionStateMixin,
-    ) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
-        polynomials = {}
-
-        for poly_row, col_data in enumerate(self.data):
-            for poly_col, poly_data in enumerate(col_data):
-                if isinstance(poly_data, (bool, np.bool_)):
-                    poly_data = int(poly_data)
-
-                if isinstance(poly_data, (int, float, np.number)):
-                    if math.isclose(poly_data, 0):
-                        polynomial = {}
-                    else:
-                        polynomial = {tuple(): poly_data}
-
-                elif isinstance(poly_data, sympy.Expr):
-                    try:
-                        poly = sympy.poly(poly_data)
-
-                    except sympy.polys.polyerrors.GeneratorsNeeded:
-                        if not math.isclose(poly_data, 0):
-                            polynomials[poly_row, poly_col] = {tuple(): poly_data}
-                        continue
-
-                    except ValueError:
-                        raise ValueError(f"{poly_data=}")
-
-                    for symbol in poly.gens:
-                        state = state.register(key=symbol, n_param=1)
-                        # print(f'{symbol}: {state.n_param}')
-
-                    polynomial = {}
-
-                    # a5 x1 x3**2 -> c=a5, m_cnt=(1, 0, 2)
-                    for value, monomial_count in zip(poly.coeffs(), poly.monoms()):
-                        if math.isclose(value, 0):
-                            continue
-
-                        # m_cnt=(1, 0, 2) -> m=(0, 2, 2)
-                        def gen_monomial():
-                            for rel_idx, p in enumerate(monomial_count):
-                                if 0 < p:
-                                    idx, _ = state.offset_dict[poly.gens[rel_idx]]
-                                    yield idx, p
-
-                        monomial = tuple(gen_monomial())
-
-                        polynomial[monomial] = value
-
-                elif isinstance(poly_data, ExpressionBaseMixin):
-                    state, instance = poly_data.apply(state)
-
-                    if not (instance.shape == (1, 1)):
-                        raise AssertionError(
-                            to_operator_exception(
-                                message=f"{instance.shape=} is not (1, 1)",
-                                stack=self.stack,
-                            )
-                        )
-
-                    polynomial = instance.get_poly(0, 0)
-
-                else:
-                    raise AssertionError(
-                        to_operator_exception(
-                            message=f"unknown data type {type(poly_data)=}",
-                            stack=self.stack,
-                        )
-                    )
-
-                polynomials[poly_row, poly_col] = polynomial
-
-        poly_matrix = init_poly_matrix(
-            data=polynomials,
-            shape=(len(self.data), len(self.data[0])),
-        )
-
-        return state, poly_matrix
diff --git a/polymatrix/expression/mixins/variablemixin.py b/polymatrix/expression/mixins/variablemixin.py
index ea1d222..c2b0bc1 100644
--- a/polymatrix/expression/mixins/variablemixin.py
+++ b/polymatrix/expression/mixins/variablemixin.py
@@ -2,12 +2,12 @@ from abc import abstractmethod
 from itertools import product
 from typing_extensions import override
 
-from polymatrix.polymatrix.abc import PolyMatrix
+from polymatrix.polymatrix.mixins import PolyMatrixMixin
 from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict, MonomialIndex, VariableIndex
 from polymatrix.polymatrix.init import init_poly_matrix
 
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
-from polymatrix.expressionstate.abc import ExpressionState
+from polymatrix.expressionstate.mixins import ExpressionStateMixin
 
 
 class VariableMixin(ExpressionBaseMixin):
@@ -25,7 +25,7 @@ class VariableMixin(ExpressionBaseMixin):
         """ Shape of the variable. """
 
     @override
-    def apply(self, state: ExpressionState) -> tuple[ExpressionState, PolyMatrix]:
+    def apply(self, state: ExpressionStateMixin) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
         """ See :py:meth:`ExpressionBaseMixin.apply`. """
         
         state, indices = state.index(self)