Comparison_extrapolation.py

import os
import re
import torch
import sympy as sp
import numpy as np
import multiprocessing
from EquationLearning.models.NNModel import NNModel
from EquationLearning.utils import get_project_root
from EquationLearning.Data.GenerateDatasets import DataLoader

np.random.seed(7)


def format_sympy_expr(expr_str):
    # try:
    # Convert to sympy expression
    exprr = sp.sympify(expr_str)

    # Convert variables like x0 to x_0
    exprr = exprr.replace(
        lambda x: isinstance(x, sp.Symbol) and re.match(r'[a-zA-Z]+\d+', str(x)),
        lambda x: sp.Symbol(re.sub(r'([a-zA-Z]+)(\d+)', r'\1_\2', str(x)))
    )

    # Convert to LaTeX string
    return f"${sp.latex(exprr, mode='plain', fold_short_frac=True, symbol_names={})}$"
    # except Exception as e:
    #     # Return raw string if parsing fails
    #     return expr_str


def round_expr(exprr):
    return exprr.xreplace({n: round(n, 3) for n in exprr.atoms(sp.Number)})


def worker(q, expr_str):
    try:
        result = sp.simplify(sp.sympify(expr_str))
        q.put(result)
    except Exception as e:
        q.put(e)


def load_NN(datasetName, data_loader):
    data = data_loader.dataset
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    folder = os.path.join(get_project_root(), "EquationLearning//saved_models//saved_NNs//" + datasetName)
    filepath = folder + "//weights-NN-" + datasetName
    nn_model = None
    if os.path.exists(filepath.replace("weights", "NNModel") + '.pth'):
        # If this file exists, it means we saved the whole model
        network = torch.load(filepath.replace("weights", "NNModel") + '.pth', map_location=device)
        nn_model = NNModel(device=device, n_features=data.n_features, loaded_NN=network)
    elif os.path.exists(filepath):
        # If this file exists, initiate a model and load the weigths
        nn_model = NNModel(device=device, n_features=data.n_features, NNtype=data_loader.modelType)
        nn_model.loadModel(filepath)
    return nn_model


def simplify_expr(model_st, timeout=30):

    q = multiprocessing.Queue()
    p = multiprocessing.Process(target=worker, args=(q, model_st))
    p.start()
    p.join(timeout)

    if p.is_alive():
        p.terminate()
        p.join()
        raise TimeoutError("Simplification exceeded time limit")

    result = q.get_nowait()
    if isinstance(result, Exception):
        raise result  # Raise the exception if any occurred during simplification

    return result


def evaluate(exp, Xs, Ys, xvars):
    exp = sp.sympify(exp)
    fs_lambda = sp.lambdify(sp.flatten(xvars), exp)
    if Xs.ndim > 1:
        ys = fs_lambda(*list(Xs.T))
    else:
        ys = fs_lambda(Xs)
    return np.mean((Ys - ys)**2)


if __name__ == "__main__":
    ####################################
    # Parameters
    ####################################
    names = ['E1', 'E2', 'E3', 'E4', 'E5', 'E6', 'E7', 'E8', 'E9', 'E10', 'E11', 'E12', 'E13', 'SB1', 'SB2', 'SB3', 'SB4']
    noise = 0
    iteration = 0
    noise_suff = '/Without-noise/'
    if noise > 0:
        noise_suff = '/Noise_' + str(noise) + '/'

    # Methods
    methods = ['NN', 'PYSR', 'TaylorGP', 'NESYMRES', 'E2E', 'TPSR', 'MST', 'uDSR', 'SETGAP']

    ####################################
    # Analyze one variable at a time
    ####################################
    mse_results = {name: {} for name in names}
    eq_results = {name: {} for name in names}

    for name in names:
        namex = name
        if name == 'E10':
            namex = 'CS1'
        elif name == 'E11':
            namex = 'CS2'
        elif name == 'E12':
            namex = 'CS3'
        elif name == 'E13':
            namex = 'CS4'
        ####################################
        # Load underlying equation
        ####################################
        dataLoader = DataLoader(name=namex, extrapolation=True)
        X, Y, var_names, expr = dataLoader.dataset.X, dataLoader.dataset.Y, dataLoader.dataset.names, dataLoader.dataset.expr
        print("Underlying function: " + str(expr))
        for im, method in enumerate(methods):
            if method != 'NN':
                # Load the expressions generated by each method
                path = str(get_project_root().parent) + "/output/LearnedEquations" + noise_suff + namex + '/' + method + ".txt"
                try:
                    with open(path, "r") as myfile:
                        expr_pred = myfile.read().splitlines()[iteration]
                        try:
                            expr_sympy = format_sympy_expr(round_expr(simplify_expr(sp.sympify(expr_pred))))
                        except:
                            expr_sympy = format_sympy_expr(round_expr((sp.sympify(expr_pred))))

                    # Evaluate original expression (this is expected to be 0, just to verify)
                    MSE_orig = evaluate(expr, X, Y, var_names)
                    # Evaluate predicted expression
                    mse_results[name][method] = evaluate(expr_pred, X, Y, var_names)
                    eq_results[name][method] = expr_sympy
                except FileNotFoundError:
                    mse_results[name][method] = '---'
                    eq_results[name][method] = "---"
            else:
                bb_model = load_NN(datasetName=namex, data_loader=dataLoader)
                ys = np.array(bb_model.evaluateFold(X, batch_size=X.shape[1]))[:, 0]
                mse_results[name][method] = np.mean((Y - ys) ** 2)
                eq_results[name][method] = "---"

    # Generate LaTeX table
    latex_table = "\\begin{table}[]\n\\resizebox{\\textwidth}{!}{%\n"
    latex_table += "\\begin{tabular}{|c|" + "c|" * len(methods) + "}\n\\hline\n"
    latex_table += "\\textbf{Problem} & " + " & ".join(methods) + " \\\\\n\\hline\n"

    for name in names:
        row = f"{name} & " + " & ".join(f"{mse_results[name][method]:.3e}"
                                        if isinstance(mse_results[name][method], float) else mse_results[name][method]
                                        for method in methods) + " \\\\\n\\hline\n"
        latex_table += row

    latex_table += "\\end{tabular}%\n}\n\\end{table}"

    print(latex_table)

    # Generate LaTeX table
    latex_table = "\\begin{table}[!ht]\n"
    latex_table += "\\caption{Comparison of predicted expressions --- Iteration 1}\n"
    latex_table += "\\label{app:tabcomparison}\n"
    latex_table += "\\centering\n"
    latex_table += "\\resizebox{\\textwidth}{!}{%\n"
    latex_table += "\\def\\arraystretch{1.45}\n"
    latex_table += "\\begin{tabular}{|c|" + "c|" * len(methods) + "}\n\\hline\n"
    latex_table += "\\textbf{Eq.} & " + " & ".join(f"\\textbf{{{m}}}" for m in methods) + " \\\\\n\\hline\n"

    for name in names:
        row = f"{name} & " + " & ".join(
            (eq_results[name][method])
            for method in methods
        ) + " \\\\\n\\hline\n"
        latex_table += row

    latex_table += "\\end{tabular}%\n}\n\\end{table}"

    print(latex_table)