update training

optuna_runs.py

@@ -0,0 +1,225 @@
+import keras
+from keras.layers import Dense, Dropout, Input,BatchNormalization
+import tensorflow as tf
+import h5py
+import numpy as np
+import pandas as pd
+import time
+import sklearn.model_selection as sk
+import matplotlib.pyplot as plt
+from sklearn.cluster import KMeans
+from sklearn.pipeline import Pipeline, make_pipeline
+from sklearn.preprocessing import StandardScaler, MinMaxScaler
+from imblearn.over_sampling import SMOTE
+from imblearn.under_sampling import RandomUnderSampler
+from imblearn.over_sampling import RandomOverSampler
+from collections import Counter
+import os
+from preprocessing import *
+from sklearn import set_config
+from importlib import reload
+set_config(transform_output = "pandas")
+
+dtype = "float32"
+activation = "relu"
+
+lr = 0.001
+batch_size = 512
+epochs = 50 # default 400 epochs
+
+lr_schedule = keras.optimizers.schedules.ExponentialDecay(
+    initial_learning_rate=lr,
+    decay_steps=2000,
+    decay_rate=0.9,
+    staircase=True
+)
+
+optimizer_simple = keras.optimizers.Adam(learning_rate=lr_schedule)
+optimizer_large = keras.optimizers.Adam(learning_rate=lr_schedule)
+optimizer_paper = keras.optimizers.Adam(learning_rate=lr_schedule)
+
+sample_fraction = 0.8
+
+# small model
+model_simple = keras.Sequential(
+    [
+        keras.Input(shape = (9,), dtype = "float32"),
+        keras.layers.Dense(units = 128, activation = "linear", dtype = "float32"),
+        # Dropout(0.2),
+        keras.layers.Dense(units = 128, activation = "elu", dtype = "float32"),
+        keras.layers.Dense(units = 9, dtype = "float32")
+    ]
+)
+
+def Safelog(val):
+    # get range of vector
+    if val > 0:
+        return np.log10(val)
+    elif val < 0:
+        return -np.log10(-val)
+    else:
+        return 0
+
+def Safeexp(val):
+    if val > 0:
+        return -10 ** -val
+    elif val < 0:
+        return 10 ** val
+    else:
+        return 0
+    
+# ? Why does the charge is using another logarithm than the other species
+
+func_dict_in = {
+    "H" : np.log1p,
+    "O" : np.log1p,
+    "Charge" : Safelog,
+    "H_0_" : np.log1p,
+    "O_0_" : np.log1p,
+    "Ba" : np.log1p,
+    "Cl" : np.log1p,
+    "S_2_" : np.log1p,
+    "S_6_" : np.log1p,
+    "Sr" : np.log1p,
+    "Barite" : np.log1p,
+    "Celestite" : np.log1p,
+}
+
+func_dict_out = {
+    "H" : np.expm1,
+    "O" : np.expm1,
+    "Charge" : Safeexp,
+    "H_0_" : np.expm1,
+    "O_0_" : np.expm1,
+    "Ba" : np.expm1,
+    "Cl" : np.expm1,
+    "S_2_" : np.expm1,
+    "S_6_" : np.expm1,
+    "Sr" : np.expm1,
+    "Barite" : np.expm1,
+    "Celestite" : np.expm1,
+}
+
+# os.chdir('/mnt/beegfs/home/signer/projects/model-training')
+data_file = h5py.File("barite_50_4_corner.h5")
+
+design = data_file["design"]
+results = data_file["result"]
+
+df_design = pd.DataFrame(np.array(design["data"]).transpose(), columns = np.array(design["names"].asstr()))
+df_results = pd.DataFrame(np.array(results["data"]).transpose(), columns = np.array(results["names"].asstr()))
+
+data_file.close()
+
+species_columns = ['H', 'O', 'Charge', 'Ba', 'Cl', 'S', 'Sr', 'Barite', 'Celestite']
+
+preprocess = preprocessing(func_dict_in=func_dict_in, func_dict_out=func_dict_out)
+X, y = preprocess.cluster(df_design[species_columns], df_results[species_columns])
+# X, y = preprocess.funcTranform(X, y)
+
+X_train, X_test, y_train, y_test = preprocess.split(X, y, ratio = 0.2)
+X_train, y_train = preprocess.balancer(X_train, y_train, strategy = "over")
+preprocess.scale_fit(X_train, y_train, scaling = "individual")
+X_train, X_test, y_train, y_test = preprocess.scale_transform(X_train, X_test, y_train, y_test)
+X_train, X_val, y_train, y_val = preprocess.split(X_train, y_train, ratio = 0.1)
+
+column_dict = {"Ba": X.columns.get_loc("Ba"), "Barite":X.columns.get_loc("Barite"), "Sr":X.columns.get_loc("Sr"), "Celestite":X.columns.get_loc("Celestite"), "H":X.columns.get_loc("H"), "H":X.columns.get_loc("H"), "O":X.columns.get_loc("O")}
+
+def custom_loss(preprocess, column_dict, h1, h2, h3, h4):
+    # extract the scaling parameters
+    scale_X = tf.convert_to_tensor(preprocess.scaler_X.scale_, dtype=tf.float32)
+    min_X = tf.convert_to_tensor(preprocess.scaler_X.min_, dtype=tf.float32)
+    scale_y = tf.convert_to_tensor(preprocess.scaler_y.scale_, dtype=tf.float32)
+    min_y = tf.convert_to_tensor(preprocess.scaler_y.min_, dtype=tf.float32)
+
+    def loss(results, predicted):
+        # inverse min/max scaling
+        predicted_inverse = predicted * scale_X + min_X
+        results_inverse = results * scale_y + min_y
+
+        # mass balance
+        dBa = tf.keras.backend.abs(
+            (predicted_inverse[:, column_dict["Ba"]] + predicted_inverse[:, column_dict["Barite"]]) -
+            (results_inverse[:, column_dict["Ba"]] + results_inverse[:, column_dict["Barite"]])
+        )
+        dSr = tf.keras.backend.abs(
+            (predicted_inverse[:, column_dict["Sr"]] + predicted_inverse[:, column_dict["Celestite"]]) -
+            (results_inverse[:, column_dict["Sr"]] + results_inverse[:, column_dict["Celestite"]])
+        )
+        
+        # H/O ratio has to be 2
+        h2o_ratio = tf.keras.backend.abs(
+            (predicted_inverse[:, column_dict["H"]] / predicted_inverse[:, column_dict["O"]]) - 2
+        )
+
+        # huber loss
+        huber_loss = tf.keras.losses.Huber()(results, predicted)
+        
+        # total loss
+        total_loss = h1 * huber_loss + h2 * dBa**2 + h3 * dSr**2 #+ h4 * h2o_ratio**2
+
+        return total_loss
+
+    return loss
+
+def mass_balance(model, X, preprocess):
+    
+    # predict the chemistry
+    columns = X.iloc[:, X.columns != "Class"].columns
+    prediction = pd.DataFrame(model.predict(X[columns]), columns=columns)
+    
+    # backtransform min/max
+    X = pd.DataFrame(preprocess.scaler_X.inverse_transform(X.iloc[:, X.columns != "Class"]), columns=columns)
+    prediction = pd.DataFrame(preprocess.scaler_y.inverse_transform(prediction), columns=columns)
+        
+    # calculate mass balance    dBa = np.abs((prediction["Ba"] + prediction["Barite"]) - (X["Ba"] + X["Barite"]))
+    dSr = np.abs((prediction["Sr"] + prediction["Celestite"]) - (X["Sr"] + X["Celestite"]))
+    
+    return dBa + dSr
+
+import optuna
+
+def create_model(model, preprocess, h1, h2, h3, h4):
+    
+    model.compile(optimizer=optimizer_simple, loss=custom_loss(preprocess, column_dict, h1, h2, h3, h4))
+    
+    return model
+
+
+def objective(trial, preprocess, X_train, y_train, X_val, y_val, X_test, y_test):
+    h1 = trial.suggest_float("h1", 0.1, 10)
+    h2 = trial.suggest_float("h2", 0.1, 10)
+    h3 = trial.suggest_float("h3", 0.1, 10)
+    h4 = trial.suggest_float("h4", 0.1, 10)
+    
+    model = create_model(model_simple, preprocess, h1, h2, h3, h4)
+    
+    callback = keras.callbacks.EarlyStopping(monitor='loss', patience=3)
+    history = model.fit(X_train.loc[:, X_train.columns != "Class"], 
+                        y_train.loc[:, y_train.columns != "Class"], 
+                        batch_size=batch_size, 
+                        epochs=50, 
+                        validation_data=(X_val.loc[:, X_val.columns != "Class"], y_val.loc[:, y_val.columns != "Class"]),
+                        callbacks=[callback])
+    
+    prediction_loss = model.evaluate(X_test.loc[:, X_test.columns != "Class"], y_test.loc[:, y_test.columns != "Class"])
+    mass_balance_results = mass_balance(model, X_test, preprocess)
+    
+    mass_balance_ratio = len(mass_balance_results[mass_balance_results < 1e-5]) / len(mass_balance_results)
+
+    return prediction_loss, mass_balance_ratio
+
+if __name__ == "__main__":
+    study  = optuna.create_study(storage="sqlite:///model_optimization.db", study_name="model_optimization", directions=["minimize", "maximize"])
+    study.optimize(lambda trial: objective(trial, preprocess, X_train, y_train, X_val, y_val, X_test, y_test), n_trials=1000)
+
+    print("Number of finished trials: ", len(study.trials))
+
+    print("Best trial:")
+    trial = study.best_trial
+
+    print("  Value: ", trial.value)
+
+    print("  Params: ")
+    for key, value in trial.params.items():
+        print("    {}: {}".format(key, value))