Update names of normalization transforms (#402)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: RemiLehe

ml/train_model.py

@@ -99,11 +99,11 @@ def connect_to_db(config_dict):
     )[db_name]
 
 
-def normalize(df, input_names, input_transform, output_names, output_transform):
+def normalize(df, input_names, input_normalization, output_names, output_normalization):
     # Apply normalization to the training data set
     norm_df = df.copy()
-    norm_df[input_names] = input_transform(torch.tensor(df[input_names].values))
-    norm_df[output_names] = output_transform(torch.tensor(df[output_names].values))
+    norm_df[input_names] = input_normalization(torch.tensor(df[input_names].values))
+    norm_df[output_names] = output_normalization(torch.tensor(df[output_names].values))
     return norm_df
 
 
@@ -130,15 +130,17 @@ def split_data(df_exp, df_sim, variables, model_type):
             return (sim_train_df[variables], sim_val_df[variables])
 
 
-def build_transforms(n_inputs, X_train, n_outputs, y_train):
-    input_transform = AffineInputTransform(
+def build_normalizations(n_inputs, X_train, n_outputs, y_train):
+    input_normalization = AffineInputTransform(
         n_inputs, coefficient=X_train.std(axis=0), offset=X_train.mean(axis=0)
     )
     # For output normalization, we need to handle potential NaN values
     y_mean = torch.nanmean(y_train, dim=0)
     y_std = torch.sqrt(torch.nanmean((y_train - y_mean) ** 2, dim=0))
-    output_transform = AffineInputTransform(n_outputs, coefficient=y_std, offset=y_mean)
-    return input_transform, output_transform
+    output_normalization = AffineInputTransform(
+        n_outputs, coefficient=y_std, offset=y_mean
+    )
+    return input_normalization, output_normalization
 
 
 def train_nn_ensemble(
@@ -218,8 +220,8 @@ def build_lume_model(
     model_type,
     input_variables,
     output_variables,
-    input_transform,
-    output_transform,
+    input_normalization,
+    output_normalization,
 ):
     # Fix mismatch in name between the config file and the expected lume-model format
     for k in input_variables:
@@ -245,8 +247,8 @@ def build_lume_model(
             model=model.cpu(),
             input_variables=input_vars,
             output_variables=distribution_output_vars,
-            input_transformers=[input_transform],
-            output_transformers=[output_transform],
+            input_transformers=[input_normalization],
+            output_transformers=[output_normalization],
         )
     else:
         # model is an ensemble list of NNs
@@ -265,10 +267,10 @@ def build_lume_model(
                     model=model_nn.cpu(),
                     input_variables=input_vars,
                     output_variables=output_vars,
-                    input_transformers=[input_transform],
+                    input_transformers=[input_normalization],
                     output_transformers=[
                         calibration_transform,
-                        output_transform,
+                        output_normalization,
                     ],  # saving calibration before normalization
                 )
             )
@@ -448,15 +450,15 @@ def register_model_to_mlflow(model, model_type, experiment, config_dict):
     # Apply normalization to the training data
     X_train = torch.tensor(df_train[input_names].values, dtype=torch.float)
     y_train = torch.tensor(df_train[output_names].values, dtype=torch.float)
-    input_transform, output_transform = build_transforms(
+    input_normalization, output_normalization = build_normalizations(
         len(input_names), X_train, len(output_names), y_train
     )
     norm_df_train = normalize(
-        df_train, input_names, input_transform, output_names, output_transform
+        df_train, input_names, input_normalization, output_names, output_normalization
     )
     if model_type != "GP":
         norm_df_val = normalize(
-            df_val, input_names, input_transform, output_names, output_transform
+            df_val, input_names, input_normalization, output_names, output_normalization
         )
 
     print("training started")
@@ -500,8 +502,8 @@ def register_model_to_mlflow(model, model_type, experiment, config_dict):
         model_type,
         input_variables,
         output_variables,
-        input_transform,
-        output_transform,
+        input_normalization,
+        output_normalization,
     )
 
     if test_mode: