From 87be1913029ccc2ad77db2e440515cd02fd41e67 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 23 Sep 2025 14:31:45 +0200 Subject: [PATCH 1/3] resolved breaks from problem.conditions datatype change --- engiopt/cgan_1d/cgan_1d.py | 5 ++--- engiopt/cgan_1d/evaluate_cgan_1d.py | 2 +- engiopt/cgan_2d/cgan_2d.py | 5 ++--- engiopt/cgan_2d/evaluate_cgan_2d.py | 2 +- engiopt/cgan_bezier/cgan_bezier.py | 8 ++++---- engiopt/cgan_cnn_2d/cgan_cnn_2d.py | 5 ++--- engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py | 2 +- engiopt/cgan_cnn_3d/cgan_cnn_3d.py | 2 +- engiopt/cgan_cnn_3d/evaluate_cgan_cnn_3d.py | 4 ++-- engiopt/cgan_vae/cgan_vae.py | 2 +- engiopt/cgan_vae/evaluate_cgan_vae.py | 6 +++--- engiopt/diffusion_2d_cond/diffusion_2d_cond.py | 6 +++--- engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py | 2 +- 13 files changed, 24 insertions(+), 27 deletions(-) diff --git a/engiopt/cgan_1d/cgan_1d.py b/engiopt/cgan_1d/cgan_1d.py index 27257f4f..016f6340 100644 --- a/engiopt/cgan_1d/cgan_1d.py +++ b/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 diff --git a/engiopt/cgan_1d/evaluate_cgan_1d.py b/engiopt/cgan_1d/evaluate_cgan_1d.py index ab6f381d..a8beb031 100644 --- a/engiopt/cgan_1d/evaluate_cgan_1d.py +++ b/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, diff --git a/engiopt/cgan_2d/cgan_2d.py b/engiopt/cgan_2d/cgan_2d.py index afd4004e..dafa7a2d 100644 --- a/engiopt/cgan_2d/cgan_2d.py +++ b/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 diff --git a/engiopt/cgan_2d/evaluate_cgan_2d.py b/engiopt/cgan_2d/evaluate_cgan_2d.py index 5f541083..c0c93d4d 100644 --- a/engiopt/cgan_2d/evaluate_cgan_2d.py +++ b/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"]) diff --git a/engiopt/cgan_bezier/cgan_bezier.py b/engiopt/cgan_bezier/cgan_bezier.py index 91742ee8..7c7c5f58 100644 --- a/engiopt/cgan_bezier/cgan_bezier.py +++ b/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([]) diff --git a/engiopt/cgan_cnn_2d/cgan_cnn_2d.py b/engiopt/cgan_cnn_2d/cgan_cnn_2d.py index a655c147..61c084d8 100644 --- a/engiopt/cgan_cnn_2d/cgan_cnn_2d.py +++ b/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 diff --git a/engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py b/engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py index 96adf9d8..5e63e1e2 100644 --- a/engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py +++ b/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 diff --git a/engiopt/cgan_cnn_3d/cgan_cnn_3d.py b/engiopt/cgan_cnn_3d/cgan_cnn_3d.py index 5cdeae7e..a271e854 100644 --- a/engiopt/cgan_cnn_3d/cgan_cnn_3d.py +++ b/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] diff --git a/engiopt/cgan_cnn_3d/evaluate_cgan_cnn_3d.py b/engiopt/cgan_cnn_3d/evaluate_cgan_cnn_3d.py index 58effa21..13602960 100644 --- a/engiopt/cgan_cnn_3d/evaluate_cgan_cnn_3d.py +++ b/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 ### diff --git a/engiopt/cgan_vae/cgan_vae.py b/engiopt/cgan_vae/cgan_vae.py index 8e8b0c00..52d1caf4 100644 --- a/engiopt/cgan_vae/cgan_vae.py +++ b/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] diff --git a/engiopt/cgan_vae/evaluate_cgan_vae.py b/engiopt/cgan_vae/evaluate_cgan_vae.py index a70953dd..fff54a88 100644 --- a/engiopt/cgan_vae/evaluate_cgan_vae.py +++ b/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 diff --git a/engiopt/diffusion_2d_cond/diffusion_2d_cond.py b/engiopt/diffusion_2d_cond/diffusion_2d_cond.py index 9da594f9..e09f838e 100644 --- a/engiopt/diffusion_2d_cond/diffusion_2d_cond.py +++ b/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 diff --git a/engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py b/engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py index fc21506c..8403908b 100644 --- a/engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py +++ b/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) From dbf4f3b3f0d53910c8dcf8d35bcd84f64a8e61f8 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 23 Sep 2025 14:37:38 +0200 Subject: [PATCH 2/3] metrics bugfix missing reset between optimize/sim calls --- engiopt/metrics.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/engiopt/metrics.py b/engiopt/metrics.py index c2723fae..d0f108b0 100644 --- a/engiopt/metrics.py +++ b/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)) From 19a2f61bb2083967f2bce2693bb904fc46f17728 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Fri, 17 Oct 2025 11:14:19 +0200 Subject: [PATCH 3/3] downgrade wandb to resolve bug --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 9463ca11..5cd07e8b 100644 --- a/pyproject.toml +++ b/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",