add skew-symmetry and gradient function

3 files changed, 259 insertions, 131 deletions

diff --git a/main.py b/main.py
index 806d1a1..4079ac5 100644
--- a/main.py
+++ b/main.py
@@ -1,5 +1,6 @@
 import sympy
 import scipy.special
+import scipy.sparse
 import numpy as np
 from polymatrix.polystruct import init_equation, init_poly_matrix, init_poly_vector
 from polymatrix.sympyutils import poly_to_data_coord, poly_to_matrix
@@ -7,6 +8,28 @@ from polymatrix.utils import variable_to_index
 
 binom = scipy.special.binom
 
+def skew_symmetric(degree, p_row, p_col):
+    if p_col < p_row:
+        return p_col, p_row, -1
+
+def zero_diagonal_for_uneven_degree(degree, p_row, p_col):
+    if degree % 2 == 1 and p_row == p_col:
+        return 0
+
+def gradient(degree, v_row, monom):
+    if degree == 1:
+        factor = sum(v_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 v_row <= monom[-1]:
+            n_v_row = v_row
+            n_monom = monom
+
+        return n_v_row, n_monom, factor
+
 x = sympy.symbols('v_dc, i_d, i_q')
 vdc, id, iq = x
 
@@ -40,17 +63,10 @@ nabla_h = init_poly_vector(subs=dh_terms)
 JR = init_poly_matrix(subs=jr_terms)
 mg = init_poly_matrix(subs=mg_terms)
 
-def skew_symmetric(idx1, idx2):
-    if idx2 < idx1:
-        return idx2, idx1, -1
-
-def zero_diagonal_for_uneven_degree(degree, idx1, idx2):
-    if degree % 2 == 1 and idx1 == idx2:
-        return 0
-
-    
-
-nabla_ha = init_poly_vector(degrees=(1,2))
+nabla_ha = init_poly_vector(
+    degrees=(1,),
+    re_index_func_2=gradient,
+)
 JRa = init_poly_matrix(
     degrees=(0,1,2), 
     re_index_func=skew_symmetric, 
@@ -58,6 +74,13 @@ JRa = init_poly_matrix(
 )
 u = init_poly_vector(degrees=(1,2))
 
+eq = init_equation(
+    terms = [(JR, nabla_ha), (JRa, nabla_ha), (JRa, nabla_h), (mg, u)],
+    n_var = 2,
+    # terms = [(JR, nabla_ha)],
+    # n_var = 2,
+)
+
 # n_var = 2
 # n_param_nabla_ha = sum(n_var*binom(n_var+p-1, p) for p in nabla_ha.degrees)
 # n_param_JRa = sum(n_var**2*binom(n_var+p-1, p) for p in JRa.non_zero_degrees)
@@ -69,33 +92,68 @@ u = init_poly_vector(degrees=(1,2))
 # print(f'{n_param_u=}')
 # print(f'{total=}')
 
-# print(n_var*binom(n_var+1-1, 1))
-# print(n_var*binom(n_var+2-1, 2))
-
 # print(binom(total+2-1, 2))
 
+# mat = init_poly_matrix(
+#     degrees=(1,), 
+#     re_index_func=skew_symmetric, 
+#     subs_func=zero_diagonal_for_uneven_degree,
+# )
+# vec = init_poly_vector(
+#     subs={0: {(0, 0): 1, (1, 0): 1}},
+# )
+
+# # mat = init_poly_matrix(
+# #     subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}}, 
+# # )
+
+# # vec = init_poly_vector(
+# #     degrees=(1,),
+# #     re_index_func_2=gradient,
+# # )
+
 # eq = init_equation(
-#     # terms = [(JR, nabla_ha), (JRa, nabla_ha), (JRa, nabla_h), (mg, u)],
-#     # n_var = 3,
-#     terms = [(JR, nabla_ha)],
+#     terms = [(mat, vec)],
 #     n_var = 2,
 # )
 
-# print(eq.create())
+eq_tuples, offset_dict, n_param = eq.create()
 
-mat = init_poly_matrix(degrees=(0,), re_index_func=skew_symmetric)
-# vec = init_poly_vector(degrees=(0,1))
-vec = init_poly_vector(subs={0: {(0, 0): 1, (1, 0): 1}})
+print(f'{list(offset_dict.values())=}')
+print(f'{n_param=}')
 
-eq = init_equation(
-    terms = [(mat, vec)],
-    n_var = 2,
-)
-eq_tuples, offset_dict = eq.create()
-
-# print(eq_tuples)
-# print(eq_tuples[1][(1, tuple())])
-# print(offset_dict[(mat, 0)])
-# print(offset_dict[(vec, 0)])
-print(offset_dict)
-# # print(variable_to_index(10, (offset_dict[(mat, 0)], offset_dict[(vec, 1)])))
+# mapping from row index to entry in eq_tuples
+rows_to_eq = list(set(key for eq_tuple_degree in eq_tuples.values() for key in eq_tuple_degree.keys()))
+eq_to_rows = {eq: idx for idx, eq in enumerate(rows_to_eq)}
+
+print(f'n_eq = {len(rows_to_eq)}')
+
+# # mapping from col index to entry in eq_tuples
+# cols_to_var = list(set(key for eq in eq_tuples[degree].values() for key in eq.keys()))
+# var_to_cols = {var: idx for idx, var in enumerate(cols_to_var)}
+
+def gen_matrix_per_degree():
+    for degree, eq_tuple_degree in eq_tuples.items():
+
+        if 0 < degree:
+            def gen_coords(eq_tuple_degree=eq_tuple_degree):
+                for eq, var_dict in eq_tuple_degree.items():
+                    zero_degree_val = eq_tuples[0][eq][0]
+                    row = eq_to_rows[eq]
+                    for var, value in var_dict.items():
+                        # col = variable_to_index(n_param, var)
+                        # col = var_to_cols[var]
+                        col = var
+                        yield row, col, value / zero_degree_val
+
+            row, col, data = list(zip(*gen_coords()))
+
+            n_col = int(binom(n_param+degree-1, degree))
+            n_row = int(max(row) + 1)
+            sparse_matrix = scipy.sparse.coo_matrix((data, (row, col)), dtype=np.float64, shape=(n_row, n_col)).toarray()
+
+            yield degree, sparse_matrix
+
+result = dict(gen_matrix_per_degree())
+
+print(result[1].shape)
diff --git a/polymatrix/polystruct.py b/polymatrix/polystruct.py
index 69a05d6..64c0f42 100644
--- a/polymatrix/polystruct.py
+++ b/polymatrix/polystruct.py
@@ -70,7 +70,9 @@ class EquationMixin(abc.ABC):
         subs: dict[PolyMatrixMixin, dict[DegreeType, dict[int, int, float]]] = None,
     ):
         if subs is None:
-            subs = {}
+            added_subs = {}
+        else:
+            added_subs = subs
 
         binom = scipy.special.binom
 
@@ -81,106 +83,102 @@ class EquationMixin(abc.ABC):
             acc = 0
 
             for struct in all_structs:
-                added_subs = subs.get(struct, {})
 
                 for degree in struct.degrees:
-                    
-                    if struct.subs is not None and degree in struct.subs:
-                        struct_subs = struct.subs[degree]
-                    
-                    elif struct.subs_func is not None:
-                        def gen_coords():
-                            for eq in range(self.n_var):
-                                
-                                if struct.is_vector:
-                                    col_indices = (0,)
-                                else:
-                                    col_indices = range(self.n_var)
 
-                                for col in col_indices:
-                                    coord = (eq, col)
+                    yield (struct, degree), acc
 
-                                    subs_val = struct.subs_func(*coord)
+                    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))
 
