add solver for roots of a set of polynomial equations

4 files changed, 769 insertions, 306 deletions

diff --git a/polymatrix/polysolver.py b/polymatrix/polysolver.py
index f611260..24c7095 100644
--- a/polymatrix/polysolver.py
+++ b/polymatrix/polysolver.py
@@ -1,41 +1,351 @@
 import numpy as np
+import sympy
 import itertools
+import scipy.linalg
+import scipy.sparse.linalg
+import collections
+from scipy.special import binom
+import more_itertools
+import time
 
-def quadratic_poly_solver(b, c, m):
-    n = b.shape[0]
-    nn = n*n
 
-    b_inv = np.linalg.inv(b)
+def substitude_x_add_a(data, a):    
 
-    p = b_inv @ np.ones((n, 1))
+    def kron_eye_and_a(n, data, a, indices):
+        permutations = more_itertools.distinct_permutations(indices)
+        # permutations = set(itertools.permutations(indices))
+        
+        if not scipy.sparse.issparse(data):
+            a_mat = np.reshape(a, (-1, 1))
+            ops = (np.eye(n), a_mat)
+
+            def acc_kron(perm):
+                *_, last = itertools.accumulate((ops[idx] for idx in perm), lambda acc, v: np.kron(acc, v))
+                return last
+
+            return sum(data @ acc_kron(perm) for perm in permutations)
+        
+        else:
+            # csr_data = data
+            csr_data = data.tocsr()
+            
+            n_eye = sum(1 for idx in indices if idx == 0)
+            n_col = n**n_eye
+            n_index = n_col * n**(len(indices) - n_eye)
+                    
+            def gen_array_per_permuatation():
+                
+                def gen_coord(perm):
+                    n_row = csr_data.shape[0]
+
+                    def acc_col_idx_and_value(acc, v):
+                        relindex, relrow, col, val = acc
 
-    def square_matrix(p):
-        return (p @ p.T).reshape(nn, 1)
+                        n_relrow = relrow / n
+                        grp_index = int(relindex / n_relrow)
+                        n_relindex = int(relindex - grp_index * n_relrow)
+                        
+                        if v == 1:
+                            n_col = col
+                            n_val = val * a[grp_index]
 
+                        else:
+                            n_col = n * col + grp_index
+                            n_val = val
+                            
+                        return n_relindex, n_relrow, n_col, n_val
 
-    def merge_matrix(p1, p2):
-        return (np.kron(np.eye(n), p1) + np.kron(p1, np.eye(n))) @ p2
+                    for row in range(n_row):
+
+                        pt = slice(csr_data.indptr[row], csr_data.indptr[row+1]) 
+
+                        for inner_idx, array_val in zip(csr_data.indices[pt], csr_data.data[pt]):
+
+                            *_, last = itertools.accumulate(
+                                perm, 
+                                acc_col_idx_and_value, 
+                                initial=(inner_idx, n_index, 0, array_val),
+                            )
+                            _, _, col, val = last
+
+                            yield col, val
+                
+                for perm in permutations:
 
+                    col, data = list(zip(*gen_coord(perm)))
+
+                    yield scipy.sparse.csr_array((data, col, csr_data.indptr), shape=(csr_data.shape[0], n_col))
+
+            return sum(gen_array_per_permuatation())
+
+    n = data[1].shape[0]
+
+    subs_data = {key: val.copy() for key, val in data.items()}
+
+    for degree, d_data in data.items():
+        for i_degree in range(degree):
+            indices = i_degree*(0,) + (degree - i_degree)*(1,)
+            update = kron_eye_and_a(n, d_data, a, indices)
+            subs_data[i_degree] += update
+    
+    first = -1/np.squeeze(subs_data[0])
+    def gen_subs_data():
+        for key, d_data in subs_data.items():
+            if scipy.sparse.issparse(d_data):
+                
+                csr_data = d_data.tocsr()
+                for row, scale in enumerate(first):
+                    csr_data.data[csr_data.indptr[row] : csr_data.indptr[row + 1]] *= scale
+                yield key, csr_data
+                
+            else:
+                d_data *= first[:, None]
+                yield key, d_data
+    return dict(gen_subs_data())
+
+
+def solve_poly_system(data, m):
+    if isinstance(data[1], np.ndarray):
+        b_num_inv = np.linalg.inv(data[1])
+    else:
+        b_num_inv = scipy.sparse.linalg.inv(data[1])
+    
+    n = b_num_inv.shape[0]
+    p = b_num_inv @ np.ones((n, 1))
+    
+    assert (data[0] + np.ones((n, 1)) < 0.01).all(), f'{data[0]=}, {data[0] + np.ones((n, 1))=}'
 
     def func(acc, _):
+        """
+        acc: a nxk matrix whose ith column represent the ith parameter of the power series for each of the n variables
+        """
+
         p = acc
-        
         k = p.shape[1]
         
-        uneven = k % 2
-        half = int((k - uneven) / 2)
+        def gen_indices(n, d):
+            """
+            n=2,d=1 = [(0,2), (1,1)]
+            n=3,d=1 = [(0,3), (1,2)]
+            n=2,d=2 = [(0,0,2), (0,1,1)]
+            """
+    
+            if d==0:
+                yield (n,)
+
+            else:
+                for idx in range(int(n/(d+1))+1):
+                    yield from ((idx,)+tuple(i+idx*(d-1) for i in indices) for indices in gen_indices(
+                        n - (2*d-1)*idx, 
+                        d-1,
+                    ))
+
+        def acc_kron(l):
+            *_, last = itertools.accumulate(l, lambda acc, v: np.kron(acc, v))
+            return last
         
