Select best fitting sparsity for SymNet and print debug info on it

.gitignore

@@ -3,9 +3,14 @@ __pycache__/
 *.py[cod]
 *$py.class
 
-# C extensions
-*.so
+# Folders
+classic_alg_exp/
 
+# Extensions
+*.so
+*.png
+*.csv
+*.docx
 # Distribution / packaging
 .Python
 pyvenv.cfg

symnet/initcoefficients.py

@@ -8,6 +8,7 @@
 
 import seaborn as sns
 import matplotlib.pyplot as plt
+from symnet.print_selection_info import ModelsInfo
 
 
 def clean_names(left_name, names: list):
@@ -33,23 +34,83 @@ def closure():
     lbfgs = torch.optim.LBFGS(params=model.parameters(), max_iter=2000, line_search_fn='strong_wolfe')
     model.train()
     lbfgs.step(closure)
-    last_step_loss = loss(model, y_train, x_train, block=1, sparsity=sparsity)
+    last_step_loss = loss(model, y_train, x_train, block=1, sparsity=0.0)
 
     return model, last_step_loss.item()
 
 
-def select_model(input_names, left_pool, u, derivs, shape, sparsity, additional_tokens):
-    models = []
-    losses = []
+
+
+
+
+
+# def right_matrices_coef(matrices, names: list[str], csym, tsym):
+#     token_matrix = {}
+#     for i in range(len(names)):
+#         token_matrix[Symbol(names[i])] = matrices[i]
+#
+#     right_side = []
+#     for i in range(len(csym)):
+#         total_mx = 1
+#         if type(tsym[i]) == Mul:
+#             if tsym[i] == Mul(Symbol("u"), Symbol("du/dx2")):
+#                 u_ux_ind = i
+#             lbls = tsym[i].args
+#             for lbl in lbls:
+#                 if type(lbl) == Symbol:
+#                     total_mx *= token_matrix.get(lbl)
+#                 else:
+#                     for j in range(lbl.args[1]):
+#                         total_mx *= token_matrix.get(lbl.args[0])
+#         elif type(tsym[i]) == Symbol:
+#             total_mx *= token_matrix.get(tsym[i])
+#         elif type(tsym[i]) == Pow:
+#             for j in range(tsym[i].args[1]):
+#                 total_mx *= token_matrix.get(tsym[i].args[0])
+#         total_mx *= csym[i]
+#         right_side.append(total_mx)
+#
+#     u_ux = 1
+#     for lbl in (Symbol("u"), Symbol("du/dx2")):
+#         u_ux *= token_matrix.get(lbl)
+#     right_u_ux = csym[u_ux_ind] * u_ux
+#     diff1 = np.fabs((np.abs(csym[u_ux_ind]) - 1) * u_ux)
+#     return right_side, right_u_ux, u_ux
+
+
+# def select_model1(input_names, left_pool, u, derivs, shape, sparsity, additional_tokens):
+#     models = []
+#     losses = []
+#     for left_side_name in left_pool:
+#         m_input_names, idx = clean_names(left_side_name, input_names)
+#         x_train, y_train = prepare_batches(u, derivs, shape, idx, additional_tokens=additional_tokens)
+#         model, last_loss = train_model(m_input_names, x_train, y_train, sparsity)
+#
+#         tsym, csym = model.coeffs(calprec=16)
+#         losses.append(last_loss)
+#         models.append(model)
+#
+#     idx = losses.index(min(losses))
+#     return models[idx], left_pool[idx]
+
+def select_model(input_names, left_pool, u, derivs, shape, additional_tokens):
+    models, losses, left_sides = [], [], []
+    info = ModelsInfo()
     for left_side_name in left_pool:
-        m_input_names, idx = clean_names(left_side_name, input_names)
-        x_train, y_train = prepare_batches(u, derivs, shape, idx, additional_tokens=additional_tokens)
-        model, last_loss = train_model(m_input_names, x_train, y_train, sparsity)
-        losses.append(last_loss)
-        models.append(model)
+        for sparsity in [0.001, 0.0000001]:
+            m_input_names, idx = clean_names(left_side_name, input_names)
+            x_train, y_train = prepare_batches(u, derivs, shape, idx, additional_tokens=additional_tokens)
+            model, last_loss = train_model(m_input_names, x_train, y_train, sparsity)
+
+            losses.append(last_loss)
+            models.append(model)
+            left_sides.append(left_side_name)
+
+            info.selection_info(model, last_loss, sparsity, left_side_name)
 
+    info.print_best()
     idx = losses.index(min(losses))
-    return models[idx], left_pool[idx]
+    return models[idx], left_sides[idx]
 
 
 def save_fig(csym, add_left=True):
@@ -68,7 +129,7 @@ def save_fig(csym, add_left=True):
     plt.savefig(f'symnet_distr{len(distr)}.png', transparent=True)
 
 
-def get_csym_tsym(u, derivs, shape, input_names, pool_names, sparsity=0.001, additional_tokens=None,
+def get_csym_tsym(u, derivs, shape, input_names, pool_names, additional_tokens=None,
                   max_deriv_order=None):
     """
     Can process only one variable! (u)

symnet/print_selection_info.py

@@ -0,0 +1,104 @@
+import numpy as np
+from decimal import Decimal
+import os
+import sys
+
+# def print_best_info(sparsity, left_side_name="du/dx1")
+class ModelsInfo():
+    def __init__(self):
+        print("\n")
+        self.sparsity_ls = []
+        self.eq_type = []
+        self.pos_ls = []
+        self.coef_ls = []
+        self.losses = []
+
+
+    def selection_info(self, model, final_loss, sparsity, left_side_name="du/dx1"):
+        tsym, csym = model.coeffs(calprec=16)
+        left_names, right_names, coeffs = init_ideal_coeffs()
+        mae, final_places = mae_and_positions(right_names, coeffs, tsym, csym, left_side_name, left_names)
+        shd = calc_shd(final_places, csym, len(right_names))
+
+        final_coefs = [csym[i] for i in final_places]
+        coef_str = ['%.2e' % Decimal(c) for c in final_coefs]
+
+        if left_side_name == "du/dx1":
+            time_der = "u_t"
+        else:
+            time_der = "u_tt"
+
+        print(f"Sparsity & time der.: {sparsity: .1e}, {time_der}; MAE: {mae:.3f}; idxs: {final_places}; coefs: {coef_str}; "
+              f"loss: {final_loss: .2e}; shd: {shd}")
+
+        self.sparsity_ls. append(sparsity)
+        self.eq_type.append(left_side_name)
+        self.pos_ls.append(final_places)
+        self.coef_ls.append(coef_str)
+        self.losses.append(final_loss)
+
+    def print_best(self):
+        idx = self.losses.index(min(self.losses))
+        print(f"\n Lambda & t.d.: {self.sparsity_ls[idx]: .1e}, {self.eq_type[idx]}; idxs: {self.pos_ls[idx]}; "
+              f"coefs: {self.coef_ls[idx]}; loss: {self.losses[idx]: .2e}")
+        print("\n")
+
+
+def mae_and_positions(right_side_name, coefficients, tsym, csym, left_calc, left_true):
+    MAE = 0
+    final_places_ls = []
+    if left_calc != left_true:
+        MAE += 2.
+    for j in range(len(tsym)):
+        in_equation = False
+        for i in range(len(right_side_name)):
+            if str(tsym[j]) == right_side_name[i]:
+                MAE += np.fabs(coefficients[i] - csym[j])
+                in_equation = True
+                final_places_ls.append(j)
+                break
+        if not in_equation:
+            MAE += np.fabs(0. - csym[j])
+    MAE = MAE / (len(tsym) + 1)
+    return MAE, final_places_ls
+
+def init_ideal_coeffs():
+    name = (sys.argv[0]).split("/")[-1]
+    if name == 'noised_wave.py' or "experiment_wave.py":
+        coeffs = [0.04]
+        right_names = ['d^2u/dx2^2']
+        left_names = "d^2u/dx1^2"
+    elif name == 'noised_burgers_sindy.py': # 256 x 101
+        coeffs = [-1., 0.1]
+        right_names = ['du/dx2*u', 'd^2u/dx2^2']
+        left_names = "du/dx1"
+    elif name == 'noised_kdv.py':
+        left_names = "du/dx1"
+        coeffs = [6., 1., 1.]
+        right_names = ['du/dx2*u', 'd^3u/dx2^3', 'cos(t)sin(x)']
+    elif name == 'noised_kdv_sindy.py': # 512 x 201
+        left_names = "du/dx1"
+        coeffs = [-6., -1.]
+        right_names = ['du/dx2*u', 'd^3u/dx2^3']
+    elif name == "noised_burgers.py":
+        coeffs = [-1.]
+        right_names = ['du/dx2*u']
+        left_names = "du/dx1"
+    else:
+        raise NameError('Wrong type of equation')
+    return left_names, right_names, coeffs
+
+
+def calc_shd(idxs: list, coefs: list, true_num_terms: int):
+    idxs.sort()
+    shd = true_num_terms - len(idxs)
+
+    set_idxs_all = set([i for i in range(len(coefs))])
+    idxs_cut = list(set_idxs_all.difference(set(idxs)))
+    for i in idxs_cut:
+        if np.abs(coefs[i]) > 1e-6:
+            shd += 1
+    for i in idxs:
+        if np.abs(coefs[i]) < 1e-6:
+            shd += 1
+    return shd