Upload run_full_experiment.py with huggingface_hub

run_full_experiment.py

@@ -0,0 +1,249 @@
+"""
+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)