-        def gen_pvector():
-            if uneven:
-                yield square_matrix(p[:, half:half+1])
-                
-            for i in range(half):
-                yield merge_matrix(p[:, i:i+1], p[:, k-i-1:k-i])
+        def gen_p():
+            for degree, d_data in data.items():
+                if 1 < degree and degree-1 <= k:
+                    indices_list = list(gen_indices(k-degree+1, degree-1))
+                    permutations = (perm for indices in indices_list for perm in more_itertools.distinct_permutations(indices))
+
+                    if not scipy.sparse.issparse(data):
+
+                        def acc_kron(perm):
+                            *_, last = itertools.accumulate(((p[:,idx:idx+1] for idx in perm)), lambda acc, v: np.kron(acc, v))
+                            return last
+
+                        yield from (d_data @ acc_kron(perm) for perm in permutations)
+
+                    else:
+                        csr_data = d_data.tocsr()
+
+                        n_eye = sum(1 for idx in indices if idx == 0)
+                        n_col = n**n_eye
+                        n_index = n_col * n**(len(indices) - n_eye)
+
+                        def gen_array_per_permuatation():
+
+                            def gen_coord(perm):
+                                n_row = csr_data.shape[0]
+
+                                def acc_col_idx_and_value(acc, v):
+                                    relindex, relrow, val = acc
+
+                                    n_relrow = relrow / n
+                                    grp_index = int(relindex / n_relrow)
+                                    n_relindex = int(relindex - grp_index * n_relrow)
+                                    n_val = val * p[grp_index, v:v+1]
+
+                                    return n_relindex, n_relrow, n_val
+
+                                for row in range(n_row):
+
+                                    pt = slice(csr_data.indptr[row], csr_data.indptr[row+1]) 
+
+                                    def gen_val_per_row():
+                                        for inner_idx, array_val in zip(csr_data.indices[pt], csr_data.data[pt]):
+
+                                            *_, last = itertools.accumulate(
+                                                perm, 
+                                                acc_col_idx_and_value, 
+                                                initial=(inner_idx, n_index, array_val),
+                                            )
+                                            _, _, val = last
+
+                                            yield val
+                                    
+                                    yield sum(gen_val_per_row())
+
+                            for perm in permutations:
+
+                                array_values = list(gen_coord(perm))
+
+                                yield scipy.sparse.csr_array((array_values, np.zeros(len(array_values)), csr_data.indptr), shape=(csr_data.shape[0], n_col))
+
+        return np.concatenate((p, -b_num_inv @ sum(gen_p())), axis=1)
+
+    *_, sol_subs = itertools.accumulate(range(m-1), func, initial=p)
+
+    # return np.asarray(sol_subs)
+    return sum(np.asarray(sol_subs).T)
+
+
+def inner_smart_solve(data, irange=None, idegree=None, a_init=None):
+    
+    if irange is None:
+        irange = np.geomspace(20, 1.1, 8)
+        
+    if idegree is None:
+        idegree = 4
+    
+    if a_init is None:
+        a_init = np.zeros(data[0].shape[0])
+    
+    def acc_func(acc, v):
+        a, prev_err = acc
+        scale = v
+
+        # print(f'iter: {index}')
+
+        # print('substitude...')
+        subs_data = substitude_x_add_a(data, a)
+
+        # print('modify scaling...')
+        def gen_data():
+            for degree, d_data in subs_data.items():
+                if 1 == degree:
+                    yield degree, d_data * scale
+                else:
+                    yield degree, d_data
+        scaled_data = dict(gen_data())
+
+        # print('solve...')
+        try:
+            sol = solve_poly_system(scaled_data, idegree)
+        except:
+            # return a, prev_err
+            raise
+
+        err = np.max(np.abs(eval_solution(scaled_data, sol)))
+
+        a = sol + a
+
+        return a, err
+    
+    a, err = list(zip(*itertools.accumulate(
+        irange,
+        acc_func, 
+        initial=(a_init, 0),
+    )))
+    
+    return a, err
+
+
+def outer_smart_solve(data, n_iter=10, a_max=1.0, irange=None, idegree=None):
+    def gen_a_err():
+        n = data[1].shape[0]
+
+        for idx_iter in range(n_iter):
+            a_init = a_max * (2 * np.random.rand(n) - 1)
+
+            a, err = inner_smart_solve(data, irange=irange, idegree=idegree, a_init=a_init)
+
+            subs_data = substitude_x_add_a(data, a[-1])
+            sol = solve_poly_system(subs_data, 6)
+            error_index = np.max(np.abs(eval_solution(subs_data, sol))) + max(a[-1])
+            # error_index = np.max(np.abs(eval_solution(subs_data, sol)))
+
+            # print(f'{error_index=}')
+
+            yield error_index, a, err
+
+            # if error_index < 1.0:
+            #     break
                 
-        p_vector = np.add.reduce(list(gen_pvector()))
-        p_new = b_inv @ c @ p_vector
-        return np.concatenate((p, p_new), axis=1)
+    _, a, err = min(gen_a_err())
+    
+    return a, err
+
+
+def eval_solution(data, x=None):
+    # if x is None:
+    #     x = sum(sol.T)
         
-    *_, sol_subs = itertools.accumulate(range(m-1), func, initial=p)
+    x = np.reshape(x, (-1, 1))
+    
+    def gen_sum():
+        n = data[1].shape[0]
+        
+        for degree, d_data in data.items():
+            if 0 < degree:
+                # *_, last = itertools.accumulate(degree*(x,), lambda acc, v: np.kron(acc, v))
+                # yield d_data @ last
+                
+                if not scipy.sparse.issparse(d_data):
+                    *_, last = itertools.accumulate(degree*(x,), lambda acc, v: np.kron(acc, v))
+                    yield d_data @ last
+
+                else:
+                    csr_data = d_data.tocsr()
+
+                    def gen_coord():
+                        n_row = csr_data.shape[0]
+                        n_index = n**degree
+
+                        def acc_col_idx_and_value(acc, v):
+                            relindex, relrow, val = acc
 
-    return np.sum(-sol_subs, axis=1, keepdims=True)
+                            n_relrow = relrow / n
+                            grp_index = int(relindex / n_relrow)
+                            n_relindex = int(relindex - grp_index * n_relrow)
+                            
+                            n_val = val * x[grp_index, 0]
+
+                            return n_relindex, n_relrow, n_val
+
+                        for row in range(n_row):
+
+                            pt = slice(csr_data.indptr[row], csr_data.indptr[row+1]) 
+
+                            for inner_idx, array_val in zip(csr_data.indices[pt], csr_data.data[pt]):
+
+                                *_, last = itertools.accumulate(
+                                    range(degree), 
+                                    acc_col_idx_and_value, 
+                                    initial=(inner_idx, n_index, array_val),
+                                )
+                                _, _, val = last
+
+                                yield val
+
+                    array_values = list(gen_coord())
+
+                    yield scipy.sparse.csr_array((array_values, np.zeros(len(array_values)), csr_data.indptr), shape=(csr_data.shape[0], 1))
+                
+            else:
+                yield d_data
+                
+    val = np.ravel(sum(gen_sum()).T)
+    return val
diff --git a/polymatrix/polystruct.py b/polymatrix/polystruct.py
index 202f6cc..ad6cb74 100644
--- a/polymatrix/polystruct.py
+++ b/polymatrix/polystruct.py
@@ -1,15 +1,15 @@
 import abc
 import collections
