run_full_experiment.py

"""
Full EQL two-phase training experiment.
Usage:
    python run_full_experiment.py --city Roskilde --steps_ahead 6 --feature wind_speed
"""
import os, sys, time
import numpy as np
from collections import OrderedDict
from tqdm import tqdm
from scipy.io import loadmat
import tensorflow as tf
import h5py
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt

from reproduce_eql import (
    SymbolicNet, Constant, Identity, Square, Sin, Sigmoid, Product,
    count_double, l12_smooth, get_folders_started, record_base_info,
    append_text_to_summary, plot_train_vs_validation, plot_histogram,
    plot_descaled_real_vs_prediction, generate_all_data, generate_variable_list,
    save_weights, network
)


def run_experiment(target_city, steps_ahead=6, feature="wind_speed"):
    target_cities = ["Esbjerg", "Odense", "Roskilde"]
    target_city_index = target_cities.index(target_city)

    filename = "{}/step{}.mat".format(feature, {6: "1", 12: "2", 18: "3", 24: "4"}[steps_ahead])
    data = loadmat('Denmark_data/{}'.format(filename))
    y_min = data["y_min_tr"][0][target_city_index]
    y_max = data["y_max_tr"][0][target_city_index]

    city_configs = {
        "Esbjerg": {"lambda_reg": 5.0, "a": 5e-3, "threshold": 8.0e-3},
        "Odense": {"lambda_reg": 5.0, "a": 5e-4, "threshold": 7.5e-3},
        "Roskilde": {"lambda_reg": 3.0, "a": 5e-3, "threshold": 7.5e-3},
    }
    city_cfg = city_configs[target_city]

    config = {
        "use_rescaled_MSE": True,
        "a_L_0.5": city_cfg["a"],
        "threshold_value": city_cfg["threshold"],
        "lambda_reg": city_cfg["lambda_reg"],
        "steps_ahead": steps_ahead,
        "feature": feature,
        "target_city": target_city,
        "epochs1": 100,
        "epochs2": 100,
        "use_phase2": True,
        "batch_size": 200,
        "phase1_lr": 1e-4,
        "phase2_lr": 1e-5,
        "eql_number_layers": 2,
        "optimizer": "rmsprop",
    }

    use_rescaled_loss = config["use_rescaled_MSE"]
    number_epochs1 = config["epochs1"]
    number_epochs2 = config["epochs2"]
    threshold_value = config["threshold_value"]
    a_L_05 = config["a_L_0.5"]
    lambda_reg = config["lambda_reg"]
    batch_size = config["batch_size"]
    first_phase_lr = config["phase1_lr"]
    second_phase_lr = config["phase2_lr"]

    x_dim = 80
    activation_funcs = [
        *[Constant()] * 2,
        *[Identity()] * 4,
        *[Square()] * 4,
        *[Sin()] * 2,
        *[Sigmoid()] * 2,
        *[Product()] * 2
    ]
    n_layers = 2
    n_double = count_double(activation_funcs)
    width = len(activation_funcs)

    init_weights = [
        tf.random.truncated_normal([x_dim, width + n_double], stddev=0.1),
        tf.random.truncated_normal([width, width + n_double], stddev=0.5),
        tf.random.truncated_normal([width, 1], stddev=1.0)
    ]
    model = SymbolicNet(n_layers, funcs=activation_funcs, initial_weights=init_weights)
    phase1_optimizer = tf.keras.optimizers.RMSprop(learning_rate=first_phase_lr)

    experiment_number = get_folders_started()
    print("\n\n" + "*" * 10 + " Beginning of Experiment {} ".format(experiment_number) + "*" * 10 + "\n\n")
    print("Feature : {}, City : {}, Steps ahead : {}".format(feature, target_city, steps_ahead))
    record_base_info(experiment_number, **config)

    x_train, y_train = generate_all_data("train", steps_ahead, feature, target_city_index)
    x_test, y_test = generate_all_data("test", steps_ahead, feature, target_city_index)
    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)
    val_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size)

    # Phase 1
    best_val_loss = float('inf')
    best_val_weights = None
    train_loss_results = []
    valid_loss_results = []
    count_loss_stagnation = 0

    for epoch in range(number_epochs1):
        epoch_loss_avg = tf.keras.metrics.MeanSquaredError()
        for x_batch, y_batch in tqdm(train_dataset, desc=f"P1 Epoch {epoch+1}", leave=False):
            with tf.GradientTape() as tape:
                y_pred = model(x_batch)
                if use_rescaled_loss:
                    y_pred_r = y_pred * (y_max - y_min) + y_min
                    y_real_r = y_batch * (y_max - y_min) + y_min
                    error = tf.keras.losses.MeanSquaredError()(y_real_r, y_pred_r)
                else:
                    error = tf.keras.losses.MeanSquaredError()(y_batch, y_pred)
                reg_loss = l12_smooth(model.get_weights(), a_L_05)
                loss = error + lambda_reg * reg_loss
            grads = tape.gradient(loss, model.get_weights())
            phase1_optimizer.apply_gradients(zip(grads, model.get_weights()))
            epoch_loss_avg.update_state(y_batch, y_pred)

        train_mse = epoch_loss_avg.result().numpy()
        train_loss_results.append(train_mse)

        val_loss_avg = tf.keras.metrics.MeanSquaredError()
        for x_val, y_val in val_dataset:
            y_pred_val = model(x_val)
            val_loss_avg.update_state(y_val, y_pred_val)
        val_mse = val_loss_avg.result().numpy()
        valid_loss_results.append(val_mse)

        if val_mse < best_val_loss:
            best_val_loss = val_mse
            best_val_weights = [w.numpy() for w in model.get_weights()]
            print(f"P1 Epoch {epoch+1}: val MSE improved to {val_mse:.4e} (train {train_mse:.4e})")
            count_loss_stagnation = 0
        else:
            count_loss_stagnation += 1
            print(f"P1 Epoch {epoch+1}: val MSE {val_mse:.4e} (train {train_mse:.4e})")

        if count_loss_stagnation >= 4 and epoch > 70:
            print("No improvement over 4 epochs, epoch > 70. Exiting phase 1...")
            break

    save_weights(best_val_weights, experiment_number, "phase1", best_val_loss)
    append_text_to_summary(experiment_number, "phase 1 best MSE validation: {}\n".format(best_val_loss))
    plot_train_vs_validation(experiment_number, number_epochs1, train_loss_results, valid_loss_results, "phase1")
    plot_histogram(experiment_number, best_val_weights[0], "phase1", "weights1", a_L_05)
    plot_histogram(experiment_number, best_val_weights[1], "phase1", "weights2", a_L_05)
    plot_histogram(experiment_number, best_val_weights[2], "phase1", "weights3", a_L_05)

    # Phase 2
    if not config["use_phase2"]:
        return experiment_number, best_val_loss, None

    print("\n--- Starting Phase 2 ---\n")
    masked_weights = []
    for w_i in best_val_weights:
        mask = tf.cast(tf.constant(tf.abs(w_i) > threshold_value), tf.float32)
        masked_weights.append(tf.multiply(w_i, mask))

    sparsity1 = 100 * np.count_nonzero(masked_weights[0] == 0) / masked_weights[0].size
    sparsity2 = 100 * np.count_nonzero(masked_weights[1] == 0) / masked_weights[1].size
    sparsity3 = 100 * np.count_nonzero(masked_weights[2] == 0) / masked_weights[2].size
    append_text_to_summary(experiment_number, f"sparsity1 phase2 after masking: {sparsity1}\n")
    append_text_to_summary(experiment_number, f"sparsity2 phase2 after masking: {sparsity2}\n")
    append_text_to_summary(experiment_number, f"sparsity3 phase2 after masking: {sparsity3}\n")

    masked_model = SymbolicNet(n_layers, funcs=activation_funcs, initial_weights=masked_weights)
    phase2_optimizer = tf.keras.optimizers.RMSprop(learning_rate=second_phase_lr)

    train_loss_results2 = []
    valid_loss_results2 = []
    best_val_loss2 = float('inf')
    best_val_weights2 = None
    count_loss_stagnation = 0

    for epoch in range(number_epochs2):
        epoch_loss_avg = tf.keras.metrics.MeanSquaredError()
        for x_batch, y_batch in tqdm(train_dataset, desc=f"P2 Epoch {epoch+1}", leave=False):
            with tf.GradientTape() as tape:
                y_pred = masked_model(x_batch)
                if use_rescaled_loss:
                    y_pred_r = y_pred * (y_max - y_min) + y_min
                    y_real_r = y_batch * (y_max - y_min) + y_min
                    error = tf.keras.losses.MeanSquaredError()(y_real_r, y_pred_r)
                else:
                    error = tf.keras.losses.MeanSquaredError()(y_batch, y_pred)
                loss = error
            grads = tape.gradient(loss, masked_model.get_weights())
            phase2_optimizer.apply_gradients(zip(grads, masked_model.get_weights()))
            epoch_loss_avg.update_state(y_batch, y_pred)

        train_mse = epoch_loss_avg.result().numpy()
        train_loss_results2.append(train_mse)

        val_loss_avg = tf.keras.metrics.MeanSquaredError()
        for x_val, y_val in val_dataset:
            y_pred_val = masked_model(x_val)
            val_loss_avg.update_state(y_val, y_pred_val)
        val_mse = val_loss_avg.result().numpy()
        valid_loss_results2.append(val_mse)

        if val_mse < best_val_loss2:
            best_val_loss2 = val_mse
            best_val_weights2 = [w.numpy() for w in masked_model.get_weights()]
            print(f"P2 Epoch {epoch+1}: val MSE improved to {val_mse:.4e} (train {train_mse:.4e})")
            count_loss_stagnation = 0
        else:
            count_loss_stagnation += 1
            print(f"P2 Epoch {epoch+1}: val MSE {val_mse:.4e} (train {train_mse:.4e})")

        if count_loss_stagnation >= 4 and epoch > 80:
            print("No improvement over 4 epochs, epoch > 80. Exiting phase 2...")
            break

    save_weights(best_val_weights2, experiment_number, "phase2", best_val_loss2)
    append_text_to_summary(experiment_number, "phase 2 best MSE validation: {}\n".format(best_val_loss2))
    plot_train_vs_validation(experiment_number, number_epochs2, train_loss_results2, valid_loss_results2, "phase2")

    # Extract formula
    var_names = generate_variable_list(26, 80)
    expr = network(best_val_weights2, activation_funcs, var_names[:x_dim], threshold=0)
    print("\nExtracted formula:\n", expr)
    append_text_to_summary(experiment_number, "Formula after phase2: {}\n".format(expr))

    # Final prediction plot
    best_model = SymbolicNet(n_layers, funcs=activation_funcs,
                             initial_weights=[tf.constant(w) for w in best_val_weights2])
    y_predicted = best_model(x_test)
    mae = plot_descaled_real_vs_prediction(experiment_number, y_test, y_predicted, y_min, y_max)
    append_text_to_summary(experiment_number, "MAE: {}\n".format(mae))
    print("\nMAE: {}".format(mae))

    print("\n\n" + "*" * 10 + " End of Experiment {} ".format(experiment_number) + "*" * 10 + "\n\n")
    return experiment_number, best_val_loss, best_val_loss2


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--city", type=str, default="Roskilde", choices=["Esbjerg", "Odense", "Roskilde"])
    parser.add_argument("--steps_ahead", type=int, default=6)
    parser.add_argument("--feature", type=str, default="wind_speed")
    args = parser.parse_args()
    run_experiment(args.city, args.steps_ahead, args.feature)