-                                    if subs_val is not None:
-                                        yield coord, subs_val
+                    acc += number_of_terms
 
-                        struct_subs = dict(gen_coords())
-                    
-                    else:
-                        struct_subs = {}
+            yield None, acc #, None
 
-                    if struct.re_index is not None and degree in struct.re_index:
-                        struct_re_index = struct.re_index[degree]
-                    
-                    elif struct.re_index_func is not None:
-                        def gen_coords():
-                            for eq in range(self.n_var):
-                                
-                                if struct.is_vector:
-                                    col_indices = (0,)
-                                else:
-                                    col_indices = range(self.n_var)
+        param_list = list(gen_n_param_per_struct())
+        n_param = param_list[-1][1]
+        offset_dict = dict(((e[0], e[1]) for e in param_list[:-1]))
+
+        def gen_substitutions():
+            for struct in all_structs:
+                added_struct_subs = added_subs.get(struct, None)
 
-                                for col in col_indices:
-                                    coord = (eq, col)
+                if added_struct_subs is not None:
+                    if struct.subs is not None:
+                        def gen_merged_subs():
+                            for degree in struct.degrees:
+                                all_subs = struct.subs.get(degree, {}) | added_struct_subs.get(degree, {})
+                                yield degree, all_subs
 
-                                    re_index_val = struct.re_index_func(*coord)
+                        all_subs = dict(gen_merged_subs())
+                    
+                    else:
+                        all_subs = added_struct_subs
 
-                                    match re_index_val:
-                                        case (row, col, value):
-                                            yield coord, ((row, col), value)
-                        struct_re_index = dict(gen_coords())
+                else:
+                    if struct.subs is not None:
+                        all_subs = struct.subs
 
                     else:
-                        struct_re_index = {}
+                        all_subs = None
+                
+                def subs_func(degree, idx_eq, idx_col, all_subs=all_subs, struct=struct):
 
-                    all_subs = struct_subs | added_subs.get(degree, {})
+                    if struct.re_index_func is not None:
+                        re_index = struct.re_index_func(degree, idx_eq, idx_col)
+                    else:
+                        re_index = None
+
+                    if re_index is None:
+                        n_coord = (idx_eq, idx_col)
+                        factor = 1
+                    else:
+                        n_coord = (re_index[0], re_index[1])
+                        factor = re_index[2]
 
-                    yield (struct, degree), acc, (all_subs, struct_re_index)
+                    if all_subs is not None and degree in all_subs:
+                        all_degree_subs = all_subs[degree]
 
-                    # def gen_predefined_coord():
-                    #     re_index = re_index_dict[struct].get(degree, {})
+                        if n_coord in all_degree_subs:
+                            subs_val = all_degree_subs[n_coord] * factor
+                        else:
+                            subs_val = None
 
-                    #     for coord_dict in (all_subs, re_index):
-                    #         for coord in coord_dict.keys():
-                    #             # in case the n_var is smaller than expected
-                    #             if max(coord) <= self.n_var:
-                    #                 yield coord
+                    elif struct.subs_func is not None:
+                        sub_value = struct.subs_func(degree, n_coord[0], n_coord[1])
 
-                    # n_predefined = len(set(gen_predefined_coord()))
+                        if sub_value is not None:
+                            subs_val = sub_value * factor
+                        else:
+                            subs_val = None
 
-                    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))
+                        subs_val = None
 
-                    acc += number_of_terms
+                    return n_coord, factor, subs_val
 
-            yield None, acc, None
+                yield struct, subs_func
 
-        param_list = list(gen_n_param_per_struct())
-        n_param = param_list[-1][1]
-        offset_dict = dict(((e[0], e[1]) for e in param_list[:-1]))
-        subs_dict = dict(((e[0], e[2][0]) for e in param_list[:-1]))
-        re_index_dict = dict(((e[0], e[2][1]) for e in param_list[:-1]))
+        subs_dict = dict(gen_substitutions())
 
         terms = collections.defaultdict(lambda: collections.defaultdict(lambda: collections.defaultdict(int)))
 
         for left, right in self.terms:
+
+            subs_1 = subs_dict[left]
+            subs_2 = subs_dict[right]
+
             for d1 in left.degrees:
 
                 n_param_1 = binom(self.n_var+d1-1, d1)
                 offset_1 = offset_dict[(left, d1)]
-                subs_1 = subs_dict[(left, d1)]
-                re_index_1 = re_index_dict[(left, d1)]
 
                 for d2 in right.degrees:
 
                     n_param_2 = binom(self.n_var+d2-1, d2)
                     offset_2 = offset_dict[(right, d2)]
-                    subs_2 = subs_dict[(right, d2)]
-                    re_index_2 = re_index_dict[(right, d2)]
 
                     total_degree = d1 + d2
 
@@ -202,67 +200,73 @@ class EquationMixin(abc.ABC):
                             if non_increasing(grp1) and non_increasing(grp2):
 
                                 left_idx = variable_to_index(self.n_var, grp1)
