[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

ml/train_model.py

@@ -129,6 +129,7 @@ def normalize(df, input_names, input_transform, output_names, output_transform):
     )
     return norm_df, norm_exp_inputs, norm_exp_outputs, norm_sim_inputs, norm_sim_outputs
 
+
 def split_data(df_exp, df_sim, variables, model_type):
     if model_type == "GP":
         if len(df_exp) > 0:
@@ -151,6 +152,7 @@ def split_data(df_exp, df_sim, variables, model_type):
         else:
             return (sim_train_df[variables], sim_val_df[variables])
 
+
 def build_transforms(n_inputs, X_data, n_outputs, y_data):
     input_transform = AffineInputTransform(
         len(input_names), coefficient=X_train.std(axis=0), offset=X_train.mean(axis=0)
@@ -161,6 +163,7 @@ def build_transforms(n_inputs, X_data, n_outputs, y_data):
     output_transform = AffineInputTransform(n_outputs, coefficient=y_std, offset=y_mean)
     return input_transform, output_transform
 
+
 def train_nn_ensemble(
     model_type,
     n_inputs,
@@ -257,10 +260,10 @@ def build_torch_model_from_nn(
         ],
     )
 
-def train_gp(
-    norm_df_train, input_names, output_names,
-    input_transform, output_transform, device):
 
+def train_gp(
+    norm_df_train, input_names, output_names, input_transform, output_transform, device
+):
     gp_models = []
 
     for i, output_name in enumerate(output_names):
@@ -345,19 +348,24 @@ def train_gp(
         ]
 
     return GPModel(
-            model=model.cpu(),
-            input_variables=[ScalarVariable(**input_variables[k]) for k in input_variables.keys()],
-            output_variables=output_variables,
-            input_transform=[input_transform],
-            output_transform=[output_transform])
+        model=model.cpu(),
+        input_variables=[
+            ScalarVariable(**input_variables[k]) for k in input_variables.keys()
+        ],
+        output_variables=output_variables,
+        input_transform=[input_transform],
+        output_transform=[output_transform],
+    )
 
 
 def write_model(model, model_type, experiment, db):
     with tempfile.TemporaryDirectory() as temp_dir:
         if model_type != "GP":
             model.dump(file=os.path.join(temp_dir, experiment + ".yml"), save_jit=True)
         else:
-            model.dump(file=os.path.join(temp_dir, experiment + ".yml"), save_models=True)
+            model.dump(
+                file=os.path.join(temp_dir, experiment + ".yml"), save_models=True
+            )
         # Upload the model to the database
         # - Load the files that were just created into a dictionary
         with open(os.path.join(temp_dir, experiment + ".yml")) as f:
@@ -412,6 +420,7 @@ def write_model(model, model_type, experiment, db):
         db["models"].insert_one(document)
         print("Model uploaded to database")
 
+
 experiment, model_type = parse_arguments()
 config_dict = load_config(experiment)
 db = connect_to_db(config_dict)
@@ -513,7 +522,13 @@ def write_model(model, model_type, experiment, db):
 else:
     # Create separate GP models for each output to handle NaN values
 
-    model = train_gp(norm_df_train, input_names, output_names,
-                     input_transform, output_transform, device)
+    model = train_gp(
+        norm_df_train,
+        input_names,
+        output_names,
+        input_transform,
+        output_transform,
+        device,
+    )
 
 write_model(model, model_type, experiment, db)