summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--main.py124
-rw-r--r--polymatrix/polystruct.py199
-rw-r--r--test_polymatrix/test_polymatrix.py67
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,
+ )