-from pickletools import int4
 import typing
 import numpy as np
 import dataclasses
 import dataclass_abc
-from scipy.special import binom
+import scipy.sparse
 import itertools
 import functools
+import more_itertools
 
-from polymatrix.utils import variable_to_index
+import polymatrix.utils
 
 DegreeType = int
 # CoordType = tuple[int, int]
@@ -42,31 +42,25 @@ class PolyMatrixMixin(abc.ABC):
 
     @property
     @abc.abstractmethod
-    def is_vector(self) -> bool:
+    def shape(self) -> tuple[int, int]:
         ...
 
-
 class EqualityConstraintMixin(abc.ABC):
     @property
     @abc.abstractmethod
     def terms(self) -> dict[int, dict[tuple[int, tuple, int], float]]:
         ...
 
-    # @property
-    # @abc.abstractmethod
-    # def offset_dict(self) -> dict[tuple[PolyMatrixMixin, int], int]:
-    #     ...
-
-    # @property
-    # @abc.abstractmethod
-    # def param_dict(self) -> dict[tuple[PolyMatrixMixin, int], int]:
-    #     ...
-
     @property
     @abc.abstractmethod
     def n_param(self) -> int:
         ...
 
+    @property
+    @abc.abstractmethod
+    def variable_to_index(self) -> typing.Callable[[int, tuple[int, ...]], int]:
+        ...
+
     @functools.cached_property
     def eq_to_row_index(self):
         rows_to_eq = list(set((eq_idx, perm) for eq_tuple_degree in self.terms.values() for (eq_idx, perm, var) in eq_tuple_degree.keys()))
@@ -77,10 +71,25 @@ class EqualityConstraintMixin(abc.ABC):
     def n_eq(self):
         return len(self.eq_to_row_index)
 
-    # @property
-    # @abc.abstractmethod
-    # def n_var(self) -> int:
-    #     ...
+    def get_constraint_matrix(self):
+        def gen_sparse_matrices():
+            for degree, degree_tuples in self.terms.items():
+                def gen_row_col_data():
+                    for (idx_eq, perm, variables), value in degree_tuples.items():
+                        row = self.eq_to_row_index[(idx_eq, perm)]
+                        col = self.variable_to_index(self.n_param, variables)
+                        yield row, col, value
+                    
+                row, col, data = list(zip(*gen_row_col_data()))
+
+                data = np.array(data, dtype=np.float)
+                
+                if degree <= 1:
+                    yield degree, scipy.sparse.coo_array((data, (row, col)), shape=(self.n_eq, self.n_param**degree)).toarray()
+                else:
+                    yield degree, scipy.sparse.coo_array((data, (row, col)), shape=(self.n_eq, self.n_param**degree))
+
+        return dict(gen_sparse_matrices())
 
     def get_constraint_func(self):
         def func(x):
@@ -100,7 +109,7 @@ class EqualityConstraintMixin(abc.ABC):
 
                     for (idx_eq, perm, variables), value in degree_tuples.items():
                         row_idx = self.eq_to_row_index[(idx_eq, perm)]
-                        vector_val = vector[variable_to_index(self.n_param, variables)]
+                        vector_val = vector[self.variable_to_index(self.n_param, variables)]
                         mat[row_idx] += value * vector_val
 
             return mat
@@ -131,7 +140,7 @@ class EqualityConstraintMixin(abc.ABC):
 
                         for var_idx, var in enumerate(variables):
                             other_variables = variables[:var_idx] + variables[var_idx+1:]
-                            vector_val = vector[variable_to_index(self.n_param, other_variables)]
+                            vector_val = vector[self.variable_to_index(self.n_param, other_variables)]
 
                             # col_idx = variable_to_index(self.n_param, (var,))
                             jac_mat[row_idx, var] += value*vector_val
@@ -168,7 +177,7 @@ class EqualityConstraintMixin(abc.ABC):
 
                             for var_idx_y, var_y in enumerate(other_variables):
                                 other_variables_2 = variables[:var_idx_y] + variables[var_idx_y+1:]
-                                vector_val = vector[variable_to_index(self.n_param, other_variables_2)]
+                                vector_val = vector[self.variable_to_index(self.n_param, other_variables_2)]
                                 hess_mat[var_x, var_y] = v[eq_idx]*value*vector_val
 
             return hess_mat
@@ -196,6 +205,11 @@ class PolyEquationMixin(abc.ABC):
 
         ...
 
