add skew-symmetry and gradient function

1 files changed, 91 insertions, 33 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)