Merge pull request #48 from IDEALLab/new_metrics

Metrics/Algo Bugfixes

Author: ffelten

engiopt/cgan_1d/cgan_1d.py

@@ -212,8 +212,7 @@ def prepare_data(problem: Problem, device: th.device) -> tuple[th.utils.data.Ten
     else:
         dummy_design, _ = problem.random_design()
         design_shape = spaces.flatten(problem.design_space, dummy_design).shape
-    conditions = problem.conditions
-    n_conds = len(conditions)
+    n_conds = len(problem.conditions_keys)
 
     # Logging
     run_name = f"{args.problem_id}__{args.algo}__{args.seed}__{int(time.time())}"
@@ -360,7 +359,7 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]:
                         ax.figure.canvas.draw()
                         img = np.array(fig.canvas.renderer.buffer_rgba())
                         axes[j].imshow(img)
-                        title = [(conditions[i][0], f"{dc[i]:.2f}") for i in range(n_conds)]
+                        title = [(problem.conditions_keys[i], f"{dc[i]:.2f}") for i in range(n_conds)]
                         title_string = "\n ".join(f"{condition}: {value}" for condition, value in title)
                         axes[j].title.set_text(title_string)  # Set title
                         axes[j].set_xticks([])  # Hide x ticks

engiopt/cgan_1d/evaluate_cgan_1d.py

@@ -98,7 +98,7 @@ def __init__(self):
 
     model = Generator(
         latent_dim=run.config["latent_dim"],
-        n_conds=len(problem.conditions),
+        n_conds=len(problem.conditions_keys),
         design_shape=design_shape,
         design_normalizer=design_normalizer,
         conds_normalizer=conds_normalizer,

engiopt/cgan_2d/cgan_2d.py

@@ -139,8 +139,7 @@ def forward(self, design: th.Tensor, conds: th.Tensor) -> th.Tensor:
     problem.reset(seed=args.seed)
 
     design_shape = problem.design_space.shape
-    conditions = problem.conditions
-    n_conds = len(conditions)
+    n_conds = len(problem.conditions_keys)
 
     # Logging
     run_name = f"{args.problem_id}__{args.algo}__{args.seed}__{int(time.time())}"
@@ -280,7 +279,7 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]:
                         img = tensor.cpu().numpy().reshape(design_shape[0], design_shape[1])  # Extract x and y coordinates
                         dc = desired_conds[j].cpu()
                         axes[j].imshow(img)  # Scatter plot
-                        title = [(conditions[i][0], f"{dc[i]:.2f}") for i in range(n_conds)]
+                        title = [(problem.conditions_keys[i], f"{dc[i]:.2f}") for i in range(n_conds)]
                         title_string = "\n ".join(f"{condition}: {value}" for condition, value in title)
                         axes[j].title.set_text(title_string)  # Set title
                         axes[j].set_xticks([])  # Hide x ticks

engiopt/cgan_2d/evaluate_cgan_2d.py

@@ -88,7 +88,7 @@ def __init__(self):
 
     model = Generator(
         latent_dim=run.config["latent_dim"],
-        n_conds=len(problem.conditions),
+        n_conds=len(problem.conditions_keys),
         design_shape=problem.design_space.shape,
     ).to(device)
     model.load_state_dict(ckpt["generator"])

engiopt/cgan_bezier/cgan_bezier.py

@@ -460,7 +460,7 @@ def denormalize(self, x: th.Tensor) -> th.Tensor:
     training_ds = th.utils.data.TensorDataset(
         th.stack(coords_set),
         th.stack(design_scalars).unsqueeze(1),
-        *[problem_dataset[key][:] for key, _ in problem.conditions],
+        *[problem_dataset[key][:] for key in problem.conditions_keys],
     )
 
     cond_tensors = th.stack(training_ds.tensors[2:])