+    @property
+    @abc.abstractmethod
+    def variable_to_index(self) -> typing.Callable[[int, tuple[int, ...]], int]:
+        ...
+
     @functools.cached_property
     def _param_list(self) -> list[tuple[PolyMatrixMixin, int], int]:
         """
@@ -206,7 +220,6 @@ class PolyEquationMixin(abc.ABC):
         all_structs = set(indexed_poly_mat for term in self.terms for indexed_poly_mat in term)
 
         def gen_n_param_per_struct():
-            # acc = 0
 
             for struct in all_structs:
 
@@ -214,11 +227,10 @@ class PolyEquationMixin(abc.ABC):
                     continue
 
                 for degree in struct.degrees:
+                    number_of_terms = int(struct.shape[0] * struct.shape[1] * (self.variable_to_index(self.n_var, degree*(self.n_var-1,)) + 1))
+                    # number_of_terms = int(struct.shape[0] * struct.shape[1] * binom(self.n_var+degree-1, degree))
 
-                    if struct.is_vector:
-                        number_of_terms = int(self.n_var * binom(self.n_var+degree-1, degree))
-                    else:
-                        number_of_terms = int(self.n_var**2 * binom(self.n_var+degree-1, degree))
+                    # print(f'{struct=}, {number_of_terms=}')
 
                     yield (struct, degree), number_of_terms
 
@@ -272,15 +284,11 @@ class PolyEquationMixin(abc.ABC):
         self,
         subs: dict[PolyMatrixMixin, dict[DegreeType, dict[int, int, float]]] = None,
     ) -> EqualityConstraintMixin:
-        n_var_2 = self.n_var**2
-
         if subs is None:
             added_subs = {}
         else:
             added_subs = subs
 
-        # binom = scipy.special.binom
-
         # create parameter offset
         all_structs = set(indexed_poly_mat for term in self.terms for indexed_poly_mat in term)
 
@@ -313,150 +321,117 @@ class PolyEquationMixin(abc.ABC):
 
         terms = collections.defaultdict(lambda: collections.defaultdict(float))
 
-        for left, right in self.terms:
+        def gen_re_indexing(term, degrees, curr_rows_col, monoms, col_defaults, d_subs, d_offsets, perm):
+            for m, degree, (poly_row, poly_col), monom, col_default, subs, offset in zip(
+                    term, degrees, itertools.pairwise(curr_rows_col), monoms, col_defaults, d_subs, d_offsets):
 
-            subs_1 = subs_dict[left]
-            subs_2 = subs_dict[right]
+                if m.re_index is not None:
+                    re_index = m.re_index(degree, poly_row, poly_col, monom)
+                else:
+                    re_index = None
 
-            for d1 in left.degrees:
+                if re_index is None:
+                    col = col_default
+                    new_poly_row = poly_row
+                    new_poly_col = poly_col
+                    factor = 1.0
 
-                if subs_1 is not None and d1 in subs_1:
-                    subs_1_d = subs_1[d1]
                 else:
-                    subs_1_d = None
+                    new_poly_row, new_poly_col, new_monom, factor = re_index
+                    col = self.variable_to_index(self.n_var, new_monom)
+
+                if subs is not None:
+                    try:
+                        subs_val = subs[(new_poly_row, new_poly_col, col)]
+                    except KeyError:
+                        subs_val = None
+                else:
+                    subs_val = None
 
-                offset_1 = self.offset_dict.get((left, d1), 0)
+                if subs_val is None:
+                    # the coefficient is selected after the reindexing
 
-                for d2 in right.degrees:
+                    row = new_poly_row + new_poly_col * m.shape[0]
 
-                    if subs_2 is not None and d2 in subs_2:
-                        subs_2_d = subs_2[d2]
-                    else:
-                        subs_2_d = None
+                    # linearize parameter matrix
+                    param_idx = int(offset + row + col * m.shape[0] * m.shape[1])
+                else:
+                    param_idx = None
+
+                # if poly_row == 0 and perm == (0,) and monom == (0,):
+                #     print(f'{new_poly_row=}')
+                #     print(f'{new_poly_col=}')
+                #     print(f'{col=}')
+                #     print(f'{monom=}')
+                #     print(f'{new_monom=}')
+
+                yield poly_row, poly_col, param_idx, factor, subs_val
+
+        for term in self.terms:
 
-                    offset_2 = self.offset_dict.get((right, d2), 0)
+            term_subs = [subs_dict[m] for m in term]
+            # n_matrices = len(term)
 
-                    total_degree = d1 + d2
+            for degrees in itertools.product(*(m.degrees for m in term)):
+                total_degree = sum(degrees)
+                d_subs = [subs[degree] if subs is not None and degree in subs else None for degree, subs in zip(degrees, term_subs)]
+                d_offsets = [self.offset_dict.get((m, degree), 0) for m, degree in zip(term, degrees)]
 
-                    for combination in itertools.combinations_with_replacement(range(self.n_var), total_degree):
-                        # n_var = 2, degree = 3
-                        # cominations = [(0, 0, 0), (0, 0, 1), (0, 1, 1), (1, 1, 1)]
+                for combination in itertools.combinations_with_replacement(range(self.n_var), total_degree):
+                    # n_var = 2, degree = 3
+                    # cominations = [(0, 0, 0), (0, 0, 1), (0, 1, 1), (1, 1, 1)]
+                    
+                    for perm in more_itertools.distinct_permutations(combination):
+
+                        def acc_func(acc, v):
+                            last, _ = acc
+                            new = last + v
+
+                            return new, perm[last:new]
+
+                        # perm=(1, 0, 1) -> monom1=(x2, x1), monom2=(x2)
+                        monoms = list(monom for _, monom in itertools.accumulate(degrees, acc_func, initial=(0, None)))[1:]
                         
-                        # def gen_coord_for_single_equation():
-                        for perm in set(itertools.permutations(combination)):
-                            
-                            # grp1 = (1, 1, 0) -> x1 x2^2
-                            grp1 = perm[0:d1]
-                            grp2 = perm[d1:]
-
-                            def non_increasing(seq):
-                                return all(y <= x for x, y in zip(seq, seq[1:]))
-
-                            # (1,0) -> x2*x1 instead of (0,1)->x1*x2
-                            if non_increasing(grp1) and non_increasing(grp2):
-
-                                left_col_default = variable_to_index(self.n_var, grp1)
-                                right_col_default = variable_to_index(self.n_var, grp2)
-
-                                # for each column of the poly matrix, and row of the poly vector
-                                for left_poly_col in range(self.n_var):
-
-                                    right_poly_row = left_poly_col
-
-                                    if right.re_index is not None:
-                                        re_index_2 = right.re_index(d2, right_poly_row, 0, grp2)
-                                    else:
-                                        re_index_2 = None
-
-                                    if re_index_2 is None:
-                                        n_right_poly_row, factor_2, right_col = right_poly_row, 1, right_col_default
-                                    else:
-                                        n_right_poly_row, _, n_grp2, factor_2 = re_index_2
-                                        right_col = variable_to_index(self.n_var, n_grp2)
-
-                                    key_2 = (n_right_poly_row, 0, right_col)
-                                    if subs_2_d is not None:
-                                        try:
-                                            subs_val_2 = subs_2_d[key_2]
-                                        except KeyError:
-                                            subs_val_2 = None
-                                    else:
-                                        subs_val_2 = None
-
-                                    if factor_2 == 0 or subs_val_2 == 0:
-                                        continue
-
-                                    right_row = n_right_poly_row
-                                    right_param_idx = int(offset_2 + right_row + right_col * self.n_var)
-
-                                    # for each polynomial equation
-                                    for left_poly_row in range(self.n_var):
-
-                                        if left.re_index is not None:
-                                            re_index_1 = left.re_index(d1, left_poly_row, left_poly_col, grp1)
-                                        else:
-                                            re_index_1 = None
-
-                                        if re_index_1 is None:
-                                            n_left_poly_row, n_left_poly_col, factor_1 = left_poly_row, left_poly_col, 1
-                                            left_col = left_col_default
-                                        else:
-                                            n_left_poly_row, n_left_poly_col, n_grp1, factor_1 = re_index_1
-                                            left_col = variable_to_index(self.n_var, n_grp1)
-
-                                        key_1 = (n_left_poly_row, n_left_poly_col, left_col)
-                                        if subs_1_d is not None:
-                                            try:
-                                                subs_val_1 = subs_1_d[key_1]
-                                            except KeyError:
-                                                subs_val_1 = None
-                                        else:
-                                            subs_val_1 = None
-
-                                        if factor_1 == 0 or subs_val_1 == 0:
-                                            continue
-
-                                        left_row = n_left_poly_row + n_left_poly_col * self.n_var
-                                        left_param_idx = int(offset_1 + left_row + left_col * n_var_2)
-
-                                        # print(left_poly_row)
-                                        # print(f'{left_row=}')
-                                        # print(f'{left_param_idx=}')
-
-                                        total_factor = factor_1 * factor_2
-
-                                        match (subs_val_1, subs_val_2):
-                                            case (None, None):
-                                                col_idx = (left_param_idx, right_param_idx)
-                                                degree = 2
-                                                value = total_factor
-
-                                                # terms[1][(idx_eq, perm)][left_param_idx] += 2*value
-                                                # terms[1][(idx_eq, perm)][right_param_idx] += 2*value
-                                                # terms[0][(idx_eq, perm)][0] += value
-                                            
-                                            case (subs_val, None):
-                                                col_idx = (right_param_idx,)
-                                                degree = 1
-                                                value = subs_val_1*total_factor
-
-                                                # terms[degree][(idx_eq, perm)][0] += value
-
-                                            case (None, subs_val):
-                                                col_idx = (left_param_idx,)
-                                                degree = 1
-                                                value = subs_val_2*total_factor
-
-                                                # terms[degree][(idx_eq, perm)][0] += value
-                                            
-                                            case _:
-                                                degree, col_idx, value = 0, tuple(), subs_val_1*subs_val_2*total_factor
-
-                                        terms[degree][left_poly_row, perm, col_idx] += value
+                        def non_increasing(seq):
+                            return all(y <= x for x, y in zip(seq, seq[1:]))
+
+                        # (1,0) -> x2*x1 instead of (0,1)->x1*x2
+                        if all(non_increasing(monom) for monom in monoms):
+
+                            col_defaults = [self.variable_to_index(self.n_var, monom) for monom in monoms]
+
+                            n_rows_col = itertools.chain((m.shape[0] for m in term), (term[-1].shape[1],))
+
+                            for curr_rows_col in itertools.product(*[range(e) for e in n_rows_col]):
+                                
+                                data = tuple(gen_re_indexing(term, degrees, curr_rows_col, monoms, col_defaults, d_subs, d_offsets, perm))
+
+                                total_factor = functools.reduce(lambda x, y: x*y, (d[3] for d in data))
+
+                                if total_factor == 0:
+                                    continue
+
+                                value = functools.reduce(lambda x, y: x*y, (d[4] for d in data if d[4] is not None), 1) * total_factor
+
+                                if value == 0:
+                                    continue
+
+                                poly_row = data[0][0]
+                                param_idx = tuple(d[2] for d in data if d[4] is None)
+                                degree = len(param_idx)
+                                
+                                # if poly_row == 0 and perm == (0,):
+                                #     print(f'{param_idx=}')
+                                #     print(f'{data=}')
+                                #     print(f'{monoms=}')
+                                #     print(f'{total_factor=}')
+
+                                terms[degree][poly_row, perm, param_idx] += value
 
         return EqualityConstraintImpl(
             terms=terms, 
             n_param=self.n_param,
+            variable_to_index=self.variable_to_index,
         )
 
     def matrix_to_poly(self, struct, x, param, tol=None):
@@ -526,47 +501,48 @@ class PolyMatrixImpl(PolyMatrix):
     subs: dict[int, dict[tuple[int, int], float]]
     re_index: typing.Callable[[int, int, int, tuple[int, ...]], tuple[int, int, int, tuple[int, ...], float]]
     is_constant: bool
-    is_vector: bool
+    shape: tuple[int, int]
 
 
 @dataclass_abc.dataclass_abc(frozen=True)
 class EqualityConstraintImpl(EqualityConstraintMixin):
     terms: dict[int, dict[tuple[int, tuple, int], float]]
-    # offset_dict: dict[tuple[PolyMatrixMixin, int], int]
-    # param_dict: dict[tuple[PolyMatrixMixin, int], int]
+    variable_to_index: typing.Callable[[int, tuple[int, ...]], int]
     n_param: int
-    # n_var: int
 
 
 @dataclass_abc.dataclass_abc(frozen=True)
 class EquationImpl(PolyEquation):
     terms: list[tuple[PolyMatrix, PolyMatrix]]
+    variable_to_index: typing.Callable[[int, tuple[int, ...]], int]
     n_var: int
 
 ########################################
 # init functions
 ########################################
 
-def init_poly_vector(
-    degrees: list[int] = None,
-    subs: dict[int, dict[tuple[int, int], float]] = None,
-    re_index: typing.Callable[[int, int, int, tuple[int, ...]], tuple[int, int, int, tuple[int, ...], float]] = None,
-    is_constant: bool = None,
-):
-    return init_poly_matrix(
-        degrees=degrees,
-        subs=subs,
-        re_index=re_index,
-        is_vector=True,
-        is_constant=is_constant,
-    )
+# def init_poly_vector(
+#     n_row: int,
+#     degrees: list[int] = None,
+#     subs: dict[int, dict[tuple[int, int], float]] = None,
+#     re_index: typing.Callable[[int, int, int, tuple[int, ...]], tuple[int, int, int, tuple[int, ...], float]] = None,
+#     is_constant: bool = None,
+# ):
+#     return init_poly_matrix(
+#         degrees=degrees,
+#         subs=subs,
+#         re_index=re_index,
+#         shape=(n_row, ,
+#         n_col=1,
+#         is_constant=is_constant,
+#     )
 
 
 def init_poly_matrix(
+    shape: tuple[int, int],
     degrees: list[int] = None,
     subs: dict[int, dict[tuple[int, int], float]] = None,
     re_index: typing.Callable[[int, int, int, tuple[int, ...]], tuple[int, int, int, tuple[int, ...], float]] = None,
-    is_vector: bool = None,
     is_constant: bool = None,
 ):
     if degrees is None:
@@ -579,25 +555,27 @@ def init_poly_matrix(
         else:
             is_constant = True
 
-    if is_vector is None:
-        is_vector = False
-
     return PolyMatrixImpl(
         degrees=degrees,
         subs=subs,
-        re_index = re_index,
-        is_constant = is_constant,
-        is_vector = is_vector,
+        re_index=re_index,
+        is_constant=is_constant,
+        shape=shape,
     )
 
 
 def init_equation(
     n_var: int,
     terms: list[tuple[PolyMatrix, PolyMatrix]],
+    variable_to_index: typing.Callable[[int, tuple[int, ...]], int] = None,
 ):
-    assert all(not left.is_vector and right.is_vector for left, right in terms)
+    # assert all(not left.is_vector and right.is_vector for left, right in terms)
+
+    if variable_to_index is None:
+        variable_to_index = polymatrix.utils.variable_to_index
 
     return EquationImpl(
         n_var=n_var,
         terms=terms,
+        variable_to_index=variable_to_index,
     )
diff --git a/polymatrix/sympyutils.py b/polymatrix/sympyutils.py
index 6fe7938..7413ba3 100644
--- a/polymatrix/sympyutils.py
+++ b/polymatrix/sympyutils.py
@@ -19,17 +19,13 @@ def poly_to_data_coord(poly_list, x, degree = None):
 
     if degree is None:
         degree = max(degree for inner_poly_list in sympy_poly_list for poly in inner_poly_list for degree in poly.degree_list())
-
-    # def all_equal(iterable):
-    #     g = itertools.groupby(iterable)
-    #     return next(g, True) and not next(g, False)
-
-    # assert all_equal((p.gens for p in poly_list)), 'all polynomials need to have identical generators'
+    
+    n = len(x)
 
     def gen_power_mat():
 
         # for all powers generate a matrix
-        for current_degree in range(degree + 1):
+        for i_degree in range(degree + 1):
         
             def gen_value_index():
 
@@ -41,45 +37,131 @@ def poly_to_data_coord(poly_list, x, degree = None):
                         # a5 x1 x3**2 -> c=a5, m=(1, 0, 2)
                         for c, m in zip(p.coeffs(), p.monoms()):
 
-                            if sum(m) == current_degree:
+                            if sum(m) == i_degree:
+                                
+                                indices = (idx for idx, p in enumerate(m) for _ in range(p))
+                                index = sum(idx*(n**level) for level, idx in enumerate(indices))
 
-                                index = variable_powers_to_index(m)
                                 yield (poly_row, poly_col, index), c
                     
             data = dict(gen_value_index()) | collections.defaultdict(float)
 
             if len(data) > 0:
-                yield current_degree, data
+                yield i_degree, data
 
     return dict(gen_power_mat())
 
 
 def poly_to_matrix(poly_list, x, power = None):
-    """
-    
-    """
-
     data_coord_dict = poly_to_data_coord(poly_list, x, power)
 
-    # n_free_symbols = len(poly_list[0][0].gens)
     n_free_symbols = len(x)
+    n_poly = len(poly_list)
+    n_eq = sum((1 for row_eq in poly_list for e in row_eq))
 
     def gen_power_mat():
+        if 0 not in data_coord_dict:
+            yield 0, np.zeros((n_free_symbols, 1))
 
         # for all powers generate a matrix
-        for current_degree, data_coord in data_coord_dict:
-            
+        for degree, data_coord in data_coord_dict.items():
+
             # empty matrix
-            shape = (len(poly_list), n_free_symbols**current_degree)
-            # m = np.zeros((len(poly_list), n_free_symbols**current_power))
+            shape = (n_eq, n_free_symbols**degree)
 
             # fill matrix
             if len(data_coord) == 0:
-                yield np.zeros((len(poly_list), n_free_symbols**current_degree))
+                yield np.zeros((len(poly_list), n_free_symbols**degree))
 
             else:
-                rows, cols, data = list(zip(*data_coord))
+                def gen_row_col_coord():
+                    for (p_row, p_col, index), value in data_coord.items():
+                        yield p_row + p_col * n_poly, index, value
+                
+                rows, cols, data = list(zip(*gen_row_col_coord()))
+                sparse_array = scipy.sparse.coo_matrix((data, (rows, cols)), dtype=np.double, shape=shape)
+                
+                if degree <= 1:
+                    yield degree, sparse_array.toarray()
+                else:
+                    yield degree, sparse_array.tocsr()
+            
+    return dict(gen_power_mat())
+
+
+# def poly_to_data_coord(poly_list, x, degree = None):
+#     """
+#     poly_list = [
+#         poly(x1*x3**2, x)
+#     ]
+#     power: up to which power 
+#     """
+
+#     sympy_poly_list = tuple(tuple(sympy.poly(p, x) for p in inner_poly_list) for inner_poly_list in poly_list)
+
+#     if degree is None:
+#         degree = max(degree for inner_poly_list in sympy_poly_list for poly in inner_poly_list for degree in poly.degree_list())
+
+#     # def all_equal(iterable):
+#     #     g = itertools.groupby(iterable)
+#     #     return next(g, True) and not next(g, False)
+
+#     # assert all_equal((p.gens for p in poly_list)), 'all polynomials need to have identical generators'
+
+#     def gen_power_mat():
+
+#         # for all powers generate a matrix
+#         for current_degree in range(degree + 1):
+        
+#             def gen_value_index():
+
+#                 # each polynomial defines a row in the matrix
+#                 for poly_row, inner_poly_list in enumerate(sympy_poly_list):
+
+#                     for poly_col, p in enumerate(inner_poly_list):
+
+#                         # a5 x1 x3**2 -> c=a5, m=(1, 0, 2)
+#                         for c, m in zip(p.coeffs(), p.monoms()):
+
+#                             if sum(m) == current_degree:
+
+#                                 index = variable_powers_to_index(m)
+#                                 yield (poly_row, poly_col, index), c
+                    
+#             data = dict(gen_value_index()) | collections.defaultdict(float)
+
+#             if len(data) > 0:
+#                 yield current_degree, data
+
+#     return dict(gen_power_mat())
+
+
+# def poly_to_matrix(poly_list, x, power = None):
+#     """
+    
+#     """
+
+#     data_coord_dict = poly_to_data_coord(poly_list, x, power)
+
+#     # n_free_symbols = len(poly_list[0][0].gens)
+#     n_free_symbols = len(x)
+
+#     def gen_power_mat():
+
+#         # for all powers generate a matrix
+#         for current_degree, data_coord in data_coord_dict:
+            
+#             # empty matrix
+#             shape = (len(poly_list), n_free_symbols**current_degree)
+#             # m = np.zeros((len(poly_list), n_free_symbols**current_power))
+
+#             # fill matrix
+#             if len(data_coord) == 0:
+#                 yield np.zeros((len(poly_list), n_free_symbols**current_degree))
+
+#             else:
+#                 rows, cols, data = list(zip(*data_coord))
 
