from_ creates new variables from str

1 files changed, 53 insertions, 37 deletions

diff --git a/polymatrix/__init__.py b/polymatrix/__init__.py
index 23fe913..64f3d63 100644
--- a/polymatrix/__init__.py
+++ b/polymatrix/__init__.py
@@ -7,10 +7,11 @@ import numpy as np
 import scipy.sparse
 import sympy
 
-from polymatrix.expression.expression import Expression
+from polymatrix.expression.mixins.expressionmixin import ExpressionMixin as Expression
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
-from polymatrix.expression.mixins.expressionmixin import ExpressionMixin
+# from polymatrix.expression.mixins.expressionmixin import ExpressionMixin
 from polymatrix.expression.mixins.parametrizeexprmixin import ParametrizeExprMixin
+from polymatrix.expression.mixins.parametrizematrixexprmixin import ParametrizeMatrixExprMixin
 from polymatrix.expressionstate.expressionstate import ExpressionState
 from polymatrix.expression.init.initblockdiagexpr import init_block_diag_expr
 from polymatrix.expression.init.initexpression import init_expression
@@ -19,7 +20,7 @@ from polymatrix.expression.init.initfromsympyexpr import init_from_sympy_expr
 from polymatrix.expression.init.initfromtermsexpr import init_from_terms_expr
 from polymatrix.expression.init.initvstackexpr import init_v_stack_expr
 from polymatrix.polymatrix.polymatrix import PolyMatrix
-from polymatrix.expression.utils.getvariableindices import get_variable_indices, get_variable_indices_from_variable
+from polymatrix.expression.utils.getvariableindices import get_variable_indices_from_variable
 from polymatrix.statemonad.init.initstatemonad import init_state_monad
 from polymatrix.statemonad.mixins.statemonadmixin import StateMonadMixin
 from polymatrix.expression.utils.monomialtoindex import monomial_to_index
@@ -52,8 +53,11 @@ def from_polymatrix(
     )
 
 def from_(
-    data: tuple[tuple[float]],
+    data: str | tuple[tuple[float]],
 ):
