Feature/1000

src/scripts/Aleatoric.py

@@ -80,6 +80,13 @@ def parse_args():
         default=DefaultsAnalysis["analysis"]["model_names_list"],
         help="Beginning of name for saved checkpoints and figures",
     )
+    parser.add_argument(
+        "--inject_type_list",
+        type=list,
+        required=False,
+        default=DefaultsAnalysis["analysis"]["inject_type_list"],
+        help="Feature and predictive",
+    )
     parser.add_argument(
         "--n_epochs",
         type=int,
@@ -133,13 +140,15 @@ def parse_args():
                 "n_models": args.n_models,
                 "n_epochs": args.n_epochs,
                 "data_prescription": args.data_prescription,
+                "data_dimension": args.data_dimension,
                 "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,
+                "inject_type_list": args.inject_type_list,
                 "plot": args.plot,
                 "savefig": args.savefig,
                 "verbose": args.verbose,
@@ -185,9 +194,6 @@ def beta_type(value):
         "analysis", "inject_type_list", "Analysis"
     )
     dim = config.get_item("model", "data_dimension", "Analysis")
-    sigma_list = []
-    for noise in noise_list:
-        sigma_list.append(DataPreparation.get_sigma(noise))
     root_dir = config.get_item("common", "dir", "Analysis")
     path_to_chk = root_dir + "checkpoints/"
     path_to_out = root_dir + "analysis/"
@@ -302,13 +308,16 @@ def beta_type(value):
                 alpha=0.25,
                 edgecolor=None,
             )
+            sigma = DataPreparation.get_sigma(
+                noise, inject_type=typei, data_dimension=dim
+            )
             ax.plot(
                 range(n_epochs),
                 al,
                 color=color_list[i],
-                label=r"$\sigma = $" + str(sigma_list[i]),
+                label=r"$\sigma = $" + str(sigma),
             )
-            ax.axhline(y=sigma_list[i], color=color_list[i], ls="--")
+            ax.axhline(y=sigma, color=color_list[i], ls="--")
         ax.set_ylabel("Aleatoric Uncertainty")
         ax.set_xlabel("Epoch")
         # if model[0:3] == "DER":

src/scripts/DeepEnsemble.py

@@ -215,6 +215,12 @@ def parse_args():
         default=DefaultsDE["model"]["n_hidden"],
         help="Number of hidden neurons in the hidden layer, default 64",
     )
+    parser.add_argument(
+        "--save_data_size",
+        action="store_true",
+        default=DefaultsDE["model"]["save_data_size"],
+        help="save chk with the number of examples in the dataset",
+    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -264,6 +270,7 @@ def parse_args():
                 "rs_list": args.rs_list,
                 "save_n_hidden": args.save_n_hidden,
                 "n_hidden": args.n_hidden,
+                "save_data_size": args.save_data_size,
                 "verbose": args.verbose,
             },
             "data": {
@@ -309,6 +316,7 @@ def beta_type(value):
 
 if __name__ == "__main__":
     config = parse_args()
+    verbose = config.get_item("model", "verbose", "DE")
     size_df = int(config.get_item("data", "size_df", "DE"))
     noise = config.get_item("data", "noise_level", "DE")
     norm = config.get_item("data", "normalize", "DE", raise_exception=False)
@@ -385,7 +393,7 @@ def beta_type(value):
         model_inputs = np.array([xs_array, ms_array, bs_array]).T
     plot_value = config.get_item("model", "plot", "DE")
     print(f"Value: {plot_value}, Type: {type(plot_value)}")
-    if plot_value:
+    if verbose:
         # briefly plot what some of the data looks like
         if dim == "0D":
             print(np.shape(xs_array), np.shape(model_outputs))
@@ -408,7 +416,7 @@ def beta_type(value):
     model_inputs, model_outputs, norm_params = DataPreparation.normalize(
         model_inputs, model_outputs, norm
     )
-    if plot_value:
+    if verbose:
         if dim == "2D":
             plt.clf()
             plt.imshow(model_inputs[0])
@@ -484,5 +492,7 @@ def beta_type(value):
         rs_list=config.get_item("model", "rs_list", "DE"),
         save_n_hidden=config.get_item("model", "save_n_hidden", "DE"),
         n_hidden=config.get_item("model", "n_hidden", "DE"),
-        verbose=config.get_item("model", "verbose", "DE"),
+        save_size_df=config.get_item("model", "save_data_size", "DE"),
+        size_df=size_df,
+        verbose=verbose,
     )

src/scripts/DeepEvidentialRegression.py

@@ -204,6 +204,12 @@ def parse_args():
         default=DefaultsDER["model"]["n_hidden"],
         help="Number of hidden neurons in the hidden layer, default 64",
     )
+    parser.add_argument(
+        "--save_data_size",
+        action="store_true",
+        default=DefaultsDER["model"]["save_data_size"],
+        help="save chk with the number of examples in the dataset",
+    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -250,6 +256,7 @@ def parse_args():
                 "rs": args.rs,
                 "save_n_hidden": args.save_n_hidden,
                 "n_hidden": args.n_hidden,
+                "save_data_size": args.save_data_size,
                 "verbose": args.verbose,
             },
             "data": {
@@ -400,7 +407,8 @@ 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
     )
-    trainData = TensorDataset(torch.Tensor(x_train), torch.Tensor(y_train))
+    trainData = TensorDataset(
+        torch.Tensor(x_train), torch.Tensor(y_train))
     trainDataLoader = DataLoader(
         trainData, batch_size=BATCH_SIZE, shuffle=True
     )
@@ -448,5 +456,7 @@ def parse_args():
         rs=config.get_item("model", "rs", "DER"),
         save_n_hidden=config.get_item("model", "save_n_hidden", "DER"),
         n_hidden=config.get_item("model", "n_hidden", "DER"),
+        save_size_df=config.get_item("model", "save_data_size", "DER"),
+        size_df=size_df,
         verbose=config.get_item("model", "verbose", "DER"),
     )