@@ -481,7 +481,7 @@ def denormalize(self, x: th.Tensor) -> th.Tensor:
     discriminator = Discriminator(
         latent_dim=args.latent_dim,
         design_scalars=len(design_scalar_keys),
-        num_conds=len(problem.conditions),
+        num_conds=len(problem.conditions_keys),
         design_shape=problem.design_space["coords"].shape,
         conds_normalizer=conds_normalizer,
         design_scalars_normalizer=design_scalars_normalizer,
@@ -490,7 +490,7 @@ def denormalize(self, x: th.Tensor) -> th.Tensor:
     generator = Generator(
         latent_dim=args.latent_dim,
         noise_dim=args.noise_dim,
-        num_conds=len(problem.conditions),
+        num_conds=len(problem.conditions_keys),
         n_control_points=bezier_control_pts,
         n_data_points=n_data_points,
         conds_normalizer=conds_normalizer,
@@ -602,7 +602,7 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor, th.Tensor]:
                     ax.figure.canvas.draw()
                     img = np.array(fig.canvas.renderer.buffer_rgba())
                     axes[j].imshow(img)
-                    title = [(problem.conditions[i - 1][0], f"{do1[i]:.2f}") for i in range(1, len(do1))]
+                    title = [(problem.conditions_keys[i - 1], f"{do1[i]:.2f}") for i in range(1, len(do1))]
                     title_string = "\n ".join(f"{condition}: {value}" for condition, value in title)
                     axes[j].title.set_text(title_string)  # Set title
                     axes[j].set_xticks([])

engiopt/cgan_cnn_2d/cgan_cnn_2d.py

@@ -235,8 +235,7 @@ def forward(self, x: th.Tensor, c: th.Tensor) -> th.Tensor:
     problem.reset(seed=args.seed)
 
     design_shape = problem.design_space.shape
-    conditions = problem.conditions
-    n_conds = len(conditions)
+    n_conds = len(problem.conditions_keys)
 
     # Logging
     run_name = f"{args.problem_id}__{args.algo}__{args.seed}__{int(time.time())}"
@@ -377,7 +376,7 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]:
                         img = tensor.cpu().numpy().reshape(design_shape[0], design_shape[1])  # Extract x and y coordinates
                         dc = desired_conds[j].cpu()
                         axes[j].imshow(img)  # Scatter plot
-                        title = [(conditions[i][0], f"{dc[i]:.2f}") for i in range(n_conds)]
+                        title = [(problem.conditions_keys[i], f"{dc[i]:.2f}") for i in range(n_conds)]
                         title_string = "\n ".join(f"{condition}: {value}" for condition, value in title)
                         axes[j].title.set_text(title_string)  # Set title
                         axes[j].set_xticks([])  # Hide x ticks

engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py

@@ -87,7 +87,7 @@ def __init__(self):
     ckpt_path = os.path.join(artifact_dir, "generator.pth")
     ckpt = th.load(ckpt_path, map_location=th.device(device))
     model = Generator(
-        latent_dim=run.config["latent_dim"], n_conds=len(problem.conditions), design_shape=problem.design_space.shape
+        latent_dim=run.config["latent_dim"], n_conds=len(problem.conditions_keys), design_shape=problem.design_space.shape
     )
     model.load_state_dict(ckpt["generator"])
     model.eval()  # Set to evaluation mode

engiopt/cgan_cnn_3d/cgan_cnn_3d.py

@@ -358,7 +358,7 @@ def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, d
     if len(design_shape) != DESIGN_SHAPE_LEN:
         raise ValueError(f"Expected 3D design shape, got {design_shape}")
 
-    conditions = problem.conditions
+    conditions = problem.conditions_keys
     n_conds = len(conditions)
     condition_names = [cond[0] for cond in conditions]
 

engiopt/cgan_cnn_3d/evaluate_cgan_cnn_3d.py

@@ -10,11 +10,11 @@
 import pandas as pd
 import torch as th
 import tyro
-import wandb
 
 from engiopt import metrics
 from engiopt.cgan_cnn_3d.cgan_cnn_3d import Generator3D
 from engiopt.dataset_sample_conditions import sample_conditions
+import wandb
 
 
 @dataclasses.dataclass
@@ -63,7 +63,7 @@ class Args:
 
     # Reshape to match the expected input shape for the model
     conditions_tensor = conditions_tensor.unsqueeze(-1).unsqueeze(-1)
-    conditions_tensor = conditions_tensor.view(args.n_samples, len(problem.conditions), 1, 1, 1)
+    conditions_tensor = conditions_tensor.view(args.n_samples, len(problem.conditions_keys), 1, 1, 1)
 
     ### Set Up Generator ###
 

engiopt/cgan_vae/cgan_vae.py

@@ -451,7 +451,7 @@ def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, d
     if len(design_shape) != DESIGN_SHAPE_LEN:
         raise ValueError(f"Expected 3D design shape, got {design_shape}")
 
-    conditions = problem.conditions
+    conditions = problem.conditions_keys
     n_conds = len(conditions)
     condition_names = [cond[0] for cond in conditions]