+    if isinstance(data, str):
+        return from_(1).parametrize(data)
+
     return from_sympy(data)
 
 def v_stack(
@@ -308,7 +312,7 @@ def shape(
 
 @dataclasses.dataclass
 class MatrixBuffer:
-    data: dict[int, ...]
+    data: dict[int, np.ndarray]
     n_row: int
     n_param: int
 
@@ -394,6 +398,9 @@ class MatrixRepresentations:
                 if isinstance(x, tuple) or isinstance(x, list):
                     x = np.array(x).reshape(-1, 1)
 
+                elif x.shape[0] == 1:
+                    x = x.reshape(-1, 1)
+
                 def acc_x_powers(acc, _):
                     next = (acc @ x.T).reshape(-1, 1)
                     return next
@@ -433,14 +440,14 @@ class MatrixRepresentations:
 
 def to_matrix_repr(
     expressions: Expression | tuple[Expression],
-    variables: Expression,
+    variables: Expression = None,
     sorted: bool = None,
 ) -> StateMonadMixin[ExpressionState, tuple[tuple[tuple[np.ndarray, ...], ...], tuple[int, ...]]]:
 
     if isinstance(expressions, Expression):
         expressions = (expressions,)
 
-    assert isinstance(variables, ExpressionBaseMixin), f'{variables=}'
+    assert isinstance(variables, ExpressionBaseMixin) or variables is None, f'{variables=}'
 
     def func(state: ExpressionState):
 
@@ -453,32 +460,36 @@ def to_matrix_repr(
             
             return state, underlying_list + (underlying,)
 
-        *_, (state, underlying_list) = tuple(itertools.accumulate(
+        *_, (state, polymatrix_list) = tuple(itertools.accumulate(
             expressions, 
             acc_underlying_application, 
             initial=(state, tuple()),
         ))
 
-        state, variable_index = get_variable_indices_from_variable(state, variables)
-
-        if sorted:
-            tagged_variable_index = tuple((offset, state.get_name_from_offset(offset)) for offset in variable_index)
-
-            sorted_variable_index = tuple(v[0] for v in sorted(tagged_variable_index, key=lambda v: v[1]))
+        if variables is None:
+            sorted_variable_index = tuple()
             
         else:
-            sorted_variable_index = variable_index
+            state, variable_index = get_variable_indices_from_variable(state, variables)
 
-        assert len(sorted_variable_index) == len(set(sorted_variable_index)), f'{sorted_variable_index=} contains repeated variables'
+            if sorted:
+                tagged_variable_index = tuple((offset, state.get_name_from_offset(offset)) for offset in variable_index)
+
+                sorted_variable_index = tuple(v[0] for v in sorted(tagged_variable_index, key=lambda v: v[1]))
+                
+            else:
+                sorted_variable_index = variable_index
+
+            assert len(sorted_variable_index) == len(set(sorted_variable_index)), f'{sorted_variable_index=} contains repeated variables. Make sure you give a unique name for each variables.'
         
         variable_index_map = {old: new for new, old in enumerate(sorted_variable_index)}
 
         n_param = len(sorted_variable_index)
 
-        def gen_underlying_matrices():
-            for underlying in underlying_list:
+        def gen_numpy_matrices():
+            for polymatrix in polymatrix_list:
 
-                n_row = underlying.shape[0]
+                n_row = polymatrix.shape[0]
 
                 buffer = MatrixBuffer(
                     data={},
@@ -487,34 +498,39 @@ def to_matrix_repr(
                 )
 
                 for row in range(n_row):
-                    underlying_terms = underlying.get_poly(row, 0)
-                    if underlying_terms is None:
+                    polymatrix_terms = polymatrix.get_poly(row, 0)
+                    
+                    if polymatrix_terms is None:
                         continue
+
+                    if len(polymatrix_terms) == 0:
+                        buffer.add(row, 0, 0, 0)
                     
-                    for monomial, value in underlying_terms.items():
+                    else:
+                        for monomial, value in polymatrix_terms.items():
 
-                        def gen_new_monomial():
-                            for var, count in monomial:
-                                try:
-                                    index = variable_index_map[var]
-                                except KeyError:
-                                    raise KeyError(f'{var=} ({state.get_key_from_offset(var)}) is incompatible with {variable_index_map=}')
+                            def gen_new_monomial():
+                                for var, count in monomial:
+                                    try:
+                                        index = variable_index_map[var]
+                                    except KeyError:
+                                        raise KeyError(f'{var=} ({state.get_key_from_offset(var)}) is incompatible with {variable_index_map=}')
 
-                                for _ in range(count):
-                                    yield index
+                                    for _ in range(count):
+                                        yield index
 
-                        new_monomial = tuple(gen_new_monomial())
+                            new_monomial = tuple(gen_new_monomial())
 
-                        cols = monomial_to_index(n_param, new_monomial)
+                            cols = monomial_to_index(n_param, new_monomial)
 
-                        col_value = value / len(cols)
+                            col_value = value / len(cols)
 
-                        for col in cols:
-                            buffer.add(row, col, sum(count for _, count in monomial), col_value)
+                            for col in cols:
+                                buffer.add(row, col, sum(count for _, count in monomial), col_value)
                 
                 yield buffer
 
-        underlying_matrices = tuple(gen_underlying_matrices())
+        underlying_matrices = tuple(gen_numpy_matrices())
 
         def gen_auxillary_equations():
             for key, monomial_terms in state.auxillary_equations.items():
@@ -641,7 +657,7 @@ def to_sympy_repr(
 
                     if isinstance(variable, sympy.core.symbol.Symbol):
                         variable_name = variable.name
-                    elif isinstance(variable, ParametrizeExprMixin):
+                    elif isinstance(variable, (ParametrizeExprMixin, ParametrizeMatrixExprMixin)):
                         variable_name = variable.name
                     elif isinstance(variable, str):
                         variable_name = variable