passing tests

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

Hide file tree

Showing 7 changed files with 9 additions and 22 deletions.
diff --git a/src/analyze/analyze.py b/src/analyze/analyze.py
@@ -45,14 +45,13 @@ def load_checkpoint(
             if load_nh_chk:
                 file_name += (f"_n_hidden_{nh}")
             file_name += ".pt"
-            checkpoint = torch.load(file_name, map_location=device)
         elif model_name[0:2] == "DE":
             file_name = (
                 str(path)
                 + f"{model_name}_noise_{noise}_beta_{BETA}_"
                   f"nmodel_{nmodel}_epoch_{epoch}.pt"
             )
-            checkpoint = torch.load(file_name, map_location=device)
+        checkpoint = torch.load(file_name, map_location=device)
         return checkpoint
 
     def ep_al_checkpoint_DE(self, checkpoint):

diff --git a/src/models/models.py b/src/models/models.py
@@ -115,7 +115,7 @@ def model_setup_DE(loss_type, DEVICE):
         # model = de_var().to(DEVICE)
         Layer = MuVarLayer
         lossFn = loss_bnll
-    model = torch.nn.Sequential(Model(2), Layer())
+    model = torch.nn.Sequential(Model(2, 64), Layer())
     model = model.to(DEVICE)
     return model, lossFn
 

diff --git a/src/scripts/Aleatoric.py b/src/scripts/Aleatoric.py
@@ -187,19 +187,10 @@ def beta_type(value):
     chk_module = AggregateCheckpoints()
     # make an empty nested dictionary with keys for
     # model names followed by noise levels
-    ep_dict = {
-        model_name: {noise: [] for noise in noise_list}
-        for model_name in model_name_list
-    }
     al_dict = {
         model_name: {noise: [] for noise in noise_list}
         for model_name in model_name_list
     }
-
-    ep_std_dict = {
-        model_name: {noise: [] for noise in noise_list}
-        for model_name in model_name_list
-    }
     al_std_dict = {
         model_name: {noise: [] for noise in noise_list}
         for model_name in model_name_list
@@ -225,12 +216,10 @@ def beta_type(value):
                     epistemic_m, aleatoric_m, e_std, a_std = (
                         chk_module.ep_al_checkpoint_DER(chk)
                     )
-                    ep_dict[model][noise].append(epistemic_m)
                     al_dict[model][noise].append(aleatoric_m)
-                    ep_std_dict[model][noise].append(e_std)
                     al_std_dict[model][noise].append(a_std)
 
-            elif model[0:3] == "DE_":
+            else:
                 n_models = config.get_item("model", "n_models", "DE")
                 for epoch in range(n_epochs):
                     list_mus = []
@@ -248,12 +237,8 @@ def beta_type(value):
                         mu_vals, var_vals = chk_module.ep_al_checkpoint_DE(chk)
                         list_mus.append(mu_vals)
                         list_vars.append(var_vals)
-                    ep_dict[model][noise].append(np.median(np.std(list_mus,
-                                                                  axis=0)))
                     al_dict[model][noise].append(np.median(np.mean(list_vars,
                                                                    axis=0)))
-                    ep_std_dict[model][noise].append(np.std(np.std(list_mus,
-                                                                   axis=0)))
                     al_std_dict[model][noise].append(np.std(np.mean(list_vars,
                                                                     axis=0)))
     # make a two-paneled plot for the different noise levels

diff --git a/src/scripts/DeepEnsemble.py b/src/scripts/DeepEnsemble.py
@@ -316,7 +316,8 @@ def beta_type(value):
                                  "model_type",
                                  "DE") + "_noise_" + noise
     model, lossFn = models.model_setup_DE(
-        config.get_item("model", "loss_type", "DE"), DEVICE
+        config.get_item("model", "loss_type", "DE"),
+        DEVICE,
     )
     print(
         "save final checkpoint has this value",

diff --git a/src/scripts/DeepEvidentialRegression.py b/src/scripts/DeepEvidentialRegression.py
@@ -298,8 +298,6 @@ def parse_args():
     x_train, x_val, y_train, y_val = DataPreparation.train_val_split(
         model_inputs, model_outputs, val_proportion=val_prop, random_state=rs
     )
-    print('train', np.shape(x_train), 'val', np.shape(x_val))
-    STOP
     trainData = TensorDataset(torch.Tensor(x_train), torch.Tensor(y_train))
     trainDataLoader = DataLoader(
         trainData,

diff --git a/test/test_Aleatoric.py b/test/test_Aleatoric.py
@@ -148,6 +148,8 @@ def create_test_config_DER(
             "savefig": False,
             "save_chk_random_seed_init": False,
             "rs": 42,
+            "save_n_hidden": False,
+            "n_hidden": 64,
             "verbose": False,
         },
         "data": {

diff --git a/test/test_DeepEvidentialRegression.py b/test/test_DeepEvidentialRegression.py
@@ -83,6 +83,8 @@ def create_test_config(
             "savefig": True,
             "save_chk_random_seed_init": False,
             "rs": 42,
+            "save_n_hidden": False,
+            "n_hidden": 64,
             "verbose": False,
         },
         "data": {