-                                right_idx = variable_to_index(self.n_var, grp2)
+                                d_right_idx = variable_to_index(self.n_var, grp2)
 
                                 # for each column of the poly matrix, and row of the poly vector
                                 for idx_col in range(self.n_var):
 
-                                    coord_2 = (idx_col, 0)
-                                    n_coord_2, factor_2 = re_index_2.get(coord_2, (coord_2, 1))
-
-                                    if factor_2 == 0:
-                                        continue
+                                    if right.re_index_func_2 is None:
+                                        re_index_2 = None
+                                    else:
+                                        re_index_2 = right.re_index_func_2(d2, idx_col, grp2)
 
-                                    if n_coord_2 in subs_2:
-                                        subs_val_2 = factor_2 * subs_2[n_coord_2]
-                                        if subs_val_2 == 0:
-                                            continue
+                                    if re_index_2 is None:
+                                        v_idx_row = idx_col
+                                        factor_22 = 1
+                                        right_idx = d_right_idx
                                     else:
-                                        subs_val_2 = None
+                                        v_idx_row, n_grp2, factor_22 = re_index_2
+                                        right_idx = variable_to_index(self.n_var, n_grp2)
+
+                                    n_coord_2, factor_21, subs_val_2 = subs_2(d2, v_idx_row, 0)
+
+                                    factor_2 = factor_21 * factor_22
+
+                                    if factor_2 == 0 or subs_val_2 == 0:
+                                        continue
 
                                     # for each polynomial equation
                                     for idx_eq in range(self.n_var):
 
-                                        coord_1 = (idx_eq, idx_col)
-                                        n_coord_1, factor_1 = re_index_1.get(coord_1, (coord_1, 1))
+                                        n_coord_1, factor_1, subs_val_1 = subs_1(d1, idx_eq, idx_col)
 
-                                        if factor_1 == 0:
+                                        if factor_1 == 0 or subs_val_1 == 0:
                                             continue
-
-                                        if n_coord_1 in subs_1:
-                                            subs_val_1 = factor_1 * subs_1[n_coord_1]
-                                            if subs_val_1 == 0:
-                                                continue
-                                        else:
-                                            subs_val_1 = None
                                         
                                         left_param_idx = int(offset_1 + left_idx + (self.n_var * n_coord_1[0] + n_coord_1[1]) * n_param_1)
                                         right_param_idx = int(offset_2 + right_idx + n_coord_2[0] * n_param_2)
 
                                         match (subs_val_1, subs_val_2):
                                             case (None, None):
-                                                row_idx = variable_to_index(n_param, (left_param_idx, right_param_idx))
+                                                col_idx = variable_to_index(n_param, (left_param_idx, right_param_idx))
+                                                # col_idx = (left_param_idx, right_param_idx)
                                                 degree = 2
                                                 value = factor_1*factor_2
+
+                                                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):
-                                                # row_idx = variable_to_index(n_param, (right_param_idx,))
-                                                row_idx = right_param_idx
+                                                col_idx = right_param_idx
+                                                # col_idx = (right_param_idx,)
                                                 degree = 1
                                                 value = subs_val_1*factor_2
 
+                                                terms[degree][(idx_eq, perm)][0] += value
+
                                             case (None, subs_val):
-                                                # row_idx = variable_to_index(n_param, (left_param_idx,))
-                                                row_idx = left_param_idx
+                                                col_idx = left_param_idx
+                                                # col_idx = (left_param_idx,)
                                                 degree = 1
                                                 value = subs_val_2*factor_1
                                             
                                             case _:
-                                                degree, row_idx, value = 0, 0, subs_val_1*subs_val_2
+                                                degree, col_idx, value = 0, tuple(), subs_val_1*subs_val_2*factor_1*factor_2
 
-                                        terms[degree][(idx_eq, perm)][row_idx] += value
+                                        terms[degree][(idx_eq, perm)][col_idx] += value
 
-        return terms, offset_dict
+        return terms, offset_dict, n_param
 
 ########################################
 # Classes
@@ -286,6 +290,7 @@ class PolyMatrixImpl(PolyMatrix):
     subs_func: typing.Callable[[int, int, int], tuple[float]]
     re_index: dict[int, dict[tuple[int, int], tuple[int, int, float]]]
     re_index_func: typing.Callable[[int, int], tuple[int, int, float]]
+    re_index_func_2: typing.Callable[[int, int, tuple[int, ...]], tuple[int, tuple[int, ...], float]]
     is_vector: bool
 
 
@@ -304,6 +309,7 @@ def init_poly_vector(
     subs_func: typing.Callable[[int, int, int], tuple[float]] = None,
     re_index: dict[int, dict[tuple[int, int], tuple[int, int, float]]] = None,
     re_index_func: typing.Callable[[int, int], tuple[int, int, float]] = None,
+    re_index_func_2: typing.Callable[[int, int], tuple[int, int, float]] = None,
 ):
     return init_poly_matrix(
         degrees=degrees,
@@ -312,6 +318,7 @@ def init_poly_vector(
         re_index=re_index,
         re_index_func=re_index_func,
         is_vector=True,
+        re_index_func_2=re_index_func_2,
     )
 
 
@@ -321,6 +328,7 @@ def init_poly_matrix(
     subs_func: typing.Callable[[int, int, int], tuple[float]] = None,
     re_index: dict[int, dict[tuple[int, int], tuple[int, int, float]]] = None,
     re_index_func: typing.Callable[[int, int], tuple[int, int, float]] = None,
+    re_index_func_2: typing.Callable[[int, int], tuple[int, int, float]] = None,
     is_vector: bool = None,
 ):
     if degrees is None:
@@ -342,6 +350,7 @@ def init_poly_matrix(
         subs_func=subs_func,
         re_index = re_index,
         re_index_func = re_index_func,
+        re_index_func_2 = re_index_func_2,
         is_vector = is_vector,
     )
 
diff --git a/test_polymatrix/test_polymatrix.py b/test_polymatrix/test_polymatrix.py
index 1f0bb73..3b239bf 100644
--- a/test_polymatrix/test_polymatrix.py
+++ b/test_polymatrix/test_polymatrix.py
@@ -135,12 +135,12 @@ class TestPolyMatrix(unittest.TestCase):
             value = 2,
         )
     
-    def test_param_matrix_const_vector_skew_symmetric(self):
+    def test_skew_symmetric_param_matrix_const_vector(self):
         """
         param = [a11 a21 a31 a41]
         """
 
-        def skew_symmetric(idx1, idx2):
+        def skew_symmetric(degree, idx1, idx2):
             if idx1 == idx2:
                 return idx1, idx2, 0
             elif idx2 < idx1:
@@ -173,4 +173,65 @@ class TestPolyMatrix(unittest.TestCase):
             eq_idx = 1,
             row_idx = offset_dict[(mat, 0)] + 1,
             value = -1,
-        )
\ No newline at end of file
+        )
+
+    def test_const_matrix_param_gradient_vector(self):
+        """
+        param = [v11 v12 v21 v22]
+        """
+
+        def gradient(degree, v_row, monom):
+            if degree == 1:
+                factor = sum(v_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 v_row <= monom[-1]:
+                    n_v_row = v_row
+                    n_monom = monom
+
+                return n_v_row, n_monom, factor
+
+        mat = init_poly_matrix(
+            subs={0: {(0, 0): 1, (1, 0): 1, (0, 1): 1, (1, 1): 1}}, 
+        )
+        vec = init_poly_vector(
+            degrees=(1,),
+            re_index_func_2=gradient,
+        )
+
+        eq = init_equation(
+            terms = [(mat, vec)],
+            n_var = 2,
+        )
+
+        terms, offset_dict = list(eq.create())
+
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 1,
+            monoms=(0,),
+            eq_idx = 0,
+            row_idx = offset_dict[(vec, 1)],
+            value = 2,
+        )
+
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 1,
+            monoms=(1,),
+            eq_idx = 0,
+            row_idx = offset_dict[(vec, 1)]+1,
+            value = 1,
+        )
+
+        self.assert_term_in_eq(
+            terms = terms,
+            degree = 1,
+            monoms=(0,),
+            eq_idx = 1,
+            row_idx = offset_dict[(vec, 1)]+1,
+            value = 1,
+        )