-                yield scipy.sparse.coo_matrix((data, (rows, cols)), dtype=np.double, shape=shape)
+#                 yield scipy.sparse.coo_matrix((data, (rows, cols)), dtype=np.double, shape=shape)
             
-    return list(gen_power_mat())
+#     return list(gen_power_mat())
diff --git a/test_polymatrix/test_polymatrix.py b/test_polymatrix/test_polymatrix.py
index 3b239bf..f693df0 100644
--- a/test_polymatrix/test_polymatrix.py
+++ b/test_polymatrix/test_polymatrix.py
@@ -1,114 +1,195 @@
 import unittest
 
-from polymatrix.polystruct import init_equation, init_poly_matrix, init_poly_vector
+from polymatrix.polystruct import init_equation, init_poly_matrix
 from polymatrix.utils import variable_to_index
 
 class TestPolyMatrix(unittest.TestCase):
     @staticmethod
-    def assert_term_in_eq(terms, degree, eq_idx, row_idx, value, monoms=None):
+    def assert_term_in_eq(result, degree, eq_idx, row_idx, value, monoms=None):
         if monoms is None:
             monoms = tuple()
 
-        assert degree in terms, f'could not find {degree} in {terms}'
-        degree_terms = terms[degree]
+        assert degree in result.terms, f'could not find {degree} in {result.terms}'
+        degree_terms = result.terms[degree]
 
-        key = eq_idx, monoms
+        key = eq_idx, monoms, row_idx
         assert key in degree_terms, f'could not find {key} in {degree_terms}'
-        eq_terms = degree_terms[key]
+        # eq_terms = degree_terms[key]
+        # assert row_idx in eq_terms, f'could not find {row_idx} in {eq_terms}'
+        assert degree_terms[key] == value, f'value {degree_terms[key]} and {value} do not match'
 
-        assert row_idx in eq_terms, f'could not find {row_idx} in {eq_terms}'
-        assert eq_terms[row_idx] == value, f'value {eq_terms[row_idx]} and {value} do not match'
-
-    def test_param_matrix_param_vector_degree_0(self):
+    def test_param_matrix_param_d0_vector_degree_d0(self):
         """
         param = [a11 a21 a31 a41 v11 v21]
         """
 
-        mat = init_poly_matrix(degrees=(0,))
-        vec = init_poly_vector(degrees=(0,))
-        n_param = 6
+        n_var = 2
+
+        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(degrees=(0,), shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         # a11 v11
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
+            degree = 2,
+            eq_idx = 0,
+            row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 0)]),
+            value = 1,
+        )
+
+        # a12 v21
+        self.assert_term_in_eq(
+            result = result,
             degree = 2,
             eq_idx = 0,
-            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 0)])),
+            row_idx = (offset_dict[(mat, 0)]+2, offset_dict[(vec, 0)]+1),
             value = 1,
         )
 
         # a21 v11
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 2,
             eq_idx = 1,
-            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)]+2, offset_dict[(vec, 0)])),
+            row_idx = (offset_dict[(mat, 0)]+1, offset_dict[(vec, 0)]),
             value = 1,
         )
 
-    def test_param_matrix_param_vector_degree_01(self):
+        # a22 v21
+        self.assert_term_in_eq(
+            result = result,
+            degree = 2,
+            eq_idx = 1,
+            row_idx = (offset_dict[(mat, 0)]+3, offset_dict[(vec, 0)]+1),
+            value = 1,
+        )
+
+    def test_param_matrix_d0_param_vector_d01(self):
         """
         param = [a11 a21 a31 a41 v011 v021 v111 v112 v121 v122]
         """
 
-        mat = init_poly_matrix(degrees=(0,))
-        vec = init_poly_vector(degrees=(0,1))
-        n_param = 10
+        n_var = 2
+
+        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(degrees=(0,1), shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         # a11 v011
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 2,
             eq_idx = 0,
-            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 0)])),
+            row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 0)]),
             value = 1,
         )
 
         # a11 v011
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 2,
             eq_idx = 0,
             monoms=(0,),
-            row_idx = variable_to_index(n_param, (offset_dict[(mat, 0)], offset_dict[(vec, 1)])),
+            row_idx = (offset_dict[(mat, 0)], offset_dict[(vec, 1)]),
+            value = 1,
+        )
+
+    def test_param_matrix_d0_const_vector_d0(self):
+        """
+        param = [a11 a21 a31 a41]
+        """
+
+        n_var = 2
+
+        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
+
+        eq = init_equation(
+            terms = [(mat, vec)],
+            n_var = n_var,
+        )
+
+        result = eq.create()
+        offset_dict = eq.offset_dict
+
+        # a11
+        self.assert_term_in_eq(
+            result = result,
+            degree = 1,
+            eq_idx = 0,
+            row_idx = (offset_dict[(mat, 0)],),
             value = 1,
         )
 
-    def test_param_matrix_const_vector_degree_0(self):
+        # a21
+        self.assert_term_in_eq(
+            result = result,
+            degree = 1,
+            eq_idx = 1,
+            row_idx = (offset_dict[(mat, 0)]+1,),
+            value = 1,
+        )
+
+    def test_param_matrix_d0_const_vector_d1(self):
         """
         param = [a11 a21 a31 a41]
         """
 
-        mat = init_poly_matrix(degrees=(0,))
-        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})
+        n_var = 2
+
+        mat = init_poly_matrix(degrees=(0,), shape=(n_var, n_var))
+        vec = init_poly_matrix(subs={1: {(0, 0, 0): 1, (0, 0, 1): 0, (1, 0, 0): 0, (1, 0, 1): 1}}, shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         # a11
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             eq_idx = 0,
-            row_idx = offset_dict[(mat, 0)],
+            monoms=(0,),
+            row_idx = (offset_dict[(mat, 0)],),
+            value = 1,
+        )
+
+        # a12
+        self.assert_term_in_eq(
+            result = result,
+            degree = 1,
+            eq_idx = 0,
+            monoms=(1,),
+            row_idx = (offset_dict[(mat, 0)]+2,),
+            value = 1,
+        )
+
+        # a21
+        self.assert_term_in_eq(
+            result = result,
+            degree = 1,
+            eq_idx = 1,
+            monoms=(0,),
+            row_idx = (offset_dict[(mat, 0)]+1,),
             value = 1,
         )
 
@@ -117,21 +198,24 @@ class TestPolyMatrix(unittest.TestCase):
         param = [a11 a21 a31 a41]
         """
 
-        mat = init_poly_matrix(subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}})
-        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})
+        n_var = 2
+
+        mat = init_poly_matrix(subs={0: {(0, 0, 0): 1, (0, 1, 0): 1, (1, 0, 0): 1, (1, 1, 0): 1}}, shape=(n_var, n_var))
+        vec = init_poly_matrix(subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 0,
             eq_idx = 0,
-            row_idx = 0,
+            row_idx = tuple(),
             value = 2,
         )
     
@@ -140,98 +224,107 @@ class TestPolyMatrix(unittest.TestCase):
         param = [a11 a21 a31 a41]
         """
 
-        def skew_symmetric(degree, idx1, idx2):
-            if idx1 == idx2:
-                return idx1, idx2, 0
-            elif idx2 < idx1:
-                return idx2, idx1, -1
+        def skew_symmetric(degree, poly_row, poly_col, monom):
+            if poly_row == poly_col:
+                return poly_row, poly_col, monom, 0
+            elif poly_col < poly_row:
+                return poly_col, poly_row, monom, -1
+
+        n_var = 2
 
         mat = init_poly_matrix(
             degrees=(0,), 
-            re_index_func=skew_symmetric,
+            re_index=skew_symmetric,
+            shape=(n_var, n_var),
         )
-        vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})
+        vec = init_poly_matrix(subs={0: {(0, 0, 0): 1, (1, 0, 0): 1}}, shape=(n_var, 1))
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             eq_idx = 0,
-            row_idx = offset_dict[(mat, 0)] + 1,
+            row_idx = (offset_dict[(mat, 0)] + 2,),
             value = 1,
         )
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             eq_idx = 1,
-            row_idx = offset_dict[(mat, 0)] + 1,
+            row_idx = (offset_dict[(mat, 0)] + 2,),
             value = -1,
         )
 
-    def test_const_matrix_param_gradient_vector(self):
+    def test_const_matrix_d0_param_gradient_vector_d1(self):
         """
         param = [v11 v12 v21 v22]
         """
 
-        def gradient(degree, v_row, monom):
+        def gradient(degree, p_row, p_col, monom):
             if degree == 1:
-                factor = sum(v_row==e for e in monom) + 1
+                factor = sum(p_row==e for e in monom) + 1
 
-                if monom[-1] < v_row:
-                    n_v_row = monom[-1]
-                    n_monom = sorted(monom + (v_row,), reverse=True)
+                if monom[-1] < p_row:
+                    n_p_row = monom[-1]
+                    n_monom = sorted(monom + (p_row,), reverse=True)
 
-                if v_row <= monom[-1]:
-                    n_v_row = v_row
+                if p_row <= monom[-1]:
+                    n_p_row = p_row
                     n_monom = monom
 
-                return n_v_row, n_monom, factor
+                return n_p_row, p_col, n_monom, factor
+
+        n_var = 2
 
         mat = init_poly_matrix(
-            subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}}, 
+            subs={0: {(0, 0, 0): 1, (0, 1, 0): 1, (1, 0, 0): 1, (1, 1, 0): 1}}, 
+            shape=(n_var, n_var),
         )
-        vec = init_poly_vector(
+        vec = init_poly_matrix(
             degrees=(1,),
-            re_index_func_2=gradient,
+            re_index=gradient,
+            shape=(n_var, 1),
         )
 
         eq = init_equation(
             terms = [(mat, vec)],
-            n_var = 2,
+            n_var = n_var,
         )
 
-        terms, offset_dict = list(eq.create())
+        result = eq.create()
+        offset_dict = eq.offset_dict
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             monoms=(0,),
             eq_idx = 0,
-            row_idx = offset_dict[(vec, 1)],
+            row_idx = (offset_dict[(vec, 1)],),
             value = 2,
         )
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             monoms=(1,),
             eq_idx = 0,
-            row_idx = offset_dict[(vec, 1)]+1,
+            row_idx = (offset_dict[(vec, 1)]+2,),
             value = 1,
         )
 
         self.assert_term_in_eq(
-            terms = terms,
+            result = result,
             degree = 1,
             monoms=(0,),
             eq_idx = 1,
-            row_idx = offset_dict[(vec, 1)]+1,
+            row_idx = (offset_dict[(vec, 1)]+2,),
             value = 1,
         )