""" 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)