flake8

deepskies · May 21, 2024 · d701989 · d701989
1 parent 5b9c12e
commit d701989
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 80 deletions.
diff --git a/src/analyze/analyze.py b/src/analyze/analyze.py
@@ -1,6 +1,5 @@
 # Contains modules to analyze the output checkpoints
 # from a trained model and make plots for the paper
-import numpy as np
 import torch
 
 
@@ -36,20 +35,19 @@ def load_checkpoint(
         """
         if model_name[0:3] == "DER":
             file_name = (
-                    str(path)
-                    + f"{model_name}_noise_{noise}_loss_{loss}"
-                    + f"_COEFF_{COEFF}_epoch_{epoch}"
-                    )
+                str(path)
+                + f"{model_name}_noise_{noise}_loss_{loss}"
+                + f"_COEFF_{COEFF}_epoch_{epoch}"
+            )
             if load_rs_chk:
-                file_name += (f"_rs_{rs}")
+                file_name += f"_rs_{rs}"
             if load_nh_chk:
-                file_name += (f"_n_hidden_{nh}")
+                file_name += f"_n_hidden_{nh}"
             file_name += ".pt"
         elif model_name[0:2] == "DE":
             file_name = (
-                str(path)
-                + f"{model_name}_noise_{noise}_beta_{BETA}_"
-                  f"nmodel_{nmodel}_epoch_{epoch}.pt"
+                str(path) + f"{model_name}_noise_{noise}_beta_{BETA}_"
+                f"nmodel_{nmodel}_epoch_{epoch}.pt"
             )
         checkpoint = torch.load(file_name, map_location=device)
         return checkpoint

diff --git a/src/models/models.py b/src/models/models.py
@@ -67,9 +67,7 @@ def forward(self, x):
         return torch.stack((gamma, nu, alpha, beta), dim=1)
 
 
-def model_setup_DER(loss_type,
-                    DEVICE,
-                    n_hidden):
+def model_setup_DER(loss_type, DEVICE, n_hidden):
     # initialize the model from scratch
     if loss_type == "SDER":
         Layer = SDERLayer
@@ -82,8 +80,7 @@ def model_setup_DER(loss_type,
 
     # from https://github.com/pasteurlabs/unreasonable_effective_der
     # /blob/main/x3_indepth.ipynb
-    model = torch.nn.Sequential(Model(4, n_hidden),
-                                Layer())
+    model = torch.nn.Sequential(Model(4, n_hidden), Layer())
     model = model.to(DEVICE)
     return model, lossFn
 
@@ -120,7 +117,6 @@ def model_setup_DE(loss_type, DEVICE):
     return model, lossFn
 
 
-
 # This following is from PasteurLabs -
 # https://github.com/pasteurlabs/unreasonable_effective_der/blob/main/models.py
 
@@ -174,14 +170,13 @@ def loss_sder(y, y_pred, coeff):
 
     # define aleatoric and epistemic uncert
     u_al = np.sqrt(
-        (beta.detach().numpy()
-         * (1 + nu.detach().numpy()))
+        (beta.detach().numpy() * (1 + nu.detach().numpy()))
         / (alpha.detach().numpy() * nu.detach().numpy())
     )
     u_ep = 1 / np.sqrt(nu.detach().numpy())
 
-    return torch.mean(torch.log(var) + (1.0 + coeff * nu) * error**2 / var), \
-        u_al, u_ep
+    return torch.mean(torch.log(var) +
+                      (1.0 + coeff * nu) * error**2 / var), u_al, u_ep
 
 
 # from martius lab

diff --git a/src/scripts/AleatoricInits.py b/src/scripts/AleatoricInits.py
@@ -23,8 +23,7 @@ def parse_args():
     # model
     # we need some info about the model to run this analysis
     # path to save the model results
-    parser.add_argument("--dir",
-                        default=DefaultsAnalysis["common"]["dir"])
+    parser.add_argument("--dir", default=DefaultsAnalysis["common"]["dir"])
     # now args for model
     parser.add_argument(
         "--n_models",
@@ -118,11 +117,13 @@ def parse_args():
         # check if args were specified in cli
         input_yaml = {
             "common": {"dir": args.dir},
-            "model": {"n_models": args.n_models,
-                      "n_epochs": args.n_epochs,
-                      "BETA": args.BETA,
-                      "COEFF": args.COEFF,
-                      "loss_type": args.loss_type},
+            "model": {
+                "n_models": args.n_models,
+                "n_epochs": args.n_epochs,
+                "BETA": args.BETA,
+                "COEFF": args.COEFF,
+                "loss_type": args.loss_type,
+            },
             "analysis": {
                 "noise_level_list": args.noise_level_list,
                 "model_names_list": args.model_names_list,
@@ -175,10 +176,10 @@ def beta_type(value):
     rs_list = [1, 2, 3, 4, 5]
     # check that this exists and if not make it
     if not os.path.isdir(path_to_out):
-        print('does not exist, making dir', path_to_out)
+        print("does not exist, making dir", path_to_out)
         os.mkdir(path_to_out)
     else:
-        print('already exists', path_to_out)
+        print("already exists", path_to_out)
     model_name_list = config.get_item("analysis",
                                       "model_names_list",
                                       "Analysis")
@@ -188,13 +189,11 @@ def beta_type(value):
     # make an empty nested dictionary with keys for
     # model names followed by noise levels
     al_dict = {
-        model_name: {noise: {rs: [] for rs in rs_list}
-                     for noise in noise_list}
+        model_name: {noise: {rs: [] for rs in rs_list} for noise in noise_list}
         for model_name in model_name_list
     }
     al_std_dict = {
-        model_name: {noise: {rs: [] for rs in rs_list}
-                     for noise in noise_list}
+        model_name: {noise: {rs: [] for rs in rs_list} for noise in noise_list}
         for model_name in model_name_list
     }
     n_epochs = config.get_item("model", "n_epochs", "Analysis")
@@ -216,13 +215,12 @@ def beta_type(value):
                             loss=loss_type,
                             load_rs_chk=True,
                             rs=rs,
-                            load_nh_chk=False
+                            load_nh_chk=False,
                         )
                         # path=path_to_chk)
                         # things to grab: 'valid_mse' and 'valid_bnll'
-                        _, aleatoric_m, _, a_std = (
+                        _, aleatoric_m, _, a_std = \
                             chk_module.ep_al_checkpoint_DER(chk)
-                        )
                         al_dict[model][noise][rs].append(aleatoric_m)
                         al_std_dict[model][noise][rs].append(a_std)
 
@@ -246,7 +244,8 @@ def beta_type(value):
                         list_vars.append(var_vals)
                         try:
                             al_dict[model][noise][nmodels + 1].append(
-                                np.mean(list_vars))
+                                np.mean(list_vars)
+                            )
                         except KeyError:
                             continue
     # make a two-paneled plot for the different noise levels
@@ -265,11 +264,11 @@ def beta_type(value):
                     al = np.array(np.sqrt(al_dict[model][noise][rs]))
                 else:
                     al = np.array(al_dict[model][noise][rs])
-                '''
-                # it doesn't really make sense to plot the std for the 
+                """
+                # it doesn't really make sense to plot the std for the
                 # case of the DE because each individual model
                 # makes up one in the ensemble
-                '''
+                """
                 if model[0:3] == "DER":
                     al_std = np.array(al_std_dict[model][noise][rs])
                     ax.fill_between(
@@ -278,7 +277,7 @@ def beta_type(value):
                         al + al_std,
                         color=color_list[n],
                         alpha=0.1,
-                        edgecolor=None
+                        edgecolor=None,
                     )
                 if r == 0:
                     ax.plot(
@@ -288,11 +287,8 @@ def beta_type(value):
                         label=r"$\sigma = $" + str(sigma_list[n]),
                     )
                 else:
-                    ax.plot(
-                        range(n_epochs),
-                        al,
-                        color=color_list[n])
-            ax.axhline(y=sigma_list[n], color=color_list[n], ls='--')
+                    ax.plot(range(n_epochs), al, color=color_list[n])
+            ax.axhline(y=sigma_list[n], color=color_list[n], ls="--")
         ax.set_ylabel("Aleatoric Uncertainty")
         ax.set_xlabel("Epoch")
         if model[0:3] == "DER":
@@ -303,12 +299,12 @@ def beta_type(value):
     plt.legend()
     if config.get_item("analysis", "savefig", "Analysis"):
         plt.savefig(
-                str(path_to_out)
-                + "aleatoric_uncertainty_n_epochs_"
-                + str(n_epochs)
-                + "_n_models_DE_"
-                + str(n_models)
-                + ".png"
-            )
+            str(path_to_out)
+            + "aleatoric_uncertainty_n_epochs_"
+            + str(n_epochs)
+            + "_n_models_DE_"
+            + str(n_models)
+            + ".png"
+        )
     if config.get_item("analysis", "plot", "Analysis"):
         plt.show()
diff --git a/src/scripts/AleatoricNHidden.py b/src/scripts/AleatoricNHidden.py
@@ -13,7 +13,8 @@
 def parse_args():
     parser = argparse.ArgumentParser(
         description="Analyzes the aleatoric uncertainty when the model \
-        architecture is jittered")
+        architecture is jittered"
+    )
     # there are three options with the parser:
     # 1) Read from a yaml
     # 2) Reads from the command line and default file
@@ -25,8 +26,7 @@ def parse_args():
     # model
     # we need some info about the model to run this analysis
     # path to save the model results
-    parser.add_argument("--dir",
-                        default=DefaultsAnalysis["common"]["dir"])
+    parser.add_argument("--dir", default=DefaultsAnalysis["common"]["dir"])
     # now args for model
     parser.add_argument(
         "--n_models",
@@ -120,11 +120,13 @@ def parse_args():
         # check if args were specified in cli
         input_yaml = {
             "common": {"dir": args.dir},
-            "model": {"n_models": args.n_models,
-                      "n_epochs": args.n_epochs,
-                      "BETA": args.BETA,
-                      "COEFF": args.COEFF,
-                      "loss_type": args.loss_type},
+            "model": {
+                "n_models": args.n_models,
+                "n_epochs": args.n_epochs,
+                "BETA": args.BETA,
+                "COEFF": args.COEFF,
+                "loss_type": args.loss_type,
+            },
             "analysis": {
                 "noise_level_list": args.noise_level_list,
                 "model_names_list": args.model_names_list,
@@ -178,10 +180,10 @@ def beta_type(value):
     rs = 1
     # check that this exists and if not make it
     if not os.path.isdir(path_to_out):
-        print('does not exist, making dir', path_to_out)
+        print("does not exist, making dir", path_to_out)
         os.mkdir(path_to_out)
     else:
-        print('already exists', path_to_out)
+        print("already exists", path_to_out)
     model_name_list = config.get_item("analysis",
                                       "model_names_list",
                                       "Analysis")
@@ -220,7 +222,7 @@ def beta_type(value):
                             load_rs_chk=True,
                             rs=rs,
                             load_nh_chk=True,
-                            nh=nh
+                            nh=nh,
                         )
                         # path=path_to_chk)
                         # things to grab: 'valid_mse' and 'valid_bnll'
@@ -250,7 +252,8 @@ def beta_type(value):
                         list_vars.append(var_vals)
                         try:
                             al_dict[model][noise][nmodels + 1].append(
-                                np.mean(list_vars))
+                                np.mean(list_vars)
+                            )
                         except KeyError:
                             continue
     # make a two-paneled plot for the different noise levels
@@ -277,9 +280,9 @@ def beta_type(value):
                     al + al_std,
                     color=color_list[n],
                     alpha=0.1,
-                    edgecolor=None
+                    edgecolor=None,
                 )
-                
+
                 if h == 0:
                     ax.plot(
                         range(n_epochs),
@@ -288,11 +291,8 @@ def beta_type(value):
                         label=r"$\sigma = $" + str(sigma_list[n]),
                     )
                 else:
-                    ax.plot(
-                        range(n_epochs),
-                        al,
-                        color=color_list[n])
-            ax.axhline(y=sigma_list[n], color=color_list[n], ls='--')
+                    ax.plot(range(n_epochs), al, color=color_list[n])
+            ax.axhline(y=sigma_list[n], color=color_list[n], ls="--")
         ax.set_ylabel("Aleatoric Uncertainty")
         ax.set_xlabel("Epoch")
         if model[0:3] == "DER":
@@ -303,12 +303,12 @@ def beta_type(value):
     plt.legend()
     if config.get_item("analysis", "savefig", "Analysis"):
         plt.savefig(
-                str(path_to_out)
-                + "aleatoric_uncertainty_n_epochs_"
-                + str(n_epochs)
-                + "_n_models_DE_"
-                + str(n_models)
-                + ".png"
-            )
+            str(path_to_out)
+            + "aleatoric_uncertainty_n_epochs_"
+            + str(n_epochs)
+            + "_n_models_DE_"
+            + str(n_models)
+            + ".png"
+        )
     if config.get_item("analysis", "plot", "Analysis"):
         plt.show()
diff --git a/src/train/train.py b/src/train/train.py
@@ -35,7 +35,6 @@ def train_DER(
     save_n_hidden=False,
     n_hidden=64,
     verbose=True,
-
 ):
     # first determine if you even need to run anything
     if not save_all_checkpoints and save_final_checkpoint: