Feature/2 d cnn

src/analyze/analyze.py

@@ -12,6 +12,7 @@ def load_checkpoint(
         model_name,
         prescription,
         inject_type,
+        data_dim,
         noise,
         epoch,
         device,
@@ -35,10 +36,11 @@ def load_checkpoint(
         :param model: PyTorch model to load the checkpoint into
         :return: Loaded model
         """
+        print(model_name)
         if model_name[0:3] == "DER":
             file_name = (
                 str(path)
-                + f"{model_name}_{prescription}_{inject_type}"
+                + f"{model_name}_{prescription}_{inject_type}_{data_dim}"
                 + f"_noise_{noise}_loss_{loss}_COEFF_{COEFF}_epoch_{epoch}"
             )
             if load_rs_chk:
@@ -48,10 +50,13 @@ def load_checkpoint(
             file_name += ".pt"
         elif model_name[0:2] == "DE":
             file_name = (
-                str(path)
-                + f"{model_name}_{prescription}_{inject_type}"
+                str(path) +
+                f"{model_name}_{prescription}_{inject_type}_{data_dim}"
                 f"_noise_{noise}_beta_{BETA}_nmodel_{nmodel}_epoch_{epoch}.pt"
             )
+        # import os
+        # print('cwd', os.getcwd())
+        # print(os.listdir(path))
         checkpoint = torch.load(file_name, map_location=device)
         return checkpoint
 

src/data/data.py

@@ -101,7 +101,7 @@ class DataPreparation:
     def __init__(self):
         self.data = None
 
-    def simulate_data_2d(
+    def image_gen(
         self,
         image_size=100,
         amplitude=10,
@@ -126,6 +126,39 @@ def simulate_data_2d(
                         )
         return image
 
+    def simulate_data_2d(self,
+                         size_df,
+                         params,
+                         image_size=32,
+                         inject_type="predictive",
+                         sigma=1):
+        image_size = 32
+        image_array = np.zeros((size_df, image_size, image_size))
+        total_brightness = []
+        for i in range(size_df):
+            image = self.image_gen(
+                image_size=image_size,
+                amplitude=params[i, 0],
+                radius=params[i, 1],
+                center_x=16,
+                center_y=16,
+                theta=params[i, 2],
+                noise_level=0)
+            if inject_type == "predictive":
+                image_array[i, :, :] = image
+                total_brightness.append(
+                    np.sum(image) + np.random.normal(
+                        loc=0, scale=sigma))
+            elif inject_type == "feature":
+                noisy_image = image + np.random.normal(
+                        loc=0, scale=sigma, size=(image_size, image_size))
+                image_array[i, :, :] = noisy_image
+                total_brightness.append(np.sum(image))
+            # we'll need the noisy image summed if we want to
+            # do a comparison of y - y':
+            # total_brightness_prop_noisy.append(np.sum(noisy_image))
+        return image_array, total_brightness
+
     def simulate_data(
         self,
         thetas,

src/models/models.py

@@ -67,7 +67,46 @@ def forward(self, x):
         return torch.stack((gamma, nu, alpha, beta), dim=1)
 
 
-def model_setup_DER(loss_type, DEVICE, n_hidden):
+class ConvLayers(nn.Module):
+    def __init__(self):
+        super(ConvLayers, self).__init__()
+        # a little strange = # of filters, usually goes from small to large
+        # double check on architecture decisions
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(10, 10, kernel_size=3, padding=1)
+        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(10, 10, kernel_size=3, padding=1)
+        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2, padding=1)
+        self.conv4 = nn.Conv2d(10, 5, kernel_size=3, padding=1)
+        self.conv5 = nn.Conv2d(5, 5, kernel_size=3, padding=1)
+        self.flatten = nn.Flatten()
+
+    def forward(self, x):
+        # print('input shape', x.shape)
+        if x.dim() == 3:  # Check if the input is of shape (batchsize, 32, 32)
+            # Add channel dimension, becomes (batchsize, 1, 32, 32)
+            x = x.unsqueeze(1)
+        # print('shape after potential unsqeeze', x.shape)
+        x = nn.functional.relu(self.conv1(x))
+        # print('shape after conv1', x.shape)
+        x = nn.functional.relu(self.conv2(x))
+        # print('shape after conv2', x.shape)
+        x = self.pool1(x)
+        # print('shape after pool1', x.shape)
+        x = nn.functional.relu(self.conv3(x))
+        # print('shape after conv3', x.shape)
+        x = self.pool2(x)
+        # print('shape after pool2', x.shape)
+        x = nn.functional.relu(self.conv4(x))
+        # print('shape after conv4', x.shape)
+        x = nn.functional.relu(self.conv5(x))
+        # print('shape after conv5', x.shape)
+        x = self.flatten(x)
+        # print('shape after flatten', x.shape)
+        return x
+
+
+def model_setup_DER(loss_type, DEVICE, n_hidden=64, data_type="0D"):
     # initialize the model from scratch
     if loss_type == "SDER":
         Layer = SDERLayer
@@ -77,10 +116,24 @@ def model_setup_DER(loss_type, DEVICE, n_hidden):
         Layer = DERLayer
         # initialize our loss function
         lossFn = loss_der
-
-    # from https://github.com/pasteurlabs/unreasonable_effective_der
-    # /blob/main/x3_indepth.ipynb
-    model = torch.nn.Sequential(Model(4, n_hidden), Layer())
+    if data_type == "2D":
+        # Define the convolutional layers
+        conv_layers = ConvLayers()
+
+        # Initialize the rest of the model
+        model = torch.nn.Sequential(
+            conv_layers,
+            Model(
+                n_hidden=n_hidden, n_input=405, n_output=4
+            ),  # Adjust input size according to the flattened output size
+            Layer(),
+        )
+    elif data_type == "0D":
+        # from https://github.com/pasteurlabs/unreasonable_effective_der
+        # /blob/main/x3_indepth.ipynb
+        model = torch.nn.Sequential(
+            Model(n_hidden=n_hidden, n_input=3, n_output=4), Layer()
+        )
     model = model.to(DEVICE)
     return model, lossFn
 
@@ -97,22 +150,37 @@ def forward(self, x):
         return torch.stack((mu, var), dim=1)
 
 
-def model_setup_DE(loss_type, DEVICE):
+def model_setup_DE(loss_type, DEVICE, n_hidden=64, data_type="0D"):
     # initialize the model from scratch
     if loss_type == "no_var_loss":
         # model = de_no_var().to(DEVICE)
         lossFn = torch.nn.MSELoss(reduction="mean")
     if loss_type == "var_loss":
         # model = de_var().to(DEVICE)
         Layer = MuVarLayer
-        lossFn = torch.nn.GaussianNLLLoss(full=False,
-                                          eps=1e-06,
-                                          reduction="mean")
+        lossFn = torch.nn.GaussianNLLLoss(
+            full=False, eps=1e-06, reduction="mean")
     if loss_type == "bnll_loss":
         # model = de_var().to(DEVICE)
         Layer = MuVarLayer
         lossFn = loss_bnll
-    model = torch.nn.Sequential(Model(2, 64), Layer())
+    if data_type == "2D":
+        # Define the convolutional layers
+        conv_layers = ConvLayers()
+        # Initialize the rest of the model
+        model = torch.nn.Sequential(
+            conv_layers,
+            Model(
+                n_hidden=n_hidden, n_input=405, n_output=2
+            ),  # Adjust input size according to the flattened output size
+            Layer(),
+        )
+    elif data_type == "0D":
+        # from https://github.com/pasteurlabs/unreasonable_effective_der
+        # /blob/main/x3_indepth.ipynb
+        model = torch.nn.Sequential(
+            Model(n_hidden=n_hidden, n_input=3, n_output=2), Layer()
+        )
     model = model.to(DEVICE)
     return model, lossFn
 
@@ -122,10 +190,10 @@ def model_setup_DE(loss_type, DEVICE):
 
 
 class Model(nn.Module):
-    def __init__(self, n_output, n_hidden):
+    def __init__(self, n_output=4, n_hidden=64, n_input=3):
         super().__init__()
         self.model = nn.Sequential(
-            nn.Linear(3, n_hidden),
+            nn.Linear(n_input, n_hidden),
             nn.ReLU(),
             nn.Linear(n_hidden, n_hidden),
             nn.ReLU(),
@@ -175,8 +243,8 @@ def loss_sder(y, y_pred, coeff):
     )
     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

src/scripts/Aleatoric.py

@@ -32,6 +32,11 @@ def parse_args():
         "-dp",
         default=DefaultsAnalysis["model"]["data_prescription"],
     )
+    parser.add_argument(
+        "--data_dimension",
+        "-dd",
+        default=DefaultsAnalysis["model"]["data_dimension"],
+    )
     parser.add_argument(
         "--n_models",
         type=int,
@@ -179,6 +184,8 @@ def beta_type(value):
         "model", "data_prescription", "Analysis")
     inject_type_list = config.get_item(
         "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))
@@ -220,6 +227,7 @@ def beta_type(value):
                         model,
                         prescription,
                         typei,
+                        dim,
                         noise,
                         epoch,
                         DEVICE,
@@ -245,6 +253,7 @@ def beta_type(value):
                             model,
                             prescription,
                             typei,
+                            dim,
                             noise,
                             epoch,
                             DEVICE,