Metrics/Algo Bugfixes

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]
 

engiopt/cgan_vae/evaluate_cgan_vae.py

@@ -10,11 +10,11 @@
 import pandas as pd
 import torch as th
 import tyro
-import wandb
 
 from engiopt import metrics
 from engiopt.cgan_vae.cgan_vae 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 ###
 
@@ -91,7 +91,7 @@ def __init__(self):
     for key in ckpt:
         print("Checkpoint key:", key)
     model = Generator3D(
-        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/diffusion_2d_cond/diffusion_2d_cond.py

@@ -264,7 +264,7 @@ def sample_timestep(
 
     # Loss function
     adversarial_loss: th.nn.Module = th.nn.MSELoss()
-    encoder_hid_dim = len(problem.conditions)
+    encoder_hid_dim = len(problem.conditions_keys)
     # Initialize UNet from Huggingface
     model = UNet2DConditionModel(
         sample_size=design_shape,
@@ -294,7 +294,7 @@ def sample_timestep(
     training_ds = th.utils.data.TensorDataset(
         filtered_ds_norm.flatten(1), *[training_ds[key][:] for key in problem.conditions_keys]
     )
-    cond_tensors = th.stack(training_ds.tensors[1 : len(problem.conditions) + 1])
+    cond_tensors = th.stack(training_ds.tensors[1 : len(problem.conditions_keys) + 1])
     conds_min = cond_tensors.amin(dim=tuple(range(1, cond_tensors.ndim)))
     conds_max = cond_tensors.amax(dim=tuple(range(1, cond_tensors.ndim)))
 
@@ -412,7 +412,7 @@ def sample_designs(model: UNet2DConditionModel, n_designs: int = 25) -> tuple[th
                         img = tensor.cpu().numpy()  # Extract x and y coordinates
                         dc = hidden_states[j, 0, :].cpu()
                         axes[j].imshow(img[0])  # image plot
-                        title = [(problem.conditions[i][0], f"{dc[i]:.2f}") for i in range(len(problem.conditions))]
+                        title = [(problem.conditions_keys[i], f"{dc[i]:.2f}") for i in range(len(problem.conditions_keys))]
                         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/diffusion_2d_cond/evaluate_diffusion_2d_cond.py

@@ -101,7 +101,7 @@ def __init__(self):
         up_block_types=("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
         layers_per_block=run.config["layers_per_block"],
         transformer_layers_per_block=1,
-        encoder_hid_dim=len(problem.conditions),
+        encoder_hid_dim=len(problem.conditions_keys),
         only_cross_attention=True,
     ).to(device)
 

engiopt/metrics.py

@@ -184,9 +184,12 @@ def metrics(
             unflattened_design = spaces.unflatten(problem.design_space, gen_designs[i])
         else:
             unflattened_design = gen_designs[i]
+        problem.reset()
         _, opt_history = problem.optimize(unflattened_design, config=conditions)
+        problem.reset()
         reference_optimum = problem.simulate(dataset_designs[i], config=conditions)
         opt_history_gaps = optimality_gap(opt_history, reference_optimum)
+        problem.reset()
 
         iog_list.append(problem.simulate(unflattened_design, config=conditions))
         cog_list.append(np.sum(opt_history_gaps))

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "numpy",
     "torch >= 2.5.0",
     "torchvision >= 0.20.1",
-    "wandb >= 0.18.7",
+    "wandb == 0.22.0",  #https://github.com/wandb/wandb/issues/10647
     "tqdm >= 4.67.1",
     "matplotlib >= 3.9.2",
     "denoising_diffusion_pytorch",