Merge pull request #23 from wkchuang/wchuang

Update for incloud tendency and activation function

config/cesm_tau_run9_process.yml

@@ -0,0 +1,16 @@
+model_path: "/glade/work/wchuang/mlmicrophysics/cesm_output/tauREV4/"
+model_file_start: "cam_ml_tauREV4_train.cam.h1"
+model_file_end: "nc"
+time_var: "time"
+time_split_interval: 1
+staggered_variables: []
+out_variables: ["depth", "row", "col", "T", "RHO_CLUBB", 
+                "CLOUD", "FREQR",
+                "QC_TAU_in", "NC_TAU_in", "QR_TAU_in", "NR_TAU_in", 
+                "QC_TAU_out", "NC_TAU_out", "QR_TAU_out", "NR_TAU_out",
+                "qctend_TAU", "nctend_TAU", "qrtend_TAU", "nrtend_TAU", "PGAM", "LAMC", "LAMR", "N0R"]
+subset_variable: ["QC_TAU_in"]
+subset_threshold: [1.0e-8]
+out_path: "/glade/derecho/scratch/wchuang/mlmicrophysics/cam_ml_tauREV4/20240110/"
+out_start: "cam_mp_data_tauREV4"
+out_format: "parquet"

config/training_tau_run10.yml

@@ -0,0 +1,53 @@
+data:
+  data_path: "/glade/derecho/scratch/wchuang/mlmicrophysics/cam_ml_tauREV4"
+  scratch_path: "/glade/derecho/scratch/wchuang/mlmicrophysics/tau_run_10/20240206"
+  out_path: "/glade/work/wchuang/mlmicrophysics/tau_run_10/20240206"
+  input_cols: ["QC_TAU_in", "QR_TAU_in", "NC_TAU_in", "NR_TAU_in", "PGAM", "LAMC", "LAMR", "N0R", "RHO_CLUBB", "CLOUD", "FREQR"]
+  output_cols: ["qctend_TAU", "qrtend_TAU", "nctend_TAU", "nrtend_TAU"]
+  qc_thresh: 1e-8
+  n_quantiles: 2500 
+  subsample: 0.01
+  random_seed: 215689
+  subset_data:
+    train_date_start: 0
+    train_date_end: 12500
+    test_date_start: 12501
+    test_date_end: 18000
+    validation_frequency: 3
+
+model:
+  hidden_layers: 1
+  hidden_neurons: 1500 
+  activation: "relu"
+  output_activation: "sigmoid"
+  loss: "mse"
+  lr: 0.0007
+  batch_size: 1024
+  epochs: 50 # default 100
+  verbose: 2
+
+callbacks:
+  EarlyStopping:
+    monitor: "loss"
+    patience: 5
+    mode: "min"
+    verbose: 0
+  ReduceLROnPlateau: 
+    monitor: "loss"
+    factor: 0.1
+    patience: 2
+    min_lr: 1.0e-12
+    min_delta: 1.0e-08
+    mode: "min"
+    verbose: 0
+  CSVLogger:
+    filename: "training_log.csv"
+    separator: ","
+    append: False
+  ModelCheckpoint:
+    filepath: "model.h5"
+    monitor: "loss"
+    save_weights: True
+    save_best_only: True
+    mode: "min"
+    verbose: 0

mlmicrophysics/data.py

@@ -96,7 +96,6 @@ def unstagger_vertical(dataset, variable, vertical_dim="lev"):
         xarray DataArray containing the vertically interpolated data
     """
     var_data = dataset[variable]
-    print("variable is ", variable)
     unstaggered_var_data = xr.DataArray(0.5 * (var_data[:, :-1].values + var_data[:, 1:].values),
                                         coords=[var_data.time, dataset[vertical_dim], var_data.lat, var_data.lon],
                                         dims=("time", vertical_dim, "lat", "lon"),

mlmicrophysics/models.py

@@ -1,5 +1,5 @@
 import tensorflow.keras.backend as K
-K.set_floatx('float64') # added by wkc to get float64
+K.set_floatx('float64') # for float64 capability
 from tensorflow.keras.layers import Input, Dense, Dropout, GaussianNoise, Activation, Concatenate, BatchNormalization, LeakyReLU
 from tensorflow.keras.models import Model
 from tensorflow.keras.regularizers import l2
@@ -11,8 +11,8 @@
 import xarray as xr
 import pandas as pd
 
-policy = tf.keras.mixed_precision.Policy("float64") # added by wkc to get float64
-tf.keras.mixed_precision.set_global_policy(policy) # added by wkc to get float64
+policy = tf.keras.mixed_precision.Policy("float64") # for float64 capability
+tf.keras.mixed_precision.set_global_policy(policy) # for float64 capability
 
 metrics_dict = {"accuracy": accuracy_score,
                  "heidke": heidke_skill_score,
@@ -95,14 +95,10 @@ def build_neural_network(self, inputs, outputs):
 
         for h in range(self.hidden_layers):
             nn_model = Dense(self.hidden_neurons,
-                             kernel_regularizer=l2(self.l2_weight),
-                             name=f"dense_{h:02d}",
-                             dtype=tf.float64)(nn_model)
-            if self.activation == "leaky":
-                nn_model = LeakyReLU(self.leaky_alpha, name="hidden_activation_{0:02d}".format(h), dtype=tf.float64)(nn_model)
-            else:
-                nn_model = Activation(self.activation, name="hidden_activation_{0:02d}".format(h), dtype=tf.float64)(nn_model)
-
+                    activation=self.activation,
+                    kernel_regularizer=l2(self.l2_weight),
+                    name=f"dense_{h:02d}",
+                    dtype=tf.float64)(nn_model)
             if self.use_dropout:
                 nn_model = Dropout(self.dropout_alpha, name=f"dropout_h_{h:02d}", dtype=tf.float64)(nn_model)
             if self.use_noise:

mlmicrophysics/objective.py

@@ -1,6 +1,7 @@
 from echo.src.base_objective import BaseObjective
 import numpy as np
 from mlmicrophysics.models import DenseNeuralNetwork
+from mlmicrophysics.callbacks import get_callbacks
 import pandas as pd
 from os.path import join
 import pickle
@@ -25,9 +26,10 @@ def train(self, trial, conf):
         with open(join(conf["data"]["out_path"], "output_quantile_transform.pkl"), "rb") as out_scaler_pickle:
             output_scaler = pickle.load(out_scaler_pickle)
         dnn = DenseNeuralNetwork(**conf["model"])
-        dnn.fit(input_quant_data["train"], output_quant_data["train"])
+        dnn.fit(input_quant_data["train"], output_quant_data["train"], callbacks=get_callbacks(conf, path_extend = conf["data"]["out_path"]))
         val_quant_preds = dnn.predict(input_quant_data["val"], batch_size=40000)
         val_preds = output_scaler.inverse_transform(val_quant_preds)
-        val_r2 = r2_score(np.log10(output_data["val"]), np.log10(val_preds))
-        results_dict = {"val_loss": val_r2}
+        val_r2 = r2_score(output_data["val"], val_preds)
+        val_quant_r2 = r2_score(output_quant_data["val"] , val_quant_preds) # r2 in quant space should be more appropriate
+        results_dict = {"val_loss": val_quant_r2} # add rmse, mae 
         return results_dict