From 646617d58a3c428c8b3e3b4441714d10183956fc Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Fri, 24 Jul 2026 13:56:14 +0200 Subject: [PATCH 01/24] Move model packages under engiopt/generators/ Mirrors engibench/problems// so the two halves of the stack are laid out the same way. Pure move: file contents are unchanged apart from the import paths the move forces (engiopt. -> engiopt.generators.). Co-Authored-By: Claude Opus 4.8 (1M context) --- engiopt/{ => generators}/cgan_1d/__init__.py | 0 engiopt/generators/cgan_1d/adapter.py | 59 ++++++++ engiopt/{ => generators}/cgan_1d/cgan_1d.py | 0 .../cgan_1d/evaluate_cgan_1d.py | 4 +- engiopt/{ => generators}/cgan_2d/__init__.py | 0 engiopt/generators/cgan_2d/adapter.py | 52 +++++++ engiopt/{ => generators}/cgan_2d/cgan_2d.py | 0 .../cgan_2d/evaluate_cgan_2d.py | 2 +- .../{ => generators}/cgan_bezier/__init__.py | 0 engiopt/generators/cgan_bezier/adapter.py | 76 ++++++++++ .../cgan_bezier/cgan_bezier.py | 0 .../{ => generators}/cgan_cnn_2d/__init__.py | 0 engiopt/generators/cgan_cnn_2d/adapter.py | 54 +++++++ .../cgan_cnn_2d/cgan_cnn_2d.py | 0 .../cgan_cnn_2d/evaluate_cgan_cnn_2d.py | 2 +- .../cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml | 44 ++++++ .../sweep_hp_tuning_heatconduction2d.yaml | 44 ++++++ .../sweep_hp_tuning_photonics2d.yaml | 44 ++++++ .../{ => generators}/cgan_cnn_3d/README.md | 0 .../{ => generators}/cgan_cnn_3d/__init__.py | 0 engiopt/generators/cgan_cnn_3d/adapter.py | 66 +++++++++ .../cgan_cnn_3d/cgan_cnn_3d.py | 0 .../cgan_cnn_3d/evaluate_cgan_cnn_3d.py | 2 +- .../cgan_cnn_3d/sweep_cgan_cnn_3d.yaml | 0 engiopt/{ => generators}/cgan_vae/README.md | 0 engiopt/{ => generators}/cgan_vae/__init__.py | 0 engiopt/generators/cgan_vae/adapter.py | 54 +++++++ engiopt/{ => generators}/cgan_vae/cgan_vae.py | 0 .../cgan_vae/evaluate_cgan_vae.py | 2 +- .../cgan_vae/sample_script.ipynb | 0 .../cgan_vae/sweep_cgan_vae.yaml | 0 .../{ => generators}/diffusion_1d/__init__.py | 0 engiopt/generators/diffusion_1d/adapter.py | 77 ++++++++++ .../diffusion_1d/diffusion_1d.py | 0 .../diffusion_1d/evaluate_diffusion_1d.py | 2 +- .../diffusion_2d_cond/__init__.py | 0 .../generators/diffusion_2d_cond/adapter.py | 116 +++++++++++++++ .../diffusion_2d_cond/diffusion_2d_cond.py | 0 .../evaluate_diffusion_2d_cond.py | 6 +- engiopt/{ => generators}/gan_1d/__init__.py | 0 engiopt/generators/gan_1d/adapter.py | 52 +++++++ .../gan_1d/evaluate_gan_1d.py | 4 +- engiopt/{ => generators}/gan_1d/gan_1d.py | 0 engiopt/{ => generators}/gan_2d/__init__.py | 0 engiopt/generators/gan_2d/adapter.py | 46 ++++++ .../gan_2d/evaluate_gan_2d.py | 2 +- engiopt/{ => generators}/gan_2d/gan_2d.py | 0 .../{ => generators}/gan_bezier/__init__.py | 0 engiopt/generators/gan_bezier/adapter.py | 79 ++++++++++ .../gan_bezier/evaluate_gan_bezier.py | 4 +- .../{ => generators}/gan_bezier/gan_bezier.py | 0 .../{ => generators}/gan_cnn_2d/__init__.py | 0 engiopt/generators/gan_cnn_2d/adapter.py | 52 +++++++ .../gan_cnn_2d/evaluate_gan_cnn_2d.py | 2 +- .../{ => generators}/gan_cnn_2d/gan_cnn_2d.py | 0 .../pixel_cnn_pp_2d/__init__.py | 0 engiopt/generators/pixel_cnn_pp_2d/adapter.py | 83 +++++++++++ .../evaluate_pixel_cnn_pp_2d.py | 4 +- .../pixel_cnn_pp_2d/pixel_cnn_pp_2d.py | 0 engiopt/{ => generators}/vqgan/__init__.py | 0 engiopt/generators/vqgan/adapter.py | 136 ++++++++++++++++++ .../{ => generators}/vqgan/evaluate_vqgan.py | 4 +- engiopt/{ => generators}/vqgan/utils.py | 0 engiopt/{ => generators}/vqgan/vqgan.py | 24 ++-- example_easy_model.ipynb | 2 +- example_hard_model.ipynb | 4 +- tests/test_cgan_output_activation.py | 8 +- tests/test_diffusion_fixes.py | 12 +- 68 files changed, 1179 insertions(+), 45 deletions(-) rename engiopt/{ => generators}/cgan_1d/__init__.py (100%) create mode 100644 engiopt/generators/cgan_1d/adapter.py rename engiopt/{ => generators}/cgan_1d/cgan_1d.py (100%) rename engiopt/{ => generators}/cgan_1d/evaluate_cgan_1d.py (97%) rename engiopt/{ => generators}/cgan_2d/__init__.py (100%) create mode 100644 engiopt/generators/cgan_2d/adapter.py rename engiopt/{ => generators}/cgan_2d/cgan_2d.py (100%) rename engiopt/{ => generators}/cgan_2d/evaluate_cgan_2d.py (98%) rename engiopt/{ => generators}/cgan_bezier/__init__.py (100%) create mode 100644 engiopt/generators/cgan_bezier/adapter.py rename engiopt/{ => generators}/cgan_bezier/cgan_bezier.py (100%) rename engiopt/{ => generators}/cgan_cnn_2d/__init__.py (100%) create mode 100644 engiopt/generators/cgan_cnn_2d/adapter.py rename engiopt/{ => generators}/cgan_cnn_2d/cgan_cnn_2d.py (100%) rename engiopt/{ => generators}/cgan_cnn_2d/evaluate_cgan_cnn_2d.py (98%) create mode 100644 engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml create mode 100644 engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_heatconduction2d.yaml create mode 100644 engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_photonics2d.yaml rename engiopt/{ => generators}/cgan_cnn_3d/README.md (100%) rename engiopt/{ => generators}/cgan_cnn_3d/__init__.py (100%) create mode 100644 engiopt/generators/cgan_cnn_3d/adapter.py rename engiopt/{ => generators}/cgan_cnn_3d/cgan_cnn_3d.py (100%) rename engiopt/{ => generators}/cgan_cnn_3d/evaluate_cgan_cnn_3d.py (98%) rename engiopt/{ => generators}/cgan_cnn_3d/sweep_cgan_cnn_3d.yaml (100%) rename engiopt/{ => generators}/cgan_vae/README.md (100%) rename engiopt/{ => generators}/cgan_vae/__init__.py (100%) create mode 100644 engiopt/generators/cgan_vae/adapter.py rename engiopt/{ => generators}/cgan_vae/cgan_vae.py (100%) rename engiopt/{ => generators}/cgan_vae/evaluate_cgan_vae.py (98%) rename engiopt/{ => generators}/cgan_vae/sample_script.ipynb (100%) rename engiopt/{ => generators}/cgan_vae/sweep_cgan_vae.yaml (100%) rename engiopt/{ => generators}/diffusion_1d/__init__.py (100%) create mode 100644 engiopt/generators/diffusion_1d/adapter.py rename engiopt/{ => generators}/diffusion_1d/diffusion_1d.py (100%) rename engiopt/{ => generators}/diffusion_1d/evaluate_diffusion_1d.py (98%) rename engiopt/{ => generators}/diffusion_2d_cond/__init__.py (100%) create mode 100644 engiopt/generators/diffusion_2d_cond/adapter.py rename engiopt/{ => generators}/diffusion_2d_cond/diffusion_2d_cond.py (100%) rename engiopt/{ => generators}/diffusion_2d_cond/evaluate_diffusion_2d_cond.py (95%) rename engiopt/{ => generators}/gan_1d/__init__.py (100%) create mode 100644 engiopt/generators/gan_1d/adapter.py rename engiopt/{ => generators}/gan_1d/evaluate_gan_1d.py (97%) rename engiopt/{ => generators}/gan_1d/gan_1d.py (100%) rename engiopt/{ => generators}/gan_2d/__init__.py (100%) create mode 100644 engiopt/generators/gan_2d/adapter.py rename engiopt/{ => generators}/gan_2d/evaluate_gan_2d.py (98%) rename engiopt/{ => generators}/gan_2d/gan_2d.py (100%) rename engiopt/{ => generators}/gan_bezier/__init__.py (100%) create mode 100644 engiopt/generators/gan_bezier/adapter.py rename engiopt/{ => generators}/gan_bezier/evaluate_gan_bezier.py (97%) rename engiopt/{ => generators}/gan_bezier/gan_bezier.py (100%) rename engiopt/{ => generators}/gan_cnn_2d/__init__.py (100%) create mode 100644 engiopt/generators/gan_cnn_2d/adapter.py rename engiopt/{ => generators}/gan_cnn_2d/evaluate_gan_cnn_2d.py (98%) rename engiopt/{ => generators}/gan_cnn_2d/gan_cnn_2d.py (100%) rename engiopt/{ => generators}/pixel_cnn_pp_2d/__init__.py (100%) create mode 100644 engiopt/generators/pixel_cnn_pp_2d/adapter.py rename engiopt/{ => generators}/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py (96%) rename engiopt/{ => generators}/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py (100%) rename engiopt/{ => generators}/vqgan/__init__.py (100%) create mode 100644 engiopt/generators/vqgan/adapter.py rename engiopt/{ => generators}/vqgan/evaluate_vqgan.py (98%) rename engiopt/{ => generators}/vqgan/utils.py (100%) rename engiopt/{ => generators}/vqgan/vqgan.py (98%) diff --git a/engiopt/cgan_1d/__init__.py b/engiopt/generators/cgan_1d/__init__.py similarity index 100% rename from engiopt/cgan_1d/__init__.py rename to engiopt/generators/cgan_1d/__init__.py diff --git a/engiopt/generators/cgan_1d/adapter.py b/engiopt/generators/cgan_1d/adapter.py new file mode 100644 index 00000000..33ce1449 --- /dev/null +++ b/engiopt/generators/cgan_1d/adapter.py @@ -0,0 +1,59 @@ +"""`Generator` contract for the conditional 1D GAN.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import design_shape_of +from engiopt.core import Generator +from engiopt.generators.cgan_1d.cgan_1d import Generator as CGAN1DNet +from engiopt.generators.cgan_1d.cgan_1d import prepare_data + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class CGAN1D(Generator): + """Conditional GAN over 1D / dict-valued designs such as airfoils. + + The normalizers are part of the model: they are rebuilt from the training + dataset rather than stored in the checkpoint, so loading touches the dataset. + """ + + algo_id = "cgan_1d" + conditional = True + design_kinds = ("1d", "dict") + checkpoint_files = ("generator.pth",) + primary_state_key = "generator" + + def __init__(self, net: CGAN1DNet, latent_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> CGAN1D: + """Load a trained cGAN-1D generator, rebuilding its normalizers.""" + config = resolved.run_config + _, conds_normalizer, design_normalizer = prepare_data(problem, device) + net = CGAN1DNet( + latent_dim=config["latent_dim"], + n_conds=len(problem.conditions_keys), + design_shape=design_shape_of(problem), + design_normalizer=design_normalizer, + conds_normalizer=conds_normalizer, + ).to(device) + net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.eval() + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Map noise plus conditions through the generator.""" + cond = conditions.require_tensor(self.algo_id) + z = th.randn((n, self.latent_dim), device=self.device) + return self.net(z, cond) diff --git a/engiopt/cgan_1d/cgan_1d.py b/engiopt/generators/cgan_1d/cgan_1d.py similarity index 100% rename from engiopt/cgan_1d/cgan_1d.py rename to engiopt/generators/cgan_1d/cgan_1d.py diff --git a/engiopt/cgan_1d/evaluate_cgan_1d.py b/engiopt/generators/cgan_1d/evaluate_cgan_1d.py similarity index 97% rename from engiopt/cgan_1d/evaluate_cgan_1d.py rename to engiopt/generators/cgan_1d/evaluate_cgan_1d.py index 85ea0837..a383bd00 100644 --- a/engiopt/cgan_1d/evaluate_cgan_1d.py +++ b/engiopt/generators/cgan_1d/evaluate_cgan_1d.py @@ -13,8 +13,8 @@ import tyro from engiopt import metrics -from engiopt.cgan_1d.cgan_1d import Generator -from engiopt.cgan_1d.cgan_1d import prepare_data +from engiopt.generators.cgan_1d.cgan_1d import Generator +from engiopt.generators.cgan_1d.cgan_1d import prepare_data from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions diff --git a/engiopt/cgan_2d/__init__.py b/engiopt/generators/cgan_2d/__init__.py similarity index 100% rename from engiopt/cgan_2d/__init__.py rename to engiopt/generators/cgan_2d/__init__.py diff --git a/engiopt/generators/cgan_2d/adapter.py b/engiopt/generators/cgan_2d/adapter.py new file mode 100644 index 00000000..ed7fc2f7 --- /dev/null +++ b/engiopt/generators/cgan_2d/adapter.py @@ -0,0 +1,52 @@ +"""`Generator` contract for the conditional MLP GAN (2D).""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.cgan_2d.cgan_2d import Generator as CGAN2DNet + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class CGAN2D(Generator): + """Conditional GAN with a fully-connected generator, for 2D problems.""" + + algo_id = "cgan_2d" + conditional = True + design_kinds = ("2d",) + checkpoint_files = ("generator.pth",) + primary_state_key = "generator" + output_clip = (1e-3, 1.0) + + def __init__(self, net: CGAN2DNet, latent_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> CGAN2D: + """Load a trained cGAN-2D generator from its checkpoint package.""" + config = resolved.run_config + net = CGAN2DNet( + latent_dim=config["latent_dim"], + n_conds=len(problem.conditions_keys), + design_shape=problem.design_space.shape, + generator_output_activation=config.get("generator_output_activation", "tanh"), + ).to(device) + net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.eval() + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Map noise plus conditions through the generator.""" + cond = conditions.require_tensor(self.algo_id) + z = th.randn((n, self.latent_dim), device=self.device) + return self.net(z, cond) diff --git a/engiopt/cgan_2d/cgan_2d.py b/engiopt/generators/cgan_2d/cgan_2d.py similarity index 100% rename from engiopt/cgan_2d/cgan_2d.py rename to engiopt/generators/cgan_2d/cgan_2d.py diff --git a/engiopt/cgan_2d/evaluate_cgan_2d.py b/engiopt/generators/cgan_2d/evaluate_cgan_2d.py similarity index 98% rename from engiopt/cgan_2d/evaluate_cgan_2d.py rename to engiopt/generators/cgan_2d/evaluate_cgan_2d.py index 3d11e717..614c13d3 100644 --- a/engiopt/cgan_2d/evaluate_cgan_2d.py +++ b/engiopt/generators/cgan_2d/evaluate_cgan_2d.py @@ -12,7 +12,7 @@ import tyro from engiopt import metrics -from engiopt.cgan_2d.cgan_2d import Generator +from engiopt.generators.cgan_2d.cgan_2d import Generator from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions diff --git a/engiopt/cgan_bezier/__init__.py b/engiopt/generators/cgan_bezier/__init__.py similarity index 100% rename from engiopt/cgan_bezier/__init__.py rename to engiopt/generators/cgan_bezier/__init__.py diff --git a/engiopt/generators/cgan_bezier/adapter.py b/engiopt/generators/cgan_bezier/adapter.py new file mode 100644 index 00000000..30f95d40 --- /dev/null +++ b/engiopt/generators/cgan_bezier/adapter.py @@ -0,0 +1,76 @@ +"""`Generator` contract for the conditional BezierGAN airfoil model.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.cgan_bezier.cgan_bezier import _EPS +from engiopt.generators.cgan_bezier.cgan_bezier import Generator as CBezierGANNet +from engiopt.generators.cgan_bezier.cgan_bezier import prepare_data +from engiopt.transforms import flatten_dict_factory + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class CGANBezier(Generator): + """Conditional BezierGAN over airfoil designs. + + Like the unconditional BezierGAN it emits Bezier control points plus a + scalar angle of attack, which are flattened into the single vector the + evaluator expects. The latent code is uniform and the noise is normal with + scale 0.5, matching training. + """ + + algo_id = "cgan_bezier" + conditional = True + design_kinds = ("dict",) + checkpoint_files = ("bezier_generator.pth",) + primary_state_key = "generator" + + def __init__(self, net: CBezierGANNet, latent_dim: int, noise_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + self.noise_dim = noise_dim + self._flatten = flatten_dict_factory(self.problem, self.device) + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> CGANBezier: + """Load a trained conditional BezierGAN, rebuilding its normalizers.""" + config = resolved.run_config + _, conds_normalizer, design_scalars_normalizer, _keys = prepare_data(problem, device) + net = CBezierGANNet( + latent_dim=config["latent_dim"], + noise_dim=config["noise_dim"], + num_conds=len(problem.conditions_keys), + n_control_points=config["bezier_control_pts"], + n_data_points=problem.design_space["coords"].shape[1], + conds_normalizer=conds_normalizer, + design_scalars_normalizer=design_scalars_normalizer, + eps=_EPS, + scalar_features=1, + ).to(device) + net.load_state_dict(th.load(resolved.files["bezier_generator.pth"], map_location=device)[cls.primary_state_key]) + net.eval() + return cls( + net=net, latent_dim=config["latent_dim"], noise_dim=config["noise_dim"], problem=problem, device=device, **base + ) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Generate conditioned Bezier airfoils and flatten the dict designs.""" + cond = conditions.require_tensor(self.algo_id) + latent = th.rand(n, self.latent_dim, device=self.device) + noise = 0.5 * th.randn(n, self.noise_dim, device=self.device) + coords, *_rest, alphas = self.net(latent, noise, cond) + designs = [ + {"coords": coord, "angle_of_attack": alpha[0]} + for coord, alpha in zip(coords.detach().cpu().numpy(), alphas.detach().cpu().numpy()) + ] + return self._flatten(designs) diff --git a/engiopt/cgan_bezier/cgan_bezier.py b/engiopt/generators/cgan_bezier/cgan_bezier.py similarity index 100% rename from engiopt/cgan_bezier/cgan_bezier.py rename to engiopt/generators/cgan_bezier/cgan_bezier.py diff --git a/engiopt/cgan_cnn_2d/__init__.py b/engiopt/generators/cgan_cnn_2d/__init__.py similarity index 100% rename from engiopt/cgan_cnn_2d/__init__.py rename to engiopt/generators/cgan_cnn_2d/__init__.py diff --git a/engiopt/generators/cgan_cnn_2d/adapter.py b/engiopt/generators/cgan_cnn_2d/adapter.py new file mode 100644 index 00000000..2adaffd1 --- /dev/null +++ b/engiopt/generators/cgan_cnn_2d/adapter.py @@ -0,0 +1,54 @@ +"""`Generator` contract for the conditional GAN with CNN backbone (2D).""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.cgan_cnn_2d.cgan_cnn_2d import Generator as CGANCNN2DNet + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class CGANCNN2D(Generator): + """Conditional GAN with a CNN generator, for 2D topology problems.""" + + algo_id = "cgan_cnn_2d" + conditional = True + design_kinds = ("2d",) + checkpoint_files = ("generator.pth",) + primary_state_key = "generator" + # Several 2D simulators are unstable at exactly zero density. + output_clip = (1e-3, 1.0) + + def __init__(self, net: CGANCNN2DNet, latent_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> CGANCNN2D: + """Load a trained cGAN-CNN generator from its checkpoint package.""" + config = resolved.run_config + net = CGANCNN2DNet( + latent_dim=config["latent_dim"], + n_conds=len(problem.conditions_keys), + design_shape=problem.design_space.shape, + generator_output_activation=config.get("generator_output_activation", "tanh"), + ) + checkpoint = th.load(resolved.files["generator.pth"], map_location=device) + net.load_state_dict(checkpoint[cls.primary_state_key]) + net.eval().to(device) + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Map noise plus conditions through the generator.""" + cond = conditions.require_tensor(self.algo_id) + z = th.randn((n, self.latent_dim, 1, 1), device=self.device, dtype=th.float) + return self.net(z, cond.reshape(n, self.n_conds, 1, 1)) diff --git a/engiopt/cgan_cnn_2d/cgan_cnn_2d.py b/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py similarity index 100% rename from engiopt/cgan_cnn_2d/cgan_cnn_2d.py rename to engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py diff --git a/engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py b/engiopt/generators/cgan_cnn_2d/evaluate_cgan_cnn_2d.py similarity index 98% rename from engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py rename to engiopt/generators/cgan_cnn_2d/evaluate_cgan_cnn_2d.py index ec1eb3c5..758f78b2 100644 --- a/engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py +++ b/engiopt/generators/cgan_cnn_2d/evaluate_cgan_cnn_2d.py @@ -12,7 +12,7 @@ import tyro from engiopt import metrics -from engiopt.cgan_cnn_2d.cgan_cnn_2d import Generator +from engiopt.generators.cgan_cnn_2d.cgan_cnn_2d import Generator from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions diff --git a/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml b/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml new file mode 100644 index 00000000..73147ee7 --- /dev/null +++ b/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml @@ -0,0 +1,44 @@ +method: grid +command: + - "python" + - "engiopt/cgan_cnn_2d/cgan_cnn_2d.py" + - "--save_model" + - "${args}" + +metric: + name: g_loss + goal: minimize +parameters: + problem_id: + value: "beams2d" + lr_gen: + distribution: log_uniform_values + min: 0.00005 + max: 0.005 + lr_disc: + distribution: log_uniform_values + min: 0.00005 + max: 0.005 + latent_dim: + distribution: categorical + values: + - 16 + - 32 + - 64 + - 128 + batch_size: + distribution: categorical + values: + - 16 + - 32 + - 64 + n_epochs: + distribution: categorical + values: + - 100 + - 200 + - 300 + seed: + value: 42 + wandb_project: + value: "engiopt" diff --git a/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_heatconduction2d.yaml b/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_heatconduction2d.yaml new file mode 100644 index 00000000..1e884788 --- /dev/null +++ b/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_heatconduction2d.yaml @@ -0,0 +1,44 @@ +method: bayes +metric: + name: g_loss + goal: minimize +run_cap: 50 +parameters: + problem_id: + value: "heatconduction2d" + lr_gen: + distribution: log_uniform_values + min: 0.00005 + max: 0.005 + lr_disc: + distribution: log_uniform_values + min: 0.00005 + max: 0.005 + latent_dim: + distribution: categorical + values: + - 16 + - 32 + - 64 + - 128 + batch_size: + distribution: categorical + values: + - 16 + - 32 + - 64 + n_epochs: + distribution: categorical + values: + - 100 + - 200 + - 300 + track: + value: true + save_model: + value: true + seed: + value: 1 + wandb_project: + value: "lv-mmd-hp-tuning" +program: engiopt/cgan_cnn_2d/cgan_cnn_2d.py diff --git a/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_photonics2d.yaml b/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_photonics2d.yaml new file mode 100644 index 00000000..cc9d3a89 --- /dev/null +++ b/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_photonics2d.yaml @@ -0,0 +1,44 @@ +method: grid +command: + - "python" + - "engiopt/cgan_cnn_2d/cgan_cnn_2d.py" + - "--save_model" + - "${args}" + +metric: + name: g_loss + goal: minimize +parameters: + problem_id: + value: "photonics2d" + lr_gen: + distribution: log_uniform_values + min: 0.00005 + max: 0.005 + lr_disc: + distribution: log_uniform_values + min: 0.00005 + max: 0.005 + latent_dim: + distribution: categorical + values: + - 16 + - 32 + - 64 + - 128 + batch_size: + distribution: categorical + values: + - 16 + - 32 + - 64 + n_epochs: + distribution: categorical + values: + - 100 + - 200 + - 300 + seed: + value: 1 + wandb_project: + value: "engiopt" diff --git a/engiopt/cgan_cnn_3d/README.md b/engiopt/generators/cgan_cnn_3d/README.md similarity index 100% rename from engiopt/cgan_cnn_3d/README.md rename to engiopt/generators/cgan_cnn_3d/README.md diff --git a/engiopt/cgan_cnn_3d/__init__.py b/engiopt/generators/cgan_cnn_3d/__init__.py similarity index 100% rename from engiopt/cgan_cnn_3d/__init__.py rename to engiopt/generators/cgan_cnn_3d/__init__.py diff --git a/engiopt/generators/cgan_cnn_3d/adapter.py b/engiopt/generators/cgan_cnn_3d/adapter.py new file mode 100644 index 00000000..a8e5dc86 --- /dev/null +++ b/engiopt/generators/cgan_cnn_3d/adapter.py @@ -0,0 +1,66 @@ +"""`Generator` contract for the conditional 3D CNN GAN.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.cgan_cnn_3d.cgan_cnn_3d import Generator3D + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +def center_crop_3d(volumes: th.Tensor, target: tuple[int, ...]) -> th.Tensor: + """Center-crop `(n, D, H, W)` volumes down to `target`. + + The 3D generators upsample to a power-of-two grid (64) and the problem's + design space is smaller (51), so the padding introduced for the network's + convenience is removed symmetrically here. + """ + cropped = volumes + for axis, size in enumerate(target, start=1): + start = (cropped.shape[axis] - size) // 2 + cropped = cropped.narrow(axis, start, size) + return cropped + + +class CGANCNN3D(Generator): + """Conditional GAN with a 3D CNN generator.""" + + algo_id = "cgan_cnn_3d" + conditional = True + design_kinds = ("3d",) + checkpoint_files = ("generator_3d.pth",) + primary_state_key = "generator" + output_clip = (1e-3, 1.0) + + def __init__(self, net: Generator3D, latent_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> CGANCNN3D: + """Load a trained 3D cGAN generator from its checkpoint package.""" + config = resolved.run_config + net = Generator3D( + latent_dim=config["latent_dim"], + n_conds=len(problem.conditions_keys), + design_shape=problem.design_space.shape, + ) + net.load_state_dict(th.load(resolved.files["generator_3d.pth"], map_location=device)[cls.primary_state_key]) + net.eval().to(device) + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Generate volumes from noise plus conditions, then trim the padding.""" + cond = conditions.require_tensor(self.algo_id).reshape(n, self.n_conds, 1, 1, 1) + z = th.randn((n, self.latent_dim, 1, 1, 1), device=self.device, dtype=th.float) + volumes = self.net(z, cond).squeeze(1) + return center_crop_3d(volumes, self.design_shape) diff --git a/engiopt/cgan_cnn_3d/cgan_cnn_3d.py b/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py similarity index 100% rename from engiopt/cgan_cnn_3d/cgan_cnn_3d.py rename to engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py diff --git a/engiopt/cgan_cnn_3d/evaluate_cgan_cnn_3d.py b/engiopt/generators/cgan_cnn_3d/evaluate_cgan_cnn_3d.py similarity index 98% rename from engiopt/cgan_cnn_3d/evaluate_cgan_cnn_3d.py rename to engiopt/generators/cgan_cnn_3d/evaluate_cgan_cnn_3d.py index 7827010b..2fda1059 100644 --- a/engiopt/cgan_cnn_3d/evaluate_cgan_cnn_3d.py +++ b/engiopt/generators/cgan_cnn_3d/evaluate_cgan_cnn_3d.py @@ -12,7 +12,7 @@ import tyro from engiopt import metrics -from engiopt.cgan_cnn_3d.cgan_cnn_3d import Generator3D +from engiopt.generators.cgan_cnn_3d.cgan_cnn_3d import Generator3D from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions diff --git a/engiopt/cgan_cnn_3d/sweep_cgan_cnn_3d.yaml b/engiopt/generators/cgan_cnn_3d/sweep_cgan_cnn_3d.yaml similarity index 100% rename from engiopt/cgan_cnn_3d/sweep_cgan_cnn_3d.yaml rename to engiopt/generators/cgan_cnn_3d/sweep_cgan_cnn_3d.yaml diff --git a/engiopt/cgan_vae/README.md b/engiopt/generators/cgan_vae/README.md similarity index 100% rename from engiopt/cgan_vae/README.md rename to engiopt/generators/cgan_vae/README.md diff --git a/engiopt/cgan_vae/__init__.py b/engiopt/generators/cgan_vae/__init__.py similarity index 100% rename from engiopt/cgan_vae/__init__.py rename to engiopt/generators/cgan_vae/__init__.py diff --git a/engiopt/generators/cgan_vae/adapter.py b/engiopt/generators/cgan_vae/adapter.py new file mode 100644 index 00000000..a50dba7c --- /dev/null +++ b/engiopt/generators/cgan_vae/adapter.py @@ -0,0 +1,54 @@ +"""`Generator` contract for the conditional multi-view 3D VAE-GAN.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.cgan_cnn_3d.adapter import center_crop_3d +from engiopt.generators.cgan_vae.cgan_vae import Generator3D + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class CGANVAE(Generator): + """Conditional VAE-GAN hybrid with a 3D decoder.""" + + algo_id = "cgan_vae" + conditional = True + design_kinds = ("3d",) + checkpoint_files = ("multiview_3d_vaegan.pth",) + primary_state_key = "generator" + output_clip = (1e-3, 1.0) + + def __init__(self, net: Generator3D, latent_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> CGANVAE: + """Load the trained 3D decoder from its checkpoint package.""" + config = resolved.run_config + net = Generator3D( + latent_dim=config["latent_dim"], + n_conds=len(problem.conditions_keys), + design_shape=problem.design_space.shape, + ) + checkpoint = th.load(resolved.files["multiview_3d_vaegan.pth"], map_location=device) + net.load_state_dict(checkpoint[cls.primary_state_key]) + net.eval().to(device) + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Decode noise plus conditions into volumes, then trim the padding.""" + cond = conditions.require_tensor(self.algo_id).reshape(n, self.n_conds, 1, 1, 1) + z = th.randn((n, self.latent_dim, 1, 1, 1), device=self.device, dtype=th.float) + volumes = self.net(z, cond).squeeze(1) + return center_crop_3d(volumes, self.design_shape) diff --git a/engiopt/cgan_vae/cgan_vae.py b/engiopt/generators/cgan_vae/cgan_vae.py similarity index 100% rename from engiopt/cgan_vae/cgan_vae.py rename to engiopt/generators/cgan_vae/cgan_vae.py diff --git a/engiopt/cgan_vae/evaluate_cgan_vae.py b/engiopt/generators/cgan_vae/evaluate_cgan_vae.py similarity index 98% rename from engiopt/cgan_vae/evaluate_cgan_vae.py rename to engiopt/generators/cgan_vae/evaluate_cgan_vae.py index 7cfc64d5..a4fcebc3 100644 --- a/engiopt/cgan_vae/evaluate_cgan_vae.py +++ b/engiopt/generators/cgan_vae/evaluate_cgan_vae.py @@ -12,7 +12,7 @@ import tyro from engiopt import metrics -from engiopt.cgan_vae.cgan_vae import Generator3D +from engiopt.generators.cgan_vae.cgan_vae import Generator3D from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions diff --git a/engiopt/cgan_vae/sample_script.ipynb b/engiopt/generators/cgan_vae/sample_script.ipynb similarity index 100% rename from engiopt/cgan_vae/sample_script.ipynb rename to engiopt/generators/cgan_vae/sample_script.ipynb diff --git a/engiopt/cgan_vae/sweep_cgan_vae.yaml b/engiopt/generators/cgan_vae/sweep_cgan_vae.yaml similarity index 100% rename from engiopt/cgan_vae/sweep_cgan_vae.yaml rename to engiopt/generators/cgan_vae/sweep_cgan_vae.yaml diff --git a/engiopt/diffusion_1d/__init__.py b/engiopt/generators/diffusion_1d/__init__.py similarity index 100% rename from engiopt/diffusion_1d/__init__.py rename to engiopt/generators/diffusion_1d/__init__.py diff --git a/engiopt/generators/diffusion_1d/adapter.py b/engiopt/generators/diffusion_1d/adapter.py new file mode 100644 index 00000000..4555b49e --- /dev/null +++ b/engiopt/generators/diffusion_1d/adapter.py @@ -0,0 +1,77 @@ +"""`Generator` contract for the unconditional 1D diffusion model.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +from denoising_diffusion_pytorch import GaussianDiffusion1D +from denoising_diffusion_pytorch import Unet1D +import numpy as np +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import design_shape_of +from engiopt.core import Generator +from engiopt.generators.diffusion_1d.diffusion_1d import prepare_data + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +def _padding_for(length: int) -> int: + """Padding needed to make a 1D design length divisible by 8, as the UNet requires.""" + return (8 - length % 8) % 8 + + +class Diffusion1D(Generator): + """Unconditional denoising diffusion over 1D / dict-valued designs. + + The UNet needs a sequence length divisible by 8, so designs are padded for + generation and the padding is trimmed off again before evaluation. + """ + + algo_id = "diffusion_1d" + conditional = False + design_kinds = ("1d", "dict") + checkpoint_files = ("model.pth",) + primary_state_key = "model" + + def __init__(self, net: GaussianDiffusion1D, design_normalizer: Any, padding_size: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.design_normalizer = design_normalizer + self.padding_size = padding_size + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> Diffusion1D: + """Rebuild the 1D UNet and its diffusion wrapper from the checkpoint.""" + config = resolved.run_config + padding_size = _padding_for(design_shape_of(problem)[0]) + padded_shape = (design_shape_of(problem)[0] + padding_size,) + _, design_normalizer = prepare_data(problem, padding_size, device) + unet = Unet1D(dim=config["unet_dim"], channels=config["n_channels"]).to(device) + diffusion = GaussianDiffusion1D( + unet, + seq_length=int(np.prod(padded_shape)), + auto_normalize=config.get("auto_norm", True), + ).to(device) + diffusion.load_state_dict(th.load(resolved.files["model.pth"], map_location=device)[cls.primary_state_key]) + diffusion.eval() + return cls( + net=diffusion, + design_normalizer=design_normalizer, + padding_size=padding_size, + problem=problem, + device=device, + **base, + ) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: # noqa: ARG002 - unconditional by design + """Sample from the diffusion model, denormalize, and trim the UNet padding.""" + designs = self.net.sample(n).squeeze(1) + designs = self.design_normalizer.denormalize(designs) + if self.padding_size > 0: + designs = designs[:, : -self.padding_size] + return designs diff --git a/engiopt/diffusion_1d/diffusion_1d.py b/engiopt/generators/diffusion_1d/diffusion_1d.py similarity index 100% rename from engiopt/diffusion_1d/diffusion_1d.py rename to engiopt/generators/diffusion_1d/diffusion_1d.py diff --git a/engiopt/diffusion_1d/evaluate_diffusion_1d.py b/engiopt/generators/diffusion_1d/evaluate_diffusion_1d.py similarity index 98% rename from engiopt/diffusion_1d/evaluate_diffusion_1d.py rename to engiopt/generators/diffusion_1d/evaluate_diffusion_1d.py index 2a286c34..20b14581 100644 --- a/engiopt/diffusion_1d/evaluate_diffusion_1d.py +++ b/engiopt/generators/diffusion_1d/evaluate_diffusion_1d.py @@ -17,7 +17,7 @@ from engiopt import metrics from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.diffusion_1d.diffusion_1d import prepare_data +from engiopt.generators.diffusion_1d.diffusion_1d import prepare_data @dataclasses.dataclass diff --git a/engiopt/diffusion_2d_cond/__init__.py b/engiopt/generators/diffusion_2d_cond/__init__.py similarity index 100% rename from engiopt/diffusion_2d_cond/__init__.py rename to engiopt/generators/diffusion_2d_cond/__init__.py diff --git a/engiopt/generators/diffusion_2d_cond/adapter.py b/engiopt/generators/diffusion_2d_cond/adapter.py new file mode 100644 index 00000000..d5da0b83 --- /dev/null +++ b/engiopt/generators/diffusion_2d_cond/adapter.py @@ -0,0 +1,116 @@ +"""`Generator` contract for the conditional 2D diffusion model.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +from diffusers import UNet2DConditionModel +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import beta_schedule +from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import denormalize_designs_from_diffusion_range +from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import DiffusionSampler + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class Diffusion2DCond(Generator): + """Conditional denoising diffusion model over 2D designs. + + Sampling runs the full reverse chain, so this is the slowest generator in the + zoo by a wide margin -- which is exactly why cost belongs on the leaderboard + next to quality. + """ + + algo_id = "diffusion_2d_cond" + conditional = True + design_kinds = ("2d",) + checkpoint_files = ("model.pth",) + primary_state_key = "model" + + def __init__( + self, + net: UNet2DConditionModel, + sampler: DiffusionSampler, + num_timesteps: int, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.net = net + self.sampler = sampler + self.num_timesteps = num_timesteps + self.design_min: th.Tensor | None = None + self.design_max: th.Tensor | None = None + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> Diffusion2DCond: + """Rebuild the UNet and its noise schedule from the checkpoint package.""" + config = resolved.run_config + checkpoint = th.load(resolved.files["model.pth"], map_location=device) + net = UNet2DConditionModel( + sample_size=problem.design_space.shape, + in_channels=1, + out_channels=1, + cross_attention_dim=64, + block_out_channels=(32, 64, 128, 256), + down_block_types=( + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "DownBlock2D", + ), + up_block_types=("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"), + layers_per_block=config["layers_per_block"], + transformer_layers_per_block=1, + encoder_hid_dim=len(problem.conditions_keys), + only_cross_attention=True, + ).to(device) + net.load_state_dict(checkpoint[cls.primary_state_key]) + net.eval() + num_timesteps = config["num_timesteps"] + betas = beta_schedule( + t=num_timesteps, + start=1e-4, + end=0.02, + scale=1.0, + options={ + "cosine": config["noise_schedule"] == "cosine", + "exp_biasing": config["noise_schedule"] == "exp", + "exp_bias_factor": 1, + }, + ) + generator = cls( + net=net, + sampler=DiffusionSampler(num_timesteps, betas), + num_timesteps=num_timesteps, + problem=problem, + device=device, + **base, + ) + if "design_min" in checkpoint and "design_max" in checkpoint: + generator.design_min = checkpoint["design_min"].to(device) + generator.design_max = checkpoint["design_max"].to(device) + generator.output_clip = ( + checkpoint["design_min"].detach().cpu().numpy(), + checkpoint["design_max"].detach().cpu().numpy(), + ) + else: + generator.output_clip = (1e-3, 1.0) + return generator + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Run the reverse diffusion chain from noise, conditioned throughout.""" + cond = conditions.require_tensor(self.algo_id).reshape(n, 1, self.n_conds) + designs = th.randn((n, 1, *self.design_shape), device=self.device) + for step in reversed(range(self.num_timesteps)): + t = th.full((n,), step, device=self.device, dtype=th.long) + designs = self.sampler.sample_timestep(self.net, designs, t, cond) + designs = designs.squeeze(1) + if self.design_min is not None and self.design_max is not None: + designs = denormalize_designs_from_diffusion_range(designs, self.design_min, self.design_max) + return designs diff --git a/engiopt/diffusion_2d_cond/diffusion_2d_cond.py b/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py similarity index 100% rename from engiopt/diffusion_2d_cond/diffusion_2d_cond.py rename to engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py diff --git a/engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py b/engiopt/generators/diffusion_2d_cond/evaluate_diffusion_2d_cond.py similarity index 95% rename from engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py rename to engiopt/generators/diffusion_2d_cond/evaluate_diffusion_2d_cond.py index b70d5a59..a17a5f3d 100644 --- a/engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py +++ b/engiopt/generators/diffusion_2d_cond/evaluate_diffusion_2d_cond.py @@ -16,9 +16,9 @@ from engiopt import metrics from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.diffusion_2d_cond.diffusion_2d_cond import beta_schedule -from engiopt.diffusion_2d_cond.diffusion_2d_cond import denormalize_designs_from_diffusion_range -from engiopt.diffusion_2d_cond.diffusion_2d_cond import DiffusionSampler +from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import beta_schedule +from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import denormalize_designs_from_diffusion_range +from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import DiffusionSampler @dataclasses.dataclass diff --git a/engiopt/gan_1d/__init__.py b/engiopt/generators/gan_1d/__init__.py similarity index 100% rename from engiopt/gan_1d/__init__.py rename to engiopt/generators/gan_1d/__init__.py diff --git a/engiopt/generators/gan_1d/adapter.py b/engiopt/generators/gan_1d/adapter.py new file mode 100644 index 00000000..ae424cbd --- /dev/null +++ b/engiopt/generators/gan_1d/adapter.py @@ -0,0 +1,52 @@ +"""`Generator` contract for the unconditional 1D GAN.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import design_shape_of +from engiopt.core import Generator +from engiopt.generators.gan_1d.gan_1d import Generator as GAN1DNet +from engiopt.generators.gan_1d.gan_1d import prepare_data + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class GAN1D(Generator): + """Unconditional GAN over 1D / dict-valued designs such as airfoils.""" + + algo_id = "gan_1d" + conditional = False + design_kinds = ("1d", "dict") + checkpoint_files = ("generator.pth",) + primary_state_key = "generator" + + def __init__(self, net: GAN1DNet, latent_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> GAN1D: + """Load a trained GAN-1D generator, rebuilding its design normalizer.""" + config = resolved.run_config + _, design_normalizer = prepare_data(problem, device) + net = GAN1DNet( + latent_dim=config["latent_dim"], + design_shape=design_shape_of(problem), + design_normalizer=design_normalizer, + ).to(device) + net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.eval() + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: # noqa: ARG002 - unconditional by design + """Map noise through the generator, ignoring conditions.""" + z = th.randn((n, self.latent_dim), device=self.device) + return self.net(z) diff --git a/engiopt/gan_1d/evaluate_gan_1d.py b/engiopt/generators/gan_1d/evaluate_gan_1d.py similarity index 97% rename from engiopt/gan_1d/evaluate_gan_1d.py rename to engiopt/generators/gan_1d/evaluate_gan_1d.py index f9a8cdaa..a977c86f 100644 --- a/engiopt/gan_1d/evaluate_gan_1d.py +++ b/engiopt/generators/gan_1d/evaluate_gan_1d.py @@ -15,8 +15,8 @@ from engiopt import metrics from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.gan_1d.gan_1d import Generator -from engiopt.gan_1d.gan_1d import prepare_data +from engiopt.generators.gan_1d.gan_1d import Generator +from engiopt.generators.gan_1d.gan_1d import prepare_data @dataclasses.dataclass diff --git a/engiopt/gan_1d/gan_1d.py b/engiopt/generators/gan_1d/gan_1d.py similarity index 100% rename from engiopt/gan_1d/gan_1d.py rename to engiopt/generators/gan_1d/gan_1d.py diff --git a/engiopt/gan_2d/__init__.py b/engiopt/generators/gan_2d/__init__.py similarity index 100% rename from engiopt/gan_2d/__init__.py rename to engiopt/generators/gan_2d/__init__.py diff --git a/engiopt/generators/gan_2d/adapter.py b/engiopt/generators/gan_2d/adapter.py new file mode 100644 index 00000000..e791e40a --- /dev/null +++ b/engiopt/generators/gan_2d/adapter.py @@ -0,0 +1,46 @@ +"""`Generator` contract for the unconditional MLP GAN (2D).""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.gan_2d.gan_2d import Generator as GAN2DNet + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class GAN2D(Generator): + """Unconditional GAN with a fully-connected generator, for 2D problems.""" + + algo_id = "gan_2d" + conditional = False + design_kinds = ("2d",) + checkpoint_files = ("generator.pth",) + primary_state_key = "generator" + output_clip = (1e-3, 1.0) + + def __init__(self, net: GAN2DNet, latent_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> GAN2D: + """Load a trained GAN-2D generator from its checkpoint package.""" + config = resolved.run_config + net = GAN2DNet(latent_dim=config["latent_dim"], design_shape=problem.design_space.shape).to(device) + net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.eval() + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: # noqa: ARG002 - unconditional by design + """Map noise through the generator, ignoring conditions.""" + z = th.randn((n, self.latent_dim), device=self.device) + return self.net(z) diff --git a/engiopt/gan_2d/evaluate_gan_2d.py b/engiopt/generators/gan_2d/evaluate_gan_2d.py similarity index 98% rename from engiopt/gan_2d/evaluate_gan_2d.py rename to engiopt/generators/gan_2d/evaluate_gan_2d.py index 984fdf76..c9721a02 100644 --- a/engiopt/gan_2d/evaluate_gan_2d.py +++ b/engiopt/generators/gan_2d/evaluate_gan_2d.py @@ -14,7 +14,7 @@ from engiopt import metrics from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.gan_2d.gan_2d import Generator +from engiopt.generators.gan_2d.gan_2d import Generator @dataclasses.dataclass diff --git a/engiopt/gan_2d/gan_2d.py b/engiopt/generators/gan_2d/gan_2d.py similarity index 100% rename from engiopt/gan_2d/gan_2d.py rename to engiopt/generators/gan_2d/gan_2d.py diff --git a/engiopt/gan_bezier/__init__.py b/engiopt/generators/gan_bezier/__init__.py similarity index 100% rename from engiopt/gan_bezier/__init__.py rename to engiopt/generators/gan_bezier/__init__.py diff --git a/engiopt/generators/gan_bezier/adapter.py b/engiopt/generators/gan_bezier/adapter.py new file mode 100644 index 00000000..fe224852 --- /dev/null +++ b/engiopt/generators/gan_bezier/adapter.py @@ -0,0 +1,79 @@ +"""`Generator` contract for the BezierGAN airfoil model.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.gan_bezier.gan_bezier import Generator as BezierGANNet +from engiopt.generators.gan_bezier.gan_bezier import prepare_data +from engiopt.transforms import flatten_dict_factory + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + +_EPS = 1e-7 +"""Numerical floor used by the Bezier parameterization, matching training.""" + + +class GANBezier(Generator): + """BezierGAN over airfoil designs. + + The network emits Bezier control points plus a scalar angle of attack, so a + design is a dict rather than an array. It is flattened here into the single + vector the evaluator and metrics expect. + + The latent code is sampled uniformly rather than normally, matching how this + model was trained and previously evaluated. + """ + + algo_id = "gan_bezier" + conditional = False + design_kinds = ("dict",) + checkpoint_files = ("bezier_generator.pth",) + primary_state_key = "generator" + + def __init__(self, net: BezierGANNet, latent_dim: int, noise_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + self.noise_dim = noise_dim + self._flatten = flatten_dict_factory(self.problem, self.device) + + @classmethod + def build( + cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, n_prepare_samples: int = 50, **base: Any + ) -> GANBezier: + """Load a trained BezierGAN generator, rebuilding its scalar normalizer.""" + config = resolved.run_config + _, design_scalars_normalizer, _ = prepare_data(problem, n_prepare_samples, device) + net = BezierGANNet( + latent_dim=config["latent_dim"], + noise_dim=config["noise_dim"], + n_control_points=config["bezier_control_pts"], + n_data_points=problem.design_space["coords"].shape[1], + design_scalars_normalizer=design_scalars_normalizer, + eps=_EPS, + scalar_features=1, + ).to(device) + net.load_state_dict(th.load(resolved.files["bezier_generator.pth"], map_location=device)[cls.primary_state_key]) + net.eval() + return cls( + net=net, latent_dim=config["latent_dim"], noise_dim=config["noise_dim"], problem=problem, device=device, **base + ) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: # noqa: ARG002 - unconditional by design + """Generate Bezier airfoils and flatten the dict-valued designs.""" + latent = th.rand(n, self.latent_dim, device=self.device) + noise = 0.5 * th.randn(n, self.noise_dim, device=self.device) + coords, *_rest, alphas = self.net(latent, noise) + designs = [ + {"coords": coord, "angle_of_attack": alpha[0]} + for coord, alpha in zip(coords.detach().cpu().numpy(), alphas.detach().cpu().numpy()) + ] + return self._flatten(designs) diff --git a/engiopt/gan_bezier/evaluate_gan_bezier.py b/engiopt/generators/gan_bezier/evaluate_gan_bezier.py similarity index 97% rename from engiopt/gan_bezier/evaluate_gan_bezier.py rename to engiopt/generators/gan_bezier/evaluate_gan_bezier.py index d20cff5d..43361dd8 100644 --- a/engiopt/gan_bezier/evaluate_gan_bezier.py +++ b/engiopt/generators/gan_bezier/evaluate_gan_bezier.py @@ -15,8 +15,8 @@ from engiopt import metrics from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.gan_bezier.gan_bezier import Generator -from engiopt.gan_bezier.gan_bezier import prepare_data +from engiopt.generators.gan_bezier.gan_bezier import Generator +from engiopt.generators.gan_bezier.gan_bezier import prepare_data from engiopt.transforms import flatten_dict_factory if TYPE_CHECKING: diff --git a/engiopt/gan_bezier/gan_bezier.py b/engiopt/generators/gan_bezier/gan_bezier.py similarity index 100% rename from engiopt/gan_bezier/gan_bezier.py rename to engiopt/generators/gan_bezier/gan_bezier.py diff --git a/engiopt/gan_cnn_2d/__init__.py b/engiopt/generators/gan_cnn_2d/__init__.py similarity index 100% rename from engiopt/gan_cnn_2d/__init__.py rename to engiopt/generators/gan_cnn_2d/__init__.py diff --git a/engiopt/generators/gan_cnn_2d/adapter.py b/engiopt/generators/gan_cnn_2d/adapter.py new file mode 100644 index 00000000..fe284e18 --- /dev/null +++ b/engiopt/generators/gan_cnn_2d/adapter.py @@ -0,0 +1,52 @@ +"""`Generator` contract for the unconditional GAN with CNN backbone (2D).""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.gan_cnn_2d.gan_cnn_2d import Generator as GANCNN2DNet + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class GANCNN2D(Generator): + """Unconditional GAN with a CNN generator, for 2D topology problems. + + Conditions are accepted and ignored: the model was never told what to + optimize for. That is precisely what conditional-adherence metrics exist to + expose, so the mismatch is left visible rather than papered over. + """ + + algo_id = "gan_cnn_2d" + conditional = False + design_kinds = ("2d",) + checkpoint_files = ("generator.pth",) + primary_state_key = "generator" + output_clip = (1e-3, 1.0) + + def __init__(self, net: GANCNN2DNet, latent_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> GANCNN2D: + """Load a trained GAN-CNN generator from its checkpoint package.""" + config = resolved.run_config + net = GANCNN2DNet(latent_dim=config["latent_dim"], design_shape=problem.design_space.shape) + checkpoint = th.load(resolved.files["generator.pth"], map_location=device) + net.load_state_dict(checkpoint[cls.primary_state_key]) + net.eval().to(device) + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: # noqa: ARG002 - unconditional by design + """Map noise through the generator, ignoring conditions.""" + z = th.randn((n, self.latent_dim, 1, 1), device=self.device, dtype=th.float) + return self.net(z) diff --git a/engiopt/gan_cnn_2d/evaluate_gan_cnn_2d.py b/engiopt/generators/gan_cnn_2d/evaluate_gan_cnn_2d.py similarity index 98% rename from engiopt/gan_cnn_2d/evaluate_gan_cnn_2d.py rename to engiopt/generators/gan_cnn_2d/evaluate_gan_cnn_2d.py index 38619be2..6c71121f 100644 --- a/engiopt/gan_cnn_2d/evaluate_gan_cnn_2d.py +++ b/engiopt/generators/gan_cnn_2d/evaluate_gan_cnn_2d.py @@ -14,7 +14,7 @@ from engiopt import metrics from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.gan_cnn_2d.gan_cnn_2d import Generator +from engiopt.generators.gan_cnn_2d.gan_cnn_2d import Generator @dataclasses.dataclass diff --git a/engiopt/gan_cnn_2d/gan_cnn_2d.py b/engiopt/generators/gan_cnn_2d/gan_cnn_2d.py similarity index 100% rename from engiopt/gan_cnn_2d/gan_cnn_2d.py rename to engiopt/generators/gan_cnn_2d/gan_cnn_2d.py diff --git a/engiopt/pixel_cnn_pp_2d/__init__.py b/engiopt/generators/pixel_cnn_pp_2d/__init__.py similarity index 100% rename from engiopt/pixel_cnn_pp_2d/__init__.py rename to engiopt/generators/pixel_cnn_pp_2d/__init__.py diff --git a/engiopt/generators/pixel_cnn_pp_2d/adapter.py b/engiopt/generators/pixel_cnn_pp_2d/adapter.py new file mode 100644 index 00000000..6c48de74 --- /dev/null +++ b/engiopt/generators/pixel_cnn_pp_2d/adapter.py @@ -0,0 +1,83 @@ +"""`Generator` contract for the conditional PixelCNN++ model (2D).""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.pixel_cnn_pp_2d.pixel_cnn_pp_2d import PixelCNNpp +from engiopt.generators.pixel_cnn_pp_2d.pixel_cnn_pp_2d import sample_from_discretized_mix_logistic + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class PixelCNNpp2D(Generator): + """Autoregressive PixelCNN++ over 2D designs. + + Sampling is one forward pass per pixel, so cost scales with the design + resolution rather than with a fixed step count. On a 50x100 grid that is + 5000 sequential passes per batch, which makes this the natural upper bound + for the cost axis of the leaderboard. + """ + + algo_id = "pixel_cnn_pp_2d" + conditional = True + design_kinds = ("2d",) + checkpoint_files = ("model.pth",) + primary_state_key = "model" + output_clip = (1e-3, 1.0) + + def __init__(self, net: PixelCNNpp, nr_logistic_mix: int, sampling_batch_size: int = 16, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.nr_logistic_mix = nr_logistic_mix + self.sampling_batch_size = sampling_batch_size + + @classmethod + def build( + cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, sampling_batch_size: int = 16, **base: Any + ) -> PixelCNNpp2D: + """Load a trained PixelCNN++ from its checkpoint package.""" + config = resolved.run_config + net = PixelCNNpp( + nr_resnet=config["nr_resnet"], + nr_filters=config["nr_filters"], + nr_logistic_mix=config["nr_logistic_mix"], + resnet_nonlinearity=config["resnet_nonlinearity"], + dropout_p=config["dropout_p"], + input_channels=1, + nr_conditions=len(problem.conditions_keys), + ) + net.load_state_dict(th.load(resolved.files["model.pth"], map_location=device)[cls.primary_state_key]) + net.eval().to(device) + return cls( + net=net, + nr_logistic_mix=config["nr_logistic_mix"], + sampling_batch_size=sampling_batch_size, + problem=problem, + device=device, + **base, + ) + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Fill the design grid one pixel at a time, in memory-bounded batches.""" + cond = conditions.require_tensor(self.algo_id).reshape(n, self.n_conds, 1, 1) + height, width = self.design_shape[0], self.design_shape[1] + batches: list[th.Tensor] = [] + for start in range(0, n, self.sampling_batch_size): + end = min(n, start + self.sampling_batch_size) + data = th.zeros((end - start, 1, height, width), device=self.device) + for i in range(height): + for j in range(width): + out = self.net(data, cond[start:end]) + data[:, :, i, j] = sample_from_discretized_mix_logistic(out, self.nr_logistic_mix).data[:, :, i, j] + # Move each finished batch off the accelerator so long runs do not + # accumulate the whole sample set in device memory. + batches.append(data.cpu()) + return th.cat(batches, dim=0) diff --git a/engiopt/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py b/engiopt/generators/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py similarity index 96% rename from engiopt/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py rename to engiopt/generators/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py index 0a6b9ebf..89d3dc55 100644 --- a/engiopt/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py +++ b/engiopt/generators/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py @@ -14,8 +14,8 @@ from engiopt import metrics from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.pixel_cnn_pp_2d.pixel_cnn_pp_2d import PixelCNNpp -from engiopt.pixel_cnn_pp_2d.pixel_cnn_pp_2d import sample_from_discretized_mix_logistic +from engiopt.generators.pixel_cnn_pp_2d.pixel_cnn_pp_2d import PixelCNNpp +from engiopt.generators.pixel_cnn_pp_2d.pixel_cnn_pp_2d import sample_from_discretized_mix_logistic @dataclass diff --git a/engiopt/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py b/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py similarity index 100% rename from engiopt/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py rename to engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py diff --git a/engiopt/vqgan/__init__.py b/engiopt/generators/vqgan/__init__.py similarity index 100% rename from engiopt/vqgan/__init__.py rename to engiopt/generators/vqgan/__init__.py diff --git a/engiopt/generators/vqgan/adapter.py b/engiopt/generators/vqgan/adapter.py new file mode 100644 index 00000000..2f0ee99b --- /dev/null +++ b/engiopt/generators/vqgan/adapter.py @@ -0,0 +1,136 @@ +"""`Generator` contract for the VQGAN + transformer model.""" + +from __future__ import annotations + +import os +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.generators.vqgan.vqgan import VQGAN +from engiopt.generators.vqgan.vqgan import VQGANTransformer +from engiopt.transforms import drop_constant +from engiopt.transforms import normalize +from engiopt.transforms import resize_to + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class VQGANGenerator(Generator): + """Discrete latent VQGAN whose codes are sampled autoregressively by a transformer. + + Conditions are themselves quantized (by a second, condition-side VQGAN) into + the transformer's start-of-sequence tokens, so this model needs the original + condition columns rather than a ready-made tensor: it may drop constant + columns and re-normalize exactly as it did during training. + """ + + algo_id = "vqgan" + conditional = True + design_kinds = ("2d",) + checkpoint_files = ("vqgan.pth", "transformer.pth") + primary_state_key = "transformer" + output_clip = (1e-3, 1.0) + + def __init__(self, net: VQGANTransformer, latent_size: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_size = latent_size + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> VQGANGenerator: + """Rebuild the design VQGAN, the condition VQGAN, and the transformer.""" + config = resolved.run_config + vqgan = VQGAN( + device=device, + is_c=False, + encoder_channels=config["encoder_channels"], + encoder_start_resolution=config["image_size"], + encoder_attn_resolutions=config["encoder_attn_resolutions"], + encoder_num_res_blocks=config["encoder_num_res_blocks"], + decoder_channels=config["decoder_channels"], + decoder_start_resolution=config["latent_size"], + decoder_attn_resolutions=config["decoder_attn_resolutions"], + decoder_num_res_blocks=config["decoder_num_res_blocks"], + image_channels=config["image_channels"], + latent_dim=config["latent_dim"], + num_codebook_vectors=config["num_codebook_vectors"], + ) + vqgan.load_state_dict(th.load(resolved.files["vqgan.pth"], map_location=device, weights_only=False)["vqgan"]) + vqgan.eval().to(device) + cvqgan = VQGAN( + device=device, + is_c=True, + cond_feature_map_dim=config["cond_feature_map_dim"], + cond_dim=config["cond_dim"], + cond_hidden_dim=config["cond_hidden_dim"], + cond_latent_dim=config["cond_latent_dim"], + cond_codebook_vectors=config["cond_codebook_vectors"], + ) + cvqgan_state = cls._load_cvqgan(resolved.root_dir, config=config, device=device) + if cvqgan_state is not None: + cvqgan.load_state_dict(cvqgan_state) + cvqgan.eval().to(device) + net = VQGANTransformer( + conditional=config["conditional"], + vqgan=vqgan, + cvqgan=cvqgan, + image_size=config["image_size"], + decoder_channels=config["decoder_channels"], + cond_feature_map_dim=config["cond_feature_map_dim"], + num_codebook_vectors=config["num_codebook_vectors"], + n_layer=config["n_layer"], + n_head=config["n_head"], + n_embd=config["n_embd"], + dropout=config["dropout"], + ) + net.load_state_dict( + th.load(resolved.files["transformer.pth"], map_location=device, weights_only=False)[cls.primary_state_key] + ) + net.eval().to(device) + return cls(net=net, latent_size=config["latent_size"], problem=problem, device=device, **base) + + @classmethod + def _load_cvqgan(cls, package_root: str, *, config: dict[str, Any], device: th.device) -> dict[str, Any] | None: + """Load the condition-side VQGAN, which ships inside the same package. + + Raises: + FileNotFoundError: If the model is conditional but the package has + no `cvqgan.pth`, which would make its conditions unusable. + """ + path = os.path.join(package_root, "cvqgan.pth") + if not os.path.exists(path): + if config["conditional"]: + raise FileNotFoundError(f"Conditional VQGAN needs cvqgan.pth, missing from {package_root}") + return None + return th.load(path, map_location=device, weights_only=False)["cvqgan"] + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Sample a full grid of latent codes autoregressively, then decode it.""" + if self.run_config["conditional"]: + start_tokens = self.net.encode_to_z(x=self._condition_tensor(conditions), is_c=True)[1] + else: + start_tokens = th.ones(n, 1, dtype=th.int64, device=self.device) * self.net.sos_token + codes = self.net.sample( + x=th.empty(n, 0, dtype=th.int64, device=self.device), + c=start_tokens, + steps=self.latent_size**2, + ) + return resize_to(data=self.net.z_to_image(codes), h=self.design_shape[0], w=self.design_shape[1]) + + def _condition_tensor(self, conditions: ConditionBatch) -> th.Tensor: + """Rebuild the condition tensor exactly as training preprocessed it.""" + dataset = conditions.dataset + if dataset is None: + return conditions.require_tensor(self.algo_id) + columns = dataset.column_names + if self.run_config["drop_constant_conditions"]: + dataset, columns = drop_constant(dataset, columns) + if self.run_config["normalize_conditions"]: + dataset, _mean, _std = normalize(dataset, columns) + return th.stack([th.as_tensor(dataset[column][:]).float() for column in columns], dim=1).to(self.device) diff --git a/engiopt/vqgan/evaluate_vqgan.py b/engiopt/generators/vqgan/evaluate_vqgan.py similarity index 98% rename from engiopt/vqgan/evaluate_vqgan.py rename to engiopt/generators/vqgan/evaluate_vqgan.py index d42f7afd..61dbc693 100644 --- a/engiopt/vqgan/evaluate_vqgan.py +++ b/engiopt/generators/vqgan/evaluate_vqgan.py @@ -17,8 +17,8 @@ from engiopt.transforms import drop_constant from engiopt.transforms import normalize from engiopt.transforms import resize_to -from engiopt.vqgan.vqgan import VQGAN -from engiopt.vqgan.vqgan import VQGANTransformer +from engiopt.generators.vqgan.vqgan import VQGAN +from engiopt.generators.vqgan.vqgan import VQGANTransformer @dataclasses.dataclass diff --git a/engiopt/vqgan/utils.py b/engiopt/generators/vqgan/utils.py similarity index 100% rename from engiopt/vqgan/utils.py rename to engiopt/generators/vqgan/utils.py diff --git a/engiopt/vqgan/vqgan.py b/engiopt/generators/vqgan/vqgan.py similarity index 98% rename from engiopt/vqgan/vqgan.py rename to engiopt/generators/vqgan/vqgan.py index 7d8872b3..fe141a68 100644 --- a/engiopt/vqgan/vqgan.py +++ b/engiopt/generators/vqgan/vqgan.py @@ -39,18 +39,18 @@ from engiopt.transforms import drop_constant from engiopt.transforms import normalize from engiopt.transforms import resize_to -from engiopt.vqgan.utils import Codebook -from engiopt.vqgan.utils import Discriminator -from engiopt.vqgan.utils import DownSampleBlock -from engiopt.vqgan.utils import GPT -from engiopt.vqgan.utils import GPTConfig -from engiopt.vqgan.utils import GreyscaleLPIPS -from engiopt.vqgan.utils import GroupNorm -from engiopt.vqgan.utils import LinearCombo -from engiopt.vqgan.utils import NonLocalBlock -from engiopt.vqgan.utils import ResidualBlock -from engiopt.vqgan.utils import Swish -from engiopt.vqgan.utils import UpSampleBlock +from engiopt.generators.vqgan.utils import Codebook +from engiopt.generators.vqgan.utils import Discriminator +from engiopt.generators.vqgan.utils import DownSampleBlock +from engiopt.generators.vqgan.utils import GPT +from engiopt.generators.vqgan.utils import GPTConfig +from engiopt.generators.vqgan.utils import GreyscaleLPIPS +from engiopt.generators.vqgan.utils import GroupNorm +from engiopt.generators.vqgan.utils import LinearCombo +from engiopt.generators.vqgan.utils import NonLocalBlock +from engiopt.generators.vqgan.utils import ResidualBlock +from engiopt.generators.vqgan.utils import Swish +from engiopt.generators.vqgan.utils import UpSampleBlock @dataclass diff --git a/example_easy_model.ipynb b/example_easy_model.ipynb index 6f37108b..67e0810f 100644 --- a/example_easy_model.ipynb +++ b/example_easy_model.ipynb @@ -246,7 +246,7 @@ "import torch as th\n", "\n", "from engiopt import metrics\n", - "from engiopt.cgan_cnn_2d.cgan_cnn_2d import Generator\n", + "from engiopt.generators.cgan_cnn_2d.cgan_cnn_2d import Generator\n", "import wandb\n", "\n", "if th.cuda.is_available():\n", diff --git a/example_hard_model.ipynb b/example_hard_model.ipynb index ba404ce5..78897e81 100644 --- a/example_hard_model.ipynb +++ b/example_hard_model.ipynb @@ -225,8 +225,8 @@ "import torch as th\n", "\n", "from engiopt import metrics\n", - "from engiopt.diffusion_2d_cond.diffusion_2d_cond import beta_schedule\n", - "from engiopt.diffusion_2d_cond.diffusion_2d_cond import DiffusionSampler\n", + "from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import beta_schedule\n", + "from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import DiffusionSampler\n", "import wandb\n", "\n", "if th.cuda.is_available():\n", diff --git a/tests/test_cgan_output_activation.py b/tests/test_cgan_output_activation.py index 2a3d788e..58d0c679 100644 --- a/tests/test_cgan_output_activation.py +++ b/tests/test_cgan_output_activation.py @@ -7,14 +7,14 @@ ("module_name", "latent_shape", "condition_shape", "design_shape", "activation_getter"), [ ( - "engiopt.cgan_2d.cgan_2d", + "engiopt.generators.cgan_2d.cgan_2d", (2, 4), (2, 1), (3, 3), lambda model: model.model[-1], ), ( - "engiopt.cgan_cnn_2d.cgan_cnn_2d", + "engiopt.generators.cgan_cnn_2d.cgan_cnn_2d", (2, 4, 1, 1), (2, 1, 1, 1), (8, 8), @@ -42,14 +42,14 @@ def test_cgan_generators_keep_tanh_as_default( ("module_name", "latent_shape", "condition_shape", "design_shape", "activation_getter"), [ ( - "engiopt.cgan_2d.cgan_2d", + "engiopt.generators.cgan_2d.cgan_2d", (2, 4), (2, 1), (3, 3), lambda model: model.model[-1], ), ( - "engiopt.cgan_cnn_2d.cgan_cnn_2d", + "engiopt.generators.cgan_cnn_2d.cgan_cnn_2d", (2, 4, 1, 1), (2, 1, 1, 1), (8, 8), diff --git a/tests/test_diffusion_fixes.py b/tests/test_diffusion_fixes.py index 8fc3f746..c87422ef 100644 --- a/tests/test_diffusion_fixes.py +++ b/tests/test_diffusion_fixes.py @@ -10,7 +10,7 @@ def test_diffusion_1d_batch_is_normalized_before_forward() -> None: """The 1D diffusion model should receive [0, 1] designs before auto-normalization.""" th = pytest.importorskip("torch") - diffusion_1d = pytest.importorskip("engiopt.diffusion_1d.diffusion_1d") + diffusion_1d = pytest.importorskip("engiopt.generators.diffusion_1d.diffusion_1d") designs = th.tensor([[2.0, 4.0], [6.0, 8.0]]) normalizer = diffusion_1d.Normalizer( @@ -18,7 +18,7 @@ def test_diffusion_1d_batch_is_normalized_before_forward() -> None: max_val=th.tensor([6.0, 8.0]), ) - batch = diffusion_1d._prepare_diffusion_batch(designs, normalizer) # noqa: SLF001 + batch = diffusion_1d._prepare_diffusion_batch(designs, normalizer) expected_shape = (2, 1, 2) assert batch.shape == expected_shape @@ -28,7 +28,7 @@ def test_diffusion_1d_batch_is_normalized_before_forward() -> None: def test_default_linear_schedule_reaches_nearly_pure_noise() -> None: """The default noisiest training step should be close to the pure-noise sampling prior.""" th = pytest.importorskip("torch") - diffusion_2d = pytest.importorskip("engiopt.diffusion_2d_cond.diffusion_2d_cond") + diffusion_2d = pytest.importorskip("engiopt.generators.diffusion_2d_cond.diffusion_2d_cond") betas = diffusion_2d.beta_schedule( t=diffusion_2d.Args().num_timesteps, @@ -46,7 +46,7 @@ def test_default_linear_schedule_reaches_nearly_pure_noise() -> None: def test_diffusion_2d_normalizer_round_trips_arbitrary_bounds() -> None: """2D diffusion should not assume designs are already in [0, 1].""" th = pytest.importorskip("torch") - diffusion_2d = pytest.importorskip("engiopt.diffusion_2d_cond.diffusion_2d_cond") + diffusion_2d = pytest.importorskip("engiopt.generators.diffusion_2d_cond.diffusion_2d_cond") designs = th.tensor([[-2.0, 1.0], [4.0, 10.0]]) design_min = designs.min() @@ -63,7 +63,7 @@ def test_diffusion_2d_normalizer_round_trips_arbitrary_bounds() -> None: def test_diffusion_2d_uses_problem_design_bounds_when_available() -> None: """EngiBench design_space bounds should be the normalization source of truth.""" th = pytest.importorskip("torch") - diffusion_2d = pytest.importorskip("engiopt.diffusion_2d_cond.diffusion_2d_cond") + diffusion_2d = pytest.importorskip("engiopt.generators.diffusion_2d_cond.diffusion_2d_cond") fallback_designs = th.tensor([[0.2, 0.8]]) problem = SimpleNamespace( @@ -82,7 +82,7 @@ def test_diffusion_2d_uses_problem_design_bounds_when_available() -> None: def test_diffusion_step_sample_uses_fresh_noise(monkeypatch: pytest.MonkeyPatch) -> None: """The reverse step should clip predicted x0 and use fresh posterior noise.""" th = pytest.importorskip("torch") - diffusion_2d = pytest.importorskip("engiopt.diffusion_2d_cond.diffusion_2d_cond") + diffusion_2d = pytest.importorskip("engiopt.generators.diffusion_2d_cond.diffusion_2d_cond") betas = diffusion_2d.beta_schedule(t=4, start=1e-4, end=0.02) sampler = diffusion_2d.DiffusionSampler(t=4, betas=betas) From 35f1041feda9126e28566b39690c041ea63eed26 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Fri, 24 Jul 2026 13:57:13 +0200 Subject: [PATCH 02/24] Add the Generator contract and make HuggingFace the only checkpoint backend engiopt/core.py defines what a model must provide to be evaluated: a few declarative attributes plus build() and _sample(). Mirrors engibench/core.py. engiopt/utils/all_generators.py mirrors engibench/utils/all_problems.py. checkpoint_store drops the W&B read-fallback and the write-only run-scoped path, which no loader ever read. Checkpoints are now filed per hyperparameter configuration at {problem}/cfg_{fingerprint}/seed_{n}, with the canonical {problem}/seed_{n} claimed only by default-hyperparameter runs so a sweep cannot redefine what a bare model name means. Evaluation metrics can be attached to a package, keeping a checkpoint self-describing. Co-Authored-By: Claude Opus 4.8 (1M context) --- engiopt/checkpoint_store.py | 314 ++++++++++------------ engiopt/core.py | 447 ++++++++++++++++++++++++++++++++ engiopt/generators/__init__.py | 8 + engiopt/utils/__init__.py | 1 + engiopt/utils/all_generators.py | 99 +++++++ 5 files changed, 696 insertions(+), 173 deletions(-) create mode 100644 engiopt/core.py create mode 100644 engiopt/generators/__init__.py create mode 100644 engiopt/utils/__init__.py create mode 100644 engiopt/utils/all_generators.py diff --git a/engiopt/checkpoint_store.py b/engiopt/checkpoint_store.py index 7ec3b88b..8fcea504 100644 --- a/engiopt/checkpoint_store.py +++ b/engiopt/checkpoint_store.py @@ -1,4 +1,14 @@ -"""Shared checkpoint save/load helpers for EngiOpt model artifacts.""" +"""Shared checkpoint save/load helpers for EngiOpt model artifacts. + +HuggingFace is the single home for anything that has to be *reloaded or +compared*: checkpoints, their run configs, and their evaluation metrics. +Weights & Biases, when enabled, hosts the things you only ever *look at* -- +loss curves, sample images, and a pointer back to the HF package. + +Each run is filed under `{problem_id}/cfg_{fingerprint}/seed_{seed}`, so every +hyperparameter setting is independently addressable, plus the canonical +`{problem_id}/seed_{seed}` when the run used the training script's defaults. +""" from __future__ import annotations @@ -15,15 +25,18 @@ import wandb +METRICS_FILE = "metrics.json" +"""Evaluation scores stored alongside the weights they describe.""" + CheckpointBackend = Literal["hf", "none"] -ModelSource = Literal["auto", "hf", "wandb", "local"] +ModelSource = Literal["auto", "hf", "local"] @dataclass(frozen=True) class ResolvedCheckpoint: """Resolved checkpoint package with local files and serialized config.""" - source: Literal["hf", "wandb", "local"] + source: Literal["hf", "local"] root_dir: str files: dict[str, str] run_config: dict[str, Any] @@ -45,14 +58,7 @@ def build_hf_package_path(problem_id: str, seed: int, extra_parts: list[str] | N return "/".join(parts) -def build_hf_run_package_path(package_path: str, wandb_run_id: str | None) -> str | None: - """Return an immutable run-specific package path when a W&B run id is available.""" - if not wandb_run_id: - return None - return f"{package_path}/run_{_sanitize_path_component(wandb_run_id)}" - - -def save_checkpoint_package( # noqa: PLR0913 +def save_checkpoint_package( *, checkpoint_backend: CheckpointBackend, hf_entity: str, @@ -66,9 +72,24 @@ def save_checkpoint_package( # noqa: PLR0913 metadata: dict[str, Any] | None = None, primary_files: list[str] | None = None, extra_path_parts: list[str] | None = None, + config_fingerprint: str | None = None, + is_default_config: bool = True, ) -> dict[str, Any]: """Save a checkpoint package to HuggingFace. + A run is filed under up to three paths in its model repo: + + - `{problem_id}/cfg_{fingerprint}/seed_{seed}` -- this exact set of + hyperparameters. Every configuration in a sweep gets its own, so runs + never overwrite each other. + - `{problem_id}/seed_{seed}` -- the canonical location that a plain + `from_pretrained(problem, seed=...)` reads. Written **only** when the run + used the training script's default hyperparameters, so a sweep cannot + redefine what the bare model name means. + Pass `config_fingerprint` and `is_default_config` together via + `engiopt.core.checkpoint_identity(args)`. Omitting them keeps the older + behaviour of writing only the canonical path. + W&B is no longer a checkpoint storage backend; the active W&B run still receives a summary pointing at the HF package for traceability. """ @@ -76,9 +97,9 @@ def save_checkpoint_package( # noqa: PLR0913 "checkpoint_backend": checkpoint_backend, "hf_repo_id": None, "hf_package_path": None, - "hf_run_package_path": None, + "hf_config_package_path": None, "hf_revision": None, - "hf_run_revision": None, + "hf_config_revision": None, } metadata_payload = _build_metadata( problem_id=problem_id, @@ -99,33 +120,42 @@ def save_checkpoint_package( # noqa: PLR0913 metadata_payload["hf_repo_id"] = repo_id metadata_payload["hf_package_path"] = package_path - run_package_path = build_hf_run_package_path(package_path, metadata_payload.get("wandb_run_id")) - info["hf_run_package_path"] = run_package_path - if run_package_path is not None: - metadata_payload["hf_run_package_path"] = run_package_path - run_revision = _upload_package_to_hf( + # Every configuration gets its own addressable location, so a sweep's + # runs cannot overwrite one another. + if config_fingerprint: + config_package_path = build_hf_package_path( + problem_id, seed, [*(extra_path_parts or []), f"cfg_{config_fingerprint}"] + ) + info["hf_config_package_path"] = config_package_path + metadata_payload["config_fingerprint"] = config_fingerprint + metadata_payload["hf_config_package_path"] = config_package_path + info["hf_config_revision"] = _upload_package_to_hf( repo_id=repo_id, hf_private=hf_private, - package_path=run_package_path, + package_path=config_package_path, checkpoint_files=checkpoint_files, run_config=run_config, metadata=metadata_payload, algo=algo, ) - info["hf_run_revision"] = run_revision - metadata_payload["hf_run_revision"] = run_revision + metadata_payload["hf_config_revision"] = info["hf_config_revision"] - revision = _upload_package_to_hf( - repo_id=repo_id, - hf_private=hf_private, - package_path=package_path, - checkpoint_files=checkpoint_files, - run_config=run_config, - metadata=metadata_payload, - algo=algo, - ) - info["hf_revision"] = revision - metadata_payload["hf_revision"] = revision + # The canonical path defines what the bare model name resolves to, so + # only a default-hyperparameter run may claim it. + if is_default_config: + revision = _upload_package_to_hf( + repo_id=repo_id, + hf_private=hf_private, + package_path=package_path, + checkpoint_files=checkpoint_files, + run_config=run_config, + metadata=metadata_payload, + algo=algo, + ) + info["hf_revision"] = revision + metadata_payload["hf_revision"] = revision + else: + info["hf_package_path"] = None if wandb.run is not None: _log_checkpoint_summary_to_wandb(metadata_payload, info) @@ -133,7 +163,7 @@ def save_checkpoint_package( # noqa: PLR0913 return info -def resolve_named_checkpoint( # noqa: PLR0913 +def resolve_named_checkpoint( *, model_source: ModelSource, problem_id: str, @@ -142,22 +172,27 @@ def resolve_named_checkpoint( # noqa: PLR0913 hf_entity: str, hf_repo_prefix: str, required_files: list[str], - wandb_project: str, - wandb_entity: str | None, - wandb_artifact_names: dict[str, str], - wandb_config_artifact_name: str | None = None, - wandb_artifact_alias: str | None = None, local_model_dir: str | None = None, extra_path_parts: list[str] | None = None, ) -> ResolvedCheckpoint: """Resolve a checkpoint package by the standard EngiOpt problem/algo/seed naming. - ``wandb_artifact_alias`` overrides the default ``f"seed_{seed}"`` alias used when - falling back to legacy W&B artifacts. Callers with custom alias schemes (e.g. - ``f"seed_{seed}_rec{r}_perf{p}"``) pass it here so the read-fallback resolves - historical artifacts that pre-date the HF cutover. + Args: + model_source: `auto` tries HuggingFace then any local directory; `hf` or + `local` restrict it to one. + problem_id: EngiBench problem the checkpoint was trained on. + algo: Model family, which selects the HF repo. + seed: Training seed. + hf_entity: HF org/user holding the checkpoint repos. + hf_repo_prefix: Prefix of the per-model-family repo. + required_files: Files the package must contain. + local_model_dir: Directory to load from instead of the Hub. + extra_path_parts: Path components identifying one configuration; see + `engiopt.core.config_path_parts`. + + Raises: + FileNotFoundError: If no backend could supply the package. """ - alias = wandb_artifact_alias or f"seed_{seed}" errors: list[str] = [] if model_source in {"auto", "hf"}: try: @@ -171,21 +206,6 @@ def resolve_named_checkpoint( # noqa: PLR0913 raise errors.append(f"hf: {exc}") - if model_source in {"auto", "wandb"}: - try: - return _resolve_wandb_package( - _required_files=required_files, - artifact_names=wandb_artifact_names, - wandb_project=wandb_project, - wandb_entity=wandb_entity, - alias=alias, - config_artifact_name=wandb_config_artifact_name, - ) - except Exception as exc: - if model_source == "wandb": - raise - errors.append(f"wandb: {exc}") - if local_model_dir is not None and model_source in {"auto", "local"}: return _resolve_local_package(local_model_dir, required_files) @@ -198,39 +218,26 @@ def resolve_checkpoint_reference( model_source: ModelSource, model_ref: str, required_files: list[str] | None = None, - active_wandb_run: wandb.sdk.wandb_run.Run | None = None, ) -> ResolvedCheckpoint: - """Resolve a checkpoint package from an explicit HF/W&B/local reference.""" + """Resolve a checkpoint package from an explicit `hf://` or local reference. + + Raises: + ValueError: If the reference cannot be interpreted. + """ inferred_source = model_source - normalized_ref = model_ref if model_source == "auto": - if model_ref.startswith("hf://"): - inferred_source = "hf" - elif model_ref.startswith("wandb://"): - inferred_source = "wandb" - elif os.path.isdir(model_ref): - inferred_source = "local" - else: - inferred_source = "wandb" + inferred_source = "local" if os.path.isdir(model_ref.removeprefix("file://")) else "hf" if inferred_source == "hf": repo_id, package_path = _parse_hf_reference(model_ref) return _resolve_hf_package(repo_id=repo_id, package_path=package_path, required_files=required_files or []) - if inferred_source == "wandb": - artifact_path = model_ref.removeprefix("wandb://") - return _resolve_wandb_reference( - artifact_path=artifact_path, - required_files=required_files or [], - active_wandb_run=active_wandb_run, - ) if inferred_source == "local": - normalized_ref = model_ref.removeprefix("file://") - return _resolve_local_package(normalized_ref, required_files or []) + return _resolve_local_package(model_ref.removeprefix("file://"), required_files or []) raise ValueError(f"Unsupported model source: {model_source}") -def _upload_package_to_hf( # noqa: PLR0913 +def _upload_package_to_hf( *, repo_id: str, hf_private: bool, @@ -279,85 +286,6 @@ def _resolve_local_package(local_model_dir: str, required_files: list[str]) -> R return _load_package_from_directory(root_dir=local_model_dir, required_files=required_files, source="local") -def _resolve_wandb_package( - *, - _required_files: list[str], - artifact_names: dict[str, str], - wandb_project: str, - wandb_entity: str | None, - alias: str, - config_artifact_name: str | None, -) -> ResolvedCheckpoint: - api = wandb.Api() - files: dict[str, str] = {} - for file_name, artifact_name in artifact_names.items(): - artifact_path = _build_wandb_artifact_path( - artifact_name=artifact_name, - wandb_project=wandb_project, - wandb_entity=wandb_entity, - alias=alias, - ) - artifact = api.artifact(artifact_path, type="model") - artifact_dir = artifact.download() - files[file_name] = os.path.join(artifact_dir, file_name) - - config_artifact = config_artifact_name or next(iter(artifact_names.values())) - config_artifact_path = _build_wandb_artifact_path( - artifact_name=config_artifact, - wandb_project=wandb_project, - wandb_entity=wandb_entity, - alias=alias, - ) - artifact = api.artifact(config_artifact_path, type="model") - run = artifact.logged_by() - if run is None or not hasattr(run, "config"): - raise ValueError("Failed to retrieve W&B run config from artifact") - - return ResolvedCheckpoint( - source="wandb", - root_dir=os.path.dirname(next(iter(files.values()))), - files=files, - run_config=dict(run.config), - metadata={ - "source": "wandb", - "artifact_paths": { - file_name: _build_wandb_artifact_path( - artifact_name=artifact_name, - wandb_project=wandb_project, - wandb_entity=wandb_entity, - alias=alias, - ) - for file_name, artifact_name in artifact_names.items() - }, - }, - ) - - -def _resolve_wandb_reference( - *, - artifact_path: str, - required_files: list[str], - active_wandb_run: wandb.sdk.wandb_run.Run | None, -) -> ResolvedCheckpoint: - artifact = ( - active_wandb_run.use_artifact(artifact_path, type="model") - if active_wandb_run is not None - else wandb.Api().artifact(artifact_path, type="model") - ) - artifact_dir = artifact.download() - files = _discover_reference_files(artifact_dir, required_files) - run = artifact.logged_by() - if run is None or not hasattr(run, "config"): - raise ValueError("Failed to retrieve W&B run config from artifact reference") - return ResolvedCheckpoint( - source="wandb", - root_dir=artifact_dir, - files=files, - run_config=dict(run.config), - metadata={"source": "wandb", "artifact_path": artifact_path}, - ) - - def _load_package_from_directory( *, root_dir: str, required_files: list[str], source: Literal["hf", "local"] ) -> ResolvedCheckpoint: @@ -417,18 +345,7 @@ def _parse_hf_reference(model_ref: str) -> tuple[str, str]: return repo_id, package_path -def _build_wandb_artifact_path( - *, - artifact_name: str, - wandb_project: str, - wandb_entity: str | None, - alias: str, -) -> str: - project_path = f"{wandb_entity}/{wandb_project}" if wandb_entity is not None else wandb_project - return f"{project_path}/{artifact_name}:{alias}" - - -def _build_metadata( # noqa: PLR0913 +def _build_metadata( *, problem_id: str, algo: str, @@ -502,16 +419,20 @@ def _repo_readme_text(algo: str) -> str: def _log_checkpoint_summary_to_wandb(metadata: dict[str, Any], info: dict[str, Any]) -> None: + """Point the W&B run at the HF package holding its weights. + + W&B stores the media -- loss curves and sample images -- while the weights + and metrics live on HuggingFace, so each side records where the other is. + """ if wandb.run is None: return wandb.summary["checkpoint_backend"] = info["checkpoint_backend"] if info["hf_repo_id"] is not None: wandb.summary["hf_repo_id"] = info["hf_repo_id"] wandb.summary["hf_package_path"] = info["hf_package_path"] + wandb.summary["hf_config_package_path"] = info["hf_config_package_path"] wandb.summary["hf_revision"] = info["hf_revision"] - if info["hf_run_package_path"] is not None: - wandb.summary["hf_run_package_path"] = info["hf_run_package_path"] - wandb.summary["hf_run_revision"] = info["hf_run_revision"] + wandb.summary["hf_config_revision"] = info["hf_config_revision"] wandb.summary["checkpoint_primary_files"] = metadata["primary_files"] @@ -527,3 +448,50 @@ def _read_json(path: str) -> dict[str, Any]: def _sanitize_path_component(value: str) -> str: return value.replace(os.sep, "_").replace(" ", "_") + + +def publish_checkpoint_metrics( + *, + hf_entity: str, + hf_repo_prefix: str, + problem_id: str, + algo: str, + seed: int, + metrics: dict[str, Any], + extra_path_parts: list[str] | None = None, + token: str | None = None, +) -> str: + """Attach evaluation metrics to a checkpoint package on HuggingFace. + + Metrics live beside the weights they describe, so a checkpoint is + self-describing: whoever downloads it can see how it scored without + consulting the leaderboard or a W&B run. + + Args: + hf_entity: HF org/user holding the checkpoint repos. + hf_repo_prefix: Prefix of the per-model-family repo. + problem_id: Problem the checkpoint was evaluated on. + algo: Model family, selecting the repo. + seed: Training seed. + metrics: Scores to record, typically one leaderboard row. + extra_path_parts: Configuration path components, from + `engiopt.core.config_path_parts`. + token: HF token; falls back to the ambient login. + + Returns: + The path written inside the repo. + """ + repo_id = build_hf_repo_id(hf_entity, hf_repo_prefix, algo) + package_path = build_hf_package_path(problem_id, seed, extra_path_parts) + api = HfApi(token=token) + with tempfile.TemporaryDirectory() as tmp: + local = Path(tmp) / METRICS_FILE + _write_json(local, metrics) + api.upload_file( + path_or_fileobj=str(local), + path_in_repo=f"{package_path}/{METRICS_FILE}", + repo_id=repo_id, + repo_type="model", + commit_message=f"Evaluation metrics for {problem_id}/seed_{seed}", + ) + return f"{package_path}/{METRICS_FILE}" diff --git a/engiopt/core.py b/engiopt/core.py new file mode 100644 index 00000000..97995ff3 --- /dev/null +++ b/engiopt/core.py @@ -0,0 +1,447 @@ +"""The `Generator` contract: what a model implements to join the EngiOpt leaderboard. + +A generator is *not* a training loop. Training scripts stay single-file and +self-contained (the CleanRL philosophy this repository is built on). A generator +is the thin, uniform surface the evaluator needs in order to treat every model +the same way: + +1. `from_pretrained` -- rebuild a trained model from a checkpoint package. +2. `sample` -- turn conditions into designs. + +Everything else (which metrics run, how conditions are drawn, how results are +stored) belongs to the evaluator, not the model. +""" + +from __future__ import annotations + +import abc +from dataclasses import dataclass +from functools import cached_property +import hashlib +import json +import time +from typing import Any, ClassVar, TYPE_CHECKING + +from gymnasium import spaces +import numpy as np +import torch as th + +from engiopt.checkpoint_store import resolve_named_checkpoint + +if TYPE_CHECKING: + from datasets import Dataset + from engibench.core import Problem + import numpy.typing as npt + + from engiopt.checkpoint_store import ModelSource + from engiopt.checkpoint_store import ResolvedCheckpoint + +DEFAULT_HF_ENTITY = "IDEALLab" +DEFAULT_HF_REPO_PREFIX = "engiopt" + + +@dataclass(frozen=True) +class ConditionBatch: + """The conditions a generator is asked to satisfy, in every form models need. + + Most models want `tensor`. Some need the original columns -- to drop constant + conditions, re-normalize them, or look them up by name -- so the raw dataset + travels alongside rather than being reconstructed from the tensor. + + Attributes: + tensor: `(n, n_conds)` float tensor on the generator's device, or None + for unconditional problems. + dataset: The sampled conditions as an EngiBench/HF dataset, if available. + keys: Condition names, in the same column order as `tensor`. + """ + + tensor: th.Tensor | None + dataset: Dataset | None = None + keys: tuple[str, ...] = () + + def __len__(self) -> int: + if self.tensor is not None: + return int(self.tensor.shape[0]) + return len(self.dataset) if self.dataset is not None else 0 + + def require_tensor(self, algo_id: str) -> th.Tensor: + """Return `tensor`, with a clear error when a conditional model got nothing. + + Raises: + ValueError: If no conditions are available. + """ + if self.tensor is None: + raise ValueError(f"{algo_id} is conditional, but no conditions were provided.") + return self.tensor + + +def pick_device() -> th.device: + """Return the best available torch device, matching the repo-wide convention.""" + if th.backends.mps.is_available(): + return th.device("mps") + if th.cuda.is_available(): + return th.device("cuda") + return th.device("cpu") + + +INFRA_CONFIG_KEYS = frozenset( + { + "algo", + "problem_id", + "seed", + "track", + "save_model", + "strict_determinism", + "checkpoint_backend", + "sample_interval", + "wandb_project", + "wandb_entity", + "hf_entity", + "hf_repo_prefix", + "hf_private", + } +) +"""Run-config keys that describe *how a run was operated*, not what the model is. + +Excluded from the configuration fingerprint so that re-running the same +architecture under a different tracking setup does not look like a new model. +""" + + +def config_fingerprint(run_config: dict[str, Any] | None) -> str: + """Short stable hash of the hyperparameters that define a model configuration. + + Two runs of the same algorithm with different `latent_dim` are different + models and must occupy different leaderboard rows; the same run re-evaluated + must occupy the same one. `algo_id` and `seed` are not enough to tell those + apart, so the remaining run-config values are fingerprinted here. + + Returns: + An 8-character hash, or `"default"` when there is no config to hash. + """ + if not run_config: + return "default" + payload = {k: v for k, v in sorted(run_config.items()) if k not in INFRA_CONFIG_KEYS} + if not payload: + return "default" + encoded = json.dumps(payload, sort_keys=True, default=str).encode() + return hashlib.sha256(encoded).hexdigest()[:8] + + +def checkpoint_identity(args: Any) -> dict[str, Any]: + """Describe a training run's configuration for checkpoint storage. + + Splat into `save_checkpoint_package` so a run is filed both under its exact + hyperparameters and, when it used the script's defaults, under the canonical + path that plain `from_pretrained(problem, seed=...)` reads:: + + save_checkpoint_package(..., run_config=vars(args), **checkpoint_identity(args)) + + Args: + args: The training script's `Args` dataclass instance. + + Returns: + `config_fingerprint` and `is_default_config`, ready to pass through. + """ + config = vars(args) + try: + defaults = vars(type(args)()) + except Exception: # noqa: BLE001 - a non-default-constructible Args is not fatal here + defaults = {} + return { + "config_fingerprint": config_fingerprint(config), + "is_default_config": bool(defaults) and config_fingerprint(config) == config_fingerprint(defaults), + } + + +def config_path_parts(config_fingerprint_value: str | None) -> list[str] | None: + """Path components addressing one configuration inside a checkpoint repo. + + `None` means the canonical (default-hyperparameter) location. + """ + if not config_fingerprint_value: + return None + return [f"cfg_{config_fingerprint_value}"] + + +def design_shape_of(problem: Problem) -> tuple[int, ...]: + """Shape of a single design for `problem`. + + For `spaces.Dict` problems (e.g. airfoil, whose design is a bundle of named + arrays) this is the *flattened* shape, since that is the form generators + emit and metrics consume. Exposed as a function so adapters can size their + networks before an instance exists. + """ + space = problem.design_space + if isinstance(space, spaces.Dict): + dummy_design, _ = problem.random_design() + # `spaces.flatten` is typed as returning any flattenable form; for a Dict + # of Boxes it is always an array, so measure it as one. + return tuple(np.asarray(spaces.flatten(space, dummy_design)).shape) + return tuple(space.shape) + + +class Generator(abc.ABC): + """Base class for every generative model that can be evaluated or ranked. + + Subclasses declare a few class attributes and implement two methods. See + `engiopt/generators/_template/adapter.py` for a minimal working example. + + Class attributes: + algo_id: Stable identifier. Used as the leaderboard key, the HF repo + suffix, and the registry key. Must match the package directory name. + conditional: Whether `sample` consumes conditions. Unconditional models + still receive them (so the call signature is uniform) but ignore + them -- which is exactly what the conditional-adherence metrics are + meant to expose. + design_kinds: Which design-space shapes this model can serve, e.g. + `("2d",)`. Used to reject nonsensical model/problem pairings early. + checkpoint_files: Files the checkpoint package must contain. + primary_state_key: Key inside the loaded checkpoint holding the + state dict, e.g. `ckpt["generator"]`. + output_clip: Optional `(low, high)` clamp applied to sampled designs. + Several simulators are unstable at exactly 0, so 2D topology models + clamp to `(1e-3, 1)`. `None` means no clamping. Bounds may be arrays, + and an instance may override the class default when the range is + recorded in the checkpoint rather than fixed by the architecture. + """ + + algo_id: ClassVar[str] + conditional: ClassVar[bool] = True + design_kinds: ClassVar[tuple[str, ...]] = ("2d",) + checkpoint_files: ClassVar[tuple[str, ...]] = ("generator.pth",) + primary_state_key: ClassVar[str] = "generator" + output_clip: tuple[Any, Any] | None = None + + def __init__( + self, + problem: Problem, + *, + problem_id: str, + seed: int, + device: th.device, + run_config: dict[str, Any] | None = None, + ) -> None: + self.problem = problem + self.problem_id = problem_id + self.seed = seed + self.device = device + self.run_config: dict[str, Any] = run_config or {} + self.last_sample_seconds: float | None = None + """Wall-clock seconds for the most recent `sample` call (cost metrics).""" + + # ------------------------------------------------------------------ + # The contract + # ------------------------------------------------------------------ + + @classmethod + @abc.abstractmethod + def build( + cls, + resolved: ResolvedCheckpoint, + problem: Problem, + device: th.device, + **base: Any, + ) -> Generator: + """Reconstruct this model from an already-fetched checkpoint package. + + Fetching the package, choosing the device, and selecting the + configuration are handled by `from_pretrained`; this method only has to + rebuild the network and hand it to the constructor:: + + config = resolved.run_config + net = MyNet(latent_dim=config["latent_dim"], design_shape=problem.design_space.shape).to(device) + net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.eval() + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + Args: + resolved: The downloaded package: `files`, `run_config`, `metadata`. + problem: The problem this checkpoint was trained on. + device: Device the model should live on. + **base: The remaining `Generator.__init__` arguments (`problem_id`, + `seed`, `run_config`); splat them into `cls(...)` unchanged. + """ + + @classmethod + def from_pretrained( + cls, + problem: Problem, + *, + problem_id: str, + seed: int = 1, + device: th.device | None = None, + model_source: ModelSource = "auto", + hf_entity: str = DEFAULT_HF_ENTITY, + hf_repo_prefix: str = DEFAULT_HF_REPO_PREFIX, + config_fingerprint: str | None = None, + **kwargs: Any, + ) -> Generator: + """Rebuild a trained generator from its checkpoint package. + + Args: + problem: The problem to load a checkpoint for. + problem_id: Its registry key, which selects the package path. + seed: Training seed to load. + device: Torch device; auto-selected when omitted. + model_source: `auto`, `hf`, or `local`. + hf_entity: HF org/user holding the checkpoint repos. + hf_repo_prefix: Prefix of the per-model-family repo. + config_fingerprint: Load one specific hyperparameter configuration; + omit for the canonical default-hyperparameter checkpoint. + **kwargs: Passed through to `build`, for model-specific options. + + Returns: + The loaded generator, ready to `sample`. + """ + device = device or pick_device() + resolved = cls.resolve_checkpoint( + problem_id=problem_id, + seed=seed, + model_source=model_source, + hf_entity=hf_entity, + hf_repo_prefix=hf_repo_prefix, + config_fingerprint=config_fingerprint, + ) + return cls.build( + resolved, + problem=problem, + device=device, + problem_id=problem_id, + seed=seed, + run_config=resolved.run_config, + **kwargs, + ) + + @abc.abstractmethod + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor | npt.NDArray[Any]: + """Draw `n` designs. Implement the model-specific sampling procedure here. + + Args: + conditions: The conditions to satisfy. Use `conditions.tensor` for + the usual `(n, n_conds)` case, or `conditions.dataset` when the + model needs the original columns. Reshaping to whatever the + network expects is the implementation's job. + n: Number of designs to draw. + + Returns: + Designs in any shape; `sample` reshapes and clamps them. + """ + + # ------------------------------------------------------------------ + # Shared behaviour + # ------------------------------------------------------------------ + + def sample( + self, + conditions: ConditionBatch | th.Tensor | npt.NDArray[Any] | None, + n: int | None = None, + seed: int | None = None, + ) -> npt.NDArray[Any]: + """Draw designs, timing the call and normalizing the output shape. + + This is the method the evaluator calls. It wraps `_sample` so that every + model returns `(n, *design_shape)` float arrays and reports its own + sampling cost, regardless of how it generates internally. + + Args: + conditions: A `ConditionBatch`, a raw tensor/array of conditions, or + None. Bare tensors are wrapped for convenience so notebooks can + call `sample(conditions_tensor)` directly. + n: Number of designs; inferred from the conditions when omitted. + seed: Optional torch seed, for reproducible sampling. + + Returns: + `(n, *design_shape)` array of designs. + """ + if seed is not None: + th.manual_seed(seed) + batch = self._as_batch(conditions) + if n is None: + n = len(batch) + if n == 0: + raise ValueError("n must be given when no conditions are provided.") + + start = time.perf_counter() + with th.no_grad(): + raw = self._sample(batch, n) + self.last_sample_seconds = time.perf_counter() - start + + designs = raw.detach().cpu().numpy() if isinstance(raw, th.Tensor) else np.asarray(raw) + designs = designs.reshape(n, *self.design_shape) + if self.output_clip is not None: + designs = np.clip(designs, *self.output_clip) + return designs + + def _as_batch(self, conditions: ConditionBatch | th.Tensor | npt.NDArray[Any] | None) -> ConditionBatch: + """Normalize whatever the caller passed into a `ConditionBatch` on this device.""" + if isinstance(conditions, ConditionBatch): + tensor = conditions.tensor + moved = None if tensor is None else tensor.to(device=self.device, dtype=th.float) + return ConditionBatch(tensor=moved, dataset=conditions.dataset, keys=conditions.keys) + if conditions is None: + return ConditionBatch(tensor=None, keys=tuple(self.problem.conditions_keys)) + tensor = conditions if isinstance(conditions, th.Tensor) else th.as_tensor(np.asarray(conditions)) + return ConditionBatch( + tensor=tensor.to(device=self.device, dtype=th.float), + keys=tuple(self.problem.conditions_keys), + ) + + @cached_property + def design_shape(self) -> tuple[int, ...]: + """Shape of a single design; flattened for dict-valued design spaces.""" + return design_shape_of(self.problem) + + @cached_property + def config_fingerprint(self) -> str: + """Fingerprint of this checkpoint's hyperparameters; see `config_fingerprint`.""" + return config_fingerprint(self.run_config) + + @property + def n_conds(self) -> int: + """Number of scalar conditioning variables for this problem.""" + return len(self.problem.conditions_keys) + + # ------------------------------------------------------------------ + # Checkpoint plumbing shared by every adapter + # ------------------------------------------------------------------ + + @classmethod + def resolve_checkpoint( + cls, + *, + problem_id: str, + seed: int, + model_source: ModelSource = "auto", + hf_entity: str = DEFAULT_HF_ENTITY, + hf_repo_prefix: str = DEFAULT_HF_REPO_PREFIX, + local_model_dir: str | None = None, + config_fingerprint: str | None = None, + ) -> ResolvedCheckpoint: + """Fetch this model's checkpoint package from HuggingFace. + + Args: + problem_id: EngiBench problem the checkpoint was trained on. + seed: Training seed to load. + model_source: Backend preference: `auto`, `hf`, `wandb`, or `local`. + hf_entity: HF org/user holding the checkpoint repos. + hf_repo_prefix: Prefix of the per-model-family HF repo. + local_model_dir: Directory to load from, for `local` sources. + config_fingerprint: Load one specific hyperparameter configuration. + Omit it to load the canonical default-hyperparameter checkpoint, + which is what a bare model name means. + """ + return resolve_named_checkpoint( + model_source=model_source, + problem_id=problem_id, + algo=cls.algo_id, + seed=seed, + hf_entity=hf_entity, + hf_repo_prefix=hf_repo_prefix, + required_files=list(cls.checkpoint_files), + local_model_dir=local_model_dir, + extra_path_parts=config_path_parts(config_fingerprint), + ) + + def __repr__(self) -> str: + return f"{type(self).__name__}(algo_id={self.algo_id!r}, problem_id={self.problem_id!r}, seed={self.seed})" diff --git a/engiopt/generators/__init__.py b/engiopt/generators/__init__.py new file mode 100644 index 00000000..7f22eb37 --- /dev/null +++ b/engiopt/generators/__init__.py @@ -0,0 +1,8 @@ +"""Generative models shipped with EngiOpt. + +Each subpackage holds one model family: a self-contained training script (kept +single-file, in the CleanRL style) plus an `adapter.py` implementing the +`Generator` contract from `engiopt.core`. + +Use `engiopt.utils.all_generators.BUILTIN_GENERATORS` to look models up by name. +""" diff --git a/engiopt/utils/__init__.py b/engiopt/utils/__init__.py new file mode 100644 index 00000000..616c4bcc --- /dev/null +++ b/engiopt/utils/__init__.py @@ -0,0 +1 @@ +"""Shared utilities for EngiOpt.""" diff --git a/engiopt/utils/all_generators.py b/engiopt/utils/all_generators.py new file mode 100644 index 00000000..00554274 --- /dev/null +++ b/engiopt/utils/all_generators.py @@ -0,0 +1,99 @@ +"""Registry of all generators in EngiOpt. + +Mirrors `engibench.utils.all_problems`, so the two halves of the stack read the +same way:: + + problem = BUILTIN_PROBLEMS["beams2d"]() + generator = BUILTIN_GENERATORS["cgan_cnn_2d"].from_pretrained(problem, problem_id="beams2d", seed=1) +""" + +from __future__ import annotations + +import importlib +import pkgutil +from typing import Any, TYPE_CHECKING + +from gymnasium import spaces + +from engiopt.core import Generator +import engiopt.generators + +if TYPE_CHECKING: + from engibench.core import Problem + +ADAPTER_MODULE = "adapter" +"""Each generator package implements the contract in this module.""" + + +def list_generators(base_module: Any = engiopt.generators) -> dict[str, type[Generator]]: + """Return a dict containing all available `Generator` classes defined in submodules of `base_module`.""" + module_path = next(iter(base_module.__path__)) + modules = pkgutil.iter_modules(path=[module_path], prefix="") + return { + m.name: extract_generator(importlib.import_module(f"{base_module.__package__}.{m.name}.{ADAPTER_MODULE}")) + for m in modules + if m.ispkg and not m.name.startswith("_") + } + + +def extract_generator(module: Any) -> type[Generator]: + """Get a `Generator` class defined in a module. + + Raises an exception if the module contains multiple `Generator` classes. + """ + generator_types = [ + o + for o in vars(module).values() + if isinstance(o, type) and issubclass(o, Generator) and o is not Generator and not _is_abstract(o) + ] + try: + (g,) = generator_types + except ValueError: + msg = f"Only one generator per adapter is allowed. Got {', '.join(t.__name__ for t in generator_types)}" + raise ValueError(msg) from None + return g + + +def _is_abstract(obj: type) -> bool: + """Whether a class still has unimplemented abstract methods.""" + return bool(getattr(obj, "__abstractmethods__", frozenset())) + + +BUILTIN_GENERATORS = list_generators() + + +def design_kind_of(problem: Problem) -> str: + """Classify a problem's design space as `1d`, `2d`, `3d`, or `dict`. + + Used to reject nonsensical pairings, such as a 2D CNN generator pointed at a + 3D problem. + """ + if isinstance(problem.design_space, spaces.Dict): + return "dict" + ndim = len(problem.design_space.shape or ()) + return {1: "1d", 2: "2d", 3: "3d"}.get(ndim, f"{ndim}d") + + +def generators_for(problem: Problem) -> dict[str, type[Generator]]: + """Return the generators whose `design_kinds` cover this problem's design space.""" + kind = design_kind_of(problem) + return {name: g for name, g in BUILTIN_GENERATORS.items() if kind in g.design_kinds} + + +def problem_id_of(problem: Problem | type[Problem]) -> str: + """Reverse-look-up a problem's registry key. + + EngiBench problems do not carry their own id, but checkpoint paths and + leaderboard rows are keyed by it. + + Raises: + ValueError: If the problem is not a built-in EngiBench problem. + """ + from engibench.utils.all_problems import BUILTIN_PROBLEMS + + problem_cls = problem if isinstance(problem, type) else type(problem) + for name, candidate in BUILTIN_PROBLEMS.items(): + if candidate is problem_cls: + return name + msg = f"{problem_cls.__name__} is not a built-in EngiBench problem; pass problem_id explicitly." + raise ValueError(msg) From a796025f8a4b03e7eac2ca6f26fe9629ba9e469a Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Fri, 24 Jul 2026 13:57:33 +0200 Subject: [PATCH 03/24] Replace per-model evaluation scripts with a shared evaluation layer The 13 evaluate_*.py scripts were ~79% identical; they differed only in which network to import, how to build it, whether to pass conditions, and a model id string. Each model now supplies a ~45-line adapter and everything else is shared: - engiopt/evaluation/: a metric registry declaring each metric's family, cost, and ranking direction; an EvaluationContext that computes the expensive optimizer pass once for every performance metric; an Evaluator; and the leaderboard, published to HuggingFace by merging rather than overwriting. - engiopt/specs//v1.json: frozen, hash-verified evaluation contracts so rows are comparable across people and across time. - python -m engiopt.evaluate replaces the deleted scripts. - engiopt/generators/_template/ is the copy-me starting point for new models. Co-Authored-By: Claude Opus 4.8 (1M context) --- engiopt/evaluate.py | 226 ++++++++++++++++ engiopt/evaluation/__init__.py | 49 ++++ engiopt/evaluation/context.py | 247 +++++++++++++++++ engiopt/evaluation/evaluator.py | 254 ++++++++++++++++++ engiopt/evaluation/leaderboard.py | 197 ++++++++++++++ engiopt/evaluation/metrics/__init__.py | 10 + engiopt/evaluation/metrics/builtin.py | 119 ++++++++ engiopt/evaluation/registry.py | 147 ++++++++++ engiopt/evaluation/spec.py | 233 ++++++++++++++++ engiopt/generators/_template/__init__.py | 5 + engiopt/generators/_template/adapter.py | 106 ++++++++ .../generators/cgan_1d/evaluate_cgan_1d.py | 135 ---------- .../generators/cgan_2d/evaluate_cgan_2d.py | 131 --------- .../cgan_cnn_2d/evaluate_cgan_cnn_2d.py | 135 ---------- .../cgan_cnn_3d/evaluate_cgan_cnn_3d.py | 141 ---------- .../generators/cgan_vae/evaluate_cgan_vae.py | 141 ---------- .../diffusion_1d/evaluate_diffusion_1d.py | 148 ---------- .../evaluate_diffusion_2d_cond.py | 175 ------------ engiopt/generators/gan_1d/evaluate_gan_1d.py | 130 --------- engiopt/generators/gan_2d/evaluate_gan_2d.py | 129 --------- .../gan_bezier/evaluate_gan_bezier.py | 153 ----------- .../gan_cnn_2d/evaluate_gan_cnn_2d.py | 131 --------- .../evaluate_pixel_cnn_pp_2d.py | 165 ------------ engiopt/generators/vqgan/evaluate_vqgan.py | 228 ---------------- engiopt/specs/beams2d/v1.json | 19 ++ engiopt/specs/heatconduction2d/v1.json | 19 ++ engiopt/specs/photonics2d/v1.json | 19 ++ engiopt/specs/thermoelastic2d/v1.json | 21 ++ 28 files changed, 1671 insertions(+), 1942 deletions(-) create mode 100644 engiopt/evaluate.py create mode 100644 engiopt/evaluation/__init__.py create mode 100644 engiopt/evaluation/context.py create mode 100644 engiopt/evaluation/evaluator.py create mode 100644 engiopt/evaluation/leaderboard.py create mode 100644 engiopt/evaluation/metrics/__init__.py create mode 100644 engiopt/evaluation/metrics/builtin.py create mode 100644 engiopt/evaluation/registry.py create mode 100644 engiopt/evaluation/spec.py create mode 100644 engiopt/generators/_template/__init__.py create mode 100644 engiopt/generators/_template/adapter.py delete mode 100644 engiopt/generators/cgan_1d/evaluate_cgan_1d.py delete mode 100644 engiopt/generators/cgan_2d/evaluate_cgan_2d.py delete mode 100644 engiopt/generators/cgan_cnn_2d/evaluate_cgan_cnn_2d.py delete mode 100644 engiopt/generators/cgan_cnn_3d/evaluate_cgan_cnn_3d.py delete mode 100644 engiopt/generators/cgan_vae/evaluate_cgan_vae.py delete mode 100644 engiopt/generators/diffusion_1d/evaluate_diffusion_1d.py delete mode 100644 engiopt/generators/diffusion_2d_cond/evaluate_diffusion_2d_cond.py delete mode 100644 engiopt/generators/gan_1d/evaluate_gan_1d.py delete mode 100644 engiopt/generators/gan_2d/evaluate_gan_2d.py delete mode 100644 engiopt/generators/gan_bezier/evaluate_gan_bezier.py delete mode 100644 engiopt/generators/gan_cnn_2d/evaluate_gan_cnn_2d.py delete mode 100644 engiopt/generators/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py delete mode 100644 engiopt/generators/vqgan/evaluate_vqgan.py create mode 100644 engiopt/specs/beams2d/v1.json create mode 100644 engiopt/specs/heatconduction2d/v1.json create mode 100644 engiopt/specs/photonics2d/v1.json create mode 100644 engiopt/specs/thermoelastic2d/v1.json diff --git a/engiopt/evaluate.py b/engiopt/evaluate.py new file mode 100644 index 00000000..9c7ccd21 --- /dev/null +++ b/engiopt/evaluate.py @@ -0,0 +1,226 @@ +"""One evaluation entry point for every model. + +Replaces the per-model `evaluate_*.py` scripts, which differed only in how they +loaded and called their generator -- now the `Generator` contract's job. + + # score three models on one problem, cheap metrics only + python -m engiopt.evaluate --problem-id beams2d --generators cgan_cnn_2d gan_cnn_2d vqgan + + # the whole zoo, several seeds, including simulator-backed metrics + python -m engiopt.evaluate --problem-id beams2d --generators all --seeds 1 2 3 --include-expensive + + # what can I run? + python -m engiopt.evaluate --list-generators --list-metrics +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import pandas as pd +import tyro + +from engiopt.checkpoint_store import publish_checkpoint_metrics +from engiopt.core import config_path_parts +from engiopt.evaluation.evaluator import Evaluator +from engiopt.evaluation.leaderboard import already_evaluated +from engiopt.evaluation.leaderboard import append_rows +from engiopt.evaluation.leaderboard import disagreement +from engiopt.evaluation.leaderboard import load_from_hub +from engiopt.evaluation.leaderboard import push_to_hub +from engiopt.evaluation.registry import METRICS +from engiopt.utils.all_generators import BUILTIN_GENERATORS +from engiopt.utils.all_generators import generators_for + +if TYPE_CHECKING: + from engiopt.core import Generator + + +@dataclasses.dataclass +class Args: + """Command-line arguments for evaluating one or more generators.""" + + problem_id: str = "beams2d" + """EngiBench problem to evaluate on.""" + generators: tuple[str, ...] = ("all",) + """Generator names, or `all` for every registered model that fits the problem.""" + seeds: tuple[int, ...] = (1,) + """Checkpoint seeds to load for each generator.""" + config_fingerprints: tuple[str, ...] = () + """Specific hyperparameter configurations to evaluate, by fingerprint. + + Empty means the canonical default-hyperparameter checkpoint -- what a bare + model name refers to. Pass several to score a sweep.""" + spec: str | None = None + """Eval spec reference, e.g. `beams2d/v1`. Defaults to `/v1`.""" + metrics: tuple[str, ...] = () + """Metric names; defaults to the spec's list.""" + include_expensive: bool = False + """Run simulator-backed metrics (COG/IOG/FOG, feasibility). Slow.""" + output_csv: str = "leaderboard_{problem_id}.csv" + """Where to append results locally; may include `{problem_id}`.""" + push_to: str | None = None + """HuggingFace dataset repo to publish into, e.g. `IDEALLab/engiopt-leaderboard`. + + The published board is downloaded and merged, so publishing one model never + recomputes or overwrites anyone else's rows.""" + skip_existing: bool = False + """Skip generator/seed pairs the published board already has rows for.""" + attach_metrics: bool = False + """Also write each model's scores into its own HF checkpoint package. + + Makes a checkpoint self-describing: whoever downloads it sees how it scored + without needing the leaderboard.""" + model_source: str = "auto" + """Checkpoint backend: `auto`, `hf`, `wandb`, or `local`.""" + hf_entity: str = "IDEALLab" + """HF org/user holding the checkpoints.""" + hf_repo_prefix: str = "engiopt" + """HF repo prefix for model-family repositories.""" + on_error: str = "skip" + """`skip` past models that fail to load, or `raise` to stop.""" + show_disagreement: bool = True + """Print each metric's ranking side by side after evaluating.""" + list_generators: bool = False + """List registered generators and exit.""" + list_metrics: bool = False + """List registered metrics and exit.""" + + +def _print_generators() -> None: + """Print every registered generator, including any that failed to import.""" + print(f"{len(BUILTIN_GENERATORS)} generators registered:\n") + for name, generator in sorted(BUILTIN_GENERATORS.items()): + kinds = "/".join(generator.design_kinds) + conditioning = "conditional" if generator.conditional else "unconditional" + print(f" {name:<20} {kinds:<10} {conditioning}") + + +def _print_metrics() -> None: + """Print every registered metric grouped by the question it answers.""" + print(f"{len(METRICS)} metrics registered:\n") + for family in sorted({spec.family for spec in METRICS.values()}): + print(f" [{family}]") + for spec in METRICS.select(family=family): + direction = {True: "higher better", False: "lower better", None: "diagnostic"}[spec.higher_is_better] + print(f" {spec.name:<10} {spec.cost:<10} {direction:<14} {spec.description}") + + +def _resolve_generator_names(requested: tuple[str, ...], problem_id: str) -> list[str]: + """Expand `all` to every registered generator compatible with the problem.""" + if "all" not in requested: + return list(requested) + from engibench.utils.all_problems import BUILTIN_PROBLEMS + + return sorted(generators_for(BUILTIN_PROBLEMS[problem_id]())) + + +def _load_generators(args: Args, evaluator: Evaluator) -> list[Generator]: + """Load every requested generator/seed pair, reporting those that fail.""" + generators: list[Generator] = [] + fingerprints: tuple[str | None, ...] = args.config_fingerprints or (None,) + for name in _resolve_generator_names(args.generators, args.problem_id): + generator_cls = BUILTIN_GENERATORS[name] + for seed in args.seeds: + for fingerprint in fingerprints: + try: + generators.append( + generator_cls.from_pretrained( + evaluator.problem, + problem_id=args.problem_id, + seed=seed, + device=evaluator.device, + model_source=args.model_source, # type: ignore[arg-type] + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + config_fingerprint=fingerprint, + ) + ) + # A missing checkpoint should not abort a whole sweep. + except Exception as exc: # noqa: PERF203 + if args.on_error != "skip": + raise + label = f"{name} seed {seed}" + (f" cfg {fingerprint}" if fingerprint else "") + print(f" could not load {label}: {exc}") + return generators + + +def _attach_metrics_to_checkpoints(args: Args, board: pd.DataFrame) -> None: + """Write each row's scores into the checkpoint package it describes.""" + for row in board.to_dict("records"): + fingerprint = row.get("config_fingerprint") + try: + path = publish_checkpoint_metrics( + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + problem_id=row["problem_id"], + algo=row["algo_id"], + seed=int(row["seed"]), + metrics={k: v for k, v in row.items() if pd.notna(v)}, + extra_path_parts=config_path_parts(fingerprint), + ) + print(f" metrics -> {row['algo_id']}: {path}") + except Exception as exc: # noqa: BLE001 - one failure must not lose the rest + print(f" could not attach metrics for {row['algo_id']}: {exc}") + + +def main(args: Args) -> None: + """Evaluate the requested generators and append the results to a CSV.""" + if args.list_generators: + _print_generators() + if args.list_metrics: + _print_metrics() + if args.list_generators or args.list_metrics: + return + + spec = args.spec or f"{args.problem_id}/v1" + evaluator = Evaluator.for_problem(args.problem_id, spec=spec) + print(f"Problem {args.problem_id} | spec {evaluator.spec.version} | n={evaluator.spec.n_samples}") + + published = load_from_hub(args.push_to) if (args.push_to and args.skip_existing) else None + generators = _load_generators(args, evaluator) + if published is not None: + generators = [ + g + for g in generators + if not already_evaluated( + published, + problem_id=args.problem_id, + algo_id=g.algo_id, + config_fingerprint=g.config_fingerprint, + seed=g.seed, + spec_version=evaluator.spec.version, + ) + ] + print(f"{len(generators)} generator(s) left after skipping already-published rows.") + if not generators: + print("No generators loaded; nothing to evaluate.") + return + + board = evaluator.leaderboard( + generators, + only=list(args.metrics) or None, + include_expensive=args.include_expensive, + on_error=args.on_error, + ) + destination = append_rows(board, args.output_csv.format(problem_id=args.problem_id)) + print(f"\n{board.to_string(index=False)}\n") + print(f"Wrote {len(board)} rows to {destination}") + + if args.push_to: + merged = push_to_hub(board, args.push_to) + print(f"Published {len(board)} row(s) to {args.push_to}; board now holds {len(merged)} rows.") + + if args.attach_metrics: + _attach_metrics_to_checkpoints(args, board) + + if args.show_disagreement and len(board) > 1: + ranked = [m for m in board.columns if m in METRICS and METRICS[m].higher_is_better is not None] + if len(ranked) > 1: + print("\nRankings by metric (1 = best) -- where these disagree is the interesting part:\n") + print(disagreement(board, ranked).to_string()) + + +if __name__ == "__main__": + main(tyro.cli(Args)) diff --git a/engiopt/evaluation/__init__.py b/engiopt/evaluation/__init__.py new file mode 100644 index 00000000..654d1b6d --- /dev/null +++ b/engiopt/evaluation/__init__.py @@ -0,0 +1,49 @@ +"""Evaluating generators and building leaderboards. + + from engiopt.evaluation import Evaluator + from engiopt.utils.all_generators import BUILTIN_GENERATORS + + ev = Evaluator.for_problem("beams2d", spec="beams2d/v1") + gen = BUILTIN_GENERATORS["cgan_cnn_2d"].from_pretrained(ev.problem, problem_id="beams2d", seed=1) + ev.score(gen) + +Metrics are registered functions rather than model methods, because a metric +compares a generated set against a reference set under a problem -- it is a +property of the comparison, not of the model. +""" + +from engiopt.evaluation.context import EvaluationContext +from engiopt.evaluation.context import OptimizationResults +from engiopt.evaluation.evaluator import Evaluator +from engiopt.evaluation.leaderboard import already_evaluated +from engiopt.evaluation.leaderboard import append_rows +from engiopt.evaluation.leaderboard import disagreement +from engiopt.evaluation.leaderboard import load_from_hub +from engiopt.evaluation.leaderboard import merge_rows +from engiopt.evaluation.leaderboard import push_to_hub +from engiopt.evaluation.leaderboard import rank +from engiopt.evaluation.registry import MetricRegistry +from engiopt.evaluation.registry import METRICS +from engiopt.evaluation.registry import MetricSpec +from engiopt.evaluation.registry import register_metric +from engiopt.evaluation.spec import EvalSpec +from engiopt.evaluation.spec import ResolvedSpec + +__all__ = [ + "METRICS", + "EvalSpec", + "EvaluationContext", + "Evaluator", + "MetricRegistry", + "MetricSpec", + "OptimizationResults", + "ResolvedSpec", + "already_evaluated", + "append_rows", + "disagreement", + "load_from_hub", + "merge_rows", + "push_to_hub", + "rank", + "register_metric", +] diff --git a/engiopt/evaluation/context.py b/engiopt/evaluation/context.py new file mode 100644 index 00000000..e955fc11 --- /dev/null +++ b/engiopt/evaluation/context.py @@ -0,0 +1,247 @@ +"""Shared state for one model-vs-reference comparison. + +The context exists so that metrics compose without recomputing. Ten metrics that +each need the flattened designs flatten once; the whole optimality-gap family +shares a single pass of the optimizer instead of running it once per metric. + +Anything expensive is a `cached_property`, so it is paid for only if some +selected metric actually asks for it. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from dataclasses import field +from functools import cached_property +from typing import Any, TYPE_CHECKING + +from gymnasium import spaces +import numpy as np + +from engiopt import metrics as metrics_mod + +if TYPE_CHECKING: + from datasets import Dataset + from engibench.core import Problem + import numpy.typing as npt + + +class MultiObjectiveScalarizationError(ValueError): + """Raised when a multi-objective problem has no declared scalarization. + + Averaging or summing objectives with different units and different intended + priorities produces a number that looks fine and means nothing, so the + evaluator refuses rather than guessing. + """ + + def __init__(self, problem_id: str, objective_names: list[str]) -> None: + super().__init__( + f"{problem_id!r} has {len(objective_names)} objectives ({', '.join(objective_names)}) but the eval " + "spec does not say how to combine them. Set `objective_weights` for fixed weights, or " + "`objective_weight_condition` when the trade-off is a per-sample condition (thermoelastic2d's " + "`weight`). Summing or averaging objectives with different units is not a valid default." + ) + + +@dataclass +class OptimizationResults: + """Per-sample outputs of the simulator/optimizer pass. + + Computing these is the expensive part of evaluation -- one optimizer run and + two simulations per sample -- so every performance metric reads from one + shared instance. + """ + + iog: list[float] = field(default_factory=list) + """Initial optimality gap: how far each generated design starts from the reference optimum.""" + cog: list[float] = field(default_factory=list) + """Sum of the optimality gaps over each design's re-optimization history.""" + fog: list[float] = field(default_factory=list) + """Final optimality gap at the end of re-optimization.""" + viol: list[bool] = field(default_factory=list) + """Whether each design missed its volume-fraction target beyond tolerance.""" + + +@dataclass +class EvaluationContext: + """Everything a metric may read about one generator's output. + + Attributes: + problem: The EngiBench problem being evaluated against. + problem_id: Registry key for the problem, used in leaderboard rows. + gen_designs: Generated designs, `(n_samples, *design_shape)`. + ref_designs: Reference (dataset-optimal) designs for the same conditions. + conditions: The conditions each design was asked to satisfy. + sigma: Gaussian-kernel bandwidth for MMD and DPP. + volfrac_tol: Tolerance used by the volume-fraction violation check. + sample_seconds: Wall-clock seconds the generator took to produce + `gen_designs`, recorded as leaderboard provenance. + objective_weights: Fixed weights over `problem.objectives`, used to + reduce a multi-objective result to one number. + objective_weight_condition: Name of a per-sample condition carrying the + trade-off instead. thermoelastic2d's `weight` splits the first two + objectives as `(w, 1 - w)`; any remaining objectives get zero. + """ + + problem: Problem + problem_id: str + gen_designs: npt.NDArray[Any] + ref_designs: npt.NDArray[Any] + conditions: Dataset | None = None + sigma: float = 10.0 + volfrac_tol: float = 0.01 + sample_seconds: float | None = None + objective_weights: tuple[float, ...] | None = None + objective_weight_condition: str | None = None + + @property + def n_samples(self) -> int: + """Number of generated designs under comparison.""" + return len(self.gen_designs) + + @property + def is_dict_space(self) -> bool: + """Whether the problem uses a `spaces.Dict` design space needing flatten/unflatten.""" + return isinstance(self.problem.design_space, spaces.Dict) + + @cached_property + def gen_flat(self) -> npt.NDArray[Any]: + """Generated designs flattened to `(n_samples, -1)`.""" + return np.asarray(self.gen_designs).reshape(self.n_samples, -1) + + @cached_property + def ref_flat(self) -> npt.NDArray[Any]: + """Reference designs flattened to `(n_samples, -1)`, unpacking dict spaces.""" + if not self.is_dict_space: + return np.asarray(self.ref_designs).reshape(len(self.ref_designs), -1) + flattened = [np.asarray(spaces.flatten(self.problem.design_space, design)) for design in self.ref_designs] + return np.asarray(flattened) + + def condition_at(self, index: int) -> dict[str, Any] | None: + """Conditions for sample `index`, or None when the problem is unconditional.""" + return self.conditions[index] if self.conditions is not None else None + + def design_for_solver(self, index: int) -> Any: + """Generated design `index` in the shape the simulator and optimizer expect.""" + if self.is_dict_space: + return spaces.unflatten(self.problem.design_space, self.gen_designs[index]) + return self.gen_designs[index] + + @property + def objective_names(self) -> list[str]: + """Names of the problem's objectives, in the order the simulator returns them.""" + return [name for name, _direction in self.problem.objectives] + + @cached_property + def objective_signs(self) -> npt.NDArray[Any]: + """`+1` for objectives to minimize, `-1` for objectives to maximize. + + A raw `objective - baseline` difference only means "worse than the + reference" when the objective is being minimized. photonics2d maximizes + `total_overlap`, so there a *positive* difference means the design is + better; flipping its sign keeps every gap pointing the same way and lets + `iog`, `cog`, and `fog` stay honestly registered as lower-is-better. + """ + signs = [] + for _name, direction in self.problem.objectives: + label = getattr(direction, "name", str(direction)) + signs.append(-1.0 if "MAXIMIZE" in label.upper() else 1.0) + return np.asarray(signs, dtype=float) + + def weights_at(self, index: int) -> npt.NDArray[Any] | None: + """Weights used to scalarize sample `index`'s objective vector. + + Returns None for single-objective problems, where no choice is needed. + + Raises: + MultiObjectiveScalarizationError: If the problem has several + objectives and the spec declares no way to combine them. + """ + names = self.objective_names + if len(names) <= 1: + return None + + if self.objective_weight_condition is not None: + conditions = self.condition_at(index) or {} + trade_off = float(conditions[self.objective_weight_condition]) + # The declared convention: w on the first objective, (1 - w) on the + # second, nothing on the rest. At w = 1 the second objective cannot + # influence the gap at all, which is the whole point. + weights = np.zeros(len(names)) + weights[0] = trade_off + weights[1] = 1.0 - trade_off + return weights + + if self.objective_weights is not None: + if len(self.objective_weights) != len(names): + msg = ( + f"{self.problem_id!r} has {len(names)} objectives but the spec declares " + f"{len(self.objective_weights)} weights." + ) + raise ValueError(msg) + return np.asarray(self.objective_weights, dtype=float) + + raise MultiObjectiveScalarizationError(self.problem_id, names) + + def scalarize_gap(self, gap: Any, index: int) -> float: + """Reduce one sample's raw `objective - baseline` vector to a single number. + + Two corrections are applied, in this order: + + 1. Objective direction, so a positive result always means "worse than + the reference optimum" even for maximized objectives. + 2. The declared weights, so objectives the sample does not care about + cannot influence its score. + + Both are linear, so applying them to the gap is equivalent to applying + them to the objectives before differencing. + """ + array = np.atleast_1d(np.asarray(gap, dtype=float)) + signed = array * self.objective_signs[: array.size] + if array.size == 1: + return float(signed[0]) + weights = self.weights_at(index) + if weights is None: # more values than objectives should not happen + raise MultiObjectiveScalarizationError(self.problem_id, self.objective_names) + return float(np.dot(weights[: array.size], signed)) + + @cached_property + def optimization(self) -> OptimizationResults: + """Run the optimizer and simulator over every sample (expensive). + + All three optimality gaps are measured against the same baseline -- the + simulated objective of the reference (dataset-optimal) design for the + same conditions: + + - `iog` is the gap *before* any re-optimization: how good the generated + design already is. + - `cog` accumulates the gap over the whole re-optimization trajectory. + - `fog` is the gap once re-optimization has finished. + + Every gap is passed through `scalarize_gap`, which flips maximized + objectives and applies the spec's declared weights. All three metrics + are therefore lower-is-better on every problem, and an objective a + sample does not care about cannot influence its score. + """ + results = OptimizationResults() + for i in range(self.n_samples): + conditions = self.condition_at(i) + design = self.design_for_solver(i) + + self.problem.reset() + _, opt_history = self.problem.optimize(design, config=conditions) + self.problem.reset() + reference_optimum = self.problem.simulate(self.ref_designs[i], config=conditions) + gaps = metrics_mod.optimality_gap(opt_history, reference_optimum) + self.problem.reset() + generated_objective = self.problem.simulate(design, config=conditions) + + results.iog.append(self.scalarize_gap(np.asarray(generated_objective) - np.asarray(reference_optimum), i)) + results.cog.append(sum(self.scalarize_gap(step_gap, i) for step_gap in gaps)) + results.fog.append(self.scalarize_gap(gaps[-1], i)) + + if conditions: + target_vol = conditions.get("volfrac") or conditions.get("volume") + if target_vol is not None: + results.viol.append(bool(np.abs(np.mean(design) - target_vol) >= self.volfrac_tol)) + return results diff --git a/engiopt/evaluation/evaluator.py b/engiopt/evaluation/evaluator.py new file mode 100644 index 00000000..ebb80572 --- /dev/null +++ b/engiopt/evaluation/evaluator.py @@ -0,0 +1,254 @@ +"""Scoring generators and assembling leaderboards. + +This replaces the per-model `evaluate_*.py` scripts. Those differed only in how +they loaded and called their model -- which is now the `Generator` contract -- +so everything else lives here once:: + + ev = Evaluator.for_problem("beams2d", spec="beams2d/v1") + row = ev.score(generator) # cheap metrics only + board = ev.leaderboard(zoo) # a DataFrame, one row per model +""" + +from __future__ import annotations + +from dataclasses import dataclass +from dataclasses import field +import datetime as dt +from typing import Any, TYPE_CHECKING + +import pandas as pd + +from engiopt.core import ConditionBatch +from engiopt.core import pick_device +from engiopt.evaluation.context import EvaluationContext +from engiopt.evaluation.registry import MetricRegistry +from engiopt.evaluation.registry import METRICS +from engiopt.evaluation.registry import MetricSpec +from engiopt.evaluation.spec import EvalSpec +from engiopt.evaluation.spec import ResolvedSpec +from engiopt.utils.all_generators import design_kind_of + +# Importing the metrics package is what populates the registry. +import engiopt.evaluation.metrics # noqa: F401 # isort: skip + +if TYPE_CHECKING: + from engibench.core import Problem + import torch as th + + from engiopt.core import Generator + +PROVENANCE_COLUMNS = ( + "problem_id", + "algo_id", + "config_fingerprint", + "seed", + "spec_version", + "n_samples", + "sample_seconds", + "evaluated_at", +) +"""Columns identifying *what was measured*, as opposed to metric values.""" + + +@dataclass +class Evaluator: + """Scores generators on one problem under one frozen evaluation spec. + + Attributes: + problem: The EngiBench problem instance. + problem_id: Its registry key. + resolved: The spec with its conditions and reference designs drawn. + device: Device passed to generators when loading and sampling. + registry: Metric registry to draw from (defaults to the global one). + """ + + problem: Problem + problem_id: str + resolved: ResolvedSpec + device: th.device + registry: MetricRegistry = field(default_factory=lambda: METRICS) + + @classmethod + def for_problem( + cls, + problem_id: str, + *, + spec: str | EvalSpec | None = None, + device: th.device | None = None, + registry: MetricRegistry | None = None, + ) -> Evaluator: + """Build an evaluator for a problem, loading its spec. + + Args: + problem_id: EngiBench problem registry key. + spec: `"/"`, an `EvalSpec`, or None to load + `"/v1"`. + device: Torch device; auto-selected when omitted. + registry: Metric registry override, useful in tests. + """ + from engibench.utils.all_problems import BUILTIN_PROBLEMS + + problem = BUILTIN_PROBLEMS[problem_id]() + eval_spec = spec if isinstance(spec, EvalSpec) else EvalSpec.load(spec or f"{problem_id}/v1") + if eval_spec.problem_id != problem_id: + raise ValueError(f"Spec is for {eval_spec.problem_id!r}, not {problem_id!r}.") + device = device or pick_device() + problem.reset(seed=eval_spec.condition_seed) + return cls( + problem=problem, + problem_id=problem_id, + resolved=eval_spec.resolve(problem, device=device), + device=device, + registry=registry or METRICS, + ) + + @property + def spec(self) -> EvalSpec: + """The frozen evaluation contract in force.""" + return self.resolved.spec + + # ------------------------------------------------------------------ + # Scoring + # ------------------------------------------------------------------ + + def context_for(self, generator: Generator) -> EvaluationContext: + """Sample from a generator and wrap the result in an evaluation context.""" + kind = design_kind_of(self.problem) + if kind not in generator.design_kinds: + raise ValueError( + f"{generator.algo_id!r} supports {generator.design_kinds} design spaces, " + f"but {self.problem_id!r} is {kind!r}." + ) + designs = generator.sample( + ConditionBatch( + tensor=self.resolved.conditions_tensor, + dataset=self.resolved.conditions, + keys=self.resolved.condition_keys, + ), + n=self.resolved.n_samples, + seed=getattr(generator, "seed", None), + ) + return EvaluationContext( + problem=self.problem, + problem_id=self.problem_id, + gen_designs=designs, + ref_designs=self.resolved.ref_designs, + conditions=self.resolved.conditions, + sigma=self.spec.sigma, + volfrac_tol=self.spec.volfrac_tol, + sample_seconds=generator.last_sample_seconds, + objective_weights=self.spec.objective_weights, + objective_weight_condition=self.spec.objective_weight_condition, + ) + + def score( + self, + generator: Generator, + *, + only: list[str] | None = None, + include_expensive: bool = False, + ) -> dict[str, Any]: + """Score one generator, returning a single leaderboard row. + + Args: + generator: A loaded generator. + only: Metric names to compute; defaults to the spec's metric list. + include_expensive: Whether to run simulator-backed metrics. Left + False, no metric can invoke the simulator or optimizer. + + Returns: + A dict of provenance columns plus one entry per metric column. + """ + ctx = self.context_for(generator) + row = self._provenance(generator, ctx) + row.update(self.score_context(ctx, only=only, include_expensive=include_expensive)) + return row + + def score_context( + self, + ctx: EvaluationContext, + *, + only: list[str] | None = None, + include_expensive: bool = False, + ) -> dict[str, Any]: + """Run the selected metrics against an existing context.""" + values: dict[str, Any] = {} + for spec in self._selected(only, include_expensive=include_expensive): + values.update(_as_columns(spec, spec.fn(ctx))) + return values + + def _selected(self, only: list[str] | None, *, include_expensive: bool) -> list[MetricSpec]: + names = list(only) if only is not None else list(self.spec.metrics) + specs = self.registry.select(names) + if not include_expensive: + specs = [spec for spec in specs if spec.cost == "cheap"] + return specs + + def _provenance(self, generator: Generator, ctx: EvaluationContext) -> dict[str, Any]: + return { + "problem_id": self.problem_id, + "algo_id": generator.algo_id, + "config_fingerprint": generator.config_fingerprint, + "seed": getattr(generator, "seed", None), + "spec_version": self.spec.version, + "n_samples": ctx.n_samples, + "sample_seconds": ctx.sample_seconds, + "evaluated_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"), + } + + # ------------------------------------------------------------------ + # Leaderboard + # ------------------------------------------------------------------ + + def leaderboard( + self, + generators: list[Generator], + *, + only: list[str] | None = None, + include_expensive: bool = False, + on_error: str = "raise", + ) -> pd.DataFrame: + """Score many generators into one table. + + Args: + generators: Loaded generators to score. + only: Metric names; defaults to the spec's list. + include_expensive: Whether to run simulator-backed metrics. + on_error: `"raise"`, or `"skip"` to drop models that fail and carry + on -- useful for a long unattended sweep where one bad + checkpoint should not lose the rest of the results. + + Returns: + A DataFrame with provenance columns first, then metric columns. + """ + rows: list[dict[str, Any]] = [] + for generator in generators: + try: + rows.append(self.score(generator, only=only, include_expensive=include_expensive)) + # One unloadable checkpoint must not lose an entire unattended sweep. + except Exception as exc: # noqa: PERF203 + if on_error != "skip": + raise + print(f"[leaderboard] skipping {generator.algo_id}: {exc}") + return order_columns(pd.DataFrame(rows)) + + +def _as_columns(spec: MetricSpec, value: float | dict[str, float]) -> dict[str, float]: + """Normalize a metric's return value into named leaderboard columns.""" + if isinstance(value, dict): + unexpected = set(value) - set(spec.columns) + if unexpected: + raise ValueError(f"Metric {spec.name!r} returned undeclared columns: {sorted(unexpected)}") + return dict(value) + if len(spec.columns) != 1: + raise ValueError(f"Metric {spec.name!r} declares {spec.columns} but returned a single value.") + return {spec.columns[0]: value} + + +def order_columns(frame: pd.DataFrame) -> pd.DataFrame: + """Put provenance columns first, metric columns after, in a stable order.""" + if frame.empty: + return frame + lead = [col for col in PROVENANCE_COLUMNS if col in frame.columns] + rest = [col for col in frame.columns if col not in lead] + return frame[lead + rest] diff --git a/engiopt/evaluation/leaderboard.py b/engiopt/evaluation/leaderboard.py new file mode 100644 index 00000000..e50fc83d --- /dev/null +++ b/engiopt/evaluation/leaderboard.py @@ -0,0 +1,197 @@ +"""Persisting and publishing leaderboard results. + +Results are rows keyed by `ROW_KEY`. Locally that is a CSV; published, it is a +HuggingFace dataset repo holding the same CSV, so a local sweep and the public +board are literally the same table. + +Everything here merges rather than replaces. Adding one model means evaluating +one model and pushing its rows -- the existing board is downloaded, the new rows +are merged on `ROW_KEY`, and the result is uploaded. Nothing else is recomputed. +""" + +from __future__ import annotations + +from pathlib import Path +import tempfile +from typing import Any, TYPE_CHECKING + +import pandas as pd + +from engiopt.evaluation.evaluator import order_columns + +if TYPE_CHECKING: + from engiopt.evaluation.registry import MetricRegistry + +ROW_KEY = ["problem_id", "algo_id", "config_fingerprint", "seed", "spec_version"] +"""Columns that uniquely identify a leaderboard row. + +`config_fingerprint` is what separates two hyperparameter settings of the same +algorithm. Without it a 50-config sweep would collapse onto one row per seed, +silently keeping only whichever ran last. +""" + +LEADERBOARD_FILE = "leaderboard.csv" +"""Filename used inside a published HuggingFace dataset repo.""" + + +def append_rows(frame: pd.DataFrame, path: str | Path, *, deduplicate: bool = True) -> Path: + """Append rows to a CSV leaderboard, keeping the newest row per key. + + Args: + frame: New rows. + path: Destination CSV. Created with a header if absent. + deduplicate: Keep only the most recent row per `ROW_KEY`. Re-evaluating + a model then supersedes its old row instead of duplicating it. + + Returns: + The path written. + """ + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + existing = pd.read_csv(destination) if destination.exists() else pd.DataFrame() + combined = merge_rows(existing, frame) if deduplicate else order_columns(pd.concat([existing, frame])) + combined.to_csv(destination, index=False) + return destination + + +def rank( + frame: pd.DataFrame, + metric: str, + *, + registry: MetricRegistry | None = None, + aggregate: str = "median", +) -> pd.DataFrame: + """Rank models by one metric, aggregating across seeds. + + Args: + frame: Leaderboard rows. + metric: Column to rank by. + registry: Registry used to look up the metric's direction. + aggregate: How to combine seeds. Defaults to `median`, because + simulation-backed objectives are heavy-tailed and a single blown-up + run otherwise decides the ranking. + + Returns: + One row per `(problem_id, algo_id)` sorted best-first, with a `rank` + column and the seed count that produced each value. + """ + from engiopt.evaluation.registry import METRICS + + spec = (registry or METRICS)[metric] + if spec.higher_is_better is None: + raise ValueError(f"Metric {metric!r} has no ranking direction; it is diagnostic only.") + + grouped = ( + frame.groupby(["problem_id", "algo_id"], as_index=False) + .agg(value=(metric, aggregate), n_seeds=(metric, "count")) + .sort_values("value", ascending=not spec.higher_is_better) + .reset_index(drop=True) + ) + grouped["rank"] = grouped.index + 1 + return grouped.rename(columns={"value": f"{metric}_{aggregate}"}) + + +def disagreement( + frame: pd.DataFrame, + metrics: list[str], + *, + registry: MetricRegistry | None = None, + aggregate: str = "median", +) -> pd.DataFrame: + """Show each metric's ranking side by side, to expose where they disagree. + + Whether metrics agree on the winner is the question a leaderboard is + supposed to answer, so this is a first-class view rather than a notebook + one-off. + + Returns: + A table indexed by model, one column of ranks per metric. + """ + boards = [] + for metric in metrics: + board = rank(frame, metric, registry=registry, aggregate=aggregate) + boards.append(board.set_index(["problem_id", "algo_id"])["rank"].rename(metric)) + return pd.concat(boards, axis=1).sort_values(metrics[0]) + + +def merge_rows(existing: pd.DataFrame, new_rows: pd.DataFrame) -> pd.DataFrame: + """Combine two result tables, letting newer rows supersede older ones. + + Rows are matched on `ROW_KEY`, so re-evaluating a checkpoint replaces its + entry while a different hyperparameter setting, seed, or spec version adds + one. + """ + combined = pd.concat([existing, new_rows], ignore_index=True) if len(existing) else new_rows + key = [col for col in ROW_KEY if col in combined.columns] + if key: + combined = combined.drop_duplicates(subset=key, keep="last") + return order_columns(combined.reset_index(drop=True)) + + +def load_from_hub(repo_id: str, *, token: str | None = None) -> pd.DataFrame: + """Download the published leaderboard, or an empty frame if none exists yet.""" + from huggingface_hub import hf_hub_download + + try: + path = hf_hub_download(repo_id=repo_id, filename=LEADERBOARD_FILE, repo_type="dataset", token=token) + except Exception: # noqa: BLE001 - a missing or private board is not an error here + return pd.DataFrame() + return pd.read_csv(path) + + +def push_to_hub( + new_rows: pd.DataFrame, + repo_id: str, + *, + token: str | None = None, + private: bool = False, + commit_message: str | None = None, +) -> pd.DataFrame: + """Merge `new_rows` into the published leaderboard and upload the result. + + The existing board is downloaded first and merged on `ROW_KEY`, so adding a + model never requires re-running anyone else's evaluation. + + Args: + new_rows: Rows to publish, as returned by `Evaluator.leaderboard`. + repo_id: HuggingFace dataset repo, e.g. `"IDEALLab/engiopt-leaderboard"`. + token: HF token; falls back to the ambient login. + private: Create the repo private if it does not exist yet. + commit_message: Defaults to a summary of what was added. + + Returns: + The full merged leaderboard as uploaded. + """ + from huggingface_hub import HfApi + + api = HfApi(token=token) + api.create_repo(repo_id=repo_id, repo_type="dataset", private=private, exist_ok=True) + + merged = merge_rows(load_from_hub(repo_id, token=token), new_rows) + + models = ", ".join(sorted(new_rows["algo_id"].unique())) if "algo_id" in new_rows else "results" + with tempfile.TemporaryDirectory() as tmp: + local = Path(tmp) / LEADERBOARD_FILE + merged.to_csv(local, index=False) + api.upload_file( + path_or_fileobj=str(local), + path_in_repo=LEADERBOARD_FILE, + repo_id=repo_id, + repo_type="dataset", + commit_message=commit_message or f"Add {len(new_rows)} row(s): {models}", + ) + return merged + + +def already_evaluated(board: pd.DataFrame, **key: Any) -> bool: + """Whether the board already holds a row for the given `ROW_KEY` values. + + Lets a sweep skip work it has already done instead of recomputing it. + """ + if board.empty: + return False + mask = pd.Series(data=True, index=board.index) + for column, value in key.items(): + if column in board.columns: + mask &= board[column] == value + return bool(mask.any()) diff --git a/engiopt/evaluation/metrics/__init__.py b/engiopt/evaluation/metrics/__init__.py new file mode 100644 index 00000000..0fb6bed7 --- /dev/null +++ b/engiopt/evaluation/metrics/__init__.py @@ -0,0 +1,10 @@ +"""Metric implementations. + +Importing this package registers every built-in metric in +`engiopt.evaluation.registry.METRICS`. Third-party metrics register themselves +the same way: decorate a function with `@register_metric` and import the module. +""" + +from engiopt.evaluation.metrics import builtin + +__all__ = ["builtin"] diff --git a/engiopt/evaluation/metrics/builtin.py b/engiopt/evaluation/metrics/builtin.py new file mode 100644 index 00000000..3d597a31 --- /dev/null +++ b/engiopt/evaluation/metrics/builtin.py @@ -0,0 +1,119 @@ +"""Built-in evaluation metrics. + +These wrap the numerics in `engiopt.metrics` and declare, for each one, the +question it answers, whether it needs a simulator, and which direction is +better. Behaviour is identical to the per-model `evaluate_*.py` scripts these +replaced, so numbers remain comparable. + +Adding a metric means adding one decorated function here (or in your own module +-- importing it is enough to register it). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +from engiopt import metrics as metrics_mod +from engiopt.evaluation.registry import register_metric + +if TYPE_CHECKING: + from engiopt.evaluation.context import EvaluationContext + + +# ---------------------------------------------------------------------- +# Distribution: does the generator match the reference design distribution? +# ---------------------------------------------------------------------- + + +@register_metric( + "mmd", + family="distribution", + cost="cheap", + higher_is_better=False, + description="Maximum Mean Discrepancy between generated and reference designs (pixel space).", +) +def mmd(ctx: EvaluationContext) -> float: + """Maximum Mean Discrepancy between generated and reference designs.""" + return float(metrics_mod.mmd(ctx.gen_flat, ctx.ref_flat, sigma=ctx.sigma)) + + +# ---------------------------------------------------------------------- +# Diversity: is the generator producing varied designs, or repeating itself? +# ---------------------------------------------------------------------- + + +@register_metric( + "dpp", + family="diversity", + cost="cheap", + higher_is_better=True, + description="Determinantal Point Process diversity of the generated set.", +) +def dpp(ctx: EvaluationContext) -> float: + """Determinantal Point Process diversity of the generated designs.""" + return float(metrics_mod.dpp_diversity(ctx.gen_flat, sigma=ctx.sigma)) + + +# ---------------------------------------------------------------------- +# Feasibility: does the design respect its stated budget? +# ---------------------------------------------------------------------- + + +@register_metric( + "viol", + family="feasibility", + cost="expensive", + higher_is_better=False, + description="Fraction of designs missing their volume-fraction target beyond tolerance.", +) +def viol(ctx: EvaluationContext) -> float: + """Fraction of designs that miss their volume-fraction target. + + Marked expensive only because it is computed during the shared + optimizer/simulator pass; the arithmetic itself is trivial. + """ + values = ctx.optimization.viol + return float(np.mean(values)) if values else float("nan") + + +# ---------------------------------------------------------------------- +# Performance: how good are the designs once physics is applied? +# ---------------------------------------------------------------------- + + +@register_metric( + "iog", + family="performance", + cost="expensive", + higher_is_better=False, + description="Mean initial optimality gap: how far generated designs start from the reference optimum.", +) +def iog(ctx: EvaluationContext) -> float: + """Mean initial optimality gap, before any re-optimization.""" + return float(np.mean(ctx.optimization.iog)) + + +@register_metric( + "cog", + family="performance", + cost="expensive", + higher_is_better=False, + description="Mean cumulative optimality gap over re-optimization from each generated design.", +) +def cog(ctx: EvaluationContext) -> float: + """Mean cumulative optimality gap.""" + return float(np.mean(ctx.optimization.cog)) + + +@register_metric( + "fog", + family="performance", + cost="expensive", + higher_is_better=False, + description="Mean final optimality gap after re-optimization.", +) +def fog(ctx: EvaluationContext) -> float: + """Mean final optimality gap.""" + return float(np.mean(ctx.optimization.fog)) diff --git a/engiopt/evaluation/registry.py b/engiopt/evaluation/registry.py new file mode 100644 index 00000000..54db5113 --- /dev/null +++ b/engiopt/evaluation/registry.py @@ -0,0 +1,147 @@ +"""Registry of evaluation metrics. + +A metric belongs to the *comparison* (generated designs vs a reference set under +a problem), not to a model -- so metrics are registered functions taking an +`EvaluationContext`, rather than methods on a generator. + +Registering a metric is one decorated function:: + + @register_metric("mmd", family="distribution", cost="cheap", higher_is_better=False) + def mmd(ctx: EvaluationContext) -> float: + return engiopt.metrics.mmd(ctx.gen_designs, ctx.ref_designs, sigma=ctx.sigma) + +The declaration is what lets the evaluator keep simulation-free metrics strictly +separate from simulator-backed ones, and what lets the leaderboard know which +direction of each column is "better". +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass +from typing import Literal, TYPE_CHECKING + +if TYPE_CHECKING: + from engiopt.evaluation.context import EvaluationContext + +MetricCost = Literal["cheap", "expensive"] +"""`cheap` metrics never invoke a simulator or optimizer; `expensive` ones may.""" + +MetricFamily = Literal["feasibility", "conditions", "performance", "distribution", "diversity", "cost"] +"""The question a metric answers. Also the grouping used by leaderboard views.""" + +MetricFn = Callable[["EvaluationContext"], "float | dict[str, float]"] + + +@dataclass(frozen=True) +class MetricSpec: + """A registered metric and everything the evaluator needs to know about it.""" + + name: str + fn: MetricFn + family: MetricFamily + cost: MetricCost + higher_is_better: bool | None + """None for metrics with no intrinsic direction (e.g. a realized volume fraction).""" + outputs: tuple[str, ...] = () + """Column names produced. Defaults to `(name,)` for single-valued metrics.""" + description: str = "" + + @property + def columns(self) -> tuple[str, ...]: + """Leaderboard columns this metric fills.""" + return self.outputs or (self.name,) + + +class MetricRegistry(Mapping[str, MetricSpec]): + """Name-to-`MetricSpec` mapping populated by the `register_metric` decorator.""" + + def __init__(self) -> None: + self._metrics: dict[str, MetricSpec] = {} + + def __getitem__(self, name: str) -> MetricSpec: + if name not in self._metrics: + available = ", ".join(sorted(self._metrics)) + raise KeyError(f"Unknown metric {name!r}. Registered: {available}") + return self._metrics[name] + + def __iter__(self) -> Iterator[str]: + return iter(self._metrics) + + def __len__(self) -> int: + return len(self._metrics) + + def add(self, spec: MetricSpec) -> None: + """Register a metric, rejecting duplicate names and duplicate output columns.""" + if spec.name in self._metrics: + raise ValueError(f"Metric {spec.name!r} is already registered.") + taken = {col: owner.name for owner in self._metrics.values() for col in owner.columns} + for column in spec.columns: + if column in taken: + raise ValueError(f"Metric {spec.name!r} emits column {column!r}, already emitted by {taken[column]!r}.") + self._metrics[spec.name] = spec + + def select( + self, + names: list[str] | None = None, + *, + cost: MetricCost | None = None, + family: MetricFamily | None = None, + ) -> list[MetricSpec]: + """Return specs filtered by name, cost, and/or family, in registration order.""" + specs = [self[name] for name in names] if names is not None else list(self._metrics.values()) + if cost is not None: + specs = [spec for spec in specs if spec.cost == cost] + if family is not None: + specs = [spec for spec in specs if spec.family == family] + return specs + + def columns(self, specs: list[MetricSpec] | None = None) -> list[str]: + """All leaderboard columns for the given specs (default: every metric).""" + return [col for spec in (specs if specs is not None else self._metrics.values()) for col in spec.columns] + + +METRICS = MetricRegistry() +"""The global metric registry.""" + + +def register_metric( + name: str, + *, + family: MetricFamily, + cost: MetricCost, + higher_is_better: bool | None, + outputs: tuple[str, ...] = (), + description: str = "", + registry: MetricRegistry | None = None, +) -> Callable[[MetricFn], MetricFn]: + """Register an evaluation metric. + + Args: + name: Unique metric name, also the default output column. + family: Which engineering question this answers. + cost: `cheap` if it never runs a simulator, `expensive` otherwise. + higher_is_better: Ranking direction, or None if the metric is diagnostic. + outputs: Column names, when the metric returns a dict of several values. + description: One-line explanation, surfaced by `engiopt.evaluate --list-metrics`. + registry: Target registry; defaults to the global one. + + Returns: + The undecorated function, so it stays directly callable and testable. + """ + + def decorator(fn: MetricFn) -> MetricFn: + (registry or METRICS).add( + MetricSpec( + name=name, + fn=fn, + family=family, + cost=cost, + higher_is_better=higher_is_better, + outputs=outputs, + description=description or (fn.__doc__ or "").strip().split("\n")[0], + ) + ) + return fn + + return decorator diff --git a/engiopt/evaluation/spec.py b/engiopt/evaluation/spec.py new file mode 100644 index 00000000..51fcf55a --- /dev/null +++ b/engiopt/evaluation/spec.py @@ -0,0 +1,233 @@ +"""Versioned evaluation specifications. + +A leaderboard is only meaningful if every row was produced under identical +conditions. An `EvalSpec` freezes that contract: which problem, which test +conditions, how many samples, which metrics, and the kernel bandwidth. Every +leaderboard row records the `spec_version` it was produced under, so a change to +the protocol shows up as a new version rather than silently invalidating history. + +Specs live in `engiopt/specs//.json` and are committed, so +the contract is reviewable in the diff. +""" + +from __future__ import annotations + +from dataclasses import asdict +from dataclasses import dataclass +import hashlib +import json +from pathlib import Path +from typing import Any, TYPE_CHECKING + +import numpy as np +import torch as th + +from engiopt.dataset_sample_conditions import sample_conditions +from engiopt.transforms import get_scalar_condition_keys + +if TYPE_CHECKING: + from datasets import Dataset + from engibench.core import Problem + import numpy.typing as npt + +SPEC_ROOT = Path(__file__).resolve().parent.parent / "specs" +"""Directory holding committed specs, `engiopt/specs//.json`.""" + + +@dataclass(frozen=True) +class EvalSpec: + """The frozen contract for evaluating models on one problem. + + Attributes: + problem_id: EngiBench problem registry key. + version: Spec version, e.g. `v1`. Bump on any change below. + n_samples: Number of conditions each model is scored on. + condition_seed: Seed used to draw the test conditions. + metrics: Metric names to compute, in leaderboard column order. + sigma: Gaussian-kernel bandwidth for MMD and DPP. + volfrac_tol: Tolerance for the volume-fraction violation check. + objective_weights: Fixed weights over `problem.objectives`, for + multi-objective problems. Summing or averaging objectives with + different units is not a valid default, so one of these two fields + is required whenever a problem has more than one objective. + objective_weight_condition: Name of a per-sample condition carrying the + trade-off instead of fixing it. thermoelastic2d's `weight` splits + the first two objectives as `(w, 1 - w)`, so a purely structural + sample (`w = 1`) is unaffected by thermal compliance. + condition_digest: Hash of the drawn conditions and reference designs. + Recomputed at evaluation time and compared, so an upstream dataset + change is caught instead of silently shifting every number. + engibench_version: EngiBench version the digest was recorded under. + notes: Free-text rationale for this version. + """ + + problem_id: str + version: str = "v1" + n_samples: int = 50 + condition_seed: int = 1 + metrics: tuple[str, ...] = ("mmd", "dpp", "viol", "iog", "cog", "fog") + sigma: float = 10.0 + volfrac_tol: float = 0.01 + objective_weights: tuple[float, ...] | None = None + objective_weight_condition: str | None = None + condition_digest: str | None = None + engibench_version: str | None = None + notes: str = "" + + def __post_init__(self) -> None: + """Normalize sequence fields to tuples. + + JSON has no tuple type, so a spec read back from disk would otherwise + carry lists and compare unequal to the spec that produced it. + """ + object.__setattr__(self, "metrics", tuple(self.metrics)) + if self.objective_weights is not None: + object.__setattr__(self, "objective_weights", tuple(self.objective_weights)) + + # ------------------------------------------------------------------ + # Persistence + # ------------------------------------------------------------------ + + @classmethod + def load(cls, reference: str, *, root: Path | None = None) -> EvalSpec: + """Load a spec from `"/"` or a path to a JSON file.""" + path = Path(reference) + if path.suffix == ".json" and path.exists(): + return cls(**json.loads(path.read_text())) + problem_id, _, version = reference.partition("/") + version = version or "v1" + spec_path = (root or SPEC_ROOT) / problem_id / f"{version}.json" + if not spec_path.exists(): + raise FileNotFoundError( + f"No eval spec at {spec_path}. Create one with `EvalSpec(problem_id={problem_id!r}).freeze(problem)`." + ) + return cls(**json.loads(spec_path.read_text())) + + def save(self, *, root: Path | None = None) -> Path: + """Write this spec to its canonical committed location.""" + spec_path = (root or SPEC_ROOT) / self.problem_id / f"{self.version}.json" + spec_path.parent.mkdir(parents=True, exist_ok=True) + spec_path.write_text(json.dumps(asdict(self), indent=2, sort_keys=True) + "\n") + return spec_path + + # ------------------------------------------------------------------ + # Resolution + # ------------------------------------------------------------------ + + def resolve(self, problem: Problem, *, device: th.device | None = None) -> ResolvedSpec: + """Draw this spec's test conditions and reference designs. + + Raises: + ValueError: If the drawn data no longer matches `condition_digest`, + meaning the upstream dataset changed and the spec must be + re-frozen under a new version. + """ + device = device or th.device("cpu") + conditions_tensor, conditions, ref_designs, indices = sample_conditions( + problem=problem, n_samples=self.n_samples, device=device, seed=self.condition_seed + ) + digest = _digest(indices, ref_designs) + if self.condition_digest is not None and digest != self.condition_digest: + raise ValueError( + f"Eval spec {self.problem_id}/{self.version} no longer reproduces its frozen conditions " + f"(expected {self.condition_digest}, got {digest}). The dataset or sampling changed: " + "freeze a new spec version rather than comparing across the change." + ) + return ResolvedSpec( + spec=self, + conditions_tensor=conditions_tensor, + conditions=conditions, + ref_designs=ref_designs, + indices=np.asarray(indices), + condition_keys=tuple(get_scalar_condition_keys(problem, problem.dataset["test"])), + ) + + def freeze(self, problem: Problem) -> EvalSpec: + """Return a copy with `condition_digest` and `engibench_version` filled in.""" + import engibench + + _, _, ref_designs, indices = sample_conditions( + problem=problem, n_samples=self.n_samples, device=th.device("cpu"), seed=self.condition_seed + ) + return EvalSpec( + **{ + **asdict(self), + "condition_digest": _digest(indices, ref_designs), + "engibench_version": getattr(engibench, "__version__", None), + } + ) + + +@dataclass(frozen=True) +class ResolvedSpec: + """An `EvalSpec` with its conditions and reference designs materialized.""" + + spec: EvalSpec + conditions_tensor: th.Tensor + conditions: Dataset + ref_designs: npt.NDArray[Any] + indices: npt.NDArray[Any] + condition_keys: tuple[str, ...] = () + """Names of the columns in `conditions_tensor`, in order. + + A subset of `problem.conditions_keys`: solver-only and array-valued + conditions cannot travel in a dense tensor. See + `engiopt.transforms.get_scalar_condition_keys`. + """ + + @property + def n_samples(self) -> int: + """Number of test conditions in this spec.""" + return self.spec.n_samples + + +def _digest(indices: npt.NDArray[Any], ref_designs: npt.NDArray[Any]) -> str: + """Stable hash of the drawn sample indices and their reference designs.""" + hasher = hashlib.sha256() + hasher.update(np.asarray(indices).astype(np.int64).tobytes()) + hasher.update(np.asarray(ref_designs, dtype=np.float64).round(8).tobytes()) + return hasher.hexdigest()[:16] + + +def freeze_spec( + problem_id: str, + *, + version: str = "v1", + n_samples: int = 50, + condition_seed: int = 1, + metrics: tuple[str, ...] = ("mmd", "dpp", "viol", "iog", "cog", "fog"), + sigma: float = 10.0, + objective_weights: tuple[float, ...] | None = None, + objective_weight_condition: str | None = None, +) -> Path: + """Draw a problem's test conditions once and commit them as a spec. + + Run this once per problem, then never again for that version -- the point of + a spec is that it stops moving. + + Returns: + Path to the written spec file. + """ + from engibench.utils.all_problems import BUILTIN_PROBLEMS + + problem = BUILTIN_PROBLEMS[problem_id]() + problem.reset(seed=condition_seed) + spec = EvalSpec( + problem_id=problem_id, + version=version, + n_samples=n_samples, + condition_seed=condition_seed, + metrics=metrics, + sigma=sigma, + objective_weights=objective_weights, + objective_weight_condition=objective_weight_condition, + ).freeze(problem) + path = spec.save() + print(f"Froze {problem_id}/{version}: digest={spec.condition_digest} n={spec.n_samples} -> {path}") + return path + + +if __name__ == "__main__": + import tyro + + tyro.cli(freeze_spec) diff --git a/engiopt/generators/_template/__init__.py b/engiopt/generators/_template/__init__.py new file mode 100644 index 00000000..4649b0c6 --- /dev/null +++ b/engiopt/generators/_template/__init__.py @@ -0,0 +1,5 @@ +"""Template package for a new generator. Copy this directory to add a model. + +The leading underscore keeps it out of the registry, so the template itself +never shows up on the leaderboard. +""" diff --git a/engiopt/generators/_template/adapter.py b/engiopt/generators/_template/adapter.py new file mode 100644 index 00000000..791cfef9 --- /dev/null +++ b/engiopt/generators/_template/adapter.py @@ -0,0 +1,106 @@ +"""Template `Generator` adapter -- copy this directory to add a model. + +Replace `my_model` throughout, point the imports at your training script, and +you are on the leaderboard. Nothing outside this directory needs to change: the +registry finds your model by walking `engiopt/generators/`. + +Checklist: + +1. Rename the directory to your `algo_id` (they must match). +2. Set the class attributes below to describe your model. +3. Implement `build` -- rebuild the network from the checkpoint package. +4. Implement `_sample` -- turn conditions into designs. +5. Run: `python -m engiopt.evaluate --problem-id beams2d --generators my_model` + +See `CONTRIBUTING_A_MODEL.md` for the full walkthrough. +""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator + +if TYPE_CHECKING: + from engibench.core import Problem + + from engiopt.checkpoint_store import ResolvedCheckpoint + + +class MyModel(Generator): + """One-line description of the model, shown by `--list-generators`.""" + + algo_id = "my_model" + """Must equal this directory's name.""" + + conditional = True + """Does `_sample` use the conditions it is given?""" + + design_kinds = ("2d",) + """Design spaces this model can serve: any of `1d`, `2d`, `3d`, `dict`.""" + + checkpoint_files = ("generator.pth",) + """Files your training script writes into the checkpoint package.""" + + primary_state_key = "generator" + """Key holding the state dict inside the loaded checkpoint.""" + + output_clip: tuple[Any, Any] | None = (1e-3, 1.0) + """Optional clamp on generated designs; `None` to leave outputs untouched.""" + + def __init__(self, net: th.nn.Module, latent_dim: int, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.net = net + self.latent_dim = latent_dim + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> MyModel: + """Rebuild the trained network from its checkpoint package. + + Fetching the package and choosing the device are already done for you: + `resolved.run_config` is your training script's `Args` exactly as saved, + and `resolved.files[name]` is a local path to each checkpoint file. + """ + config = resolved.run_config + + # Replace this block with your own network, e.g.: + # + # from engiopt.generators.my_model.my_model import Generator as MyNet + # + # net = MyNet( + # latent_dim=config["latent_dim"], + # design_shape=design_shape_of(problem), + # ).to(device) + # state = th.load(resolved.files["generator.pth"], map_location=device) + # net.load_state_dict(state[cls.primary_state_key]) + # net.eval() + net = cls._build_network(config, problem, device, resolved) + + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) + + @staticmethod + def _build_network( + config: dict[str, Any], + problem: Problem, + device: th.device, + resolved: ResolvedCheckpoint, + ) -> th.nn.Module: + """Construct the network and load its weights. Replace this entirely.""" + raise NotImplementedError("Build your network here.") + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + """Draw `n` designs. + + `sample` (the public method) already handles seeding, timing, reshaping + to `(n, *design_shape)`, and clamping -- so return whatever shape is + natural for your network. + + Use `conditions.require_tensor(self.algo_id)` for the usual `(n, n_conds)` + tensor, or `conditions.dataset` if you need the original columns. + """ + cond = conditions.require_tensor(self.algo_id) + z = th.randn((n, self.latent_dim), device=self.device) + return self.net(z, cond) diff --git a/engiopt/generators/cgan_1d/evaluate_cgan_1d.py b/engiopt/generators/cgan_1d/evaluate_cgan_1d.py deleted file mode 100644 index a383bd00..00000000 --- a/engiopt/generators/cgan_1d/evaluate_cgan_1d.py +++ /dev/null @@ -1,135 +0,0 @@ -"""Evaluation for the CGAN 1D.""" - -from __future__ import annotations - -import dataclasses -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -from gymnasium import spaces -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.generators.cgan_1d.cgan_1d import Generator -from engiopt.generators.cgan_1d.cgan_1d import prepare_data -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions - - -@dataclasses.dataclass -class Args: - """Command-line arguments.""" - - problem_id: str = "airfoil" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 10 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "cgan_1d_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Seeding for reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - # Select device - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - if isinstance(problem.design_space, spaces.Box): - design_shape = problem.design_space.shape - else: - dummy_design, _ = problem.random_design() - design_shape = spaces.flatten(problem.design_space, dummy_design).shape - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, - n_samples=args.n_samples, - device=device, - seed=seed, - ) - - ### Set Up Generator ### - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="cgan_1d", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["generator.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"generator.pth": f"{args.problem_id}_cgan_1d_generator"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["generator.pth"] - ckpt = th.load(ckpt_path, map_location=device) - - _, conds_normalizer, design_normalizer = prepare_data(problem, device) - - model = Generator( - latent_dim=run_config["latent_dim"], - n_conds=len(problem.conditions_keys), - design_shape=design_shape, - design_normalizer=design_normalizer, - conds_normalizer=conds_normalizer, - ).to(device) - model.load_state_dict(ckpt["generator"]) - model.eval() - - # Sample noise and generate designs - z = th.randn((args.n_samples, run_config["latent_dim"]), device=device) - gen_designs = model(z, conditions_tensor) - gen_designs_np = gen_designs.detach().cpu().numpy() - print(gen_designs_np.shape) - - fail_ratio = metrics.simulate_failure_ratio( - problem=problem, - gen_designs=gen_designs_np, - sampled_conditions=sampled_conditions, - ) - - # Append result row to CSV - results_dict = { - "problem_id": args.problem_id, - "model_id": "cgan_1d", - "seed": seed, - "n_samples": args.n_samples, - "fail_ratio": fail_ratio, - } - metrics_df = pd.DataFrame(results_dict, index=[0]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/cgan_2d/evaluate_cgan_2d.py b/engiopt/generators/cgan_2d/evaluate_cgan_2d.py deleted file mode 100644 index 614c13d3..00000000 --- a/engiopt/generators/cgan_2d/evaluate_cgan_2d.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Evaluation for the CGAN 2D.""" - -from __future__ import annotations - -import dataclasses -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.generators.cgan_2d.cgan_2d import Generator -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions - - -@dataclasses.dataclass -class Args: - """Command-line arguments for a single-seed CGAN 2D evaluation.""" - - problem_id: str = "beams2d" - """Problem identifier (e.g. beams2d).""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name (if any).""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "cgan_2d_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - # Select device - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, - n_samples=args.n_samples, - device=device, - seed=seed, - ) - - ### Set Up Generator ### - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="cgan_2d", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["generator.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"generator.pth": f"{args.problem_id}_cgan_2d_generator"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["generator.pth"] - ckpt = th.load(ckpt_path, map_location=device) - - model = Generator( - latent_dim=run_config["latent_dim"], - n_conds=len(problem.conditions_keys), - design_shape=problem.design_space.shape, - generator_output_activation=run_config.get("generator_output_activation", "tanh"), - ).to(device) - model.load_state_dict(ckpt["generator"]) - model.eval() - - # Sample noise and generate designs - z = th.randn((args.n_samples, run_config["latent_dim"]), device=device) - gen_designs = model(z, conditions_tensor) - gen_designs_np = gen_designs.detach().cpu().numpy() - gen_designs_np = np.clip(gen_designs_np, 1e-3, 1.0) - - # Compute metrics - metrics_dict = metrics.metrics( - problem, - gen_designs_np, - sampled_designs_np, - sampled_conditions, - sigma=args.sigma, - ) - - # Add metadata to metrics - metrics_dict.update( - { - "seed": seed, - "problem_id": args.problem_id, - "model_id": "cgan_2d", - "n_samples": args.n_samples, - "sigma": args.sigma, - } - ) - - # Append result row to CSV - metrics_df = pd.DataFrame([metrics_dict]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/cgan_cnn_2d/evaluate_cgan_cnn_2d.py b/engiopt/generators/cgan_cnn_2d/evaluate_cgan_cnn_2d.py deleted file mode 100644 index 758f78b2..00000000 --- a/engiopt/generators/cgan_cnn_2d/evaluate_cgan_cnn_2d.py +++ /dev/null @@ -1,135 +0,0 @@ -"""Evaluation for the CGAN 2D w/ CNN.""" - -from __future__ import annotations - -import dataclasses -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.generators.cgan_cnn_2d.cgan_cnn_2d import Generator -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions - - -@dataclasses.dataclass -class Args: - """Command-line arguments for a single-seed cGAN CNN 2D evaluation.""" - - problem_id: str = "beams2d" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "cgan_cnn_2d_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, n_samples=args.n_samples, device=device, seed=seed - ) - - # Reshape to match the expected input shape for the model - conditions_tensor = conditions_tensor.unsqueeze(-1).unsqueeze(-1) - - ### Set Up Generator ### - - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="cgan_cnn_2d", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["generator.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"generator.pth": f"{args.problem_id}_cgan_cnn_2d_generator"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["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_keys), - design_shape=problem.design_space.shape, - generator_output_activation=run_config.get("generator_output_activation", "tanh"), - ) - model.load_state_dict(ckpt["generator"]) - model.eval() # Set to evaluation mode - model.to(device) - - # Sample noise as generator input - z = th.randn((args.n_samples, run_config["latent_dim"], 1, 1), device=device, dtype=th.float) - - # Generate a batch of designs - gen_designs = model(z, conditions_tensor) - gen_designs_np = gen_designs.detach().cpu().numpy() - gen_designs_np = gen_designs_np.reshape(args.n_samples, *problem.design_space.shape) - - # Clip to boundaries for running THIS IS PROBLEM DEPENDENT - gen_designs_np = np.clip(gen_designs_np, 1e-3, 1) - - # Compute metrics - metrics_dict = metrics.metrics( - problem, - gen_designs_np, - sampled_designs_np, - sampled_conditions, - sigma=args.sigma, - ) - - metrics_dict.update( - { - "seed": seed, - "problem_id": args.problem_id, - "model_id": "cgan_cnn_2d", - "n_samples": args.n_samples, - "sigma": args.sigma, - } - ) - - # Append result row to CSV - metrics_df = pd.DataFrame([metrics_dict]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/cgan_cnn_3d/evaluate_cgan_cnn_3d.py b/engiopt/generators/cgan_cnn_3d/evaluate_cgan_cnn_3d.py deleted file mode 100644 index 2fda1059..00000000 --- a/engiopt/generators/cgan_cnn_3d/evaluate_cgan_cnn_3d.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Evaluation for the CGAN 3D w/ CNN.""" - -from __future__ import annotations - -import dataclasses -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.generators.cgan_cnn_3d.cgan_cnn_3d import Generator3D -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions - - -@dataclasses.dataclass -class Args: - """Command-line arguments for a single-seed cGAN CNN 3D evaluation.""" - - problem_id: str = "heatconduction3d" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "cgan_cnn_3d_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, n_samples=args.n_samples, device=device, seed=seed - ) - - # 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_keys), 1, 1, 1) - - ### Set Up Generator ### - - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="cgan_cnn_3d", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["generator_3d.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"generator_3d.pth": f"{args.problem_id}_cgan_cnn_3d_generator_3d"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["generator_3d.pth"] - ckpt = th.load(ckpt_path, map_location=th.device(device)) - # Safer debug output - 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 - ) - model.load_state_dict(ckpt["generator"]) - model.eval() # Set to evaluation mode - model.to(device) - - # Sample noise as generator input - z = th.randn((args.n_samples, run_config["latent_dim"], 1, 1, 1), device=device, dtype=th.float) - - # Generate a batch of designs - gen_designs = model(z, conditions_tensor) - print("gen_designs.shape:", gen_designs.shape) - gen_designs_np = gen_designs.squeeze(1).detach().cpu().numpy() - # Removal of Padding - crop_start = (64 - 51) // 2 # 6 - crop_end = crop_start + 51 # 6 + 51 = 57 - - gen_designs_np = gen_designs_np[:, crop_start:crop_end, crop_start:crop_end, crop_start:crop_end] - - # Clip to boundaries for running THIS IS PROBLEM DEPENDENT - gen_designs_np = np.clip(gen_designs_np, 1e-3, 1) - - # Compute metrics - metrics_dict = metrics.metrics( - problem, - gen_designs_np, - sampled_designs_np, - sampled_conditions, - sigma=args.sigma, - ) - - metrics_dict.update( - { - "seed": seed, - "problem_id": args.problem_id, - "model_id": "cgan_cnn_3d", - "n_samples": args.n_samples, - "sigma": args.sigma, - } - ) - - # Append result row to CSV - metrics_df = pd.DataFrame([metrics_dict]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/cgan_vae/evaluate_cgan_vae.py b/engiopt/generators/cgan_vae/evaluate_cgan_vae.py deleted file mode 100644 index a4fcebc3..00000000 --- a/engiopt/generators/cgan_vae/evaluate_cgan_vae.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Evaluation for the Multiview VAE cGAN 3D.""" - -from __future__ import annotations - -import dataclasses -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.generators.cgan_vae.cgan_vae import Generator3D -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions - - -@dataclasses.dataclass -class Args: - """Command-line arguments for a single-seed Multiview VAE cGAN 3D evaluation.""" - - problem_id: str = "heatconduction3d" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "cgan_vae_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, n_samples=args.n_samples, device=device, seed=seed - ) - - # 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_keys), 1, 1, 1) - - ### Set Up Generator ### - - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="cgan_vae", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["multiview_3d_vaegan.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"multiview_3d_vaegan.pth": f"{args.problem_id}_cgan_vae_models"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["multiview_3d_vaegan.pth"] - ckpt = th.load(ckpt_path, map_location=th.device(device)) - # Safer debug output - for key in ckpt: - print("Checkpoint key:", key) - model = Generator3D( - 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 - model.to(device) - - # Sample noise as generator input - z = th.randn((args.n_samples, run_config["latent_dim"], 1, 1, 1), device=device, dtype=th.float) - - # Generate a batch of designs - gen_designs = model(z, conditions_tensor) - print("gen_designs.shape:", gen_designs.shape) - gen_designs_np = gen_designs.squeeze(1).detach().cpu().numpy() - # Removal of Padding - crop_start = (64 - 51) // 2 # 6 - crop_end = crop_start + 51 # 6 + 51 = 57 - - gen_designs_np = gen_designs_np[:, crop_start:crop_end, crop_start:crop_end, crop_start:crop_end] - - # Clip to boundaries for running THIS IS PROBLEM DEPENDENT - gen_designs_np = np.clip(gen_designs_np, 1e-3, 1) - - # Compute metrics - metrics_dict = metrics.metrics( - problem, - gen_designs_np, - sampled_designs_np, - sampled_conditions, - sigma=args.sigma, - ) - - metrics_dict.update( - { - "seed": seed, - "problem_id": args.problem_id, - "model_id": "cgan_vae", - "n_samples": args.n_samples, - "sigma": args.sigma, - } - ) - - # Append result row to CSV - metrics_df = pd.DataFrame([metrics_dict]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/diffusion_1d/evaluate_diffusion_1d.py b/engiopt/generators/diffusion_1d/evaluate_diffusion_1d.py deleted file mode 100644 index 20b14581..00000000 --- a/engiopt/generators/diffusion_1d/evaluate_diffusion_1d.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Evaluation for the Diffusion 1D.""" - -from __future__ import annotations - -import dataclasses -import os - -from denoising_diffusion_pytorch import GaussianDiffusion1D -from denoising_diffusion_pytorch import Unet1D -from engibench.utils.all_problems import BUILTIN_PROBLEMS -from gymnasium import spaces -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.generators.diffusion_1d.diffusion_1d import prepare_data - - -@dataclasses.dataclass -class Args: - """Command-line arguments.""" - - problem_id: str = "airfoil" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 10 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "diffusion_1d_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Seeding for reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - # Select device - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - if isinstance(problem.design_space, spaces.Box): - design_shape = problem.design_space.shape - else: - dummy_design, _ = problem.random_design() - flattened = spaces.flatten(problem.design_space, dummy_design) - design_shape = np.array(flattened).shape - - # Add padding for the UNet (1D requires the input to be divisible by 8) - padding_size = (8 - design_shape[0] % 8) % 8 # Only pad if needed - padded_size = design_shape[0] + padding_size - if padding_size > 0: - print(f"Padding design from {design_shape[0]} to {padded_size} dimensions") - design_shape = (padded_size,) - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, - n_samples=args.n_samples, - device=device, - seed=seed, - ) - - ### Load Diffusion Model ### - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="diffusion_1d", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["model.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"model.pth": f"{args.problem_id}_diffusion_1d_model"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["model.pth"] - ckpt = th.load(ckpt_path, map_location=device) - - _, design_normalizer = prepare_data(problem, padding_size, device) - - model = Unet1D( - dim=run_config["unet_dim"], # Used for the sinusoidal positional embeddings - channels=run_config["n_channels"], # Number of channels in the input - ).to(device) - - diffusion = GaussianDiffusion1D( - model, - seq_length=np.prod(design_shape), - auto_normalize=run_config.get("auto_norm", True), - ).to(device) - - diffusion.load_state_dict(ckpt["model"]) - - # Sample noise and generate designs - gen_designs = diffusion.sample(args.n_samples).squeeze(1) - gen_designs = design_normalizer.denormalize(gen_designs) - gen_designs_np = gen_designs.detach().cpu().numpy() - if padding_size > 0: - gen_designs_np = gen_designs_np[:, :-padding_size] - - fail_ratio = metrics.simulate_failure_ratio( - problem=problem, - gen_designs=gen_designs_np, - sampled_conditions=sampled_conditions, - ) - - # Append result row to CSV - results_dict = { - "problem_id": args.problem_id, - "model_id": "diffusion_1d", - "seed": seed, - "n_samples": args.n_samples, - "fail_ratio": fail_ratio, - } - metrics_df = pd.DataFrame(results_dict, index=[0]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/diffusion_2d_cond/evaluate_diffusion_2d_cond.py b/engiopt/generators/diffusion_2d_cond/evaluate_diffusion_2d_cond.py deleted file mode 100644 index a17a5f3d..00000000 --- a/engiopt/generators/diffusion_2d_cond/evaluate_diffusion_2d_cond.py +++ /dev/null @@ -1,175 +0,0 @@ -"""Evaluation for the Diffusion 2d_cond w/ seed looping and CSV saving.""" - -from __future__ import annotations - -import dataclasses -import os - -from diffusers import UNet2DConditionModel -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro -import wandb - -from engiopt import metrics -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import beta_schedule -from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import denormalize_designs_from_diffusion_range -from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import DiffusionSampler - - -@dataclasses.dataclass -class Args: - """Command-line arguments for a single-seed Diffusion 2D Conditional evaluation.""" - - problem_id: str = "beams2d" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "diffusion_2d_cond_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Seeding for reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - # Select device - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, - n_samples=args.n_samples, - device=device, - seed=seed, - ) - # Add channel dim - conditions_tensor = conditions_tensor.unsqueeze(1) - - ### Set Up Diffusion Model ### - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="diffusion_2d_cond", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["model.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"model.pth": f"{args.problem_id}_diffusion_2d_cond_model"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["model.pth"] - ckpt = th.load(ckpt_path, map_location=device) - - # Build UNet - model = UNet2DConditionModel( - sample_size=problem.design_space.shape, - in_channels=1, - out_channels=1, - cross_attention_dim=64, - block_out_channels=(32, 64, 128, 256), - down_block_types=("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D"), - 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_keys), - only_cross_attention=True, - ).to(device) - - # Noise schedule - options = { - "cosine": run_config["noise_schedule"] == "cosine", - "exp_biasing": run_config["noise_schedule"] == "exp", - "exp_bias_factor": 1, - } - betas = beta_schedule( - t=run_config["num_timesteps"], - start=1e-4, - end=0.02, - scale=1.0, - options=options, - ) - ddm_sampler = DiffusionSampler(run_config["num_timesteps"], betas) - - model.load_state_dict(ckpt["model"]) - model.eval() - - # Generate and reshape - design_shape: tuple = problem.design_space.shape - gen_designs = th.randn((args.n_samples, 1, *design_shape), device=device) - assert run_config["num_timesteps"] is not None - for i in reversed(range(run_config["num_timesteps"])): - t = th.full((args.n_samples,), i, device=device, dtype=th.long) - gen_designs = ddm_sampler.sample_timestep(model, gen_designs, t, conditions_tensor) - - gen_designs = gen_designs.squeeze(1) - if "design_min" in ckpt and "design_max" in ckpt: - design_min = ckpt["design_min"].to(device) - design_max = ckpt["design_max"].to(device) - gen_designs = denormalize_designs_from_diffusion_range(gen_designs, design_min, design_max) - gen_designs_np = gen_designs.detach().cpu().numpy().reshape(args.n_samples, *problem.design_space.shape) - if "design_min" in ckpt and "design_max" in ckpt: - design_min_np = ckpt["design_min"].detach().cpu().numpy() - design_max_np = ckpt["design_max"].detach().cpu().numpy() - gen_designs_np = np.clip(gen_designs_np, design_min_np, design_max_np) - else: - gen_designs_np = np.clip(gen_designs_np, 1e-3, 1.0) - - # Compute metrics - metrics_dict = metrics.metrics( - problem, - gen_designs_np, - sampled_designs_np, - sampled_conditions, - sigma=args.sigma, - ) - # Add metadata to metrics - metrics_dict.update( - { - "seed": seed, - "problem_id": args.problem_id, - "model_id": "diffusion_2d_cond", - "n_samples": args.n_samples, - "sigma": args.sigma, - } - ) - - # Append result row to CSV - metrics_df = pd.DataFrame([metrics_dict]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/gan_1d/evaluate_gan_1d.py b/engiopt/generators/gan_1d/evaluate_gan_1d.py deleted file mode 100644 index a977c86f..00000000 --- a/engiopt/generators/gan_1d/evaluate_gan_1d.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Evaluation for the GAN 1D.""" - -from __future__ import annotations - -import dataclasses -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -from gymnasium import spaces -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.generators.gan_1d.gan_1d import Generator -from engiopt.generators.gan_1d.gan_1d import prepare_data - - -@dataclasses.dataclass -class Args: - """Command-line arguments.""" - - problem_id: str = "airfoil" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "gan_1d_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Seeding for reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - # Select device - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - if isinstance(problem.design_space, spaces.Box): - design_shape = problem.design_space.shape - else: - dummy_design, _ = problem.random_design() - design_shape = spaces.flatten(problem.design_space, dummy_design).shape - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, - n_samples=args.n_samples, - device=device, - seed=seed, - ) - - ### Set Up Generator ### - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="gan_1d", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["generator.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"generator.pth": f"{args.problem_id}_gan_1d_generator"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["generator.pth"] - ckpt = th.load(ckpt_path, map_location=device) - - _, design_normalizer = prepare_data(problem, device) - - model = Generator( - latent_dim=run_config["latent_dim"], design_shape=design_shape, design_normalizer=design_normalizer - ).to(device) - model.load_state_dict(ckpt["generator"]) - model.eval() - - # Sample noise and generate designs - z = th.randn((args.n_samples, run_config["latent_dim"]), device=device) - gen_designs = model(z) - gen_designs_np = gen_designs.detach().cpu().numpy() - - fail_ratio = metrics.simulate_failure_ratio( - problem=problem, - gen_designs=gen_designs_np, - sampled_conditions=sampled_conditions, - ) - - # Append result row to CSV - results_dict = { - "problem_id": args.problem_id, - "model_id": "gan_1d", - "seed": seed, - "n_samples": args.n_samples, - "fail_ratio": fail_ratio, - } - metrics_df = pd.DataFrame(results_dict, index=[0]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/gan_2d/evaluate_gan_2d.py b/engiopt/generators/gan_2d/evaluate_gan_2d.py deleted file mode 100644 index c9721a02..00000000 --- a/engiopt/generators/gan_2d/evaluate_gan_2d.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Evaluation for the GAN 2D.""" - -from __future__ import annotations - -import dataclasses -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.generators.gan_2d.gan_2d import Generator - - -@dataclasses.dataclass -class Args: - """Command-line arguments.""" - - problem_id: str = "beams2d" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "gan_2d_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Seeding for reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - # Select device - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, - n_samples=args.n_samples, - device=device, - seed=seed, - ) - - ### Set Up Generator ### - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="gan_2d", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["generator.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"generator.pth": f"{args.problem_id}_gan_2d_generator"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["generator.pth"] - ckpt = th.load(ckpt_path, map_location=device) - - model = Generator( - latent_dim=run_config["latent_dim"], - design_shape=problem.design_space.shape, - ).to(device) - model.load_state_dict(ckpt["generator"]) - model.eval() - - # Sample noise and generate designs - z = th.randn((args.n_samples, run_config["latent_dim"]), device=device) - gen_designs = model(z) - gen_designs_np = gen_designs.detach().cpu().numpy() - gen_designs_np = np.clip(gen_designs_np, 1e-3, 1.0) - - # Compute metrics - metrics_dict = metrics.metrics( - problem, - gen_designs_np, - sampled_designs_np, - sampled_conditions, - sigma=args.sigma, - ) - - # Add metadata to metrics - metrics_dict.update( - { - "seed": seed, - "problem_id": args.problem_id, - "model_id": "gan_2d", - "n_samples": args.n_samples, - "sigma": args.sigma, - } - ) - - # Append result row to CSV - metrics_df = pd.DataFrame([metrics_dict]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/gan_bezier/evaluate_gan_bezier.py b/engiopt/generators/gan_bezier/evaluate_gan_bezier.py deleted file mode 100644 index 43361dd8..00000000 --- a/engiopt/generators/gan_bezier/evaluate_gan_bezier.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Evaluation for the GAN Bezier.""" - -from __future__ import annotations - -import dataclasses -import os -from typing import TYPE_CHECKING - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.generators.gan_bezier.gan_bezier import Generator -from engiopt.generators.gan_bezier.gan_bezier import prepare_data -from engiopt.transforms import flatten_dict_factory - -if TYPE_CHECKING: - from gymnasium import spaces - -_EPS = 1e-7 - - -@dataclasses.dataclass -class Args: - """Command-line arguments.""" - - problem_id: str = "airfoil" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "gan_bezier_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Seeding for reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - # Select device - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - coords_space: spaces.Box = problem.design_space["coords"] - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, _ = sample_conditions( - problem=problem, - n_samples=args.n_samples, - device=device, - seed=seed, - ) - - ### Set Up Generator ### - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="gan_bezier", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["bezier_generator.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"bezier_generator.pth": f"{args.problem_id}_gan_bezier_generator"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["bezier_generator.pth"] - ckpt = th.load(ckpt_path, map_location=device) - - _, design_scalars_normalizer, _ = prepare_data(problem, args.n_samples, device) - - model = Generator( - latent_dim=run_config["latent_dim"], - noise_dim=run_config["noise_dim"], - n_control_points=run_config["bezier_control_pts"], - n_data_points=coords_space.shape[1], - design_scalars_normalizer=design_scalars_normalizer, - eps=_EPS, - scalar_features=1, - ).to(device) - model.load_state_dict(ckpt["generator"]) - model.eval() - - # Sample noise and generate designs - bounds = (0.0, 1.0) # Bounds for angle of attack - c = (bounds[1] - bounds[0]) * th.rand(args.n_samples, run_config["latent_dim"], device=device) + bounds[0] - z = 0.5 * th.randn(args.n_samples, run_config["noise_dim"], device=device) - gen_designs, _, _, _, _, alphas = model(c, z) - - gen_designs_np = gen_designs.detach().cpu().numpy() - alphas_np = alphas.detach().cpu().numpy() - - # Reshape as dict - gen_designs_dict = [] - for idx, design in enumerate(gen_designs_np): - d = {"coords": design, "angle_of_attack": alphas_np[idx][0]} - gen_designs_dict.append(d) - - # Flatten dict for metrics - transform = flatten_dict_factory(problem, device) - transformed_gen_designs = transform(gen_designs_dict) - transformed_gen_designs = transformed_gen_designs.cpu().numpy() - - fail_ratio = metrics.simulate_failure_ratio( - problem=problem, - gen_designs=transformed_gen_designs, - sampled_conditions=sampled_conditions, - ) - - # Append result row to CSV - results_dict = { - "problem_id": args.problem_id, - "model_id": "gan_bezier", - "seed": seed, - "n_samples": args.n_samples, - "fail_ratio": fail_ratio, - } - metrics_df = pd.DataFrame(results_dict, index=[0]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/gan_cnn_2d/evaluate_gan_cnn_2d.py b/engiopt/generators/gan_cnn_2d/evaluate_gan_cnn_2d.py deleted file mode 100644 index 6c71121f..00000000 --- a/engiopt/generators/gan_cnn_2d/evaluate_gan_cnn_2d.py +++ /dev/null @@ -1,131 +0,0 @@ -"""Evaluation for the CGAN 2D w/ CNN.""" - -from __future__ import annotations - -import dataclasses -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.generators.gan_cnn_2d.gan_cnn_2d import Generator - - -@dataclasses.dataclass -class Args: - """Command-line arguments.""" - - problem_id: str = "beams2d" - """Problem identifier (e.g. beams2d).""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name (if any).""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "gan_cnn_2d_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Seeding - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - ### Set up testing conditions ### - conditions_tensor, sampled_conditions, sampled_designs_np, selected_indices = sample_conditions( - problem=problem, n_samples=args.n_samples, device=device, seed=seed - ) - - # Reshape to match the expected input shape for the model - conditions_tensor = conditions_tensor.unsqueeze(-1).unsqueeze(-1) - - ### Set Up Generator ### - - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="gan_cnn_2d", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["generator.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"generator.pth": f"{args.problem_id}_gan_cnn_2d_generator"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["generator.pth"] - ckpt = th.load(ckpt_path, map_location=th.device(device)) - model = Generator(latent_dim=run_config["latent_dim"], design_shape=problem.design_space.shape) - model.load_state_dict(ckpt["generator"]) - model.eval() # Set to evaluation mode - model.to(device) - - # Sample noise as generator input - z = th.randn((args.n_samples, run_config["latent_dim"], 1, 1), device=device, dtype=th.float) - - # Generate a batch of designs - gen_designs = model(z) - gen_designs_np = gen_designs.detach().cpu().numpy() - gen_designs_np = gen_designs_np.reshape(args.n_samples, *problem.design_space.shape) - - # Clip to boundaries for running THIS IS PROBLEM DEPENDENT - gen_designs_np = np.clip(gen_designs_np, 1e-3, 1) - - # Compute metrics - metrics_dict = metrics.metrics( - problem, - gen_designs_np, - sampled_designs_np, - sampled_conditions, - sigma=args.sigma, - ) - - # Add metadata to metrics - metrics_dict.update( - { - "seed": seed, - "problem_id": args.problem_id, - "model_id": "gan_cnn_2d", - "n_samples": args.n_samples, - "sigma": args.sigma, - } - ) - - # Append result row to CSV - metrics_df = pd.DataFrame([metrics_dict]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py b/engiopt/generators/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py deleted file mode 100644 index 89d3dc55..00000000 --- a/engiopt/generators/pixel_cnn_pp_2d/evaluate_pixel_cnn_pp_2d.py +++ /dev/null @@ -1,165 +0,0 @@ -"""Evaluation of the PixelCNN++ model.""" - -from __future__ import annotations - -from dataclasses import dataclass -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.generators.pixel_cnn_pp_2d.pixel_cnn_pp_2d import PixelCNNpp -from engiopt.generators.pixel_cnn_pp_2d.pixel_cnn_pp_2d import sample_from_discretized_mix_logistic - - -@dataclass -class Args: - """Command-line arguments for a single-seed PixelCNN++ 2D evaluation.""" - - problem_id: str = "beams2d" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - sampling_batch_size: int = 5 - """Batch size to use during sampling.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF org/user where checkpoints are stored.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used for model-family repositories.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "pixel_cnn_pp_2d_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Seeding for reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - # Select device - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - # Set up testing conditions - conditions_tensor, sampled_conditions, sampled_designs_np, selected_indices = sample_conditions( - problem=problem, - n_samples=args.n_samples, - device=device, - seed=seed, - ) - - # adapt to PixelCNN++ input shape requirements - conditions_tensor = conditions_tensor.unsqueeze(-1).unsqueeze(-1) - design_shape = (problem.design_space.shape[0], problem.design_space.shape[1]) - - # Set up PixelCNN++ model - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="pixel_cnn_pp_2d", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["model.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"model.pth": f"{args.problem_id}_pixel_cnn_pp_2d_model"}, - ) - run_config = resolved.run_config - - ckpt_path = resolved.files["model.pth"] - ckpt = th.load(ckpt_path, map_location=device) - - model = PixelCNNpp( - nr_resnet=run_config["nr_resnet"], - nr_filters=run_config["nr_filters"], - nr_logistic_mix=run_config["nr_logistic_mix"], - resnet_nonlinearity=run_config["resnet_nonlinearity"], - dropout_p=run_config["dropout_p"], - input_channels=1, - nr_conditions=conditions_tensor.shape[1], - ) - - model.load_state_dict(ckpt["model"]) - model.eval() - model.to(device) - - # Sample designs in batches - all_batches: list[th.Tensor] = [] - - for start in range(0, args.n_samples, args.sampling_batch_size): - end = min(args.n_samples, start + args.sampling_batch_size) - b = end - start - - # prepare batch-local tensors on the same device as the model - batch_conds = conditions_tensor[start:end] - data = th.zeros((b, 1, *design_shape), device=device) - - # Autoregressive pixel sampling for this batch - for i in range(design_shape[0]): - for j in range(design_shape[1]): - out = model(data, batch_conds) - out_sample = sample_from_discretized_mix_logistic(out, run_config["nr_logistic_mix"]) - data[:, :, i, j] = out_sample.data[:, :, i, j] - - # move completed batch to CPU to free GPU memory and store - all_batches.append(data.cpu()) - - # concatenate all batches on CPU - gen_designs = th.cat(all_batches, dim=0) - - gen_designs_np = gen_designs.detach().cpu().numpy() - gen_designs_np = gen_designs_np.reshape(args.n_samples, *problem.design_space.shape) # remove channel dim - gen_designs_np = np.clip(gen_designs_np, 1e-3, 1.0) - - # Compute metrics - metrics_dict = metrics.metrics( - problem, - gen_designs_np, - sampled_designs_np, - sampled_conditions, - sigma=args.sigma, - ) - - # Add metadata to metrics - metrics_dict.update( - { - "seed": seed, - "problem_id": args.problem_id, - "model_id": "pixel_cnn_pp_2d", - "n_samples": args.n_samples, - "sigma": args.sigma, - } - ) - - # Append result row to CSV - metrics_df = pd.DataFrame([metrics_dict]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/generators/vqgan/evaluate_vqgan.py b/engiopt/generators/vqgan/evaluate_vqgan.py deleted file mode 100644 index 61dbc693..00000000 --- a/engiopt/generators/vqgan/evaluate_vqgan.py +++ /dev/null @@ -1,228 +0,0 @@ -"""Evaluation for the VQGAN.""" - -from __future__ import annotations - -import dataclasses -import os - -from engibench.utils.all_problems import BUILTIN_PROBLEMS -import numpy as np -import pandas as pd -import torch as th -import tyro - -from engiopt import metrics -from engiopt.checkpoint_store import resolve_named_checkpoint -from engiopt.dataset_sample_conditions import sample_conditions -from engiopt.transforms import drop_constant -from engiopt.transforms import normalize -from engiopt.transforms import resize_to -from engiopt.generators.vqgan.vqgan import VQGAN -from engiopt.generators.vqgan.vqgan import VQGANTransformer - - -@dataclasses.dataclass -class Args: - """Command-line arguments for a single-seed VQGAN 2D evaluation.""" - - problem_id: str = "beams2d" - """Problem identifier.""" - seed: int = 1 - """Random seed to run.""" - wandb_project: str = "engiopt" - """Wandb project name.""" - wandb_entity: str | None = None - """Wandb entity name.""" - hf_entity: str = "IDEALLab" - """HF organization or user for checkpoint storage.""" - hf_repo_prefix: str = "engiopt" - """HF repo prefix used to build per-family model repos.""" - n_samples: int = 50 - """Number of generated samples per seed.""" - sigma: float = 10.0 - """Kernel bandwidth for MMD and DPP metrics.""" - output_csv: str = "vqgan_{problem_id}_metrics.csv" - """Output CSV path template; may include {problem_id}.""" - - -if __name__ == "__main__": - args = tyro.cli(Args) - - seed = args.seed - problem = BUILTIN_PROBLEMS[args.problem_id]() - problem.reset(seed=seed) - - # Reproducibility - th.manual_seed(seed) - rng = np.random.default_rng(seed) - th.backends.cudnn.deterministic = True - - if th.backends.mps.is_available(): - device = th.device("mps") - elif th.cuda.is_available(): - device = th.device("cuda") - else: - device = th.device("cpu") - - ### Set Up Transformer ### - - resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="vqgan", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["vqgan.pth", "transformer.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={ - "vqgan.pth": f"{args.problem_id}_vqgan_vqgan", - "transformer.pth": f"{args.problem_id}_vqgan_transformer", - }, - wandb_config_artifact_name=f"{args.problem_id}_vqgan_transformer", - ) - run_config = resolved.run_config - - ckpt_path_cvqgan = os.path.join(resolved.root_dir, "cvqgan.pth") - if run_config["conditional"] and not os.path.exists(ckpt_path_cvqgan): - cvqgan_resolved = resolve_named_checkpoint( - model_source="auto", - problem_id=args.problem_id, - algo="vqgan", - seed=seed, - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - required_files=["cvqgan.pth"], - wandb_project=args.wandb_project, - wandb_entity=args.wandb_entity, - wandb_artifact_names={"cvqgan.pth": f"{args.problem_id}_vqgan_cvqgan"}, - wandb_config_artifact_name=f"{args.problem_id}_vqgan_transformer", - ) - ckpt_path_cvqgan = cvqgan_resolved.files["cvqgan.pth"] - - ckpt_path_vqgan = resolved.files["vqgan.pth"] - ckpt_path_transformer = resolved.files["transformer.pth"] - ckpt_cvqgan = None - if os.path.exists(ckpt_path_cvqgan): - ckpt_cvqgan = th.load(ckpt_path_cvqgan, map_location=th.device(device), weights_only=False) - elif run_config["conditional"]: - raise FileNotFoundError("Conditional VQGAN evaluation requires cvqgan.pth, but no checkpoint was found.") - ckpt_vqgan = th.load(ckpt_path_vqgan, map_location=th.device(device), weights_only=False) - ckpt_transformer = th.load(ckpt_path_transformer, map_location=th.device(device), weights_only=False) - - vqgan = VQGAN( - device=device, - is_c=False, - encoder_channels=run_config["encoder_channels"], - encoder_start_resolution=run_config["image_size"], - encoder_attn_resolutions=run_config["encoder_attn_resolutions"], - encoder_num_res_blocks=run_config["encoder_num_res_blocks"], - decoder_channels=run_config["decoder_channels"], - decoder_start_resolution=run_config["latent_size"], - decoder_attn_resolutions=run_config["decoder_attn_resolutions"], - decoder_num_res_blocks=run_config["decoder_num_res_blocks"], - image_channels=run_config["image_channels"], - latent_dim=run_config["latent_dim"], - num_codebook_vectors=run_config["num_codebook_vectors"], - ) - vqgan.load_state_dict(ckpt_vqgan["vqgan"]) - vqgan.eval() # Set to evaluation mode - vqgan.to(device) - - cvqgan = VQGAN( - device=device, - is_c=True, - cond_feature_map_dim=run_config["cond_feature_map_dim"], - cond_dim=run_config["cond_dim"], - cond_hidden_dim=run_config["cond_hidden_dim"], - cond_latent_dim=run_config["cond_latent_dim"], - cond_codebook_vectors=run_config["cond_codebook_vectors"], - ) - if ckpt_cvqgan is not None: - cvqgan.load_state_dict(ckpt_cvqgan["cvqgan"]) - cvqgan.eval() # Set to evaluation mode - cvqgan.to(device) - - model = VQGANTransformer( - conditional=run_config["conditional"], - vqgan=vqgan, - cvqgan=cvqgan, - image_size=run_config["image_size"], - decoder_channels=run_config["decoder_channels"], - cond_feature_map_dim=run_config["cond_feature_map_dim"], - num_codebook_vectors=run_config["num_codebook_vectors"], - n_layer=run_config["n_layer"], - n_head=run_config["n_head"], - n_embd=run_config["n_embd"], - dropout=run_config["dropout"], - ) - model.load_state_dict(ckpt_transformer["transformer"]) - model.eval() # Set to evaluation mode - model.to(device) - - ### Set up testing conditions ### - _, sampled_conditions, sampled_designs_np, _sampled_designs_tensor = sample_conditions( - problem=problem, n_samples=args.n_samples, device=device, seed=seed - ) - - # Clean up conditions based on model training settings and convert back to tensor - sampled_conditions_new = sampled_conditions.select(range(len(sampled_conditions))) - conditions = sampled_conditions_new.column_names - - # Drop constant condition columns if enabled - if run_config["drop_constant_conditions"]: - sampled_conditions_new, conditions = drop_constant(sampled_conditions_new, sampled_conditions_new.column_names) - - # Normalize condition columns if enabled - if run_config["normalize_conditions"]: - sampled_conditions_new, mean, std = normalize(sampled_conditions_new, conditions) - - # Convert to tensor - conditions_tensor = th.stack([th.as_tensor(sampled_conditions_new[c][:]).float() for c in conditions], dim=1).to(device) - - # Set the start-of-sequence tokens for the transformer using the CVQGAN to discretize the conditions if enabled - if run_config["conditional"]: - c = model.encode_to_z(x=conditions_tensor, is_c=True)[1] - else: - c = th.ones(args.n_samples, 1, dtype=th.int64, device=device) * model.sos_token - - # Generate a batch of designs - latent_designs = model.sample( - x=th.empty(args.n_samples, 0, dtype=th.int64, device=device), c=c, steps=(run_config["latent_size"] ** 2) - ) - gen_designs = resize_to( - data=model.z_to_image(latent_designs), h=problem.design_space.shape[0], w=problem.design_space.shape[1] - ) - gen_designs_np = gen_designs.detach().cpu().numpy() - gen_designs_np = gen_designs_np.reshape(args.n_samples, *problem.design_space.shape) - - # Clip to boundaries for running THIS IS PROBLEM DEPENDENT - gen_designs_np = np.clip(gen_designs_np, 1e-3, 1) - - # Compute metrics - metrics_dict = metrics.metrics( - problem, - gen_designs_np, - sampled_designs_np, - sampled_conditions, - sigma=args.sigma, - ) - - metrics_dict.update( - { - "seed": seed, - "problem_id": args.problem_id, - "model_id": "vqgan", - "n_samples": args.n_samples, - "sigma": args.sigma, - } - ) - - # Append result row to CSV - metrics_df = pd.DataFrame([metrics_dict]) - out_path = args.output_csv.format(problem_id=args.problem_id) - write_header = not os.path.exists(out_path) - metrics_df.to_csv(out_path, mode="a", header=write_header, index=False) - - print(f"Seed {seed} done; appended to {out_path}") diff --git a/engiopt/specs/beams2d/v1.json b/engiopt/specs/beams2d/v1.json new file mode 100644 index 00000000..7a1b666e --- /dev/null +++ b/engiopt/specs/beams2d/v1.json @@ -0,0 +1,19 @@ +{ + "condition_digest": "67f84c4e50c355cf", + "condition_seed": 1, + "engibench_version": "0.2.0", + "metrics": [ + "mmd", + "dpp", + "viol", + "iog", + "cog", + "fog" + ], + "n_samples": 50, + "notes": "", + "problem_id": "beams2d", + "sigma": 10.0, + "version": "v1", + "volfrac_tol": 0.01 +} diff --git a/engiopt/specs/heatconduction2d/v1.json b/engiopt/specs/heatconduction2d/v1.json new file mode 100644 index 00000000..60595964 --- /dev/null +++ b/engiopt/specs/heatconduction2d/v1.json @@ -0,0 +1,19 @@ +{ + "condition_digest": "50961cce837117ed", + "condition_seed": 1, + "engibench_version": "0.2.0", + "metrics": [ + "mmd", + "dpp", + "viol", + "iog", + "cog", + "fog" + ], + "n_samples": 50, + "notes": "", + "problem_id": "heatconduction2d", + "sigma": 10.0, + "version": "v1", + "volfrac_tol": 0.01 +} diff --git a/engiopt/specs/photonics2d/v1.json b/engiopt/specs/photonics2d/v1.json new file mode 100644 index 00000000..8eaefcde --- /dev/null +++ b/engiopt/specs/photonics2d/v1.json @@ -0,0 +1,19 @@ +{ + "condition_digest": "ade3d18084b31b5e", + "condition_seed": 1, + "engibench_version": "0.2.0", + "metrics": [ + "mmd", + "dpp", + "viol", + "iog", + "cog", + "fog" + ], + "n_samples": 50, + "notes": "", + "problem_id": "photonics2d", + "sigma": 10.0, + "version": "v1", + "volfrac_tol": 0.01 +} diff --git a/engiopt/specs/thermoelastic2d/v1.json b/engiopt/specs/thermoelastic2d/v1.json new file mode 100644 index 00000000..291704db --- /dev/null +++ b/engiopt/specs/thermoelastic2d/v1.json @@ -0,0 +1,21 @@ +{ + "condition_digest": "3a1d2ae6c45add7f", + "condition_seed": 1, + "engibench_version": "0.2.0", + "metrics": [ + "mmd", + "dpp", + "viol", + "iog", + "cog", + "fog" + ], + "n_samples": 50, + "notes": "", + "objective_weight_condition": "weight", + "objective_weights": null, + "problem_id": "thermoelastic2d", + "sigma": 10.0, + "version": "v1", + "volfrac_tol": 0.01 +} From 2b9c58dc27d96df6eacfb2dc3a5ded9a9569ffab Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Fri, 24 Jul 2026 13:58:54 +0200 Subject: [PATCH 04/24] Decouple training from W&B, and fix four evaluation correctness bugs Training scripts: - 10 of 14 only saved checkpoints when --track was on, so training without W&B silently produced nothing on HuggingFace. VQGAN additionally gated early stopping -- real training logic -- behind it, so under default settings its transformer checkpoint was never written at all. - Each run is now filed under its own hyperparameter configuration via checkpoint_identity(args). Evaluation correctness: - iog stored the raw simulated objective rather than a gap against the reference optimum, unlike cog and fog. An oracle generator returning the reference design now scores ~0 where it previously reported raw compliance. - Condition sampling could not serve photonics2d (whose conditions include solver settings absent from the dataset) or thermoelastic2d (whose boundary conditions are 65x65 matrices). Scalar conditions now feed the model tensor while the simulator still receives every condition. - Multi-objective gaps were summed or averaged, letting an objective a sample does not prioritize dominate its score. Scalarization is declared in the eval spec, and the evaluator refuses to guess when it is absent. - Gaps now respect ObjectiveDirection, so photonics2d -- which maximizes total_overlap -- is no longer ranked in reverse. metrics.metrics() and simulate_failure_ratio() are removed: both became dead once the evaluator landed, and a second evaluation path only invites drift. Co-Authored-By: Claude Opus 4.8 (1M context) --- CONTRIBUTING_A_MODEL.md | 180 ++++++++++ README.md | 42 +-- docs/checkpoint_layout.md | 82 +++++ docs/checkpoint_migration_playbook.md | 96 ----- engiopt/dataset_sample_conditions.py | 28 +- engiopt/generators/cgan_1d/cgan_1d.py | 66 ++-- engiopt/generators/cgan_2d/cgan_2d.py | 66 ++-- engiopt/generators/cgan_bezier/cgan_bezier.py | 74 ++-- engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py | 72 ++-- engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py | 6 +- engiopt/generators/cgan_vae/cgan_vae.py | 6 +- .../generators/diffusion_1d/diffusion_1d.py | 48 +-- .../diffusion_2d_cond/diffusion_2d_cond.py | 58 +-- engiopt/generators/gan_1d/gan_1d.py | 66 ++-- engiopt/generators/gan_2d/gan_2d.py | 66 ++-- engiopt/generators/gan_bezier/gan_bezier.py | 8 +- engiopt/generators/gan_cnn_2d/gan_cnn_2d.py | 66 ++-- .../pixel_cnn_pp_2d/pixel_cnn_pp_2d.py | 58 +-- engiopt/generators/vqgan/utils.py | 4 +- engiopt/generators/vqgan/vqgan.py | 159 +++++---- engiopt/metrics.py | 178 +--------- engiopt/transforms.py | 43 +++ pyproject.toml | 5 + tests/conftest.py | 105 ++++++ tests/test_checkpoint_identity.py | 76 ++++ tests/test_checkpoint_store_layout.py | 56 +++ tests/test_conditions_and_gaps.py | 270 ++++++++++++++ tests/test_evaluation.py | 329 ++++++++++++++++++ tests/test_generator_registry.py | 115 ++++++ 29 files changed, 1752 insertions(+), 676 deletions(-) create mode 100644 CONTRIBUTING_A_MODEL.md create mode 100644 docs/checkpoint_layout.md delete mode 100644 docs/checkpoint_migration_playbook.md create mode 100644 tests/conftest.py create mode 100644 tests/test_checkpoint_identity.py create mode 100644 tests/test_checkpoint_store_layout.py create mode 100644 tests/test_conditions_and_gaps.py create mode 100644 tests/test_evaluation.py create mode 100644 tests/test_generator_registry.py diff --git a/CONTRIBUTING_A_MODEL.md b/CONTRIBUTING_A_MODEL.md new file mode 100644 index 00000000..85de95c5 --- /dev/null +++ b/CONTRIBUTING_A_MODEL.md @@ -0,0 +1,180 @@ +# Adding a model to EngiOpt + +EngiOpt keeps training and evaluation deliberately separate: + +- **Training** stays a single self-contained script per model, in the CleanRL + style. Write your algorithm however it reads best. There is no base class to + inherit and no framework to learn. +- **Evaluation** is shared. You describe your model once, through the + `Generator` contract, and every metric, leaderboard, and comparison works. + +The contract is two methods and a handful of class attributes. + +## 1. Create the package + +```bash +cp -r engiopt/generators/_template engiopt/generators/my_model +``` + +The directory name **is** your `algo_id`: it keys the registry, the leaderboard, +and your HuggingFace repo. Use lowercase with underscores. + +## 2. Write the training script + +Put it at `engiopt/generators/my_model/my_model.py`. Follow any existing model +for the shape of it — `cgan_cnn_2d/cgan_cnn_2d.py` is a good short one. + +The only requirement is how you save. Use the shared checkpoint helper so your +model lands where the loader expects: + +```python +from engiopt.checkpoint_store import save_checkpoint_package + +save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, # must equal your algo_id + seed=args.seed, + checkpoint_files={"generator.pth": "generator.pth"}, + run_config=vars(args), # everything needed to rebuild the model + primary_files=["generator.pth"], + **checkpoint_identity(args), # files this run under its own hyperparameters +) +``` + +`checkpoint_identity(args)` is what lets many hyperparameter settings coexist. +Each run is stored at `{problem_id}/cfg_{fingerprint}/seed_{seed}`, and a run +using your script's *default* hyperparameters additionally claims the canonical +`{problem_id}/seed_{seed}` -- which is what the bare model name resolves to. A +sweep therefore can never redefine what your model means. + +`run_config` matters: it is what `from_pretrained` reads to reconstruct your +network, so every architectural hyperparameter must be in `Args`. + +Also reuse the shared helpers rather than rewriting them: + +| Need | Use | +|---|---| +| Seeding | `engiopt.reproducibility.seed_training` | +| Deterministic dataloading | `engiopt.reproducibility.make_dataloader_generator` | +| Device selection | `engiopt.core.pick_device` | +| Design shape (incl. dict spaces) | `engiopt.core.design_shape_of` | + +## 3. Fill in the adapter + +Edit `engiopt/generators/my_model/adapter.py`. Declare what your model is: + +```python +class MyModel(Generator): + algo_id = "my_model" # == directory name + conditional = True # does _sample use the conditions? + design_kinds = ("2d",) # 1d / 2d / 3d / dict + checkpoint_files = ("generator.pth",) + primary_state_key = "generator" + output_clip = (1e-3, 1.0) # or None +``` + +Then implement the two methods: + +- **`build(resolved, problem, device, **base)`** — rebuild the network from + `resolved.run_config` and `resolved.files`, then hand it to `cls(...)`. + Fetching the package, choosing the device, and selecting the configuration + are already done for you. +- **`_sample(conditions, n)`** — produce `n` designs. Seeding, timing, reshaping + and clamping are handled for you; return whatever shape is natural. + +```python +@classmethod +def build(cls, resolved, problem, device, **base): + config = resolved.run_config + net = MyNet(latent_dim=config["latent_dim"], design_shape=problem.design_space.shape).to(device) + net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.eval() + return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) +``` + +Exactly one `Generator` subclass per adapter module — the same rule EngiBench +applies to `Problem` classes, so the name-to-model mapping stays unambiguous. + +### Conditions + +`_sample` receives a `ConditionBatch`, not a bare tensor: + +```python +cond = conditions.require_tensor(self.algo_id) # (n, n_conds) on your device +conditions.dataset # original columns, if you need them +conditions.keys # condition names, in column order +``` + +Most models want `require_tensor`. Reach for `dataset` only when your model +preprocesses conditions the way `vqgan` does (dropping constant columns, +re-normalizing). Unconditional models ignore the argument entirely — and the +conditional-adherence metrics are what will show that. + +## 4. Check it + +```bash +python -m engiopt.evaluate --list-generators # your model should appear +python -m engiopt.evaluate --problem-id beams2d --generators my_model +``` + +Adding `--include-expensive` also runs the simulator-backed metrics (COG, IOG, +FOG, feasibility). Those are slow; leave them off while iterating. + +## 5. Adding a metric instead + +Metrics are registered functions, not model methods — a metric compares a +generated set against a reference set under a problem, so it belongs to the +comparison rather than to any one model: + +```python +from engiopt.evaluation import register_metric + +@register_metric("my_metric", family="diversity", cost="cheap", higher_is_better=True) +def my_metric(ctx) -> float: + """One line, shown by --list-metrics.""" + return float(...) # ctx.gen_flat, ctx.ref_flat, ctx.conditions, ... +``` + +Declare `cost="expensive"` if it touches `ctx.optimization` (the simulator or +optimizer). The registry enforces the split, so a cheap run can never +accidentally launch a simulation. + +Importing the module is what registers it — which means you can define a metric +in a notebook cell and it will appear in the next leaderboard you build. + +## 6. Where things live + +| Artifact | Home | +|---|---| +| Weights, `run_config.json`, `metadata.json` | HuggingFace model repo | +| Evaluation metrics (`metrics.json`) | HuggingFace, beside the weights | +| The leaderboard | HuggingFace dataset repo | +| Loss curves, sample images | Weights & Biases, when `--track` is on | + +HuggingFace holds everything that must be reloaded or compared; W&B holds what +you only look at. Nothing in the evaluation path requires W&B -- training with +`--track false` produces exactly the same checkpoints and scores. + +To attach a model's scores to its own checkpoint: + +```bash +python -m engiopt.evaluate --problem-id beams2d --generators my_model --attach-metrics +``` + +## 7. Evaluation specs + +A leaderboard only means something if every row was measured the same way. Each +problem has a committed spec (`engiopt/specs//v1.json`) freezing the +test conditions, sample count, metric list, and kernel bandwidth. Every result +records the `spec_version` it came from. + +Do not edit a published spec. If the protocol has to change, freeze a new +version: + +```bash +python -m engiopt.evaluation.spec --problem-id my_problem --version v2 +``` diff --git a/README.md b/README.md index 83466098..349fe7d9 100644 --- a/README.md +++ b/README.md @@ -22,24 +22,24 @@ As much as we can, we follow the [CleanRL](https://github.com/vwxyzjn/cleanrl) p **Algorithm** | **Class** | **Dimensions** | **Conditional?** | **Model** --- | --- | --- | --- | --- -[cgan_1d](engiopt/cgan_1d/) | Inverse Design | 1D | ✅ | GAN MLP -[cgan_2d](engiopt/cgan_2d/) | Inverse Design | 2D | ✅ | GAN MLP -[cgan_bezier](engiopt/cgan_bezier/) | Inverse Design | 1D | ✅ | GAN + Bezier layer -[cgan_cnn_2d](engiopt/cgan_cnn_2d/) | Inverse Design | 2D | ✅ | GAN + CNN -[cgan_cnn_3d](engiopt/cgan_cnn_3d/) | Inverse Design | 3D | ✅ | GAN + 3D CNN -[cgan_vae](engiopt/cgan_vae/) | Inverse Design | 3D | ✅ | MultiView GAN + VAE -[diffusion_1d](engiopt/diffusion_1d/) | Inverse Design | 1D | ❌ | Diffusion -[diffusion_2d_cond](engiopt/diffusion_2d_cond/) | Inverse Design | 2D | ✅ | Diffusion -[gan_1d](engiopt/gan_1d/) | Inverse Design | 1D | ❌ | GAN MLP -[gan_2d](engiopt/gan_2d/) | Inverse Design | 2D | ❌ | GAN MLP -[gan_bezier](engiopt/gan_bezier/) | Inverse Design | 1D | ❌ | GAN + Bezier layer -[gan_cnn_2d](engiopt/gan_cnn_2d/) | Inverse Design | 2D | ❌ | GAN + CNN +[cgan_1d](engiopt/generators/cgan_1d/) | Inverse Design | 1D | ✅ | GAN MLP +[cgan_2d](engiopt/generators/cgan_2d/) | Inverse Design | 2D | ✅ | GAN MLP +[cgan_bezier](engiopt/generators/cgan_bezier/) | Inverse Design | 1D | ✅ | GAN + Bezier layer +[cgan_cnn_2d](engiopt/generators/cgan_cnn_2d/) | Inverse Design | 2D | ✅ | GAN + CNN +[cgan_cnn_3d](engiopt/generators/cgan_cnn_3d/) | Inverse Design | 3D | ✅ | GAN + 3D CNN +[cgan_vae](engiopt/generators/cgan_vae/) | Inverse Design | 3D | ✅ | MultiView GAN + VAE +[diffusion_1d](engiopt/generators/diffusion_1d/) | Inverse Design | 1D | ❌ | Diffusion +[diffusion_2d_cond](engiopt/generators/diffusion_2d_cond/) | Inverse Design | 2D | ✅ | Diffusion +[gan_1d](engiopt/generators/gan_1d/) | Inverse Design | 1D | ❌ | GAN MLP +[gan_2d](engiopt/generators/gan_2d/) | Inverse Design | 2D | ❌ | GAN MLP +[gan_bezier](engiopt/generators/gan_bezier/) | Inverse Design | 1D | ❌ | GAN + Bezier layer +[gan_cnn_2d](engiopt/generators/gan_cnn_2d/) | Inverse Design | 2D | ❌ | GAN + CNN [surrogate_model](engiopt/surrogate_model/) | Surrogate Model | 1D | ❌ | MLP -[vqgan](engiopt/vqgan) | Inverse Design | 2D | ✅ | VQVAE + Transformer -[pixel_cnn_pp_2d](engiopt/pixel_cnn_pp_2d) | Inverse Design | 2D | ✅ | PixelCNN++ Autoregressive Model +[vqgan](engiopt/generators/vqgan) | Inverse Design | 2D | ✅ | VQVAE + Transformer +[pixel_cnn_pp_2d](engiopt/generators/pixel_cnn_pp_2d) | Inverse Design | 2D | ✅ | PixelCNN++ Autoregressive Model ## Dashboards -The integration with WandB allows us to access live dashboards of our runs (on the cluster or not). New checkpoint packages are stored on the Hugging Face Hub by default, while WandB keeps experiment tracking, metadata, and links back to the canonical checkpoint location. Historical WandB model artifacts remain supported for backward compatibility. You can access some of our runs at https://wandb.ai/engibench/engiopt. +HuggingFace hosts everything that has to be reloaded or compared -- model weights, run configs, evaluation metrics, and the leaderboard. WandB hosts what you only look at: loss curves and sample images. Nothing in the evaluation path requires WandB, so training with `--track false` produces exactly the same checkpoints and scores. You can access some of our runs at https://wandb.ai/engibench/engiopt. WandB dashboards @@ -77,7 +77,7 @@ Usually, we provide two scripts per algorithm: one to train the model, and one t To train a model, you can run (for example): ``` -python engiopt/cgan_cnn_2d/cgan_cnn_2d.py --problem-id "beams2d" --track --wandb-entity None --save-model --n-epochs 200 --seed 1 +python engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py --problem-id "beams2d" --track --wandb-entity None --save-model --n-epochs 200 --seed 1 ``` This trains a CGAN 2D w/ CNN on `beams2d`. The flags mirror W&B's: `--track` enables W&B logging, `--wandb-entity`/`--wandb-project` say where the run goes, `--save-model` uploads the checkpoint to HuggingFace, and `--hf-entity`/`--hf-repo-prefix` say where the checkpoint goes. @@ -89,23 +89,23 @@ wandb login # for tracking huggingface-cli login # for checkpoints (or: export HF_TOKEN=...) ``` -The defaults (`--hf-entity IDEALLab --hf-repo-prefix engiopt`) push to `huggingface.co/IDEALLab/engiopt-cgan-cnn-2d/beams2d/seed_1/`. The W&B run summary records the HF path for traceability. Each checkpoint package contains the model files plus `run_config.json` and `metadata.json`, so evaluation needs no live W&B state. +The defaults (`--hf-entity IDEALLab --hf-repo-prefix engiopt`) push to `huggingface.co/IDEALLab/engiopt-cgan-cnn-2d/beams2d/cfg_/seed_1/`, one location per hyperparameter configuration. A run using the script's default hyperparameters additionally claims `beams2d/seed_1/`, which is what the bare model name resolves to. The W&B run summary records the HF path for traceability. Each checkpoint package contains the model files plus `run_config.json` and `metadata.json`, so evaluation needs no live W&B state. For reproducible debugging runs, you can additionally enable strict deterministic mode: ``` -python engiopt/cgan_cnn_2d/cgan_cnn_2d.py --problem-id "beams2d" --seed 1 --strict-determinism +python engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py --problem-id "beams2d" --seed 1 --strict-determinism ``` For new cGAN density-field runs, you can emit designs natively in the EngiBench `[0, 1]` density range while preserving older `tanh` checkpoint behavior by default: ``` -python engiopt/cgan_cnn_2d/cgan_cnn_2d.py --problem-id "beams2d" --generator-output-activation sigmoid +python engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py --problem-id "beams2d" --generator-output-activation sigmoid ``` Then evaluate: ``` -python engiopt/cgan_cnn_2d/evaluate_cgan_cnn_2d.py --problem-id "beams2d" --seed 1 --n-samples 10 +python -m engiopt.evaluate --problem-id "beams2d" --generators cgan_cnn_2d --seeds 1 ``` -Evaluation pulls the checkpoint from HF automatically. For runs trained before the HF cutover, evaluation transparently falls back to the legacy W&B artifact. Pass `--hf-entity` / `--hf-repo-prefix` to point at a different HF repo. +Evaluation pulls the checkpoint from HF automatically. Pass `--hf-entity` / `--hf-repo-prefix` to point at a different HF repo, and `--config-fingerprints` to score specific hyperparameter configurations instead of the default one. ### Surrogate model diff --git a/docs/checkpoint_layout.md b/docs/checkpoint_layout.md new file mode 100644 index 00000000..41ab9724 --- /dev/null +++ b/docs/checkpoint_layout.md @@ -0,0 +1,82 @@ +# Checkpoint layout + +HuggingFace is the single home for anything that has to be reloaded or compared: +model weights, the run config needed to rebuild them, evaluation metrics, and +the leaderboard. Weights & Biases, when enabled, hosts what you only look at: +loss curves and sample images. Nothing in the evaluation path requires W&B. + +## Where a checkpoint goes + +One repo per model family, named from the algorithm: + +``` +{hf_entity}/{hf_repo_prefix}-{algo} e.g. IDEALLab/engiopt-cgan-cnn-2d +``` + +Inside it, one package per hyperparameter configuration and seed: + +``` +{problem_id}/cfg_{fingerprint}/seed_{seed}/ every configuration, always written +{problem_id}/seed_{seed}/ the canonical default, see below +``` + +`fingerprint` is a short hash of the run config with infrastructure settings +(seed, tracking flags, HF/W&B destinations) excluded, so a sweep of fifty +configurations produces fifty independently addressable packages rather than +overwriting one. Computed by `engiopt.core.config_fingerprint`. + +The canonical `{problem_id}/seed_{seed}` path is what a bare model name resolves +to. It is written **only** when a run used the training script's default +hyperparameters, so a sweep can never redefine what `cgan_cnn_2d` means. + +Each package contains the model files plus: + +| file | contents | +|---|---| +| `run_config.json` | the training script's `Args`, enough to rebuild the model | +| `metadata.json` | problem, algo, seed, fingerprint, primary files, W&B run URL | +| `metrics.json` | evaluation scores, when `--attach-metrics` was used | + +## Writing one + +Training scripts call `save_checkpoint_package` and splat +`engiopt.core.checkpoint_identity(args)`, which supplies the fingerprint and +whether the run used default hyperparameters: + +```python +save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"generator.pth": "generator.pth"}, + run_config=vars(args), + **checkpoint_identity(args), +) +``` + +## Reading one + +```python +# the default-hyperparameter checkpoint -- what the bare model name means +gen = BUILTIN_GENERATORS["cgan_cnn_2d"].from_pretrained(problem, problem_id="beams2d", seed=1) + +# one specific configuration from a sweep +gen = BUILTIN_GENERATORS["cgan_cnn_2d"].from_pretrained( + problem, problem_id="beams2d", seed=1, config_fingerprint="3ab84748") +``` + +## The leaderboard + +Published as a HuggingFace **dataset** repo holding `leaderboard.csv`, one row +per `(problem_id, algo_id, config_fingerprint, seed, spec_version)`. Pushing +merges into the existing board rather than replacing it, so adding one model +never recomputes or disturbs anyone else's rows: + +```bash +python -m engiopt.evaluate --problem-id beams2d --generators my_model \ + --push-to IDEALLab/engiopt-leaderboard --attach-metrics +``` diff --git a/docs/checkpoint_migration_playbook.md b/docs/checkpoint_migration_playbook.md deleted file mode 100644 index 243d2a6d..00000000 --- a/docs/checkpoint_migration_playbook.md +++ /dev/null @@ -1,96 +0,0 @@ -# Checkpoint Migration Playbook - -This playbook describes how to migrate historical EngiOpt checkpoints from Weights & Biases (W&B) to the Hugging Face Hub (HF) without breaking backward compatibility. - -This is guidance only for phase 1 of the HF checkpoint backend rollout. - -What this phase does not do: -- delete historical W&B artifacts -- mutate historical W&B artifacts in place -- assume every historical run should be migrated - -## Goals - -1. Stop creating new long-lived checkpoint pressure in W&B by making HF the default backend for new saved models. -2. Keep historical W&B checkpoints loadable while HF-backed packages roll out. -3. Create a safe path to reclaim W&B storage later, after validation. - -## Recommended Scope - -Start with the official or release subset of checkpoints first, not the full historical backlog. - -Good early candidates: -- checkpoints referenced in papers, benchmarks, or public notebooks -- checkpoints linked from README examples -- checkpoints used by downstream evaluation scripts or case studies - -## Migration Package Format - -Each migrated HF checkpoint package should be self-contained and include: -- the original model weight file or files -- `run_config.json` -- `metadata.json` - -`metadata.json` should record at least: -- `problem_id` -- `algo` -- `seed` -- original W&B project, run id, and artifact names -- HF repo id and package path -- primary file list - -## Recommended Procedure - -1. Inventory the target artifacts. - Create a manifest with artifact name, aliases, run id, problem, algorithm, seed, expected files, and whether the artifact is part of the official/release subset. - -2. Freeze the migration manifest before uploads. - Use the manifest as the source of truth so the migration is reproducible and reviewable. - -3. Download the original W&B artifacts. - For each target artifact, download the stored files and extract the associated W&B run config needed to rebuild the model outside W&B. - -4. Build the HF package. - Upload the weights together with `run_config.json` and `metadata.json` into the per-model-family HF repo under the deterministic package path: - `problem_id[/extra_parts]/seed_` - -5. Record the mapping. - For every migrated checkpoint, store a durable mapping from the W&B artifact alias to the HF repo id, revision, and package path. - -6. Validate restores before cleanup. - Test a representative sample end-to-end with EngiOpt’s evaluation or surrogate-model loading path and confirm the HF package reproduces the expected model restore behavior. - -7. Add pointers back into W&B metadata. - Once validated, update run metadata, summaries, or notes so the WandB run points to the canonical HF checkpoint location. - -8. Keep an overlap period. - Retain both HF and W&B copies long enough to verify that downstream users and scripts are not relying on the old blob storage path. - -9. Only then define deletion policy. - Deletion or retention changes should happen in a separate maintenance pass, using the manifest as the authoritative record. - -## Validation Checklist - -Before any cleanup is considered for a migrated checkpoint: - -1. The HF package contains all required weight files. -2. `run_config.json` is present and sufficient to rebuild the model. -3. `metadata.json` is present and points back to the original W&B lineage. -4. A real EngiOpt load path has succeeded against the HF package. -5. The corresponding W&B run contains the HF pointer or mapping information. - -## Cleanup Guidance for Later - -When the project is ready to reclaim W&B storage: - -1. Start with official checkpoints only. -2. Delete in small batches, not all at once. -3. Confirm the manifest entry is complete before each deletion. -4. Re-run a small restore audit after each batch. -5. Keep at least one validated overlap window where both copies coexist. - -## Notes - -- W&B remains part of the lineage story even after HF becomes the canonical checkpoint host. -- HF should be treated as the long-lived storage backend for public or durable checkpoints. -- Backward compatibility matters more than immediate cleanup. diff --git a/engiopt/dataset_sample_conditions.py b/engiopt/dataset_sample_conditions.py index 1741e453..94766429 100644 --- a/engiopt/dataset_sample_conditions.py +++ b/engiopt/dataset_sample_conditions.py @@ -8,30 +8,46 @@ import numpy as np import torch as th +from engiopt.transforms import get_scalar_condition_keys + def sample_conditions( problem: Problem, n_samples: int, device: th.device, seed: int ) -> tuple[th.Tensor, Dataset, np.ndarray, np.ndarray]: """Samples conditions and designs from the dataset and prepares tensors for the generator. + The two outputs serve different consumers and deliberately differ: + + - `sampled_conditions` keeps every condition column the dataset provides, + including array-valued ones, because that is what the simulator and + optimizer need to reproduce the scenario. + - `conditions_tensor` keeps only the scalar conditions, because that is all a + generator can take as a dense `(n_samples, n_conds)` input. Conditions that + are solver settings rather than dataset columns (photonics2d) or that are + images rather than scalars (thermoelastic2d) are excluded; see + `engiopt.transforms.get_scalar_condition_keys`. + Args: problem (Problem): The problem containing the dataset with conditions and designs. n_samples (int): Number of samples to draw. - device (th.device): The device (e.g., 'cpu', 'mps', 'cuda') to place the tensors on. + device (th.device): The device (e.g. 'cpu', 'mps', 'cuda') to place the tensors on. seed (int): Random seed for reproducibility. Returns: - conditions_tensor: A PyTorch tensor of sampled conditions, reshaped for the generator. - sampled_conditions: A Hugging Face Dataset object of sampled conditions. + conditions_tensor: Scalar conditions as a (n_samples, n_scalar_conds) tensor. + sampled_conditions: A Hugging Face Dataset of every available condition column. sampled_designs_np: A NumPy array of sampled optimal designs. selected_indices: The indices of the sampled conditions and designs. """ ### Set up testing conditions ### rng = np.random.default_rng(seed) - # Extract the conditions + # Extract the conditions the dataset actually carries; the simulator wants all + # of them, the generator only the scalar ones. dataset = problem.dataset["test"] - conditions_ds = dataset.select_columns(problem.conditions_keys) + available_keys = [key for key in problem.conditions_keys if key in dataset.column_names] + scalar_keys = get_scalar_condition_keys(problem, dataset) + conditions_ds = dataset.select_columns(available_keys) # Sample conditions and test_ds designs at random indices selected_indices = rng.choice(len(dataset), n_samples, replace=True) @@ -39,6 +55,6 @@ def sample_conditions( sampled_designs_np = np.array(dataset["optimal_design"])[selected_indices] # Create tensor for conditions to be used in the generator - conditions_tensor = th.tensor(list(sampled_conditions[:].values()), dtype=th.float32, device=device).T + conditions_tensor = th.tensor([sampled_conditions[key] for key in scalar_keys], dtype=th.float32, device=device).T return conditions_tensor, sampled_conditions, sampled_designs_np, selected_indices diff --git a/engiopt/generators/cgan_1d/cgan_1d.py b/engiopt/generators/cgan_1d/cgan_1d.py index 8f28e6bb..b0e163b8 100644 --- a/engiopt/generators/cgan_1d/cgan_1d.py +++ b/engiopt/generators/cgan_1d/cgan_1d.py @@ -20,13 +20,14 @@ from torchvision import transforms import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import flatten_dict_factory +import wandb if TYPE_CHECKING: from engibench.utils.problem import Problem @@ -334,8 +335,8 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader) + i if args.track: - batches_done = epoch * len(dataloader) + i wandb.log( { "d_loss": d_loss.item(), @@ -385,35 +386,36 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: - ckpt_gen = { - "epoch": epoch, - "batches_done": batches_done, - "generator": generator.state_dict(), - "optimizer_generator": optimizer_generator.state_dict(), - "loss": g_loss.item(), - } - ckpt_disc = { - "epoch": epoch, - "batches_done": batches_done, - "discriminator": discriminator.state_dict(), - "optimizer_discriminator": optimizer_discriminator.state_dict(), - "loss": d_loss.item(), - } - - th.save(ckpt_gen, "generator.pth") - th.save(ckpt_disc, "discriminator.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, - run_config=vars(args), - primary_files=["generator.pth"], - ) + if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: + ckpt_gen = { + "epoch": epoch, + "batches_done": batches_done, + "generator": generator.state_dict(), + "optimizer_generator": optimizer_generator.state_dict(), + "loss": g_loss.item(), + } + ckpt_disc = { + "epoch": epoch, + "batches_done": batches_done, + "discriminator": discriminator.state_dict(), + "optimizer_discriminator": optimizer_discriminator.state_dict(), + "loss": d_loss.item(), + } + + th.save(ckpt_gen, "generator.pth") + th.save(ckpt_disc, "discriminator.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, + run_config=vars(args), + **checkpoint_identity(args), + primary_files=["generator.pth"], + ) wandb.finish() diff --git a/engiopt/generators/cgan_2d/cgan_2d.py b/engiopt/generators/cgan_2d/cgan_2d.py index cc238dc0..43b82180 100644 --- a/engiopt/generators/cgan_2d/cgan_2d.py +++ b/engiopt/generators/cgan_2d/cgan_2d.py @@ -18,12 +18,13 @@ from torch import nn import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +import wandb GeneratorOutputActivation = Literal["tanh", "sigmoid"] @@ -283,8 +284,8 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader) + i if args.track: - batches_done = epoch * len(dataloader) + i wandb.log( { "d_loss": d_loss.item(), @@ -326,35 +327,36 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: - ckpt_gen = { - "epoch": epoch, - "batches_done": batches_done, - "generator": generator.state_dict(), - "optimizer_generator": optimizer_generator.state_dict(), - "loss": g_loss.item(), - } - ckpt_disc = { - "epoch": epoch, - "batches_done": batches_done, - "discriminator": discriminator.state_dict(), - "optimizer_discriminator": optimizer_discriminator.state_dict(), - "loss": d_loss.item(), - } - - th.save(ckpt_gen, "generator.pth") - th.save(ckpt_disc, "discriminator.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, - run_config=vars(args), - primary_files=["generator.pth"], - ) + if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: + ckpt_gen = { + "epoch": epoch, + "batches_done": batches_done, + "generator": generator.state_dict(), + "optimizer_generator": optimizer_generator.state_dict(), + "loss": g_loss.item(), + } + ckpt_disc = { + "epoch": epoch, + "batches_done": batches_done, + "discriminator": discriminator.state_dict(), + "optimizer_discriminator": optimizer_discriminator.state_dict(), + "loss": d_loss.item(), + } + + th.save(ckpt_gen, "generator.pth") + th.save(ckpt_disc, "discriminator.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, + run_config=vars(args), + **checkpoint_identity(args), + primary_files=["generator.pth"], + ) wandb.finish() diff --git a/engiopt/generators/cgan_bezier/cgan_bezier.py b/engiopt/generators/cgan_bezier/cgan_bezier.py index d684c9c7..f5ae1aaf 100644 --- a/engiopt/generators/cgan_bezier/cgan_bezier.py +++ b/engiopt/generators/cgan_bezier/cgan_bezier.py @@ -20,6 +20,7 @@ import tyro from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training @@ -28,6 +29,8 @@ if TYPE_CHECKING: from collections.abc import Callable + from engibench.core import Problem + _EPS = 1e-7 @@ -120,7 +123,7 @@ def forward(self, x: th.Tensor) -> th.Tensor: class Deconv1DCombo(nn.Module): - def __init__( # noqa: PLR0913 + def __init__( self, in_channels: int, out_channels: int, @@ -241,7 +244,7 @@ class Generator(nn.Module): 3) Bezier => final design """ - def __init__( # noqa: PLR0913 + def __init__( self, latent_dim: int, noise_dim: int, @@ -301,7 +304,7 @@ def forward( class Discriminator(nn.Module): """Bezier GAN discriminator.""" - def __init__( # noqa: PLR0913 + def __init__( self, latent_dim: int, design_scalars: int, @@ -431,6 +434,49 @@ def denormalize(self, x: th.Tensor) -> th.Tensor: return x * (self.max_val - self.min_val + self.eps) + self.min_val +def prepare_data( + problem: Problem, device: th.device +) -> tuple[th.utils.data.TensorDataset, Normalizer, Normalizer, list[str]]: + """Build the training tensors and the normalizers derived from them. + + Extracted from the training loop so that evaluation can rebuild the exact + same normalizers: they are computed from dataset statistics rather than + stored in the checkpoint, so a generator cannot be reconstructed without + them. Matches the `prepare_data` helpers in the sibling 1D models. + + Args: + problem: The problem whose training split defines the statistics. + device: Device the normalizer bounds should live on. + + Returns: + The training dataset, the conditions normalizer, the design-scalars + normalizer, and the design scalar key names. + """ + problem_dataset = problem.dataset.with_format("torch")["train"] + design_scalar_keys = list(problem_dataset["optimal_design"][0].keys()) + design_scalar_keys.remove("coords") + coords_set = [problem_dataset[i]["optimal_design"]["coords"][:] for i in range(len(problem_dataset))] + design_scalars = [example["optimal_design"][key] for example in problem_dataset for key in design_scalar_keys] + training_ds = th.utils.data.TensorDataset( + th.stack(coords_set), + th.stack(design_scalars).unsqueeze(1), + *[problem_dataset[key][:] for key in problem.conditions_keys], + ) + + cond_tensors = th.stack(training_ds.tensors[2:]) + conds_min = cond_tensors.amin(dim=tuple(range(1, cond_tensors.ndim))).to(device) + conds_max = cond_tensors.amax(dim=tuple(range(1, cond_tensors.ndim))).to(device) + design_scalars_min = training_ds.tensors[1].amin(dim=0).to(device) + design_scalars_max = training_ds.tensors[1].amax(dim=0).to(device) + + return ( + training_ds, + Normalizer(conds_min, conds_max), + Normalizer(design_scalars_min, design_scalars_max), + design_scalar_keys, + ) + + if __name__ == "__main__": args = tyro.cli(Args) th.autograd.set_detect_anomaly(True) @@ -460,25 +506,8 @@ def denormalize(self, x: th.Tensor) -> th.Tensor: n_data_points = problem.design_space["coords"].shape[1] # for airfoil, 192 # The Discriminator uses shape [N, 2, #points]. - problem_dataset = problem.dataset.with_format("torch")["train"] - design_scalar_keys = list(problem_dataset["optimal_design"][0].keys()) - design_scalar_keys.remove("coords") - coords_set = [problem_dataset[i]["optimal_design"]["coords"][:] for i in range(len(problem_dataset))] - design_scalars = [example["optimal_design"][key] for example in problem_dataset for key in design_scalar_keys] - training_ds = th.utils.data.TensorDataset( - th.stack(coords_set), - th.stack(design_scalars).unsqueeze(1), - *[problem_dataset[key][:] for key in problem.conditions_keys], - ) - - cond_tensors = th.stack(training_ds.tensors[2:]) - conds_min = cond_tensors.amin(dim=tuple(range(1, cond_tensors.ndim))).to(device) - conds_max = cond_tensors.amax(dim=tuple(range(1, cond_tensors.ndim))).to(device) - design_scalars_min = training_ds.tensors[1].amin(dim=0).to(device) - design_scalars_max = training_ds.tensors[1].amax(dim=0).to(device) - - conds_normalizer = Normalizer(conds_min, conds_max) - design_scalars_normalizer = Normalizer(design_scalars_min, design_scalars_max) + training_ds, conds_normalizer, design_scalars_normalizer, design_scalar_keys = prepare_data(problem, device) + conds_min, conds_max = conds_normalizer.min_val, conds_normalizer.max_val dataloader = th.utils.data.DataLoader( training_ds, @@ -658,6 +687,7 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor, th.Tensor]: "bezier_discriminator.pth": "bezier_discriminator.pth", }, run_config=vars(args), + **checkpoint_identity(args), primary_files=["bezier_generator.pth"], ) diff --git a/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py b/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py index 00a882cb..d94092c2 100644 --- a/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py +++ b/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py @@ -17,12 +17,13 @@ from torchvision import transforms import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +import wandb GeneratorOutputActivation = Literal["tanh", "sigmoid"] @@ -379,8 +380,8 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader) + i if args.track: - batches_done = epoch * len(dataloader) + i wandb.log( { "d_loss": d_loss.item(), @@ -422,38 +423,39 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: - ckpt_gen = { - "epoch": epoch, - "batches_done": batches_done, - "generator": generator.state_dict(), - "optimizer_generator": optimizer_generator.state_dict(), - "loss": g_loss.item(), - } - ckpt_disc = { - "epoch": epoch, - "batches_done": batches_done, - "discriminator": discriminator.state_dict(), - "optimizer_discriminator": optimizer_discriminator.state_dict(), - "loss": d_loss.item(), - } - - th.save(ckpt_gen, "generator.pth") - th.save(ckpt_disc, "discriminator.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={ - "generator.pth": "generator.pth", - "discriminator.pth": "discriminator.pth", - }, - run_config=vars(args), - primary_files=["generator.pth"], - ) + if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: + ckpt_gen = { + "epoch": epoch, + "batches_done": batches_done, + "generator": generator.state_dict(), + "optimizer_generator": optimizer_generator.state_dict(), + "loss": g_loss.item(), + } + ckpt_disc = { + "epoch": epoch, + "batches_done": batches_done, + "discriminator": discriminator.state_dict(), + "optimizer_discriminator": optimizer_discriminator.state_dict(), + "loss": d_loss.item(), + } + + th.save(ckpt_gen, "generator.pth") + th.save(ckpt_disc, "discriminator.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={ + "generator.pth": "generator.pth", + "discriminator.pth": "discriminator.pth", + }, + run_config=vars(args), + **checkpoint_identity(args), + primary_files=["generator.pth"], + ) wandb.finish() diff --git a/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py b/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py index 6e3281af..40a7848e 100644 --- a/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py +++ b/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py @@ -20,14 +20,15 @@ from torch.nn import functional import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.metrics import dpp_diversity from engiopt.metrics import mmd from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +import wandb @dataclass @@ -330,7 +331,7 @@ def forward(self, x: th.Tensor, c: th.Tensor) -> th.Tensor: return self.final_conv(h) # -> (B, out_channels, 1, 1, 1) -def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, device, lambda_gp=20.0): # noqa: PLR0913 +def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, device, lambda_gp=20.0): """Calculates the gradient penalty loss for WGAN GP.""" batch_size = real_samples.size(0) # Random weight term for interpolation between real and fake samples @@ -696,6 +697,7 @@ def sample_3d_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: seed=args.seed, checkpoint_files={"generator_3d.pth": "generator_3d.pth", "discriminator_3d.pth": "discriminator_3d.pth"}, run_config=vars(args), + **checkpoint_identity(args), primary_files=["generator_3d.pth"], ) diff --git a/engiopt/generators/cgan_vae/cgan_vae.py b/engiopt/generators/cgan_vae/cgan_vae.py index e6fce8de..49c92955 100644 --- a/engiopt/generators/cgan_vae/cgan_vae.py +++ b/engiopt/generators/cgan_vae/cgan_vae.py @@ -20,14 +20,15 @@ from torch.nn import functional import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.metrics import dpp_diversity from engiopt.metrics import mmd from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +import wandb @dataclass @@ -423,7 +424,7 @@ def forward(self, x: th.Tensor, c: th.Tensor) -> th.Tensor: return self.final_conv(h) # -> (B, out_channels, 1, 1, 1) -def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, device, lambda_gp=20.0): # noqa: PLR0913 +def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, device, lambda_gp=20.0): """Calculates the gradient penalty loss for WGAN GP.""" batch_size = real_samples.size(0) # Random weight term for interpolation between real and fake samples @@ -823,6 +824,7 @@ def sample_3d_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: seed=args.seed, checkpoint_files={"multiview_3d_vaegan.pth": "multiview_3d_vaegan.pth"}, run_config=vars(args), + **checkpoint_identity(args), primary_files=["multiview_3d_vaegan.pth"], ) diff --git a/engiopt/generators/diffusion_1d/diffusion_1d.py b/engiopt/generators/diffusion_1d/diffusion_1d.py index 887a7209..9c2eeae7 100644 --- a/engiopt/generators/diffusion_1d/diffusion_1d.py +++ b/engiopt/generators/diffusion_1d/diffusion_1d.py @@ -20,13 +20,14 @@ from torchvision import transforms import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import flatten_dict_factory +import wandb if TYPE_CHECKING: from engibench.utils.problem import Problem @@ -228,8 +229,8 @@ class Args: # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader) + i if args.track: - batches_done = epoch * len(dataloader) + i wandb.log({"loss": loss.item(), "epoch": epoch, "batch": batches_done}) print(f"[Epoch {epoch}/{args.n_epochs}] [Batch {i}/{len(dataloader)}] [loss: {loss.item()}]") @@ -274,26 +275,27 @@ class Args: # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: - ckpt = { - "epoch": epoch, - "batches_done": batches_done, - "model": diffusion.state_dict(), - "loss": loss.item(), - } - - th.save(ckpt, "model.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={"model.pth": "model.pth"}, - run_config=vars(args), - primary_files=["model.pth"], - ) + if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: + ckpt = { + "epoch": epoch, + "batches_done": batches_done, + "model": diffusion.state_dict(), + "loss": loss.item(), + } + + th.save(ckpt, "model.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"model.pth": "model.pth"}, + run_config=vars(args), + **checkpoint_identity(args), + primary_files=["model.pth"], + ) wandb.finish() diff --git a/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py b/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py index c15af485..57e1f115 100644 --- a/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py +++ b/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py @@ -16,12 +16,13 @@ from torch.nn import functional import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +import wandb if TYPE_CHECKING: from collections.abc import Callable @@ -442,8 +443,8 @@ def sample_designs(model: UNet2DConditionModel, n_designs: int = 25) -> tuple[th # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader) + i if args.track: - batches_done = epoch * len(dataloader) + i wandb.log( { "loss": loss.item(), @@ -486,31 +487,32 @@ def sample_designs(model: UNet2DConditionModel, n_designs: int = 25) -> tuple[th # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: - ckpt_model = { - "epoch": epoch, - "batches_done": batches_done, - "model": model.state_dict(), - "optimizer_generator": optimizer.state_dict(), - "loss": loss.item(), - "design_min": filtered_ds_min.detach().cpu(), - "design_max": filtered_ds_max.detach().cpu(), - "diffusion_sample_min": DIFFUSION_SAMPLE_MIN, - "diffusion_sample_max": DIFFUSION_SAMPLE_MAX, - } - - th.save(ckpt_model, "model.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={"model.pth": "model.pth"}, - run_config=vars(args), - primary_files=["model.pth"], - ) + if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: + ckpt_model = { + "epoch": epoch, + "batches_done": batches_done, + "model": model.state_dict(), + "optimizer_generator": optimizer.state_dict(), + "loss": loss.item(), + "design_min": filtered_ds_min.detach().cpu(), + "design_max": filtered_ds_max.detach().cpu(), + "diffusion_sample_min": DIFFUSION_SAMPLE_MIN, + "diffusion_sample_max": DIFFUSION_SAMPLE_MAX, + } + + th.save(ckpt_model, "model.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"model.pth": "model.pth"}, + run_config=vars(args), + **checkpoint_identity(args), + primary_files=["model.pth"], + ) wandb.finish() diff --git a/engiopt/generators/gan_1d/gan_1d.py b/engiopt/generators/gan_1d/gan_1d.py index 446525a5..44858e42 100644 --- a/engiopt/generators/gan_1d/gan_1d.py +++ b/engiopt/generators/gan_1d/gan_1d.py @@ -19,13 +19,14 @@ from torchvision import transforms import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import flatten_dict_factory +import wandb if TYPE_CHECKING: from engibench.utils.problem import Problem @@ -279,8 +280,8 @@ def prepare_data(problem: Problem, device: th.device) -> tuple[th.utils.data.Ten # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader) + i if args.track: - batches_done = epoch * len(dataloader) + i wandb.log( { "d_loss": d_loss.item(), @@ -326,35 +327,36 @@ def prepare_data(problem: Problem, device: th.device) -> tuple[th.utils.data.Ten # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: - ckpt_gen = { - "epoch": epoch, - "batches_done": batches_done, - "generator": generator.state_dict(), - "optimizer_generator": optimizer_generator.state_dict(), - "loss": g_loss.item(), - } - ckpt_disc = { - "epoch": epoch, - "batches_done": batches_done, - "discriminator": discriminator.state_dict(), - "optimizer_discriminator": optimizer_discriminator.state_dict(), - "loss": d_loss.item(), - } - - th.save(ckpt_gen, "generator.pth") - th.save(ckpt_disc, "discriminator.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, - run_config=vars(args), - primary_files=["generator.pth"], - ) + if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: + ckpt_gen = { + "epoch": epoch, + "batches_done": batches_done, + "generator": generator.state_dict(), + "optimizer_generator": optimizer_generator.state_dict(), + "loss": g_loss.item(), + } + ckpt_disc = { + "epoch": epoch, + "batches_done": batches_done, + "discriminator": discriminator.state_dict(), + "optimizer_discriminator": optimizer_discriminator.state_dict(), + "loss": d_loss.item(), + } + + th.save(ckpt_gen, "generator.pth") + th.save(ckpt_disc, "discriminator.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, + run_config=vars(args), + **checkpoint_identity(args), + primary_files=["generator.pth"], + ) wandb.finish() diff --git a/engiopt/generators/gan_2d/gan_2d.py b/engiopt/generators/gan_2d/gan_2d.py index c2ed6d22..91723fc0 100644 --- a/engiopt/generators/gan_2d/gan_2d.py +++ b/engiopt/generators/gan_2d/gan_2d.py @@ -16,12 +16,13 @@ from torch import nn import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +import wandb @dataclass @@ -225,8 +226,8 @@ def sample_designs(n_designs: int) -> th.Tensor: # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader) + i if args.track: - batches_done = epoch * len(dataloader) + i wandb.log( { "d_loss": d_loss.item(), @@ -264,35 +265,36 @@ def sample_designs(n_designs: int) -> th.Tensor: # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: - ckpt_gen = { - "epoch": epoch, - "batches_done": batches_done, - "generator": generator.state_dict(), - "optimizer_generator": optimizer_generator.state_dict(), - "loss": g_loss.item(), - } - ckpt_disc = { - "epoch": epoch, - "batches_done": batches_done, - "discriminator": discriminator.state_dict(), - "optimizer_discriminator": optimizer_discriminator.state_dict(), - "loss": d_loss.item(), - } - - th.save(ckpt_gen, "generator.pth") - th.save(ckpt_disc, "discriminator.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, - run_config=vars(args), - primary_files=["generator.pth"], - ) + if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: + ckpt_gen = { + "epoch": epoch, + "batches_done": batches_done, + "generator": generator.state_dict(), + "optimizer_generator": optimizer_generator.state_dict(), + "loss": g_loss.item(), + } + ckpt_disc = { + "epoch": epoch, + "batches_done": batches_done, + "discriminator": discriminator.state_dict(), + "optimizer_discriminator": optimizer_discriminator.state_dict(), + "loss": d_loss.item(), + } + + th.save(ckpt_gen, "generator.pth") + th.save(ckpt_disc, "discriminator.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, + run_config=vars(args), + **checkpoint_identity(args), + primary_files=["generator.pth"], + ) wandb.finish() diff --git a/engiopt/generators/gan_bezier/gan_bezier.py b/engiopt/generators/gan_bezier/gan_bezier.py index 4c980920..ab10e522 100644 --- a/engiopt/generators/gan_bezier/gan_bezier.py +++ b/engiopt/generators/gan_bezier/gan_bezier.py @@ -20,6 +20,7 @@ import tyro from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training @@ -120,7 +121,7 @@ def forward(self, x: th.Tensor) -> th.Tensor: class Deconv1DCombo(nn.Module): - def __init__( # noqa: PLR0913 + def __init__( self, in_channels: int, out_channels: int, @@ -243,7 +244,7 @@ class Generator(nn.Module): 3) Bezier => final design """ - def __init__( # noqa: PLR0913 + def __init__( self, latent_dim: int, noise_dim: int, @@ -297,7 +298,7 @@ def forward( class Discriminator(nn.Module): """Bezier GAN discriminator.""" - def __init__( # noqa: PLR0913 + def __init__( self, latent_dim: int, design_scalars: int, @@ -637,6 +638,7 @@ def prepare_data(problem, batch_size, device, seed=None): "bezier_discriminator.pth": "bezier_discriminator.pth", }, run_config=vars(args), + **checkpoint_identity(args), primary_files=["bezier_generator.pth"], ) diff --git a/engiopt/generators/gan_cnn_2d/gan_cnn_2d.py b/engiopt/generators/gan_cnn_2d/gan_cnn_2d.py index b578ccbc..32c28151 100644 --- a/engiopt/generators/gan_cnn_2d/gan_cnn_2d.py +++ b/engiopt/generators/gan_cnn_2d/gan_cnn_2d.py @@ -16,12 +16,13 @@ from torchvision import transforms import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +import wandb @dataclass @@ -280,8 +281,8 @@ def sample_designs(n_designs: int) -> th.Tensor: # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader) + i if args.track: - batches_done = epoch * len(dataloader) + i wandb.log( { "d_loss": d_loss.item(), @@ -319,35 +320,36 @@ def sample_designs(n_designs: int) -> th.Tensor: # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: - ckpt_gen = { - "epoch": epoch, - "batches_done": batches_done, - "generator": generator.state_dict(), - "optimizer_generator": optimizer_generator.state_dict(), - "loss": g_loss.item(), - } - ckpt_disc = { - "epoch": epoch, - "batches_done": batches_done, - "discriminator": discriminator.state_dict(), - "optimizer_discriminator": optimizer_discriminator.state_dict(), - "loss": d_loss.item(), - } - - th.save(ckpt_gen, "generator.pth") - th.save(ckpt_disc, "discriminator.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, - run_config=vars(args), - primary_files=["generator.pth"], - ) + if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: + ckpt_gen = { + "epoch": epoch, + "batches_done": batches_done, + "generator": generator.state_dict(), + "optimizer_generator": optimizer_generator.state_dict(), + "loss": g_loss.item(), + } + ckpt_disc = { + "epoch": epoch, + "batches_done": batches_done, + "discriminator": discriminator.state_dict(), + "optimizer_discriminator": optimizer_discriminator.state_dict(), + "loss": d_loss.item(), + } + + th.save(ckpt_gen, "generator.pth") + th.save(ckpt_disc, "discriminator.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, + run_config=vars(args), + **checkpoint_identity(args), + primary_files=["generator.pth"], + ) wandb.finish() diff --git a/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py b/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py index 2bd30e7f..6fa3571a 100644 --- a/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py +++ b/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py @@ -23,12 +23,13 @@ from torch.nn.utils.parametrizations import weight_norm import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +import wandb @dataclass @@ -109,7 +110,7 @@ def forward(self, x: th.Tensor) -> th.Tensor: class GatedResnet(nn.Module): - def __init__( # noqa: PLR0913 + def __init__( self, nr_filters: int, conv_op: nn.Module, @@ -299,7 +300,7 @@ def forward(self, x: th.Tensor, output_padding: tuple | None = None) -> th.Tenso class PixelCNNpp(nn.Module): - def __init__( # noqa: PLR0913 + def __init__( self, nr_resnet: int, nr_filters: int, @@ -765,7 +766,7 @@ def sample_from_discretized_mix_logistic(l: th.Tensor, nr_mix: int) -> th.Tensor optimizer = th.optim.Adam(model.parameters(), lr=args.lr, betas=(args.b1, args.b2)) @th.no_grad() - def sample_designs( # noqa: PLR0913 + def sample_designs( model: PixelCNNpp, design_shape: tuple[int, int, int], conditions: th.Tensor, @@ -835,8 +836,8 @@ def sample_designs( # noqa: PLR0913 # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader) + i if args.track: - batches_done = epoch * len(dataloader) + i wandb.log( { "loss": loss.item(), @@ -885,28 +886,29 @@ def sample_designs( # noqa: PLR0913 # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: - # if args.save_model and (((epoch + 1) % args.model_storage_interval == 0) or (epoch == args.n_epochs - 1)) and i == len(dataloader) - 1: - ckpt_model = { - "epoch": epoch, - "batches_done": batches_done, - "model": model.state_dict(), - "optimizer": optimizer.state_dict(), - "loss": loss.item(), - } - - th.save(ckpt_model, "model.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={"model.pth": "model.pth"}, - run_config=vars(args), - primary_files=["model.pth"], - ) + if args.save_model and epoch == args.n_epochs - 1 and i == len(dataloader) - 1: + # if args.save_model and (((epoch + 1) % args.model_storage_interval == 0) or (epoch == args.n_epochs - 1)) and i == len(dataloader) - 1: + ckpt_model = { + "epoch": epoch, + "batches_done": batches_done, + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "loss": loss.item(), + } + + th.save(ckpt_model, "model.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"model.pth": "model.pth"}, + run_config=vars(args), + **checkpoint_identity(args), + primary_files=["model.pth"], + ) wandb.finish() diff --git a/engiopt/generators/vqgan/utils.py b/engiopt/generators/vqgan/utils.py index ba03c7d3..72eaacea 100644 --- a/engiopt/generators/vqgan/utils.py +++ b/engiopt/generators/vqgan/utils.py @@ -37,7 +37,7 @@ class Codebook(nn.Module): init (bool): if true, the codebook has been initialized """ - def __init__( # noqa: PLR0913 + def __init__( self, *, num_codebook_vectors: int, @@ -450,7 +450,7 @@ class GreyscaleLPIPS(nn.Module): ckpt_name: Key in URL_MAP/CKPT_MAP for loading LPIPS heads. """ - def __init__( # noqa: PLR0913 + def __init__( self, *, use_raw: bool = True, diff --git a/engiopt/generators/vqgan/vqgan.py b/engiopt/generators/vqgan/vqgan.py index fe141a68..59166689 100644 --- a/engiopt/generators/vqgan/vqgan.py +++ b/engiopt/generators/vqgan/vqgan.py @@ -30,15 +30,9 @@ from torch.nn import functional as f import tqdm import tyro -import wandb from engiopt.checkpoint_store import save_checkpoint_package -from engiopt.reproducibility import enable_strict_determinism -from engiopt.reproducibility import make_dataloader_generator -from engiopt.reproducibility import seed_training -from engiopt.transforms import drop_constant -from engiopt.transforms import normalize -from engiopt.transforms import resize_to +from engiopt.core import checkpoint_identity from engiopt.generators.vqgan.utils import Codebook from engiopt.generators.vqgan.utils import Discriminator from engiopt.generators.vqgan.utils import DownSampleBlock @@ -51,6 +45,13 @@ from engiopt.generators.vqgan.utils import ResidualBlock from engiopt.generators.vqgan.utils import Swish from engiopt.generators.vqgan.utils import UpSampleBlock +from engiopt.reproducibility import enable_strict_determinism +from engiopt.reproducibility import make_dataloader_generator +from engiopt.reproducibility import seed_training +from engiopt.transforms import drop_constant +from engiopt.transforms import normalize +from engiopt.transforms import resize_to +import wandb @dataclass @@ -158,7 +159,7 @@ class Args: n_epochs_transformer: int = 100 """number of epochs of training""" early_stopping: bool = True - """whether to use early stopping for the transformer; if True requires args.track to be True""" + """whether to use early stopping for the transformer based on held-out validation loss""" early_stopping_patience: int = 3 """number of epochs with no improvement after which training will be stopped""" early_stopping_delta: float = 1e-3 @@ -198,7 +199,7 @@ class Encoder(nn.Module): latent_dim (int): dimensionality of the latent space """ - def __init__( # noqa: PLR0913 + def __init__( self, encoder_channels: tuple[int, ...], encoder_start_resolution: int, @@ -278,7 +279,7 @@ class Decoder(nn.Module): latent_dim (int): dimensionality of the latent space """ - def __init__( # noqa: PLR0913 + def __init__( self, decoder_channels: tuple[int, ...], decoder_start_resolution: int, @@ -371,7 +372,7 @@ class VQGAN(nn.Module): num_codebook_vectors (int): Number of codebook vectors. """ - def __init__( # noqa: PLR0913 + def __init__( self, *, device: th.device, @@ -484,7 +485,7 @@ class VQGANTransformer(nn.Module): bias (bool): If True, use bias terms in the Transformer layers. """ - def __init__( # noqa: PLR0913 + def __init__( self, *, conditional: bool = True, @@ -962,8 +963,8 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader_cvqgan) + i if args.track: - batches_done = epoch * len(dataloader_cvqgan) + i wandb.log({"cvqgan_loss": cvq_loss.item(), "cvqgan_step": batches_done}) wandb.log({"epoch_cvqgan": epoch, "cvqgan_step": batches_done}) print( @@ -973,29 +974,30 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # -------------- # Save model # -------------- - if args.save_model and epoch == args.n_epochs_cvqgan - 1 and i == len(dataloader_cvqgan) - 1: - ckpt_cvq = { - "epoch": epoch, - "batches_done": batches_done, - "cvqgan": cvqgan.state_dict(), - "optimizer_cvqgan": opt_cvq.state_dict(), - "loss": cvq_loss.item(), - } - - th.save(ckpt_cvq, "cvqgan.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={"cvqgan.pth": "cvqgan.pth"}, - run_config=vars(args), - metadata={"stage": "cvqgan"}, - primary_files=["cvqgan.pth"], - ) + if args.save_model and epoch == args.n_epochs_cvqgan - 1 and i == len(dataloader_cvqgan) - 1: + ckpt_cvq = { + "epoch": epoch, + "batches_done": batches_done, + "cvqgan": cvqgan.state_dict(), + "optimizer_cvqgan": opt_cvq.state_dict(), + "loss": cvq_loss.item(), + } + + th.save(ckpt_cvq, "cvqgan.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={"cvqgan.pth": "cvqgan.pth"}, + run_config=vars(args), + **checkpoint_identity(args), + metadata={"stage": "cvqgan"}, + primary_files=["cvqgan.pth"], + ) # Freeze CVQGAN for later use in Stage 2 Transformer for p in cvqgan.parameters(): @@ -1044,8 +1046,8 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader_vqgan) + i if args.track: - batches_done = epoch * len(dataloader_vqgan) + i wandb.log({"vqgan_loss": vq_loss.item(), "vqgan_step": batches_done}) wandb.log({"discriminator_loss": gan_loss.item(), "vqgan_step": batches_done}) wandb.log({"epoch_vqgan": epoch, "vqgan_step": batches_done}) @@ -1081,40 +1083,41 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # -------------- # Save models # -------------- - if args.save_model and epoch == args.n_epochs_vqgan - 1 and i == len(dataloader_vqgan) - 1: - ckpt_vq = { - "epoch": epoch, - "batches_done": batches_done, - "vqgan": vqgan.state_dict(), - "optimizer_vqgan": opt_vq.state_dict(), - "loss": vq_loss.item(), - } - ckpt_disc = { - "epoch": epoch, - "batches_done": batches_done, - "discriminator": discriminator.state_dict(), - "optimizer_discriminator": opt_disc.state_dict(), - "loss": gan_loss.item(), - } - - th.save(ckpt_vq, "vqgan.pth") - th.save(ckpt_disc, "discriminator.pth") - save_checkpoint_package( - checkpoint_backend="hf", - hf_entity=args.hf_entity, - hf_repo_prefix=args.hf_repo_prefix, - hf_private=False, - problem_id=args.problem_id, - algo=args.algo, - seed=args.seed, - checkpoint_files={ - "vqgan.pth": "vqgan.pth", - "discriminator.pth": "discriminator.pth", - }, - run_config=vars(args), - metadata={"stage": "vqgan"}, - primary_files=["vqgan.pth"], - ) + if args.save_model and epoch == args.n_epochs_vqgan - 1 and i == len(dataloader_vqgan) - 1: + ckpt_vq = { + "epoch": epoch, + "batches_done": batches_done, + "vqgan": vqgan.state_dict(), + "optimizer_vqgan": opt_vq.state_dict(), + "loss": vq_loss.item(), + } + ckpt_disc = { + "epoch": epoch, + "batches_done": batches_done, + "discriminator": discriminator.state_dict(), + "optimizer_discriminator": opt_disc.state_dict(), + "loss": gan_loss.item(), + } + + th.save(ckpt_vq, "vqgan.pth") + th.save(ckpt_disc, "discriminator.pth") + save_checkpoint_package( + checkpoint_backend="hf", + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + hf_private=False, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + checkpoint_files={ + "vqgan.pth": "vqgan.pth", + "discriminator.pth": "discriminator.pth", + }, + run_config=vars(args), + **checkpoint_identity(args), + metadata={"stage": "vqgan"}, + primary_files=["vqgan.pth"], + ) # Freeze VQGAN for later use in Stage 2 Transformer for p in vqgan.parameters(): @@ -1148,8 +1151,8 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # ---------- # Logging # ---------- + batches_done = epoch * len(dataloader_transformer) + i if args.track: - batches_done = epoch * len(dataloader_transformer) + i wandb.log({"transformer_loss": loss.item(), "transformer_step": batches_done}) wandb.log({"epoch_transformer": epoch, "transformer_step": batches_done}) print( @@ -1185,8 +1188,10 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: plt.close() wandb.log({"designs_transformer": wandb.Image(img_fname)}) - # Early stopping based on held-out validation loss - if args.track and args.early_stopping: + # Early stopping based on held-out validation loss. This is training + # logic -- it decides when to stop and which checkpoint to keep -- so it + # runs regardless of W&B; only the metric *logging* below is track-gated. + if args.early_stopping: transformer.eval() val_losses = [] with th.no_grad(): @@ -1197,7 +1202,8 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: val_loss = f.cross_entropy(val_logits.reshape(-1, val_logits.size(-1)), val_targets.reshape(-1)) val_losses.append(val_loss.item()) val_loss = sum(val_losses) / len(val_losses) - wandb.log({"transformer_val_loss": val_loss, "transformer_step": batches_done}) + if args.track: + wandb.log({"transformer_val_loss": val_loss, "transformer_step": batches_done}) if val_loss < best_val - args.early_stopping_delta: best_val = val_loss @@ -1224,7 +1230,7 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: # -------------- # Save model # -------------- - if args.track and args.save_model: + if args.save_model: if not args.early_stopping: ckpt_tr = { "epoch": epoch, @@ -1253,6 +1259,7 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: seed=args.seed, checkpoint_files=checkpoint_files, run_config=vars(args), + **checkpoint_identity(args), metadata={"stage": "transformer"}, primary_files=["transformer.pth"], ) diff --git a/engiopt/metrics.py b/engiopt/metrics.py index a2e121eb..de8b56d1 100644 --- a/engiopt/metrics.py +++ b/engiopt/metrics.py @@ -1,32 +1,19 @@ -"""This module provides metrics for evaluating generative model designs. +"""Numerical primitives used by the evaluation metrics. -Maximum Mean Discrepancy (MMD), Determinantal Point Process (DPP) diversity, and -optimality gap calculations. +Kernel-based distribution and diversity measures, plus the optimality-gap +definition. The evaluator in `engiopt.evaluation` composes these; see +`engiopt/evaluation/metrics/builtin.py` for the registered metrics themselves. """ from __future__ import annotations -import multiprocessing -import os -import sys -import traceback -from typing import Any, TYPE_CHECKING +from typing import TYPE_CHECKING -from gymnasium import spaces import numpy as np -import numpy.typing as npt from scipy.spatial.distance import cdist if TYPE_CHECKING: - from datasets import Dataset from engibench import OptiStep - from engibench.core import Problem - - -if sys.platform != "win32": # only set fork on non-Windows - multiprocessing.set_start_method("fork", force=True) -else: - multiprocessing.set_start_method("spawn", force=True) def mmd(x: np.ndarray, y: np.ndarray, sigma: float = 1.0) -> float: @@ -86,157 +73,4 @@ def optimality_gap(opt_history: list[OptiStep], baseline: float) -> list[float]: return [opt.obj_values - baseline for opt in opt_history] -def simulate_failure_ratio( # noqa: C901 - problem: Problem, - gen_designs: npt.NDArray, - sampled_conditions: Dataset | None = None, -) -> float: - """Compute the failure ratio of generated designs. This is designed for the airfoil problem. - - Args: - problem: The optimization problem to evaluate. - gen_designs (np.ndarray): Array of shape (n_samples, l, w) for generative model designs. - sampled_conditions (Dataset): Dataset of sampled conditions for optimization. If None, no conditions are used. - - Returns: - float: The failure ratio of generated designs. - """ - failure_count = 0 - for idx, design in enumerate(gen_designs): - if isinstance(problem.design_space, spaces.Dict): - # Need to unflatten the design to be used for optimization or simulation - unflattened_design = spaces.unflatten(problem.design_space, design) - unflattened_design["angle_of_attack"] = unflattened_design["angle_of_attack"][0] - else: - unflattened_design = design - - def worker(idx, config, return_queue): - try: - objs = problem.simulate(unflattened_design, config=config, mpicores=10) # noqa: B023 - if np.isnan(objs[0]) or np.isnan(objs[1]): - print(f"Simulation returned NaN values for design {idx}") - raise Exception("Simulation returned NaN values") # noqa: TRY002, TRY301 - return_queue.put(("ok", objs)) - except Exception: # noqa: BLE001 - return_queue.put(("error", traceback.format_exc())) - - # Attempt to simulate the design - def run_with_timeout(idx, timeout=30): - config = sampled_conditions[idx] if sampled_conditions is not None else None - q = multiprocessing.Queue() - p = multiprocessing.Process(target=worker, args=(idx, config, q)) - p.start() - p.join(timeout) - - if p.is_alive(): - p.terminate() # force-kill the child process - p.join() - os.system("docker stop machaero") # noqa: S605 - raise RuntimeError(f"Simulation timed out for design {idx}") - - if not q.empty(): - status, payload = q.get() - if status == "ok": - print(f"Simulation successful for design {idx}: {payload}") - else: - raise RuntimeError(f"Simulation error for design {idx}:\n{payload}") - else: - raise RuntimeError("Simulation process ended without reporting back") - - try: - run_with_timeout(idx, timeout=30) - except RuntimeError as e: - failure_count += 1 - print(e) - - return failure_count / len(gen_designs) # Return the failure ratio - - -def metrics( - problem: Problem, - gen_designs: npt.NDArray, - dataset_designs: npt.NDArray, - sampled_conditions: Dataset | None = None, - sigma: float = 1.0, -) -> dict[str, Any]: - """Compute various metrics for evaluating generative model designs. - - Args: - problem: The optimization problem to evaluate. - gen_designs (np.ndarray): Array of shape (n_samples, l, w) for generative model designs (potentially flattened for dict spaces). - dataset_designs (np.ndarray): Array of shape (n_samples, l, w) for dataset designs (these are not flattened). - sampled_conditions (Dataset): Dataset of sampled conditions for optimization. If None, no conditions are used. - sigma (float): Bandwidth parameter for the Gaussian kernel (in mmd and dpp calculation). - - Returns: - dict[str, Any]: A dictionary containing the computed metrics: - - "iog": Average Initial Optimality Gap (float). - - "cog": Average Cumulative Optimality Gap (float). - - "fog": Average Final Optimality Gap (float). - - "mmd": Maximum Mean Discrepancy (float). - - "dpp": Determinantal Point Process diversity (float). - - "viol": Average violation ratio (float). - """ - n_samples = len(gen_designs) - - cog_list = [] - iog_list = [] - fog_list = [] - viol_list = [] - for i in range(n_samples): - conditions = sampled_conditions[i] if sampled_conditions is not None else None - if isinstance(problem.design_space, spaces.Dict): - # Need to unflatten the design to be used for optimization or simulation - 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)) - fog_list.append(opt_history_gaps[-1]) - - # Check if conditions dict has 'volfrac' or 'volume' key and compare with design mean - if conditions: - tol = 0.01 # Tolerance for equality constraint deviation - target_vol = conditions.get("volfrac") or conditions.get("volume") - if target_vol is not None: - viol = np.abs(np.mean(unflattened_design) - target_vol) >= tol - viol_list.append(viol) - - # Compute the average Initial Optimality Gap (IOG), Cumulative Optimality Gap (COG), and Final Optimality Gap (FOG) - average_iog: float = float(np.mean(iog_list)) # Average of initial optimality gaps - average_cog: float = float(np.mean(cog_list)) # Average of cumulative optimality gaps - average_fog: float = float(np.mean(fog_list)) # Average of final optimality gaps - average_viol: float = float(np.mean(viol_list)) # Average of violation ratios - - # Compute the Maximum Mean Discrepancy (MMD) between generated and dataset designs - # We compute the MMD on the flattened designs - flattened_ds_designs: list[npt.NDArray] = [] - for design in dataset_designs: - if isinstance(problem.design_space, spaces.Dict): - flattened = spaces.flatten(problem.design_space, design) - flattened_ds_designs.append(np.array(flattened)) - else: - flattened_ds_designs.append(design) - flattened_ds_designs_array: npt.NDArray = np.array(flattened_ds_designs) - mmd_value: float = mmd(gen_designs, flattened_ds_designs_array, sigma=sigma) - - # Compute the Determinantal Point Process (DPP) diversity for generated designs - # We compute the DPP on the flattened designs - dpp_value: float = dpp_diversity(gen_designs, sigma=sigma) - - # Return all computed metrics as a dictionary - return { - "iog": average_iog, - "cog": average_cog, - "fog": average_fog, - "mmd": mmd_value, - "dpp": dpp_value, - "viol": average_viol, - } +# Return the failure ratio diff --git a/engiopt/transforms.py b/engiopt/transforms.py index c26d02f3..7f6de89d 100644 --- a/engiopt/transforms.py +++ b/engiopt/transforms.py @@ -5,10 +5,53 @@ from datasets import Dataset from engibench.core import Problem from gymnasium import spaces +import numpy as np import torch as th import torch.nn.functional as f +def get_scalar_condition_keys(problem: Problem, dataset: Dataset, *, drop_constants: bool = False) -> list[str]: + """Return the condition keys usable as a dense model input. + + `problem.conditions_keys` is the full contract, which is broader than what a + generator can consume as a `(n, n_conds)` tensor. Two cases are excluded: + + 1. Keys absent from the dataset. Some are solver settings rather than + per-sample conditions -- photonics2d declares `num_elems_x`, + `num_elems_y`, and `num_optimization_steps`, none of which vary per row. + 2. Array-valued keys. thermoelastic2d encodes boundary conditions as 65x65 + matrices, which cannot be stacked alongside scalars. + + Args: + problem: An EngiBench problem instance. + dataset: A dataset split, e.g. `problem.dataset["test"]`. + drop_constants: Also drop columns with zero standard deviation. + + Returns: + Condition names, in `conditions_keys` order. + """ + scalar_keys = [ + key for key in problem.conditions_keys if key in dataset.column_names and np.asarray(dataset[0][key]).ndim == 0 + ] + + if drop_constants and scalar_keys: + conds = th.stack([th.as_tensor(dataset[c][:]).float() for c in scalar_keys], dim=1) + std = conds.std(dim=0) + scalar_keys = [c for i, c in enumerate(scalar_keys) if std[i] > 0] + + return scalar_keys + + +def get_image_condition_keys(problem: Problem, dataset: Dataset) -> list[str]: + """Return the array-valued condition keys present in the dataset. + + The complement of `get_scalar_condition_keys`, e.g. thermoelastic2d's 65x65 + boundary matrices. These still reach the simulator through the conditions + dataset; they simply cannot travel in the dense condition tensor. + """ + return [key for key in problem.conditions_keys if key in dataset.column_names and np.asarray(dataset[0][key]).ndim > 0] + + def flatten_dict_factory(problem: Problem, device: th.device) -> Callable: """Factory function to create a flatten_dict function.""" diff --git a/pyproject.toml b/pyproject.toml index 96d89803..c0733cf7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -172,12 +172,17 @@ order-by-type = false "D104", # D104 - Package dostrings not required in test files. "ARG", # ARG - Unused args are common in tests with mock patches and mock functions. "D", # Pydocstyle + "PLR2004", # Magic values are the point of an assertion in a test. + "SLF001", # Tests may exercise private helpers directly. ] "*/__init__.py" = [ "F401", # Module imported but unused "D104", # D104 - Package dostrings not required in init files. "E402" ] +"engiopt/generators/_template/*.py" = [ + "ERA001", # The commented-out example code is the point of the template. +] "docs/**" = [ "D", # Pydocstyle ] diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..39000ccf --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,105 @@ +"""Shared fixtures. + +The evaluation tests deliberately avoid EngiBench problems: they need a design +space and some conditions, not a physics solver, and pulling a real dataset +would make the suite slow and network-dependent. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from gymnasium import spaces +import numpy as np +import pytest + + +@dataclass +class FakeOptiStep: + """Stand-in for `engibench.OptiStep`, which carries one objective value.""" + + obj_values: Any + + +class FakeProblem: + """A minimal stand-in exposing the parts of `Problem` the evaluator touches. + + The fake objective is `mean(design)`, so the optimum is known analytically + and the optimality gaps can be asserted exactly. + """ + + def __init__( + self, + shape: tuple[int, ...] = (8, 10), + conditions_keys: tuple[str, ...] = ("volfrac", "rmin"), + *, + n_objectives: int = 1, + directions: tuple[str, ...] | None = None, + ): + self.design_space = spaces.Box(low=0.0, high=1.0, shape=shape, dtype=np.float64) + self.conditions_keys = list(conditions_keys) + self.n_objectives = n_objectives + # Mirrors `Problem.objectives`: (name, direction) pairs in simulate order. + directions = directions or ("MINIMIZE",) * n_objectives + self.objectives = tuple((f"objective_{i}", directions[i]) for i in range(n_objectives)) + self.reset_calls = 0 + self.optimize_calls = 0 + + def reset(self, seed: int | None = None) -> None: + """Record that the evaluator reset us between solver calls.""" + self.reset_calls += 1 + + def random_design(self) -> tuple[Any, dict[str, Any]]: + """Return an arbitrary in-space design.""" + return self.design_space.sample(), {} + + def simulate(self, design: Any, config: Any = None) -> np.ndarray: + """Objective is the design's mean, repeated per objective.""" + return np.full(self.n_objectives, float(np.mean(np.asarray(design)))) + + def optimize(self, design: Any, config: Any = None) -> tuple[Any, list[FakeOptiStep]]: + """Walk the objective down toward zero in three steps.""" + self.optimize_calls += 1 + start = float(np.mean(np.asarray(design))) + history = [FakeOptiStep(np.full(self.n_objectives, start * factor)) for factor in (1.0, 0.5, 0.25)] + return design, history + + +@pytest.fixture +def fake_problem() -> FakeProblem: + """A small 2D problem with two scalar conditions.""" + return FakeProblem() + + +class FakeDataset: + """A tiny stand-in for a HuggingFace dataset split.""" + + def __init__(self, columns: dict[str, list[Any]]): + self._columns = columns + self.column_names = list(columns) + + def __getitem__(self, item: Any) -> Any: + if isinstance(item, str): + return self._columns[item] + return {name: values[item] for name, values in self._columns.items()} + + def __len__(self) -> int: + return len(next(iter(self._columns.values()))) + + +@pytest.fixture +def mixed_condition_dataset() -> FakeDataset: + """A dataset mixing scalar, array-valued, and constant conditions. + + Mirrors the two problems that previously broke condition sampling: + photonics2d (solver-only keys absent from the dataset) and thermoelastic2d + (array-valued boundary conditions). + """ + return FakeDataset( + { + "volfrac": [0.3, 0.4, 0.5], + "rmin": [2.0, 2.0, 2.0], # constant + "fixed_elements": [np.zeros((4, 4)) for _ in range(3)], # array-valued + } + ) diff --git a/tests/test_checkpoint_identity.py b/tests/test_checkpoint_identity.py new file mode 100644 index 00000000..3c0f3ce6 --- /dev/null +++ b/tests/test_checkpoint_identity.py @@ -0,0 +1,76 @@ +"""Tests for how checkpoints are filed by hyperparameter configuration. + +Two requirements pull against each other and both must hold: + +1. Many hyperparameter settings of one algorithm must coexist on the Hub and be + evaluated separately. +2. Naming a model with no extra qualification must mean its *default* + hyperparameters -- a sweep must never redefine that. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from engiopt.checkpoint_store import build_hf_package_path +from engiopt.core import checkpoint_identity +from engiopt.core import config_fingerprint +from engiopt.core import config_path_parts + + +@dataclass +class _Args: + """Stand-in for a training script's Args.""" + + latent_dim: int = 32 + lr_gen: float = 1e-4 + seed: int = 1 + wandb_entity: str | None = None + track: bool = True + + +def test_default_run_is_marked_as_the_standard() -> None: + """An untouched Args is the default configuration, so it owns the plain path.""" + identity = checkpoint_identity(_Args()) + assert identity["is_default_config"] is True + + +def test_tuned_run_is_not_the_standard() -> None: + """A swept configuration must not claim the canonical path.""" + identity = checkpoint_identity(_Args(latent_dim=128)) + assert identity["is_default_config"] is False + + +def test_infrastructure_settings_do_not_change_the_configuration() -> None: + """Seeds and tracking are how a run was operated, not what the model is.""" + assert checkpoint_identity(_Args(seed=7, wandb_entity="someone", track=False))["is_default_config"] is True + + +def test_each_hyperparameter_setting_gets_its_own_path() -> None: + """The case that previously overwrote itself: same algo, same seed, different HP.""" + small = config_fingerprint(vars(_Args(latent_dim=32))) + large = config_fingerprint(vars(_Args(latent_dim=128))) + assert small != large + + small_path = build_hf_package_path("beams2d", 1, config_path_parts(small)) + large_path = build_hf_package_path("beams2d", 1, config_path_parts(large)) + assert small_path != large_path + assert small_path == f"beams2d/cfg_{small}/seed_1" + + +def test_seeds_of_one_configuration_stay_distinct() -> None: + """Seeds are separate checkpoints, not overwrites of each other.""" + fingerprint = config_fingerprint(vars(_Args())) + parts = config_path_parts(fingerprint) + assert build_hf_package_path("beams2d", 1, parts) != build_hf_package_path("beams2d", 2, parts) + + +def test_bare_model_name_resolves_to_the_canonical_path() -> None: + """No fingerprint means the default checkpoint, which is what a bare name means.""" + assert config_path_parts(None) is None + assert build_hf_package_path("beams2d", 1, config_path_parts(None)) == "beams2d/seed_1" + + +def test_identical_hyperparameters_fingerprint_identically() -> None: + """Re-running the same configuration must land on the same path, not a new one.""" + assert config_fingerprint(vars(_Args(latent_dim=64))) == config_fingerprint(vars(_Args(latent_dim=64))) diff --git a/tests/test_checkpoint_store_layout.py b/tests/test_checkpoint_store_layout.py new file mode 100644 index 00000000..c63ed745 --- /dev/null +++ b/tests/test_checkpoint_store_layout.py @@ -0,0 +1,56 @@ +"""Tests for where checkpoints, metrics, and media live. + +The split is deliberate: HuggingFace hosts anything that must be reloaded or +compared (weights, run configs, metrics, the leaderboard); Weights & Biases, +when enabled, hosts what you only look at (loss curves, sample images). +""" + +from __future__ import annotations + +import inspect + +from engiopt import checkpoint_store +from engiopt.core import Generator + + +def test_wandb_is_not_a_checkpoint_backend() -> None: + """HuggingFace and local directories are the only sources of weights.""" + assert set(checkpoint_store.ModelSource.__args__) == {"auto", "hf", "local"} + + +def test_wandb_checkpoint_resolvers_are_gone() -> None: + """The legacy W&B read-fallback was dead weight once HF became the only host.""" + for removed in ("_resolve_wandb_package", "_resolve_wandb_reference", "_build_wandb_artifact_path"): + assert not hasattr(checkpoint_store, removed), f"{removed} should have been removed" + + +def test_write_only_run_scoped_path_is_gone() -> None: + """That path was never readable by any loader, so nothing could use it.""" + assert not hasattr(checkpoint_store, "build_hf_run_package_path") + + +def test_generators_no_longer_carry_wandb_naming() -> None: + """The contract should not mention a backend it cannot read from.""" + assert not hasattr(Generator, "wandb_artifact_names") + assert not hasattr(Generator, "_wandb_artifact_names") + + +def test_resolve_named_checkpoint_takes_no_wandb_arguments() -> None: + """Its signature is the clearest statement of what backends exist.""" + params = set(inspect.signature(checkpoint_store.resolve_named_checkpoint).parameters) + assert not {p for p in params if "wandb" in p} + assert {"hf_entity", "hf_repo_prefix", "extra_path_parts"} <= params + + +def test_metrics_are_publishable_next_to_the_weights() -> None: + """A checkpoint should be self-describing without the leaderboard.""" + assert checkpoint_store.METRICS_FILE == "metrics.json" + params = set(inspect.signature(checkpoint_store.publish_checkpoint_metrics).parameters) + assert {"metrics", "problem_id", "algo", "seed", "extra_path_parts"} <= params + + +def test_wandb_still_receives_a_pointer_to_the_hf_package() -> None: + """Each side records where the other is, so curves and weights stay linked.""" + source = inspect.getsource(checkpoint_store._log_checkpoint_summary_to_wandb) + assert "hf_repo_id" in source + assert "hf_config_package_path" in source diff --git a/tests/test_conditions_and_gaps.py b/tests/test_conditions_and_gaps.py new file mode 100644 index 00000000..312333d5 --- /dev/null +++ b/tests/test_conditions_and_gaps.py @@ -0,0 +1,270 @@ +"""Tests for condition selection and the optimality-gap definitions. + +Both were real defects: condition sampling could not handle two of the four +benchmark problems, and `iog` reported a raw objective rather than a gap. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest + +from engiopt.evaluation.context import EvaluationContext +from engiopt.evaluation.context import MultiObjectiveScalarizationError +from engiopt.transforms import get_image_condition_keys +from engiopt.transforms import get_scalar_condition_keys + + +class _Problem: + """Just enough of `Problem` for the condition-key helpers.""" + + def __init__(self, conditions_keys: list[str]): + self.conditions_keys = conditions_keys + + +# ---------------------------------------------------------------------- +# Condition selection +# ---------------------------------------------------------------------- + + +def test_scalar_keys_drop_conditions_absent_from_the_dataset(mixed_condition_dataset: Any) -> None: + """photonics2d declares solver settings that are not dataset columns.""" + problem = _Problem(["volfrac", "num_elems_x", "num_optimization_steps"]) + assert get_scalar_condition_keys(problem, mixed_condition_dataset) == ["volfrac"] + + +def test_scalar_keys_drop_array_valued_conditions(mixed_condition_dataset: Any) -> None: + """thermoelastic2d's boundary conditions are matrices; they cannot be stacked.""" + problem = _Problem(["volfrac", "fixed_elements"]) + assert get_scalar_condition_keys(problem, mixed_condition_dataset) == ["volfrac"] + + +def test_scalar_keys_preserve_conditions_key_order(mixed_condition_dataset: Any) -> None: + """Column order defines the tensor layout, so it must follow the contract.""" + problem = _Problem(["rmin", "volfrac"]) + assert get_scalar_condition_keys(problem, mixed_condition_dataset) == ["rmin", "volfrac"] + + +def test_scalar_keys_can_drop_constant_columns(mixed_condition_dataset: Any) -> None: + """A constant condition carries no signal for a conditional model.""" + problem = _Problem(["volfrac", "rmin"]) + assert get_scalar_condition_keys(problem, mixed_condition_dataset, drop_constants=True) == ["volfrac"] + + +def test_image_keys_are_the_complement(mixed_condition_dataset: Any) -> None: + """Array conditions still reach the simulator, so they must be identifiable.""" + problem = _Problem(["volfrac", "rmin", "fixed_elements"]) + scalars = get_scalar_condition_keys(problem, mixed_condition_dataset) + images = get_image_condition_keys(problem, mixed_condition_dataset) + assert images == ["fixed_elements"] + assert not set(scalars) & set(images) + + +# ---------------------------------------------------------------------- +# Optimality gaps +# ---------------------------------------------------------------------- + + +def _context(problem: Any, gen: np.ndarray) -> EvaluationContext: + return EvaluationContext(problem=problem, problem_id="fake", gen_designs=gen, ref_designs=gen.copy()) + + +def test_iog_is_zero_when_the_generated_design_is_the_reference(fake_problem: Any) -> None: + """The defining property: no gap between a design and itself. + + Before the fix `iog` stored the raw simulated objective, so this returned the + design's objective (~0.5 here) instead of 0. + """ + designs = np.full((3, *fake_problem.design_space.shape), 0.5) + assert np.allclose(_context(fake_problem, designs).optimization.iog, 0.0) + + +def test_iog_grows_as_the_design_gets_worse(fake_problem: Any) -> None: + """A worse design must show a larger initial gap.""" + shape = fake_problem.design_space.shape + reference = np.full((2, *shape), 0.2) + near = EvaluationContext( + problem=fake_problem, problem_id="fake", gen_designs=np.full((2, *shape), 0.3), ref_designs=reference + ) + far = EvaluationContext( + problem=fake_problem, problem_id="fake", gen_designs=np.full((2, *shape), 0.9), ref_designs=reference + ) + assert np.mean(far.optimization.iog) > np.mean(near.optimization.iog) > 0 + + +def test_all_three_gaps_share_one_baseline(fake_problem: Any) -> None: + """iog, cog, and fog must all be measured against the reference optimum. + + With the fake solver the objective decays 1.0 -> 0.5 -> 0.25 of its start, so + for a design at 0.8 against a reference at 0.4 the gaps are exactly known. + """ + shape = fake_problem.design_space.shape + ctx = EvaluationContext( + problem=fake_problem, + problem_id="fake", + gen_designs=np.full((1, *shape), 0.8), + ref_designs=np.full((1, *shape), 0.4), + ) + results = ctx.optimization + assert results.iog[0] == pytest.approx(0.8 - 0.4) + assert results.cog[0] == pytest.approx((0.8 - 0.4) + (0.4 - 0.4) + (0.2 - 0.4)) + assert results.fog[0] == pytest.approx(0.2 - 0.4) + + +def test_optimization_is_computed_once_and_cached(fake_problem: Any) -> None: + """Every performance metric shares one solver pass; that is the point of the context.""" + designs = np.full((2, *fake_problem.design_space.shape), 0.5) + ctx = _context(fake_problem, designs) + ctx.optimization + ctx.optimization + assert fake_problem.optimize_calls == 2 # one per sample, not per access + + +# ---------------------------------------------------------------------- +# Multi-objective scalarization +# ---------------------------------------------------------------------- + + +def _multi_context(**kwargs: Any) -> EvaluationContext: + """A 3-objective context whose objective vector is `mean(design)` repeated.""" + from tests.conftest import FakeDataset + from tests.conftest import FakeProblem + + problem = FakeProblem(n_objectives=3, conditions_keys=("weight",)) + designs = np.full((1, *problem.design_space.shape), 0.8) + return EvaluationContext( + problem=problem, + problem_id="thermo_like", + gen_designs=designs, + ref_designs=np.full_like(designs, 0.5), + conditions=FakeDataset({"weight": [kwargs.pop("weight", 0.5)]}), + **kwargs, + ) + + +def test_multi_objective_without_a_declared_weighting_is_refused() -> None: + """Silently averaging objectives with different units produces a meaningless number.""" + with pytest.raises(MultiObjectiveScalarizationError, match="does not say how to combine"): + _multi_context().optimization + + +def test_a_purely_structural_sample_ignores_the_thermal_objective() -> None: + """With `weight = 1` the second objective must not influence the gap at all. + + This is the property that makes the weighted scalarization correct and both + summing and averaging wrong: at w=1 a change in thermal compliance has to + leave every gap untouched. + """ + ctx = _multi_context(weight=1.0, objective_weight_condition="weight") + weights = ctx.weights_at(0) + assert weights is not None + assert weights[0] == pytest.approx(1.0) + assert weights[1] == pytest.approx(0.0) + + baseline = ctx.scalarize_gap(np.array([0.3, 7.0, 99.0]), 0) + perturbed = ctx.scalarize_gap(np.array([0.3, 1000.0, -5.0]), 0) + assert baseline == pytest.approx(perturbed) + assert baseline == pytest.approx(0.3) + + +def test_a_purely_thermal_sample_ignores_the_structural_objective() -> None: + """The mirror case: at `weight = 0` only the second objective counts.""" + ctx = _multi_context(weight=0.0, objective_weight_condition="weight") + assert ctx.scalarize_gap(np.array([500.0, 0.7, 3.0]), 0) == pytest.approx(0.7) + + +def test_weighted_gaps_use_the_per_sample_trade_off(fake_problem: Any) -> None: + """Each gap is scalarized with that sample's own weight, not a global one.""" + ctx = _multi_context(weight=0.25, objective_weight_condition="weight") + results = ctx.optimization + # Objective vector is mean(design) repeated, so every component's gap is 0.3. + # Weights (0.25, 0.75, 0.0) sum to 1, leaving the gap at 0.3. + assert results.iog[0] == pytest.approx(0.8 - 0.5) + + +def test_fixed_objective_weights_are_honoured() -> None: + """A spec may declare constant weights instead of a per-sample condition.""" + ctx = _multi_context(objective_weights=(0.5, 0.5, 0.0)) + assert ctx.scalarize_gap(np.array([2.0, 4.0, 1000.0]), 0) == pytest.approx(3.0) + + +def test_mismatched_weight_count_is_rejected() -> None: + """A weight vector that does not line up with the objectives is a spec bug.""" + ctx = _multi_context(objective_weights=(1.0,)) + with pytest.raises(ValueError, match="3 objectives but the spec declares 1 weights"): + ctx.weights_at(0) + + +def test_single_objective_problems_need_no_weighting(fake_problem: Any) -> None: + """The common case stays configuration-free.""" + ctx = EvaluationContext( + problem=fake_problem, + problem_id="fake", + gen_designs=np.full((1, *fake_problem.design_space.shape), 0.5), + ref_designs=np.full((1, *fake_problem.design_space.shape), 0.5), + ) + assert ctx.weights_at(0) is None + assert ctx.scalarize_gap(np.array([2.5]), 0) == pytest.approx(2.5) + + +# ---------------------------------------------------------------------- +# Objective direction +# ---------------------------------------------------------------------- + + +def _directed_context(direction: str, gen: float, ref: float) -> EvaluationContext: + from tests.conftest import FakeProblem + + problem = FakeProblem(n_objectives=1, directions=(direction,)) + return EvaluationContext( + problem=problem, + problem_id="fake", + gen_designs=np.full((1, *problem.design_space.shape), gen), + ref_designs=np.full((1, *problem.design_space.shape), ref), + ) + + +def test_minimized_objective_keeps_its_sign() -> None: + """Worse (higher) than the reference must give a positive gap.""" + assert _directed_context("MINIMIZE", gen=0.8, ref=0.5).optimization.iog[0] == pytest.approx(0.3) + + +def test_maximized_objective_flips_sign() -> None: + """photonics2d maximizes total_overlap, so a *higher* value is better. + + Without the flip its gap would be +0.3 for a design that beats the + reference, and every lower-is-better metric would rank it last. + """ + assert _directed_context("MAXIMIZE", gen=0.8, ref=0.5).optimization.iog[0] == pytest.approx(-0.3) + + +def test_worse_designs_score_higher_under_both_directions() -> None: + """The invariant the leaderboard depends on: bigger gap == worse design.""" + minimize_worse = _directed_context("MINIMIZE", gen=0.9, ref=0.5).optimization.iog[0] + minimize_better = _directed_context("MINIMIZE", gen=0.6, ref=0.5).optimization.iog[0] + maximize_worse = _directed_context("MAXIMIZE", gen=0.1, ref=0.5).optimization.iog[0] + maximize_better = _directed_context("MAXIMIZE", gen=0.4, ref=0.5).optimization.iog[0] + assert minimize_worse > minimize_better + assert maximize_worse > maximize_better + + +def test_signs_are_per_objective_for_mixed_directions() -> None: + """A problem may mix directions; each objective flips independently.""" + from tests.conftest import FakeDataset + from tests.conftest import FakeProblem + + problem = FakeProblem(n_objectives=2, conditions_keys=("weight",), directions=("MINIMIZE", "MAXIMIZE")) + designs = np.zeros((1, *problem.design_space.shape)) + ctx = EvaluationContext( + problem=problem, + problem_id="mixed", + gen_designs=designs, + ref_designs=designs, + conditions=FakeDataset({"weight": [0.5]}), + objective_weight_condition="weight", + ) + assert list(ctx.objective_signs) == [1.0, -1.0] + # Both objectives beat the reference by 1.0, so both should read as -0.5. + assert ctx.scalarize_gap(np.array([-1.0, 1.0]), 0) == pytest.approx(-1.0) diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py new file mode 100644 index 00000000..90abe46f --- /dev/null +++ b/tests/test_evaluation.py @@ -0,0 +1,329 @@ +"""Tests for the metric registry, evaluation context, and leaderboard. + +Metric correctness is checked by properties that must hold for any correct +implementation (MMD of a set against itself is zero; noise is maximally +diverse), rather than by pinning historical numbers. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pandas as pd +import pytest + +from engiopt.evaluation.context import EvaluationContext +from engiopt.evaluation.evaluator import _as_columns +from engiopt.evaluation.evaluator import order_columns +from engiopt.evaluation.leaderboard import append_rows +from engiopt.evaluation.leaderboard import disagreement +from engiopt.evaluation.leaderboard import rank +from engiopt.evaluation.leaderboard import ROW_KEY +from engiopt.evaluation.registry import MetricRegistry +from engiopt.evaluation.registry import METRICS +from engiopt.evaluation.registry import MetricSpec +from engiopt.evaluation.registry import register_metric +from engiopt.evaluation.spec import EvalSpec + +# Importing the metrics package registers the built-ins. +import engiopt.evaluation.metrics # noqa: F401 # isort: skip + + +def _context(problem: Any, gen: np.ndarray, ref: np.ndarray, **kwargs: Any) -> EvaluationContext: + return EvaluationContext(problem=problem, problem_id="fake", gen_designs=gen, ref_designs=ref, **kwargs) + + +# ---------------------------------------------------------------------- +# Metric properties +# ---------------------------------------------------------------------- + + +def test_mmd_of_identical_sets_is_zero(fake_problem: Any) -> None: + """A generator that reproduced the reference set exactly should score 0.""" + rng = np.random.default_rng(0) + designs = rng.random((12, *fake_problem.design_space.shape)) + ctx = _context(fake_problem, designs, designs.copy()) + assert METRICS["mmd"].fn(ctx) == pytest.approx(0.0, abs=1e-9) + + +def test_mmd_is_invariant_to_sample_order(fake_problem: Any) -> None: + """MMD compares distributions, so shuffling the samples must not change it.""" + rng = np.random.default_rng(1) + gen = rng.random((10, *fake_problem.design_space.shape)) + ref = rng.random((10, *fake_problem.design_space.shape)) + baseline = METRICS["mmd"].fn(_context(fake_problem, gen, ref)) + shuffled = METRICS["mmd"].fn(_context(fake_problem, rng.permutation(gen), ref)) + assert baseline == pytest.approx(shuffled, rel=1e-9) + + +def test_mmd_grows_as_distributions_separate(fake_problem: Any) -> None: + """Shifting the generated set away from the reference must raise MMD.""" + rng = np.random.default_rng(2) + ref = rng.random((15, *fake_problem.design_space.shape)) + near = METRICS["mmd"].fn(_context(fake_problem, ref + 0.01, ref, sigma=1.0)) + far = METRICS["mmd"].fn(_context(fake_problem, ref + 0.5, ref, sigma=1.0)) + assert far > near + + +def test_dpp_prefers_varied_designs_over_duplicates(fake_problem: Any) -> None: + """A collapsed generator must score below a varied one.""" + rng = np.random.default_rng(3) + varied = rng.random((8, *fake_problem.design_space.shape)) + collapsed = np.repeat(varied[:1], 8, axis=0) + ref = rng.random((8, *fake_problem.design_space.shape)) + assert METRICS["dpp"].fn(_context(fake_problem, varied, ref)) > METRICS["dpp"].fn( + _context(fake_problem, collapsed, ref) + ) + + +# ---------------------------------------------------------------------- +# Cheap / expensive separation +# ---------------------------------------------------------------------- + + +def test_builtin_metrics_declare_their_cost() -> None: + """Simulator-backed metrics must be marked expensive, and the rest cheap.""" + assert {spec.name for spec in METRICS.select(cost="cheap")} == {"mmd", "dpp"} + assert {spec.name for spec in METRICS.select(cost="expensive")} == {"viol", "iog", "cog", "fog"} + + +def test_cheap_metrics_never_touch_the_solver(fake_problem: Any) -> None: + """Running every cheap metric must not call `problem.reset`, which only the solver path does.""" + rng = np.random.default_rng(4) + ctx = _context( + fake_problem, + rng.random((6, *fake_problem.design_space.shape)), + rng.random((6, *fake_problem.design_space.shape)), + ) + for spec in METRICS.select(cost="cheap"): + spec.fn(ctx) + assert fake_problem.reset_calls == 0 + + +# ---------------------------------------------------------------------- +# Registry behaviour +# ---------------------------------------------------------------------- + + +def test_registry_rejects_duplicate_names() -> None: + """Two metrics with one name would make the leaderboard ambiguous.""" + registry = MetricRegistry() + register_metric("dup", family="diversity", cost="cheap", higher_is_better=True, registry=registry)(lambda ctx: 1.0) + with pytest.raises(ValueError, match="already registered"): + register_metric("dup", family="diversity", cost="cheap", higher_is_better=True, registry=registry)(lambda ctx: 2.0) + + +def test_registry_rejects_colliding_output_columns() -> None: + """Two metrics writing the same column would silently overwrite each other.""" + registry = MetricRegistry() + register_metric( + "first", family="diversity", cost="cheap", higher_is_better=True, outputs=("shared",), registry=registry + )(lambda ctx: {"shared": 1.0}) + with pytest.raises(ValueError, match="already emitted"): + register_metric( + "second", family="diversity", cost="cheap", higher_is_better=True, outputs=("shared",), registry=registry + )(lambda ctx: {"shared": 2.0}) + + +def test_multi_output_metric_must_declare_its_columns() -> None: + """Undeclared columns are rejected so the leaderboard schema stays predictable.""" + spec = MetricSpec( + name="multi", + fn=lambda ctx: {"a": 1.0}, + family="diversity", + cost="cheap", + higher_is_better=True, + outputs=("a", "b"), + ) + assert _as_columns(spec, {"a": 1.0, "b": 2.0}) == {"a": 1.0, "b": 2.0} + with pytest.raises(ValueError, match="undeclared columns"): + _as_columns(spec, {"a": 1.0, "c": 3.0}) + + +# ---------------------------------------------------------------------- +# Leaderboard +# ---------------------------------------------------------------------- + + +def _board() -> pd.DataFrame: + return pd.DataFrame( + [ + { + "problem_id": "p", + "algo_id": "a", + "config_fingerprint": "cfg_a", + "seed": 1, + "spec_version": "v1", + "mmd": 0.1, + "dpp": 5.0, + }, + { + "problem_id": "p", + "algo_id": "a", + "config_fingerprint": "cfg_a", + "seed": 2, + "spec_version": "v1", + "mmd": 0.3, + "dpp": 7.0, + }, + { + "problem_id": "p", + "algo_id": "b", + "config_fingerprint": "cfg_b", + "seed": 1, + "spec_version": "v1", + "mmd": 0.05, + "dpp": 1.0, + }, + ] + ) + + +def test_rank_respects_each_metric_direction() -> None: + """Lower MMD ranks first; higher DPP ranks first.""" + assert rank(_board(), "mmd").iloc[0]["algo_id"] == "b" + assert rank(_board(), "dpp").iloc[0]["algo_id"] == "a" + + +def test_rank_aggregates_seeds_with_the_median_by_default() -> None: + """Median across seeds, so one blown-up run cannot decide the ranking.""" + ranked = rank(_board(), "mmd").set_index("algo_id") + assert ranked.loc["a", "mmd_median"] == pytest.approx(0.2) + assert ranked.loc["a", "n_seeds"] == 2 + + +def test_disagreement_surfaces_conflicting_rankings() -> None: + """The two metrics crown different winners; the view must show both.""" + table = disagreement(_board(), ["mmd", "dpp"]) + assert table.loc[("p", "b"), "mmd"] == 1 + assert table.loc[("p", "a"), "dpp"] == 1 + + +def test_append_rows_supersedes_a_rerun_instead_of_duplicating(tmp_path: Any) -> None: + """Re-evaluating a model replaces its row rather than adding a second one.""" + destination = tmp_path / "board.csv" + append_rows(_board(), destination) + rerun = _board().iloc[[0]].copy() + rerun.loc[:, "mmd"] = 0.999 + append_rows(rerun, destination) + + stored = pd.read_csv(destination) + assert len(stored) == 3 + match = stored.set_index(ROW_KEY).loc[("p", "a", "cfg_a", 1, "v1")] + assert match["mmd"] == pytest.approx(0.999) + + +def test_order_columns_puts_provenance_first() -> None: + """Leaderboard CSVs should read identity-first, then measurements.""" + scrambled = pd.DataFrame([{"mmd": 0.1, "algo_id": "a", "problem_id": "p"}]) + assert list(order_columns(scrambled).columns)[:2] == ["problem_id", "algo_id"] + + +# ---------------------------------------------------------------------- +# Eval spec +# ---------------------------------------------------------------------- + + +def test_spec_round_trips_through_disk(tmp_path: Any) -> None: + """A spec must reload byte-identically, since it is the published contract.""" + spec = EvalSpec(problem_id="fake", version="v3", n_samples=7, condition_digest="abc123") + spec.save(root=tmp_path) + assert EvalSpec.load("fake/v3", root=tmp_path) == spec + + +def test_missing_spec_explains_how_to_create_one(tmp_path: Any) -> None: + """The error should tell you the next command, not just that a file is absent.""" + with pytest.raises(FileNotFoundError, match="freeze"): + EvalSpec.load("nope/v1", root=tmp_path) + + +def test_committed_specs_are_frozen() -> None: + """Every committed spec must carry a digest, or it is not pinning anything.""" + from engiopt.evaluation.spec import SPEC_ROOT + + specs = list(SPEC_ROOT.glob("*/*.json")) + assert specs, "expected at least one committed eval spec" + for path in specs: + spec = EvalSpec.load(str(path)) + assert spec.condition_digest, f"{path} is not frozen" + assert spec.metrics + + +# ---------------------------------------------------------------------- +# Incremental leaderboard updates +# ---------------------------------------------------------------------- + + +def _row(**overrides: Any) -> dict[str, Any]: + base = { + "problem_id": "beams2d", + "algo_id": "cgan_cnn_2d", + "config_fingerprint": "aaaa1111", + "seed": 1, + "spec_version": "v1", + "mmd": 0.10, + } + base.update(overrides) + return base + + +def test_different_hyperparameters_get_their_own_rows() -> None: + """A sweep of one algorithm must not collapse onto a single row. + + This is the case that silently discarded 49 of 50 sweep configs before + `config_fingerprint` joined the key. + """ + from engiopt.evaluation.leaderboard import merge_rows + + board = merge_rows( + pd.DataFrame([_row(config_fingerprint="aaaa1111", mmd=0.10)]), + pd.DataFrame([_row(config_fingerprint="bbbb2222", mmd=0.42)]), + ) + assert len(board) == 2 + assert set(board["mmd"]) == {0.10, 0.42} + + +def test_reevaluating_the_same_checkpoint_supersedes_it() -> None: + """Same model, same config, same seed, same spec -> one row, newest wins.""" + from engiopt.evaluation.leaderboard import merge_rows + + board = merge_rows(pd.DataFrame([_row(mmd=0.10)]), pd.DataFrame([_row(mmd=0.99)])) + assert len(board) == 1 + assert board.iloc[0]["mmd"] == pytest.approx(0.99) + + +def test_adding_a_model_leaves_existing_rows_untouched() -> None: + """Publishing one model must not disturb anyone else's results.""" + from engiopt.evaluation.leaderboard import merge_rows + + existing = pd.DataFrame([_row(algo_id="cgan_cnn_2d", mmd=0.10), _row(algo_id="vqgan", mmd=0.20)]) + board = merge_rows(existing, pd.DataFrame([_row(algo_id="my_model", mmd=0.05)])) + assert len(board) == 3 + assert board.set_index("algo_id").loc["cgan_cnn_2d", "mmd"] == pytest.approx(0.10) + assert board.set_index("algo_id").loc["vqgan", "mmd"] == pytest.approx(0.20) + + +def test_seeds_are_separate_rows_and_aggregate_at_ranking_time() -> None: + """The board stores one row per seed; `rank` collapses them.""" + from engiopt.evaluation.leaderboard import merge_rows + + board = merge_rows( + pd.DataFrame([_row(seed=1, mmd=0.10)]), + pd.DataFrame([_row(seed=2, mmd=0.30), _row(seed=3, mmd=0.20)]), + ) + assert len(board) == 3 + ranked = rank(board, "mmd") + assert ranked.iloc[0]["n_seeds"] == 3 + assert ranked.iloc[0]["mmd_median"] == pytest.approx(0.20) + + +def test_already_evaluated_detects_published_work() -> None: + """`--skip-existing` relies on this to avoid recomputing published rows.""" + from engiopt.evaluation.leaderboard import already_evaluated + + board = pd.DataFrame([_row()]) + assert already_evaluated(board, algo_id="cgan_cnn_2d", config_fingerprint="aaaa1111", seed=1) + assert not already_evaluated(board, algo_id="cgan_cnn_2d", config_fingerprint="bbbb2222", seed=1) + assert not already_evaluated(board, algo_id="my_model", config_fingerprint="aaaa1111", seed=1) + assert not already_evaluated(pd.DataFrame(), algo_id="anything") diff --git a/tests/test_generator_registry.py b/tests/test_generator_registry.py new file mode 100644 index 00000000..69a73966 --- /dev/null +++ b/tests/test_generator_registry.py @@ -0,0 +1,115 @@ +"""Tests for generator discovery and the `Generator` contract.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import torch as th + +from engiopt.core import ConditionBatch +from engiopt.core import Generator +from engiopt.utils.all_generators import BUILTIN_GENERATORS +from engiopt.utils.all_generators import extract_generator + + +def test_every_generator_directory_has_an_adapter() -> None: + """Discovery must find an adapter for every model package.""" + packages = {p.name for p in Path("engiopt/generators").iterdir() if p.is_dir() and not p.name.startswith("_")} + assert set(BUILTIN_GENERATORS) == packages + + +def test_algo_ids_match_directory_names() -> None: + """`algo_id` keys the leaderboard and the HF repo, so it must match the package.""" + for name, generator in BUILTIN_GENERATORS.items(): + assert generator.algo_id == name + + +def test_template_is_not_registered() -> None: + """The copy-me template must never appear as a real model.""" + assert "_template" not in BUILTIN_GENERATORS + + +def test_every_generator_declares_the_files_it_needs() -> None: + """`checkpoint_files` drives resolution, so it must be non-empty and concrete.""" + for name, generator in BUILTIN_GENERATORS.items(): + assert generator.checkpoint_files, f"{name} declares no checkpoint files" + assert all(f.endswith(".pth") for f in generator.checkpoint_files), name + + +def test_extract_generator_rejects_ambiguous_modules() -> None: + """Two generators in one adapter module would make the registry ambiguous.""" + + class _Module: + First = type("First", (_StubGenerator,), {"algo_id": "a"}) + Second = type("Second", (_StubGenerator,), {"algo_id": "b"}) + __name__ = "fake.adapter" + + with pytest.raises(ValueError, match="Only one generator per adapter"): + extract_generator(_Module) + + +class _StubGenerator(Generator): + """Minimal in-memory generator, standing in for a trained model.""" + + algo_id = "stub" + conditional = True + design_kinds = ("2d",) + + @classmethod + def build(cls, resolved: Any, problem: Any, device: Any, **base: Any) -> _StubGenerator: + raise NotImplementedError + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + # Deliberately the wrong shape and out of range, to prove `sample` fixes both. + return th.rand(n, int(np.prod(self.design_shape))) * 4.0 - 2.0 + + +def _stub(problem: Any) -> _StubGenerator: + return _StubGenerator(problem, problem_id="fake", seed=1, device=th.device("cpu")) + + +def test_sample_normalizes_shape_and_records_cost(fake_problem: Any) -> None: + """`sample` reshapes to the design shape and times the call.""" + generator = _stub(fake_problem) + designs = generator.sample(th.zeros(5, 2), n=5) + assert designs.shape == (5, *fake_problem.design_space.shape) + assert generator.last_sample_seconds is not None + assert generator.last_sample_seconds >= 0 + + +def test_sample_applies_output_clip(fake_problem: Any) -> None: + """`output_clip` bounds designs that a simulator would otherwise reject.""" + generator = _stub(fake_problem) + generator.output_clip = (0.1, 0.9) + designs = generator.sample(th.zeros(4, 2), n=4) + assert designs.min() >= 0.1 + assert designs.max() <= 0.9 + + +def test_sample_infers_n_from_conditions(fake_problem: Any) -> None: + """Omitting `n` uses the number of supplied conditions.""" + designs = _stub(fake_problem).sample(th.zeros(7, 2)) + assert designs.shape[0] == 7 + + +def test_sample_is_reproducible_for_a_fixed_seed(fake_problem: Any) -> None: + """Same seed, same designs -- the basis of every reproducibility claim.""" + generator = _stub(fake_problem) + first = generator.sample(th.zeros(3, 2), n=3, seed=123) + second = generator.sample(th.zeros(3, 2), n=3, seed=123) + np.testing.assert_array_equal(first, second) + + +def test_condition_batch_reports_a_clear_error_when_empty() -> None: + """A conditional model given nothing should say so by name.""" + with pytest.raises(ValueError, match="my_algo is conditional"): + ConditionBatch(tensor=None).require_tensor("my_algo") + + +def test_registry_is_a_plain_mapping_like_builtin_problems() -> None: + """`BUILTIN_GENERATORS` mirrors `BUILTIN_PROBLEMS`: a dict of name -> class.""" + assert isinstance(BUILTIN_GENERATORS, dict) + assert all(isinstance(g, type) and issubclass(g, Generator) for g in BUILTIN_GENERATORS.values()) From a7a1b497825e72ba55f2d9e04d2157af3d197d3d Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Fri, 24 Jul 2026 14:25:11 +0200 Subject: [PATCH 05/24] Fix CI: unbreak the surrogate model loader and pin down new ruff rules Two failures on the PR, both surfaced only by CI because it runs the whole repository with the latest ruff while local runs used 0.11. - surrogate_model/run_pe_optimization.py passed active_wandb_run to resolve_checkpoint_reference, which lost that parameter when W&B stopped being a checkpoint source. Caught by mypy; the loader now takes only a model reference. - ruff 0.16 formats Python inside Markdown fences, which reformatted the two new docs, and select = ["ALL"] means new rules enable themselves on every upgrade. PLR0917 is ignored for the same reason PLR0913 already is, and PLC0415 because lazy imports are used deliberately to keep module import cost down and avoid cycles. The stale PLR0913 noqa directives those rules made redundant are removed. Co-Authored-By: Claude Opus 4.8 (1M context) --- CONTRIBUTING_A_MODEL.md | 19 +++++++------- docs/checkpoint_layout.md | 3 ++- engiopt/surrogate_model/mlp_tabular_only.py | 4 +-- engiopt/surrogate_model/model_pipeline.py | 2 +- .../surrogate_model/run_pe_optimization.py | 25 +++++-------------- engiopt/surrogate_model/training_utils.py | 2 +- pyproject.toml | 2 ++ 7 files changed, 24 insertions(+), 33 deletions(-) diff --git a/CONTRIBUTING_A_MODEL.md b/CONTRIBUTING_A_MODEL.md index 85de95c5..8c1e2968 100644 --- a/CONTRIBUTING_A_MODEL.md +++ b/CONTRIBUTING_A_MODEL.md @@ -36,12 +36,12 @@ save_checkpoint_package( hf_repo_prefix=args.hf_repo_prefix, hf_private=False, problem_id=args.problem_id, - algo=args.algo, # must equal your algo_id + algo=args.algo, # must equal your algo_id seed=args.seed, checkpoint_files={"generator.pth": "generator.pth"}, - run_config=vars(args), # everything needed to rebuild the model + run_config=vars(args), # everything needed to rebuild the model primary_files=["generator.pth"], - **checkpoint_identity(args), # files this run under its own hyperparameters + **checkpoint_identity(args), # files this run under its own hyperparameters ) ``` @@ -69,12 +69,12 @@ Edit `engiopt/generators/my_model/adapter.py`. Declare what your model is: ```python class MyModel(Generator): - algo_id = "my_model" # == directory name - conditional = True # does _sample use the conditions? - design_kinds = ("2d",) # 1d / 2d / 3d / dict + algo_id = "my_model" # == directory name + conditional = True # does _sample use the conditions? + design_kinds = ("2d",) # 1d / 2d / 3d / dict checkpoint_files = ("generator.pth",) primary_state_key = "generator" - output_clip = (1e-3, 1.0) # or None + output_clip = (1e-3, 1.0) # or None ``` Then implement the two methods: @@ -105,8 +105,8 @@ applies to `Problem` classes, so the name-to-model mapping stays unambiguous. ```python cond = conditions.require_tensor(self.algo_id) # (n, n_conds) on your device -conditions.dataset # original columns, if you need them -conditions.keys # condition names, in column order +conditions.dataset # original columns, if you need them +conditions.keys # condition names, in column order ``` Most models want `require_tensor`. Reach for `dataset` only when your model @@ -133,6 +133,7 @@ comparison rather than to any one model: ```python from engiopt.evaluation import register_metric + @register_metric("my_metric", family="diversity", cost="cheap", higher_is_better=True) def my_metric(ctx) -> float: """One line, shown by --list-metrics.""" diff --git a/docs/checkpoint_layout.md b/docs/checkpoint_layout.md index 41ab9724..37b30db2 100644 --- a/docs/checkpoint_layout.md +++ b/docs/checkpoint_layout.md @@ -66,7 +66,8 @@ gen = BUILTIN_GENERATORS["cgan_cnn_2d"].from_pretrained(problem, problem_id="bea # one specific configuration from a sweep gen = BUILTIN_GENERATORS["cgan_cnn_2d"].from_pretrained( - problem, problem_id="beams2d", seed=1, config_fingerprint="3ab84748") + problem, problem_id="beams2d", seed=1, config_fingerprint="3ab84748" +) ``` ## The leaderboard diff --git a/engiopt/surrogate_model/mlp_tabular_only.py b/engiopt/surrogate_model/mlp_tabular_only.py index 8be76e82..0731d284 100644 --- a/engiopt/surrogate_model/mlp_tabular_only.py +++ b/engiopt/surrogate_model/mlp_tabular_only.py @@ -20,7 +20,6 @@ import torch from torch.utils.data import DataLoader import tyro -import wandb from engiopt.args_utils import parse_list_from_single_item_list from engiopt.args_utils import parse_list_from_string @@ -33,6 +32,7 @@ from engiopt.surrogate_model.training_utils import get_device from engiopt.surrogate_model.training_utils import PlainTabularDataset from engiopt.surrogate_model.training_utils import train_one_model +import wandb @dataclass @@ -123,7 +123,7 @@ def __post_init__(self) -> None: self.params_cols = parse_list_from_single_item_list(self.params_cols, "--params_cols") -def scale_data( # noqa: PLR0913 +def scale_data( x_train: npt.NDArray, x_val: npt.NDArray, x_test: npt.NDArray, diff --git a/engiopt/surrogate_model/model_pipeline.py b/engiopt/surrogate_model/model_pipeline.py index a53dd85d..273191d5 100644 --- a/engiopt/surrogate_model/model_pipeline.py +++ b/engiopt/surrogate_model/model_pipeline.py @@ -322,7 +322,7 @@ class ModelPipeline: Provides convenience methods for prediction, evaluation, saving, and loading. """ - def __init__( # noqa: PLR0913 + def __init__( self, models: list[torch.nn.Module], scalers: dict[str, Any], diff --git a/engiopt/surrogate_model/run_pe_optimization.py b/engiopt/surrogate_model/run_pe_optimization.py index c6154def..92359b38 100644 --- a/engiopt/surrogate_model/run_pe_optimization.py +++ b/engiopt/surrogate_model/run_pe_optimization.py @@ -34,12 +34,12 @@ from pymoo.optimize import minimize from pymoo.termination import get_termination import tyro -import wandb from engiopt.checkpoint_store import resolve_checkpoint_reference from engiopt.surrogate_model.model_pipeline import ModelPipeline from engiopt.surrogate_model.pymoo_pe_problem import PymooPowerElecProblem from engiopt.surrogate_model.training_utils import get_device +import wandb if TYPE_CHECKING: from pymoo.core.algorithm import Algorithm @@ -172,16 +172,11 @@ def save_front(res: Result, output_dir: str) -> tuple[str, str, str, str, str]: return evals_csv, designs_csv, pareto_csv, evals_txt, designs_txt -def load_model_from_reference( - model_ref: str, - *, - active_wandb_run: wandb.sdk.wandb_run.Run | None, -) -> ModelPipeline: - """Load a model pipeline from a W&B artifact, HF package, or local directory. +def load_model_from_reference(model_ref: str) -> ModelPipeline: + """Load a model pipeline from an HF package or local directory. Args: - model_ref: Reference to the stored model package or artifact. - active_wandb_run: Optional active W&B run for artifact access. + model_ref: Reference to the stored model package. Returns: Loaded model pipeline. @@ -190,7 +185,6 @@ def load_model_from_reference( model_source="auto", model_ref=model_ref, required_files=[], - active_wandb_run=active_wandb_run, ) model_file = next(file_path for file_path in resolved.files.values() if file_path.endswith(".pkl")) return ModelPipeline.load(model_file) @@ -216,15 +210,8 @@ def main(args: Args) -> None: wandb.define_metric("generation") wandb.define_metric("*", step_metric="generation") - active_wandb_run = wandb.run if args.track else None - pipeline_g = load_model_from_reference( - args.model_gain_path, - active_wandb_run=active_wandb_run, - ) - pipeline_r = load_model_from_reference( - args.model_ripple_path, - active_wandb_run=active_wandb_run, - ) + pipeline_g = load_model_from_reference(args.model_gain_path) + pipeline_r = load_model_from_reference(args.model_ripple_path) problem = PymooPowerElecProblem( pipeline_r=pipeline_r, diff --git a/engiopt/surrogate_model/training_utils.py b/engiopt/surrogate_model/training_utils.py index b2bcc270..25f1173f 100644 --- a/engiopt/surrogate_model/training_utils.py +++ b/engiopt/surrogate_model/training_utils.py @@ -136,7 +136,7 @@ def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: return self.x[idx], self.y[idx] -def _create_mlp( # noqa: PLR0913 +def _create_mlp( in_dim: int, hidden_size: int, hidden_layers: int, diff --git a/pyproject.toml b/pyproject.toml index c0733cf7..64c87e15 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -128,7 +128,9 @@ ignore = [ "EM102", # f-string-in-exception "E741", # ambiguous-variable-name "FIX002", # flake8-fixme (flake8-todos is enough) + "PLC0415", # import-outside-top-level (used deliberately for lazy/heavy imports) "PLR0913", # too-many-arguments (common in scientific/ML configs) + "PLR0917", # too-many-positional-arguments (same rationale as PLR0913) "PTH", # flake8-use-pathlib "S101", # flake8-bandit: assert "S301", # suspicious-pickle-usage From 3a2568ed22a00efc923d2702fde84f64fcfdd7c7 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Wed, 5 Aug 2026 15:30:06 +0200 Subject: [PATCH 06/24] Address PR #75 review: condition schema, spec pinning, and leaderboard integrity Resolves the thirteen review comments on the generator-contract PR. Three were correctness bugs; the rest close gaps that would have corrupted a shared leaderboard once more than one person wrote to it. Condition schema. `n_conds` counted every entry of `problem.conditions_keys`, including the array-valued and solver-only ones that never reach a dense condition tensor. A cGAN built for thermoelastic2d's seven declared conditions received three columns and failed on reshape. `transforms.condition_keys` is now the single definition of what a model consumes, used by training, loading, and sampling alike, and each checkpoint records the schema it was trained under so a model keeps loading after a problem gains a condition. Handing a model conditions from a different schema now raises rather than silently conditioning on the wrong numbers. Eval specs. The digest covered only the drawn indices, so an upstream edit to a condition value changed every score without tripping the check; it now covers the condition names and values too. Specs also record and load a pinned dataset revision, since an EngiBench version does not identify a dataset. All four specs are re-frozen; photonics2d and thermoelastic2d, which no longer resolved, now do. Feasibility. `viol` only existed for problems with a `volfrac` or `volume` condition, returned NaN for photonics2d, and skipped a legitimate target of 0.0. It is now `problem.check_constraints` plus a volume budget the spec names explicitly. Two further defects surfaced while implementing it: `check_constraints` raised on thermoelastic2d, whose boundary matrices arrive from HuggingFace as lists, and float64 designs were judged infeasible against a float32 design space on dtype alone -- which marked that problem's own dataset optima as violations. Leaderboard integrity. Ranking pooled different hyperparameter configurations and different spec versions into one score; it now groups by both. A failed download read as an empty board, so a transient error could turn the next publish into a deletion; only genuine not-found cases are caught now. Publishing is conditional on the revision it read and re-merges on conflict, rather than overwriting a concurrent publisher. Rows carry the checkpoint revision, a hash of the weights themselves, and the code version, so a score is traceable to the model that earned it. Usability. `local_model_dir` reaches `from_pretrained` and the CLI, so the advertised local source is usable. The CLI exits non-zero when nothing could be evaluated, instead of reporting an empty batch as success. Sample timing synchronises the device before stopping the clock. Each uploaded checkpoint package carries metadata describing itself rather than a path the run may not have written. Testing. The repository ran only ruff and mypy, so a pytest job is added along with tests for the condition schema, publishing races, local package loading, and a network-marked check that every committed spec still resolves against its pinned dataset. Also removes the sweep configs and the per-algorithm notebook that should not have been committed, and two dead helpers this branch had added. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/lint.yaml | 13 + CONTRIBUTING_A_MODEL.md | 30 ++ engiopt/checkpoint_store.py | 111 +++++--- engiopt/core.py | 91 +++++- engiopt/evaluate.py | 88 ++++-- engiopt/evaluation/context.py | 63 ++++- engiopt/evaluation/evaluator.py | 44 ++- engiopt/evaluation/leaderboard.py | 116 ++++++-- engiopt/evaluation/metrics/builtin.py | 10 +- engiopt/evaluation/spec.py | 98 ++++++- engiopt/generators/_template/adapter.py | 4 + engiopt/generators/cgan_1d/adapter.py | 3 +- engiopt/generators/cgan_1d/cgan_1d.py | 11 +- engiopt/generators/cgan_2d/adapter.py | 3 +- engiopt/generators/cgan_2d/cgan_2d.py | 11 +- engiopt/generators/cgan_bezier/adapter.py | 3 +- engiopt/generators/cgan_bezier/cgan_bezier.py | 13 +- engiopt/generators/cgan_cnn_2d/adapter.py | 3 +- engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py | 11 +- .../cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml | 44 --- .../sweep_hp_tuning_heatconduction2d.yaml | 44 --- .../sweep_hp_tuning_photonics2d.yaml | 44 --- engiopt/generators/cgan_cnn_3d/README.md | 24 +- engiopt/generators/cgan_cnn_3d/adapter.py | 3 +- engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py | 8 +- .../cgan_cnn_3d/sweep_cgan_cnn_3d.yaml | 40 --- engiopt/generators/cgan_vae/README.md | 24 +- engiopt/generators/cgan_vae/adapter.py | 3 +- engiopt/generators/cgan_vae/cgan_vae.py | 8 +- .../generators/cgan_vae/sample_script.ipynb | 16 -- .../generators/cgan_vae/sweep_cgan_vae.yaml | 56 ---- .../generators/diffusion_2d_cond/adapter.py | 3 +- .../diffusion_2d_cond/diffusion_2d_cond.py | 15 +- engiopt/generators/gan_1d/gan_1d.py | 3 +- engiopt/generators/pixel_cnn_pp_2d/adapter.py | 3 +- .../pixel_cnn_pp_2d/pixel_cnn_pp_2d.py | 10 +- engiopt/generators/vqgan/vqgan.py | 10 +- engiopt/specs/beams2d/v1.json | 11 +- engiopt/specs/heatconduction2d/v1.json | 11 +- engiopt/specs/photonics2d/v1.json | 11 +- engiopt/specs/thermoelastic2d/v1.json | 9 +- engiopt/transforms.py | 20 ++ engiopt/utils/all_generators.py | 19 -- pyproject.toml | 5 + tests/conftest.py | 28 ++ tests/test_checkpoint_packages.py | 263 ++++++++++++++++++ tests/test_condition_schema.py | 96 +++++++ tests/test_conditions_and_gaps.py | 78 ++++++ tests/test_evaluation.py | 46 ++- tests/test_leaderboard_publishing.py | 169 +++++++++++ tests/test_specs.py | 148 ++++++++++ 51 files changed, 1534 insertions(+), 464 deletions(-) delete mode 100644 engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml delete mode 100644 engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_heatconduction2d.yaml delete mode 100644 engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_photonics2d.yaml delete mode 100644 engiopt/generators/cgan_cnn_3d/sweep_cgan_cnn_3d.yaml delete mode 100644 engiopt/generators/cgan_vae/sample_script.ipynb delete mode 100644 engiopt/generators/cgan_vae/sweep_cgan_vae.yaml create mode 100644 tests/test_checkpoint_packages.py create mode 100644 tests/test_condition_schema.py create mode 100644 tests/test_leaderboard_publishing.py create mode 100644 tests/test_specs.py diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index 13851858..fe29057b 100644 --- a/.github/workflows/lint.yaml +++ b/.github/workflows/lint.yaml @@ -27,3 +27,16 @@ jobs: - uses: actions/checkout@v4 - run: pip install mypy numpy pytest - run: mypy + + pytest: + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + - uses: actions/setup-python@v4 + with: + python-version: "3.12" + - uses: actions/checkout@v4 + - run: pip install -e ".[testing]" + # Every committed eval spec must still resolve against its pinned dataset + # revision, so a spec cannot silently stop reproducing its own conditions. + - run: pytest tests/ -v diff --git a/CONTRIBUTING_A_MODEL.md b/CONTRIBUTING_A_MODEL.md index 8c1e2968..106ca238 100644 --- a/CONTRIBUTING_A_MODEL.md +++ b/CONTRIBUTING_A_MODEL.md @@ -41,6 +41,7 @@ save_checkpoint_package( checkpoint_files={"generator.pth": "generator.pth"}, run_config=vars(args), # everything needed to rebuild the model primary_files=["generator.pth"], + condition_keys=cond_keys, # the condition schema you trained against **checkpoint_identity(args), # files this run under its own hyperparameters ) ``` @@ -54,6 +55,35 @@ sweep therefore can never redefine what your model means. `run_config` matters: it is what `from_pretrained` reads to reconstruct your network, so every architectural hyperparameter must be in `Args`. +### How many conditions does your model have? + +Not `len(problem.conditions_keys)`. That is the full contract, and it is wider +than a dense condition tensor can carry: thermoelastic2d declares seven +conditions, four of which are 65x65 boundary matrices, and photonics2d declares +solver settings that are not dataset columns at all. Size your network from the +*scalar* conditions instead: + +```python +from engiopt.transforms import condition_keys + +cond_keys = condition_keys(problem) # e.g. ("volume_fraction_target", "rmin", "weight") +n_conds = len(cond_keys) +``` + +Use `cond_keys` for your dataloader columns too, and pass it to +`save_checkpoint_package` as above. Your adapter then rebuilds the network for +exactly the columns the run saw: + +```python +from engiopt.core import condition_keys_for + +n_conds = len(condition_keys_for(problem, resolved)) +``` + +That is what keeps a checkpoint loadable after a problem gains a condition, and +what makes the evaluator refuse -- with a clear message -- to feed your model +conditions it was not trained on. + Also reuse the shared helpers rather than rewriting them: | Need | Use | diff --git a/engiopt/checkpoint_store.py b/engiopt/checkpoint_store.py index 8fcea504..c55cd9f1 100644 --- a/engiopt/checkpoint_store.py +++ b/engiopt/checkpoint_store.py @@ -13,6 +13,7 @@ from __future__ import annotations from dataclasses import dataclass +import hashlib import json import os from pathlib import Path @@ -41,6 +42,15 @@ class ResolvedCheckpoint: files: dict[str, str] run_config: dict[str, Any] metadata: dict[str, Any] + revision: str | None = None + """Repo commit the package was downloaded at; None for local packages.""" + content_hash: str | None = None + """Hash of the weight files themselves, identifying these exact weights. + + Two runs of the same configuration and seed produce different weights, and + the repo path cannot tell them apart. This can, so a leaderboard row is + traceable to the model that produced it. + """ def build_hf_repo_id(hf_entity: str, hf_repo_prefix: str, algo: str) -> str: @@ -74,6 +84,7 @@ def save_checkpoint_package( extra_path_parts: list[str] | None = None, config_fingerprint: str | None = None, is_default_config: bool = True, + condition_keys: list[str] | tuple[str, ...] | None = None, ) -> dict[str, Any]: """Save a checkpoint package to HuggingFace. @@ -90,6 +101,13 @@ def save_checkpoint_package( `engiopt.core.checkpoint_identity(args)`. Omitting them keeps the older behaviour of writing only the canonical path. + Each uploaded package carries its own `metadata.json` describing *that* + package, so a `metadata.json` never points at a path this run did not write. + + Pass `condition_keys` (from `engiopt.transforms.condition_keys`) so the + checkpoint records the condition schema it was trained under; loading then + rebuilds the network for exactly those columns. + W&B is no longer a checkpoint storage backend; the active W&B run still receives a summary pointing at the HF package for traceability. """ @@ -101,7 +119,7 @@ def save_checkpoint_package( "hf_revision": None, "hf_config_revision": None, } - metadata_payload = _build_metadata( + base_metadata = _build_metadata( problem_id=problem_id, algo=algo, seed=seed, @@ -109,56 +127,49 @@ def save_checkpoint_package( checkpoint_files=checkpoint_files, primary_files=primary_files, metadata=metadata, + condition_keys=condition_keys, ) - metadata_payload.update(_build_wandb_run_metadata()) + base_metadata.update(_build_wandb_run_metadata()) if checkpoint_backend == "hf": repo_id = build_hf_repo_id(hf_entity, hf_repo_prefix, algo) - package_path = build_hf_package_path(problem_id, seed, extra_path_parts) + canonical_path = build_hf_package_path(problem_id, seed, extra_path_parts) info["hf_repo_id"] = repo_id - info["hf_package_path"] = package_path - metadata_payload["hf_repo_id"] = repo_id - metadata_payload["hf_package_path"] = package_path + base_metadata["hf_repo_id"] = repo_id + if config_fingerprint: + base_metadata["config_fingerprint"] = config_fingerprint # Every configuration gets its own addressable location, so a sweep's # runs cannot overwrite one another. if config_fingerprint: - config_package_path = build_hf_package_path( - problem_id, seed, [*(extra_path_parts or []), f"cfg_{config_fingerprint}"] - ) - info["hf_config_package_path"] = config_package_path - metadata_payload["config_fingerprint"] = config_fingerprint - metadata_payload["hf_config_package_path"] = config_package_path + config_path = build_hf_package_path(problem_id, seed, [*(extra_path_parts or []), f"cfg_{config_fingerprint}"]) + info["hf_config_package_path"] = config_path info["hf_config_revision"] = _upload_package_to_hf( repo_id=repo_id, hf_private=hf_private, - package_path=config_package_path, + package_path=config_path, checkpoint_files=checkpoint_files, run_config=run_config, - metadata=metadata_payload, + metadata={**base_metadata, "hf_package_path": config_path}, algo=algo, ) - metadata_payload["hf_config_revision"] = info["hf_config_revision"] # The canonical path defines what the bare model name resolves to, so # only a default-hyperparameter run may claim it. if is_default_config: - revision = _upload_package_to_hf( + info["hf_package_path"] = canonical_path + info["hf_revision"] = _upload_package_to_hf( repo_id=repo_id, hf_private=hf_private, - package_path=package_path, + package_path=canonical_path, checkpoint_files=checkpoint_files, run_config=run_config, - metadata=metadata_payload, + metadata={**base_metadata, "hf_package_path": canonical_path}, algo=algo, ) - info["hf_revision"] = revision - metadata_payload["hf_revision"] = revision - else: - info["hf_package_path"] = None if wandb.run is not None: - _log_checkpoint_summary_to_wandb(metadata_payload, info) + _log_checkpoint_summary_to_wandb(base_metadata, info) return info @@ -277,7 +288,11 @@ def _resolve_hf_package(*, repo_id: str, package_path: str, required_files: list root_dir = os.path.join(repo_snapshot, package_path) if not os.path.isdir(root_dir): raise FileNotFoundError(f"HF package path not found: {repo_id}/{package_path}") - return _load_package_from_directory(root_dir=root_dir, required_files=required_files, source="hf") + # `snapshot_download` returns `.../snapshots/`, which is the + # revision the files were actually taken from. + return _load_package_from_directory( + root_dir=root_dir, required_files=required_files, source="hf", revision=os.path.basename(repo_snapshot) + ) def _resolve_local_package(local_model_dir: str, required_files: list[str]) -> ResolvedCheckpoint: @@ -286,8 +301,26 @@ def _resolve_local_package(local_model_dir: str, required_files: list[str]) -> R return _load_package_from_directory(root_dir=local_model_dir, required_files=required_files, source="local") +def hash_checkpoint_files(files: dict[str, str]) -> str: + """Content hash of a package's weight files, identifying these exact weights. + + Args: + files: Package file name to local path, as on `ResolvedCheckpoint`. + + Returns: + A 16-character hash over the file names and their bytes. + """ + hasher = hashlib.sha256() + for file_name in sorted(files): + hasher.update(file_name.encode()) + with open(files[file_name], "rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + hasher.update(chunk) + return hasher.hexdigest()[:16] + + def _load_package_from_directory( - *, root_dir: str, required_files: list[str], source: Literal["hf", "local"] + *, root_dir: str, required_files: list[str], source: Literal["hf", "local"], revision: str | None = None ) -> ResolvedCheckpoint: run_config_path = os.path.join(root_dir, "run_config.json") metadata_path = os.path.join(root_dir, "metadata.json") @@ -300,22 +333,15 @@ def _load_package_from_directory( for file_name, file_path in files.items(): if not os.path.exists(file_path): raise FileNotFoundError(f"Missing checkpoint file {file_name} in {root_dir}") - return ResolvedCheckpoint(source=source, root_dir=root_dir, files=files, run_config=run_config, metadata=metadata) - - -def _discover_reference_files(root_dir: str, required_files: list[str]) -> dict[str, str]: - if required_files: - files = {file_name: os.path.join(root_dir, file_name) for file_name in required_files} - for file_name, file_path in files.items(): - if not os.path.exists(file_path): - raise FileNotFoundError(f"Missing checkpoint file {file_name} in {root_dir}") - return files - - discovered = [entry for entry in os.listdir(root_dir) if os.path.isfile(os.path.join(root_dir, entry))] - if len(discovered) != 1: - raise ValueError(f"Expected exactly one file in {root_dir}, found {discovered}") - file_name = discovered[0] - return {file_name: os.path.join(root_dir, file_name)} + return ResolvedCheckpoint( + source=source, + root_dir=root_dir, + files=files, + run_config=run_config, + metadata=metadata, + revision=revision, + content_hash=hash_checkpoint_files(files), + ) def _discover_package_files(root_dir: str, metadata: dict[str, Any]) -> list[str]: @@ -354,6 +380,7 @@ def _build_metadata( checkpoint_files: dict[str, str], primary_files: list[str] | None, metadata: dict[str, Any] | None, + condition_keys: list[str] | tuple[str, ...] | None = None, ) -> dict[str, Any]: payload = dict(metadata or {}) payload.update( @@ -366,6 +393,8 @@ def _build_metadata( "primary_files": primary_files or sorted(checkpoint_files), } ) + if condition_keys is not None: + payload["condition_keys"] = list(condition_keys) return payload diff --git a/engiopt/core.py b/engiopt/core.py index 97995ff3..328f6e4a 100644 --- a/engiopt/core.py +++ b/engiopt/core.py @@ -27,6 +27,7 @@ import torch as th from engiopt.checkpoint_store import resolve_named_checkpoint +from engiopt.transforms import condition_keys as scalar_condition_keys if TYPE_CHECKING: from datasets import Dataset @@ -164,6 +165,32 @@ def config_path_parts(config_fingerprint_value: str | None) -> list[str] | None: return [f"cfg_{config_fingerprint_value}"] +CONDITION_KEYS_METADATA_FIELD = "condition_keys" +"""Checkpoint-metadata field recording which conditions a model was trained on.""" + + +def condition_keys_for(problem: Problem, resolved: ResolvedCheckpoint | None = None) -> tuple[str, ...]: + """Which conditions a generator consumes, in the order its input tensor holds them. + + A checkpoint records the exact condition schema it was trained under, so a + network is always rebuilt for the columns it actually saw. Checkpoints + written before that was recorded fall back to the problem's current scalar + conditions, which is what those runs used. + + Args: + problem: The problem the checkpoint was trained on. + resolved: A downloaded checkpoint package, when one is available. + + Returns: + Condition names, in `problem.conditions_keys` order. + """ + if resolved is not None: + recorded = resolved.metadata.get(CONDITION_KEYS_METADATA_FIELD) + if recorded: + return tuple(str(key) for key in recorded) + return tuple(scalar_condition_keys(problem)) + + def design_shape_of(problem: Problem) -> tuple[int, ...]: """Shape of a single design for `problem`. @@ -221,12 +248,20 @@ def __init__( seed: int, device: th.device, run_config: dict[str, Any] | None = None, + condition_keys: tuple[str, ...] | None = None, + checkpoint_revision: str | None = None, + checkpoint_hash: str | None = None, ) -> None: self.problem = problem self.problem_id = problem_id self.seed = seed self.device = device self.run_config: dict[str, Any] = run_config or {} + self._condition_keys = condition_keys + self.checkpoint_revision = checkpoint_revision + """Commit the checkpoint was downloaded at, when it came from the Hub.""" + self.checkpoint_hash = checkpoint_hash + """Content hash of the weight files, identifying the exact weights scored.""" self.last_sample_seconds: float | None = None """Wall-clock seconds for the most recent `sample` call (cost metrics).""" @@ -274,6 +309,7 @@ def from_pretrained( model_source: ModelSource = "auto", hf_entity: str = DEFAULT_HF_ENTITY, hf_repo_prefix: str = DEFAULT_HF_REPO_PREFIX, + local_model_dir: str | None = None, config_fingerprint: str | None = None, **kwargs: Any, ) -> Generator: @@ -287,6 +323,8 @@ def from_pretrained( model_source: `auto`, `hf`, or `local`. hf_entity: HF org/user holding the checkpoint repos. hf_repo_prefix: Prefix of the per-model-family repo. + local_model_dir: Directory holding an unpacked checkpoint package, + for `local` (or as the fallback under `auto`). config_fingerprint: Load one specific hyperparameter configuration; omit for the canonical default-hyperparameter checkpoint. **kwargs: Passed through to `build`, for model-specific options. @@ -301,6 +339,7 @@ def from_pretrained( model_source=model_source, hf_entity=hf_entity, hf_repo_prefix=hf_repo_prefix, + local_model_dir=local_model_dir, config_fingerprint=config_fingerprint, ) return cls.build( @@ -310,6 +349,9 @@ def from_pretrained( problem_id=problem_id, seed=seed, run_config=resolved.run_config, + condition_keys=condition_keys_for(problem, resolved), + checkpoint_revision=resolved.revision, + checkpoint_hash=resolved.content_hash, **kwargs, ) @@ -365,6 +407,9 @@ def sample( start = time.perf_counter() with th.no_grad(): raw = self._sample(batch, n) + # CUDA and MPS queue work asynchronously, so the elapsed time is only + # the sampling cost once the device has actually finished. + self._synchronize_device() self.last_sample_seconds = time.perf_counter() - start designs = raw.detach().cpu().numpy() if isinstance(raw, th.Tensor) else np.asarray(raw) @@ -373,19 +418,40 @@ def sample( designs = np.clip(designs, *self.output_clip) return designs + def _synchronize_device(self) -> None: + """Wait for queued device work, so timings measure sampling rather than dispatch.""" + if self.device.type == "cuda": + th.cuda.synchronize(self.device) + elif self.device.type == "mps": + th.mps.synchronize() + def _as_batch(self, conditions: ConditionBatch | th.Tensor | npt.NDArray[Any] | None) -> ConditionBatch: """Normalize whatever the caller passed into a `ConditionBatch` on this device.""" if isinstance(conditions, ConditionBatch): + self._check_condition_keys(conditions.keys) tensor = conditions.tensor moved = None if tensor is None else tensor.to(device=self.device, dtype=th.float) return ConditionBatch(tensor=moved, dataset=conditions.dataset, keys=conditions.keys) if conditions is None: - return ConditionBatch(tensor=None, keys=tuple(self.problem.conditions_keys)) + return ConditionBatch(tensor=None, keys=self.condition_keys) tensor = conditions if isinstance(conditions, th.Tensor) else th.as_tensor(np.asarray(conditions)) - return ConditionBatch( - tensor=tensor.to(device=self.device, dtype=th.float), - keys=tuple(self.problem.conditions_keys), - ) + return ConditionBatch(tensor=tensor.to(device=self.device, dtype=th.float), keys=self.condition_keys) + + def _check_condition_keys(self, keys: tuple[str, ...]) -> None: + """Refuse conditions whose columns are not the ones this model was trained on. + + Silently accepting them produces either a shape error deep inside the + network or, worse, a design conditioned on the wrong numbers. + + Raises: + ValueError: If the caller's condition columns differ from the + checkpoint's. + """ + if keys and self.condition_keys and tuple(keys) != self.condition_keys: + raise ValueError( + f"{self.algo_id} was trained on conditions {self.condition_keys} but was given {tuple(keys)}. " + "Re-train against the current problem, or evaluate the checkpoint that matches it." + ) @cached_property def design_shape(self) -> tuple[int, ...]: @@ -397,10 +463,21 @@ def config_fingerprint(self) -> str: """Fingerprint of this checkpoint's hyperparameters; see `config_fingerprint`.""" return config_fingerprint(self.run_config) + @cached_property + def condition_keys(self) -> tuple[str, ...]: + """Condition columns this model consumes, in input-tensor order. + + Taken from the checkpoint when it recorded them, so a model keeps + working even if the problem later gains a condition. + """ + if self._condition_keys is not None: + return tuple(self._condition_keys) + return condition_keys_for(self.problem) + @property def n_conds(self) -> int: - """Number of scalar conditioning variables for this problem.""" - return len(self.problem.conditions_keys) + """Number of scalar conditioning variables this model consumes.""" + return len(self.condition_keys) # ------------------------------------------------------------------ # Checkpoint plumbing shared by every adapter diff --git a/engiopt/evaluate.py b/engiopt/evaluate.py index 9c7ccd21..74952b4d 100644 --- a/engiopt/evaluate.py +++ b/engiopt/evaluate.py @@ -16,6 +16,7 @@ from __future__ import annotations import dataclasses +import sys from typing import TYPE_CHECKING import pandas as pd @@ -73,7 +74,13 @@ class Args: Makes a checkpoint self-describing: whoever downloads it sees how it scored without needing the leaderboard.""" model_source: str = "auto" - """Checkpoint backend: `auto`, `hf`, `wandb`, or `local`.""" + """Checkpoint backend: `auto`, `hf`, or `local`.""" + local_model_dir: str | None = None + """Directory holding an unpacked checkpoint package, for `--model-source local`. + + The package layout is the one `save_checkpoint_package` writes: the weight + files plus `run_config.json` and `metadata.json`. `{algo}` and `{seed}` are + substituted, so one directory template can serve a multi-model run.""" hf_entity: str = "IDEALLab" """HF org/user holding the checkpoints.""" hf_repo_prefix: str = "engiopt" @@ -134,6 +141,7 @@ def _load_generators(args: Args, evaluator: Evaluator) -> list[Generator]: model_source=args.model_source, # type: ignore[arg-type] hf_entity=args.hf_entity, hf_repo_prefix=args.hf_repo_prefix, + local_model_dir=_local_dir_for(args, name, seed), config_fingerprint=fingerprint, ) ) @@ -146,6 +154,34 @@ def _load_generators(args: Args, evaluator: Evaluator) -> list[Generator]: return generators +def _local_dir_for(args: Args, algo: str, seed: int) -> str | None: + """Fill `{algo}` / `{seed}` into `--local-model-dir`, so one template serves a run.""" + if args.local_model_dir is None: + return None + return args.local_model_dir.format(algo=algo, seed=seed, problem_id=args.problem_id) + + +def _drop_already_published(args: Args, evaluator: Evaluator, generators: list[Generator]) -> list[Generator]: + """Drop generators the published board already holds rows for, under `--skip-existing`.""" + if not (args.push_to and args.skip_existing): + return generators + published = load_from_hub(args.push_to) + remaining = [ + generator + for generator in generators + if not already_evaluated( + published, + problem_id=args.problem_id, + algo_id=generator.algo_id, + config_fingerprint=generator.config_fingerprint, + seed=generator.seed, + spec_version=evaluator.spec.version, + ) + ] + print(f"{len(remaining)} generator(s) left after skipping already-published rows.") + return remaining + + def _attach_metrics_to_checkpoints(args: Args, board: pd.DataFrame) -> None: """Write each row's scores into the checkpoint package it describes.""" for row in board.to_dict("records"): @@ -165,38 +201,37 @@ def _attach_metrics_to_checkpoints(args: Args, board: pd.DataFrame) -> None: print(f" could not attach metrics for {row['algo_id']}: {exc}") -def main(args: Args) -> None: - """Evaluate the requested generators and append the results to a CSV.""" +def main(args: Args) -> int: + """Evaluate the requested generators and append the results to a CSV. + + Returns: + A process exit status: non-zero when nothing could be evaluated, so a + batch run that loaded no checkpoints does not look like a clean sweep. + A partially successful run still succeeds -- one bad checkpoint must not + discard the rest of the results. + """ if args.list_generators: _print_generators() if args.list_metrics: _print_metrics() if args.list_generators or args.list_metrics: - return + return 0 spec = args.spec or f"{args.problem_id}/v1" evaluator = Evaluator.for_problem(args.problem_id, spec=spec) print(f"Problem {args.problem_id} | spec {evaluator.spec.version} | n={evaluator.spec.n_samples}") - published = load_from_hub(args.push_to) if (args.push_to and args.skip_existing) else None generators = _load_generators(args, evaluator) - if published is not None: - generators = [ - g - for g in generators - if not already_evaluated( - published, - problem_id=args.problem_id, - algo_id=g.algo_id, - config_fingerprint=g.config_fingerprint, - seed=g.seed, - spec_version=evaluator.spec.version, - ) - ] - print(f"{len(generators)} generator(s) left after skipping already-published rows.") + loaded = len(generators) + generators = _drop_already_published(args, evaluator, generators) if not generators: - print("No generators loaded; nothing to evaluate.") - return + # Having skipped everything already on the board is a finished job; + # having loaded nothing at all is a failed one. + if loaded: + print("Everything requested is already published; nothing to evaluate.") + return 0 + print("No generators could be loaded; nothing was evaluated.") + return 1 board = evaluator.leaderboard( generators, @@ -204,10 +239,19 @@ def main(args: Args) -> None: include_expensive=args.include_expensive, on_error=args.on_error, ) + if board.empty: + print("Every generator failed to evaluate; no rows produced.") + return 1 destination = append_rows(board, args.output_csv.format(problem_id=args.problem_id)) print(f"\n{board.to_string(index=False)}\n") print(f"Wrote {len(board)} rows to {destination}") + _publish(args, board) + return 0 + + +def _publish(args: Args, board: pd.DataFrame) -> None: + """Push the results wherever the flags asked, then print the ranking comparison.""" if args.push_to: merged = push_to_hub(board, args.push_to) print(f"Published {len(board)} row(s) to {args.push_to}; board now holds {len(merged)} rows.") @@ -223,4 +267,4 @@ def main(args: Args) -> None: if __name__ == "__main__": - main(tyro.cli(Args)) + sys.exit(main(tyro.cli(Args))) diff --git a/engiopt/evaluation/context.py b/engiopt/evaluation/context.py index e955fc11..a3edfdb7 100644 --- a/engiopt/evaluation/context.py +++ b/engiopt/evaluation/context.py @@ -26,6 +26,16 @@ import numpy.typing as npt +def _as_arrays(conditions: dict[str, Any] | None) -> dict[str, Any]: + """Give array-valued conditions back their array type before constraint checking. + + A HuggingFace dataset hands back thermoelastic2d's boundary matrices as + nested lists, and EngiBench's bound checks compare them numerically, which + raises `TypeError: '<=' not supported between 'float' and 'list'`. + """ + return {key: np.asarray(value) if isinstance(value, list) else value for key, value in (conditions or {}).items()} + + class MultiObjectiveScalarizationError(ValueError): """Raised when a multi-objective problem has no declared scalarization. @@ -59,7 +69,7 @@ class OptimizationResults: fog: list[float] = field(default_factory=list) """Final optimality gap at the end of re-optimization.""" viol: list[bool] = field(default_factory=list) - """Whether each design missed its volume-fraction target beyond tolerance.""" + """Whether each design is infeasible; see `EvaluationContext.is_infeasible`.""" @dataclass @@ -74,6 +84,10 @@ class EvaluationContext: conditions: The conditions each design was asked to satisfy. sigma: Gaussian-kernel bandwidth for MMD and DPP. volfrac_tol: Tolerance used by the volume-fraction violation check. + volume_condition: Name of the condition holding the volume-fraction + budget a design must hit, when the problem has one. Declared by the + spec rather than guessed, so a problem without a volume budget + (photonics2d) is not silently scored against a missing column. sample_seconds: Wall-clock seconds the generator took to produce `gen_designs`, recorded as leaderboard provenance. objective_weights: Fixed weights over `problem.objectives`, used to @@ -90,6 +104,7 @@ class EvaluationContext: conditions: Dataset | None = None sigma: float = 10.0 volfrac_tol: float = 0.01 + volume_condition: str | None = None sample_seconds: float | None = None objective_weights: tuple[float, ...] | None = None objective_weight_condition: str | None = None @@ -239,9 +254,45 @@ def optimization(self) -> OptimizationResults: results.iog.append(self.scalarize_gap(np.asarray(generated_objective) - np.asarray(reference_optimum), i)) results.cog.append(sum(self.scalarize_gap(step_gap, i) for step_gap in gaps)) results.fog.append(self.scalarize_gap(gaps[-1], i)) - - if conditions: - target_vol = conditions.get("volfrac") or conditions.get("volume") - if target_vol is not None: - results.viol.append(bool(np.abs(np.mean(design) - target_vol) >= self.volfrac_tol)) + results.viol.append(self.is_infeasible(design, conditions)) return results + + def is_infeasible(self, design: Any, conditions: dict[str, Any] | None) -> bool: + """Whether one generated design fails the problem's declared feasibility. + + Two sources, both of which the problem itself defines: + + 1. `problem.check_constraints`, EngiBench's own contract -- the design + must lie in the design space and satisfy every declared constraint. + This is what a new problem gets for free. + 2. The volume-fraction budget named by the spec's `volume_condition`, + when the problem has one. Missing that target is a design failing to + honour its brief rather than an invalid design, and no EngiBench + constraint covers it. + + A problem with neither -- photonics2d has no volume budget -- is scored + on (1) alone, rather than reporting NaN. + """ + violations = self.problem.check_constraints(self._as_space_dtype(design), _as_arrays(conditions)) + if getattr(violations, "violations", None): + return True + if self.volume_condition is not None and conditions is not None: + target = conditions.get(self.volume_condition) + # An explicit None test, so a legitimate target of 0.0 still counts. + if target is not None: + return bool(np.abs(np.mean(design) - float(target)) >= self.volfrac_tol) + return False + + def _as_space_dtype(self, design: Any) -> Any: + """Cast a design to the design space's own dtype before the membership check. + + `Box.contains` requires a safely castable dtype, so a float64 design is + rejected by a float32 space no matter what its values are -- which would + mark thermoelastic2d's own dataset-optimal designs infeasible. The cast + makes the check about the values, which is what it is meant to test. + """ + space = self.problem.design_space + dtype = getattr(space, "dtype", None) + if self.is_dict_space or dtype is None: + return design + return np.asarray(design).astype(dtype, copy=False) diff --git a/engiopt/evaluation/evaluator.py b/engiopt/evaluation/evaluator.py index ebb80572..a3ac4540 100644 --- a/engiopt/evaluation/evaluator.py +++ b/engiopt/evaluation/evaluator.py @@ -14,6 +14,10 @@ from dataclasses import dataclass from dataclasses import field import datetime as dt +import functools +import importlib.metadata +from pathlib import Path +import subprocess from typing import Any, TYPE_CHECKING import pandas as pd @@ -45,9 +49,18 @@ "spec_version", "n_samples", "sample_seconds", + "checkpoint_revision", + "checkpoint_hash", + "code_version", "evaluated_at", ) -"""Columns identifying *what was measured*, as opposed to metric values.""" +"""Columns identifying *what was measured*, as opposed to metric values. + +`config_fingerprint` and `seed` say which configuration produced the row; +`checkpoint_hash` says which *weights* did. Re-training the same configuration +and seed, or changing the training code, yields different weights under the same +name, so without it a row cannot be traced back to the model that earned it. +""" @dataclass @@ -136,6 +149,7 @@ def context_for(self, generator: Generator) -> EvaluationContext: conditions=self.resolved.conditions, sigma=self.spec.sigma, volfrac_tol=self.spec.volfrac_tol, + volume_condition=self.spec.volume_condition, sample_seconds=generator.last_sample_seconds, objective_weights=self.spec.objective_weights, objective_weight_condition=self.spec.objective_weight_condition, @@ -193,6 +207,9 @@ def _provenance(self, generator: Generator, ctx: EvaluationContext) -> dict[str, "spec_version": self.spec.version, "n_samples": ctx.n_samples, "sample_seconds": ctx.sample_seconds, + "checkpoint_revision": getattr(generator, "checkpoint_revision", None), + "checkpoint_hash": getattr(generator, "checkpoint_hash", None), + "code_version": code_version(), "evaluated_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"), } @@ -233,6 +250,31 @@ def leaderboard( return order_columns(pd.DataFrame(rows)) +@functools.cache +def code_version() -> str: + """Which EngiOpt produced a result: the installed version, plus a git sha in a checkout. + + Two evaluations of the same checkpoint can differ if the evaluation code + changed between them, so the row records which code it was. + """ + try: + version = importlib.metadata.version("engiopt") + except importlib.metadata.PackageNotFoundError: + version = "unknown" + try: + sha = subprocess.run( + ["git", "-C", str(Path(__file__).resolve().parent), "rev-parse", "--short", "HEAD"], + capture_output=True, + text=True, + check=True, + timeout=5, + ).stdout.strip() + # Outside a checkout (an installed wheel, a container) the version is all there is. + except (OSError, subprocess.SubprocessError): + return version + return f"{version}+{sha}" if sha else version + + def _as_columns(spec: MetricSpec, value: float | dict[str, float]) -> dict[str, float]: """Normalize a metric's return value into named leaderboard columns.""" if isinstance(value, dict): diff --git a/engiopt/evaluation/leaderboard.py b/engiopt/evaluation/leaderboard.py index e50fc83d..e8272057 100644 --- a/engiopt/evaluation/leaderboard.py +++ b/engiopt/evaluation/leaderboard.py @@ -30,6 +30,15 @@ silently keeping only whichever ran last. """ +ENTRY_KEY = ["problem_id", "algo_id", "config_fingerprint", "spec_version"] +"""What a *ranked entry* is: one configuration of one algorithm under one spec. + +`ROW_KEY` minus `seed`, because seeds are what a ranking aggregates over. +Everything else must match: averaging one configuration's score with another's, +or a score under `v1` with a score under `v2`, produces a number that describes +no model and no protocol. +""" + LEADERBOARD_FILE = "leaderboard.csv" """Filename used inside a published HuggingFace dataset repo.""" @@ -72,8 +81,8 @@ def rank( run otherwise decides the ranking. Returns: - One row per `(problem_id, algo_id)` sorted best-first, with a `rank` - column and the seed count that produced each value. + One row per `ENTRY_KEY` sorted best-first, with a `rank` column and the + seed count that produced each value. """ from engiopt.evaluation.registry import METRICS @@ -82,7 +91,7 @@ def rank( raise ValueError(f"Metric {metric!r} has no ranking direction; it is diagnostic only.") grouped = ( - frame.groupby(["problem_id", "algo_id"], as_index=False) + frame.groupby(entry_key(frame), as_index=False) .agg(value=(metric, aggregate), n_seeds=(metric, "count")) .sort_values("value", ascending=not spec.higher_is_better) .reset_index(drop=True) @@ -91,6 +100,19 @@ def rank( return grouped.rename(columns={"value": f"{metric}_{aggregate}"}) +def entry_key(frame: pd.DataFrame) -> list[str]: + """The `ENTRY_KEY` columns present in `frame`, for grouping. + + Raises: + ValueError: If none of them are, since grouping would then silently + pool every result in the table into one number. + """ + key = [col for col in ENTRY_KEY if col in frame.columns] + if not key: + raise ValueError(f"Leaderboard is missing all of {ENTRY_KEY}; cannot tell entries apart.") + return key + + def disagreement( frame: pd.DataFrame, metrics: list[str], @@ -105,12 +127,13 @@ def disagreement( one-off. Returns: - A table indexed by model, one column of ranks per metric. + A table indexed by `ENTRY_KEY`, one column of ranks per metric. """ + key = entry_key(frame) boards = [] for metric in metrics: board = rank(frame, metric, registry=registry, aggregate=aggregate) - boards.append(board.set_index(["problem_id", "algo_id"])["rank"].rename(metric)) + boards.append(board.set_index(key)["rank"].rename(metric)) return pd.concat(boards, axis=1).sort_values(metrics[0]) @@ -129,14 +152,34 @@ def merge_rows(existing: pd.DataFrame, new_rows: pd.DataFrame) -> pd.DataFrame: def load_from_hub(repo_id: str, *, token: str | None = None) -> pd.DataFrame: - """Download the published leaderboard, or an empty frame if none exists yet.""" + """Download the published leaderboard, or an empty frame if none exists yet. + + Raises: + HfHubHTTPError: On any failure other than "no board published yet". + Treating an auth or network error as an empty board would let the + next `push_to_hub` upload only the new rows, deleting everyone + else's. + """ + return _load_from_hub(repo_id, token=token)[0] + + +def _load_from_hub(repo_id: str, *, token: str | None = None) -> tuple[pd.DataFrame, str | None]: + """Download the board together with the repo commit it was read at.""" from huggingface_hub import hf_hub_download + from huggingface_hub import HfApi + from huggingface_hub.errors import EntryNotFoundError + from huggingface_hub.errors import RepositoryNotFoundError + from huggingface_hub.errors import RevisionNotFoundError try: - path = hf_hub_download(repo_id=repo_id, filename=LEADERBOARD_FILE, repo_type="dataset", token=token) - except Exception: # noqa: BLE001 - a missing or private board is not an error here - return pd.DataFrame() - return pd.read_csv(path) + revision = HfApi(token=token).repo_info(repo_id=repo_id, repo_type="dataset").sha + path = hf_hub_download( + repo_id=repo_id, filename=LEADERBOARD_FILE, repo_type="dataset", revision=revision, token=token + ) + # Only "there is nothing published yet" is a normal, empty-board outcome. + except (EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError): + return pd.DataFrame(), None + return pd.read_csv(path), revision def push_to_hub( @@ -146,41 +189,66 @@ def push_to_hub( token: str | None = None, private: bool = False, commit_message: str | None = None, + max_attempts: int = 3, ) -> pd.DataFrame: """Merge `new_rows` into the published leaderboard and upload the result. The existing board is downloaded first and merged on `ROW_KEY`, so adding a model never requires re-running anyone else's evaluation. + The upload is conditional on the revision the board was read at. If another + job published in between, the commit is rejected rather than overwriting it, + and the read-merge-upload cycle is retried against the newer board. + Args: new_rows: Rows to publish, as returned by `Evaluator.leaderboard`. repo_id: HuggingFace dataset repo, e.g. `"IDEALLab/engiopt-leaderboard"`. token: HF token; falls back to the ambient login. private: Create the repo private if it does not exist yet. commit_message: Defaults to a summary of what was added. + max_attempts: How many times to retry after losing a race. Returns: The full merged leaderboard as uploaded. + + Raises: + RuntimeError: If every attempt lost the race to a concurrent publisher. """ from huggingface_hub import HfApi + from huggingface_hub.errors import HfHubHTTPError api = HfApi(token=token) api.create_repo(repo_id=repo_id, repo_type="dataset", private=private, exist_ok=True) - - merged = merge_rows(load_from_hub(repo_id, token=token), new_rows) - models = ", ".join(sorted(new_rows["algo_id"].unique())) if "algo_id" in new_rows else "results" - with tempfile.TemporaryDirectory() as tmp: - local = Path(tmp) / LEADERBOARD_FILE - merged.to_csv(local, index=False) - api.upload_file( - path_or_fileobj=str(local), - path_in_repo=LEADERBOARD_FILE, - repo_id=repo_id, - repo_type="dataset", - commit_message=commit_message or f"Add {len(new_rows)} row(s): {models}", - ) - return merged + + for attempt in range(max_attempts): + existing, revision = _load_from_hub(repo_id, token=token) + merged = merge_rows(existing, new_rows) + with tempfile.TemporaryDirectory() as tmp: + local = Path(tmp) / LEADERBOARD_FILE + merged.to_csv(local, index=False) + try: + api.upload_file( + path_or_fileobj=str(local), + path_in_repo=LEADERBOARD_FILE, + repo_id=repo_id, + repo_type="dataset", + commit_message=commit_message or f"Add {len(new_rows)} row(s): {models}", + parent_commit=revision, + ) + except HfHubHTTPError as exc: + if not _is_stale_commit(exc) or attempt == max_attempts - 1: + raise + print(f"[leaderboard] {repo_id} moved under us; re-merging (attempt {attempt + 2}/{max_attempts}).") + continue + return merged + raise RuntimeError(f"Could not publish to {repo_id} in {max_attempts} attempts: it kept changing underneath.") + + +def _is_stale_commit(exc: Exception) -> bool: + """Whether an upload failed because `parent_commit` was no longer the head.""" + status = getattr(getattr(exc, "response", None), "status_code", None) + return status == 412 or "parent_commit" in str(exc).lower() # noqa: PLR2004 - HTTP 412 Precondition Failed def already_evaluated(board: pd.DataFrame, **key: Any) -> bool: diff --git a/engiopt/evaluation/metrics/builtin.py b/engiopt/evaluation/metrics/builtin.py index 3d597a31..6a60ee7a 100644 --- a/engiopt/evaluation/metrics/builtin.py +++ b/engiopt/evaluation/metrics/builtin.py @@ -57,7 +57,7 @@ def dpp(ctx: EvaluationContext) -> float: # ---------------------------------------------------------------------- -# Feasibility: does the design respect its stated budget? +# Feasibility: is the design admissible, and does it respect its budget? # ---------------------------------------------------------------------- @@ -66,10 +66,14 @@ def dpp(ctx: EvaluationContext) -> float: family="feasibility", cost="expensive", higher_is_better=False, - description="Fraction of designs missing their volume-fraction target beyond tolerance.", + description="Fraction of designs violating the problem's constraints or their volume budget.", ) def viol(ctx: EvaluationContext) -> float: - """Fraction of designs that miss their volume-fraction target. + """Fraction of infeasible designs; see `EvaluationContext.is_infeasible`. + + Defined for every problem: `problem.check_constraints` always applies, and + the spec's `volume_condition` adds the volume-fraction budget for problems + that have one. Marked expensive only because it is computed during the shared optimizer/simulator pass; the arithmetic itself is trivial. diff --git a/engiopt/evaluation/spec.py b/engiopt/evaluation/spec.py index 51fcf55a..d4bf8642 100644 --- a/engiopt/evaluation/spec.py +++ b/engiopt/evaluation/spec.py @@ -46,6 +46,10 @@ class EvalSpec: metrics: Metric names to compute, in leaderboard column order. sigma: Gaussian-kernel bandwidth for MMD and DPP. volfrac_tol: Tolerance for the volume-fraction violation check. + volume_condition: Name of the condition holding the volume-fraction + budget a design must hit, e.g. beams2d's `volfrac`. Feasibility is + otherwise `problem.check_constraints` alone, which is the right + answer for problems with no volume budget (photonics2d). objective_weights: Fixed weights over `problem.objectives`, for multi-objective problems. Summing or averaging objectives with different units is not a valid default, so one of these two fields @@ -54,10 +58,18 @@ class EvalSpec: trade-off instead of fixing it. thermoelastic2d's `weight` splits the first two objectives as `(w, 1 - w)`, so a purely structural sample (`w = 1`) is unaffected by thermal compliance. - condition_digest: Hash of the drawn conditions and reference designs. - Recomputed at evaluation time and compared, so an upstream dataset - change is caught instead of silently shifting every number. + condition_digest: Hash of the drawn indices, condition values, and + reference designs. Recomputed at evaluation time and compared, so an + upstream dataset change is caught instead of silently shifting every + number. + dataset_id: HuggingFace dataset the conditions were drawn from. + dataset_revision: The dataset commit they were drawn at. A spec loads + the dataset at exactly this revision, so a new upload to the dataset + repo cannot change what a leaderboard row means. Leave it empty to + follow the dataset's current main branch. engibench_version: EngiBench version the digest was recorded under. + Provenance only -- a package version does not identify a dataset, + which is what `dataset_revision` is for. notes: Free-text rationale for this version. """ @@ -68,9 +80,12 @@ class EvalSpec: metrics: tuple[str, ...] = ("mmd", "dpp", "viol", "iog", "cog", "fog") sigma: float = 10.0 volfrac_tol: float = 0.01 + volume_condition: str | None = None objective_weights: tuple[float, ...] | None = None objective_weight_condition: str | None = None condition_digest: str | None = None + dataset_id: str | None = None + dataset_revision: str | None = None engibench_version: str | None = None notes: str = "" @@ -123,10 +138,11 @@ def resolve(self, problem: Problem, *, device: th.device | None = None) -> Resol re-frozen under a new version. """ device = device or th.device("cpu") + self.pin_dataset(problem) conditions_tensor, conditions, ref_designs, indices = sample_conditions( problem=problem, n_samples=self.n_samples, device=device, seed=self.condition_seed ) - digest = _digest(indices, ref_designs) + digest = _digest(indices, conditions, ref_designs) if self.condition_digest is not None and digest != self.condition_digest: raise ValueError( f"Eval spec {self.problem_id}/{self.version} no longer reproduces its frozen conditions " @@ -142,17 +158,37 @@ def resolve(self, problem: Problem, *, device: th.device | None = None) -> Resol condition_keys=tuple(get_scalar_condition_keys(problem, problem.dataset["test"])), ) + def pin_dataset(self, problem: Problem) -> None: + """Point `problem.dataset` at this spec's frozen dataset revision. + + `Problem.dataset` lazily loads `dataset_id` from the Hub's main branch, + with no hook for a revision, so the spec loads the pinned revision and + seeds the problem's cache before anything reads it. + """ + if not self.dataset_revision: + return + from datasets import load_dataset + + dataset_id = self.dataset_id or getattr(problem, "dataset_id", None) + if dataset_id is None: + raise ValueError(f"Spec {self.problem_id}/{self.version} pins a dataset revision but names no dataset.") + problem._dataset = load_dataset(dataset_id, revision=self.dataset_revision) # noqa: SLF001 + def freeze(self, problem: Problem) -> EvalSpec: - """Return a copy with `condition_digest` and `engibench_version` filled in.""" + """Return a copy with the digest, dataset revision, and versions filled in.""" import engibench - _, _, ref_designs, indices = sample_conditions( + dataset_id = self.dataset_id or getattr(problem, "dataset_id", None) + revision = self.dataset_revision or _dataset_revision(dataset_id) + frozen = EvalSpec(**{**asdict(self), "dataset_id": dataset_id, "dataset_revision": revision}) + frozen.pin_dataset(problem) + _, conditions, ref_designs, indices = sample_conditions( problem=problem, n_samples=self.n_samples, device=th.device("cpu"), seed=self.condition_seed ) return EvalSpec( **{ - **asdict(self), - "condition_digest": _digest(indices, ref_designs), + **asdict(frozen), + "condition_digest": _digest(indices, conditions, ref_designs), "engibench_version": getattr(engibench, "__version__", None), } ) @@ -181,14 +217,46 @@ def n_samples(self) -> int: return self.spec.n_samples -def _digest(indices: npt.NDArray[Any], ref_designs: npt.NDArray[Any]) -> str: - """Stable hash of the drawn sample indices and their reference designs.""" +def _digest(indices: npt.NDArray[Any], conditions: Dataset, ref_designs: npt.NDArray[Any]) -> str: + """Stable hash of everything a model is scored against. + + Covers the drawn indices, the condition columns *and their values*, and the + reference designs. Hashing only the indices would let an upstream edit + change a condition value -- and so every score -- without tripping the + check, since row 7 is still row 7. + """ hasher = hashlib.sha256() hasher.update(np.asarray(indices).astype(np.int64).tobytes()) + for name in sorted(conditions.column_names): + hasher.update(name.encode()) + hasher.update(_condition_bytes(conditions[name])) hasher.update(np.asarray(ref_designs, dtype=np.float64).round(8).tobytes()) return hasher.hexdigest()[:16] +def _condition_bytes(values: Any) -> bytes: + """Stable byte representation of one condition column. + + Numeric columns (including the array-valued ones, e.g. thermoelastic2d's + boundary matrices) are rounded before hashing so that float noise below the + tolerance the metrics themselves use cannot invalidate a spec; anything + non-numeric falls back to its JSON form. + """ + array = np.asarray(values) + if array.dtype.kind in "fiub": + return np.ascontiguousarray(array, dtype=np.float64).round(8).tobytes() + return json.dumps(values, sort_keys=True, default=str).encode() + + +def _dataset_revision(dataset_id: str | None) -> str | None: + """Current commit of a HuggingFace dataset repo, or None if it cannot be read.""" + if dataset_id is None: + return None + from huggingface_hub import HfApi + + return HfApi().dataset_info(dataset_id).sha + + def freeze_spec( problem_id: str, *, @@ -199,11 +267,13 @@ def freeze_spec( sigma: float = 10.0, objective_weights: tuple[float, ...] | None = None, objective_weight_condition: str | None = None, + notes: str = "", ) -> Path: """Draw a problem's test conditions once and commit them as a spec. Run this once per problem, then never again for that version -- the point of - a spec is that it stops moving. + a spec is that it stops moving. The dataset revision in force is recorded, + so later uploads to the dataset repo cannot change what the spec means. Returns: Path to the written spec file. @@ -221,9 +291,13 @@ def freeze_spec( sigma=sigma, objective_weights=objective_weights, objective_weight_condition=objective_weight_condition, + notes=notes, ).freeze(problem) path = spec.save() - print(f"Froze {problem_id}/{version}: digest={spec.condition_digest} n={spec.n_samples} -> {path}") + print( + f"Froze {problem_id}/{version}: digest={spec.condition_digest} n={spec.n_samples} " + f"dataset={spec.dataset_id}@{spec.dataset_revision} -> {path}" + ) return path diff --git a/engiopt/generators/_template/adapter.py b/engiopt/generators/_template/adapter.py index 791cfef9..24085703 100644 --- a/engiopt/generators/_template/adapter.py +++ b/engiopt/generators/_template/adapter.py @@ -72,6 +72,10 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device # # net = MyNet( # latent_dim=config["latent_dim"], + # # Size the conditioning from the schema the run was trained on -- + # # not `len(problem.conditions_keys)`, which counts array-valued + # # and solver-only conditions your tensor will never carry. + # n_conds=len(condition_keys_for(problem, resolved)), # design_shape=design_shape_of(problem), # ).to(device) # state = th.load(resolved.files["generator.pth"], map_location=device) diff --git a/engiopt/generators/cgan_1d/adapter.py b/engiopt/generators/cgan_1d/adapter.py index 33ce1449..3327c296 100644 --- a/engiopt/generators/cgan_1d/adapter.py +++ b/engiopt/generators/cgan_1d/adapter.py @@ -6,6 +6,7 @@ import torch as th +from engiopt.core import condition_keys_for from engiopt.core import ConditionBatch from engiopt.core import design_shape_of from engiopt.core import Generator @@ -43,7 +44,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device _, conds_normalizer, design_normalizer = prepare_data(problem, device) net = CGAN1DNet( latent_dim=config["latent_dim"], - n_conds=len(problem.conditions_keys), + n_conds=len(condition_keys_for(problem, resolved)), design_shape=design_shape_of(problem), design_normalizer=design_normalizer, conds_normalizer=conds_normalizer, diff --git a/engiopt/generators/cgan_1d/cgan_1d.py b/engiopt/generators/cgan_1d/cgan_1d.py index b0e163b8..b841f6f1 100644 --- a/engiopt/generators/cgan_1d/cgan_1d.py +++ b/engiopt/generators/cgan_1d/cgan_1d.py @@ -26,6 +26,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys from engiopt.transforms import flatten_dict_factory import wandb @@ -192,7 +193,7 @@ def prepare_data(problem: Problem, device: th.device) -> tuple[th.utils.data.Ten training_ds = th.utils.data.TensorDataset( transform(training_ds["optimal_design"][:]), - *[training_ds[key][:] for key in problem.conditions_keys], + *[training_ds[key][:] for key in cond_keys], ) # Create condition normalizer @@ -223,7 +224,10 @@ 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 - n_conds = len(problem.conditions_keys) + # The scalar conditions the generator is conditioned on; array-valued and + # solver-only entries of `problem.conditions_keys` cannot enter a dense tensor. + cond_keys = condition_keys(problem) + n_conds = len(cond_keys) # Logging run_name = f"{args.problem_id}__{args.algo}__{args.seed}__{int(time.time())}" @@ -370,7 +374,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 = [(problem.conditions_keys[i], f"{dc[i]:.2f}") for i in range(n_conds)] + title = [(cond_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 @@ -415,6 +419,7 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, run_config=vars(args), **checkpoint_identity(args), + condition_keys=cond_keys, primary_files=["generator.pth"], ) diff --git a/engiopt/generators/cgan_2d/adapter.py b/engiopt/generators/cgan_2d/adapter.py index ed7fc2f7..e3b8fb37 100644 --- a/engiopt/generators/cgan_2d/adapter.py +++ b/engiopt/generators/cgan_2d/adapter.py @@ -6,6 +6,7 @@ import torch as th +from engiopt.core import condition_keys_for from engiopt.core import ConditionBatch from engiopt.core import Generator from engiopt.generators.cgan_2d.cgan_2d import Generator as CGAN2DNet @@ -37,7 +38,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device config = resolved.run_config net = CGAN2DNet( latent_dim=config["latent_dim"], - n_conds=len(problem.conditions_keys), + n_conds=len(condition_keys_for(problem, resolved)), design_shape=problem.design_space.shape, generator_output_activation=config.get("generator_output_activation", "tanh"), ).to(device) diff --git a/engiopt/generators/cgan_2d/cgan_2d.py b/engiopt/generators/cgan_2d/cgan_2d.py index 43b82180..46e21be9 100644 --- a/engiopt/generators/cgan_2d/cgan_2d.py +++ b/engiopt/generators/cgan_2d/cgan_2d.py @@ -24,6 +24,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys import wandb GeneratorOutputActivation = Literal["tanh", "sigmoid"] @@ -172,7 +173,10 @@ def forward(self, design: th.Tensor, conds: th.Tensor) -> th.Tensor: problem.reset(seed=args.seed) design_shape = problem.design_space.shape - n_conds = len(problem.conditions_keys) + # The scalar conditions the generator is conditioned on; array-valued and + # solver-only entries of `problem.conditions_keys` cannot enter a dense tensor. + cond_keys = condition_keys(problem) + n_conds = len(cond_keys) # Logging run_name = f"{args.problem_id}__{args.algo}__{args.seed}__{int(time.time())}" @@ -208,7 +212,7 @@ def forward(self, design: th.Tensor, conds: th.Tensor) -> th.Tensor: training_ds = problem.dataset.with_format("torch", device=device)["train"] training_ds = th.utils.data.TensorDataset( - training_ds["optimal_design"][:].flatten(1), *[training_ds[key][:] for key in problem.conditions_keys] + training_ds["optimal_design"][:].flatten(1), *[training_ds[key][:] for key in cond_keys] ) dataloader = th.utils.data.DataLoader( training_ds, @@ -312,7 +316,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 = [(problem.conditions_keys[i], f"{dc[i]:.2f}") for i in range(n_conds)] + title = [(cond_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 @@ -356,6 +360,7 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: checkpoint_files={"generator.pth": "generator.pth", "discriminator.pth": "discriminator.pth"}, run_config=vars(args), **checkpoint_identity(args), + condition_keys=cond_keys, primary_files=["generator.pth"], ) diff --git a/engiopt/generators/cgan_bezier/adapter.py b/engiopt/generators/cgan_bezier/adapter.py index 30f95d40..f2947d70 100644 --- a/engiopt/generators/cgan_bezier/adapter.py +++ b/engiopt/generators/cgan_bezier/adapter.py @@ -6,6 +6,7 @@ import torch as th +from engiopt.core import condition_keys_for from engiopt.core import ConditionBatch from engiopt.core import Generator from engiopt.generators.cgan_bezier.cgan_bezier import _EPS @@ -49,7 +50,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device net = CBezierGANNet( latent_dim=config["latent_dim"], noise_dim=config["noise_dim"], - num_conds=len(problem.conditions_keys), + num_conds=len(condition_keys_for(problem, resolved)), n_control_points=config["bezier_control_pts"], n_data_points=problem.design_space["coords"].shape[1], conds_normalizer=conds_normalizer, diff --git a/engiopt/generators/cgan_bezier/cgan_bezier.py b/engiopt/generators/cgan_bezier/cgan_bezier.py index f5ae1aaf..f104bbd3 100644 --- a/engiopt/generators/cgan_bezier/cgan_bezier.py +++ b/engiopt/generators/cgan_bezier/cgan_bezier.py @@ -24,6 +24,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys import wandb if TYPE_CHECKING: @@ -460,7 +461,7 @@ def prepare_data( training_ds = th.utils.data.TensorDataset( th.stack(coords_set), th.stack(design_scalars).unsqueeze(1), - *[problem_dataset[key][:] for key in problem.conditions_keys], + *[problem_dataset[key][:] for key in condition_keys(problem)], ) cond_tensors = th.stack(training_ds.tensors[2:]) @@ -504,6 +505,9 @@ def prepare_data( bezier_control_pts = args.bezier_control_pts n_data_points = problem.design_space["coords"].shape[1] # for airfoil, 192 + # The scalar conditions the generator is conditioned on; array-valued and + # solver-only entries of `problem.conditions_keys` cannot enter a dense tensor. + cond_keys = condition_keys(problem) # The Discriminator uses shape [N, 2, #points]. training_ds, conds_normalizer, design_scalars_normalizer, design_scalar_keys = prepare_data(problem, device) @@ -519,7 +523,7 @@ def prepare_data( discriminator = Discriminator( latent_dim=args.latent_dim, design_scalars=len(design_scalar_keys), - num_conds=len(problem.conditions_keys), + num_conds=len(cond_keys), design_shape=problem.design_space["coords"].shape, conds_normalizer=conds_normalizer, design_scalars_normalizer=design_scalars_normalizer, @@ -528,7 +532,7 @@ def prepare_data( generator = Generator( latent_dim=args.latent_dim, noise_dim=args.noise_dim, - num_conds=len(problem.conditions_keys), + num_conds=len(cond_keys), n_control_points=bezier_control_pts, n_data_points=n_data_points, conds_normalizer=conds_normalizer, @@ -640,7 +644,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_keys[i - 1], f"{do1[i]:.2f}") for i in range(1, len(do1))] + title = [(cond_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([]) @@ -688,6 +692,7 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor, th.Tensor]: }, run_config=vars(args), **checkpoint_identity(args), + condition_keys=cond_keys, primary_files=["bezier_generator.pth"], ) diff --git a/engiopt/generators/cgan_cnn_2d/adapter.py b/engiopt/generators/cgan_cnn_2d/adapter.py index 2adaffd1..5120de20 100644 --- a/engiopt/generators/cgan_cnn_2d/adapter.py +++ b/engiopt/generators/cgan_cnn_2d/adapter.py @@ -6,6 +6,7 @@ import torch as th +from engiopt.core import condition_keys_for from engiopt.core import ConditionBatch from engiopt.core import Generator from engiopt.generators.cgan_cnn_2d.cgan_cnn_2d import Generator as CGANCNN2DNet @@ -38,7 +39,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device config = resolved.run_config net = CGANCNN2DNet( latent_dim=config["latent_dim"], - n_conds=len(problem.conditions_keys), + n_conds=len(condition_keys_for(problem, resolved)), design_shape=problem.design_space.shape, generator_output_activation=config.get("generator_output_activation", "tanh"), ) diff --git a/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py b/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py index d94092c2..6699f45f 100644 --- a/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py +++ b/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py @@ -23,6 +23,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys import wandb GeneratorOutputActivation = Literal["tanh", "sigmoid"] @@ -262,7 +263,10 @@ def forward(self, x: th.Tensor, c: th.Tensor) -> th.Tensor: problem.reset(seed=args.seed) design_shape = problem.design_space.shape - n_conds = len(problem.conditions_keys) + # The scalar conditions the generator is conditioned on; array-valued and + # solver-only entries of `problem.conditions_keys` cannot enter a dense tensor. + cond_keys = condition_keys(problem) + n_conds = len(cond_keys) # Logging run_name = f"{args.problem_id}__{args.algo}__{args.seed}__{int(time.time())}" @@ -302,7 +306,7 @@ def forward(self, x: th.Tensor, c: th.Tensor) -> th.Tensor: # Configure data loader training_ds = problem.dataset.with_format("torch", device=device)["train"] training_ds = th.utils.data.TensorDataset( - training_ds["optimal_design"][:].flatten(1), *[training_ds[key][:] for key in problem.conditions_keys] + training_ds["optimal_design"][:].flatten(1), *[training_ds[key][:] for key in cond_keys] ) dataloader = th.utils.data.DataLoader( training_ds, @@ -408,7 +412,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 = [(problem.conditions_keys[i], f"{dc[i]:.2f}") for i in range(n_conds)] + title = [(cond_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 @@ -455,6 +459,7 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: }, run_config=vars(args), **checkpoint_identity(args), + condition_keys=cond_keys, primary_files=["generator.pth"], ) diff --git a/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml b/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml deleted file mode 100644 index 73147ee7..00000000 --- a/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_beams2d.yaml +++ /dev/null @@ -1,44 +0,0 @@ -method: grid -command: - - "python" - - "engiopt/cgan_cnn_2d/cgan_cnn_2d.py" - - "--save_model" - - "${args}" - -metric: - name: g_loss - goal: minimize -parameters: - problem_id: - value: "beams2d" - lr_gen: - distribution: log_uniform_values - min: 0.00005 - max: 0.005 - lr_disc: - distribution: log_uniform_values - min: 0.00005 - max: 0.005 - latent_dim: - distribution: categorical - values: - - 16 - - 32 - - 64 - - 128 - batch_size: - distribution: categorical - values: - - 16 - - 32 - - 64 - n_epochs: - distribution: categorical - values: - - 100 - - 200 - - 300 - seed: - value: 42 - wandb_project: - value: "engiopt" diff --git a/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_heatconduction2d.yaml b/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_heatconduction2d.yaml deleted file mode 100644 index 1e884788..00000000 --- a/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_heatconduction2d.yaml +++ /dev/null @@ -1,44 +0,0 @@ -method: bayes -metric: - name: g_loss - goal: minimize -run_cap: 50 -parameters: - problem_id: - value: "heatconduction2d" - lr_gen: - distribution: log_uniform_values - min: 0.00005 - max: 0.005 - lr_disc: - distribution: log_uniform_values - min: 0.00005 - max: 0.005 - latent_dim: - distribution: categorical - values: - - 16 - - 32 - - 64 - - 128 - batch_size: - distribution: categorical - values: - - 16 - - 32 - - 64 - n_epochs: - distribution: categorical - values: - - 100 - - 200 - - 300 - track: - value: true - save_model: - value: true - seed: - value: 1 - wandb_project: - value: "lv-mmd-hp-tuning" -program: engiopt/cgan_cnn_2d/cgan_cnn_2d.py diff --git a/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_photonics2d.yaml b/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_photonics2d.yaml deleted file mode 100644 index cc9d3a89..00000000 --- a/engiopt/generators/cgan_cnn_2d/sweep_hp_tuning_photonics2d.yaml +++ /dev/null @@ -1,44 +0,0 @@ -method: grid -command: - - "python" - - "engiopt/cgan_cnn_2d/cgan_cnn_2d.py" - - "--save_model" - - "${args}" - -metric: - name: g_loss - goal: minimize -parameters: - problem_id: - value: "photonics2d" - lr_gen: - distribution: log_uniform_values - min: 0.00005 - max: 0.005 - lr_disc: - distribution: log_uniform_values - min: 0.00005 - max: 0.005 - latent_dim: - distribution: categorical - values: - - 16 - - 32 - - 64 - - 128 - batch_size: - distribution: categorical - values: - - 16 - - 32 - - 64 - n_epochs: - distribution: categorical - values: - - 100 - - 200 - - 300 - seed: - value: 1 - wandb_project: - value: "engiopt" diff --git a/engiopt/generators/cgan_cnn_3d/README.md b/engiopt/generators/cgan_cnn_3d/README.md index 1e5d7a5c..ec208d3d 100644 --- a/engiopt/generators/cgan_cnn_3d/README.md +++ b/engiopt/generators/cgan_cnn_3d/README.md @@ -23,31 +23,21 @@ This implementation extends the classic 2D cDCGAN architecture to 3D volumetric ### Basic Training ```bash # Train with default parameters -python cgan_cnn_3d.py --problem_id heatconduction3d --n_epochs 300 +python engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py --problem-id heatconduction3d --n-epochs 300 -# Save model for later evaluation -python cgan_cnn_3d.py --save_model --track +# Save the checkpoint for later evaluation +python engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py --save-model --track ``` ### Evaluation ```bash -# Evaluate trained model using EngiBench metrics -python evaluate_cgan_cnn_3d.py --seed 1 --n_samples 50 +# Score the trained checkpoint under the problem's frozen eval spec +python -m engiopt.evaluate --problem-id heatconduction3d --generators cgan_cnn_3d --seeds 1 -# Batch evaluation across multiple seeds -for seed in {1..5}; do - python evaluate_cgan_cnn_3d.py --seed $seed --n_samples 100 -done +# Several seeds at once, including the simulator-backed metrics +python -m engiopt.evaluate --problem-id heatconduction3d --generators cgan_cnn_3d --seeds 1 2 3 --include-expensive ``` -### Hyperparameter Sweep -```bash -# Run Bayesian optimization sweep -wandb sweep sweep_cgan_cnn_3d.yaml -wandb agent -``` -!! Change the path of cgan_cnn_3d.py to your project path - ## Key Parameters ### Core Training diff --git a/engiopt/generators/cgan_cnn_3d/adapter.py b/engiopt/generators/cgan_cnn_3d/adapter.py index a8e5dc86..1cf3ed90 100644 --- a/engiopt/generators/cgan_cnn_3d/adapter.py +++ b/engiopt/generators/cgan_cnn_3d/adapter.py @@ -6,6 +6,7 @@ import torch as th +from engiopt.core import condition_keys_for from engiopt.core import ConditionBatch from engiopt.core import Generator from engiopt.generators.cgan_cnn_3d.cgan_cnn_3d import Generator3D @@ -51,7 +52,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device config = resolved.run_config net = Generator3D( latent_dim=config["latent_dim"], - n_conds=len(problem.conditions_keys), + n_conds=len(condition_keys_for(problem, resolved)), design_shape=problem.design_space.shape, ) net.load_state_dict(th.load(resolved.files["generator_3d.pth"], map_location=device)[cls.primary_state_key]) diff --git a/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py b/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py index 40a7848e..984beaf7 100644 --- a/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py +++ b/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py @@ -28,6 +28,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys import wandb @@ -369,7 +370,9 @@ 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_keys + # The scalar conditions the generator is conditioned on; array-valued and + # solver-only entries of `problem.conditions_keys` cannot enter a dense tensor. + conditions = condition_keys(problem) n_conds = len(conditions) condition_names = [cond[0] for cond in conditions] @@ -419,7 +422,7 @@ def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, d # Extract 3D designs and conditions designs_3d = training_ds["optimal_design"][:] # Should be (N, D, H, W) - condition_tensors = [training_ds[key][:] for key in problem.conditions_keys] + condition_tensors = [training_ds[key][:] for key in conditions] training_ds = th.utils.data.TensorDataset(designs_3d, *condition_tensors) dataloader = th.utils.data.DataLoader( @@ -698,6 +701,7 @@ def sample_3d_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: checkpoint_files={"generator_3d.pth": "generator_3d.pth", "discriminator_3d.pth": "discriminator_3d.pth"}, run_config=vars(args), **checkpoint_identity(args), + condition_keys=conditions, primary_files=["generator_3d.pth"], ) diff --git a/engiopt/generators/cgan_cnn_3d/sweep_cgan_cnn_3d.yaml b/engiopt/generators/cgan_cnn_3d/sweep_cgan_cnn_3d.yaml deleted file mode 100644 index 7423a5c8..00000000 --- a/engiopt/generators/cgan_cnn_3d/sweep_cgan_cnn_3d.yaml +++ /dev/null @@ -1,40 +0,0 @@ -method: bayes -metric: - name: log_dpp - goal: maximize -early_terminate: - type: hyperband - max_iter: 500 - eta: 3 -parameters: - batch_size: - distribution: int_uniform - max: 16 - min: 4 - latent_dim: - distribution: int_uniform - max: 128 - min: 32 - lr_disc: - distribution: log_uniform_values - min: 0.00005 - max: 0.001 - lr_gen: - distribution: log_uniform_values - min: 0.0001 - max: 0.003 - n_epochs: - distribution: int_uniform - min: 150 - max: 500 - gen_iters: - distribution: categorical - values: - - 1 - - 2 - discrim_iters: - distribution: categorical - values: - - 1 - - 2 -program: /cgan_cnn_3d/cgan_cnn_3d.py diff --git a/engiopt/generators/cgan_vae/README.md b/engiopt/generators/cgan_vae/README.md index c9cffdd5..84037d38 100644 --- a/engiopt/generators/cgan_vae/README.md +++ b/engiopt/generators/cgan_vae/README.md @@ -24,31 +24,21 @@ This implementation combines VAE and GAN architectures to generate 3D engineerin ### Basic Training ```bash # Train with default parameters -python cgan_vae.py --problem_id heatconduction3d --n_epochs 300 +python engiopt/generators/cgan_vae/cgan_vae.py --problem-id heatconduction3d --n-epochs 300 -# Save model for later evaluation -python cgan_vae.py --save_model --track +# Save the checkpoint for later evaluation +python engiopt/generators/cgan_vae/cgan_vae.py --save-model --track ``` ### Evaluation ```bash -# Evaluate trained model using EngiBench metrics -python evaluate_cgan_vae.py --seed 1 --n_samples 50 +# Score the trained checkpoint under the problem's frozen eval spec +python -m engiopt.evaluate --problem-id heatconduction3d --generators cgan_vae --seeds 1 -# Batch evaluation across multiple seeds -for seed in {1..5}; do - python evaluate_cgan_vae.py --seed $seed --n_samples 100 -done +# Several seeds at once, including the simulator-backed metrics +python -m engiopt.evaluate --problem-id heatconduction3d --generators cgan_vae --seeds 1 2 3 --include-expensive ``` -### Hyperparameter Sweep -```bash -# Run Bayesian optimization sweep -wandb sweep sweep_cgan_vae.yaml -wandb agent -``` -!! Change the path of cgan_vae.py to your project path - ## Key Parameters ### Core Training diff --git a/engiopt/generators/cgan_vae/adapter.py b/engiopt/generators/cgan_vae/adapter.py index a50dba7c..8323e900 100644 --- a/engiopt/generators/cgan_vae/adapter.py +++ b/engiopt/generators/cgan_vae/adapter.py @@ -6,6 +6,7 @@ import torch as th +from engiopt.core import condition_keys_for from engiopt.core import ConditionBatch from engiopt.core import Generator from engiopt.generators.cgan_cnn_3d.adapter import center_crop_3d @@ -38,7 +39,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device config = resolved.run_config net = Generator3D( latent_dim=config["latent_dim"], - n_conds=len(problem.conditions_keys), + n_conds=len(condition_keys_for(problem, resolved)), design_shape=problem.design_space.shape, ) checkpoint = th.load(resolved.files["multiview_3d_vaegan.pth"], map_location=device) diff --git a/engiopt/generators/cgan_vae/cgan_vae.py b/engiopt/generators/cgan_vae/cgan_vae.py index 49c92955..bf606894 100644 --- a/engiopt/generators/cgan_vae/cgan_vae.py +++ b/engiopt/generators/cgan_vae/cgan_vae.py @@ -28,6 +28,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys import wandb @@ -462,7 +463,9 @@ 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_keys + # The scalar conditions the generator is conditioned on; array-valued and + # solver-only entries of `problem.conditions_keys` cannot enter a dense tensor. + conditions = condition_keys(problem) n_conds = len(conditions) condition_names = [cond[0] for cond in conditions] @@ -520,7 +523,7 @@ def compute_gradient_penalty(discriminator, real_samples, fake_samples, conds, d # Extract 3d designs and conditions designs_3d = training_ds["optimal_design"][:] - condition_tensors = [training_ds[key][:] for key in problem.conditions_keys] + condition_tensors = [training_ds[key][:] for key in conditions] training_ds = th.utils.data.TensorDataset(designs_3d, *condition_tensors) dataloader = th.utils.data.DataLoader( @@ -825,6 +828,7 @@ def sample_3d_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: checkpoint_files={"multiview_3d_vaegan.pth": "multiview_3d_vaegan.pth"}, run_config=vars(args), **checkpoint_identity(args), + condition_keys=conditions, primary_files=["multiview_3d_vaegan.pth"], ) diff --git a/engiopt/generators/cgan_vae/sample_script.ipynb b/engiopt/generators/cgan_vae/sample_script.ipynb deleted file mode 100644 index 505d6823..00000000 --- a/engiopt/generators/cgan_vae/sample_script.ipynb +++ /dev/null @@ -1,16 +0,0 @@ -{ - "cells": [], - "metadata": { - "kernelspec": { - "display_name": "engibench", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.11.6" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/engiopt/generators/cgan_vae/sweep_cgan_vae.yaml b/engiopt/generators/cgan_vae/sweep_cgan_vae.yaml deleted file mode 100644 index 057644d8..00000000 --- a/engiopt/generators/cgan_vae/sweep_cgan_vae.yaml +++ /dev/null @@ -1,56 +0,0 @@ -method: bayes -metric: - name: log_dpp - goal: maximize -early_terminate: - type: hyperband - max_iter: 500 - eta: 3 -parameters: - batch_size: - distribution: int_uniform - min: 4 - max: 16 - latent_dim: - distribution: int_uniform - min: 32 - max: 128 - lr_disc: - distribution: log_uniform_values - min: 0.00005 - max: 0.001 - lr_gen: - distribution: log_uniform_values - min: 0.0001 - max: 0.003 - lr_enc: - distribution: log_uniform_values - min: 0.0001 - max: 0.003 - n_epochs: - distribution: int_uniform - min: 150 - max: 500 - kl_weight: - distribution: log_uniform_values - min: 0.001 - max: 0.1 - recon_weight: - distribution: log_uniform_values - min: 0.5 - max: 2.0 - n_slices: - distribution: int_uniform - min: 3 - max: 9 - gen_iters: - distribution: categorical - values: - - 1 - - 2 - discrim_iters: - distribution: categorical - values: - - 1 - - 2 -program: /cgan_vae/cgan_vae.py diff --git a/engiopt/generators/diffusion_2d_cond/adapter.py b/engiopt/generators/diffusion_2d_cond/adapter.py index d5da0b83..ac005c0e 100644 --- a/engiopt/generators/diffusion_2d_cond/adapter.py +++ b/engiopt/generators/diffusion_2d_cond/adapter.py @@ -7,6 +7,7 @@ from diffusers import UNet2DConditionModel import torch as th +from engiopt.core import condition_keys_for from engiopt.core import ConditionBatch from engiopt.core import Generator from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import beta_schedule @@ -67,7 +68,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device up_block_types=("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"), layers_per_block=config["layers_per_block"], transformer_layers_per_block=1, - encoder_hid_dim=len(problem.conditions_keys), + encoder_hid_dim=len(condition_keys_for(problem, resolved)), only_cross_attention=True, ).to(device) net.load_state_dict(checkpoint[cls.primary_state_key]) diff --git a/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py b/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py index 57e1f115..bd5a755b 100644 --- a/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py +++ b/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py @@ -22,6 +22,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys import wandb if TYPE_CHECKING: @@ -323,7 +324,10 @@ def sample_timestep( # Loss function adversarial_loss: th.nn.Module = th.nn.MSELoss() - encoder_hid_dim = len(problem.conditions_keys) + # The scalar conditions the generator is conditioned on; array-valued and + # solver-only entries of `problem.conditions_keys` cannot enter a dense tensor. + cond_keys = condition_keys(problem) + encoder_hid_dim = len(cond_keys) # Initialize UNet from Huggingface model = UNet2DConditionModel( sample_size=design_shape, @@ -349,10 +353,8 @@ def sample_timestep( filtered_ds[i] = training_ds[i]["optimal_design"][:].reshape(1, design_shape[0], design_shape[1]) filtered_ds_min, filtered_ds_max = get_design_bounds(problem, filtered_ds, device) filtered_ds_norm = normalize_designs_to_diffusion_range(filtered_ds, filtered_ds_min, filtered_ds_max) - 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_keys) + 1]) + training_ds = th.utils.data.TensorDataset(filtered_ds_norm.flatten(1), *[training_ds[key][:] for key in cond_keys]) + cond_tensors = th.stack(training_ds.tensors[1 : len(cond_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))) @@ -472,7 +474,7 @@ def sample_designs(model: UNet2DConditionModel, n_designs: int = 25) -> tuple[th img = design.cpu().numpy() # Extract x and y coordinates dc = hidden_states[j, 0, :].cpu() axes[j].imshow(img[0]) # image plot - title = [(problem.conditions_keys[i], f"{dc[i]:.2f}") for i in range(len(problem.conditions_keys))] + title = [(cond_keys[i], f"{dc[i]:.2f}") for i in range(len(cond_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 @@ -512,6 +514,7 @@ def sample_designs(model: UNet2DConditionModel, n_designs: int = 25) -> tuple[th checkpoint_files={"model.pth": "model.pth"}, run_config=vars(args), **checkpoint_identity(args), + condition_keys=cond_keys, primary_files=["model.pth"], ) diff --git a/engiopt/generators/gan_1d/gan_1d.py b/engiopt/generators/gan_1d/gan_1d.py index 44858e42..5082443f 100644 --- a/engiopt/generators/gan_1d/gan_1d.py +++ b/engiopt/generators/gan_1d/gan_1d.py @@ -25,6 +25,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys from engiopt.transforms import flatten_dict_factory import wandb @@ -164,7 +165,7 @@ def prepare_data(problem: Problem, device: th.device) -> tuple[th.utils.data.Ten training_ds = th.utils.data.TensorDataset( transform(training_ds["optimal_design"][:]), - *[training_ds[key][:] for key in problem.conditions_keys], + *[training_ds[key][:] for key in condition_keys(problem)], ) # Create design normalizer diff --git a/engiopt/generators/pixel_cnn_pp_2d/adapter.py b/engiopt/generators/pixel_cnn_pp_2d/adapter.py index 6c48de74..16b4fa22 100644 --- a/engiopt/generators/pixel_cnn_pp_2d/adapter.py +++ b/engiopt/generators/pixel_cnn_pp_2d/adapter.py @@ -6,6 +6,7 @@ import torch as th +from engiopt.core import condition_keys_for from engiopt.core import ConditionBatch from engiopt.core import Generator from engiopt.generators.pixel_cnn_pp_2d.pixel_cnn_pp_2d import PixelCNNpp @@ -52,7 +53,7 @@ def build( resnet_nonlinearity=config["resnet_nonlinearity"], dropout_p=config["dropout_p"], input_channels=1, - nr_conditions=len(problem.conditions_keys), + nr_conditions=len(condition_keys_for(problem, resolved)), ) net.load_state_dict(th.load(resolved.files["model.pth"], map_location=device)[cls.primary_state_key]) net.eval().to(device) diff --git a/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py b/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py index 6fa3571a..65a850f2 100644 --- a/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py +++ b/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py @@ -29,6 +29,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys import wandb @@ -711,7 +712,9 @@ def sample_from_discretized_mix_logistic(l: th.Tensor, nr_mix: int) -> th.Tensor design_shape = problem.design_space.shape - conditions = problem.conditions_keys + # The scalar conditions the generator is conditioned on; array-valued and + # solver-only entries of `problem.conditions_keys` cannot enter a dense tensor. + conditions = condition_keys(problem) nr_conditions = len(conditions) # Logging @@ -751,7 +754,7 @@ def sample_from_discretized_mix_logistic(l: th.Tensor, nr_mix: int) -> th.Tensor # Configure data loader training_ds = problem.dataset.with_format("torch", device=device)["train"] - condition_tensors = [training_ds[key][:] for key in problem.conditions_keys] + condition_tensors = [training_ds[key][:] for key in conditions] training_ds = th.utils.data.TensorDataset(training_ds["optimal_design"][:], *condition_tensors) @@ -871,7 +874,7 @@ def sample_designs( dc = desired_conds[j].cpu().squeeze() # Extract design conditions axes[j].imshow(img) # image plot - title = [(problem.conditions_keys[i][0], f"{dc[i]:.2f}") for i in range(nr_conditions)] + title = [(conditions[i], f"{dc[i]:.2f}") for i in range(nr_conditions)] 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 @@ -908,6 +911,7 @@ def sample_designs( checkpoint_files={"model.pth": "model.pth"}, run_config=vars(args), **checkpoint_identity(args), + condition_keys=conditions, primary_files=["model.pth"], ) diff --git a/engiopt/generators/vqgan/vqgan.py b/engiopt/generators/vqgan/vqgan.py index 59166689..6030608e 100644 --- a/engiopt/generators/vqgan/vqgan.py +++ b/engiopt/generators/vqgan/vqgan.py @@ -48,6 +48,7 @@ from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training +from engiopt.transforms import condition_keys from engiopt.transforms import drop_constant from engiopt.transforms import normalize from engiopt.transforms import resize_to @@ -691,7 +692,9 @@ def log_images(self, x: th.Tensor, c: th.Tensor, top_k: int | None = None) -> tu image_channels = training_ds["optimal_upsampled"][:].shape[1] latent_size = args.image_size // (2 ** (len(args.encoder_channels) - 2)) - conditions = problem.conditions_keys + # Only the scalar conditions can be stacked into the condition tensor; + # array-valued and solver-only keys are excluded. + conditions = condition_keys(problem) # Optionally drop condition columns that are constant like overhang_constraint in beams2d if args.drop_constant_conditions: training_ds, conditions = drop_constant(training_ds, conditions) @@ -743,7 +746,7 @@ def log_images(self, x: th.Tensor, c: th.Tensor, top_k: int | None = None) -> tu # Optionally drop condition columns that are constant like overhang_constraint in beams2d if args.drop_constant_conditions: - to_drop = [c for c in problem.conditions_keys if c not in conditions] + to_drop = [c for c in condition_keys(problem) if c not in conditions] if to_drop: val_ds = val_ds.remove_columns(to_drop) @@ -995,6 +998,7 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: checkpoint_files={"cvqgan.pth": "cvqgan.pth"}, run_config=vars(args), **checkpoint_identity(args), + condition_keys=conditions, metadata={"stage": "cvqgan"}, primary_files=["cvqgan.pth"], ) @@ -1115,6 +1119,7 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: }, run_config=vars(args), **checkpoint_identity(args), + condition_keys=conditions, metadata={"stage": "vqgan"}, primary_files=["vqgan.pth"], ) @@ -1260,6 +1265,7 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: checkpoint_files=checkpoint_files, run_config=vars(args), **checkpoint_identity(args), + condition_keys=conditions, metadata={"stage": "transformer"}, primary_files=["transformer.pth"], ) diff --git a/engiopt/specs/beams2d/v1.json b/engiopt/specs/beams2d/v1.json index 7a1b666e..b233c964 100644 --- a/engiopt/specs/beams2d/v1.json +++ b/engiopt/specs/beams2d/v1.json @@ -1,6 +1,8 @@ { - "condition_digest": "67f84c4e50c355cf", + "condition_digest": "f65e2c02551e7e08", "condition_seed": 1, + "dataset_id": "IDEALLab/beams_2d_50_100_v0", + "dataset_revision": "ccb64d4f09cf66a5ed9ba2081438d695b9e2438d", "engibench_version": "0.2.0", "metrics": [ "mmd", @@ -11,9 +13,12 @@ "fog" ], "n_samples": 50, - "notes": "", + "notes": "Feasibility = EngiBench constraints plus the volfrac budget.", + "objective_weight_condition": null, + "objective_weights": null, "problem_id": "beams2d", "sigma": 10.0, "version": "v1", - "volfrac_tol": 0.01 + "volfrac_tol": 0.01, + "volume_condition": "volfrac" } diff --git a/engiopt/specs/heatconduction2d/v1.json b/engiopt/specs/heatconduction2d/v1.json index 60595964..2f182309 100644 --- a/engiopt/specs/heatconduction2d/v1.json +++ b/engiopt/specs/heatconduction2d/v1.json @@ -1,6 +1,8 @@ { - "condition_digest": "50961cce837117ed", + "condition_digest": "40030d1b2482e003", "condition_seed": 1, + "dataset_id": "IDEALLab/heat_conduction_2d_v0", + "dataset_revision": "9f07c1e70d244f783e34a71c96db621173871dc1", "engibench_version": "0.2.0", "metrics": [ "mmd", @@ -11,9 +13,12 @@ "fog" ], "n_samples": 50, - "notes": "", + "notes": "Feasibility = EngiBench constraints plus the volume budget.", + "objective_weight_condition": null, + "objective_weights": null, "problem_id": "heatconduction2d", "sigma": 10.0, "version": "v1", - "volfrac_tol": 0.01 + "volfrac_tol": 0.01, + "volume_condition": "volume" } diff --git a/engiopt/specs/photonics2d/v1.json b/engiopt/specs/photonics2d/v1.json index 8eaefcde..dfaf548d 100644 --- a/engiopt/specs/photonics2d/v1.json +++ b/engiopt/specs/photonics2d/v1.json @@ -1,6 +1,8 @@ { - "condition_digest": "ade3d18084b31b5e", + "condition_digest": "6b8330bee5a8718c", "condition_seed": 1, + "dataset_id": "IDEALLab/photonics_2d_120_120_v1", + "dataset_revision": "be811ff69d67cd8c79148e29f2a68d293ec00955", "engibench_version": "0.2.0", "metrics": [ "mmd", @@ -11,9 +13,12 @@ "fog" ], "n_samples": 50, - "notes": "", + "notes": "No volume budget: feasibility is EngiBench's declared constraints alone.", + "objective_weight_condition": null, + "objective_weights": null, "problem_id": "photonics2d", "sigma": 10.0, "version": "v1", - "volfrac_tol": 0.01 + "volfrac_tol": 0.01, + "volume_condition": null } diff --git a/engiopt/specs/thermoelastic2d/v1.json b/engiopt/specs/thermoelastic2d/v1.json index 291704db..892239b6 100644 --- a/engiopt/specs/thermoelastic2d/v1.json +++ b/engiopt/specs/thermoelastic2d/v1.json @@ -1,6 +1,8 @@ { - "condition_digest": "3a1d2ae6c45add7f", + "condition_digest": "818d93c2a2d5ac25", "condition_seed": 1, + "dataset_id": "IDEALLab/thermoelastic_2d_v1", + "dataset_revision": "7d5b4a694bd1a8f3f25fd745ac9f669764e29ea5", "engibench_version": "0.2.0", "metrics": [ "mmd", @@ -11,11 +13,12 @@ "fog" ], "n_samples": 50, - "notes": "", + "notes": "Objectives combined per-sample by the `weight` condition. Feasibility = EngiBench constraints plus the volume_fraction_target budget.", "objective_weight_condition": "weight", "objective_weights": null, "problem_id": "thermoelastic2d", "sigma": 10.0, "version": "v1", - "volfrac_tol": 0.01 + "volfrac_tol": 0.01, + "volume_condition": "volume_fraction_target" } diff --git a/engiopt/transforms.py b/engiopt/transforms.py index 7f6de89d..189d21fe 100644 --- a/engiopt/transforms.py +++ b/engiopt/transforms.py @@ -42,6 +42,26 @@ def get_scalar_condition_keys(problem: Problem, dataset: Dataset, *, drop_consta return scalar_keys +def condition_keys(problem: Problem, split: str = "train") -> list[str]: + """The scalar condition columns a generator is conditioned on, in tensor order. + + This is the one definition of "how many conditions does this problem have" + that training, checkpoint loading, and sampling all share. Using + `len(problem.conditions_keys)` instead builds a network for columns that + never reach it: thermoelastic2d declares 7 conditions, of which 4 are 65x65 + boundary matrices, so a generator sized for 7 fails on a 3-column tensor. + + Args: + problem: An EngiBench problem instance. + split: Dataset split to inspect; the schema is the same in all of them. + + Returns: + Condition names, in `conditions_keys` order. + """ + dataset = problem.dataset + return get_scalar_condition_keys(problem, dataset[split] if split in dataset else next(iter(dataset.values()))) + + def get_image_condition_keys(problem: Problem, dataset: Dataset) -> list[str]: """Return the array-valued condition keys present in the dataset. diff --git a/engiopt/utils/all_generators.py b/engiopt/utils/all_generators.py index 00554274..fec02949 100644 --- a/engiopt/utils/all_generators.py +++ b/engiopt/utils/all_generators.py @@ -78,22 +78,3 @@ def generators_for(problem: Problem) -> dict[str, type[Generator]]: """Return the generators whose `design_kinds` cover this problem's design space.""" kind = design_kind_of(problem) return {name: g for name, g in BUILTIN_GENERATORS.items() if kind in g.design_kinds} - - -def problem_id_of(problem: Problem | type[Problem]) -> str: - """Reverse-look-up a problem's registry key. - - EngiBench problems do not carry their own id, but checkpoint paths and - leaderboard rows are keyed by it. - - Raises: - ValueError: If the problem is not a built-in EngiBench problem. - """ - from engibench.utils.all_problems import BUILTIN_PROBLEMS - - problem_cls = problem if isinstance(problem, type) else type(problem) - for name, candidate in BUILTIN_PROBLEMS.items(): - if candidate is problem_cls: - return name - msg = f"{problem_cls.__name__} is not a built-in EngiBench problem; pass problem_id explicitly." - raise ValueError(msg) diff --git a/pyproject.toml b/pyproject.toml index 64c87e15..b79812ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,11 @@ engiopt = [ ############################################# Linters and Test tools ####################################################### +[tool.pytest.ini_options] +markers = [ + "network: needs the HuggingFace Hub (downloads a dataset or checkpoint). Skip with -m 'not network'.", +] + [tool.ruff] # Exclude a variety of commonly ignored directories. exclude = [ diff --git a/tests/conftest.py b/tests/conftest.py index 39000ccf..91127ea3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -22,6 +22,14 @@ class FakeOptiStep: obj_values: Any +@dataclass +class FakeViolations: + """Stand-in for `engibench.constraint.Violations`.""" + + violations: list[str] + n_constraints: int = 1 + + class FakeProblem: """A minimal stand-in exposing the parts of `Problem` the evaluator touches. @@ -36,6 +44,7 @@ def __init__( *, n_objectives: int = 1, directions: tuple[str, ...] | None = None, + infeasible: bool = False, ): self.design_space = spaces.Box(low=0.0, high=1.0, shape=shape, dtype=np.float64) self.conditions_keys = list(conditions_keys) @@ -43,9 +52,28 @@ def __init__( # Mirrors `Problem.objectives`: (name, direction) pairs in simulate order. directions = directions or ("MINIMIZE",) * n_objectives self.objectives = tuple((f"objective_{i}", directions[i]) for i in range(n_objectives)) + self.infeasible = infeasible + """Whether `check_constraints` reports every design as violating.""" self.reset_calls = 0 self.optimize_calls = 0 + @property + def dataset(self) -> dict[str, Any]: + """One split whose columns are the declared conditions, all scalar.""" + rows = 3 + return {"train": FakeDataset({key: [0.5] * rows for key in self.conditions_keys})} + + def check_constraints(self, design: Any, config: dict[str, Any]) -> FakeViolations: + """Mirror `Problem.check_constraints`: declared constraints plus design-space membership. + + Like the real one, `design` is passed to `contains` exactly as given -- + dtype included, since that is what `Box.contains` is strict about. + """ + violations = ["forced"] if self.infeasible else [] + if not self.design_space.contains(design): + violations.append("design not in design_space") + return FakeViolations(violations=violations) + def reset(self, seed: int | None = None) -> None: """Record that the evaluator reset us between solver calls.""" self.reset_calls += 1 diff --git a/tests/test_checkpoint_packages.py b/tests/test_checkpoint_packages.py new file mode 100644 index 00000000..5458e704 --- /dev/null +++ b/tests/test_checkpoint_packages.py @@ -0,0 +1,263 @@ +"""Tests for reading and writing checkpoint packages. + +A package is the weight files plus `run_config.json` and `metadata.json`. What +matters here is that a package can be loaded from a plain directory (no Hub +account required to try a model), that its metadata describes the package it is +actually in, and that a loaded package identifies the exact weights it holds. +""" + +from __future__ import annotations + +import inspect +import json +from pathlib import Path +from typing import Any, ClassVar + +import pytest +import torch as th + +from engiopt import checkpoint_store +from engiopt.core import condition_keys_for +from engiopt.core import Generator + + +def _write_package(root: Path, *, weights: bytes = b"weights", metadata: dict[str, Any] | None = None) -> Path: + root.mkdir(parents=True, exist_ok=True) + (root / "generator.pth").write_bytes(weights) + (root / "run_config.json").write_text(json.dumps({"latent_dim": 32})) + (root / "metadata.json").write_text(json.dumps(metadata or {"algo": "demo", "seed": 1})) + return root + + +# ---------------------------------------------------------------------- +# Loading from a local directory +# ---------------------------------------------------------------------- + + +def test_local_package_loads_without_the_hub(tmp_path: Path) -> None: + """`--model-source local` has to be reachable, not just advertised.""" + package = _write_package(tmp_path / "pkg") + + resolved = checkpoint_store.resolve_named_checkpoint( + model_source="local", + problem_id="beams2d", + algo="demo", + seed=1, + hf_entity="none", + hf_repo_prefix="none", + required_files=["generator.pth"], + local_model_dir=str(package), + ) + + assert resolved.source == "local" + assert resolved.run_config == {"latent_dim": 32} + assert Path(resolved.files["generator.pth"]).read_bytes() == b"weights" + + +def test_from_pretrained_passes_the_local_directory_through() -> None: + """The public API must accept the directory `resolve_checkpoint` needs.""" + assert "local_model_dir" in inspect.signature(Generator.from_pretrained).parameters + assert "local_model_dir" in inspect.signature(Generator.resolve_checkpoint).parameters + + +def test_the_cli_exposes_the_local_directory() -> None: + """An option nobody can point at a directory is not a usable source.""" + from engiopt.evaluate import Args + + assert "local_model_dir" in set(Args.__dataclass_fields__) + + +def test_a_missing_local_directory_says_so(tmp_path: Path) -> None: + """A clear error beats a confusing failure deeper in loading.""" + with pytest.raises(FileNotFoundError, match="Local checkpoint directory not found"): + checkpoint_store.resolve_named_checkpoint( + model_source="local", + problem_id="beams2d", + algo="demo", + seed=1, + hf_entity="none", + hf_repo_prefix="none", + required_files=["generator.pth"], + local_model_dir=str(tmp_path / "absent"), + ) + + +# ---------------------------------------------------------------------- +# Identifying the exact weights +# ---------------------------------------------------------------------- + + +def test_package_records_a_content_hash(tmp_path: Path) -> None: + """Config and seed name a *configuration*; only a content hash names the weights.""" + first = checkpoint_store._load_package_from_directory( + root_dir=str(_write_package(tmp_path / "a", weights=b"one")), required_files=["generator.pth"], source="local" + ) + second = checkpoint_store._load_package_from_directory( + root_dir=str(_write_package(tmp_path / "b", weights=b"two")), required_files=["generator.pth"], source="local" + ) + same = checkpoint_store._load_package_from_directory( + root_dir=str(_write_package(tmp_path / "c", weights=b"one")), required_files=["generator.pth"], source="local" + ) + + assert first.content_hash != second.content_hash + assert first.content_hash == same.content_hash + + +def test_retrained_weights_get_a_new_hash_under_the_same_name(tmp_path: Path) -> None: + """Re-training the same config and seed replaces a leaderboard row; the hash says so.""" + package = _write_package(tmp_path / "pkg", weights=b"first training") + before = checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth"], source="local" + ) + (package / "generator.pth").write_bytes(b"second training") + after = checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth"], source="local" + ) + assert before.content_hash != after.content_hash + + +def test_evaluation_rows_carry_the_checkpoint_identity() -> None: + """Provenance has to reach the leaderboard, not just the loader.""" + from engiopt.evaluation.evaluator import PROVENANCE_COLUMNS + + assert {"checkpoint_revision", "checkpoint_hash", "code_version"} <= set(PROVENANCE_COLUMNS) + + +# ---------------------------------------------------------------------- +# Metadata describes the package it sits in +# ---------------------------------------------------------------------- + + +def test_each_uploaded_package_describes_itself(monkeypatch: pytest.MonkeyPatch) -> None: + """A non-default run writes only the config path, so its metadata must name that path. + + Sharing one metadata object made a config-only run's `metadata.json` point + at a canonical package it never uploaded. + """ + uploads: list[dict[str, Any]] = [] + + def fake_upload(**kwargs: Any) -> str: + uploads.append({"path": kwargs["package_path"], "metadata": kwargs["metadata"]}) + return "sha" + + monkeypatch.setattr(checkpoint_store, "_upload_package_to_hf", fake_upload) + monkeypatch.setattr(checkpoint_store.wandb, "run", None) + + info = checkpoint_store.save_checkpoint_package( + checkpoint_backend="hf", + hf_entity="org", + hf_repo_prefix="engiopt", + hf_private=False, + problem_id="beams2d", + algo="demo", + seed=1, + checkpoint_files={}, + run_config={"latent_dim": 64}, + config_fingerprint="abc123", + is_default_config=False, + ) + + assert len(uploads) == 1 + assert uploads[0]["path"] == "beams2d/cfg_abc123/seed_1" + assert uploads[0]["metadata"]["hf_package_path"] == "beams2d/cfg_abc123/seed_1" + assert info["hf_package_path"] is None + + +def test_a_default_run_claims_both_paths_with_matching_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + """Each of the two packages describes itself, not the other.""" + uploads: list[dict[str, Any]] = [] + + def fake_upload(**kwargs: Any) -> str: + uploads.append({"path": kwargs["package_path"], "metadata": kwargs["metadata"]}) + return "sha" + + monkeypatch.setattr(checkpoint_store, "_upload_package_to_hf", fake_upload) + monkeypatch.setattr(checkpoint_store.wandb, "run", None) + + checkpoint_store.save_checkpoint_package( + checkpoint_backend="hf", + hf_entity="org", + hf_repo_prefix="engiopt", + hf_private=False, + problem_id="beams2d", + algo="demo", + seed=1, + checkpoint_files={}, + run_config={}, + config_fingerprint="abc123", + is_default_config=True, + ) + + assert [upload["path"] for upload in uploads] == ["beams2d/cfg_abc123/seed_1", "beams2d/seed_1"] + for upload in uploads: + assert upload["metadata"]["hf_package_path"] == upload["path"] + + +def test_condition_schema_travels_with_the_checkpoint(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """A model is rebuilt for the conditions it saw, not the problem's current list.""" + uploads: list[dict[str, Any]] = [] + + def fake_upload(**kwargs: Any) -> str: + uploads.append(kwargs["metadata"]) + return "sha" + + monkeypatch.setattr(checkpoint_store, "_upload_package_to_hf", fake_upload) + monkeypatch.setattr(checkpoint_store.wandb, "run", None) + + checkpoint_store.save_checkpoint_package( + checkpoint_backend="hf", + hf_entity="org", + hf_repo_prefix="engiopt", + hf_private=False, + problem_id="thermoelastic2d", + algo="demo", + seed=1, + checkpoint_files={}, + run_config={}, + condition_keys=["volume_fraction_target", "rmin", "weight"], + ) + + assert uploads[0]["condition_keys"] == ["volume_fraction_target", "rmin", "weight"] + + +def test_recorded_schema_survives_the_problem_gaining_a_condition(tmp_path: Path) -> None: + """An old checkpoint keeps loading when the problem later adds a condition.""" + package = _write_package(tmp_path / "pkg", metadata={"condition_keys": ["volfrac", "rmin"]}) + resolved = checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth"], source="local" + ) + + class _Problem: + conditions_keys: ClassVar[list[str]] = ["volfrac", "rmin", "brand_new"] + + assert condition_keys_for(_Problem(), resolved) == ("volfrac", "rmin") + + +# ---------------------------------------------------------------------- +# Sampling cost is measured, not dispatched +# ---------------------------------------------------------------------- + + +def test_sample_timing_waits_for_the_device(fake_problem: Any) -> None: + """CUDA and MPS queue work asynchronously; the timer must not stop mid-flight.""" + synchronized: list[bool] = [] + + class _Timed(Generator): + algo_id = "timed" + conditional = False + + @classmethod + def build(cls, *_args: Any, **_kwargs: Any) -> _Timed: + raise NotImplementedError + + def _synchronize_device(self) -> None: + synchronized.append(True) + + def _sample(self, conditions: Any, n: int) -> th.Tensor: + return th.zeros((n, *self.design_shape)) + + generator = _Timed(problem=fake_problem, problem_id="p", seed=1, device=th.device("cpu")) + generator.sample(None, n=2) + + assert synchronized == [True] + assert generator.last_sample_seconds is not None diff --git a/tests/test_condition_schema.py b/tests/test_condition_schema.py new file mode 100644 index 00000000..40ad4151 --- /dev/null +++ b/tests/test_condition_schema.py @@ -0,0 +1,96 @@ +"""Tests for how many conditions a generator is built and sampled for. + +`problem.conditions_keys` is the full contract and is broader than what a dense +condition tensor can carry: thermoelastic2d declares seven conditions, four of +which are 65x65 boundary matrices. A network sized from the full list therefore +expects seven columns and receives three. These tests pin the one definition +that training, loading, and sampling all share. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest +import torch as th + +from engiopt.checkpoint_store import ResolvedCheckpoint +from engiopt.core import condition_keys_for +from engiopt.core import ConditionBatch +from engiopt.core import Generator + +# thermoelastic2d's real schema: three scalars among seven declared conditions. +THERMOELASTIC_KEYS = ( + "fixed_elements", + "force_elements_x", + "force_elements_y", + "heatsink_elements", + "volume_fraction_target", + "rmin", + "weight", +) +THERMOELASTIC_SCALARS = ("volume_fraction_target", "rmin", "weight") + + +class _EchoGenerator(Generator): + """A generator that returns its condition tensor, so shape errors surface.""" + + algo_id = "echo" + conditional = True + + @classmethod + def build(cls, resolved: ResolvedCheckpoint, problem: Any, device: th.device, **base: Any) -> _EchoGenerator: + """Unused: these tests construct the generator directly.""" + raise NotImplementedError + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + cond = conditions.require_tensor(self.algo_id) + assert cond.shape[1] == self.n_conds, f"got {cond.shape[1]} condition columns, built for {self.n_conds}" + return th.zeros((n, *self.design_shape)) + + +def _resolved(metadata: dict[str, Any]) -> ResolvedCheckpoint: + return ResolvedCheckpoint(source="local", root_dir=".", files={}, run_config={}, metadata=metadata) + + +def _generator(problem: Any, keys: tuple[str, ...]) -> _EchoGenerator: + return _EchoGenerator( + problem=problem, problem_id="thermoelastic2d", seed=1, device=th.device("cpu"), condition_keys=keys + ) + + +def test_checkpoint_condition_schema_is_preferred_over_the_problem(fake_problem: Any) -> None: + """A checkpoint states the columns it was trained on, and that wins.""" + resolved = _resolved({"condition_keys": list(THERMOELASTIC_SCALARS)}) + assert condition_keys_for(fake_problem, resolved) == THERMOELASTIC_SCALARS + + +def test_matrix_valued_conditions_do_not_count_toward_n_conds(fake_problem: Any) -> None: + """The reviewer's case: 7 declared conditions, 3 columns in the tensor.""" + fake_problem.conditions_keys = list(THERMOELASTIC_KEYS) + generator = _generator(fake_problem, THERMOELASTIC_SCALARS) + assert generator.n_conds == len(THERMOELASTIC_SCALARS) + assert generator.n_conds != len(fake_problem.conditions_keys) + + +def test_sampling_succeeds_with_the_scalar_condition_schema(fake_problem: Any) -> None: + """A model built for the scalar columns samples from the scalar columns.""" + fake_problem.conditions_keys = list(THERMOELASTIC_KEYS) + generator = _generator(fake_problem, THERMOELASTIC_SCALARS) + batch = ConditionBatch(tensor=th.zeros((4, 3)), keys=THERMOELASTIC_SCALARS) + assert generator.sample(batch, n=4).shape == (4, *fake_problem.design_space.shape) + + +def test_conditions_from_a_different_schema_are_refused(fake_problem: Any) -> None: + """Silently accepting them conditions the design on the wrong numbers.""" + generator = _generator(fake_problem, THERMOELASTIC_SCALARS) + mismatched = ConditionBatch(tensor=th.zeros((4, 2)), keys=("volfrac", "rmin")) + with pytest.raises(ValueError, match="was trained on conditions"): + generator.sample(mismatched, n=4) + + +def test_unlabelled_conditions_are_accepted(fake_problem: Any) -> None: + """A bare tensor carries no column names, so there is nothing to contradict.""" + generator = _generator(fake_problem, THERMOELASTIC_SCALARS) + assert generator.sample(np.zeros((4, 3)), n=4).shape == (4, *fake_problem.design_space.shape) diff --git a/tests/test_conditions_and_gaps.py b/tests/test_conditions_and_gaps.py index 312333d5..a0012891 100644 --- a/tests/test_conditions_and_gaps.py +++ b/tests/test_conditions_and_gaps.py @@ -8,6 +8,7 @@ from typing import Any +from gymnasium import spaces import numpy as np import pytest @@ -15,6 +16,7 @@ from engiopt.evaluation.context import MultiObjectiveScalarizationError from engiopt.transforms import get_image_condition_keys from engiopt.transforms import get_scalar_condition_keys +from tests.conftest import FakeViolations class _Problem: @@ -268,3 +270,79 @@ def test_signs_are_per_objective_for_mixed_directions() -> None: assert list(ctx.objective_signs) == [1.0, -1.0] # Both objectives beat the reference by 1.0, so both should read as -0.5. assert ctx.scalarize_gap(np.array([-1.0, 1.0]), 0) == pytest.approx(-1.0) + + +# ---------------------------------------------------------------------- +# Feasibility +# ---------------------------------------------------------------------- + + +def _feasibility_context(problem: Any, design_value: float, **kwargs: Any) -> EvaluationContext: + designs = np.full((1, *problem.design_space.shape), design_value) + return EvaluationContext( + problem=problem, + problem_id="fake", + gen_designs=designs, + ref_designs=designs.copy(), + **kwargs, + ) + + +def test_feasibility_falls_back_to_the_problems_own_constraints(fake_problem: Any) -> None: + """Every problem has `check_constraints`, so `viol` is defined even without a volume budget. + + photonics2d has no volume-fraction condition; before this it returned NaN. + """ + ctx = _feasibility_context(fake_problem, 0.5) + assert ctx.is_infeasible(np.full(fake_problem.design_space.shape, 0.5), {"lambda1": 1.5}) is False + + fake_problem.infeasible = True + assert ctx.is_infeasible(np.full(fake_problem.design_space.shape, 0.5), {"lambda1": 1.5}) is True + + +def test_volume_budget_is_checked_when_the_spec_names_one(fake_problem: Any) -> None: + """Missing the declared volume target is a violation of the design's brief.""" + ctx = _feasibility_context(fake_problem, 0.5, volume_condition="volfrac", volfrac_tol=0.01) + design = np.full(fake_problem.design_space.shape, 0.5) + assert ctx.is_infeasible(design, {"volfrac": 0.5}) is False + assert ctx.is_infeasible(design, {"volfrac": 0.9}) is True + + +def test_a_zero_volume_target_is_still_a_target(fake_problem: Any) -> None: + """`conditions.get(a) or conditions.get(b)` silently skipped a legitimate 0.0.""" + ctx = _feasibility_context(fake_problem, 0.5, volume_condition="volfrac", volfrac_tol=0.01) + assert ctx.is_infeasible(np.full(fake_problem.design_space.shape, 0.5), {"volfrac": 0.0}) is True + + +def test_an_unnamed_volume_condition_is_not_guessed(fake_problem: Any) -> None: + """Without a declared `volume_condition`, only the problem's constraints apply.""" + ctx = _feasibility_context(fake_problem, 0.5) + assert ctx.is_infeasible(np.full(fake_problem.design_space.shape, 0.5), {"volfrac": 0.9}) is False + + +def test_a_design_is_judged_on_its_values_not_its_dtype(fake_problem: Any) -> None: + """`Box.contains` rejects an uncastable dtype outright. + + thermoelastic2d's space is float32 while its designs are float64, so without + the cast every design -- including the dataset's own optima -- is infeasible. + """ + fake_problem.design_space = spaces.Box(low=0.0, high=1.0, shape=(8, 10), dtype=np.float32) + ctx = _feasibility_context(fake_problem, 0.5) + assert ctx.is_infeasible(np.full(fake_problem.design_space.shape, 0.5, dtype=np.float64), {}) is False + assert ctx.is_infeasible(np.full(fake_problem.design_space.shape, 5.0, dtype=np.float64), {}) is True + + +def test_array_valued_conditions_survive_the_constraint_check(fake_problem: Any) -> None: + """A HuggingFace dataset hands matrices back as nested lists, which bound checks cannot compare.""" + seen: dict[str, Any] = {} + + def record(design: Any, config: dict[str, Any]) -> Any: + seen.update(config) + return FakeViolations(violations=[]) + + fake_problem.check_constraints = record + ctx = _feasibility_context(fake_problem, 0.5) + ctx.is_infeasible(np.full(fake_problem.design_space.shape, 0.5), {"fixed_elements": [[0, 1], [1, 0]], "rmin": 2.0}) + + assert isinstance(seen["fixed_elements"], np.ndarray) + assert seen["rmin"] == 2.0 diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py index 90abe46f..4faa10f6 100644 --- a/tests/test_evaluation.py +++ b/tests/test_evaluation.py @@ -196,8 +196,50 @@ def test_rank_aggregates_seeds_with_the_median_by_default() -> None: def test_disagreement_surfaces_conflicting_rankings() -> None: """The two metrics crown different winners; the view must show both.""" table = disagreement(_board(), ["mmd", "dpp"]) - assert table.loc[("p", "b"), "mmd"] == 1 - assert table.loc[("p", "a"), "dpp"] == 1 + assert table.loc[("p", "b", "cfg_b", "v1"), "mmd"] == 1 + assert table.loc[("p", "a", "cfg_a", "v1"), "dpp"] == 1 + + +def _two_configs_board() -> pd.DataFrame: + """One algorithm, two hyperparameter settings, both evaluated at seed 1.""" + return pd.DataFrame( + [ + {"problem_id": "p", "algo_id": "a", "config_fingerprint": "cfg_a", "seed": 1, "spec_version": "v1", + "mmd": 0.1}, + {"problem_id": "p", "algo_id": "a", "config_fingerprint": "cfg_b", "seed": 1, "spec_version": "v1", + "mmd": 0.9}, + ] + ) # fmt: skip + + +def test_rank_keeps_hyperparameter_configurations_apart() -> None: + """Two configs are two entries, not one averaged score over a phantom two seeds.""" + ranked = rank(_two_configs_board(), "mmd") + assert len(ranked) == 2 + assert set(ranked["config_fingerprint"]) == {"cfg_a", "cfg_b"} + assert list(ranked["n_seeds"]) == [1, 1] + assert ranked.iloc[0]["mmd_median"] == pytest.approx(0.1) + + +def test_rank_keeps_spec_versions_apart() -> None: + """A score under v1 and one under v2 measure different protocols.""" + frame = pd.DataFrame( + [ + {"problem_id": "p", "algo_id": "a", "config_fingerprint": "cfg", "seed": 1, "spec_version": "v1", + "mmd": 0.1}, + {"problem_id": "p", "algo_id": "a", "config_fingerprint": "cfg", "seed": 1, "spec_version": "v2", + "mmd": 0.9}, + ] + ) # fmt: skip + ranked = rank(frame, "mmd") + assert len(ranked) == 2 + assert set(ranked["spec_version"]) == {"v1", "v2"} + + +def test_disagreement_keeps_hyperparameter_configurations_apart() -> None: + """The disagreement view groups exactly as `rank` does.""" + table = disagreement(_two_configs_board(), ["mmd"]) + assert len(table) == 2 def test_append_rows_supersedes_a_rerun_instead_of_duplicating(tmp_path: Any) -> None: diff --git a/tests/test_leaderboard_publishing.py b/tests/test_leaderboard_publishing.py new file mode 100644 index 00000000..5e82c4a5 --- /dev/null +++ b/tests/test_leaderboard_publishing.py @@ -0,0 +1,169 @@ +"""Tests for publishing to the shared leaderboard. + +Publishing is read-merge-write against a table other people are also writing to, +so the failure modes that matter are the ones that silently *delete* rows: a +download error mistaken for an empty board, and a concurrent publisher whose +commit gets overwritten. +""" + +from __future__ import annotations + +from typing import Any + +from huggingface_hub.errors import EntryNotFoundError +from huggingface_hub.errors import HfHubHTTPError +from huggingface_hub.errors import RepositoryNotFoundError +import pandas as pd +import pytest + +from engiopt.evaluation import leaderboard as lb + + +def _row(algo: str, mmd: float) -> dict[str, Any]: + return { + "problem_id": "p", + "algo_id": algo, + "config_fingerprint": "default", + "seed": 1, + "spec_version": "v1", + "mmd": mmd, + } + + +class _FakeApi: + """Records uploads and can fail the first one, like a lost commit race.""" + + def __init__(self, *, revision: str | None = "sha1", fail_first: bool = False): + self.revision = revision + self.fail_first = fail_first + self.uploads: list[dict[str, Any]] = [] + + def __call__(self, *_args: Any, **_kwargs: Any) -> _FakeApi: + return self + + def create_repo(self, **_kwargs: Any) -> None: + return None + + def repo_info(self, **_kwargs: Any) -> Any: + if self.revision is None: + raise RepositoryNotFoundError("no such repo") + return type("Info", (), {"sha": self.revision})() + + def upload_file(self, **kwargs: Any) -> None: + self.uploads.append({**kwargs, "content": pd.read_csv(kwargs["path_or_fileobj"])}) + if self.fail_first and len(self.uploads) == 1: + # What the Hub returns when `parent_commit` is no longer the head. + self.revision = "sha2" + raise HfHubHTTPError("412 Client Error: parent_commit is out of date") + + +@pytest.fixture +def published(monkeypatch: pytest.MonkeyPatch, tmp_path: Any) -> Any: + """Patch the Hub so `push_to_hub` runs against a local CSV.""" + board_file = tmp_path / "leaderboard.csv" + pd.DataFrame([_row("existing", 0.5)]).to_csv(board_file, index=False) + + def fake_download(**_kwargs: Any) -> str: + return str(board_file) + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download) + return board_file + + +# ---------------------------------------------------------------------- +# Reading +# ---------------------------------------------------------------------- + + +def test_missing_board_reads_as_empty(monkeypatch: pytest.MonkeyPatch) -> None: + """The first model to publish finds no board, and that is not an error.""" + api = _FakeApi(revision=None) + monkeypatch.setattr("huggingface_hub.HfApi", api) + assert lb.load_from_hub("org/board").empty + + +def test_missing_file_in_an_existing_repo_reads_as_empty(monkeypatch: pytest.MonkeyPatch) -> None: + """A repo created but not yet populated is also just an empty board.""" + monkeypatch.setattr("huggingface_hub.HfApi", _FakeApi()) + + def fake_download(**_kwargs: Any) -> str: + raise EntryNotFoundError("no leaderboard.csv") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download) + assert lb.load_from_hub("org/board").empty + + +def test_transient_download_failure_is_not_an_empty_board(monkeypatch: pytest.MonkeyPatch) -> None: + """Swallowing this would let the next push upload only the new rows, deleting the rest.""" + monkeypatch.setattr("huggingface_hub.HfApi", _FakeApi()) + + def fake_download(**_kwargs: Any) -> str: + raise HfHubHTTPError("503 Server Error: Service Unavailable") + + monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download) + with pytest.raises(HfHubHTTPError): + lb.load_from_hub("org/board") + + +def test_auth_failure_is_not_an_empty_board(monkeypatch: pytest.MonkeyPatch) -> None: + """Same reasoning: a private board we cannot read is not a board with no rows.""" + api = _FakeApi() + + def raise_auth(**_kwargs: Any) -> Any: + raise HfHubHTTPError("401 Client Error: Unauthorized") + + api.repo_info = raise_auth # type: ignore[method-assign] + monkeypatch.setattr("huggingface_hub.HfApi", api) + with pytest.raises(HfHubHTTPError): + lb.load_from_hub("org/board") + + +# ---------------------------------------------------------------------- +# Writing +# ---------------------------------------------------------------------- + + +def test_push_merges_instead_of_replacing(monkeypatch: pytest.MonkeyPatch, published: Any) -> None: + """Publishing one model must not require re-running anyone else's evaluation.""" + api = _FakeApi() + monkeypatch.setattr("huggingface_hub.HfApi", api) + + merged = lb.push_to_hub(pd.DataFrame([_row("newcomer", 0.1)]), "org/board") + + assert set(merged["algo_id"]) == {"existing", "newcomer"} + assert set(api.uploads[0]["content"]["algo_id"]) == {"existing", "newcomer"} + + +def test_push_is_conditional_on_the_revision_it_read(monkeypatch: pytest.MonkeyPatch, published: Any) -> None: + """Without `parent_commit`, two concurrent jobs silently drop one another's rows.""" + api = _FakeApi() + monkeypatch.setattr("huggingface_hub.HfApi", api) + + lb.push_to_hub(pd.DataFrame([_row("newcomer", 0.1)]), "org/board") + + assert api.uploads[0]["parent_commit"] == "sha1" + + +def test_push_retries_against_the_newer_board_after_losing_a_race(monkeypatch: pytest.MonkeyPatch, published: Any) -> None: + """A rejected commit means someone else published; re-merge rather than overwrite.""" + api = _FakeApi(fail_first=True) + monkeypatch.setattr("huggingface_hub.HfApi", api) + + merged = lb.push_to_hub(pd.DataFrame([_row("newcomer", 0.1)]), "org/board") + + assert len(api.uploads) == 2 + assert api.uploads[1]["parent_commit"] == "sha2" + assert set(merged["algo_id"]) == {"existing", "newcomer"} + + +def test_push_gives_up_rather_than_forcing_a_write(monkeypatch: pytest.MonkeyPatch, published: Any) -> None: + """Losing every attempt must surface, not resolve into a clobbering write.""" + + class _AlwaysStale(_FakeApi): + def upload_file(self, **kwargs: Any) -> None: + self.uploads.append(kwargs) + raise HfHubHTTPError("412 Client Error: parent_commit is out of date") + + monkeypatch.setattr("huggingface_hub.HfApi", _AlwaysStale()) + with pytest.raises(HfHubHTTPError): + lb.push_to_hub(pd.DataFrame([_row("newcomer", 0.1)]), "org/board", max_attempts=2) diff --git a/tests/test_specs.py b/tests/test_specs.py new file mode 100644 index 00000000..0129f546 --- /dev/null +++ b/tests/test_specs.py @@ -0,0 +1,148 @@ +"""Tests for the committed evaluation specs. + +A spec is the promise that every leaderboard row was produced under identical +conditions. These check the promise is well-formed offline, and -- when the +datasets are reachable -- that every committed spec still reproduces exactly the +conditions it was frozen with. +""" + +from __future__ import annotations + +import dataclasses +import json +from typing import Any + +import numpy as np +import pytest + +from engiopt.evaluation.registry import METRICS +from engiopt.evaluation.spec import _digest +from engiopt.evaluation.spec import EvalSpec +from engiopt.evaluation.spec import SPEC_ROOT + +SPEC_PATHS = sorted(SPEC_ROOT.glob("*/*.json")) +SPEC_IDS = [f"{path.parent.name}/{path.stem}" for path in SPEC_PATHS] + + +class _FakeConditions: + """A stand-in for the sampled-conditions dataset the digest hashes.""" + + def __init__(self, columns: dict[str, Any]): + self._columns = columns + self.column_names = list(columns) + + def __getitem__(self, name: str) -> Any: + return self._columns[name] + + +# ---------------------------------------------------------------------- +# Offline: the committed files are well-formed +# ---------------------------------------------------------------------- + + +@pytest.mark.parametrize("path", SPEC_PATHS, ids=SPEC_IDS) +def test_committed_spec_loads_and_round_trips(path: Any) -> None: + """A spec must parse into the current dataclass, with no stale fields.""" + spec = EvalSpec.load(str(path)) + assert spec.problem_id == path.parent.name + assert spec.version == path.stem + assert dataclasses.asdict(spec) == dataclasses.asdict(EvalSpec(**json.loads(path.read_text()))) + + +@pytest.mark.parametrize("path", SPEC_PATHS, ids=SPEC_IDS) +def test_committed_spec_requests_registered_metrics(path: Any) -> None: + """A spec asking for a metric nobody registered would fail at evaluation time.""" + for metric in EvalSpec.load(str(path)).metrics: + assert metric in METRICS, f"{path.name} requests unregistered metric {metric!r}" + + +@pytest.mark.parametrize("path", SPEC_PATHS, ids=SPEC_IDS) +def test_committed_spec_pins_its_dataset(path: Any) -> None: + """A package version does not identify a dataset; a revision does. + + Without this, an upstream dataset upload silently changes what every + published score means. + """ + spec = EvalSpec.load(str(path)) + assert spec.condition_digest, f"{path.name} was never frozen" + assert spec.dataset_id, f"{path.name} does not name its dataset" + assert spec.dataset_revision, f"{path.name} does not pin a dataset revision" + + +@pytest.mark.parametrize("path", SPEC_PATHS, ids=SPEC_IDS) +def test_multi_objective_specs_declare_a_scalarization(path: Any) -> None: + """Combining objectives with different units needs a declared rule, not a default.""" + spec = EvalSpec.load(str(path)) + if spec.objective_weights is None and spec.objective_weight_condition is None: + return + assert not (spec.objective_weights and spec.objective_weight_condition), ( + f"{path.name} declares two conflicting scalarizations" + ) + + +# ---------------------------------------------------------------------- +# The digest covers everything a score depends on +# ---------------------------------------------------------------------- + + +def test_digest_changes_when_a_condition_value_changes() -> None: + """Hashing only the row indices would miss an edited condition: row 7 is still row 7.""" + indices = np.array([0, 1, 2]) + designs = np.zeros((3, 4)) + before = _digest(indices, _FakeConditions({"volfrac": [0.3, 0.4, 0.5]}), designs) + after = _digest(indices, _FakeConditions({"volfrac": [0.3, 0.4, 0.9]}), designs) + assert before != after + + +def test_digest_changes_when_a_condition_is_renamed() -> None: + """The column names are part of the contract, not just the numbers.""" + indices = np.array([0, 1, 2]) + designs = np.zeros((3, 4)) + before = _digest(indices, _FakeConditions({"volfrac": [0.3, 0.4, 0.5]}), designs) + after = _digest(indices, _FakeConditions({"volume": [0.3, 0.4, 0.5]}), designs) + assert before != after + + +def test_digest_is_stable_across_column_order() -> None: + """Two datasets holding the same conditions must agree, whatever their column order.""" + indices = np.array([0, 1]) + designs = np.zeros((2, 4)) + one = _digest(indices, _FakeConditions({"a": [1.0, 2.0], "b": [3.0, 4.0]}), designs) + other = _digest(indices, _FakeConditions({"b": [3.0, 4.0], "a": [1.0, 2.0]}), designs) + assert one == other + + +def test_digest_covers_array_valued_conditions() -> None: + """thermoelastic2d's boundary matrices decide the answer as much as the scalars do.""" + indices = np.array([0]) + designs = np.zeros((1, 4)) + before = _digest(indices, _FakeConditions({"fixed_elements": [np.zeros((2, 2))]}), designs) + after = _digest(indices, _FakeConditions({"fixed_elements": [np.eye(2)]}), designs) + assert before != after + + +# ---------------------------------------------------------------------- +# Online: every committed spec still resolves +# ---------------------------------------------------------------------- + + +@pytest.mark.network +@pytest.mark.parametrize("path", SPEC_PATHS, ids=SPEC_IDS) +def test_committed_spec_reproduces_its_frozen_conditions(path: Any) -> None: + """The check the reviewer ran by hand: every spec must resolve, not just parse. + + `resolve` raises if the digest no longer matches, so this fails loudly when + a dataset moves out from under a published spec. + """ + from engibench.utils.all_problems import BUILTIN_PROBLEMS + + spec = EvalSpec.load(str(path)) + problem = BUILTIN_PROBLEMS[spec.problem_id]() + problem.reset(seed=spec.condition_seed) + + resolved = spec.resolve(problem) + + assert len(resolved.ref_designs) == spec.n_samples + assert resolved.conditions_tensor.shape == (spec.n_samples, len(resolved.condition_keys)) + if spec.volume_condition is not None: + assert spec.volume_condition in problem.conditions_keys From 810529f03941ebcd1f25bd0ab0f4fad357108d72 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Wed, 5 Aug 2026 15:51:39 +0200 Subject: [PATCH 07/24] Pin the EngiBench the specs were frozen against, and fix the CI breaks Follow-up to the review-response commit, addressing what CI surfaced. The two specs Soheyl reported as broken were never a case of the dataset drifting upstream. They were frozen against EngiBench main, whose photonics2d and thermoelastic2d read the v1 datasets, while the 0.2.0 release on PyPI still points those two at v0. Sampling a different dataset is why the digests could not match, and why CI could not reproduce them either. So the specs now record the problem definition they were frozen against, not just the dataset: the full condition list plus the EngiBench commit. That is checked before the digest, which turns the failure from two hashes that will never agree into a message naming the conditions and datasets that differ. CI installs EngiBench from git at a pinned commit, since no release carries the v1 problems yet, and the README says so for anyone running an evaluation. Two of my own breaks are fixed here too. A Python block inside CONTRIBUTING_A_MODEL.md was misformatted, which the newer ruff in CI checks and my older local one does not. And the leaderboard publishing tests constructed huggingface_hub errors in a way that only works on older releases, where `response` is optional and unread; they now build a real requests.Response, which every version accepts. Verified against hub 1.26.0 as well as 0.34.4. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/lint.yaml | 11 +++- CONTRIBUTING_A_MODEL.md | 2 +- README.md | 7 ++- engiopt/evaluation/spec.py | 85 +++++++++++++++++++++++++- engiopt/specs/beams2d/v1.json | 8 ++- engiopt/specs/heatconduction2d/v1.json | 6 +- engiopt/specs/photonics2d/v1.json | 10 ++- engiopt/specs/thermoelastic2d/v1.json | 11 +++- tests/test_leaderboard_publishing.py | 33 ++++++++-- tests/test_specs.py | 53 +++++++++++++++- 10 files changed, 208 insertions(+), 18 deletions(-) diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index fe29057b..28b99266 100644 --- a/.github/workflows/lint.yaml +++ b/.github/workflows/lint.yaml @@ -37,6 +37,13 @@ jobs: python-version: "3.12" - uses: actions/checkout@v4 - run: pip install -e ".[testing]" - # Every committed eval spec must still resolve against its pinned dataset - # revision, so a spec cannot silently stop reproducing its own conditions. + # Installed last so it wins: the committed eval specs are frozen against + # the current EngiBench, whose photonics2d and thermoelastic2d read the v1 + # datasets. The newest release on PyPI (0.2.0) still points those two at + # v0, so the PyPI build draws different conditions and no spec could + # reproduce itself. Pinned to a commit rather than a branch so CI cannot + # drift; bump it together with the specs' `engibench_version`. + - run: pip install "engibench[all] @ git+https://github.com/IDEALLab/EngiBench.git@0a028c03d02b266cb4394968c386e45781f4605b" + # Every committed spec must still resolve: same problem definition, same + # pinned dataset revision, same digest. - run: pytest tests/ -v diff --git a/CONTRIBUTING_A_MODEL.md b/CONTRIBUTING_A_MODEL.md index 106ca238..f9667577 100644 --- a/CONTRIBUTING_A_MODEL.md +++ b/CONTRIBUTING_A_MODEL.md @@ -66,7 +66,7 @@ solver settings that are not dataset columns at all. Size your network from the ```python from engiopt.transforms import condition_keys -cond_keys = condition_keys(problem) # e.g. ("volume_fraction_target", "rmin", "weight") +cond_keys = condition_keys(problem) # e.g. ("volume_fraction_target", "rmin", "weight") n_conds = len(cond_keys) ``` diff --git a/README.md b/README.md index 349fe7d9..24d5c375 100644 --- a/README.md +++ b/README.md @@ -52,12 +52,17 @@ pip install -e . You might want to install a specific PyTorch version, e.g., with CUDA on top of it, see [PyTorch install](https://pytorch.org/get-started/locally/). -If you're modifying EngiBench, you can install it from source and as editable: +**Evaluation needs EngiBench from source.** The committed evaluation specs are frozen +against the current EngiBench, whose `photonics2d` and `thermoelastic2d` read the `v1` +datasets; the newest PyPI release (0.2.0) still points those two at `v0`. Installing +EngiBench from PyPI therefore draws different conditions, and the specs report a clear +mismatch rather than scoring against the wrong data: ``` git clone git@github.com:IDEALLab/EngiBench.git cd EngiBench/ pip install -e ".[all]" ``` +This is also what CI installs, pinned to a commit. Training is unaffected. ## Running the code diff --git a/engiopt/evaluation/spec.py b/engiopt/evaluation/spec.py index d4bf8642..84457050 100644 --- a/engiopt/evaluation/spec.py +++ b/engiopt/evaluation/spec.py @@ -34,6 +34,43 @@ """Directory holding committed specs, `engiopt/specs//.json`.""" +class ProblemDefinitionMismatchError(ValueError): + """Raised when the installed EngiBench defines a problem differently than a spec expects. + + Distinct from a digest mismatch: the data did not change underneath us, the + *problem* did. Nothing about the spec can be reproduced until the EngiBench + versions agree, so this says which side moved rather than reporting two + hashes that will not match. + """ + + def __init__( + self, + *, + problem_id: str, + version: str, + expected_conditions: tuple[str, ...], + found_conditions: tuple[str, ...], + expected_dataset: str | None, + found_dataset: str | None, + engibench_version: str | None, + ) -> None: + self.expected_conditions = expected_conditions + self.found_conditions = found_conditions + details = [ + f"Eval spec {problem_id}/{version} was frozen against a different definition of {problem_id!r}.", + f" conditions expected: {list(expected_conditions)}", + f" conditions found: {list(found_conditions)}", + ] + if expected_dataset != found_dataset: + details.append(f" dataset expected: {expected_dataset}") + details.append(f" dataset found: {found_dataset}") + details.append( + f"The spec was frozen under EngiBench {engibench_version or 'unknown'}. Install an EngiBench whose " + f"{problem_id!r} matches, or freeze a new spec version against the one you have." + ) + super().__init__("\n".join(details)) + + @dataclass(frozen=True) class EvalSpec: """The frozen contract for evaluating models on one problem. @@ -67,9 +104,16 @@ class EvalSpec: the dataset at exactly this revision, so a new upload to the dataset repo cannot change what a leaderboard row means. Leave it empty to follow the dataset's current main branch. - engibench_version: EngiBench version the digest was recorded under. - Provenance only -- a package version does not identify a dataset, - which is what `dataset_revision` is for. + problem_conditions: The problem's full condition list when the spec was + frozen. Pinning the dataset is not enough on its own: which + conditions exist, and which dataset a problem points at, are defined + by the installed EngiBench, so a spec frozen against a different + EngiBench draws different columns. Checked before the digest, so + that case reports itself instead of surfacing as an opaque hash + mismatch. + engibench_version: EngiBench version the spec was frozen under. Recorded + for provenance; `problem_conditions` is what actually gets checked, + since a version string does not pin a problem definition either. notes: Free-text rationale for this version. """ @@ -86,6 +130,7 @@ class EvalSpec: condition_digest: str | None = None dataset_id: str | None = None dataset_revision: str | None = None + problem_conditions: tuple[str, ...] | None = None engibench_version: str | None = None notes: str = "" @@ -98,6 +143,8 @@ def __post_init__(self) -> None: object.__setattr__(self, "metrics", tuple(self.metrics)) if self.objective_weights is not None: object.__setattr__(self, "objective_weights", tuple(self.objective_weights)) + if self.problem_conditions is not None: + object.__setattr__(self, "problem_conditions", tuple(self.problem_conditions)) # ------------------------------------------------------------------ # Persistence @@ -138,6 +185,7 @@ def resolve(self, problem: Problem, *, device: th.device | None = None) -> Resol re-frozen under a new version. """ device = device or th.device("cpu") + self.check_problem_definition(problem) self.pin_dataset(problem) conditions_tensor, conditions, ref_designs, indices = sample_conditions( problem=problem, n_samples=self.n_samples, device=device, seed=self.condition_seed @@ -158,6 +206,36 @@ def resolve(self, problem: Problem, *, device: th.device | None = None) -> Resol condition_keys=tuple(get_scalar_condition_keys(problem, problem.dataset["test"])), ) + def check_problem_definition(self, problem: Problem) -> None: + """Verify the installed EngiBench defines this problem the way the spec was frozen against. + + A pinned dataset revision fixes the *data*; it does not fix which + conditions the problem declares or which dataset it points at, both of + which live in EngiBench. A spec frozen against a different EngiBench + draws different condition columns and so computes a different digest -- + which, without this check, surfaces only as an unexplained hash + mismatch. + + Raises: + ProblemDefinitionMismatchError: If the installed problem declares + different conditions, or points at a different dataset. + """ + if self.problem_conditions is None: + return + current = tuple(problem.conditions_keys) + current_dataset = getattr(problem, "dataset_id", None) + dataset_moved = self.dataset_id is not None and current_dataset is not None and current_dataset != self.dataset_id + if current != self.problem_conditions or dataset_moved: + raise ProblemDefinitionMismatchError( + problem_id=self.problem_id, + version=self.version, + expected_conditions=self.problem_conditions, + found_conditions=current, + expected_dataset=self.dataset_id, + found_dataset=current_dataset, + engibench_version=self.engibench_version, + ) + def pin_dataset(self, problem: Problem) -> None: """Point `problem.dataset` at this spec's frozen dataset revision. @@ -189,6 +267,7 @@ def freeze(self, problem: Problem) -> EvalSpec: **{ **asdict(frozen), "condition_digest": _digest(indices, conditions, ref_designs), + "problem_conditions": tuple(problem.conditions_keys), "engibench_version": getattr(engibench, "__version__", None), } ) diff --git a/engiopt/specs/beams2d/v1.json b/engiopt/specs/beams2d/v1.json index b233c964..deb48742 100644 --- a/engiopt/specs/beams2d/v1.json +++ b/engiopt/specs/beams2d/v1.json @@ -3,7 +3,7 @@ "condition_seed": 1, "dataset_id": "IDEALLab/beams_2d_50_100_v0", "dataset_revision": "ccb64d4f09cf66a5ed9ba2081438d695b9e2438d", - "engibench_version": "0.2.0", + "engibench_version": "0.2.0+0a028c03d02b", "metrics": [ "mmd", "dpp", @@ -16,6 +16,12 @@ "notes": "Feasibility = EngiBench constraints plus the volfrac budget.", "objective_weight_condition": null, "objective_weights": null, + "problem_conditions": [ + "volfrac", + "rmin", + "forcedist", + "overhang_constraint" + ], "problem_id": "beams2d", "sigma": 10.0, "version": "v1", diff --git a/engiopt/specs/heatconduction2d/v1.json b/engiopt/specs/heatconduction2d/v1.json index 2f182309..56875d8b 100644 --- a/engiopt/specs/heatconduction2d/v1.json +++ b/engiopt/specs/heatconduction2d/v1.json @@ -3,7 +3,7 @@ "condition_seed": 1, "dataset_id": "IDEALLab/heat_conduction_2d_v0", "dataset_revision": "9f07c1e70d244f783e34a71c96db621173871dc1", - "engibench_version": "0.2.0", + "engibench_version": "0.2.0+0a028c03d02b", "metrics": [ "mmd", "dpp", @@ -16,6 +16,10 @@ "notes": "Feasibility = EngiBench constraints plus the volume budget.", "objective_weight_condition": null, "objective_weights": null, + "problem_conditions": [ + "volume", + "length" + ], "problem_id": "heatconduction2d", "sigma": 10.0, "version": "v1", diff --git a/engiopt/specs/photonics2d/v1.json b/engiopt/specs/photonics2d/v1.json index dfaf548d..62a8f73b 100644 --- a/engiopt/specs/photonics2d/v1.json +++ b/engiopt/specs/photonics2d/v1.json @@ -3,7 +3,7 @@ "condition_seed": 1, "dataset_id": "IDEALLab/photonics_2d_120_120_v1", "dataset_revision": "be811ff69d67cd8c79148e29f2a68d293ec00955", - "engibench_version": "0.2.0", + "engibench_version": "0.2.0+0a028c03d02b", "metrics": [ "mmd", "dpp", @@ -16,6 +16,14 @@ "notes": "No volume budget: feasibility is EngiBench's declared constraints alone.", "objective_weight_condition": null, "objective_weights": null, + "problem_conditions": [ + "lambda1", + "lambda2", + "blur_radius", + "num_elems_x", + "num_elems_y", + "num_optimization_steps" + ], "problem_id": "photonics2d", "sigma": 10.0, "version": "v1", diff --git a/engiopt/specs/thermoelastic2d/v1.json b/engiopt/specs/thermoelastic2d/v1.json index 892239b6..4dee1208 100644 --- a/engiopt/specs/thermoelastic2d/v1.json +++ b/engiopt/specs/thermoelastic2d/v1.json @@ -3,7 +3,7 @@ "condition_seed": 1, "dataset_id": "IDEALLab/thermoelastic_2d_v1", "dataset_revision": "7d5b4a694bd1a8f3f25fd745ac9f669764e29ea5", - "engibench_version": "0.2.0", + "engibench_version": "0.2.0+0a028c03d02b", "metrics": [ "mmd", "dpp", @@ -16,6 +16,15 @@ "notes": "Objectives combined per-sample by the `weight` condition. Feasibility = EngiBench constraints plus the volume_fraction_target budget.", "objective_weight_condition": "weight", "objective_weights": null, + "problem_conditions": [ + "fixed_elements", + "force_elements_x", + "force_elements_y", + "heatsink_elements", + "volume_fraction_target", + "rmin", + "weight" + ], "problem_id": "thermoelastic2d", "sigma": 10.0, "version": "v1", diff --git a/tests/test_leaderboard_publishing.py b/tests/test_leaderboard_publishing.py index 5e82c4a5..25b8760f 100644 --- a/tests/test_leaderboard_publishing.py +++ b/tests/test_leaderboard_publishing.py @@ -15,10 +15,31 @@ from huggingface_hub.errors import RepositoryNotFoundError import pandas as pd import pytest +import requests from engiopt.evaluation import leaderboard as lb +def hub_error(error_cls: type[Exception], message: str, status: int = 500) -> Exception: + """Build a `huggingface_hub` error the way the library itself would. + + Versions disagree about `response`: newer ones require it as a keyword-only + argument, and they read attributes off it (`request`, `headers`). A stub + would raise `TypeError`/`AttributeError` in place of the error under test, + turning a real assertion into a confusing failure -- so this passes a + genuine `requests.Response`, which satisfies every version. + """ + response = requests.Response() + response.status_code = status + response.reason = message + response._content = message.encode() + response.request = requests.Request(method="GET", url="https://huggingface.co/api").prepare() + try: + return error_cls(message, response=response) # type: ignore[call-arg] + except TypeError: + return error_cls(message) + + def _row(algo: str, mmd: float) -> dict[str, Any]: return { "problem_id": "p", @@ -46,7 +67,7 @@ def create_repo(self, **_kwargs: Any) -> None: def repo_info(self, **_kwargs: Any) -> Any: if self.revision is None: - raise RepositoryNotFoundError("no such repo") + raise hub_error(RepositoryNotFoundError, "no such repo", status=404) return type("Info", (), {"sha": self.revision})() def upload_file(self, **kwargs: Any) -> None: @@ -54,7 +75,7 @@ def upload_file(self, **kwargs: Any) -> None: if self.fail_first and len(self.uploads) == 1: # What the Hub returns when `parent_commit` is no longer the head. self.revision = "sha2" - raise HfHubHTTPError("412 Client Error: parent_commit is out of date") + raise hub_error(HfHubHTTPError, "412 Client Error: parent_commit is out of date", status=412) @pytest.fixture @@ -87,7 +108,7 @@ def test_missing_file_in_an_existing_repo_reads_as_empty(monkeypatch: pytest.Mon monkeypatch.setattr("huggingface_hub.HfApi", _FakeApi()) def fake_download(**_kwargs: Any) -> str: - raise EntryNotFoundError("no leaderboard.csv") + raise hub_error(EntryNotFoundError, "no leaderboard.csv", status=404) monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download) assert lb.load_from_hub("org/board").empty @@ -98,7 +119,7 @@ def test_transient_download_failure_is_not_an_empty_board(monkeypatch: pytest.Mo monkeypatch.setattr("huggingface_hub.HfApi", _FakeApi()) def fake_download(**_kwargs: Any) -> str: - raise HfHubHTTPError("503 Server Error: Service Unavailable") + raise hub_error(HfHubHTTPError, "503 Server Error: Service Unavailable", status=503) monkeypatch.setattr("huggingface_hub.hf_hub_download", fake_download) with pytest.raises(HfHubHTTPError): @@ -110,7 +131,7 @@ def test_auth_failure_is_not_an_empty_board(monkeypatch: pytest.MonkeyPatch) -> api = _FakeApi() def raise_auth(**_kwargs: Any) -> Any: - raise HfHubHTTPError("401 Client Error: Unauthorized") + raise hub_error(HfHubHTTPError, "401 Client Error: Unauthorized", status=401) api.repo_info = raise_auth # type: ignore[method-assign] monkeypatch.setattr("huggingface_hub.HfApi", api) @@ -162,7 +183,7 @@ def test_push_gives_up_rather_than_forcing_a_write(monkeypatch: pytest.MonkeyPat class _AlwaysStale(_FakeApi): def upload_file(self, **kwargs: Any) -> None: self.uploads.append(kwargs) - raise HfHubHTTPError("412 Client Error: parent_commit is out of date") + raise hub_error(HfHubHTTPError, "412 Client Error: parent_commit is out of date", status=412) monkeypatch.setattr("huggingface_hub.HfApi", _AlwaysStale()) with pytest.raises(HfHubHTTPError): diff --git a/tests/test_specs.py b/tests/test_specs.py index 0129f546..673d36f8 100644 --- a/tests/test_specs.py +++ b/tests/test_specs.py @@ -10,7 +10,7 @@ import dataclasses import json -from typing import Any +from typing import Any, ClassVar import numpy as np import pytest @@ -18,6 +18,7 @@ from engiopt.evaluation.registry import METRICS from engiopt.evaluation.spec import _digest from engiopt.evaluation.spec import EvalSpec +from engiopt.evaluation.spec import ProblemDefinitionMismatchError from engiopt.evaluation.spec import SPEC_ROOT SPEC_PATHS = sorted(SPEC_ROOT.glob("*/*.json")) @@ -133,12 +134,21 @@ def test_committed_spec_reproduces_its_frozen_conditions(path: Any) -> None: `resolve` raises if the digest no longer matches, so this fails loudly when a dataset moves out from under a published spec. + + A spec frozen against a *different definition of the problem* is skipped + rather than failed: the installed EngiBench cannot draw those conditions at + all, so there is nothing here to verify. The skip names the mismatch, and a + digest mismatch under a matching definition still fails. """ from engibench.utils.all_problems import BUILTIN_PROBLEMS spec = EvalSpec.load(str(path)) problem = BUILTIN_PROBLEMS[spec.problem_id]() problem.reset(seed=spec.condition_seed) + try: + spec.check_problem_definition(problem) + except ProblemDefinitionMismatchError as mismatch: + pytest.skip(str(mismatch)) resolved = spec.resolve(problem) @@ -146,3 +156,44 @@ def test_committed_spec_reproduces_its_frozen_conditions(path: Any) -> None: assert resolved.conditions_tensor.shape == (spec.n_samples, len(resolved.condition_keys)) if spec.volume_condition is not None: assert spec.volume_condition in problem.conditions_keys + + +def test_a_differently_defined_problem_reports_itself() -> None: + """A pinned dataset does not pin the problem definition, which also lives in EngiBench. + + Without this check the case shows up only as two hashes that will never + match, giving no hint that the EngiBench versions disagree. + """ + + class _Problem: + conditions_keys: ClassVar[list[str]] = ["volfrac", "rmin", "weight"] + dataset_id = "IDEALLab/thermoelastic_2d_v0" + + spec = EvalSpec( + problem_id="thermoelastic2d", + problem_conditions=("volume_fraction_target", "rmin", "weight"), + dataset_id="IDEALLab/thermoelastic_2d_v1", + engibench_version="0.2.0", + ) + with pytest.raises(ProblemDefinitionMismatchError) as caught: + spec.check_problem_definition(_Problem()) + + message = str(caught.value) + assert "volume_fraction_target" in message + assert "volfrac" in message + assert "thermoelastic_2d_v0" in message + + +def test_a_matching_problem_definition_passes() -> None: + """The check must not fire when the installed EngiBench agrees with the spec.""" + + class _Problem: + conditions_keys: ClassVar[list[str]] = ["volfrac", "rmin"] + dataset_id = "IDEALLab/beams_2d_50_100_v0" + + spec = EvalSpec( + problem_id="beams2d", + problem_conditions=("volfrac", "rmin"), + dataset_id="IDEALLab/beams_2d_50_100_v0", + ) + spec.check_problem_definition(_Problem()) From 7f2088e75af5fc85ecb0bdb4be95543e3e5505a3 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Thu, 6 Aug 2026 10:29:22 +0200 Subject: [PATCH 08/24] Address review round 2: make CI actually verify, and replay training-time conditions Nine follow-up comments, eight of them correct and four describing defects the previous round introduced. The important one: CI was green because it skipped the two specs the job exists to check. The pinned EngiBench install was a no-op -- the git build reports the same version string as the PyPI one, so pip considered the requirement already satisfied and kept the wrong copy. It is now a --force-reinstall with an explicit assertion that the pin took effect, and a spec frozen against a different EngiBench definition fails rather than skipping, since skipping would go green exactly when the pin has broken. `freeze()` also captures the EngiBench source revision itself, so it reproduces the values already committed in the specs instead of relying on how they were written by hand. Conditions are now presented to a model the way its training data looked. The evaluator supplies every scalar condition, unscaled; a model that trained on a subset, a different order, or a rescaling gets that view reconstructed from what its checkpoint recorded. Previously the contract demanded exact equality, which rejected VQGAN outright -- it drops constant columns during training -- and VQGAN worked around it by re-deriving its preprocessing from the evaluation sample. That was silently wrong: fifty evaluation rows do not have the training split's mean, and "columns that never vary" is a different set in fifty rows than in forty thousand. Training now records the statistics it fitted, sampling replays them, and the adapter's bespoke preprocessing is deleted rather than patched. Any two-stage model that preprocesses conditions gets the same treatment without writing code for it. Feasibility no longer rides along with the optimizer pass. It describes the design as generated, so it is judged before any solver runs -- which means it still reports when the optimizer refuses to start from an invalid design, the case where the answer matters most, and asking for `viol` alone now costs a constraint check instead of a full optimization. It is a cheap metric as a result. Smaller corrections: ranks restart within each problem and spec version rather than running across the whole table; a private repo is no longer mistaken for an empty leaderboard, and a repo whose leaderboard file is merely absent keeps its revision so two first publishers cannot overwrite each other; `--skip-existing` compares the checkpoint hash, so retrained weights are re-scored; and the package hash covers every weight file present rather than only the declared ones, since VQGAN loads a condition encoder it cannot declare. Not adopted: pinning the EngiBench revision as a hard requirement at resolution, or adding it to leaderboard ranking identity. Recording it is right and is done; enforcing it would break every contributor whose EngiBench differs by any commit, and ranking on it would partition the board per build so that two models could not be compared unless scored on byte-identical EngiBench. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/lint.yaml | 19 ++-- engiopt/checkpoint_store.py | 37 ++++++-- engiopt/core.py | 91 +++++++++++++++---- engiopt/evaluate.py | 1 + engiopt/evaluation/context.py | 15 +++- engiopt/evaluation/leaderboard.py | 66 ++++++++++++-- engiopt/evaluation/metrics/builtin.py | 10 ++- engiopt/evaluation/spec.py | 31 ++++++- engiopt/generators/vqgan/adapter.py | 25 ++---- engiopt/generators/vqgan/vqgan.py | 9 +- tests/test_checkpoint_packages.py | 61 +++++++++++++ tests/test_condition_schema.py | 83 +++++++++++++++++ tests/test_conditions_and_gaps.py | 30 +++++++ tests/test_evaluation.py | 10 ++- tests/test_leaderboard_publishing.py | 124 ++++++++++++++++++++++++++ tests/test_specs.py | 15 ++-- 16 files changed, 547 insertions(+), 80 deletions(-) diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index 28b99266..30866aad 100644 --- a/.github/workflows/lint.yaml +++ b/.github/workflows/lint.yaml @@ -37,13 +37,18 @@ jobs: python-version: "3.12" - uses: actions/checkout@v4 - run: pip install -e ".[testing]" - # Installed last so it wins: the committed eval specs are frozen against - # the current EngiBench, whose photonics2d and thermoelastic2d read the v1 - # datasets. The newest release on PyPI (0.2.0) still points those two at - # v0, so the PyPI build draws different conditions and no spec could - # reproduce itself. Pinned to a commit rather than a branch so CI cannot - # drift; bump it together with the specs' `engibench_version`. - - run: pip install "engibench[all] @ git+https://github.com/IDEALLab/EngiBench.git@0a028c03d02b266cb4394968c386e45781f4605b" + # The committed eval specs are frozen against the current EngiBench, whose + # photonics2d and thermoelastic2d read the v1 datasets. The newest release + # on PyPI (0.2.0) still points those two at v0, so the PyPI build draws + # different conditions and no spec could reproduce itself. + # + # --force-reinstall is required, not decorative: the git build reports the + # same version string (0.2.0) as the PyPI one, so a plain install is a + # no-op and the specs would be verified against the wrong EngiBench. + # --no-deps keeps it from re-resolving the whole stack. + - run: pip install --force-reinstall --no-deps "engibench @ git+https://github.com/IDEALLab/EngiBench.git@0a028c03d02b266cb4394968c386e45781f4605b" + # Fail loudly if the pin silently stopped taking effect. + - run: python -c "import engibench.problems.thermoelastic2d.v0 as m; assert m.ThermoElastic2D.dataset_id.endswith('_v1'), m.ThermoElastic2D.dataset_id" # Every committed spec must still resolve: same problem definition, same # pinned dataset revision, same digest. - run: pytest tests/ -v diff --git a/engiopt/checkpoint_store.py b/engiopt/checkpoint_store.py index c55cd9f1..3ea619f6 100644 --- a/engiopt/checkpoint_store.py +++ b/engiopt/checkpoint_store.py @@ -85,6 +85,7 @@ def save_checkpoint_package( config_fingerprint: str | None = None, is_default_config: bool = True, condition_keys: list[str] | tuple[str, ...] | None = None, + condition_stats: dict[str, list[float]] | None = None, ) -> dict[str, Any]: """Save a checkpoint package to HuggingFace. @@ -106,7 +107,10 @@ def save_checkpoint_package( Pass `condition_keys` (from `engiopt.transforms.condition_keys`) so the checkpoint records the condition schema it was trained under; loading then - rebuilds the network for exactly those columns. + rebuilds the network for exactly those columns. A model that also *rescales* + its conditions passes `condition_stats={"mean": [...], "std": [...]}` so the + evaluator can reproduce the same scale, instead of refitting it on the far + smaller evaluation sample. W&B is no longer a checkpoint storage backend; the active W&B run still receives a summary pointing at the HF package for traceability. @@ -128,6 +132,7 @@ def save_checkpoint_package( primary_files=primary_files, metadata=metadata, condition_keys=condition_keys, + condition_stats=condition_stats, ) base_metadata.update(_build_wandb_run_metadata()) @@ -301,19 +306,34 @@ def _resolve_local_package(local_model_dir: str, required_files: list[str]) -> R return _load_package_from_directory(root_dir=local_model_dir, required_files=required_files, source="local") -def hash_checkpoint_files(files: dict[str, str]) -> str: - """Content hash of a package's weight files, identifying these exact weights. +def hash_package_contents(root_dir: str) -> str: + """Content hash of every weight file in a package, identifying these exact weights. + + Hashes what the package *contains* rather than what the model declared it + needs. A model may load a file it did not list -- VQGAN's conditional + variant reads `cvqgan.pth`, which is not in its `checkpoint_files` because + the unconditional variant has none -- and two packages differing only in + that file generate differently. Hashing the declared list would give them + the same identity. + + `run_config.json` and `metadata.json` are excluded: they describe the + package rather than being weights, and `metadata.json` records paths and + revisions that differ between the two locations one run writes. Args: - files: Package file name to local path, as on `ResolvedCheckpoint`. + root_dir: Local directory holding the package. Returns: A 16-character hash over the file names and their bytes. """ + described = {"run_config.json", "metadata.json"} + weight_files = sorted( + entry for entry in os.listdir(root_dir) if entry not in described and os.path.isfile(os.path.join(root_dir, entry)) + ) hasher = hashlib.sha256() - for file_name in sorted(files): + for file_name in weight_files: hasher.update(file_name.encode()) - with open(files[file_name], "rb") as handle: + with open(os.path.join(root_dir, file_name), "rb") as handle: for chunk in iter(lambda: handle.read(1 << 20), b""): hasher.update(chunk) return hasher.hexdigest()[:16] @@ -340,7 +360,7 @@ def _load_package_from_directory( run_config=run_config, metadata=metadata, revision=revision, - content_hash=hash_checkpoint_files(files), + content_hash=hash_package_contents(root_dir), ) @@ -381,6 +401,7 @@ def _build_metadata( primary_files: list[str] | None, metadata: dict[str, Any] | None, condition_keys: list[str] | tuple[str, ...] | None = None, + condition_stats: dict[str, list[float]] | None = None, ) -> dict[str, Any]: payload = dict(metadata or {}) payload.update( @@ -395,6 +416,8 @@ def _build_metadata( ) if condition_keys is not None: payload["condition_keys"] = list(condition_keys) + if condition_stats is not None: + payload["condition_stats"] = {key: [float(v) for v in values] for key, values in condition_stats.items()} return payload diff --git a/engiopt/core.py b/engiopt/core.py index 328f6e4a..527e357d 100644 --- a/engiopt/core.py +++ b/engiopt/core.py @@ -168,6 +168,15 @@ def config_path_parts(config_fingerprint_value: str | None) -> list[str] | None: CONDITION_KEYS_METADATA_FIELD = "condition_keys" """Checkpoint-metadata field recording which conditions a model was trained on.""" +CONDITION_STATS_METADATA_FIELD = "condition_stats" +"""Checkpoint-metadata field recording how a model rescaled its conditions. + +`{"mean": [...], "std": [...]}`, aligned to `condition_keys`. A model that +normalized its conditions during training must see the same scale at evaluation, +and only the training split's statistics give that -- recomputing them from the +evaluation sample silently feeds the network numbers it was never trained on. +""" + def condition_keys_for(problem: Problem, resolved: ResolvedCheckpoint | None = None) -> tuple[str, ...]: """Which conditions a generator consumes, in the order its input tensor holds them. @@ -191,6 +200,20 @@ def condition_keys_for(problem: Problem, resolved: ResolvedCheckpoint | None = N return tuple(scalar_condition_keys(problem)) +def condition_stats_for(resolved: ResolvedCheckpoint | None) -> tuple[list[float], list[float]] | None: + """The `(mean, std)` a checkpoint rescaled its conditions by, if it recorded any. + + Returns None for the great majority of models, which feed conditions to the + network unscaled and need no replay. + """ + if resolved is None: + return None + stats = resolved.metadata.get(CONDITION_STATS_METADATA_FIELD) + if not stats or "mean" not in stats or "std" not in stats: + return None + return [float(v) for v in stats["mean"]], [float(v) for v in stats["std"]] + + def design_shape_of(problem: Problem) -> tuple[int, ...]: """Shape of a single design for `problem`. @@ -249,6 +272,7 @@ def __init__( device: th.device, run_config: dict[str, Any] | None = None, condition_keys: tuple[str, ...] | None = None, + condition_stats: tuple[list[float], list[float]] | None = None, checkpoint_revision: str | None = None, checkpoint_hash: str | None = None, ) -> None: @@ -258,6 +282,7 @@ def __init__( self.device = device self.run_config: dict[str, Any] = run_config or {} self._condition_keys = condition_keys + self._condition_stats = condition_stats self.checkpoint_revision = checkpoint_revision """Commit the checkpoint was downloaded at, when it came from the Hub.""" self.checkpoint_hash = checkpoint_hash @@ -350,6 +375,7 @@ def from_pretrained( seed=seed, run_config=resolved.run_config, condition_keys=condition_keys_for(problem, resolved), + condition_stats=condition_stats_for(resolved), checkpoint_revision=resolved.revision, checkpoint_hash=resolved.content_hash, **kwargs, @@ -426,32 +452,67 @@ def _synchronize_device(self) -> None: th.mps.synchronize() def _as_batch(self, conditions: ConditionBatch | th.Tensor | npt.NDArray[Any] | None) -> ConditionBatch: - """Normalize whatever the caller passed into a `ConditionBatch` on this device.""" + """Present the conditions to the model exactly as its training data looked. + + The evaluator hands every model the same batch: all of the problem's + scalar conditions, unscaled. A model that trained on some other view of + them -- a subset, a different order, a rescaling -- gets that view + reconstructed here from what its checkpoint recorded, rather than + re-deriving it from the evaluation sample. Fifty evaluation rows do not + have the training split's mean, and "columns that never vary" is a + different set in fifty rows than in forty thousand. + """ if isinstance(conditions, ConditionBatch): - self._check_condition_keys(conditions.keys) - tensor = conditions.tensor - moved = None if tensor is None else tensor.to(device=self.device, dtype=th.float) - return ConditionBatch(tensor=moved, dataset=conditions.dataset, keys=conditions.keys) + projected = self._select_trained_columns(conditions) + return ConditionBatch( + tensor=self._rescale(self._to_device(projected.tensor)), + dataset=conditions.dataset, + keys=projected.keys, + ) if conditions is None: return ConditionBatch(tensor=None, keys=self.condition_keys) + # A bare tensor carries no column names, so it is taken as already in + # this model's own layout. tensor = conditions if isinstance(conditions, th.Tensor) else th.as_tensor(np.asarray(conditions)) - return ConditionBatch(tensor=tensor.to(device=self.device, dtype=th.float), keys=self.condition_keys) + return ConditionBatch(tensor=self._to_device(tensor), keys=self.condition_keys) + + def _to_device(self, tensor: th.Tensor | None) -> th.Tensor | None: + """Move a condition tensor onto this generator's device, passing None through.""" + return None if tensor is None else tensor.to(device=self.device, dtype=th.float) - def _check_condition_keys(self, keys: tuple[str, ...]) -> None: - """Refuse conditions whose columns are not the ones this model was trained on. + def _select_trained_columns(self, batch: ConditionBatch) -> ConditionBatch: + """Pick out the columns this checkpoint was trained on, in its own order. - Silently accepting them produces either a shape error deep inside the - network or, worse, a design conditioned on the wrong numbers. + The checkpoint's schema is authoritative, so a caller supplying *more* + conditions than the model uses is fine -- they are projected away. Only + a column the model needs and did not get is an error, because that one + cannot be reconstructed. Raises: - ValueError: If the caller's condition columns differ from the - checkpoint's. + ValueError: If a condition the checkpoint was trained on is absent. """ - if keys and self.condition_keys and tuple(keys) != self.condition_keys: + if not batch.keys or not self.condition_keys or tuple(batch.keys) == self.condition_keys: + return batch + missing = [key for key in self.condition_keys if key not in batch.keys] + if missing: raise ValueError( - f"{self.algo_id} was trained on conditions {self.condition_keys} but was given {tuple(keys)}. " - "Re-train against the current problem, or evaluate the checkpoint that matches it." + f"{self.algo_id} was trained on conditions {self.condition_keys}, but {missing} " + f"are missing from the {list(batch.keys)} it was given. Re-train against the current " + "problem, or evaluate the checkpoint that matches it." ) + if batch.tensor is None: + return ConditionBatch(tensor=None, dataset=batch.dataset, keys=self.condition_keys) + columns = [batch.keys.index(key) for key in self.condition_keys] + return ConditionBatch(tensor=batch.tensor[:, columns], dataset=batch.dataset, keys=self.condition_keys) + + def _rescale(self, tensor: th.Tensor | None) -> th.Tensor | None: + """Apply the checkpoint's recorded condition normalization, if it recorded any.""" + if tensor is None or self._condition_stats is None: + return tensor + mean, std = self._condition_stats + mean_t = th.tensor(mean, device=tensor.device, dtype=tensor.dtype) + std_t = th.tensor(std, device=tensor.device, dtype=tensor.dtype).clamp(min=1e-8) + return (tensor - mean_t) / std_t @cached_property def design_shape(self) -> tuple[int, ...]: diff --git a/engiopt/evaluate.py b/engiopt/evaluate.py index 74952b4d..e72246b3 100644 --- a/engiopt/evaluate.py +++ b/engiopt/evaluate.py @@ -176,6 +176,7 @@ def _drop_already_published(args: Args, evaluator: Evaluator, generators: list[G config_fingerprint=generator.config_fingerprint, seed=generator.seed, spec_version=evaluator.spec.version, + checkpoint_hash=generator.checkpoint_hash, ) ] print(f"{len(remaining)} generator(s) left after skipping already-published rows.") diff --git a/engiopt/evaluation/context.py b/engiopt/evaluation/context.py index a3edfdb7..d341bd5c 100644 --- a/engiopt/evaluation/context.py +++ b/engiopt/evaluation/context.py @@ -68,8 +68,6 @@ class OptimizationResults: """Sum of the optimality gaps over each design's re-optimization history.""" fog: list[float] = field(default_factory=list) """Final optimality gap at the end of re-optimization.""" - viol: list[bool] = field(default_factory=list) - """Whether each design is infeasible; see `EvaluationContext.is_infeasible`.""" @dataclass @@ -254,9 +252,20 @@ def optimization(self) -> OptimizationResults: results.iog.append(self.scalarize_gap(np.asarray(generated_objective) - np.asarray(reference_optimum), i)) results.cog.append(sum(self.scalarize_gap(step_gap, i) for step_gap in gaps)) results.fog.append(self.scalarize_gap(gaps[-1], i)) - results.viol.append(self.is_infeasible(design, conditions)) return results + @cached_property + def feasibility(self) -> list[bool]: + """Whether each generated design is infeasible, judged before any solver runs. + + Deliberately independent of `optimization`. Feasibility is a property of + the design as generated, so it must survive the optimizer refusing to + start from an invalid one -- exactly the case where the answer matters + most. Keeping it separate also means asking for `viol` alone costs a + constraint check rather than a full optimization pass. + """ + return [self.is_infeasible(self.design_for_solver(i), self.condition_at(i)) for i in range(self.n_samples)] + def is_infeasible(self, design: Any, conditions: dict[str, Any] | None) -> bool: """Whether one generated design fails the problem's declared feasibility. diff --git a/engiopt/evaluation/leaderboard.py b/engiopt/evaluation/leaderboard.py index e8272057..4fcc6256 100644 --- a/engiopt/evaluation/leaderboard.py +++ b/engiopt/evaluation/leaderboard.py @@ -30,6 +30,14 @@ silently keeping only whichever ran last. """ +RANK_PARTITION = ["problem_id", "spec_version"] +"""Within which group a rank of 1 means "best". + +Two models are only comparable if they were scored on the same problem under the +same protocol, so ranks restart per partition rather than running across the +whole table. +""" + ENTRY_KEY = ["problem_id", "algo_id", "config_fingerprint", "spec_version"] """What a *ranked entry* is: one configuration of one algorithm under one spec. @@ -82,7 +90,10 @@ def rank( Returns: One row per `ENTRY_KEY` sorted best-first, with a `rank` column and the - seed count that produced each value. + seed count that produced each value. Ranks restart at 1 within each + `(problem_id, spec_version)`: a beams2d score and a photonics2d score + measure different things, so placing them in one ordering would invent a + comparison that does not exist. """ from engiopt.evaluation.registry import METRICS @@ -96,7 +107,8 @@ def rank( .sort_values("value", ascending=not spec.higher_is_better) .reset_index(drop=True) ) - grouped["rank"] = grouped.index + 1 + partition = [col for col in RANK_PARTITION if col in grouped.columns] + grouped["rank"] = grouped.groupby(partition).cumcount() + 1 if partition else grouped.index + 1 return grouped.rename(columns={"value": f"{metric}_{aggregate}"}) @@ -164,24 +176,49 @@ def load_from_hub(repo_id: str, *, token: str | None = None) -> pd.DataFrame: def _load_from_hub(repo_id: str, *, token: str | None = None) -> tuple[pd.DataFrame, str | None]: - """Download the board together with the repo commit it was read at.""" + """Download the board together with the repo commit it was read at. + + The two "not found" cases are different and must stay different: + + - **No repo yet.** Nothing exists, so there is no revision either, and the + first publish creates it. + - **Repo exists, no `leaderboard.csv` yet.** The repo already has a head + commit, and returning it matters: it becomes the `parent_commit` of the + first publish, so two jobs racing to publish first cannot overwrite each + other. + + A missing repo is also what the Hub reports for a private repo the caller + cannot see, so that case is separated by status code rather than by + exception type -- treating a 401 as an empty board would let the next + publish delete everything. + """ from huggingface_hub import hf_hub_download from huggingface_hub import HfApi from huggingface_hub.errors import EntryNotFoundError from huggingface_hub.errors import RepositoryNotFoundError - from huggingface_hub.errors import RevisionNotFoundError try: revision = HfApi(token=token).repo_info(repo_id=repo_id, repo_type="dataset").sha + except RepositoryNotFoundError as exc: + if _status_code(exc) in {401, 403}: + raise + return pd.DataFrame(), None + + try: path = hf_hub_download( repo_id=repo_id, filename=LEADERBOARD_FILE, repo_type="dataset", revision=revision, token=token ) - # Only "there is nothing published yet" is a normal, empty-board outcome. - except (EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError): - return pd.DataFrame(), None + # The repo is there but empty; keep its revision for the first publish. + except EntryNotFoundError: + return pd.DataFrame(), revision return pd.read_csv(path), revision +def _status_code(exc: Exception) -> int | None: + """HTTP status behind a `huggingface_hub` error, when it carries one.""" + return getattr(getattr(exc, "response", None), "status_code", None) + + def push_to_hub( new_rows: pd.DataFrame, repo_id: str, @@ -252,14 +289,25 @@ def _is_stale_commit(exc: Exception) -> bool: def already_evaluated(board: pd.DataFrame, **key: Any) -> bool: - """Whether the board already holds a row for the given `ROW_KEY` values. + """Whether the board already holds a row for exactly this model. Lets a sweep skip work it has already done instead of recomputing it. + + Pass `checkpoint_hash` and the skip becomes about the *weights*, not just + the name: re-training the same configuration and seed produces different + weights under the same identity, and those deserve a fresh score rather + than inheriting the old row's. A row predating the hash column has no value + to compare, so it is treated as a match and still skipped. """ if board.empty: return False mask = pd.Series(data=True, index=board.index) for column, value in key.items(): - if column in board.columns: + if column not in board.columns: + continue + if column == "checkpoint_hash": + # Older rows carry no hash; do not force them to be re-evaluated. + mask &= board[column].isna() | (board[column] == value) + else: mask &= board[column] == value return bool(mask.any()) diff --git a/engiopt/evaluation/metrics/builtin.py b/engiopt/evaluation/metrics/builtin.py index 6a60ee7a..2f912f71 100644 --- a/engiopt/evaluation/metrics/builtin.py +++ b/engiopt/evaluation/metrics/builtin.py @@ -64,7 +64,7 @@ def dpp(ctx: EvaluationContext) -> float: @register_metric( "viol", family="feasibility", - cost="expensive", + cost="cheap", higher_is_better=False, description="Fraction of designs violating the problem's constraints or their volume budget.", ) @@ -75,10 +75,12 @@ def viol(ctx: EvaluationContext) -> float: the spec's `volume_condition` adds the volume-fraction budget for problems that have one. - Marked expensive only because it is computed during the shared - optimizer/simulator pass; the arithmetic itself is trivial. + Cheap, and deliberately so. Feasibility describes the design as generated, + so it is judged before any solver runs -- which means it still reports when + the optimizer refuses to start from an invalid design, the case where the + answer matters most. """ - values = ctx.optimization.viol + values = ctx.feasibility return float(np.mean(values)) if values else float("nan") diff --git a/engiopt/evaluation/spec.py b/engiopt/evaluation/spec.py index 84457050..09535a55 100644 --- a/engiopt/evaluation/spec.py +++ b/engiopt/evaluation/spec.py @@ -17,6 +17,7 @@ import hashlib import json from pathlib import Path +import subprocess from typing import Any, TYPE_CHECKING import numpy as np @@ -254,8 +255,6 @@ def pin_dataset(self, problem: Problem) -> None: def freeze(self, problem: Problem) -> EvalSpec: """Return a copy with the digest, dataset revision, and versions filled in.""" - import engibench - dataset_id = self.dataset_id or getattr(problem, "dataset_id", None) revision = self.dataset_revision or _dataset_revision(dataset_id) frozen = EvalSpec(**{**asdict(self), "dataset_id": dataset_id, "dataset_revision": revision}) @@ -268,7 +267,7 @@ def freeze(self, problem: Problem) -> EvalSpec: **asdict(frozen), "condition_digest": _digest(indices, conditions, ref_designs), "problem_conditions": tuple(problem.conditions_keys), - "engibench_version": getattr(engibench, "__version__", None), + "engibench_version": _engibench_version(), } ) @@ -327,6 +326,32 @@ def _condition_bytes(values: Any) -> bytes: return json.dumps(values, sort_keys=True, default=str).encode() +def _engibench_version() -> str: + """The EngiBench that produced a spec: its version, plus a git sha from a source checkout. + + The release version alone does not identify a problem definition -- 0.2.0 on + PyPI and 0.2.0 from `main` point `thermoelastic2d` at different datasets. A + change to `simulate` or `optimize` moves the optimality gaps without + touching any condition name, and only the revision records that. + """ + import engibench + + version = getattr(engibench, "__version__", "unknown") + source_root = Path(engibench.__file__).resolve().parent.parent + try: + sha = subprocess.run( + ["git", "-C", str(source_root), "rev-parse", "--short=12", "HEAD"], + capture_output=True, + text=True, + check=True, + timeout=5, + ).stdout.strip() + # Installed as a wheel rather than a checkout: the version is all there is. + except (OSError, subprocess.SubprocessError): + return version + return f"{version}+{sha}" if sha else version + + def _dataset_revision(dataset_id: str | None) -> str | None: """Current commit of a HuggingFace dataset repo, or None if it cannot be read.""" if dataset_id is None: diff --git a/engiopt/generators/vqgan/adapter.py b/engiopt/generators/vqgan/adapter.py index 2f0ee99b..67fbd730 100644 --- a/engiopt/generators/vqgan/adapter.py +++ b/engiopt/generators/vqgan/adapter.py @@ -11,8 +11,6 @@ from engiopt.core import Generator from engiopt.generators.vqgan.vqgan import VQGAN from engiopt.generators.vqgan.vqgan import VQGANTransformer -from engiopt.transforms import drop_constant -from engiopt.transforms import normalize from engiopt.transforms import resize_to if TYPE_CHECKING: @@ -25,15 +23,18 @@ class VQGANGenerator(Generator): """Discrete latent VQGAN whose codes are sampled autoregressively by a transformer. Conditions are themselves quantized (by a second, condition-side VQGAN) into - the transformer's start-of-sequence tokens, so this model needs the original - condition columns rather than a ready-made tensor: it may drop constant - columns and re-normalize exactly as it did during training. + the transformer's start-of-sequence tokens. Training may drop constant + condition columns and rescale the rest; both are recorded in the checkpoint, + so `Generator.sample` hands `_sample` the conditions already in that form and + this adapter does no preprocessing of its own. """ algo_id = "vqgan" conditional = True design_kinds = ("2d",) checkpoint_files = ("vqgan.pth", "transformer.pth") + """`cvqgan.pth` is loaded by `build` too, but only conditional runs write one, + so it cannot be required here. The package content hash covers it regardless.""" primary_state_key = "transformer" output_clip = (1e-3, 1.0) @@ -113,7 +114,7 @@ def _load_cvqgan(cls, package_root: str, *, config: dict[str, Any], device: th.d def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: """Sample a full grid of latent codes autoregressively, then decode it.""" if self.run_config["conditional"]: - start_tokens = self.net.encode_to_z(x=self._condition_tensor(conditions), is_c=True)[1] + start_tokens = self.net.encode_to_z(x=conditions.require_tensor(self.algo_id), is_c=True)[1] else: start_tokens = th.ones(n, 1, dtype=th.int64, device=self.device) * self.net.sos_token codes = self.net.sample( @@ -122,15 +123,3 @@ def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: steps=self.latent_size**2, ) return resize_to(data=self.net.z_to_image(codes), h=self.design_shape[0], w=self.design_shape[1]) - - def _condition_tensor(self, conditions: ConditionBatch) -> th.Tensor: - """Rebuild the condition tensor exactly as training preprocessed it.""" - dataset = conditions.dataset - if dataset is None: - return conditions.require_tensor(self.algo_id) - columns = dataset.column_names - if self.run_config["drop_constant_conditions"]: - dataset, columns = drop_constant(dataset, columns) - if self.run_config["normalize_conditions"]: - dataset, _mean, _std = normalize(dataset, columns) - return th.stack([th.as_tensor(dataset[column][:]).float() for column in columns], dim=1).to(self.device) diff --git a/engiopt/generators/vqgan/vqgan.py b/engiopt/generators/vqgan/vqgan.py index 6030608e..7b6bc578 100644 --- a/engiopt/generators/vqgan/vqgan.py +++ b/engiopt/generators/vqgan/vqgan.py @@ -699,9 +699,13 @@ def log_images(self, x: th.Tensor, c: th.Tensor, top_k: int | None = None) -> tu if args.drop_constant_conditions: training_ds, conditions = drop_constant(training_ds, conditions) - # Optionally normalize condition columns + # Optionally normalize condition columns. The statistics are recorded in the + # checkpoint so evaluation reproduces this exact scale rather than refitting + # on its own much smaller sample. + condition_stats = None if args.normalize_conditions: training_ds, mean, std = normalize(training_ds, conditions) + condition_stats = {"mean": mean.tolist(), "std": std.tolist()} n_conds = len(conditions) args.cond_dim = n_conds @@ -999,6 +1003,7 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: run_config=vars(args), **checkpoint_identity(args), condition_keys=conditions, + condition_stats=condition_stats, metadata={"stage": "cvqgan"}, primary_files=["cvqgan.pth"], ) @@ -1120,6 +1125,7 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: run_config=vars(args), **checkpoint_identity(args), condition_keys=conditions, + condition_stats=condition_stats, metadata={"stage": "vqgan"}, primary_files=["vqgan.pth"], ) @@ -1266,6 +1272,7 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: run_config=vars(args), **checkpoint_identity(args), condition_keys=conditions, + condition_stats=condition_stats, metadata={"stage": "transformer"}, primary_files=["transformer.pth"], ) diff --git a/tests/test_checkpoint_packages.py b/tests/test_checkpoint_packages.py index 5458e704..c55ba65d 100644 --- a/tests/test_checkpoint_packages.py +++ b/tests/test_checkpoint_packages.py @@ -261,3 +261,64 @@ def _sample(self, conditions: Any, n: int) -> th.Tensor: assert synchronized == [True] assert generator.last_sample_seconds is not None + + +def test_the_hash_covers_files_the_model_did_not_declare(tmp_path: Path) -> None: + """VQGAN's conditional variant loads `cvqgan.pth`, which is not in `checkpoint_files`. + + Hashing only the declared list would give two packages with different + condition encoders the same identity, even though they generate differently. + """ + package = _write_package(tmp_path / "pkg") + (package / "cvqgan.pth").write_bytes(b"encoder one") + before = checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth"], source="local" + ) + + (package / "cvqgan.pth").write_bytes(b"encoder two") + after = checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth"], source="local" + ) + + assert before.content_hash != after.content_hash + + +def test_the_hash_ignores_the_files_describing_the_package(tmp_path: Path) -> None: + """`metadata.json` records the path and revision, which differ between the two + locations one run writes -- so the same weights must not hash differently.""" + package = _write_package(tmp_path / "pkg", metadata={"hf_package_path": "beams2d/seed_1"}) + canonical = checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth"], source="local" + ) + + (package / "metadata.json").write_text(json.dumps({"hf_package_path": "beams2d/cfg_abc/seed_1"})) + config_scoped = checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth"], source="local" + ) + + assert canonical.content_hash == config_scoped.content_hash + + +def test_condition_statistics_travel_with_the_checkpoint(monkeypatch: pytest.MonkeyPatch) -> None: + """A model that rescales its conditions must record the scale it used.""" + uploads: list[dict[str, Any]] = [] + monkeypatch.setattr( + checkpoint_store, "_upload_package_to_hf", lambda **kwargs: uploads.append(kwargs["metadata"]) or "sha" + ) + monkeypatch.setattr(checkpoint_store.wandb, "run", None) + + checkpoint_store.save_checkpoint_package( + checkpoint_backend="hf", + hf_entity="org", + hf_repo_prefix="engiopt", + hf_private=False, + problem_id="beams2d", + algo="demo", + seed=1, + checkpoint_files={}, + run_config={}, + condition_keys=["volfrac"], + condition_stats={"mean": [0.4], "std": [0.2]}, + ) + + assert uploads[0]["condition_stats"] == {"mean": [0.4], "std": [0.2]} diff --git a/tests/test_condition_schema.py b/tests/test_condition_schema.py index 40ad4151..32bf34ab 100644 --- a/tests/test_condition_schema.py +++ b/tests/test_condition_schema.py @@ -94,3 +94,86 @@ def test_unlabelled_conditions_are_accepted(fake_problem: Any) -> None: """A bare tensor carries no column names, so there is nothing to contradict.""" generator = _generator(fake_problem, THERMOELASTIC_SCALARS) assert generator.sample(np.zeros((4, 3)), n=4).shape == (4, *fake_problem.design_space.shape) + + +# ---------------------------------------------------------------------- +# A model sees the conditions the way its training data looked +# ---------------------------------------------------------------------- + + +class _RecordingGenerator(_EchoGenerator): + """Captures the tensor `_sample` actually received.""" + + algo_id = "recorder" + received: th.Tensor | None = None + + def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: + type(self).received = conditions.require_tensor(self.algo_id) + return th.zeros((n, *self.design_shape)) + + +def _recorder(problem: Any, **kwargs: Any) -> _RecordingGenerator: + type(_RecordingGenerator).received = None + return _RecordingGenerator(problem=problem, problem_id="beams2d", seed=1, device=th.device("cpu"), **kwargs) + + +def test_a_model_trained_on_fewer_conditions_gets_them_projected(fake_problem: Any) -> None: + """VQGAN drops constant columns during training; the evaluator supplies all of them. + + The checkpoint's schema is authoritative, so the extra columns are projected + away rather than rejected. + """ + generator = _recorder(fake_problem, condition_keys=("volfrac",)) + batch = ConditionBatch(tensor=th.tensor([[0.3, 2.0], [0.5, 2.0]]), keys=("volfrac", "rmin")) + + generator.sample(batch, n=2) + + assert _RecordingGenerator.received is not None + assert _RecordingGenerator.received.shape == (2, 1) + assert _RecordingGenerator.received.flatten().tolist() == pytest.approx([0.3, 0.5]) + + +def test_columns_are_reordered_into_the_trained_order(fake_problem: Any) -> None: + """Channel k must be the condition the network learned in slot k.""" + generator = _recorder(fake_problem, condition_keys=("rmin", "volfrac")) + batch = ConditionBatch(tensor=th.tensor([[0.3, 2.0]]), keys=("volfrac", "rmin")) + + generator.sample(batch, n=1) + + assert _RecordingGenerator.received is not None + assert _RecordingGenerator.received.flatten().tolist() == pytest.approx([2.0, 0.3]) + + +def test_a_condition_the_model_needs_but_did_not_get_is_an_error(fake_problem: Any) -> None: + """A missing column cannot be reconstructed, so this one still refuses.""" + generator = _recorder(fake_problem, condition_keys=("volfrac", "brand_new")) + batch = ConditionBatch(tensor=th.tensor([[0.3, 2.0]]), keys=("volfrac", "rmin")) + + with pytest.raises(ValueError, match="brand_new"): + generator.sample(batch, n=1) + + +def test_recorded_normalization_is_replayed_not_refitted(fake_problem: Any) -> None: + """The training split's statistics, not the evaluation sample's. + + Refitting on the evaluation rows hands the network a different scale than it + was trained on -- wrong designs, no error. + """ + generator = _recorder(fake_problem, condition_keys=("volfrac",), condition_stats=([0.4], [0.2])) + batch = ConditionBatch(tensor=th.tensor([[0.6], [0.4]]), keys=("volfrac",)) + + generator.sample(batch, n=2) + + # (0.6 - 0.4) / 0.2 = 1.0 ; (0.4 - 0.4) / 0.2 = 0.0 -- and note the evaluation + # sample's own mean is 0.5, which would have produced a different answer. + assert _RecordingGenerator.received is not None + assert _RecordingGenerator.received.flatten().tolist() == pytest.approx([1.0, 0.0]) + + +def test_a_model_without_recorded_statistics_is_left_alone(fake_problem: Any) -> None: + """Most models feed conditions through unscaled, and must keep doing so.""" + generator = _recorder(fake_problem, condition_keys=("volfrac",)) + generator.sample(ConditionBatch(tensor=th.tensor([[0.6]]), keys=("volfrac",)), n=1) + + assert _RecordingGenerator.received is not None + assert _RecordingGenerator.received.flatten().tolist() == pytest.approx([0.6]) diff --git a/tests/test_conditions_and_gaps.py b/tests/test_conditions_and_gaps.py index a0012891..1c7e21fa 100644 --- a/tests/test_conditions_and_gaps.py +++ b/tests/test_conditions_and_gaps.py @@ -13,11 +13,15 @@ import pytest from engiopt.evaluation.context import EvaluationContext +from engiopt.evaluation.registry import METRICS from engiopt.evaluation.context import MultiObjectiveScalarizationError from engiopt.transforms import get_image_condition_keys from engiopt.transforms import get_scalar_condition_keys from tests.conftest import FakeViolations +# Importing the metrics package registers the built-ins. +import engiopt.evaluation.metrics # noqa: F401 # isort: skip + class _Problem: """Just enough of `Problem` for the condition-key helpers.""" @@ -346,3 +350,29 @@ def record(design: Any, config: dict[str, Any]) -> Any: assert isinstance(seen["fixed_elements"], np.ndarray) assert seen["rmin"] == 2.0 + + +def test_feasibility_survives_an_optimizer_that_refuses_the_design(fake_problem: Any) -> None: + """An invalid design is exactly when `viol` matters, and when `optimize` may raise. + + Computing feasibility inside the optimizer pass meant the metric was lost on + precisely the designs it was meant to flag. + """ + + def refuse(design: Any, config: Any = None) -> Any: + raise RuntimeError("optimizer rejected an invalid starting point") + + fake_problem.optimize = refuse + fake_problem.infeasible = True + ctx = _feasibility_context(fake_problem, 0.5) + + assert METRICS["viol"].fn(ctx) == 1.0 + + +def test_asking_only_for_viol_does_not_run_the_optimizer(fake_problem: Any) -> None: + """Feasibility is a constraint check, so it should not cost a full optimization.""" + ctx = _feasibility_context(fake_problem, 0.5) + + METRICS["viol"].fn(ctx) + + assert fake_problem.optimize_calls == 0 diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py index 4faa10f6..809b7f91 100644 --- a/tests/test_evaluation.py +++ b/tests/test_evaluation.py @@ -83,9 +83,13 @@ def test_dpp_prefers_varied_designs_over_duplicates(fake_problem: Any) -> None: def test_builtin_metrics_declare_their_cost() -> None: - """Simulator-backed metrics must be marked expensive, and the rest cheap.""" - assert {spec.name for spec in METRICS.select(cost="cheap")} == {"mmd", "dpp"} - assert {spec.name for spec in METRICS.select(cost="expensive")} == {"viol", "iog", "cog", "fog"} + """Simulator-backed metrics must be marked expensive, and the rest cheap. + + `viol` is cheap: feasibility describes the design as generated, so it is + judged by a constraint check rather than by running the optimizer. + """ + assert {spec.name for spec in METRICS.select(cost="cheap")} == {"mmd", "dpp", "viol"} + assert {spec.name for spec in METRICS.select(cost="expensive")} == {"iog", "cog", "fog"} def test_cheap_metrics_never_touch_the_solver(fake_problem: Any) -> None: diff --git a/tests/test_leaderboard_publishing.py b/tests/test_leaderboard_publishing.py index 25b8760f..917163ab 100644 --- a/tests/test_leaderboard_publishing.py +++ b/tests/test_leaderboard_publishing.py @@ -188,3 +188,127 @@ def upload_file(self, **kwargs: Any) -> None: monkeypatch.setattr("huggingface_hub.HfApi", _AlwaysStale()) with pytest.raises(HfHubHTTPError): lb.push_to_hub(pd.DataFrame([_row("newcomer", 0.1)]), "org/board", max_attempts=2) + + +# ---------------------------------------------------------------------- +# Distinguishing "no repo" from "no access" and "no file" +# ---------------------------------------------------------------------- + + +def test_a_private_repo_is_not_an_empty_board(monkeypatch: pytest.MonkeyPatch) -> None: + """The Hub reports a repo you cannot see as not-found; that must not read as empty.""" + api = _FakeApi() + + def raise_private(**_kwargs: Any) -> Any: + raise hub_error(RepositoryNotFoundError, "401 Client Error: Unauthorized", status=401) + + api.repo_info = raise_private # type: ignore[method-assign] + monkeypatch.setattr("huggingface_hub.HfApi", api) + with pytest.raises(RepositoryNotFoundError): + lb.load_from_hub("org/board") + + +def test_first_publish_into_an_existing_repo_keeps_its_revision(monkeypatch: pytest.MonkeyPatch) -> None: + """A repo with no leaderboard.csv still has a head commit, and it guards the first write. + + Discarding it means two jobs both publishing for the first time can each + overwrite the other. + """ + api = _FakeApi(revision="sha1") + monkeypatch.setattr("huggingface_hub.HfApi", api) + + def missing_file(**_kwargs: Any) -> str: + raise hub_error(EntryNotFoundError, "no leaderboard.csv", status=404) + + monkeypatch.setattr("huggingface_hub.hf_hub_download", missing_file) + + lb.push_to_hub(pd.DataFrame([_row("first", 0.1)]), "org/board") + + assert api.uploads[0]["parent_commit"] == "sha1" + + +# ---------------------------------------------------------------------- +# Ranking stays inside a problem +# ---------------------------------------------------------------------- + + +def test_ranks_restart_within_each_problem() -> None: + """A beams2d score and a photonics2d score are not competitors.""" + frame = pd.DataFrame( + [ + {"problem_id": "beams2d", "algo_id": "a", "config_fingerprint": "c", "seed": 1, + "spec_version": "v1", "mmd": 0.1}, + {"problem_id": "beams2d", "algo_id": "b", "config_fingerprint": "c", "seed": 1, + "spec_version": "v1", "mmd": 0.2}, + {"problem_id": "photonics2d", "algo_id": "a", "config_fingerprint": "c", "seed": 1, + "spec_version": "v1", "mmd": 0.3}, + {"problem_id": "photonics2d", "algo_id": "b", "config_fingerprint": "c", "seed": 1, + "spec_version": "v1", "mmd": 0.4}, + ] + ) # fmt: skip + ranked = lb.rank(frame, "mmd") + for problem in ("beams2d", "photonics2d"): + assert sorted(ranked[ranked["problem_id"] == problem]["rank"]) == [1, 2] + + +def test_ranks_restart_within_each_spec_version() -> None: + """Two protocols are two boards, even for the same problem.""" + frame = pd.DataFrame( + [ + {"problem_id": "p", "algo_id": "a", "config_fingerprint": "c", "seed": 1, "spec_version": "v1", + "mmd": 0.1}, + {"problem_id": "p", "algo_id": "b", "config_fingerprint": "c", "seed": 1, "spec_version": "v2", + "mmd": 0.2}, + ] + ) # fmt: skip + assert list(lb.rank(frame, "mmd")["rank"]) == [1, 1] + + +# ---------------------------------------------------------------------- +# Skipping work already done +# ---------------------------------------------------------------------- + + +def _published_row(**overrides: Any) -> pd.DataFrame: + row = {**_row("a", 0.1), "checkpoint_hash": "hash_old"} + row.update(overrides) + return pd.DataFrame([row]) + + +def test_retrained_weights_are_not_skipped() -> None: + """Same config and seed, different weights: the old row does not describe them.""" + assert not lb.already_evaluated( + _published_row(), + problem_id="p", + algo_id="a", + config_fingerprint="default", + seed=1, + spec_version="v1", + checkpoint_hash="hash_new", + ) + + +def test_the_same_weights_are_skipped() -> None: + """Re-running an unchanged checkpoint is the work `--skip-existing` exists to avoid.""" + assert lb.already_evaluated( + _published_row(), + problem_id="p", + algo_id="a", + config_fingerprint="default", + seed=1, + spec_version="v1", + checkpoint_hash="hash_old", + ) + + +def test_rows_predating_the_hash_column_are_still_skipped() -> None: + """An older board has no hash to compare, so it should not force a re-run.""" + assert lb.already_evaluated( + _published_row(checkpoint_hash=None), + problem_id="p", + algo_id="a", + config_fingerprint="default", + seed=1, + spec_version="v1", + checkpoint_hash="hash_new", + ) diff --git a/tests/test_specs.py b/tests/test_specs.py index 673d36f8..e2b7f30b 100644 --- a/tests/test_specs.py +++ b/tests/test_specs.py @@ -132,23 +132,18 @@ def test_digest_covers_array_valued_conditions() -> None: def test_committed_spec_reproduces_its_frozen_conditions(path: Any) -> None: """The check the reviewer ran by hand: every spec must resolve, not just parse. - `resolve` raises if the digest no longer matches, so this fails loudly when - a dataset moves out from under a published spec. + `resolve` raises if the digest no longer matches, or if the installed + EngiBench defines the problem differently than the spec was frozen against. - A spec frozen against a *different definition of the problem* is skipped - rather than failed: the installed EngiBench cannot draw those conditions at - all, so there is nothing here to verify. The skip names the mismatch, and a - digest mismatch under a matching definition still fails. + Both are failures, not skips. Skipping the mismatch would mean this job goes + green precisely when the EngiBench pin has stopped taking effect -- which is + the one thing it exists to catch. """ from engibench.utils.all_problems import BUILTIN_PROBLEMS spec = EvalSpec.load(str(path)) problem = BUILTIN_PROBLEMS[spec.problem_id]() problem.reset(seed=spec.condition_seed) - try: - spec.check_problem_definition(problem) - except ProblemDefinitionMismatchError as mismatch: - pytest.skip(str(mismatch)) resolved = spec.resolve(problem) From 2ed07436c729ae6e73e6ca5e5a50d20b4383e8a4 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Thu, 6 Aug 2026 10:35:24 +0200 Subject: [PATCH 09/24] Say so when a model cannot see a condition that varies in the evaluation set Projecting the evaluator's conditions onto a checkpoint's recorded schema is right when the dropped column is constant -- it carries no information, which is why training dropped it. beams2d's overhang_constraint is 0 across all fifty test cases, so VQGAN loses nothing the other models have. That is a property of today's data, not a guarantee. A column constant in a training split can vary in a future evaluation set, and then the model is blind to a requirement its competitors can see while nothing says so. Not an error: a model is allowed to ignore conditions, and the conditional-adherence metrics exist to expose exactly that. But it should never be silent. Co-Authored-By: Claude Opus 5 (1M context) --- engiopt/core.py | 21 +++++++++++++++++++++ tests/test_condition_schema.py | 24 ++++++++++++++++++++++++ tests/test_conditions_and_gaps.py | 2 +- 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/engiopt/core.py b/engiopt/core.py index 527e357d..2a6cff56 100644 --- a/engiopt/core.py +++ b/engiopt/core.py @@ -502,9 +502,30 @@ def _select_trained_columns(self, batch: ConditionBatch) -> ConditionBatch: ) if batch.tensor is None: return ConditionBatch(tensor=None, dataset=batch.dataset, keys=self.condition_keys) + self._warn_about_ignored_variation(batch) columns = [batch.keys.index(key) for key in self.condition_keys] return ConditionBatch(tensor=batch.tensor[:, columns], dataset=batch.dataset, keys=self.condition_keys) + def _warn_about_ignored_variation(self, batch: ConditionBatch) -> None: + """Say so when a projected-away condition actually varies across the test set. + + Dropping a condition that is constant everywhere costs nothing -- it + carries no information, which is why training dropped it. Dropping one + that *varies in the evaluation set* is different: this model cannot see a + requirement its competitors can, and its scores should be read knowing + that. Not an error, because a model is allowed to ignore conditions and + the conditional-adherence metrics exist to expose it -- but never silent. + """ + if batch.tensor is None: + return + ignored = [key for key in batch.keys if key not in self.condition_keys] + varying = [key for key in ignored if float(batch.tensor[:, batch.keys.index(key)].float().std().nan_to_num()) > 0] + if varying: + print( + f"[{self.algo_id}] was not trained on {varying}, which vary across this evaluation's " + "conditions. It cannot respond to them; its conditional-adherence scores reflect that." + ) + def _rescale(self, tensor: th.Tensor | None) -> th.Tensor | None: """Apply the checkpoint's recorded condition normalization, if it recorded any.""" if tensor is None or self._condition_stats is None: diff --git a/tests/test_condition_schema.py b/tests/test_condition_schema.py index 32bf34ab..f18ac209 100644 --- a/tests/test_condition_schema.py +++ b/tests/test_condition_schema.py @@ -177,3 +177,27 @@ def test_a_model_without_recorded_statistics_is_left_alone(fake_problem: Any) -> assert _RecordingGenerator.received is not None assert _RecordingGenerator.received.flatten().tolist() == pytest.approx([0.6]) + + +def test_dropping_a_constant_condition_is_silent(fake_problem: Any, capsys: Any) -> None: + """A condition that never varies carries no information, so losing it costs nothing. + + beams2d's `overhang_constraint` is 0 across every test case, which is why + VQGAN's training dropped it in the first place. + """ + generator = _recorder(fake_problem, condition_keys=("volfrac",)) + batch = ConditionBatch(tensor=th.tensor([[0.3, 2.0], [0.5, 2.0]]), keys=("volfrac", "constant")) + + generator.sample(batch, n=2) + + assert "constant" not in capsys.readouterr().out + + +def test_dropping_a_varying_condition_says_so(fake_problem: Any, capsys: Any) -> None: + """Here the model cannot see a requirement its competitors can; that must be visible.""" + generator = _recorder(fake_problem, condition_keys=("volfrac",)) + batch = ConditionBatch(tensor=th.tensor([[0.3, 1.0], [0.5, 9.0]]), keys=("volfrac", "rmin")) + + generator.sample(batch, n=2) + + assert "rmin" in capsys.readouterr().out diff --git a/tests/test_conditions_and_gaps.py b/tests/test_conditions_and_gaps.py index 1c7e21fa..05d93526 100644 --- a/tests/test_conditions_and_gaps.py +++ b/tests/test_conditions_and_gaps.py @@ -13,8 +13,8 @@ import pytest from engiopt.evaluation.context import EvaluationContext -from engiopt.evaluation.registry import METRICS from engiopt.evaluation.context import MultiObjectiveScalarizationError +from engiopt.evaluation.registry import METRICS from engiopt.transforms import get_image_condition_keys from engiopt.transforms import get_scalar_condition_keys from tests.conftest import FakeViolations From 43207eb2f1012764f8c8f67fc2e9070768229fe7 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Mon, 10 Aug 2026 13:56:33 +0200 Subject: [PATCH 10/24] Keep a checkpoint's identity a property of its weights Three ways a package could misrepresent itself, all of which reach the leaderboard: `metrics.json` was inside the content hash. It is written by `publish_checkpoint_metrics` into the package it describes, so attaching a score changed the identity of the thing scored: the next `--skip-existing` run would see a new hash, re-evaluate, rewrite the metrics, and change the hash again. No package has metrics attached yet, so this was latent. It is now in DESCRIPTIVE_FILES alongside run_config.json and metadata.json, which also stops it being served back as a loadable checkpoint file. A multi-stage model that dies partway leaves a package that looks complete. VQGAN uploads after each of its three stages so a crash does not lose the earlier work, which means an interrupted run publishes a run_config.json and a metadata.json with no usable weights -- indistinguishable, from the Hub, from a finished run. 4 of the pool's 102 VQGAN runs ended this way, holding vqgan.pth but no transformer.pth. `package_complete` is the distinction, and the load error now says which stage it reached and what the package holds. The same gap applies to the canonical path, which an intermediate stage of a default-hyperparameter run could claim. That publishes the bare model name pointing at something unloadable for as long as the remaining stages take -- hours, for VQGAN -- and permanently if the run dies first. Only a finished run may claim it now. Co-Authored-By: Claude Opus 5 (1M context) --- engiopt/checkpoint_store.py | 267 +++++++++++++++++++++++++++-- engiopt/generators/vqgan/vqgan.py | 7 +- tests/test_checkpoint_integrity.py | 216 +++++++++++++++++++++++ 3 files changed, 472 insertions(+), 18 deletions(-) create mode 100644 tests/test_checkpoint_integrity.py diff --git a/engiopt/checkpoint_store.py b/engiopt/checkpoint_store.py index 3ea619f6..0b27dc9b 100644 --- a/engiopt/checkpoint_store.py +++ b/engiopt/checkpoint_store.py @@ -23,12 +23,57 @@ from huggingface_hub import HfApi from huggingface_hub import snapshot_download - import wandb METRICS_FILE = "metrics.json" """Evaluation scores stored alongside the weights they describe.""" +RUN_CONFIG_FILE = "run_config.json" +"""Hyperparameters the run was launched with, enough to rebuild the network.""" + +METADATA_FILE = "metadata.json" +"""What this package is: paths, revisions, condition schema, completeness.""" + +PACKAGE_COMPLETE_FIELD = "package_complete" +"""Metadata flag: whether the run that wrote this package finished. + +Multi-stage models (VQGAN trains a condition autoencoder, then a VQGAN, then a +transformer) upload after each stage so a crash does not lose the earlier work. +That means an interrupted run leaves a package carrying a `run_config.json` and +a `metadata.json` -- indistinguishable, from the Hub, from a finished one. This +field is the distinction; single-stage models write `True` and never think about +it. +""" + +CONDITION_NORMALIZER_METADATA_FIELD = "condition_normalizer" +"""Metadata field recording the min/max a model scaled its *conditions* by.""" + +DESIGN_NORMALIZER_METADATA_FIELD = "design_normalizer" +"""Metadata field recording the min/max a model scaled its *designs* by. + +Several 1D models normalize designs into `[0, 1]` for training and denormalize +on the way out, with bounds fitted on the training split. Those bounds are part +of the model -- a design decoded against different bounds is a different design +-- but a plain `Normalizer` is not an `nn.Module`, so they never reached the +state dict. Recording them here is what makes such a checkpoint reproducible. +""" + +_MAX_LISTED_PACKAGES = 8 +"""How many sibling packages an error message names before summarizing. + +A sweep publishes hundreds; a wall of them is not more helpful than a handful +plus a count. +""" + +DESCRIPTIVE_FILES = frozenset({RUN_CONFIG_FILE, METADATA_FILE, METRICS_FILE}) +"""Package files that describe the checkpoint rather than being part of it. + +Neither hashed as weights nor served as loadable checkpoint files. `metrics.json` +belongs here for a reason worth stating: it is written *after* evaluation, into +the package it describes, so counting it would let scoring a checkpoint change +that checkpoint's identity. +""" + CheckpointBackend = Literal["hf", "none"] ModelSource = Literal["auto", "hf", "local"] @@ -86,6 +131,9 @@ def save_checkpoint_package( is_default_config: bool = True, condition_keys: list[str] | tuple[str, ...] | None = None, condition_stats: dict[str, list[float]] | None = None, + condition_normalizer: dict[str, Any] | None = None, + design_normalizer: dict[str, Any] | None = None, + package_complete: bool = True, ) -> dict[str, Any]: """Save a checkpoint package to HuggingFace. @@ -112,6 +160,16 @@ def save_checkpoint_package( evaluator can reproduce the same scale, instead of refitting it on the far smaller evaluation sample. + A model whose preprocessing is a min/max `Normalizer` rather than a mean/std + rescaling records it through `condition_normalizer` / `design_normalizer` + (see `engiopt.transforms.normalizer_state`). The principle is the same one: + whatever a model fitted on the training split travels with the weights, so + loading never has to re-derive it from whatever dataset happens to be current. + + Multi-stage models pass `package_complete=False` on every upload but the + last, so an interrupted run leaves a package that says it is unfinished + rather than one that merely fails to load. + W&B is no longer a checkpoint storage backend; the active W&B run still receives a summary pointing at the HF package for traceability. """ @@ -133,6 +191,9 @@ def save_checkpoint_package( metadata=metadata, condition_keys=condition_keys, condition_stats=condition_stats, + condition_normalizer=condition_normalizer, + design_normalizer=design_normalizer, + package_complete=package_complete, ) base_metadata.update(_build_wandb_run_metadata()) @@ -160,8 +221,11 @@ def save_checkpoint_package( ) # The canonical path defines what the bare model name resolves to, so - # only a default-hyperparameter run may claim it. - if is_default_config: + # only a finished default-hyperparameter run may claim it. Otherwise a + # multi-stage model publishes the bare model name at stage 0 and leaves + # it unloadable for as long as the remaining stages take -- hours, for + # VQGAN -- or permanently if the run dies, which 4 of 102 did. + if is_default_config and package_complete: info["hf_package_path"] = canonical_path info["hf_revision"] = _upload_package_to_hf( repo_id=repo_id, @@ -210,13 +274,19 @@ def resolve_named_checkpoint( FileNotFoundError: If no backend could supply the package. """ errors: list[str] = [] + repo_id = build_hf_repo_id(hf_entity, hf_repo_prefix, algo) if model_source in {"auto", "hf"}: try: return _resolve_hf_package( - repo_id=build_hf_repo_id(hf_entity, hf_repo_prefix, algo), + repo_id=repo_id, package_path=build_hf_package_path(problem_id, seed, extra_path_parts), required_files=required_files, ) + except FileNotFoundError as exc: + hint = _sibling_packages_hint(repo_id, problem_id, seed, extra_path_parts) + if model_source == "hf": + raise FileNotFoundError(f"{exc}{hint}") from exc + errors.append(f"hf: {exc}{hint}") except Exception as exc: if model_source == "hf": raise @@ -229,6 +299,133 @@ def resolve_named_checkpoint( raise FileNotFoundError(f"Unable to resolve checkpoint for {algo}/{problem_id}/seed_{seed}: {attempted}") +def list_packages(repo_id: str, problem_id: str | None = None) -> list[str]: + """Every checkpoint package path in a model repo, optionally for one problem. + + A package is a directory holding a `run_config.json`, which is what every + upload writes and what loading requires. + + Args: + repo_id: HF model repo, e.g. `IDEALLab/engiopt-cgan-cnn-2d`. + problem_id: Restrict to one problem's packages. + + Returns: + Sorted package paths, e.g. `["beams2d/cfg_023dd1fb/seed_42", ...]`. + """ + api = HfApi() + prefix = f"{_sanitize_path_component(problem_id)}/" if problem_id else "" + return sorted( + entry[: -len(f"/{RUN_CONFIG_FILE}")] + for entry in api.list_repo_files(repo_id=repo_id, repo_type="model") + if entry.endswith(f"/{RUN_CONFIG_FILE}") and entry.startswith(prefix) + ) + + +def _sibling_packages_hint(repo_id: str, problem_id: str, seed: int, extra_path_parts: list[str] | None) -> str: + """Name the packages that *do* exist when the requested one does not. + + The common failure is asking for the canonical `{problem}/seed_N` when only + sweep configurations were ever published -- every arm of a sweep varies some + hyperparameter, so none of them is the default run that claims the canonical + path. "Not found" alone sends the reader looking for a broken path; listing + the neighbours shows them what to pass to `--config-fingerprints`, or that + the canonical run is simply missing and needs training or promoting. + """ + try: + available = list_packages(repo_id, problem_id) + except Exception: # noqa: BLE001 - a hint must never replace the original error + return "" + if not available: + return f" No packages for {problem_id!r} exist in {repo_id}." + shown = ", ".join(available[:_MAX_LISTED_PACKAGES]) + ( + f", ... ({len(available)} total)" if len(available) > _MAX_LISTED_PACKAGES else "" + ) + hint = f" That repo does hold: {shown}." + if not extra_path_parts: + hint += ( + f" Nothing claims the canonical {problem_id}/seed_{seed}, which only a run using the training " + "script's default hyperparameters writes. Pass --config-fingerprints to score one of the above, " + "or run `python -m engiopt.promote_checkpoint` to make one of them canonical." + ) + return hint + + +def promote_to_canonical( + *, + hf_entity: str, + hf_repo_prefix: str, + problem_id: str, + algo: str, + seed: int, + config_fingerprint: str, + token: str | None = None, +) -> str: + """Copy a configuration's package to the canonical `{problem_id}/seed_{seed}` path. + + The canonical path is what a bare model name resolves to, and normally only + a default-hyperparameter run writes it. A sweep has no such run -- every arm + varies something -- so a sweep alone leaves the canonical path empty and the + documented `--seeds 1` command resolving nothing. This promotes one already + trained arm into that role without retraining it. + + The promoted package keeps its own `config_fingerprint` in metadata, so a + leaderboard row still records which configuration actually earned the score; + only the *address* changes. + + Args: + hf_entity: HF org/user holding the checkpoint repos. + hf_repo_prefix: Prefix of the per-model-family repo. + problem_id: Problem the checkpoint was trained on. + algo: Model family, selecting the repo. + seed: Training seed; the promoted package keeps it. + config_fingerprint: Which configuration to promote. + token: HF token; falls back to the ambient login. + + Returns: + The canonical path written. + + Raises: + FileNotFoundError: If the source package is absent or incomplete. + """ + repo_id = build_hf_repo_id(hf_entity, hf_repo_prefix, algo) + source_path = build_hf_package_path(problem_id, seed, [f"cfg_{config_fingerprint}"]) + canonical_path = build_hf_package_path(problem_id, seed) + + snapshot = snapshot_download(repo_id=repo_id, repo_type="model", allow_patterns=[f"{source_path}/*"], token=token) + source_dir = Path(snapshot) / source_path + if not source_dir.is_dir(): + raise FileNotFoundError(f"Cannot promote {repo_id}/{source_path}: it does not exist.") + + metadata_path = source_dir / METADATA_FILE + metadata = _read_json(str(metadata_path)) if metadata_path.exists() else {} + if metadata.get(PACKAGE_COMPLETE_FIELD) is False: + raise FileNotFoundError( + f"Refusing to promote {repo_id}/{source_path}: it is marked incomplete, so the canonical " + "path would resolve to a package that cannot be loaded." + ) + + api = HfApi(token=token) + with tempfile.TemporaryDirectory() as tmpdir: + stage_dir = Path(tmpdir) + for entry in source_dir.iterdir(): + # Metrics describe the score of the package at its old address; the + # promoted copy has not been evaluated under its new name yet. + if entry.is_file() and entry.name != METRICS_FILE: + shutil.copy2(entry, stage_dir / entry.name) + _write_json( + stage_dir / METADATA_FILE, + {**metadata, "hf_package_path": canonical_path, "promoted_from": source_path}, + ) + api.upload_folder( + repo_id=repo_id, + repo_type="model", + folder_path=str(stage_dir), + path_in_repo=canonical_path, + commit_message=f"Promote cfg_{config_fingerprint} to canonical {canonical_path}", + ) + return canonical_path + + def resolve_checkpoint_reference( *, model_source: ModelSource, @@ -316,9 +513,16 @@ def hash_package_contents(root_dir: str) -> str: that file generate differently. Hashing the declared list would give them the same identity. - `run_config.json` and `metadata.json` are excluded: they describe the - package rather than being weights, and `metadata.json` records paths and - revisions that differ between the two locations one run writes. + Files that *describe* the package rather than constitute it are excluded: + + - `run_config.json`, `metadata.json` -- metadata records the package path and + revision, which differ between the two locations one run writes, so + including it would give identical weights two different hashes. + - `metrics.json` -- written by `publish_checkpoint_metrics` *after* the + weights, into this same directory. Including it would make attaching a + score change the identity of the thing scored: the next `--skip-existing` + run would see a new hash, re-evaluate, rewrite the metrics, and change the + hash again. The hash must be a property of the weights alone. Args: root_dir: Local directory holding the package. @@ -326,9 +530,10 @@ def hash_package_contents(root_dir: str) -> str: Returns: A 16-character hash over the file names and their bytes. """ - described = {"run_config.json", "metadata.json"} weight_files = sorted( - entry for entry in os.listdir(root_dir) if entry not in described and os.path.isfile(os.path.join(root_dir, entry)) + entry + for entry in os.listdir(root_dir) + if entry not in DESCRIPTIVE_FILES and os.path.isfile(os.path.join(root_dir, entry)) ) hasher = hashlib.sha256() for file_name in weight_files: @@ -342,17 +547,17 @@ def hash_package_contents(root_dir: str) -> str: def _load_package_from_directory( *, root_dir: str, required_files: list[str], source: Literal["hf", "local"], revision: str | None = None ) -> ResolvedCheckpoint: - run_config_path = os.path.join(root_dir, "run_config.json") - metadata_path = os.path.join(root_dir, "metadata.json") + run_config_path = os.path.join(root_dir, RUN_CONFIG_FILE) + metadata_path = os.path.join(root_dir, METADATA_FILE) if not os.path.exists(run_config_path): - raise FileNotFoundError(f"Missing run_config.json in {root_dir}") + raise FileNotFoundError(f"Missing {RUN_CONFIG_FILE} in {root_dir}") run_config = _read_json(run_config_path) metadata = _read_json(metadata_path) if os.path.exists(metadata_path) else {} package_files = required_files or _discover_package_files(root_dir, metadata) files = {file_name: os.path.join(root_dir, file_name) for file_name in package_files} - for file_name, file_path in files.items(): - if not os.path.exists(file_path): - raise FileNotFoundError(f"Missing checkpoint file {file_name} in {root_dir}") + missing = [file_name for file_name, file_path in files.items() if not os.path.exists(file_path)] + if missing: + raise FileNotFoundError(_missing_files_message(root_dir, missing, metadata)) return ResolvedCheckpoint( source=source, root_dir=root_dir, @@ -364,16 +569,36 @@ def _load_package_from_directory( ) +def _missing_files_message(root_dir: str, missing: list[str], metadata: dict[str, Any]) -> str: + """Explain a package that is missing weights, saying so when the run never finished. + + A multi-stage model uploads after each stage, so a run killed partway leaves + a package that has a `run_config.json` and a `metadata.json` and therefore + looks like a checkpoint from the outside. `package_complete` is what tells + the two apart, and saying which stage it stopped at turns "file not found" + into something the reader can act on. + """ + present = sorted(entry for entry in os.listdir(root_dir) if entry not in DESCRIPTIVE_FILES) + lines = [f"Missing checkpoint file(s) {missing} in {root_dir}; the package holds {present or 'no weight files'}."] + if metadata.get(PACKAGE_COMPLETE_FIELD) is False: + stage = metadata.get("stage") + reached = f" It reached the {stage!r} stage." if stage else "" + lines.append( + f"This package is marked incomplete: the training run that wrote it did not finish.{reached} " + "Re-run training to completion; a partial package cannot be loaded." + ) + return " ".join(lines) + + def _discover_package_files(root_dir: str, metadata: dict[str, Any]) -> list[str]: primary_files = metadata.get("primary_files") if isinstance(primary_files, list) and primary_files: return [str(file_name) for file_name in primary_files] - excluded_names = {"metadata.json", "run_config.json"} discovered = [ entry for entry in os.listdir(root_dir) - if os.path.isfile(os.path.join(root_dir, entry)) and entry not in excluded_names + if os.path.isfile(os.path.join(root_dir, entry)) and entry not in DESCRIPTIVE_FILES ] if not discovered: raise FileNotFoundError(f"No checkpoint files found in {root_dir}") @@ -402,6 +627,9 @@ def _build_metadata( metadata: dict[str, Any] | None, condition_keys: list[str] | tuple[str, ...] | None = None, condition_stats: dict[str, list[float]] | None = None, + condition_normalizer: dict[str, Any] | None = None, + design_normalizer: dict[str, Any] | None = None, + package_complete: bool = True, ) -> dict[str, Any]: payload = dict(metadata or {}) payload.update( @@ -412,12 +640,17 @@ def _build_metadata( "checkpoint_backend": checkpoint_backend, "checkpoint_files": sorted(checkpoint_files), "primary_files": primary_files or sorted(checkpoint_files), + PACKAGE_COMPLETE_FIELD: package_complete, } ) if condition_keys is not None: payload["condition_keys"] = list(condition_keys) if condition_stats is not None: payload["condition_stats"] = {key: [float(v) for v in values] for key, values in condition_stats.items()} + if condition_normalizer is not None: + payload[CONDITION_NORMALIZER_METADATA_FIELD] = condition_normalizer + if design_normalizer is not None: + payload[DESIGN_NORMALIZER_METADATA_FIELD] = design_normalizer return payload diff --git a/engiopt/generators/vqgan/vqgan.py b/engiopt/generators/vqgan/vqgan.py index 7b6bc578..a637127b 100644 --- a/engiopt/generators/vqgan/vqgan.py +++ b/engiopt/generators/vqgan/vqgan.py @@ -30,6 +30,7 @@ from torch.nn import functional as f import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -52,7 +53,6 @@ from engiopt.transforms import drop_constant from engiopt.transforms import normalize from engiopt.transforms import resize_to -import wandb @dataclass @@ -1005,6 +1005,9 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: condition_keys=conditions, condition_stats=condition_stats, metadata={"stage": "cvqgan"}, + # Two stages still to run; a package this partial cannot be + # loaded and must not look like one that can. + package_complete=False, primary_files=["cvqgan.pth"], ) @@ -1127,6 +1130,8 @@ def sample_designs_transformer(n_designs: int) -> tuple[th.Tensor, th.Tensor]: condition_keys=conditions, condition_stats=condition_stats, metadata={"stage": "vqgan"}, + # The transformer stage has not run yet. + package_complete=False, primary_files=["vqgan.pth"], ) diff --git a/tests/test_checkpoint_integrity.py b/tests/test_checkpoint_integrity.py new file mode 100644 index 00000000..27b8fdd4 --- /dev/null +++ b/tests/test_checkpoint_integrity.py @@ -0,0 +1,216 @@ +"""Tests for what a checkpoint package's identity covers, and for partial packages. + +Two properties are load-bearing for the leaderboard: + +- A checkpoint's hash identifies its *weights*. Attaching a score to a package + must not change the identity of the thing that was scored, or `--skip-existing` + chases its own tail forever. +- A package written by a run that never finished must say so. Multi-stage models + upload after each stage, so an interrupted run leaves something that has every + file a loader looks for except the weights. +""" + +from __future__ import annotations + +import json +from typing import Any, TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +from engiopt import checkpoint_store + + +def _write_package( + root: Path, + *, + weights: bytes = b"weights", + files: dict[str, bytes] | None = None, + metadata: dict[str, Any] | None = None, +) -> Path: + root.mkdir(parents=True, exist_ok=True) + (root / "generator.pth").write_bytes(weights) + for name, payload in (files or {}).items(): + (root / name).write_bytes(payload) + (root / "run_config.json").write_text(json.dumps({"latent_dim": 32})) + (root / "metadata.json").write_text(json.dumps(metadata or {"algo": "demo", "seed": 1})) + return root + + +# ---------------------------------------------------------------------- +# Attaching metrics must not change what the weights hash to +# ---------------------------------------------------------------------- + + +def test_attaching_metrics_does_not_change_the_checkpoint_hash(tmp_path: Path) -> None: + """`publish_checkpoint_metrics` writes into the package it describes. + + If that file counted as package content, evaluating a checkpoint would change + its hash, `--skip-existing` would see new weights, re-evaluate, rewrite the + metrics, and change the hash again -- an evaluation loop with no fixed point. + """ + package = _write_package(tmp_path / "pkg") + before = checkpoint_store.hash_package_contents(str(package)) + + (package / checkpoint_store.METRICS_FILE).write_text(json.dumps({"mmd": 0.04, "dpp": 1e-29})) + after = checkpoint_store.hash_package_contents(str(package)) + + assert before == after + + # And a second, different evaluation still does not move it. + (package / checkpoint_store.METRICS_FILE).write_text(json.dumps({"mmd": 0.99})) + assert checkpoint_store.hash_package_contents(str(package)) == before + + +def test_the_hash_still_covers_undeclared_weight_files(tmp_path: Path) -> None: + """Excluding metrics must not weaken the guarantee that weights are covered.""" + package = _write_package(tmp_path / "pkg", files={"cvqgan.pth": b"conditional encoder"}) + before = checkpoint_store.hash_package_contents(str(package)) + + (package / "cvqgan.pth").write_bytes(b"a different conditional encoder") + + assert checkpoint_store.hash_package_contents(str(package)) != before + + +def test_metrics_are_not_served_as_a_checkpoint_file(tmp_path: Path) -> None: + """A package that has been evaluated must not offer `metrics.json` as a weight file.""" + package = _write_package(tmp_path / "pkg") + (package / checkpoint_store.METRICS_FILE).write_text(json.dumps({"mmd": 0.04})) + + resolved = checkpoint_store._load_package_from_directory(root_dir=str(package), required_files=[], source="local") + + assert checkpoint_store.METRICS_FILE not in resolved.files + + +# ---------------------------------------------------------------------- +# A package from an unfinished run has to say so +# ---------------------------------------------------------------------- + + +def test_an_incomplete_package_explains_itself(tmp_path: Path) -> None: + """The VQGAN case: stage 0 uploaded, the run died, the package looks loadable.""" + package = tmp_path / "pkg" + package.mkdir() + (package / "cvqgan.pth").write_bytes(b"condition autoencoder only") + (package / "run_config.json").write_text(json.dumps({"conditional": True})) + (package / "metadata.json").write_text(json.dumps({"stage": "cvqgan", checkpoint_store.PACKAGE_COMPLETE_FIELD: False})) + + with pytest.raises(FileNotFoundError) as excinfo: + checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["vqgan.pth", "transformer.pth"], source="local" + ) + + message = str(excinfo.value) + assert "incomplete" in message + assert "did not finish" in message + assert "cvqgan" in message + # It also says what the package does contain, so the reader can tell how far it got. + assert "cvqgan.pth" in message + + +def test_a_complete_package_missing_a_file_reports_what_is_there(tmp_path: Path) -> None: + """A finished run missing a file is a different problem and must not blame the run.""" + package = _write_package(tmp_path / "pkg") + + with pytest.raises(FileNotFoundError) as excinfo: + checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth", "absent.pth"], source="local" + ) + + message = str(excinfo.value) + assert "absent.pth" in message + assert "generator.pth" in message + assert "incomplete" not in message + + +def test_an_incomplete_stage_never_claims_the_canonical_path(monkeypatch: pytest.MonkeyPatch) -> None: + """The canonical path is what a bare model name resolves to; a partial run must not own it. + + Without this, a multi-stage default-hyperparameter run publishes the bare + model name at stage 0, and it stays unloadable for as long as the remaining + stages run -- hours, for VQGAN -- or permanently if the run dies before the + last one, as 4 of the pool's 102 VQGAN runs did. + """ + uploaded: list[str] = [] + + def fake_upload(*, package_path: str, **_: Any) -> str: + uploaded.append(package_path) + return "deadbeef" + + monkeypatch.setattr(checkpoint_store, "_upload_package_to_hf", fake_upload) + monkeypatch.setattr(checkpoint_store.wandb, "run", None) + + info = checkpoint_store.save_checkpoint_package( + checkpoint_backend="hf", + hf_entity="TestOrg", + hf_repo_prefix="engiopt", + hf_private=False, + problem_id="beams2d", + algo="vqgan", + seed=42, + checkpoint_files={"cvqgan.pth": "cvqgan.pth"}, + run_config={"latent_dim": 16}, + config_fingerprint="abc12345", + is_default_config=True, + package_complete=False, + ) + + assert uploaded == ["beams2d/cfg_abc12345/seed_42"] + assert info["hf_package_path"] is None + + +def test_a_finished_default_run_still_claims_the_canonical_path(monkeypatch: pytest.MonkeyPatch) -> None: + """The completeness guard must not cost ordinary single-stage models their canonical path.""" + uploaded: list[str] = [] + + def fake_upload(*, package_path: str, **_: Any) -> str: + uploaded.append(package_path) + return "deadbeef" + + monkeypatch.setattr(checkpoint_store, "_upload_package_to_hf", fake_upload) + monkeypatch.setattr(checkpoint_store.wandb, "run", None) + + checkpoint_store.save_checkpoint_package( + checkpoint_backend="hf", + hf_entity="TestOrg", + hf_repo_prefix="engiopt", + hf_private=False, + problem_id="beams2d", + algo="cgan_cnn_2d", + seed=1, + checkpoint_files={"generator.pth": "generator.pth"}, + run_config={"latent_dim": 16}, + config_fingerprint="abc12345", + is_default_config=True, + ) + + assert uploaded == ["beams2d/cfg_abc12345/seed_1", "beams2d/seed_1"] + + +def test_metadata_records_completeness_for_ordinary_models(monkeypatch: pytest.MonkeyPatch) -> None: + """A single-stage model should not have to opt in to being considered finished.""" + written: list[dict[str, Any]] = [] + + def fake_upload(*, metadata: dict[str, Any], **_: Any) -> str: + written.append(metadata) + return "deadbeef" + + monkeypatch.setattr(checkpoint_store, "_upload_package_to_hf", fake_upload) + monkeypatch.setattr(checkpoint_store.wandb, "run", None) + + checkpoint_store.save_checkpoint_package( + checkpoint_backend="hf", + hf_entity="TestOrg", + hf_repo_prefix="engiopt", + hf_private=False, + problem_id="beams2d", + algo="cgan_cnn_2d", + seed=1, + checkpoint_files={"generator.pth": "generator.pth"}, + run_config={}, + is_default_config=True, + ) + + assert all(entry[checkpoint_store.PACKAGE_COMPLETE_FIELD] is True for entry in written) From 562feb2b66489e2bef0e7dd70afccebd23787fb8 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Fri, 7 Aug 2026 15:18:01 +0200 Subject: [PATCH 11/24] Make the five remaining adapters replay their fitted preprocessing VQGAN records the statistics it normalized its conditions by, and the contract replays them. cgan_1d, cgan_bezier, gan_1d, gan_bezier, and diffusion_1d did not: they rebuilt their normalizers from whatever `problem.dataset["train"]` currently holds, at load time. cgan_1d's docstring said so outright. Those bounds are as much a part of the model as its weights -- a design decoded against different bounds is a different design -- but `Normalizer` is a plain class rather than an nn.Module, so nothing ever put them in the state dict. The result is that the same checkpoint, with the same content hash, produces different designs after a dataset revision, with nothing to show it happened. The eval spec's pinned revision does not help: it governs the test conditions, not the training split these read. `normalizer_state` / `load_normalizer_state` record and restore the fitted bounds, duck-typed on min_val/max_val/eps so the five identical copies of `Normalizer` do not have to be unified first -- that is a change to training code this does not need. A checkpoint with nothing recorded is left to fit from the dataset exactly as before, which is what those runs really did. Co-Authored-By: Claude Opus 5 (1M context) --- engiopt/core.py | 28 +++- engiopt/generators/cgan_1d/adapter.py | 11 +- engiopt/generators/cgan_1d/cgan_1d.py | 5 +- engiopt/generators/cgan_bezier/adapter.py | 7 + engiopt/generators/cgan_bezier/cgan_bezier.py | 5 +- engiopt/generators/diffusion_1d/adapter.py | 3 + .../generators/diffusion_1d/diffusion_1d.py | 4 +- engiopt/generators/gan_1d/adapter.py | 3 + engiopt/generators/gan_1d/gan_1d.py | 4 +- engiopt/generators/gan_bezier/adapter.py | 5 + engiopt/generators/gan_bezier/gan_bezier.py | 4 +- engiopt/transforms.py | 65 +++++++- tests/test_normalizer_replay.py | 141 ++++++++++++++++++ 13 files changed, 274 insertions(+), 11 deletions(-) create mode 100644 tests/test_normalizer_replay.py diff --git a/engiopt/core.py b/engiopt/core.py index 2a6cff56..c257aa7b 100644 --- a/engiopt/core.py +++ b/engiopt/core.py @@ -26,6 +26,8 @@ import numpy as np import torch as th +from engiopt.checkpoint_store import CONDITION_NORMALIZER_METADATA_FIELD +from engiopt.checkpoint_store import DESIGN_NORMALIZER_METADATA_FIELD from engiopt.checkpoint_store import resolve_named_checkpoint from engiopt.transforms import condition_keys as scalar_condition_keys @@ -214,6 +216,30 @@ def condition_stats_for(resolved: ResolvedCheckpoint | None) -> tuple[list[float return [float(v) for v in stats["mean"]], [float(v) for v in stats["std"]] +def recorded_condition_normalizer(resolved: ResolvedCheckpoint | None) -> dict[str, Any] | None: + """The min/max bounds a checkpoint scaled its conditions by, if it recorded any.""" + return _recorded_normalizer(resolved, CONDITION_NORMALIZER_METADATA_FIELD) + + +def recorded_design_normalizer(resolved: ResolvedCheckpoint | None) -> dict[str, Any] | None: + """The min/max bounds a checkpoint scaled its designs by, if it recorded any.""" + return _recorded_normalizer(resolved, DESIGN_NORMALIZER_METADATA_FIELD) + + +def _recorded_normalizer(resolved: ResolvedCheckpoint | None, field: str) -> dict[str, Any] | None: + """Read one recorded normalizer, returning None for checkpoints predating it. + + None means "fall back to fitting from the dataset", which reproduces exactly + what those older runs did. + """ + if resolved is None: + return None + state = resolved.metadata.get(field) + if not isinstance(state, dict) or "min" not in state or "max" not in state: + return None + return state + + def design_shape_of(problem: Problem) -> tuple[int, ...]: """Shape of a single design for `problem`. @@ -582,7 +608,7 @@ def resolve_checkpoint( Args: problem_id: EngiBench problem the checkpoint was trained on. seed: Training seed to load. - model_source: Backend preference: `auto`, `hf`, `wandb`, or `local`. + model_source: Backend preference: `auto`, `hf`, or `local`. hf_entity: HF org/user holding the checkpoint repos. hf_repo_prefix: Prefix of the per-model-family HF repo. local_model_dir: Directory to load from, for `local` sources. diff --git a/engiopt/generators/cgan_1d/adapter.py b/engiopt/generators/cgan_1d/adapter.py index 3327c296..a106c776 100644 --- a/engiopt/generators/cgan_1d/adapter.py +++ b/engiopt/generators/cgan_1d/adapter.py @@ -10,8 +10,11 @@ from engiopt.core import ConditionBatch from engiopt.core import design_shape_of from engiopt.core import Generator +from engiopt.core import recorded_condition_normalizer +from engiopt.core import recorded_design_normalizer from engiopt.generators.cgan_1d.cgan_1d import Generator as CGAN1DNet from engiopt.generators.cgan_1d.cgan_1d import prepare_data +from engiopt.transforms import load_normalizer_state if TYPE_CHECKING: from engibench.core import Problem @@ -22,8 +25,10 @@ class CGAN1D(Generator): """Conditional GAN over 1D / dict-valued designs such as airfoils. - The normalizers are part of the model: they are rebuilt from the training - dataset rather than stored in the checkpoint, so loading touches the dataset. + The normalizers are part of the model, so the checkpoint records the bounds + fitted during training and loading replays them. They are still fitted from + the dataset first, which is what supplies their shape and what checkpoints + written before the bounds were recorded fall back to. """ algo_id = "cgan_1d" @@ -42,6 +47,8 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device """Load a trained cGAN-1D generator, rebuilding its normalizers.""" config = resolved.run_config _, conds_normalizer, design_normalizer = prepare_data(problem, device) + conds_normalizer = load_normalizer_state(conds_normalizer, recorded_condition_normalizer(resolved), device) + design_normalizer = load_normalizer_state(design_normalizer, recorded_design_normalizer(resolved), device) net = CGAN1DNet( latent_dim=config["latent_dim"], n_conds=len(condition_keys_for(problem, resolved)), diff --git a/engiopt/generators/cgan_1d/cgan_1d.py b/engiopt/generators/cgan_1d/cgan_1d.py index b841f6f1..9a81231e 100644 --- a/engiopt/generators/cgan_1d/cgan_1d.py +++ b/engiopt/generators/cgan_1d/cgan_1d.py @@ -20,6 +20,7 @@ from torchvision import transforms import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -28,7 +29,7 @@ from engiopt.reproducibility import seed_training from engiopt.transforms import condition_keys from engiopt.transforms import flatten_dict_factory -import wandb +from engiopt.transforms import normalizer_state if TYPE_CHECKING: from engibench.utils.problem import Problem @@ -421,6 +422,8 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor]: **checkpoint_identity(args), condition_keys=cond_keys, primary_files=["generator.pth"], + condition_normalizer=normalizer_state(conds_normalizer), + design_normalizer=normalizer_state(design_normalizer), ) wandb.finish() diff --git a/engiopt/generators/cgan_bezier/adapter.py b/engiopt/generators/cgan_bezier/adapter.py index f2947d70..2fb1d839 100644 --- a/engiopt/generators/cgan_bezier/adapter.py +++ b/engiopt/generators/cgan_bezier/adapter.py @@ -9,10 +9,13 @@ from engiopt.core import condition_keys_for from engiopt.core import ConditionBatch from engiopt.core import Generator +from engiopt.core import recorded_condition_normalizer +from engiopt.core import recorded_design_normalizer from engiopt.generators.cgan_bezier.cgan_bezier import _EPS from engiopt.generators.cgan_bezier.cgan_bezier import Generator as CBezierGANNet from engiopt.generators.cgan_bezier.cgan_bezier import prepare_data from engiopt.transforms import flatten_dict_factory +from engiopt.transforms import load_normalizer_state if TYPE_CHECKING: from engibench.core import Problem @@ -47,6 +50,10 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device """Load a trained conditional BezierGAN, rebuilding its normalizers.""" config = resolved.run_config _, conds_normalizer, design_scalars_normalizer, _keys = prepare_data(problem, device) + conds_normalizer = load_normalizer_state(conds_normalizer, recorded_condition_normalizer(resolved), device) + design_scalars_normalizer = load_normalizer_state( + design_scalars_normalizer, recorded_design_normalizer(resolved), device + ) net = CBezierGANNet( latent_dim=config["latent_dim"], noise_dim=config["noise_dim"], diff --git a/engiopt/generators/cgan_bezier/cgan_bezier.py b/engiopt/generators/cgan_bezier/cgan_bezier.py index f104bbd3..6216d569 100644 --- a/engiopt/generators/cgan_bezier/cgan_bezier.py +++ b/engiopt/generators/cgan_bezier/cgan_bezier.py @@ -18,6 +18,7 @@ from torch import nn import torch.nn.functional as f import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -25,7 +26,7 @@ from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import condition_keys -import wandb +from engiopt.transforms import normalizer_state if TYPE_CHECKING: from collections.abc import Callable @@ -694,6 +695,8 @@ def sample_designs(n_designs: int) -> tuple[th.Tensor, th.Tensor, th.Tensor]: **checkpoint_identity(args), condition_keys=cond_keys, primary_files=["bezier_generator.pth"], + condition_normalizer=normalizer_state(conds_normalizer), + design_normalizer=normalizer_state(design_scalars_normalizer), ) if args.track: diff --git a/engiopt/generators/diffusion_1d/adapter.py b/engiopt/generators/diffusion_1d/adapter.py index 4555b49e..d6eed741 100644 --- a/engiopt/generators/diffusion_1d/adapter.py +++ b/engiopt/generators/diffusion_1d/adapter.py @@ -12,7 +12,9 @@ from engiopt.core import ConditionBatch from engiopt.core import design_shape_of from engiopt.core import Generator +from engiopt.core import recorded_design_normalizer from engiopt.generators.diffusion_1d.diffusion_1d import prepare_data +from engiopt.transforms import load_normalizer_state if TYPE_CHECKING: from engibench.core import Problem @@ -51,6 +53,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device padding_size = _padding_for(design_shape_of(problem)[0]) padded_shape = (design_shape_of(problem)[0] + padding_size,) _, design_normalizer = prepare_data(problem, padding_size, device) + design_normalizer = load_normalizer_state(design_normalizer, recorded_design_normalizer(resolved), device) unet = Unet1D(dim=config["unet_dim"], channels=config["n_channels"]).to(device) diffusion = GaussianDiffusion1D( unet, diff --git a/engiopt/generators/diffusion_1d/diffusion_1d.py b/engiopt/generators/diffusion_1d/diffusion_1d.py index 9c2eeae7..fac4ff4d 100644 --- a/engiopt/generators/diffusion_1d/diffusion_1d.py +++ b/engiopt/generators/diffusion_1d/diffusion_1d.py @@ -20,6 +20,7 @@ from torchvision import transforms import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -27,7 +28,7 @@ from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import flatten_dict_factory -import wandb +from engiopt.transforms import normalizer_state if TYPE_CHECKING: from engibench.utils.problem import Problem @@ -296,6 +297,7 @@ class Args: run_config=vars(args), **checkpoint_identity(args), primary_files=["model.pth"], + design_normalizer=normalizer_state(design_normalizer), ) wandb.finish() diff --git a/engiopt/generators/gan_1d/adapter.py b/engiopt/generators/gan_1d/adapter.py index ae424cbd..8aebf264 100644 --- a/engiopt/generators/gan_1d/adapter.py +++ b/engiopt/generators/gan_1d/adapter.py @@ -9,8 +9,10 @@ from engiopt.core import ConditionBatch from engiopt.core import design_shape_of from engiopt.core import Generator +from engiopt.core import recorded_design_normalizer from engiopt.generators.gan_1d.gan_1d import Generator as GAN1DNet from engiopt.generators.gan_1d.gan_1d import prepare_data +from engiopt.transforms import load_normalizer_state if TYPE_CHECKING: from engibench.core import Problem @@ -37,6 +39,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device """Load a trained GAN-1D generator, rebuilding its design normalizer.""" config = resolved.run_config _, design_normalizer = prepare_data(problem, device) + design_normalizer = load_normalizer_state(design_normalizer, recorded_design_normalizer(resolved), device) net = GAN1DNet( latent_dim=config["latent_dim"], design_shape=design_shape_of(problem), diff --git a/engiopt/generators/gan_1d/gan_1d.py b/engiopt/generators/gan_1d/gan_1d.py index 5082443f..ba0da9d0 100644 --- a/engiopt/generators/gan_1d/gan_1d.py +++ b/engiopt/generators/gan_1d/gan_1d.py @@ -19,6 +19,7 @@ from torchvision import transforms import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -27,7 +28,7 @@ from engiopt.reproducibility import seed_training from engiopt.transforms import condition_keys from engiopt.transforms import flatten_dict_factory -import wandb +from engiopt.transforms import normalizer_state if TYPE_CHECKING: from engibench.utils.problem import Problem @@ -358,6 +359,7 @@ def prepare_data(problem: Problem, device: th.device) -> tuple[th.utils.data.Ten run_config=vars(args), **checkpoint_identity(args), primary_files=["generator.pth"], + design_normalizer=normalizer_state(design_normalizer), ) wandb.finish() diff --git a/engiopt/generators/gan_bezier/adapter.py b/engiopt/generators/gan_bezier/adapter.py index fe224852..062d7ae2 100644 --- a/engiopt/generators/gan_bezier/adapter.py +++ b/engiopt/generators/gan_bezier/adapter.py @@ -8,9 +8,11 @@ from engiopt.core import ConditionBatch from engiopt.core import Generator +from engiopt.core import recorded_design_normalizer from engiopt.generators.gan_bezier.gan_bezier import Generator as BezierGANNet from engiopt.generators.gan_bezier.gan_bezier import prepare_data from engiopt.transforms import flatten_dict_factory +from engiopt.transforms import load_normalizer_state if TYPE_CHECKING: from engibench.core import Problem @@ -52,6 +54,9 @@ def build( """Load a trained BezierGAN generator, rebuilding its scalar normalizer.""" config = resolved.run_config _, design_scalars_normalizer, _ = prepare_data(problem, n_prepare_samples, device) + design_scalars_normalizer = load_normalizer_state( + design_scalars_normalizer, recorded_design_normalizer(resolved), device + ) net = BezierGANNet( latent_dim=config["latent_dim"], noise_dim=config["noise_dim"], diff --git a/engiopt/generators/gan_bezier/gan_bezier.py b/engiopt/generators/gan_bezier/gan_bezier.py index ab10e522..fc8ba86e 100644 --- a/engiopt/generators/gan_bezier/gan_bezier.py +++ b/engiopt/generators/gan_bezier/gan_bezier.py @@ -18,13 +18,14 @@ from torch import nn import torch.nn.functional as f import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training -import wandb +from engiopt.transforms import normalizer_state if TYPE_CHECKING: from collections.abc import Callable @@ -640,6 +641,7 @@ def prepare_data(problem, batch_size, device, seed=None): run_config=vars(args), **checkpoint_identity(args), primary_files=["bezier_generator.pth"], + design_normalizer=normalizer_state(design_scalars_normalizer), ) if args.track: diff --git a/engiopt/transforms.py b/engiopt/transforms.py index 189d21fe..dc083890 100644 --- a/engiopt/transforms.py +++ b/engiopt/transforms.py @@ -1,14 +1,20 @@ """Transformations for the data.""" -from collections.abc import Callable +from __future__ import annotations + +from typing import Any, TYPE_CHECKING -from datasets import Dataset -from engibench.core import Problem from gymnasium import spaces import numpy as np import torch as th import torch.nn.functional as f +if TYPE_CHECKING: + from collections.abc import Callable + + from datasets import Dataset + from engibench.core import Problem + def get_scalar_condition_keys(problem: Problem, dataset: Dataset, *, drop_constants: bool = False) -> list[str]: """Return the condition keys usable as a dense model input. @@ -129,3 +135,56 @@ def drop_constant(ds: Dataset, condition_names: list[str]) -> tuple[Dataset, lis ds = ds.remove_columns(dropped) return ds, kept + + +def normalizer_state(normalizer: Any) -> dict[str, Any]: + """Serialize a min/max `Normalizer`'s fitted bounds for a checkpoint. + + Several 1D and Bezier models scale designs or conditions into `[0, 1]` using + bounds fitted on the training split. Those bounds are as much a part of the + model as its weights -- decode against different bounds and you get a + different design -- but `Normalizer` is a plain class rather than an + `nn.Module`, so nothing put them in the state dict. Recording them alongside + the weights is what stops a later dataset revision from silently changing + what a fixed checkpoint produces. + + Duck-typed on `min_val`/`max_val`/`eps` because the five training scripts + each define their own identical `Normalizer`, and unifying them is a change + to training code that this does not need. + + Args: + normalizer: Any object exposing `min_val`, `max_val`, and optionally `eps`. + + Returns: + A JSON-serializable dict of the fitted bounds. + """ + return { + "min": th.as_tensor(normalizer.min_val).detach().cpu().flatten().tolist(), + "max": th.as_tensor(normalizer.max_val).detach().cpu().flatten().tolist(), + "eps": float(getattr(normalizer, "eps", 1e-7)), + } + + +def load_normalizer_state(normalizer: Any, state: dict[str, Any] | None, device: th.device) -> Any: + """Restore a normalizer's bounds from what the checkpoint recorded. + + Returns `normalizer` unchanged when the checkpoint predates this being + recorded, which leaves those packages loading exactly as they did before -- + fitted from the current dataset, which is what those runs actually used. + + Args: + normalizer: A freshly fitted normalizer to overwrite in place. + state: The recorded bounds, or None for a checkpoint without them. + device: Device the restored tensors should live on. + + Returns: + The same object, carrying the recorded bounds when there were any. + """ + if not state or "min" not in state or "max" not in state: + return normalizer + reference = th.as_tensor(normalizer.min_val) + normalizer.min_val = th.tensor(state["min"], device=device, dtype=reference.dtype).reshape(reference.shape) + normalizer.max_val = th.tensor(state["max"], device=device, dtype=reference.dtype).reshape(reference.shape) + if "eps" in state: + normalizer.eps = float(state["eps"]) + return normalizer diff --git a/tests/test_normalizer_replay.py b/tests/test_normalizer_replay.py new file mode 100644 index 00000000..9de4f6cb --- /dev/null +++ b/tests/test_normalizer_replay.py @@ -0,0 +1,141 @@ +"""Tests that a checkpoint's preprocessing travels with its weights. + +Several 1D and Bezier models scale designs or conditions into `[0, 1]` using +bounds fitted on the training split. `Normalizer` is a plain class rather than an +`nn.Module`, so those bounds never reached the state dict, and loading refitted +them from whatever dataset was current. Same checkpoint, same content hash, +different dataset revision, different outputs -- with nothing to show it happened. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +import torch as th + +from engiopt.core import recorded_condition_normalizer +from engiopt.core import recorded_design_normalizer +from engiopt.transforms import load_normalizer_state +from engiopt.transforms import normalizer_state + + +class _Normalizer: + """Stand-in with the same surface as the five training scripts' copies.""" + + def __init__(self, min_val: th.Tensor, max_val: th.Tensor, eps: float = 1e-7) -> None: + self.min_val = min_val + self.max_val = max_val + self.eps = eps + + def normalize(self, x: th.Tensor) -> th.Tensor: + return (x - self.min_val) / (self.max_val - self.min_val + self.eps) + + def denormalize(self, x: th.Tensor) -> th.Tensor: + return x * (self.max_val - self.min_val + self.eps) + self.min_val + + +class _Resolved: + """Minimal `ResolvedCheckpoint` stand-in carrying only metadata.""" + + def __init__(self, metadata: dict[str, Any]) -> None: + self.metadata = metadata + + +def test_recorded_bounds_survive_the_dataset_moving() -> None: + """The property the whole change exists for. + + A checkpoint trained when designs spanned `[0, 10]` must keep decoding + against `[0, 10]`, even when today's dataset spans `[0, 100]`. + """ + trained = _Normalizer(th.tensor([0.0]), th.tensor([10.0])) + recorded = normalizer_state(trained) + + # Later: the dataset revision moved, so loading fits very different bounds. + refitted = _Normalizer(th.tensor([0.0]), th.tensor([100.0])) + restored = load_normalizer_state(refitted, recorded, th.device("cpu")) + + sample = th.tensor([[0.5]]) + assert th.allclose(restored.denormalize(sample), trained.denormalize(sample)) + assert not th.allclose( + restored.denormalize(sample), _Normalizer(th.tensor([0.0]), th.tensor([100.0])).denormalize(sample) + ) + + +def test_state_round_trips_through_json_shaped_values() -> None: + """Metadata is JSON, so the recorded form has to be plain lists and floats.""" + normalizer = _Normalizer(th.tensor([1.0, 2.0]), th.tensor([3.0, 4.0]), eps=1e-5) + state = normalizer_state(normalizer) + + assert state == {"min": [1.0, 2.0], "max": [3.0, 4.0], "eps": 1e-5} + + restored = load_normalizer_state(_Normalizer(th.zeros(2), th.ones(2)), state, th.device("cpu")) + assert th.allclose(restored.min_val, th.tensor([1.0, 2.0])) + assert th.allclose(restored.max_val, th.tensor([3.0, 4.0])) + assert restored.eps == pytest.approx(1e-5) + + +def test_multidimensional_bounds_keep_their_shape() -> None: + """Bezier scalar normalizers are not one-dimensional; flattening must be reversed.""" + normalizer = _Normalizer(th.zeros(2, 3), th.ones(2, 3)) + state = normalizer_state(normalizer) + + restored = load_normalizer_state(_Normalizer(th.zeros(2, 3), th.zeros(2, 3)), state, th.device("cpu")) + + assert restored.max_val.shape == (2, 3) + + +def test_a_checkpoint_without_recorded_bounds_is_left_alone() -> None: + """Packages predating this must load exactly as before -- fitted from the dataset. + + Those runs really did use the current dataset's bounds, so refitting is the + faithful reconstruction, not a fallback that quietly changes their meaning. + """ + fitted = _Normalizer(th.tensor([0.0]), th.tensor([7.0])) + + unchanged = load_normalizer_state(fitted, None, th.device("cpu")) + + assert th.allclose(unchanged.max_val, th.tensor([7.0])) + + +def test_partial_metadata_is_ignored_rather_than_half_applied() -> None: + """Half a normalizer is worse than none: it would scale by mismatched bounds.""" + fitted = _Normalizer(th.tensor([0.0]), th.tensor([7.0])) + + unchanged = load_normalizer_state(fitted, {"min": [1.0]}, th.device("cpu")) + + assert th.allclose(unchanged.min_val, th.tensor([0.0])) + assert th.allclose(unchanged.max_val, th.tensor([7.0])) + + +# ---------------------------------------------------------------------- +# Reading the recorded state back off a checkpoint +# ---------------------------------------------------------------------- + + +def test_recorded_accessors_read_their_own_field() -> None: + """Condition and design bounds are different numbers and must not be crossed.""" + resolved = _Resolved( + { + "condition_normalizer": {"min": [0.0], "max": [1.0]}, + "design_normalizer": {"min": [-5.0], "max": [5.0]}, + } + ) + + assert recorded_condition_normalizer(resolved) == {"min": [0.0], "max": [1.0]} + assert recorded_design_normalizer(resolved) == {"min": [-5.0], "max": [5.0]} + + +@pytest.mark.parametrize( + "metadata", + [{}, {"design_normalizer": None}, {"design_normalizer": {"min": [0.0]}}, {"design_normalizer": "nonsense"}], +) +def test_unusable_recorded_state_reads_as_absent(metadata: dict[str, Any]) -> None: + """Anything that is not a complete pair of bounds means "fall back to fitting".""" + assert recorded_design_normalizer(_Resolved(metadata)) is None + + +def test_no_checkpoint_reads_as_absent() -> None: + """Adapters call these before a package exists in some code paths.""" + assert recorded_design_normalizer(None) is None + assert recorded_condition_normalizer(None) is None From 082ef1fa25da88e975e858991c991d2eb774ca8f Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Fri, 7 Aug 2026 15:18:20 +0200 Subject: [PATCH 12/24] Say what scored a row, and stop skipping rows of unknown identity `already_evaluated` treated a row with no `checkpoint_hash` as a match. You are right that it should not: an absent hash means the identity of those weights is unknown, and unknown is not equal. One hashless row could suppress every future evaluation of that configuration and seed, retrained weights included, until somebody deleted it by hand. Since this schema has never shipped on `main` there is no historical board whose re-evaluation cost the leniency was protecting, so the legacy-row test is updated rather than kept. Rows now also record `engibench_version`: the EngiBench that ran the evaluation, as opposed to the one the spec was frozen against. This is the compromise from the earlier thread -- recorded, not enforced, and not part of ranking identity, so contributors on a slightly different commit can still be compared while the difference stays visible. Deriving that version had the bug you predicted. `git -C` searches parent directories, so a wheel unpacked into a virtualenv inside a checkout would report the *enclosing* repository's sha as EngiBench's -- a wrong sha, worse than none. It now reads PEP 610 `direct_url.json` first, which is where a non-editable `pip install git+...` records the commit and is the case CI hits, and only accepts a git answer when the repository root is the package's own directory and the path is not an installed-package path. Separately, `--config-fingerprints` was applied as a cross product over generators. A fingerprint hashes one algorithm's hyperparameters, so evaluating two models against a flat list asked for packages that were never going to exist and buried the run in load failures that were not failures. `algo:fingerprint` scopes an entry to its owner. Co-Authored-By: Claude Opus 5 (1M context) --- engiopt/evaluate.py | 39 +++++++- engiopt/evaluation/evaluator.py | 24 +++++ engiopt/evaluation/leaderboard.py | 16 ++-- engiopt/evaluation/spec.py | 58 ++++++++++-- tests/test_engibench_provenance.py | 129 +++++++++++++++++++++++++++ tests/test_leaderboard_identity.py | 103 +++++++++++++++++++++ tests/test_leaderboard_publishing.py | 13 ++- 7 files changed, 363 insertions(+), 19 deletions(-) create mode 100644 tests/test_engibench_provenance.py create mode 100644 tests/test_leaderboard_identity.py diff --git a/engiopt/evaluate.py b/engiopt/evaluate.py index e72246b3..9415e618 100644 --- a/engiopt/evaluate.py +++ b/engiopt/evaluate.py @@ -52,7 +52,14 @@ class Args: """Specific hyperparameter configurations to evaluate, by fingerprint. Empty means the canonical default-hyperparameter checkpoint -- what a bare - model name refers to. Pass several to score a sweep.""" + model name refers to. Pass several to score a sweep. + + Scope an entry to one model with `algo:fingerprint`. A fingerprint hashes + one algorithm's hyperparameters, so evaluating several models against a flat + list tries every fingerprint against every model and reports the mismatches + as load failures. `--generators cgan_cnn_2d gan_cnn_2d --config-fingerprints + cgan_cnn_2d:023dd1fb gan_cnn_2d:06d9a9a1` asks for exactly the two that + exist.""" spec: str | None = None """Eval spec reference, e.g. `beams2d/v1`. Defaults to `/v1`.""" metrics: tuple[str, ...] = () @@ -123,14 +130,40 @@ def _resolve_generator_names(requested: tuple[str, ...], problem_id: str) -> lis return sorted(generators_for(BUILTIN_PROBLEMS[problem_id]())) +def _fingerprints_for(requested: tuple[str, ...], algo: str) -> tuple[str | None, ...]: + """Which configurations to try for one algorithm. + + A fingerprint hashes one algorithm's hyperparameters, so it is meaningless + against a different algorithm: evaluating two models against the union of + both their fingerprints asks for packages that were never going to exist and + buries the run in load failures that are not failures. `algo:fingerprint` + scopes an entry to its owner; a bare fingerprint still applies to every + algorithm, which is what a single-model run wants. + + Args: + requested: Raw `--config-fingerprints` values. + algo: The algorithm being loaded. + + Returns: + Fingerprints to try, or `(None,)` meaning the canonical checkpoint. + """ + if not requested: + return (None,) + selected = [ + entry.split(":", 1)[1] if ":" in entry else entry + for entry in requested + if ":" not in entry or entry.split(":", 1)[0] == algo + ] + return tuple(selected) + + def _load_generators(args: Args, evaluator: Evaluator) -> list[Generator]: """Load every requested generator/seed pair, reporting those that fail.""" generators: list[Generator] = [] - fingerprints: tuple[str | None, ...] = args.config_fingerprints or (None,) for name in _resolve_generator_names(args.generators, args.problem_id): generator_cls = BUILTIN_GENERATORS[name] for seed in args.seeds: - for fingerprint in fingerprints: + for fingerprint in _fingerprints_for(args.config_fingerprints, name): try: generators.append( generator_cls.from_pretrained( diff --git a/engiopt/evaluation/evaluator.py b/engiopt/evaluation/evaluator.py index a3ac4540..daf7bd28 100644 --- a/engiopt/evaluation/evaluator.py +++ b/engiopt/evaluation/evaluator.py @@ -52,6 +52,7 @@ "checkpoint_revision", "checkpoint_hash", "code_version", + "engibench_version", "evaluated_at", ) """Columns identifying *what was measured*, as opposed to metric values. @@ -60,6 +61,15 @@ `checkpoint_hash` says which *weights* did. Re-training the same configuration and seed, or changing the training code, yields different weights under the same name, so without it a row cannot be traced back to the model that earned it. + +`engibench_version` is the EngiBench that *ran* the evaluation, as opposed to the +one the spec was frozen against. A change to `simulate` or `optimize` moves the +optimality gaps without touching any condition name or dataset, so the spec's +digest cannot catch it. It is recorded rather than enforced: refusing to evaluate +unless the sha matches would lock out every contributor whose EngiBench differs +by any commit, and partitioning ranks by it would mean two models could never be +compared unless they were scored on identical builds. Recording it keeps the +difference visible without either cost. """ @@ -210,6 +220,7 @@ def _provenance(self, generator: Generator, ctx: EvaluationContext) -> dict[str, "checkpoint_revision": getattr(generator, "checkpoint_revision", None), "checkpoint_hash": getattr(generator, "checkpoint_hash", None), "code_version": code_version(), + "engibench_version": engibench_version(), "evaluated_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"), } @@ -250,6 +261,19 @@ def leaderboard( return order_columns(pd.DataFrame(rows)) +@functools.cache +def engibench_version() -> str: + """Which EngiBench *ran* this evaluation, however it was installed. + + Shares its implementation with the value `EvalSpec.freeze` records, so the + spec's "frozen against" and the row's "evaluated on" are the same kind of + string and can be compared directly. + """ + from engiopt.evaluation.spec import engibench_version as _spec_engibench_version + + return _spec_engibench_version() + + @functools.cache def code_version() -> str: """Which EngiOpt produced a result: the installed version, plus a git sha in a checkout. diff --git a/engiopt/evaluation/leaderboard.py b/engiopt/evaluation/leaderboard.py index 4fcc6256..b9848df0 100644 --- a/engiopt/evaluation/leaderboard.py +++ b/engiopt/evaluation/leaderboard.py @@ -296,8 +296,14 @@ def already_evaluated(board: pd.DataFrame, **key: Any) -> bool: Pass `checkpoint_hash` and the skip becomes about the *weights*, not just the name: re-training the same configuration and seed produces different weights under the same identity, and those deserve a fresh score rather - than inheriting the old row's. A row predating the hash column has no value - to compare, so it is treated as a match and still skipped. + than inheriting the old row's. + + A row with no hash does not match. An absent hash means the identity of the + weights behind that row is unknown, and unknown is not the same as equal -- + treating it as a match would let one hashless row suppress every future + re-evaluation of that configuration and seed, including genuinely new + weights. This schema has never shipped, so there is no historical board + whose re-evaluation cost that leniency would be protecting. """ if board.empty: return False @@ -305,9 +311,5 @@ def already_evaluated(board: pd.DataFrame, **key: Any) -> bool: for column, value in key.items(): if column not in board.columns: continue - if column == "checkpoint_hash": - # Older rows carry no hash; do not force them to be re-evaluated. - mask &= board[column].isna() | (board[column] == value) - else: - mask &= board[column] == value + mask &= board[column] == value return bool(mask.any()) diff --git a/engiopt/evaluation/spec.py b/engiopt/evaluation/spec.py index 09535a55..238b4623 100644 --- a/engiopt/evaluation/spec.py +++ b/engiopt/evaluation/spec.py @@ -267,7 +267,7 @@ def freeze(self, problem: Problem) -> EvalSpec: **asdict(frozen), "condition_digest": _digest(indices, conditions, ref_designs), "problem_conditions": tuple(problem.conditions_keys), - "engibench_version": _engibench_version(), + "engibench_version": engibench_version(), } ) @@ -326,7 +326,7 @@ def _condition_bytes(values: Any) -> bytes: return json.dumps(values, sort_keys=True, default=str).encode() -def _engibench_version() -> str: +def engibench_version() -> str: """The EngiBench that produced a spec: its version, plus a git sha from a source checkout. The release version alone does not identify a problem definition -- 0.2.0 on @@ -337,8 +337,54 @@ def _engibench_version() -> str: import engibench version = getattr(engibench, "__version__", "unknown") - source_root = Path(engibench.__file__).resolve().parent.parent + sha = _installed_vcs_commit() or _source_checkout_commit(Path(engibench.__file__).resolve().parent.parent) + return f"{version}+{sha}" if sha else version + + +def _installed_vcs_commit() -> str | None: + """The commit recorded by `pip install "engibench @ git+..."`, if it was installed that way. + + A non-editable VCS install leaves no `.git` directory, but pip records the + exact commit in the distribution's `direct_url.json` (PEP 610). That is the + authoritative answer whenever it exists, and it is the case CI hits. + """ + from importlib import metadata + try: + raw = metadata.distribution("engibench").read_text("direct_url.json") + except (metadata.PackageNotFoundError, OSError): + return None + if not raw: + return None + try: + commit = json.loads(raw).get("vcs_info", {}).get("commit_id") + except json.JSONDecodeError: + return None + return str(commit)[:12] if commit else None + + +def _source_checkout_commit(source_root: Path) -> str | None: + """The commit of an EngiBench *source checkout*, or None if it is not one. + + `git -C` searches parent directories, so asking it about a wheel unpacked + into a virtualenv inside another repository answers with *that* repository's + commit. A wheel in `EngiOpt/.venv/.../site-packages/engibench` would report + EngiOpt's sha as EngiBench's, which is worse than reporting nothing. Two + guards prevent it: an installed-package path is never a checkout, and the + repository root must be the package's own parent rather than some ancestor. + """ + if "site-packages" in source_root.parts or "dist-packages" in source_root.parts: + return None + try: + toplevel = subprocess.run( + ["git", "-C", str(source_root), "rev-parse", "--show-toplevel"], + capture_output=True, + text=True, + check=True, + timeout=5, + ).stdout.strip() + if not toplevel or Path(toplevel).resolve() != source_root: + return None sha = subprocess.run( ["git", "-C", str(source_root), "rev-parse", "--short=12", "HEAD"], capture_output=True, @@ -346,10 +392,10 @@ def _engibench_version() -> str: check=True, timeout=5, ).stdout.strip() - # Installed as a wheel rather than a checkout: the version is all there is. + # Installed as a plain wheel rather than a checkout: the version is all there is. except (OSError, subprocess.SubprocessError): - return version - return f"{version}+{sha}" if sha else version + return None + return sha or None def _dataset_revision(dataset_id: str | None) -> str | None: diff --git a/tests/test_engibench_provenance.py b/tests/test_engibench_provenance.py new file mode 100644 index 00000000..1ad294f3 --- /dev/null +++ b/tests/test_engibench_provenance.py @@ -0,0 +1,129 @@ +"""Tests for identifying which EngiBench produced a spec and which one scored a row. + +A version string does not pin a problem definition: 0.2.0 from PyPI and 0.2.0 +from `main` point `thermoelastic2d` at different datasets. The spec's condition +and dataset checks catch that class of difference. What they cannot catch is a +change to `simulate` or `optimize`, which moves IOG/COG/FOG without touching any +condition name -- so the revision is recorded on both sides. +""" + +from __future__ import annotations + +import json +from typing import Any, TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + + import pytest + +from engiopt.evaluation import spec as spec_mod +from engiopt.evaluation.evaluator import PROVENANCE_COLUMNS + + +def test_a_vcs_install_reports_its_recorded_commit(monkeypatch: pytest.MonkeyPatch) -> None: + """`pip install "engibench @ git+..."` leaves no `.git`, but records the commit. + + This is the case CI hits, and the one that has to work: the spec files carry + `0.2.0+0a028c03d02b`, and `freeze()` must be able to reproduce that string. + """ + payload = json.dumps({"url": "https://github.com/IDEALLab/EngiBench", "vcs_info": {"commit_id": "0a028c03d02bfeed"}}) + + class _Distribution: + @staticmethod + def read_text(name: str) -> str | None: + return payload if name == "direct_url.json" else None + + import importlib.metadata as importlib_metadata + + monkeypatch.setattr(importlib_metadata, "distribution", lambda _name: _Distribution()) + + assert spec_mod._installed_vcs_commit() == "0a028c03d02b" + + +def test_a_wheel_without_vcs_info_reports_nothing(monkeypatch: pytest.MonkeyPatch) -> None: + """A plain PyPI wheel has a `direct_url.json` only sometimes, and never a commit.""" + + class _Distribution: + @staticmethod + def read_text(_name: str) -> str | None: + return json.dumps({"url": "file:///tmp/engibench.whl", "archive_info": {}}) + + import importlib.metadata as importlib_metadata + + monkeypatch.setattr(importlib_metadata, "distribution", lambda _name: _Distribution()) + + assert spec_mod._installed_vcs_commit() is None + + +def test_an_installed_package_is_never_read_as_a_checkout(tmp_path: Path) -> None: + """The bug this guards: `git -C` searches *parent* directories. + + A wheel unpacked into a virtualenv inside the EngiOpt repository would make + `git -C .../site-packages/..` answer with EngiOpt's commit and record it as + EngiBench's -- a wrong sha, which is worse than no sha. + """ + site_packages = tmp_path / ".venv" / "lib" / "python3.12" / "site-packages" + site_packages.mkdir(parents=True) + + assert spec_mod._source_checkout_commit(site_packages) is None + + +def test_a_directory_inside_an_unrelated_repo_is_not_that_repo(tmp_path: Path) -> None: + """Only the repository *root* counts as the package's own checkout.""" + import subprocess + + repo = tmp_path / "some_repo" + (repo / "nested" / "engibench").mkdir(parents=True) + subprocess.run(["git", "init", "-q", str(repo)], check=True) + + # `nested` sits inside a git repo but is not its root, so it is not a checkout + # of the package -- exactly the shape a vendored install takes. + assert spec_mod._source_checkout_commit(repo / "nested") is None + + +def test_freeze_and_the_row_report_the_same_kind_of_string() -> None: + """The spec's "frozen against" and the row's "evaluated on" must be comparable.""" + from engiopt.evaluation.evaluator import engibench_version as row_version + + assert row_version() == spec_mod.engibench_version() + + +def test_the_runtime_engibench_is_recorded_on_every_row() -> None: + """Recorded, not enforced. + + Enforcing an exact match would refuse to evaluate for any contributor whose + EngiBench differs by a commit; partitioning ranks by it would mean two models + could never be compared unless scored on identical builds. Recording keeps + the difference visible at neither cost. + """ + assert "engibench_version" in PROVENANCE_COLUMNS + + +def test_the_version_carries_a_sha_when_one_is_identifiable(monkeypatch: pytest.MonkeyPatch) -> None: + """The composed value is `+`, which is what the specs record.""" + monkeypatch.setattr(spec_mod, "_installed_vcs_commit", lambda: "abcdef123456") + + assert spec_mod.engibench_version().endswith("+abcdef123456") + + +def test_the_version_degrades_to_the_bare_release(monkeypatch: pytest.MonkeyPatch) -> None: + """With nothing identifying the source, the release version is all there is.""" + monkeypatch.setattr(spec_mod, "_installed_vcs_commit", lambda: None) + monkeypatch.setattr(spec_mod, "_source_checkout_commit", lambda _root: None) + + assert "+" not in spec_mod.engibench_version() + + +def test_the_committed_specs_record_a_resolvable_engibench(tmp_path: Path) -> None: + """Every committed spec should say which EngiBench it was frozen against.""" + spec_files = sorted(spec_mod.SPEC_ROOT.glob("*/*.json")) + assert spec_files, "no committed specs found" + + missing: list[str] = [] + for path in spec_files: + payload: dict[str, Any] = json.loads(path.read_text()) + if not payload.get("engibench_version"): + missing.append(f"{path.parent.name}/{path.stem}") + + assert not missing, f"specs with no recorded EngiBench: {missing}" diff --git a/tests/test_leaderboard_identity.py b/tests/test_leaderboard_identity.py new file mode 100644 index 00000000..622ed807 --- /dev/null +++ b/tests/test_leaderboard_identity.py @@ -0,0 +1,103 @@ +"""Tests for what identifies a leaderboard row and what may be skipped. + +`--skip-existing` decides whether work has already been done. Getting it wrong in +either direction is expensive: too strict re-runs a finished sweep, too lenient +silently keeps a stale score for weights that have since been retrained. +""" + +from __future__ import annotations + +import pandas as pd + +from engiopt.evaluate import _fingerprints_for +from engiopt.evaluation.leaderboard import already_evaluated + + +def _board(**overrides: object) -> pd.DataFrame: + row = { + "problem_id": "beams2d", + "algo_id": "cgan_cnn_2d", + "config_fingerprint": "023dd1fb", + "seed": 1, + "spec_version": "v1", + "checkpoint_hash": "aaaaaaaaaaaaaaaa", + } + row.update(overrides) + return pd.DataFrame([row]) + + +_KEY = { + "problem_id": "beams2d", + "algo_id": "cgan_cnn_2d", + "config_fingerprint": "023dd1fb", + "seed": 1, + "spec_version": "v1", +} + + +def test_the_same_weights_are_skipped() -> None: + """The point of the flag: a finished sweep must not re-run.""" + assert already_evaluated(_board(), **_KEY, checkpoint_hash="aaaaaaaaaaaaaaaa") + + +def test_retrained_weights_are_not_skipped() -> None: + """Same config and seed, new training run, different weights: score them.""" + assert not already_evaluated(_board(), **_KEY, checkpoint_hash="bbbbbbbbbbbbbbbb") + + +def test_a_row_without_a_hash_does_not_count_as_a_match() -> None: + """Unknown identity is not the same as equal identity. + + Treating a hashless row as a match lets one such row suppress every future + evaluation of that configuration and seed, including genuinely new weights, + until somebody deletes it by hand. + """ + assert not already_evaluated(_board(checkpoint_hash=None), **_KEY, checkpoint_hash="aaaaaaaaaaaaaaaa") + + +def test_a_different_seed_is_not_skipped() -> None: + """Sanity: the ordinary key columns still discriminate.""" + assert not already_evaluated(_board(), **{**_KEY, "seed": 2}, checkpoint_hash="aaaaaaaaaaaaaaaa") + + +def test_an_empty_board_skips_nothing() -> None: + assert not already_evaluated(pd.DataFrame(), **_KEY, checkpoint_hash="aaaaaaaaaaaaaaaa") + + +def test_a_board_predating_the_hash_column_still_matches_on_the_rest() -> None: + """A board with no hash column at all is a schema difference, not a mismatch.""" + board = _board().drop(columns=["checkpoint_hash"]) + + assert already_evaluated(board, **_KEY, checkpoint_hash="aaaaaaaaaaaaaaaa") + + +# ---------------------------------------------------------------------- +# Config fingerprints belong to one algorithm +# ---------------------------------------------------------------------- + + +def test_a_bare_fingerprint_applies_to_every_model() -> None: + """Single-model runs are the common case and should not need scoping.""" + assert _fingerprints_for(("023dd1fb",), "cgan_cnn_2d") == ("023dd1fb",) + assert _fingerprints_for(("023dd1fb",), "gan_cnn_2d") == ("023dd1fb",) + + +def test_a_scoped_fingerprint_reaches_only_its_own_model() -> None: + """Otherwise a two-model run asks for packages that were never going to exist.""" + requested = ("cgan_cnn_2d:023dd1fb", "gan_cnn_2d:06d9a9a1") + + assert _fingerprints_for(requested, "cgan_cnn_2d") == ("023dd1fb",) + assert _fingerprints_for(requested, "gan_cnn_2d") == ("06d9a9a1",) + + +def test_scoped_and_bare_entries_mix() -> None: + """A shared fingerprint alongside a scoped one is a reasonable thing to ask for.""" + requested = ("cgan_cnn_2d:023dd1fb", "deadbeef") + + assert _fingerprints_for(requested, "cgan_cnn_2d") == ("023dd1fb", "deadbeef") + assert _fingerprints_for(requested, "gan_cnn_2d") == ("deadbeef",) + + +def test_no_fingerprints_means_the_canonical_checkpoint() -> None: + """`None` is what `from_pretrained` reads as "the default-hyperparameter package".""" + assert _fingerprints_for((), "cgan_cnn_2d") == (None,) diff --git a/tests/test_leaderboard_publishing.py b/tests/test_leaderboard_publishing.py index 917163ab..668d84a8 100644 --- a/tests/test_leaderboard_publishing.py +++ b/tests/test_leaderboard_publishing.py @@ -301,9 +301,16 @@ def test_the_same_weights_are_skipped() -> None: ) -def test_rows_predating_the_hash_column_are_still_skipped() -> None: - """An older board has no hash to compare, so it should not force a re-run.""" - assert lb.already_evaluated( +def test_a_row_without_a_hash_is_not_treated_as_a_match() -> None: + """A missing hash means the weights behind that row are unknown, not identical. + + Skipping on it lets one hashless row suppress every future evaluation of that + configuration and seed -- including genuinely retrained weights -- until + somebody deletes the row by hand. This schema has never shipped on `main`, so + there is no historical board whose re-evaluation cost the old leniency was + protecting. + """ + assert not lb.already_evaluated( _published_row(checkpoint_hash=None), problem_id="p", algo_id="a", From 4215c102f6a9eba6b5a02624c95195ad34d097c1 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Mon, 10 Aug 2026 13:56:58 +0200 Subject: [PATCH 13/24] Give a swept pool a way to reach the canonical checkpoint path `{problem_id}/seed_{seed}` is what a bare model name resolves to: what `from_pretrained(problem, seed=1)` reads and what the README's evaluation command uses. Only a run using the training script's default hyperparameters writes it, and a sweep has no such run, because every arm sets at least one flag. The consequence is currently live on the Hub. cgan_cnn_2d and gan_cnn_2d have 138 published packages each and not one canonical package for any problem, so the documented `--generators cgan_cnn_2d --seeds 1` fails with a bare "not found". vqgan and diffusion_2d_cond do have canonical packages, but only because their swept center configuration happens to coincide with the script defaults -- luck rather than design, and precisely what makes the bare model name unreliable. Two changes. A failed canonical resolution now lists the configurations that do exist, so the reader learns the package was never written rather than suspecting a broken path. And `python -m engiopt.promote_checkpoint` copies a chosen configuration to the canonical path, so an existing sweep arm can become the default without spending a few hundred GPU-hours retraining a model we already have. The promoted package keeps its own `config_fingerprint`, so a leaderboard row still records which configuration earned the score; only the address changes. An incomplete package is refused. Co-Authored-By: Claude Opus 5 (1M context) --- engiopt/promote_checkpoint.py | 89 ++++++++++++++ tests/test_canonical_promotion.py | 189 ++++++++++++++++++++++++++++++ 2 files changed, 278 insertions(+) create mode 100644 engiopt/promote_checkpoint.py create mode 100644 tests/test_canonical_promotion.py diff --git a/engiopt/promote_checkpoint.py b/engiopt/promote_checkpoint.py new file mode 100644 index 00000000..f1de0bf9 --- /dev/null +++ b/engiopt/promote_checkpoint.py @@ -0,0 +1,89 @@ +"""Promote a swept configuration to the canonical checkpoint path. + +`{problem_id}/seed_{seed}` is what a bare model name resolves to -- what +`from_pretrained(problem, seed=1)` and `python -m engiopt.evaluate --seeds 1` +read. Only a run using the training script's default hyperparameters writes it, +which means a hyperparameter sweep never does: every arm varies something, so a +sweep can publish hundreds of packages and still leave that path empty. + +This is the escape hatch. Pick the arm that should be the default and promote it, +instead of spending GPU hours retraining a configuration you already have:: + + python -m engiopt.promote_checkpoint --algo cgan_cnn_2d --problem-id beams2d \ + --seed 1 --config-fingerprint 023dd1fb + +To see what is available first:: + + python -m engiopt.promote_checkpoint --algo cgan_cnn_2d --problem-id beams2d --list +""" + +from __future__ import annotations + +from dataclasses import dataclass +import sys + +import tyro + +from engiopt.checkpoint_store import build_hf_repo_id +from engiopt.checkpoint_store import list_packages +from engiopt.checkpoint_store import promote_to_canonical +from engiopt.core import DEFAULT_HF_ENTITY +from engiopt.core import DEFAULT_HF_REPO_PREFIX + + +@dataclass +class Args: + """Command-line arguments for promoting a checkpoint.""" + + algo: str + """Model family, e.g. `cgan_cnn_2d`. Selects the HF repo.""" + problem_id: str + """EngiBench problem the checkpoint was trained on.""" + seed: int = 1 + """Training seed. The promoted package keeps it.""" + config_fingerprint: str | None = None + """Which configuration to promote. Required unless `--list` is passed.""" + list: bool = False + """Print the packages this repo holds for the problem, then exit.""" + hf_entity: str = DEFAULT_HF_ENTITY + """HF org/user holding the checkpoint repos.""" + hf_repo_prefix: str = DEFAULT_HF_REPO_PREFIX + """Prefix of the per-model-family repo.""" + + +def main(args: Args) -> int: + """Promote one configuration, or list what is available. + + Returns: + Process exit status: 0 on success, 1 when there is nothing to promote. + """ + repo_id = build_hf_repo_id(args.hf_entity, args.hf_repo_prefix, args.algo) + available = list_packages(repo_id, args.problem_id) + + if args.list or args.config_fingerprint is None: + if not available: + print(f"No packages for {args.problem_id!r} in {repo_id}.") + return 1 + print(f"{len(available)} package(s) in {repo_id}:") + for path in available: + print(f" {path}") + if not args.list: + print("\nPass --config-fingerprint to promote one of these to the canonical path.") + return 1 + return 0 + + canonical = promote_to_canonical( + hf_entity=args.hf_entity, + hf_repo_prefix=args.hf_repo_prefix, + problem_id=args.problem_id, + algo=args.algo, + seed=args.seed, + config_fingerprint=args.config_fingerprint, + ) + print(f"Promoted cfg_{args.config_fingerprint} to {repo_id}/{canonical}") + print(f"`--generators {args.algo} --seeds {args.seed}` now resolves.") + return 0 + + +if __name__ == "__main__": + sys.exit(main(tyro.cli(Args))) diff --git a/tests/test_canonical_promotion.py b/tests/test_canonical_promotion.py new file mode 100644 index 00000000..bb0ecdc8 --- /dev/null +++ b/tests/test_canonical_promotion.py @@ -0,0 +1,189 @@ +"""Tests for reaching the canonical checkpoint path after a sweep. + +`{problem_id}/seed_{seed}` is what a bare model name resolves to. Only a run +using the training script's default hyperparameters writes it, and a sweep has no +such run -- every arm varies something. So a 500-run sweep can publish hundreds of +packages and leave the documented `--seeds 1` command resolving nothing, which is +exactly the state the IDEALLab repos were in. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +from engiopt import checkpoint_store + + +def _package(root: Path, *, metadata: dict[str, Any] | None = None) -> Path: + root.mkdir(parents=True, exist_ok=True) + (root / "generator.pth").write_bytes(b"weights") + (root / "run_config.json").write_text(json.dumps({"latent_dim": 32})) + (root / "metadata.json").write_text(json.dumps(metadata or {"algo": "cgan_cnn_2d", "seed": 1})) + return root + + +def test_a_missing_canonical_package_names_the_ones_that_exist(monkeypatch: pytest.MonkeyPatch) -> None: + """ "Not found" alone sends the reader looking for a broken path. + + The actual situation is that the canonical package was never written, and the + configurations that *were* published are sitting right there. Say so. + """ + monkeypatch.setattr( + checkpoint_store, + "list_packages", + lambda _repo, _problem=None: ["beams2d/cfg_023dd1fb/seed_1", "beams2d/cfg_1763fb14/seed_1"], + ) + + hint = checkpoint_store._sibling_packages_hint("IDEALLab/engiopt-cgan-cnn-2d", "beams2d", 1, None) + + assert "cfg_023dd1fb" in hint + assert "canonical beams2d/seed_1" in hint + assert "--config-fingerprints" in hint + assert "promote_checkpoint" in hint + + +def test_the_hint_is_quiet_when_a_configuration_was_asked_for(monkeypatch: pytest.MonkeyPatch) -> None: + """Asking for a specific fingerprint that is absent is a different mistake.""" + monkeypatch.setattr(checkpoint_store, "list_packages", lambda _repo, _problem=None: ["beams2d/cfg_023dd1fb/seed_1"]) + + hint = checkpoint_store._sibling_packages_hint("IDEALLab/engiopt-cgan-cnn-2d", "beams2d", 1, ["cfg_deadbeef"]) + + assert "cfg_023dd1fb" in hint + assert "promote_checkpoint" not in hint + + +def test_an_empty_repo_says_so(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(checkpoint_store, "list_packages", lambda _repo, _problem=None: []) + + hint = checkpoint_store._sibling_packages_hint("IDEALLab/engiopt-vqgan", "photonics2d", 1, None) + + assert "No packages for 'photonics2d'" in hint + + +def test_the_hint_never_replaces_the_real_error(monkeypatch: pytest.MonkeyPatch) -> None: + """A hint that itself fails must not mask what the caller was actually told.""" + + def explode(_repo: str, _problem: str | None = None) -> list[str]: + raise ConnectionError("the Hub is down") + + monkeypatch.setattr(checkpoint_store, "list_packages", explode) + + assert checkpoint_store._sibling_packages_hint("repo", "beams2d", 1, None) == "" + + +# ---------------------------------------------------------------------- +# Promotion +# ---------------------------------------------------------------------- + + +def test_promotion_copies_a_configuration_to_the_canonical_path(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """The escape hatch: make an already-trained arm the default without retraining.""" + source = _package( + tmp_path / "snapshot" / "beams2d" / "cfg_023dd1fb" / "seed_1", + metadata={"algo": "cgan_cnn_2d", "seed": 1, "config_fingerprint": "023dd1fb"}, + ) + (source / checkpoint_store.METRICS_FILE).write_text(json.dumps({"mmd": 0.04})) + + uploads: list[dict[str, Any]] = [] + + class _Api: + def __init__(self, **_: Any) -> None: + pass + + def upload_folder(self, **kwargs: Any) -> None: + staged = sorted(entry.name for entry in Path(kwargs["folder_path"]).iterdir()) + uploads.append({"path_in_repo": kwargs["path_in_repo"], "files": staged}) + + monkeypatch.setattr(checkpoint_store, "HfApi", _Api) + monkeypatch.setattr(checkpoint_store, "snapshot_download", lambda **_: str(tmp_path / "snapshot")) + + canonical = checkpoint_store.promote_to_canonical( + hf_entity="IDEALLab", + hf_repo_prefix="engiopt", + problem_id="beams2d", + algo="cgan_cnn_2d", + seed=1, + config_fingerprint="023dd1fb", + ) + + assert canonical == "beams2d/seed_1" + assert uploads[0]["path_in_repo"] == "beams2d/seed_1" + assert "generator.pth" in uploads[0]["files"] + # The old address's scores do not describe the promoted copy. + assert checkpoint_store.METRICS_FILE not in uploads[0]["files"] + + +def test_promotion_records_where_it_came_from(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """A leaderboard row must still be traceable to the configuration that earned it.""" + _package( + tmp_path / "snapshot" / "beams2d" / "cfg_023dd1fb" / "seed_1", + metadata={"algo": "cgan_cnn_2d", "seed": 1, "config_fingerprint": "023dd1fb"}, + ) + written: dict[str, Any] = {} + + class _Api: + def __init__(self, **_: Any) -> None: + pass + + def upload_folder(self, **kwargs: Any) -> None: + written.update(json.loads((Path(kwargs["folder_path"]) / "metadata.json").read_text())) + + monkeypatch.setattr(checkpoint_store, "HfApi", _Api) + monkeypatch.setattr(checkpoint_store, "snapshot_download", lambda **_: str(tmp_path / "snapshot")) + + checkpoint_store.promote_to_canonical( + hf_entity="IDEALLab", + hf_repo_prefix="engiopt", + problem_id="beams2d", + algo="cgan_cnn_2d", + seed=1, + config_fingerprint="023dd1fb", + ) + + assert written["config_fingerprint"] == "023dd1fb" + assert written["promoted_from"] == "beams2d/cfg_023dd1fb/seed_1" + assert written["hf_package_path"] == "beams2d/seed_1" + + +def test_an_incomplete_package_is_never_promoted(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """Promoting a torn package would put an unloadable thing at the canonical path. + + That is precisely the state VQGAN was already in; promotion must not be a + second way to reach it. + """ + _package( + tmp_path / "snapshot" / "beams2d" / "cfg_1151406c" / "seed_42", + metadata={"stage": "cvqgan", checkpoint_store.PACKAGE_COMPLETE_FIELD: False}, + ) + + monkeypatch.setattr(checkpoint_store, "HfApi", lambda **_: None) + monkeypatch.setattr(checkpoint_store, "snapshot_download", lambda **_: str(tmp_path / "snapshot")) + + with pytest.raises(FileNotFoundError, match="marked incomplete"): + checkpoint_store.promote_to_canonical( + hf_entity="IDEALLab", + hf_repo_prefix="engiopt", + problem_id="beams2d", + algo="vqgan", + seed=42, + config_fingerprint="1151406c", + ) + + +def test_promoting_an_absent_configuration_says_so(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setattr(checkpoint_store, "HfApi", lambda **_: None) + monkeypatch.setattr(checkpoint_store, "snapshot_download", lambda **_: str(tmp_path / "snapshot")) + + with pytest.raises(FileNotFoundError, match="does not exist"): + checkpoint_store.promote_to_canonical( + hf_entity="IDEALLab", + hf_repo_prefix="engiopt", + problem_id="beams2d", + algo="cgan_cnn_2d", + seed=1, + config_fingerprint="absent00", + ) From e6231bc197a407bbc1fd59de593d3c56c59fdd61 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Fri, 7 Aug 2026 15:19:13 +0200 Subject: [PATCH 14/24] Make the docs describe the migration that actually happened, and lint mean it The README still advertised W&B artifact refs for the power-electronics optimizer and linked docs/checkpoint_migration_playbook.md, which this branch deletes; `run_pe_optimization.py` still described its arguments as W&B artifacts and used one in its usage example; and `Generator.resolve_checkpoint` still listed `wandb` as a model source. None of those paths exist. The README also documented no leaderboard workflow at all -- `--push-to`, `--skip-existing`, and `--attach-metrics` were undocumented, as was the fact that a sweep leaves the canonical path empty. Both now have a section, along with how an outside contributor submits a model. The ruff job ran `ruff check --fix`, which repairs violations inside the runner and then reports success, so the job was green on code that does not pass. Dropping `--fix` makes it mean what it says, and the 20 pre-existing I001 violations it was hiding are fixed here. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/lint.yaml | 4 ++- README.md | 25 +++++++++++++++++-- engiopt/generators/cgan_2d/cgan_2d.py | 2 +- engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py | 2 +- engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py | 2 +- engiopt/generators/cgan_vae/cgan_vae.py | 2 +- .../diffusion_2d_cond/diffusion_2d_cond.py | 2 +- engiopt/generators/gan_2d/gan_2d.py | 2 +- engiopt/generators/gan_cnn_2d/gan_cnn_2d.py | 2 +- .../pixel_cnn_pp_2d/pixel_cnn_pp_2d.py | 2 +- engiopt/surrogate_model/mlp_tabular_only.py | 2 +- .../surrogate_model/run_pe_optimization.py | 10 ++++---- engiopt/surrogate_model/training_utils.py | 1 - example_easy_model.ipynb | 2 +- example_hard_model.ipynb | 2 +- 15 files changed, 42 insertions(+), 20 deletions(-) diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index 30866aad..77a7b853 100644 --- a/.github/workflows/lint.yaml +++ b/.github/workflows/lint.yaml @@ -16,7 +16,9 @@ jobs: with: args: "--version" - run: ruff format --check - - run: ruff check --fix --output-format=github . + # No --fix: repairing violations inside the runner makes the job pass on + # code that does not, so import-order drift never reaches a review. + - run: ruff check --output-format=github . mypy: runs-on: ubuntu-latest diff --git a/README.md b/README.md index 24d5c375..609683fe 100644 --- a/README.md +++ b/README.md @@ -112,6 +112,28 @@ python -m engiopt.evaluate --problem-id "beams2d" --generators cgan_cnn_2d --see ``` Evaluation pulls the checkpoint from HF automatically. Pass `--hf-entity` / `--hf-repo-prefix` to point at a different HF repo, and `--config-fingerprints` to score specific hyperparameter configurations instead of the default one. +A bare `--seeds 1` resolves the **canonical** package, `{problem_id}/seed_1`, which only a run using the training script's default hyperparameters writes. A hyperparameter sweep never writes it — every arm varies something — so after a sweep the canonical path can be empty while hundreds of `cfg_*` packages exist. When that happens the error lists what the repo does hold. To make one of those arms the default without retraining it: + +``` +python -m engiopt.promote_checkpoint --algo cgan_cnn_2d --problem-id beams2d --list +python -m engiopt.promote_checkpoint --algo cgan_cnn_2d --problem-id beams2d --seed 1 --config-fingerprint 023dd1fb +``` + +### Leaderboard + +Results are rows in a CSV, keyed by problem, algorithm, config fingerprint, seed, and spec version. Locally that CSV is `--output-csv`; published, it is the same table in a HuggingFace dataset repo: + +``` +python -m engiopt.evaluate --problem-id beams2d --generators cgan_cnn_2d --seeds 1 \ + --push-to IDEALLab/engiopt-leaderboard +``` + +Publishing downloads the existing board, merges on the row key, and uploads the result conditional on the revision it read, so adding one model never recomputes or overwrites anyone else's rows. Ranks restart at 1 within each `(problem_id, spec_version)` — a beams2d score and a photonics2d score measure different things and are never placed in one ordering. `--skip-existing` skips rows the board already holds for the exact weights, compared by `checkpoint_hash`. + +Every row records what produced it: the checkpoint revision and content hash, the EngiOpt version, and the EngiBench version that ran the evaluation. The spec records what it was frozen against, including the pinned dataset revision, so a change to either side is visible rather than silently shifting every number. + +**Contributing a model.** Train against a committed spec, publish the checkpoint package to your own HF repo, then open an issue with the repo path and the seed. Evaluation runs on our side and writes the row, so the leaderboard is only ever written by the code that produced the scores. Point `--hf-entity` / `--hf-repo-prefix` at your own repos to reproduce any row locally first. + ### Surrogate model The current surrogate model comprises several steps: @@ -123,11 +145,10 @@ The current surrogate model comprises several steps: See this [notebook](https://github.com/IDEALLab/EngiOpt/blob/main/engiopt/surrogate_model/case_study_pe_notebook.ipynb) for an example. Surrogate-model optimization paths now use the same checkpoint abstraction. For example, the power-electronics optimizer can consume: -* legacy WandB artifact refs * HF package refs such as `hf://IDEALLab/engiopt-mlp-tabular-only/power_electronics/DcGain/seed_42` * local checkpoint package directories -For migration guidance on moving historical checkpoint subsets from WandB to the IDEALLab HF organization later, see [docs/checkpoint_migration_playbook.md](docs/checkpoint_migration_playbook.md). +HuggingFace is the only checkpoint backend. W&B artifacts are no longer a model source anywhere in the codebase; see [docs/checkpoint_layout.md](docs/checkpoint_layout.md) for the package layout. diff --git a/engiopt/generators/cgan_2d/cgan_2d.py b/engiopt/generators/cgan_2d/cgan_2d.py index 46e21be9..c6626689 100644 --- a/engiopt/generators/cgan_2d/cgan_2d.py +++ b/engiopt/generators/cgan_2d/cgan_2d.py @@ -18,6 +18,7 @@ from torch import nn import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -25,7 +26,6 @@ from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import condition_keys -import wandb GeneratorOutputActivation = Literal["tanh", "sigmoid"] diff --git a/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py b/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py index 6699f45f..4818d170 100644 --- a/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py +++ b/engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py @@ -17,6 +17,7 @@ from torchvision import transforms import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -24,7 +25,6 @@ from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import condition_keys -import wandb GeneratorOutputActivation = Literal["tanh", "sigmoid"] diff --git a/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py b/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py index 984beaf7..f9ef1c87 100644 --- a/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py +++ b/engiopt/generators/cgan_cnn_3d/cgan_cnn_3d.py @@ -20,6 +20,7 @@ from torch.nn import functional import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -29,7 +30,6 @@ from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import condition_keys -import wandb @dataclass diff --git a/engiopt/generators/cgan_vae/cgan_vae.py b/engiopt/generators/cgan_vae/cgan_vae.py index bf606894..d436f45f 100644 --- a/engiopt/generators/cgan_vae/cgan_vae.py +++ b/engiopt/generators/cgan_vae/cgan_vae.py @@ -20,6 +20,7 @@ from torch.nn import functional import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -29,7 +30,6 @@ from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import condition_keys -import wandb @dataclass diff --git a/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py b/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py index bd5a755b..92ba33c1 100644 --- a/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py +++ b/engiopt/generators/diffusion_2d_cond/diffusion_2d_cond.py @@ -16,6 +16,7 @@ from torch.nn import functional import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -23,7 +24,6 @@ from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import condition_keys -import wandb if TYPE_CHECKING: from collections.abc import Callable diff --git a/engiopt/generators/gan_2d/gan_2d.py b/engiopt/generators/gan_2d/gan_2d.py index 91723fc0..5b930dff 100644 --- a/engiopt/generators/gan_2d/gan_2d.py +++ b/engiopt/generators/gan_2d/gan_2d.py @@ -16,13 +16,13 @@ from torch import nn import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training -import wandb @dataclass diff --git a/engiopt/generators/gan_cnn_2d/gan_cnn_2d.py b/engiopt/generators/gan_cnn_2d/gan_cnn_2d.py index 32c28151..2007ea2b 100644 --- a/engiopt/generators/gan_cnn_2d/gan_cnn_2d.py +++ b/engiopt/generators/gan_cnn_2d/gan_cnn_2d.py @@ -16,13 +16,13 @@ from torchvision import transforms import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity from engiopt.reproducibility import enable_strict_determinism from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training -import wandb @dataclass diff --git a/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py b/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py index 65a850f2..b5727fa5 100644 --- a/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py +++ b/engiopt/generators/pixel_cnn_pp_2d/pixel_cnn_pp_2d.py @@ -23,6 +23,7 @@ from torch.nn.utils.parametrizations import weight_norm import tqdm import tyro +import wandb from engiopt.checkpoint_store import save_checkpoint_package from engiopt.core import checkpoint_identity @@ -30,7 +31,6 @@ from engiopt.reproducibility import make_dataloader_generator from engiopt.reproducibility import seed_training from engiopt.transforms import condition_keys -import wandb @dataclass diff --git a/engiopt/surrogate_model/mlp_tabular_only.py b/engiopt/surrogate_model/mlp_tabular_only.py index 0731d284..13021b34 100644 --- a/engiopt/surrogate_model/mlp_tabular_only.py +++ b/engiopt/surrogate_model/mlp_tabular_only.py @@ -20,6 +20,7 @@ import torch from torch.utils.data import DataLoader import tyro +import wandb from engiopt.args_utils import parse_list_from_single_item_list from engiopt.args_utils import parse_list_from_string @@ -32,7 +33,6 @@ from engiopt.surrogate_model.training_utils import get_device from engiopt.surrogate_model.training_utils import PlainTabularDataset from engiopt.surrogate_model.training_utils import train_one_model -import wandb @dataclass diff --git a/engiopt/surrogate_model/run_pe_optimization.py b/engiopt/surrogate_model/run_pe_optimization.py index 92359b38..f1902fa3 100644 --- a/engiopt/surrogate_model/run_pe_optimization.py +++ b/engiopt/surrogate_model/run_pe_optimization.py @@ -11,8 +11,8 @@ Example usage ------------- python run_pe_optimization.py \ - --model_gain_path "my_entity/engiopt/your_run_name_model:latest" \ - --model_ripple_path "my_entity/engiopt/your_other_run_name_model:latest" \ + --model_gain_path "hf://IDEALLab/engiopt-mlp-tabular-only/power_electronics/DcGain/seed_1" \ + --model_ripple_path "hf://IDEALLab/engiopt-mlp-tabular-only/power_electronics/VoltageRipple/seed_1" \ --device mps \ --pop_size 500 \ --n_gen 100 \ @@ -34,12 +34,12 @@ from pymoo.optimize import minimize from pymoo.termination import get_termination import tyro +import wandb from engiopt.checkpoint_store import resolve_checkpoint_reference from engiopt.surrogate_model.model_pipeline import ModelPipeline from engiopt.surrogate_model.pymoo_pe_problem import PymooPowerElecProblem from engiopt.surrogate_model.training_utils import get_device -import wandb if TYPE_CHECKING: from pymoo.core.algorithm import Algorithm @@ -53,9 +53,9 @@ class Args: # Surrogate pipelines model_gain_path: str - """Model ref for the gain surrogate (W&B artifact, HF package ref, or local package directory).""" + """Model ref for the gain surrogate (HF package ref, or local package directory).""" model_ripple_path: str - """Model ref for the ripple surrogate (W&B artifact, HF package ref, or local package directory).""" + """Model ref for the ripple surrogate (HF package ref, or local package directory).""" # Optimisation hyperparameters seed: int diff --git a/engiopt/surrogate_model/training_utils.py b/engiopt/surrogate_model/training_utils.py index 25f1173f..ec755151 100644 --- a/engiopt/surrogate_model/training_utils.py +++ b/engiopt/surrogate_model/training_utils.py @@ -10,7 +10,6 @@ import torch.nn.functional as F # noqa: N812 from torch.utils.data import DataLoader from torch.utils.data import Dataset - import wandb if TYPE_CHECKING: diff --git a/example_easy_model.ipynb b/example_easy_model.ipynb index 67e0810f..000f9ba6 100644 --- a/example_easy_model.ipynb +++ b/example_easy_model.ipynb @@ -244,10 +244,10 @@ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import torch as th\n", + "import wandb\n", "\n", "from engiopt import metrics\n", "from engiopt.generators.cgan_cnn_2d.cgan_cnn_2d import Generator\n", - "import wandb\n", "\n", "if th.cuda.is_available():\n", " device = th.device(\"cuda\")\n", diff --git a/example_hard_model.ipynb b/example_hard_model.ipynb index 78897e81..4753a420 100644 --- a/example_hard_model.ipynb +++ b/example_hard_model.ipynb @@ -223,11 +223,11 @@ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import torch as th\n", + "import wandb\n", "\n", "from engiopt import metrics\n", "from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import beta_schedule\n", "from engiopt.generators.diffusion_2d_cond.diffusion_2d_cond import DiffusionSampler\n", - "import wandb\n", "\n", "if th.cuda.is_available():\n", " device = th.device(\"cuda\")\n", From c0ea5ce50b86d979045d7b75b493381b411ca7ca Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 11 Aug 2026 14:40:04 +0200 Subject: [PATCH 15/24] Make the board publishable by strangers, and say when a score is retrieval An internal leaderboard can assume its rows were produced in good faith by someone with no reason to fabricate them. A public one cannot assume either, and the two failure modes need separate answers. Fabrication. A row recorded a checkpoint hash but never a repo, and every default pointed at IDEALLab, so a contributor could neither publish weights nor produce a row anyone could trace. Rows now carry repo, path, revision and hash; `python -m engiopt.verify` follows that address, checks the package still hashes to what was scored, and re-scores it. Rows land unverified and unranked until a runner reproduces them, and because the address is public the runner is just the first auditor rather than a trusted oracle. Verification publishes its own numbers instead of a pass/fail, since one seed on different hardware draws different designs and any tolerance loose enough for that is loose enough for fudging. Gaming. The protocol is public: the spec names the scored conditions and the dataset supplies each optimal design, so `mmd`, `iog` and `fog` are by definition maximised by returning them. Measured on beams2d, a lookup table scores mmd=0.000 and viol=0.00 against the trained cGAN's 0.561 and 0.92 -- it wins every headline metric. `novelty`/`copy_rate` catch it (copy_rate=1.00), and `cond_sens` catches a model that ignores conditions it declares. Both are diagnostic, not ranked: ranking novelty would seat pure noise in first place, which is closing one vector by opening another. Flagged rows are published and left out of the ordering, because deleting a submission is moderation and declining to rank one is a statement about what its number measures. Also closes the ways a score could drift from the thing scored: `required_seeds` stops a best-three-of-twenty median; `ROW_KEY` gains `checkpoint_repo` so two contributors' identically-configured models cannot overwrite each other; the content hash now covers fitted preprocessing, which changes outputs but left the hash untouched; a complete upload clears stale files, while a staged one deliberately does not; an incomplete package is refused before file discovery rather than while explaining a missing file; and `code_version` stops reporting an enclosing repository's sha as its own. Docs: LEADERBOARD.md for the trust model and the copying hole that a v2 spec scored on off-dataset conditions would actually close. README and CONTRIBUTING corrected -- the documented commands now name packages that exist, `--list-generators --check-availability` reports which of the fourteen registered models have published weights, and the guidance to re-fit preprocessing from `conditions.dataset` is replaced by recording it in the checkpoint. 51 new tests. Verified against the live Hub: a real cgan_cnn_2d checkpoint scores end to end with its address recorded and cond_sens=0.171, and all four committed specs still reproduce their frozen digests. Co-Authored-By: Claude Opus 5 (1M context) --- CONTRIBUTING_A_MODEL.md | 62 ++- LEADERBOARD.md | 165 ++++++++ README.md | 35 +- engiopt/checkpoint_store.py | 218 ++++++++-- engiopt/core.py | 14 + engiopt/evaluate.py | 113 ++++- engiopt/evaluation/context.py | 76 ++++ engiopt/evaluation/evaluator.py | 131 +++++- engiopt/evaluation/leaderboard.py | 106 ++++- engiopt/evaluation/metrics/builtin.py | 84 ++++ engiopt/evaluation/registry.py | 10 +- engiopt/evaluation/spec.py | 71 +++- engiopt/evaluation/submission.py | 203 +++++++++ engiopt/evaluation/verify.py | 318 ++++++++++++++ engiopt/specs/beams2d/v1.json | 10 + engiopt/specs/heatconduction2d/v1.json | 10 + engiopt/specs/photonics2d/v1.json | 10 + engiopt/specs/thermoelastic2d/v1.json | 10 + engiopt/verify.py | 132 ++++++ tests/conftest.py | 20 +- tests/test_checkpoint_integrity.py | 125 ++++++ tests/test_engibench_provenance.py | 14 +- tests/test_evaluation.py | 16 +- tests/test_leaderboard_integrity.py | 558 +++++++++++++++++++++++++ tests/test_leaderboard_publishing.py | 10 + tests/test_verification.py | 254 +++++++++++ 26 files changed, 2669 insertions(+), 106 deletions(-) create mode 100644 LEADERBOARD.md create mode 100644 engiopt/evaluation/submission.py create mode 100644 engiopt/evaluation/verify.py create mode 100644 engiopt/verify.py create mode 100644 tests/test_leaderboard_integrity.py create mode 100644 tests/test_verification.py diff --git a/CONTRIBUTING_A_MODEL.md b/CONTRIBUTING_A_MODEL.md index f9667577..35bf7254 100644 --- a/CONTRIBUTING_A_MODEL.md +++ b/CONTRIBUTING_A_MODEL.md @@ -139,22 +139,64 @@ conditions.dataset # original columns, if you need them conditions.keys # condition names, in column order ``` -Most models want `require_tensor`. Reach for `dataset` only when your model -preprocesses conditions the way `vqgan` does (dropping constant columns, -re-normalizing). Unconditional models ignore the argument entirely — and the -conditional-adherence metrics are what will show that. +Most models want `require_tensor`. Reach for `dataset` only for conditions that +**cannot travel in the dense tensor at all** — thermoelastic2d's 65x65 boundary +matrices, for instance, which `condition_keys` excludes for exactly that reason. + +Do **not** use it to re-derive preprocessing. Dropping constant columns or +re-fitting a normalizer against `dataset` at evaluation time changes the scale +your network sees, because the 50 evaluation rows have neither the training +split's statistics nor the same set of columns that never vary. Anything you +fitted during training travels with the weights instead, and is replayed on load: + +```python +save_checkpoint_package( + ..., + condition_stats={"mean": [...], "std": [...]}, # if you rescaled conditions + condition_normalizer=normalizer_state(cond_norm), # if you used a min/max Normalizer + design_normalizer=normalizer_state(design_norm), +) +``` + +These are part of your model's identity: they are folded into the checkpoint +content hash, so changing a bound produces a different checkpoint rather than +inheriting the previous one's leaderboard row. + +Unconditional models ignore the argument entirely, which is allowed. They are +still handed conditions so that `cond_sens` can measure it — see +[LEADERBOARD.md](LEADERBOARD.md). ## 4. Check it ```bash python -m engiopt.evaluate --list-generators # your model should appear -python -m engiopt.evaluate --problem-id beams2d --generators my_model +python -m engiopt.evaluate --problem-id beams2d --generators my_model --hf-entity my-hf-username ``` -Adding `--include-expensive` also runs the simulator-backed metrics (COG, IOG, -FOG, feasibility). Those are slow; leave them off while iterating. +The default pass runs the cheap metrics: `mmd`, `dpp`, `viol`, and the two +integrity checks (`novelty`, `cond_sens`). Feasibility is among them because it +describes the design as generated, so it is a constraint check rather than a +solver run. `--include-expensive` adds the optimality gaps (`iog`, `cog`, +`fog`), which do run the optimizer and are slow — leave them off while +iterating. + +Watch two columns while you develop: + +- **`cond_sens`** should be greater than zero if you declared `conditional = + True`. Exactly zero means your conditions are not reaching the network, which + is a wiring bug far more often than a modelling choice. +- **`copy_rate`** should be near zero. High means your model is reproducing + training designs rather than generating, and the board will publish it without + ranking it. + +## 5. Publish it + +See [LEADERBOARD.md](LEADERBOARD.md) for the full submission path. The short +version: train and evaluate under your own `--hf-entity`, run seeds 1, 2 and 3, +and push. Your rows appear immediately and are ranked once a runner has +re-fetched the weights and reproduced the scores. -## 5. Adding a metric instead +## 6. Adding a metric instead Metrics are registered functions, not model methods — a metric compares a generated set against a reference set under a problem, so it belongs to the @@ -177,7 +219,7 @@ accidentally launch a simulation. Importing the module is what registers it — which means you can define a metric in a notebook cell and it will appear in the next leaderboard you build. -## 6. Where things live +## 7. Where things live | Artifact | Home | |---|---| @@ -196,7 +238,7 @@ To attach a model's scores to its own checkpoint: python -m engiopt.evaluate --problem-id beams2d --generators my_model --attach-metrics ``` -## 7. Evaluation specs +## 8. Evaluation specs A leaderboard only means something if every row was measured the same way. Each problem has a committed spec (`engiopt/specs//v1.json`) freezing the diff --git a/LEADERBOARD.md b/LEADERBOARD.md new file mode 100644 index 00000000..612b4e7f --- /dev/null +++ b/LEADERBOARD.md @@ -0,0 +1,165 @@ +# The EngiOpt leaderboard + +A public board has two problems an internal one does not, and they are +independent. Someone can report numbers no model produced, and someone can +report perfectly honest numbers from a model that games the metric instead of +solving the problem. Neither defence helps with the other. + +- Against **fabrication**: re-execution. Nothing is ranked until a runner has + fetched the weights itself and reproduced the score. +- Against **gaming**: measurement. The scoring protocol is public and cannot be + un-published, so the board measures the exploits it cannot prevent and + declines to rank what trips them. + +## Submitting + +Nothing here needs write access to `IDEALLab`. Publish your checkpoints to your +own HuggingFace account: + +```bash +python engiopt/generators/my_model/my_model.py \ + --problem-id beams2d --seed 1 --hf-entity my-hf-username + +python -m engiopt.evaluate \ + --problem-id beams2d --generators my_model --seeds 1 2 3 \ + --hf-entity my-hf-username \ + --include-expensive \ + --push-to IDEALLab/engiopt-leaderboard +``` + +Your rows land immediately, marked `verified=false`. They are visible on the +board and absent from the ranking until a runner re-scores them. + +Three things are worth knowing before you run it. + +**Your checkpoints must be on the Hub.** Every row records the repo, path, +revision, and content hash of the weights it was scored on, and a row without +that address is refused at publish time — not because it is untrusted, but +because nobody, including you, could ever re-run it. Evaluating with +`--model-source local` produces exactly such a row. + +**Run the seeds the spec asks for.** `EvalSpec.required_seeds` is `(1, 2, 3)`, +and an entry missing any of them is published but not ranked. This is the +cherry-picking rule: a count would still let you run twenty seeds and publish +your best three, which turns a median into a maximum while every individual row +stays honest. Naming the seeds removes the choice. + +**You cannot mark your own row verified.** `push_to_hub` clears the column. + +## Verification + +```bash +# audit anyone's rows; needs no credentials and writes nothing +python -m engiopt.verify --board IDEALLab/engiopt-leaderboard + +# the official runner +python -m engiopt.verify --board IDEALLab/engiopt-leaderboard \ + --verifier ideallab-ci --include-expensive --publish +``` + +Verification fetches the package **at the revision the row recorded**, checks +that it still hashes to what was scored, rebuilds the model through its +registered adapter, and scores it again. What gets published is the runner's own +numbers, not a verdict on yours. + +That distinction is deliberate. Sampling from one seed on different hardware +genuinely produces different designs — CUDA and CPU draw different values from +the same generator state — so a pass/fail comparison would need a tolerance loose +enough to wave through real fudging. Re-scoring sidesteps it: the runner's number +is the number, and yours becomes a claim reported as corroborated or not. + +Four outcomes: + +| Status | Meaning | +|---|---| +| `verified` | Fetched, re-scored. The row now carries the runner's numbers. | +| `weights_moved` | The package no longer hashes to what the row recorded. Not re-scored — whatever is there now is not what earned the score. | +| `unresolvable` | The address led nowhere, or the package would not load. | +| `unknown_generator` | No registered adapter can rebuild the model. | + +The runner is not privileged. It is the first auditor, and because every row +carries a full address, the same command run by anyone else produces the same +answer — which is what keeps the runner honest too. + +## What gets ranked + +A row is published if it can be checked, and ranked if checking it went well. +Those are different bars on purpose: deleting a submission is moderation, and +declining to rank one is a statement about what its number measures. Only the +second scales. + +To be ranked, a row must be `verified`, carry no disqualifying flag, and belong +to an entry covering every required seed. + +| Flag | Trips when | +|---|---| +| `unverified` | No runner has reproduced it yet. | +| `memorized` | `copy_rate` exceeds the spec's `max_copy_rate`. | +| `ignores_conditions` | `cond_sens` is exactly zero for a model registered as conditional. | + +## The copying problem + +This is the one worth understanding before reading any number on the board. + +The spec is public. It names the condition seed, the sample count, and a digest +of the drawn conditions, so anyone can recompute exactly which 50 test rows are +scored — and the dataset supplies the optimal design for each of them. A lookup +table keyed on the condition vector returns those designs and therefore posts a +perfect `mmd` and a near-zero `iog` and `fog`. + +That is not an implementation flaw. Those metrics are *defined* as closeness to +the reference designs, and a retrieval system is closest. No amount of care in +the evaluator changes it, and it cannot be fixed by hiding which rows are +scored: the whole dataset is public, so a memorizer memorizes all of it. + +So the board measures retrieval instead: + +- **`novelty`** — mean per-element RMS distance from each generated design to the + nearest design the model could have copied: the training split it was fitted + on, plus the reference designs the protocol names. +- **`copy_rate`** — the share of the batch closer than `copy_tol`. This is the + one that flags an entry. + +Both are diagnostic, with no ranking direction, and that is not an oversight. +Ranking on novelty would put pure noise in first place — zero means retrieval, +but large means only "unlike the data", which a broken model also achieves. +Closing one gaming vector by opening another is not progress. The same reasoning +applies to `cond_sens`: an unconditional model is a legitimate thing to build, +and responding to conditions *wrongly* also moves the output, so a large value is +not by itself a good one. + +Read them next to `mmd` and `viol`, never on their own. + +### What would actually close it + +Scoring on conditions that have no published solution. A v2 spec would draw +conditions from the condition space rather than from dataset rows, and carry the +reference *objective* for each — computed once by whoever freezes the spec, with +no design released. A lookup table then has nothing to look up, and `mmd` moves +to comparing against the training distribution rather than against 50 named +designs. + +That is the right benchmark for a generative design model, and it is not built +here: freezing such a spec means an optimizer run per condition. The mechanism +this PR adds — flags, eligibility, and a spec version that carries its own +thresholds — is what a v2 would plug into. + +## Reading the board + +```python +import pandas as pd +from engiopt.evaluation.leaderboard import disagreement, load_from_hub, rank +from engiopt.evaluation.spec import EvalSpec + +board = load_from_hub("IDEALLab/engiopt-leaderboard") +spec = EvalSpec.load("beams2d/v1") + +rank(board, "fog", eval_spec=spec) # the public ordering +rank(board, "fog", eval_spec=spec, eligible_only=False) # everything, including claims +disagreement(board, ["mmd", "dpp", "fog"], eval_spec=spec) +``` + +`rank` refuses a diagnostic metric outright rather than inventing a direction +for it. `disagreement` is the view worth looking at first: where the metrics +disagree about the winner is the only place the board is telling you something a +single number could not. diff --git a/README.md b/README.md index 609683fe..b93e9275 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ This repository contains the code for optimization and machine learning algorith ## Coding Philosophy As much as we can, we follow the [CleanRL](https://github.com/vwxyzjn/cleanrl) philosophy: single-file, high-quality implementations with research-friendly features: -* Single-file implementation: every training detail is in one file, so you can easily understand and modify the code. There is usually another file that contains evaluation code. +* Single-file implementation: every training detail is in one file, so you can easily understand and modify the code. Evaluation is not per-model — every model is scored by one shared evaluator through the `Generator` contract in `adapter.py`, so a metric is added once rather than thirteen times. * High-quality: we use type hints, docstrings, and comments to make the code easy to understand. We also rely on linters for formatting and checking our code. * Logging: we use experiment tracking tools like [Weights & Biases](https://wandb.ai/site) to log the results of our experiments. All our "official" runs are logged in the [EngiOpt project](https://wandb.ai/engibench/engiopt). * Reproducibility: we seed all the random number generators, make PyTorch deterministic, report the hyperparameters and code in WandB. @@ -108,31 +108,52 @@ python engiopt/generators/cgan_cnn_2d/cgan_cnn_2d.py --problem-id "beams2d" --ge Then evaluate: ``` -python -m engiopt.evaluate --problem-id "beams2d" --generators cgan_cnn_2d --seeds 1 +python -m engiopt.evaluate --problem-id "beams2d" --generators cgan_cnn_2d --seeds 1 2 3 \ + --config-fingerprints cgan_cnn_2d:825831f6 ``` Evaluation pulls the checkpoint from HF automatically. Pass `--hf-entity` / `--hf-repo-prefix` to point at a different HF repo, and `--config-fingerprints` to score specific hyperparameter configurations instead of the default one. -A bare `--seeds 1` resolves the **canonical** package, `{problem_id}/seed_1`, which only a run using the training script's default hyperparameters writes. A hyperparameter sweep never writes it — every arm varies something — so after a sweep the canonical path can be empty while hundreds of `cfg_*` packages exist. When that happens the error lists what the repo does hold. To make one of those arms the default without retraining it: +The fingerprint is needed here rather than optional. A bare `--seeds 1` resolves the **canonical** package, `{problem_id}/seed_1`, which only a run using the training script's default hyperparameters writes — and a hyperparameter sweep never writes one, because every arm varies something. `IDEALLab/engiopt-cgan-cnn-2d` currently holds 46 beams2d packages and no canonical path, so a bare `--seeds 1` there fails and lists what the repo does hold. `diffusion_2d_cond` and `vqgan` do have canonical packages and can be evaluated by seed alone: + +``` +python -m engiopt.evaluate --problem-id beams2d --generators diffusion_2d_cond --seeds 1 2 3 +``` + +To give a swept arm the canonical path without retraining it: ``` python -m engiopt.promote_checkpoint --algo cgan_cnn_2d --problem-id beams2d --list -python -m engiopt.promote_checkpoint --algo cgan_cnn_2d --problem-id beams2d --seed 1 --config-fingerprint 023dd1fb +python -m engiopt.promote_checkpoint --algo cgan_cnn_2d --problem-id beams2d --seed 1 --config-fingerprint 825831f6 ``` +`825831f6` is the one cgan_cnn_2d configuration trained on seeds 1–10; the rest of the sweep is seed 42 only, so promoting any other arm to `seed_1` would name a package that does not exist. + ### Leaderboard Results are rows in a CSV, keyed by problem, algorithm, config fingerprint, seed, and spec version. Locally that CSV is `--output-csv`; published, it is the same table in a HuggingFace dataset repo: ``` -python -m engiopt.evaluate --problem-id beams2d --generators cgan_cnn_2d --seeds 1 \ +python -m engiopt.evaluate --problem-id beams2d --generators cgan_cnn_2d --seeds 1 2 3 \ + --config-fingerprints cgan_cnn_2d:825831f6 \ --push-to IDEALLab/engiopt-leaderboard ``` Publishing downloads the existing board, merges on the row key, and uploads the result conditional on the revision it read, so adding one model never recomputes or overwrites anyone else's rows. Ranks restart at 1 within each `(problem_id, spec_version)` — a beams2d score and a photonics2d score measure different things and are never placed in one ordering. `--skip-existing` skips rows the board already holds for the exact weights, compared by `checkpoint_hash`. -Every row records what produced it: the checkpoint revision and content hash, the EngiOpt version, and the EngiBench version that ran the evaluation. The spec records what it was frozen against, including the pinned dataset revision, so a change to either side is visible rather than silently shifting every number. +Every row records what produced it: **which repo, path, revision, and content hash** the weights came from, the EngiOpt version, and the EngiBench version that ran the evaluation. The spec records what it was frozen against, including the pinned dataset revision, so a change to either side is visible rather than silently shifting every number. + +**Anyone may publish; nothing is ranked until it is re-run.** Rows land `verified=false`, and a runner re-fetches the checkpoint at its recorded revision and scores it itself before they enter the ranking: + +``` +python -m engiopt.verify --board IDEALLab/engiopt-leaderboard # audit; writes nothing +python -m engiopt.verify --board ... --verifier ideallab-ci --publish # the official runner +``` + +Because the row carries the full address, that audit is not privileged — anyone can run the same command and get the same answer. + +**Two integrity metrics decide whether a score means what it looks like.** The evaluation protocol is public, so a lookup table keyed on the condition vector returns the dataset-optimal designs and posts a perfect `mmd` and a zero `viol` — measured on beams2d, it beats a trained cGAN on every headline metric. `novelty` / `copy_rate` catch it (`copy_rate=1.00`), and `cond_sens` catches a model that ignores the conditions it claims to use. Both are diagnostic rather than ranked, because ranking on them would just reward the opposite extreme. Flagged rows are published and left out of the ordering. -**Contributing a model.** Train against a committed spec, publish the checkpoint package to your own HF repo, then open an issue with the repo path and the seed. Evaluation runs on our side and writes the row, so the leaderboard is only ever written by the code that produced the scores. Point `--hf-entity` / `--hf-repo-prefix` at your own repos to reproduce any row locally first. +See **[LEADERBOARD.md](LEADERBOARD.md)** for the submission path, the flags, and what would actually close the copying hole. ### Surrogate model diff --git a/engiopt/checkpoint_store.py b/engiopt/checkpoint_store.py index 0b27dc9b..5a397f58 100644 --- a/engiopt/checkpoint_store.py +++ b/engiopt/checkpoint_store.py @@ -68,12 +68,60 @@ DESCRIPTIVE_FILES = frozenset({RUN_CONFIG_FILE, METADATA_FILE, METRICS_FILE}) """Package files that describe the checkpoint rather than being part of it. -Neither hashed as weights nor served as loadable checkpoint files. `metrics.json` +Neither hashed as raw bytes nor served as loadable checkpoint files. `metrics.json` belongs here for a reason worth stating: it is written *after* evaluation, into the package it describes, so counting it would let scoring a checkpoint change that checkpoint's identity. + +`metadata.json` is excluded as a *file* but not as *content*: the fields in it +that change what the model computes are folded into the hash separately, by +`identity_metadata`. Hashing the file whole would fail the other way, since it +records the package's own path and revision, which differ between the two +locations one run writes. +""" + +ADDRESS_METADATA_FIELDS = frozenset( + { + "hf_repo_id", + "hf_package_path", + "hf_config_package_path", + "hf_revision", + "hf_config_revision", + "promoted_from", + "wandb_entity", + "wandb_project", + "wandb_run_id", + "wandb_run_url", + PACKAGE_COMPLETE_FIELD, + } +) +"""Metadata fields describing *where a package sits*, not *what it computes*. + +Excluded from the content hash. Everything else in `metadata.json` is included, +and that direction is deliberate: a field that changes the model's output but is +missed by the hash lets a behaviourally different checkpoint inherit an old +score, while a field that is merely bookkeeping but gets hashed only causes a +harmless re-evaluation. An allowlist would fail the dangerous way round every +time someone adds a field and forgets to list it. + +`package_complete` is here because an incomplete package is now refused at load, +so it can never reach a leaderboard row to be confused with a complete one. """ + +def identity_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + """The part of a package's metadata that changes what the model computes. + + Preprocessing state lives here -- `condition_stats`, `condition_normalizer`, + `design_normalizer`, `condition_keys`. Those are fitted on the training split + and replayed at load, so two packages with byte-identical weights and + different normalizer bounds *decode differently* and are different models. + Leaving them out of the hash would let `--skip-existing` hand one of them the + other's leaderboard row. + """ + return {key: value for key, value in metadata.items() if key not in ADDRESS_METADATA_FIELDS} + + CheckpointBackend = Literal["hf", "none"] ModelSource = Literal["auto", "hf", "local"] @@ -96,6 +144,22 @@ class ResolvedCheckpoint: the repo path cannot tell them apart. This can, so a leaderboard row is traceable to the model that produced it. """ + repo_id: str | None = None + """HF repo the package came from; None for local packages. + + Recorded so a leaderboard row can name where its weights live. Without it a + row published by anyone outside the default entity points at nothing + fetchable, and no third party can re-run the evaluation that produced it. + """ + package_path: str | None = None + """Path of the package inside `repo_id`, e.g. `beams2d/cfg_023dd1fb/seed_1`.""" + + @property + def reference(self) -> str | None: + """A `hf://entity/repo/path` reference that round-trips through `resolve_checkpoint_reference`.""" + if self.repo_id is None or self.package_path is None: + return None + return f"hf://{self.repo_id}/{self.package_path}" def build_hf_repo_id(hf_entity: str, hf_repo_prefix: str, algo: str) -> str: @@ -422,6 +486,11 @@ def promote_to_canonical( folder_path=str(stage_dir), path_in_repo=canonical_path, commit_message=f"Promote cfg_{config_fingerprint} to canonical {canonical_path}", + # A promotion copies a whole package, so it replaces whatever held + # the canonical path before. Leaving the previous occupant's files + # in place would blend two configurations under the name that a bare + # `--seeds 1` resolves to. + delete_patterns=["*"], ) return canonical_path @@ -431,9 +500,19 @@ def resolve_checkpoint_reference( model_source: ModelSource, model_ref: str, required_files: list[str] | None = None, + revision: str | None = None, ) -> ResolvedCheckpoint: """Resolve a checkpoint package from an explicit `hf://` or local reference. + Args: + model_source: `auto` infers `local` for an existing directory, else `hf`. + model_ref: `hf:////`, or a local directory. + required_files: Files the package must contain. + revision: Repo commit to fetch. Pass the revision a leaderboard row + recorded to re-read exactly the package it was scored on; without it + the current head is fetched, which is a different question and can + be a different model. + Raises: ValueError: If the reference cannot be interpreted. """ @@ -443,7 +522,9 @@ def resolve_checkpoint_reference( if inferred_source == "hf": repo_id, package_path = _parse_hf_reference(model_ref) - return _resolve_hf_package(repo_id=repo_id, package_path=package_path, required_files=required_files or []) + return _resolve_hf_package( + repo_id=repo_id, package_path=package_path, required_files=required_files or [], revision=revision + ) if inferred_source == "local": return _resolve_local_package(model_ref.removeprefix("file://"), required_files or []) @@ -477,15 +558,41 @@ def _upload_package_to_hf( folder_path=tmpdir, path_in_repo=package_path, commit_message=f"Upload checkpoint for {algo} {package_path}", + delete_patterns=_stale_file_patterns(metadata), ) return getattr(commit_info, "oid", None) -def _resolve_hf_package(*, repo_id: str, package_path: str, required_files: list[str]) -> ResolvedCheckpoint: +def _stale_file_patterns(metadata: dict[str, Any]) -> list[str] | None: + """What to delete from a package path that this upload is not replacing. + + `upload_folder` overwrites the files it carries and leaves everything else + alone, so re-training a configuration into a path that already held a + *larger* set of files produces a package mixing two runs' weights. Discovery + then serves that mixture, and it loads: the filenames are all present, so + nothing complains. + + Which files are stale depends on whether this upload is the whole package: + + - **Complete** (`package_complete=True`): these files *are* the package, so + anything else at the path is left over from a previous run and goes. That + includes `metrics.json`, which described weights that no longer exist. + - **Incomplete**: a stage of a multi-stage run, deliberately additive. + VQGAN's second stage uploads `vqgan.pth` and `discriminator.pth` without + re-uploading the `cvqgan.pth` its first stage wrote, so deleting here + would destroy the earlier stage's work. + """ + return ["*"] if metadata.get(PACKAGE_COMPLETE_FIELD, True) else None + + +def _resolve_hf_package( + *, repo_id: str, package_path: str, required_files: list[str], revision: str | None = None +) -> ResolvedCheckpoint: repo_snapshot = snapshot_download( repo_id=repo_id, repo_type="model", allow_patterns=[f"{package_path}/*"], + revision=revision, ) root_dir = os.path.join(repo_snapshot, package_path) if not os.path.isdir(root_dir): @@ -493,7 +600,12 @@ def _resolve_hf_package(*, repo_id: str, package_path: str, required_files: list # `snapshot_download` returns `.../snapshots/`, which is the # revision the files were actually taken from. return _load_package_from_directory( - root_dir=root_dir, required_files=required_files, source="hf", revision=os.path.basename(repo_snapshot) + root_dir=root_dir, + required_files=required_files, + source="hf", + revision=os.path.basename(repo_snapshot), + repo_id=repo_id, + package_path=package_path, ) @@ -504,31 +616,40 @@ def _resolve_local_package(local_model_dir: str, required_files: list[str]) -> R def hash_package_contents(root_dir: str) -> str: - """Content hash of every weight file in a package, identifying these exact weights. + """Content hash of a package: everything that decides what it generates. + + This is the identity a leaderboard row is traced back to, so the property it + must have is that *any* change to what the model computes changes the hash. + Two things decide that, and both are covered: - Hashes what the package *contains* rather than what the model declared it - needs. A model may load a file it did not list -- VQGAN's conditional - variant reads `cvqgan.pth`, which is not in its `checkpoint_files` because - the unconditional variant has none -- and two packages differing only in - that file generate differently. Hashing the declared list would give them - the same identity. + **The weight files.** Hashed by what the package *contains* rather than what + the model declared it needs. A model may load a file it did not list -- + VQGAN's conditional variant reads `cvqgan.pth`, which is not in its + `checkpoint_files` because the unconditional variant has none -- and two + packages differing only in that file generate differently. Hashing the + declared list would give them the same identity. - Files that *describe* the package rather than constitute it are excluded: + **The fitted preprocessing.** Normalizer bounds and condition statistics are + replayed at load and rescale both what goes into the network and what comes + out, so they are part of the model even though they never reached a state + dict. See `identity_metadata` for exactly which fields count. - - `run_config.json`, `metadata.json` -- metadata records the package path and - revision, which differ between the two locations one run writes, so - including it would give identical weights two different hashes. + Two files are excluded outright: + + - `run_config.json` -- the hyperparameters are already what + `config_fingerprint` identifies, and it carries operational keys (W&B + entity, tracking flags) that do not change a single output value. - `metrics.json` -- written by `publish_checkpoint_metrics` *after* the weights, into this same directory. Including it would make attaching a score change the identity of the thing scored: the next `--skip-existing` run would see a new hash, re-evaluate, rewrite the metrics, and change the - hash again. The hash must be a property of the weights alone. + hash again -- an evaluation loop with no fixed point. Args: root_dir: Local directory holding the package. Returns: - A 16-character hash over the file names and their bytes. + A 16-character hash over the weight bytes and the identity metadata. """ weight_files = sorted( entry @@ -541,11 +662,23 @@ def hash_package_contents(root_dir: str) -> str: with open(os.path.join(root_dir, file_name), "rb") as handle: for chunk in iter(lambda: handle.read(1 << 20), b""): hasher.update(chunk) + + metadata_path = os.path.join(root_dir, METADATA_FILE) + metadata = _read_json(metadata_path) if os.path.exists(metadata_path) else {} + identity = identity_metadata(metadata) + if identity: + hasher.update(json.dumps(identity, sort_keys=True, default=str).encode()) return hasher.hexdigest()[:16] def _load_package_from_directory( - *, root_dir: str, required_files: list[str], source: Literal["hf", "local"], revision: str | None = None + *, + root_dir: str, + required_files: list[str], + source: Literal["hf", "local"], + revision: str | None = None, + repo_id: str | None = None, + package_path: str | None = None, ) -> ResolvedCheckpoint: run_config_path = os.path.join(root_dir, RUN_CONFIG_FILE) metadata_path = os.path.join(root_dir, METADATA_FILE) @@ -553,6 +686,13 @@ def _load_package_from_directory( raise FileNotFoundError(f"Missing {RUN_CONFIG_FILE} in {root_dir}") run_config = _read_json(run_config_path) metadata = _read_json(metadata_path) if os.path.exists(metadata_path) else {} + # Checked before the files are even looked at. A stage that happens to have + # written every file the *next* stage's loader asks for is still a package + # from a run that did not finish, and it must not load just because the + # filenames line up -- that is precisely the case where a silently wrong + # model reaches a leaderboard row. + if metadata.get(PACKAGE_COMPLETE_FIELD) is False: + raise FileNotFoundError(_incomplete_package_message(root_dir, metadata)) package_files = required_files or _discover_package_files(root_dir, metadata) files = {file_name: os.path.join(root_dir, file_name) for file_name in package_files} missing = [file_name for file_name, file_path in files.items() if not os.path.exists(file_path)] @@ -566,11 +706,18 @@ def _load_package_from_directory( metadata=metadata, revision=revision, content_hash=hash_package_contents(root_dir), + repo_id=repo_id, + package_path=package_path, ) -def _missing_files_message(root_dir: str, missing: list[str], metadata: dict[str, Any]) -> str: - """Explain a package that is missing weights, saying so when the run never finished. +def _present_files(root_dir: str) -> list[str]: + """The package's weight files, for error messages that say how far a run got.""" + return sorted(entry for entry in os.listdir(root_dir) if entry not in DESCRIPTIVE_FILES) + + +def _incomplete_package_message(root_dir: str, metadata: dict[str, Any]) -> str: + """Explain a package whose run never finished. A multi-stage model uploads after each stage, so a run killed partway leaves a package that has a `run_config.json` and a `metadata.json` and therefore @@ -578,16 +725,27 @@ def _missing_files_message(root_dir: str, missing: list[str], metadata: dict[str the two apart, and saying which stage it stopped at turns "file not found" into something the reader can act on. """ - present = sorted(entry for entry in os.listdir(root_dir) if entry not in DESCRIPTIVE_FILES) - lines = [f"Missing checkpoint file(s) {missing} in {root_dir}; the package holds {present or 'no weight files'}."] - if metadata.get(PACKAGE_COMPLETE_FIELD) is False: - stage = metadata.get("stage") - reached = f" It reached the {stage!r} stage." if stage else "" - lines.append( - f"This package is marked incomplete: the training run that wrote it did not finish.{reached} " - "Re-run training to completion; a partial package cannot be loaded." - ) - return " ".join(lines) + present = _present_files(root_dir) + stage = metadata.get("stage") + reached = f" It reached the {stage!r} stage." if stage else "" + return ( + f"Refusing to load {root_dir}: it is marked incomplete, so the training run that wrote it " + f"did not finish.{reached} The package holds {present or 'no weight files'}. " + "Re-run training to completion; a partial package cannot be loaded." + ) + + +def _missing_files_message(root_dir: str, missing: list[str], metadata: dict[str, Any]) -> str: + """Explain a finished package that is nevertheless missing weights. + + Reached only for packages marked complete -- an unfinished one is refused + before file discovery, by `_incomplete_package_message` -- so this must not + blame the run for stopping early. Something else is wrong: the wrong + `checkpoint_files` for this model, or an upload that lost a file. + """ + del metadata # completeness is decided before this point + present = _present_files(root_dir) + return f"Missing checkpoint file(s) {missing} in {root_dir}; the package holds {present or 'no weight files'}." def _discover_package_files(root_dir: str, metadata: dict[str, Any]) -> list[str]: diff --git a/engiopt/core.py b/engiopt/core.py index c257aa7b..68f7c0c7 100644 --- a/engiopt/core.py +++ b/engiopt/core.py @@ -301,6 +301,8 @@ def __init__( condition_stats: tuple[list[float], list[float]] | None = None, checkpoint_revision: str | None = None, checkpoint_hash: str | None = None, + checkpoint_repo: str | None = None, + checkpoint_path: str | None = None, ) -> None: self.problem = problem self.problem_id = problem_id @@ -313,6 +315,16 @@ def __init__( """Commit the checkpoint was downloaded at, when it came from the Hub.""" self.checkpoint_hash = checkpoint_hash """Content hash of the weight files, identifying the exact weights scored.""" + self.checkpoint_repo = checkpoint_repo + """HF repo holding the weights, e.g. `someone/engiopt-gan-cnn-2d`. + + Together with `checkpoint_path` and `checkpoint_revision` this is the + address a leaderboard row is audited from: it is what lets a third party + fetch the exact weights behind a score and re-run it. A row that cannot + name its repo cannot be verified, and an unverifiable row is a claim. + """ + self.checkpoint_path = checkpoint_path + """Package path inside `checkpoint_repo`, e.g. `beams2d/cfg_023dd1fb/seed_1`.""" self.last_sample_seconds: float | None = None """Wall-clock seconds for the most recent `sample` call (cost metrics).""" @@ -404,6 +416,8 @@ def from_pretrained( condition_stats=condition_stats_for(resolved), checkpoint_revision=resolved.revision, checkpoint_hash=resolved.content_hash, + checkpoint_repo=resolved.repo_id, + checkpoint_path=resolved.package_path, **kwargs, ) diff --git a/engiopt/evaluate.py b/engiopt/evaluate.py index 9415e618..88e8c0b4 100644 --- a/engiopt/evaluate.py +++ b/engiopt/evaluate.py @@ -31,6 +31,10 @@ from engiopt.evaluation.leaderboard import load_from_hub from engiopt.evaluation.leaderboard import push_to_hub from engiopt.evaluation.registry import METRICS +from engiopt.evaluation.submission import FLAG_IGNORES_CONDITIONS +from engiopt.evaluation.submission import FLAG_MEMORIZED +from engiopt.evaluation.submission import FLAG_UNVERIFIED +from engiopt.evaluation.submission import integrity_flags from engiopt.utils.all_generators import BUILTIN_GENERATORS from engiopt.utils.all_generators import generators_for @@ -98,17 +102,75 @@ class Args: """Print each metric's ranking side by side after evaluating.""" list_generators: bool = False """List registered generators and exit.""" + check_availability: bool = False + """With `--list-generators`, also report which have published checkpoints. + + One Hub request per generator. Being registered means an adapter exists in + this repository; being available additionally means somebody published + weights, and only the second lets you evaluate anything.""" list_metrics: bool = False """List registered metrics and exit.""" -def _print_generators() -> None: - """Print every registered generator, including any that failed to import.""" +def _print_generators(args: Args) -> None: + """Print every registered generator, and say which ones can actually be loaded. + + Registered and available are different things, and the gap is wide enough to + mislead: an adapter is code in this repository, while a usable model also + needs published weights. Listing all fourteen as though they were + interchangeable sends people to `--generators all` and a wall of load + failures that look like bugs. + """ print(f"{len(BUILTIN_GENERATORS)} generators registered:\n") + published = _published_packages(args) for name, generator in sorted(BUILTIN_GENERATORS.items()): kinds = "/".join(generator.design_kinds) conditioning = "conditional" if generator.conditional else "unconditional" - print(f" {name:<20} {kinds:<10} {conditioning}") + availability = "" if published is None else f" {_availability_label(published.get(name))}" + print(f" {name:<20} {kinds:<10} {conditioning:<15}{availability}") + if published is not None: + missing = sorted(name for name, count in published.items() if not count) + if missing: + print( + f"\n{len(missing)} generator(s) have no published {args.problem_id} checkpoint under " + f"{args.hf_entity}: {', '.join(missing)}.\n" + "They can be trained and evaluated, but `--generators all` will report them as load " + "failures until someone publishes weights. W&B is not a checkpoint source." + ) + else: + print( + f"\nPass --check-availability to also query {args.hf_entity} for which of these have " + f"published {args.problem_id} checkpoints." + ) + + +def _published_packages(args: Args) -> dict[str, int] | None: + """How many packages each generator has for this problem, or None if not asked. + + Behind a flag because it is one Hub request per generator, and `-h`-adjacent + commands should not depend on the network. + """ + if not args.check_availability: + return None + from engiopt.checkpoint_store import build_hf_repo_id + from engiopt.checkpoint_store import list_packages + + counts: dict[str, int] = {} + for name in BUILTIN_GENERATORS: + repo = build_hf_repo_id(args.hf_entity, args.hf_repo_prefix, name) + try: + counts[name] = len(list_packages(repo, args.problem_id)) + # A repo that does not exist is the answer, not an error. + except Exception: # noqa: BLE001 + counts[name] = 0 + return counts + + +def _availability_label(count: int | None) -> str: + """One-word availability, with the package count when there is one.""" + if not count: + return "no published checkpoints" + return f"{count} package(s)" def _print_metrics() -> None: @@ -245,7 +307,7 @@ def main(args: Args) -> int: discard the rest of the results. """ if args.list_generators: - _print_generators() + _print_generators(args) if args.list_metrics: _print_metrics() if args.list_generators or args.list_metrics: @@ -280,24 +342,59 @@ def main(args: Args) -> int: print(f"\n{board.to_string(index=False)}\n") print(f"Wrote {len(board)} rows to {destination}") - _publish(args, board) + _publish(args, evaluator, board) return 0 -def _publish(args: Args, board: pd.DataFrame) -> None: +def _publish(args: Args, evaluator: Evaluator, board: pd.DataFrame) -> None: """Push the results wherever the flags asked, then print the ranking comparison.""" if args.push_to: - merged = push_to_hub(board, args.push_to) + merged = push_to_hub(board, args.push_to, eval_spec=evaluator.spec) print(f"Published {len(board)} row(s) to {args.push_to}; board now holds {len(merged)} rows.") + print( + "These rows are unverified. They will not be ranked until a runner re-fetches the " + f"checkpoints and reproduces the scores: `python -m engiopt.verify --board {args.push_to}`." + ) if args.attach_metrics: _attach_metrics_to_checkpoints(args, board) + _print_integrity_warnings(board) + if args.show_disagreement and len(board) > 1: ranked = [m for m in board.columns if m in METRICS and METRICS[m].higher_is_better is not None] if len(ranked) > 1: + # Freshly computed rows are unverified by construction, so the + # eligibility filter would empty this table. It is a local preview + # of one run, not the public ordering. print("\nRankings by metric (1 = best) -- where these disagree is the interesting part:\n") - print(disagreement(board, ranked).to_string()) + print(disagreement(board, ranked, eligible_only=False).to_string()) + + +def _print_integrity_warnings(board: pd.DataFrame) -> None: + """Say so, locally and immediately, when a row will not be rankable. + + Better here than on the board: the submitter finds out while they can still + do something about it, rather than after publishing and wondering why their + model never appears in the ordering. + """ + for row in board.to_dict("records"): + flags = [flag for flag in integrity_flags(row) if flag != FLAG_UNVERIFIED] + if not flags: + continue + label = f"{row.get('algo_id')} seed {row.get('seed')}" + if FLAG_MEMORIZED in flags: + print( + f"\n [{label}] copy_rate={row.get('copy_rate'):.2f}: most of this batch reproduces designs " + "from the dataset rather than generating them. Its distribution and performance scores " + "measure retrieval, and it will be published but not ranked." + ) + if FLAG_IGNORES_CONDITIONS in flags: + print( + f"\n [{label}] cond_sens=0: output did not change at all when the conditions were shuffled, " + "though this model declares itself conditional. Its conditions are most likely not reaching " + "the network. It will be published but not ranked." + ) if __name__ == "__main__": diff --git a/engiopt/evaluation/context.py b/engiopt/evaluation/context.py index d341bd5c..578e17bb 100644 --- a/engiopt/evaluation/context.py +++ b/engiopt/evaluation/context.py @@ -21,6 +21,8 @@ from engiopt import metrics as metrics_mod if TYPE_CHECKING: + from collections.abc import Callable + from datasets import Dataset from engibench.core import Problem import numpy.typing as npt @@ -93,6 +95,17 @@ class EvaluationContext: objective_weight_condition: Name of a per-sample condition carrying the trade-off instead. thermoelastic2d's `weight` splits the first two objectives as `(w, 1 - w)`; any remaining objectives get zero. + copy_corpus_fn: Returns the designs a model could plausibly have copied + -- the dataset designs it may have trained on, plus the reference + optima it is being scored against. Deferred behind a callable so + that fetching a training split is paid for only when a memorization + metric actually runs, and shared across every model in a sweep. + copy_tol: RMS per-element distance below which a generated design counts + as a copy of something in the corpus. + resample_permuted: Re-runs the generator on a permutation of the same + conditions, with the same latent draw. `None` when the comparison is + unavailable -- an unconditional problem, or a context built directly + from arrays rather than from a live generator. """ problem: Problem @@ -106,6 +119,9 @@ class EvaluationContext: sample_seconds: float | None = None objective_weights: tuple[float, ...] | None = None objective_weight_condition: str | None = None + copy_corpus_fn: Callable[[], npt.NDArray[Any]] | None = None + copy_tol: float = 0.01 + resample_permuted: Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None = None @property def n_samples(self) -> int: @@ -130,6 +146,66 @@ def ref_flat(self) -> npt.NDArray[Any]: flattened = [np.asarray(spaces.flatten(self.problem.design_space, design)) for design in self.ref_designs] return np.asarray(flattened) + @cached_property + def copy_corpus(self) -> npt.NDArray[Any] | None: + """Flattened designs a generator could have memorized, or None if there are none. + + The scored reference designs are always in here, even though they are + also `mmd`'s comparison target. They are the single most attractive + thing to copy precisely *because* the protocol is public and names them, + so a memorization check that omitted them would miss the easiest attack + on this board. `copy_corpus_fn` widens the corpus to the training split + the model was actually fitted on. + + Parts whose flattened width does not match the generated designs are + dropped rather than raising: a mismatched corpus makes the check + unavailable, and that is not a reason to fail an evaluation. + """ + width = self.gen_flat.shape[1] + parts = [part for part in (self.ref_flat,) if part.shape[1] == width] + if self.copy_corpus_fn is not None: + extra = np.asarray(self.copy_corpus_fn()) + if len(extra): + flattened = extra.reshape(len(extra), -1) + if flattened.shape[1] == width: + parts.append(flattened) + return np.concatenate(parts) if parts else None + + @cached_property + def nearest_corpus_distance(self) -> npt.NDArray[Any] | None: + """Per-design RMS distance to the closest design in `copy_corpus`. + + Normalized by the square root of the design dimension so the number is a + *per-element* deviation. That makes one tolerance meaningful across + problems of different resolution, where a raw L2 norm would not be. + """ + corpus = self.copy_corpus + if corpus is None or not len(corpus): + return None + from scipy.spatial.distance import cdist + + distances = cdist(self.gen_flat, corpus, "euclidean") + return np.asarray(distances.min(axis=1) / np.sqrt(self.gen_flat.shape[1])) + + @cached_property + def permuted_designs(self) -> npt.NDArray[Any] | None: + """Designs the generator produces when the conditions are shuffled between samples. + + Same latent draw, same model, different brief. A model that reads its + conditions produces something different; a model that ignores them + produces the identical batch in a different order at best, and an + identical batch outright at worst. Comparing the two is what turns "this + model declares itself conditional" into a measurement. + """ + if self.resample_permuted is None or self.n_samples < 2: # noqa: PLR2004 - a permutation needs two rows + return None + # Derangement by rotation: every sample is scored against a condition + # that belongs to some *other* sample, with no fixed points to dilute + # the signal, and no RNG to make the metric irreproducible. + permutation = np.roll(np.arange(self.n_samples), 1) + permuted = self.resample_permuted(permutation) + return None if permuted is None else np.asarray(permuted).reshape(self.n_samples, -1) + def condition_at(self, index: int) -> dict[str, Any] | None: """Conditions for sample `index`, or None when the problem is unconditional.""" return self.conditions[index] if self.conditions is not None else None diff --git a/engiopt/evaluation/evaluator.py b/engiopt/evaluation/evaluator.py index daf7bd28..81aa227d 100644 --- a/engiopt/evaluation/evaluator.py +++ b/engiopt/evaluation/evaluator.py @@ -17,9 +17,9 @@ import functools import importlib.metadata from pathlib import Path -import subprocess from typing import Any, TYPE_CHECKING +import numpy as np import pandas as pd from engiopt.core import ConditionBatch @@ -30,13 +30,17 @@ from engiopt.evaluation.registry import MetricSpec from engiopt.evaluation.spec import EvalSpec from engiopt.evaluation.spec import ResolvedSpec +from engiopt.evaluation.spec import source_checkout_commit from engiopt.utils.all_generators import design_kind_of # Importing the metrics package is what populates the registry. import engiopt.evaluation.metrics # noqa: F401 # isort: skip if TYPE_CHECKING: + from collections.abc import Callable + from engibench.core import Problem + import numpy.typing as npt import torch as th from engiopt.core import Generator @@ -49,11 +53,17 @@ "spec_version", "n_samples", "sample_seconds", + "checkpoint_repo", + "checkpoint_path", "checkpoint_revision", "checkpoint_hash", "code_version", "engibench_version", "evaluated_at", + "verified", + "verified_by", + "verified_at", + "flags", ) """Columns identifying *what was measured*, as opposed to metric values. @@ -62,6 +72,17 @@ and seed, or changing the training code, yields different weights under the same name, so without it a row cannot be traced back to the model that earned it. +`checkpoint_repo` and `checkpoint_path` complete that into an *address*. A hash +proves two rows describe the same weights; only the address lets someone else go +and fetch them. Those four columns together are what make a row a reproducible +claim rather than an assertion, and they are the input to `engiopt.verify`. + +`verified` is never set by whoever computed the row. It is stamped by a runner +that re-fetched the checkpoint and reproduced the numbers; see +`engiopt.evaluation.verify`. `flags` carries the integrity checks a row tripped, +which is how a retrieval system stays visible in the table while staying out of +the ranking. + `engibench_version` is the EngiBench that *ran* the evaluation, as opposed to the one the spec was frozen against. A change to `simulate` or `optimize` moves the optimality gaps without touching any condition name or dataset, so the spec's @@ -142,15 +163,7 @@ def context_for(self, generator: Generator) -> EvaluationContext: f"{generator.algo_id!r} supports {generator.design_kinds} design spaces, " f"but {self.problem_id!r} is {kind!r}." ) - designs = generator.sample( - ConditionBatch( - tensor=self.resolved.conditions_tensor, - dataset=self.resolved.conditions, - keys=self.resolved.condition_keys, - ), - n=self.resolved.n_samples, - seed=getattr(generator, "seed", None), - ) + designs = self._sample(generator) return EvaluationContext( problem=self.problem, problem_id=self.problem_id, @@ -163,8 +176,74 @@ def context_for(self, generator: Generator) -> EvaluationContext: sample_seconds=generator.last_sample_seconds, objective_weights=self.spec.objective_weights, objective_weight_condition=self.spec.objective_weight_condition, + copy_corpus_fn=self.copy_corpus, + copy_tol=self.spec.copy_tol, + resample_permuted=self._permuted_sampler(generator), + ) + + def _sample(self, generator: Generator, order: npt.NDArray[Any] | None = None) -> npt.NDArray[Any]: + """Draw this spec's batch from a generator, optionally reordering the conditions. + + The seed is the generator's own, so re-drawing with a different + condition order reuses the identical latent sample and isolates the + model's response to the conditions themselves. + """ + conditions = self.resolved.conditions + tensor = self.resolved.conditions_tensor + if order is not None: + conditions = conditions.select(order) if conditions is not None else None + tensor = tensor[order] if tensor is not None else None + return generator.sample( + ConditionBatch(tensor=tensor, dataset=conditions, keys=self.resolved.condition_keys), + n=self.resolved.n_samples, + seed=getattr(generator, "seed", None), ) + def _permuted_sampler(self, generator: Generator) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None: + """A callable re-drawing from `generator` under a permutation, or None if meaningless. + + Withheld when the problem supplies no conditions at all: shuffling + nothing measures nothing, and a metric that reported 0 there would say + "this model ignores its conditions" about a model that was never given + any. + """ + if self.resolved.conditions_tensor is None and self.resolved.conditions is None: + return None + return lambda order: self._sample(generator, order) + + @functools.cached_property + def copy_corpus(self) -> Callable[[], npt.NDArray[Any]]: + """Designs from the training split that a model on this problem could have memorized. + + Built once per evaluator and shared by every model in a sweep, and + deferred behind a callable so a run that selects no memorization metric + never pays for the fetch. + """ + + @functools.cache + def corpus() -> npt.NDArray[Any]: + return self._draw_copy_corpus() + + return corpus + + def _draw_copy_corpus(self) -> npt.NDArray[Any]: + """Subsample the training split's optimal designs, deterministically. + + Drawn with the spec's own `condition_seed`, so the corpus a model is + checked against is as reproducible as the conditions it is scored on -- + an audit that drew a different corpus could reach a different verdict. + """ + try: + train = self.problem.dataset["train"] + designs = np.asarray(train["optimal_design"]) + # A problem with no training split simply has no wider corpus; the + # reference designs still are one, and they are the case that matters. + except (KeyError, TypeError, AttributeError): + return np.empty((0, 0)) + size = min(self.spec.copy_corpus_size, len(designs)) + rng = np.random.default_rng(self.spec.condition_seed) + return designs[rng.choice(len(designs), size, replace=False)] + def score( self, generator: Generator, @@ -209,6 +288,13 @@ def _selected(self, only: list[str] | None, *, include_expensive: bool) -> list[ return specs def _provenance(self, generator: Generator, ctx: EvaluationContext) -> dict[str, Any]: + """The row's identity and audit trail, minus the verification stamp. + + `verified` is deliberately False here and cannot be set from this side. + The whole point of the column is that it records someone *else* having + re-fetched these weights and reproduced these numbers, so a value + written by the process that computed them would mean nothing. + """ return { "problem_id": self.problem_id, "algo_id": generator.algo_id, @@ -217,11 +303,17 @@ def _provenance(self, generator: Generator, ctx: EvaluationContext) -> dict[str, "spec_version": self.spec.version, "n_samples": ctx.n_samples, "sample_seconds": ctx.sample_seconds, + "checkpoint_repo": getattr(generator, "checkpoint_repo", None), + "checkpoint_path": getattr(generator, "checkpoint_path", None), "checkpoint_revision": getattr(generator, "checkpoint_revision", None), "checkpoint_hash": getattr(generator, "checkpoint_hash", None), "code_version": code_version(), "engibench_version": engibench_version(), "evaluated_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"), + "verified": False, + "verified_by": None, + "verified_at": None, + "flags": "", } # ------------------------------------------------------------------ @@ -280,22 +372,19 @@ def code_version() -> str: Two evaluations of the same checkpoint can differ if the evaluation code changed between them, so the row records which code it was. + + The sha is taken with the same guards as EngiBench's, and for the same + reason: `git -C` searches upward, so an EngiOpt wheel installed into a + virtualenv inside some *other* repository would otherwise report that + repository's commit as the code that produced the score. A row claiming a + commit it did not run is worse than a row claiming none, because verifying + it means checking out the wrong tree. """ try: version = importlib.metadata.version("engiopt") except importlib.metadata.PackageNotFoundError: version = "unknown" - try: - sha = subprocess.run( - ["git", "-C", str(Path(__file__).resolve().parent), "rev-parse", "--short", "HEAD"], - capture_output=True, - text=True, - check=True, - timeout=5, - ).stdout.strip() - # Outside a checkout (an installed wheel, a container) the version is all there is. - except (OSError, subprocess.SubprocessError): - return version + sha = source_checkout_commit(Path(__file__).resolve().parent.parent.parent, distribution="engiopt") return f"{version}+{sha}" if sha else version diff --git a/engiopt/evaluation/leaderboard.py b/engiopt/evaluation/leaderboard.py index b9848df0..9fdfaa33 100644 --- a/engiopt/evaluation/leaderboard.py +++ b/engiopt/evaluation/leaderboard.py @@ -21,13 +21,22 @@ if TYPE_CHECKING: from engiopt.evaluation.registry import MetricRegistry + from engiopt.evaluation.spec import EvalSpec -ROW_KEY = ["problem_id", "algo_id", "config_fingerprint", "seed", "spec_version"] +ROW_KEY = ["problem_id", "algo_id", "config_fingerprint", "checkpoint_repo", "seed", "spec_version"] """Columns that uniquely identify a leaderboard row. `config_fingerprint` is what separates two hyperparameter settings of the same algorithm. Without it a 50-config sweep would collapse onto one row per seed, silently keeping only whichever ran last. + +`checkpoint_repo` is what separates two *people*. On a public board the other +four columns are not unique: two contributors who both train `cgan_cnn_2d` with +this repository's default hyperparameters on seed 1 produce the same key and +genuinely different weights. Merging on that key would let either one's push +silently replace the other's verified row -- no malice required, and no trace +left afterwards. Their checkpoints live in their own repos, so that is the +column that tells them apart. """ RANK_PARTITION = ["problem_id", "spec_version"] @@ -38,13 +47,15 @@ whole table. """ -ENTRY_KEY = ["problem_id", "algo_id", "config_fingerprint", "spec_version"] -"""What a *ranked entry* is: one configuration of one algorithm under one spec. +ENTRY_KEY = ["problem_id", "algo_id", "config_fingerprint", "checkpoint_repo", "spec_version"] +"""What a *ranked entry* is: one person's configuration of one algorithm under one spec. `ROW_KEY` minus `seed`, because seeds are what a ranking aggregates over. Everything else must match: averaging one configuration's score with another's, or a score under `v1` with a score under `v2`, produces a number that describes -no model and no protocol. +no model and no protocol. `checkpoint_repo` is included for the same reason it is +in `ROW_KEY` -- two contributors' identically-configured models are two entries, +and pooling their seeds would produce a median over a mixture. """ LEADERBOARD_FILE = "leaderboard.csv" @@ -77,6 +88,8 @@ def rank( *, registry: MetricRegistry | None = None, aggregate: str = "median", + eval_spec: EvalSpec | None = None, + eligible_only: bool = True, ) -> pd.DataFrame: """Rank models by one metric, aggregating across seeds. @@ -87,6 +100,11 @@ def rank( aggregate: How to combine seeds. Defaults to `median`, because simulation-backed objectives are heavy-tailed and a single blown-up run otherwise decides the ranking. + eval_spec: The spec these rows were scored under, for `required_seeds`. + Without it the seed-coverage rule cannot be applied and is skipped. + eligible_only: Rank only rows a public board should rank; see + `eligible_rows`. Turn it off to see the raw table, including rows + nobody has re-run. Returns: One row per `ENTRY_KEY` sorted best-first, with a `rank` column and the @@ -99,10 +117,18 @@ def rank( spec = (registry or METRICS)[metric] if spec.higher_is_better is None: - raise ValueError(f"Metric {metric!r} has no ranking direction; it is diagnostic only.") + raise ValueError( + f"Metric {metric!r} has no ranking direction; it is diagnostic only. " + "Metrics like `novelty` and `cond_sens` say whether a score means what it looks like, " + "and ranking on them would just reward whichever extreme happens to be unoccupied." + ) + + ranked = eligible_rows(frame, eval_spec) if eligible_only else frame + if ranked.empty: + return ranked.assign(**{f"{metric}_{aggregate}": [], "n_seeds": [], "rank": []}) grouped = ( - frame.groupby(entry_key(frame), as_index=False) + ranked.groupby(entry_key(ranked), as_index=False) .agg(value=(metric, aggregate), n_seeds=(metric, "count")) .sort_values("value", ascending=not spec.higher_is_better) .reset_index(drop=True) @@ -112,6 +138,51 @@ def rank( return grouped.rename(columns={"value": f"{metric}_{aggregate}"}) +def eligible_rows(frame: pd.DataFrame, eval_spec: EvalSpec | None = None) -> pd.DataFrame: + """The rows a public ranking may draw on. + + Three conditions, each answering a different way a table of self-reported + numbers goes wrong: + + 1. **Verified.** Somebody re-fetched these weights and reproduced these + numbers. Without this the ranking is ordering claims. + 2. **Unflagged.** A retrieval system posts an excellent `fog`, and it is not + wrong to publish that -- it is wrong to call it first place. Flagged rows + stay in the table and out of the order. + 3. **Complete seeds.** An entry must cover every seed in + `EvalSpec.required_seeds`. A submitter who runs twenty seeds and publishes + their best three otherwise gets a median computed over a maximum, which is + the cheapest way to climb a board and leaves no trace in any single row. + + Rows are dropped, never edited, so the board itself keeps everything. + """ + from engiopt.evaluation.submission import is_flagged + + eligible = frame + if "verified" in eligible.columns: + eligible = eligible[eligible["verified"].fillna(value=False).astype(bool)] + if "flags" in eligible.columns: + eligible = eligible[~eligible["flags"].map(is_flagged)] + if eval_spec is not None and eval_spec.required_seeds and "seed" in eligible.columns: + eligible = _entries_covering_seeds(eligible, eval_spec.required_seeds) + return eligible + + +def _entries_covering_seeds(frame: pd.DataFrame, required: tuple[int, ...]) -> pd.DataFrame: + """Keep only entries that supply every required seed. + + Applied per `ENTRY_KEY` rather than per row, because seed coverage is a + property of the group a ranking aggregates over -- no single row can be + inspected to tell whether its siblings exist. + """ + if frame.empty: + return frame + key = entry_key(frame) + needed = set(required) + covered = frame.groupby(key)["seed"].transform(lambda seeds: needed.issubset(set(seeds.astype(int)))) + return frame[covered.astype(bool)] + + def entry_key(frame: pd.DataFrame) -> list[str]: """The `ENTRY_KEY` columns present in `frame`, for grouping. @@ -131,6 +202,8 @@ def disagreement( *, registry: MetricRegistry | None = None, aggregate: str = "median", + eval_spec: EvalSpec | None = None, + eligible_only: bool = True, ) -> pd.DataFrame: """Show each metric's ranking side by side, to expose where they disagree. @@ -144,7 +217,9 @@ def disagreement( key = entry_key(frame) boards = [] for metric in metrics: - board = rank(frame, metric, registry=registry, aggregate=aggregate) + board = rank( + frame, metric, registry=registry, aggregate=aggregate, eval_spec=eval_spec, eligible_only=eligible_only + ) boards.append(board.set_index(key)["rank"].rename(metric)) return pd.concat(boards, axis=1).sort_values(metrics[0]) @@ -227,6 +302,8 @@ def push_to_hub( private: bool = False, commit_message: str | None = None, max_attempts: int = 3, + eval_spec: EvalSpec | None = None, + check_admission: bool = True, ) -> pd.DataFrame: """Merge `new_rows` into the published leaderboard and upload the result. @@ -237,6 +314,11 @@ def push_to_hub( job published in between, the commit is rejected rather than overwriting it, and the read-merge-upload cycle is retried against the newer board. + Rows are put through `prepare_submission` first. That refuses anything with + no fetchable checkpoint -- a row nobody can re-run is not a result -- and + clears any verification stamp, since publishing is exactly the moment a + submitter would like to award themselves one. + Args: new_rows: Rows to publish, as returned by `Evaluator.leaderboard`. repo_id: HuggingFace dataset repo, e.g. `"IDEALLab/engiopt-leaderboard"`. @@ -244,16 +326,26 @@ def push_to_hub( private: Create the repo private if it does not exist yet. commit_message: Defaults to a summary of what was added. max_attempts: How many times to retry after losing a race. + eval_spec: Spec the rows were scored under, for the flag thresholds. + check_admission: Publish rows unchecked. Only for a verification runner + writing back rows it produced itself, which is the one caller whose + `verified=True` is not a self-assertion. Returns: The full merged leaderboard as uploaded. Raises: RuntimeError: If every attempt lost the race to a concurrent publisher. + SubmissionRejectedError: If a row cannot be published; see `submission`. """ from huggingface_hub import HfApi from huggingface_hub.errors import HfHubHTTPError + from engiopt.evaluation.submission import prepare_submission + + if check_admission: + new_rows = prepare_submission(new_rows, eval_spec) + api = HfApi(token=token) api.create_repo(repo_id=repo_id, repo_type="dataset", private=private, exist_ok=True) models = ", ".join(sorted(new_rows["algo_id"].unique())) if "algo_id" in new_rows else "results" diff --git a/engiopt/evaluation/metrics/builtin.py b/engiopt/evaluation/metrics/builtin.py index 2f912f71..a7c089b8 100644 --- a/engiopt/evaluation/metrics/builtin.py +++ b/engiopt/evaluation/metrics/builtin.py @@ -56,6 +56,90 @@ def dpp(ctx: EvaluationContext) -> float: return float(metrics_mod.dpp_diversity(ctx.gen_flat, sigma=ctx.sigma)) +# ---------------------------------------------------------------------- +# Memorization: did the model generate this, or retrieve it? +# ---------------------------------------------------------------------- + + +@register_metric( + "novelty", + family="memorization", + cost="cheap", + higher_is_better=None, + outputs=("novelty", "copy_rate"), + description="Distance from generated designs to the nearest design the model could have memorized.", +) +def novelty(ctx: EvaluationContext) -> dict[str, float]: + """How far the generated designs sit from the corpus a model could copy from. + + The evaluation protocol is public: which conditions are scored, and the + dataset-optimal design for each one, can be recomputed by anyone from the + committed spec. A lookup table keyed on the condition vector therefore tops + `mmd`, `iog`, and `fog` -- not by cheating the implementation, but because + those metrics are *defined* as closeness to exactly the designs it returns. + No amount of care in the evaluator changes that; the only defence available + to a public board is to measure retrieval and say so. + + Two columns, because they answer different questions: + + - `novelty` -- mean per-element RMS distance to the nearest corpus design. + Diagnostic, deliberately: it has no good direction. Zero means the model + is a retrieval system, but large means only that the output is far from + the data, which pure noise also achieves. Read it next to `mmd` and + `viol`, never on its own. + - `copy_rate` -- the fraction of designs closer than `copy_tol`, i.e. the + share of this batch that is a reproduction rather than a generation. This + is the one that gates a submission. + + Returns NaN when no corpus is available, rather than claiming novelty that + was never checked. + """ + distances = ctx.nearest_corpus_distance + if distances is None: + return {"novelty": float("nan"), "copy_rate": float("nan")} + return { + "novelty": float(np.mean(distances)), + "copy_rate": float(np.mean(distances < ctx.copy_tol)), + } + + +# ---------------------------------------------------------------------- +# Conditions: does the model actually read the brief it was given? +# ---------------------------------------------------------------------- + + +@register_metric( + "cond_sens", + family="conditions", + cost="cheap", + higher_is_better=None, + description="How much the generated design changes when the conditions are shuffled.", +) +def cond_sens(ctx: EvaluationContext) -> float: + """Mean per-element RMS change in output when each sample is given another's conditions. + + Sampled from the same seed, so the latent draw is held fixed and the only + thing that varies is the brief. A model that ignores its conditions returns + the identical batch and scores 0; a model that responds to them scores the + size of that response. + + Diagnostic rather than ranked. Being unconditional is a legitimate thing for + a model to be -- the contract says so, and such models are still handed + conditions so that this can be measured -- while a large response is not by + itself a good one, since responding *wrongly* also moves the output. + Its job is to stop an unconditional model quietly collecting a conditional + model's `mmd` score, not to be maximized. + + Returns NaN when there is nothing to compare: an unconditional problem, or + fewer than two samples. + """ + permuted = ctx.permuted_designs + if permuted is None: + return float("nan") + deltas = np.linalg.norm(ctx.gen_flat - permuted, axis=1) / np.sqrt(ctx.gen_flat.shape[1]) + return float(np.mean(deltas)) + + # ---------------------------------------------------------------------- # Feasibility: is the design admissible, and does it respect its budget? # ---------------------------------------------------------------------- diff --git a/engiopt/evaluation/registry.py b/engiopt/evaluation/registry.py index 54db5113..eea3cce9 100644 --- a/engiopt/evaluation/registry.py +++ b/engiopt/evaluation/registry.py @@ -27,8 +27,14 @@ def mmd(ctx: EvaluationContext) -> float: MetricCost = Literal["cheap", "expensive"] """`cheap` metrics never invoke a simulator or optimizer; `expensive` ones may.""" -MetricFamily = Literal["feasibility", "conditions", "performance", "distribution", "diversity", "cost"] -"""The question a metric answers. Also the grouping used by leaderboard views.""" +MetricFamily = Literal["feasibility", "conditions", "performance", "distribution", "diversity", "memorization", "cost"] +"""The question a metric answers. Also the grouping used by leaderboard views. + +`memorization` and `conditions` are integrity families: they do not say how good +a model is, they say whether its other scores mean what they appear to. A +retrieval system and an unconditional model both post excellent `mmd`, and only +these two families distinguish them from a model that earned it. +""" MetricFn = Callable[["EvaluationContext"], "float | dict[str, float]"] diff --git a/engiopt/evaluation/spec.py b/engiopt/evaluation/spec.py index 238b4623..faf6c36d 100644 --- a/engiopt/evaluation/spec.py +++ b/engiopt/evaluation/spec.py @@ -96,6 +96,21 @@ class EvalSpec: trade-off instead of fixing it. thermoelastic2d's `weight` splits the first two objectives as `(w, 1 - w)`, so a purely structural sample (`w = 1`) is unaffected by thermal compliance. + required_seeds: Which training seeds an entry must supply before it can + be ranked. Named rather than counted, because a count only stops a + submitter running one seed -- it does not stop them running twenty + and publishing their best three, which produces a median over a + maximum. Fixing *which* seeds removes the choice. + copy_tol: Per-element RMS distance below which a generated design counts + as a copy of a design the model could have memorized. See the + `novelty` metric. + max_copy_rate: The share of copied designs above which an entry is + flagged and left out of the ranking. Set to 1.0 to disable the gate + and report `copy_rate` without acting on it. + copy_corpus_size: How many dataset designs to draw as the memorization + corpus. The scored reference designs are always included on top of + these, since the public spec names them and they are the most + attractive thing to copy. condition_digest: Hash of the drawn indices, condition values, and reference designs. Recomputed at evaluation time and compared, so an upstream dataset change is caught instead of silently shifting every @@ -122,12 +137,16 @@ class EvalSpec: version: str = "v1" n_samples: int = 50 condition_seed: int = 1 - metrics: tuple[str, ...] = ("mmd", "dpp", "viol", "iog", "cog", "fog") + metrics: tuple[str, ...] = ("mmd", "dpp", "novelty", "cond_sens", "viol", "iog", "cog", "fog") sigma: float = 10.0 volfrac_tol: float = 0.01 volume_condition: str | None = None objective_weights: tuple[float, ...] | None = None objective_weight_condition: str | None = None + required_seeds: tuple[int, ...] = (1, 2, 3) + copy_tol: float = 0.01 + max_copy_rate: float = 0.5 + copy_corpus_size: int = 512 condition_digest: str | None = None dataset_id: str | None = None dataset_revision: str | None = None @@ -142,6 +161,7 @@ def __post_init__(self) -> None: carry lists and compare unequal to the spec that produced it. """ object.__setattr__(self, "metrics", tuple(self.metrics)) + object.__setattr__(self, "required_seeds", tuple(int(seed) for seed in self.required_seeds)) if self.objective_weights is not None: object.__setattr__(self, "objective_weights", tuple(self.objective_weights)) if self.problem_conditions is not None: @@ -337,12 +357,29 @@ def engibench_version() -> str: import engibench version = getattr(engibench, "__version__", "unknown") - sha = _installed_vcs_commit() or _source_checkout_commit(Path(engibench.__file__).resolve().parent.parent) + sha = source_checkout_commit(Path(engibench.__file__).resolve().parent.parent, distribution="engibench") return f"{version}+{sha}" if sha else version -def _installed_vcs_commit() -> str | None: - """The commit recorded by `pip install "engibench @ git+..."`, if it was installed that way. +def source_checkout_commit(source_root: Path, *, distribution: str) -> str | None: + """The git commit a package was installed from, or None if that cannot be established. + + Two ways a commit can be known, tried in that order: pip recorded it at + install time (PEP 610), or the package is being imported out of a source + checkout. Anything else -- a plain wheel from PyPI, a vendored copy -- + honestly has no commit, and says so. + + Args: + source_root: The directory that would be the repository root if this + package were a source checkout, i.e. the parent of the package + directory itself. + distribution: Installed distribution name, for the PEP 610 lookup. + """ + return _installed_vcs_commit(distribution) or _checkout_commit(source_root) + + +def _installed_vcs_commit(distribution: str) -> str | None: + """The commit recorded by `pip install " @ git+..."`, if it was installed that way. A non-editable VCS install leaves no `.git` directory, but pip records the exact commit in the distribution's `direct_url.json` (PEP 610). That is the @@ -351,7 +388,7 @@ def _installed_vcs_commit() -> str | None: from importlib import metadata try: - raw = metadata.distribution("engibench").read_text("direct_url.json") + raw = metadata.distribution(distribution).read_text("direct_url.json") except (metadata.PackageNotFoundError, OSError): return None if not raw: @@ -363,8 +400,8 @@ def _installed_vcs_commit() -> str | None: return str(commit)[:12] if commit else None -def _source_checkout_commit(source_root: Path) -> str | None: - """The commit of an EngiBench *source checkout*, or None if it is not one. +def _checkout_commit(source_root: Path) -> str | None: + """The commit of a *source checkout*, or None if `source_root` is not one. `git -C` searches parent directories, so asking it about a wheel unpacked into a virtualenv inside another repository answers with *that* repository's @@ -413,10 +450,16 @@ def freeze_spec( version: str = "v1", n_samples: int = 50, condition_seed: int = 1, - metrics: tuple[str, ...] = ("mmd", "dpp", "viol", "iog", "cog", "fog"), + metrics: tuple[str, ...] = EvalSpec.metrics, sigma: float = 10.0, + volume_condition: str | None = None, + volfrac_tol: float = EvalSpec.volfrac_tol, objective_weights: tuple[float, ...] | None = None, objective_weight_condition: str | None = None, + required_seeds: tuple[int, ...] = EvalSpec.required_seeds, + copy_tol: float = EvalSpec.copy_tol, + max_copy_rate: float = EvalSpec.max_copy_rate, + copy_corpus_size: int = EvalSpec.copy_corpus_size, notes: str = "", ) -> Path: """Draw a problem's test conditions once and commit them as a spec. @@ -425,6 +468,12 @@ def freeze_spec( a spec is that it stops moving. The dataset revision in force is recorded, so later uploads to the dataset repo cannot change what the spec means. + Every field of the contract is settable here. That is not merely convenient: + a parameter this entry point omits is a field it silently resets to the + dataclass default, so freezing a new version of an existing spec would + quietly hand it a different feasibility rule or a different copy gate than + the version it was meant to succeed. + Returns: Path to the written spec file. """ @@ -439,8 +488,14 @@ def freeze_spec( condition_seed=condition_seed, metrics=metrics, sigma=sigma, + volume_condition=volume_condition, + volfrac_tol=volfrac_tol, objective_weights=objective_weights, objective_weight_condition=objective_weight_condition, + required_seeds=required_seeds, + copy_tol=copy_tol, + max_copy_rate=max_copy_rate, + copy_corpus_size=copy_corpus_size, notes=notes, ).freeze(problem) path = spec.save() diff --git a/engiopt/evaluation/submission.py b/engiopt/evaluation/submission.py new file mode 100644 index 00000000..32b5003d --- /dev/null +++ b/engiopt/evaluation/submission.py @@ -0,0 +1,203 @@ +"""What a public leaderboard accepts, and what it accepts but does not rank. + +A board that anyone may write to needs two different gates, and conflating them +is the usual way these go wrong: + +- **Admission.** A row that cannot be checked by anybody is not a weak result, + it is not a result. Without a fetchable checkpoint address there is nothing to + re-run, so the number is an assertion and the table is a comment section. + Rows that fail admission are refused at publish time. +- **Integrity flags.** A row that *is* checkable but whose score does not mean + what it appears to -- a retrieval system, an unconditional model on a + conditional problem, a single cherry-picked seed. These are published and + visible, and left out of the ranking. Hiding them would throw away the most + useful thing the board knows. + +Deleting a submission is a moderation action; declining to rank one is a +statement about what the number measures. Only the second scales. +""" + +from __future__ import annotations + +import math +from typing import Any, TYPE_CHECKING + +import pandas as pd + +if TYPE_CHECKING: + from engiopt.evaluation.spec import EvalSpec + +CHECKPOINT_ADDRESS_COLUMNS = ("checkpoint_repo", "checkpoint_path", "checkpoint_revision", "checkpoint_hash") +"""What it takes to fetch the exact weights a row was scored on. + +Repo and path say where to look, revision pins the commit so a later upload +cannot change what the row refers to, and the hash detects it if one does +anyway. +""" + +IDENTITY_COLUMNS = ("problem_id", "algo_id", "config_fingerprint", "seed", "spec_version") +"""What it takes to know which model, and under which protocol, a row describes.""" + +FLAG_MEMORIZED = "memorized" +"""The batch is largely reproductions of designs the model could have seen.""" + +FLAG_UNVERIFIED = "unverified" +"""No runner has re-fetched these weights and reproduced these numbers.""" + +FLAG_INCOMPLETE_SEEDS = "incomplete_seeds" +"""The entry does not cover every seed the spec requires; see `EvalSpec.required_seeds`.""" + +FLAG_IGNORES_CONDITIONS = "ignores_conditions" +"""Output did not move when the conditions did, on a problem where they vary.""" + +DISQUALIFYING_FLAGS = frozenset({FLAG_MEMORIZED, FLAG_IGNORES_CONDITIONS}) +"""Flags that keep a row out of the ranking while leaving it on the board. + +`unverified` is not among them only because ranking already requires a +verification stamp, so listing it would be saying the same thing twice. +""" + + +def is_flagged(flags: Any) -> bool: + """Whether a row's `flags` cell carries anything disqualifying.""" + if _is_blank(flags): + return False + return bool(DISQUALIFYING_FLAGS.intersection(part.strip() for part in str(flags).split(","))) + + +def admission_problems(row: dict[str, Any]) -> list[str]: + """Why this row cannot join a public board at all, or an empty list. + + Deliberately short. Everything here is about whether the row can be + *checked*, never about whether it scored well -- a board that refuses bad + results is a board that only publishes good ones. + """ + problems = [f"missing {column}" for column in IDENTITY_COLUMNS if _is_blank(row.get(column))] + missing_address = [column for column in CHECKPOINT_ADDRESS_COLUMNS if _is_blank(row.get(column))] + if missing_address: + problems.append( + f"no fetchable checkpoint ({', '.join(missing_address)} not set) -- nobody, including you, " + "could re-run this row. Train with --checkpoint-backend hf and evaluate the published package " + "rather than a local directory." + ) + if not _is_blank(row.get("verified")) and bool(row.get("verified")): + problems.append( + "verified=True was set by the submitter. Verification records that somebody else re-ran these " + "weights; it is stamped by `python -m engiopt.verify` and cannot be self-asserted." + ) + return problems + + +def integrity_flags(row: dict[str, Any], spec: EvalSpec | None = None) -> list[str]: + """Which integrity checks this row trips. + + Row-level only. Whether an entry covers the seeds the spec requires is a + property of the *entry*, which is several rows, so it is applied during + ranking instead -- see `leaderboard.rank`. + """ + flags: list[str] = [] + copy_rate = _as_float(row.get("copy_rate")) + max_copy_rate = spec.max_copy_rate if spec is not None else 0.5 + if copy_rate is not None and copy_rate > max_copy_rate: + flags.append(FLAG_MEMORIZED) + if _declares_conditioning_it_does_not_use(row): + flags.append(FLAG_IGNORES_CONDITIONS) + if not bool(row.get("verified")): + flags.append(FLAG_UNVERIFIED) + return flags + + +def _declares_conditioning_it_does_not_use(row: dict[str, Any]) -> bool: + """Whether a model claiming to be conditional produced identical output for different briefs. + + No tolerance to choose here, and deliberately so: the comparison holds the + latent draw fixed, so a model that genuinely reads its conditions cannot + return a bit-identical batch under a derangement of them. Exactly zero is + the signature of the conditions never reaching the network. + + An honestly unconditional model is not flagged. The contract permits those, + and they are handed conditions only so that this can be measured. + """ + sensitivity = _as_float(row.get("cond_sens")) + if sensitivity is None or sensitivity > 0: + return False + from engiopt.utils.all_generators import BUILTIN_GENERATORS + + generator = BUILTIN_GENERATORS.get(str(row.get("algo_id"))) + return generator is not None and generator.conditional + + +def prepare_submission(frame: pd.DataFrame, spec: EvalSpec | None = None) -> pd.DataFrame: + """Check a table of new rows and stamp their flags, ready to publish. + + Args: + frame: Rows as produced by `Evaluator.leaderboard`. + spec: The spec they were scored under, for its gate thresholds. + + Returns: + The same rows with `verified` forced False and `flags` filled in. + + Raises: + SubmissionRejectedError: If any row fails admission. All of them are reported + at once, since fixing them one round-trip at a time is miserable. + """ + if frame.empty: + return frame + prepared = frame.copy() + if "verified" not in prepared.columns: + prepared["verified"] = False + + rejected: dict[int, list[str]] = {} + for position, record in enumerate(prepared.to_dict("records")): + problems = admission_problems(record) + if problems: + rejected[position] = problems + if rejected: + raise SubmissionRejectedError(rejected) + + # Publishing never carries a verification stamp: the point of the column is + # that somebody other than the scorer put it there. + prepared["verified"] = False + prepared["verified_by"] = None + prepared["verified_at"] = None + prepared["flags"] = [",".join(integrity_flags(record, spec)) for record in prepared.to_dict("records")] + return prepared + + +class SubmissionRejectedError(ValueError): + """Raised when rows cannot join a public board, listing every reason found.""" + + def __init__(self, problems_by_row: dict[int, list[str]]) -> None: + self.problems_by_row = problems_by_row + lines = [f"{len(problems_by_row)} row(s) cannot be published:"] + lines.extend(f" row {position}: {'; '.join(problems)}" for position, problems in sorted(problems_by_row.items())) + super().__init__("\n".join(lines)) + + +def _as_float(value: Any) -> float | None: + """A cell as a usable number, or None when it carries no measurement. + + NaN counts as no measurement, which matters: `cond_sens` is NaN on an + unconditional problem, and treating that as zero would flag every model on + it for ignoring conditions that were never supplied. + """ + if _is_blank(value): + return None + try: + number = float(value) + except (TypeError, ValueError): + return None + return None if math.isnan(number) else number + + +def _is_blank(value: Any) -> bool: + """Whether a cell carries no usable value, covering None, NaN, and empty strings.""" + if value is None: + return True + if isinstance(value, str): + return not value.strip() + try: + return bool(pd.isna(value)) + # Arrays and other non-scalars are never blank. + except (TypeError, ValueError): + return False diff --git a/engiopt/evaluation/verify.py b/engiopt/evaluation/verify.py new file mode 100644 index 00000000..0ea6c6e1 --- /dev/null +++ b/engiopt/evaluation/verify.py @@ -0,0 +1,318 @@ +"""Re-running a published row from its checkpoint address. + +The board cannot trust the numbers it is given. Whoever computed a row also +chose what to report, and no amount of schema design changes that -- the defence +against a fabricated score is not a stricter column, it is running the model +again. + +So a row is a *claim plus an address*. Verification follows the address, fetches +the weights at the revision the row named, and scores them itself. What it +publishes is **its own numbers**, not a verdict on the submitted ones. That +distinction matters more than it looks: sampling from the same seed on different +hardware genuinely produces different designs, because CUDA and CPU draw +different values from the same generator state. A pass/fail comparison would +therefore have to carry a tolerance loose enough to wave through real fudging. +Re-scoring instead makes the runner's result the result, and demotes the +submitted numbers to a claim that is reported as corroborated or not. + +Nothing here is privileged. The runner is simply the first auditor, and because +every row carries a full address, anyone can run the same command and get the +same answer:: + + python -m engiopt.verify --board IDEALLab/engiopt-leaderboard --problem-id beams2d +""" + +from __future__ import annotations + +from dataclasses import dataclass +from dataclasses import field +import datetime as dt +import math +from typing import Any, Literal, TYPE_CHECKING + +import pandas as pd + +from engiopt.checkpoint_store import resolve_checkpoint_reference +from engiopt.core import condition_keys_for +from engiopt.core import condition_stats_for +from engiopt.core import pick_device +from engiopt.evaluation.evaluator import Evaluator +from engiopt.evaluation.leaderboard import ROW_KEY +from engiopt.evaluation.registry import METRICS +from engiopt.evaluation.submission import integrity_flags +from engiopt.utils.all_generators import BUILTIN_GENERATORS + +if TYPE_CHECKING: + import torch as th + + from engiopt.core import Generator + +VerificationStatus = Literal["verified", "weights_moved", "unresolvable", "unknown_generator"] +"""Outcome of trying to reproduce one row. + +- `verified` -- the weights were fetched and re-scored. The row now carries the + runner's numbers. +- `weights_moved` -- the package at that address no longer hashes to what the + row recorded. Whatever is there now is not what earned the score, so the row + is not verified and must not be silently re-scored into place either. +- `unresolvable` -- the address did not lead to a loadable package. +- `unknown_generator` -- no registered adapter can rebuild this model, so + nothing here can run it. +""" + +CORROBORATION_TOLERANCE = 0.05 +"""Relative gap above which a submitted value is reported as uncorroborated. + +Advisory only -- it decides what the note says, never whether the row is +verified, because the published numbers are the runner's own either way. Set +loosely on purpose: this is the scale of honest hardware-to-hardware variation, +so anything inside it is not evidence of anything. +""" + + +@dataclass +class VerificationResult: + """What happened when one row was re-run.""" + + key: dict[str, Any] + status: VerificationStatus + detail: str = "" + row: dict[str, Any] | None = None + """The re-scored row, present only when `status` is `verified`.""" + uncorroborated: dict[str, tuple[float, float]] = field(default_factory=dict) + """Metric -> (submitted, recomputed) for values that did not survive re-scoring.""" + + @property + def ok(self) -> bool: + """Whether this row was successfully reproduced.""" + return self.status == "verified" + + +def verify_row( + row: dict[str, Any], + *, + verifier: str, + evaluator: Evaluator | None = None, + device: th.device | None = None, + include_expensive: bool = False, +) -> VerificationResult: + """Fetch a row's checkpoint at its recorded revision and score it again. + + Args: + row: A published leaderboard row. + verifier: Who is running this, recorded in `verified_by`. A verification + is only worth what its author is, so it is signed rather than + anonymous. + evaluator: Reuse an evaluator when sweeping a whole board; one is built + from the row's `problem_id` and `spec_version` otherwise. + device: Torch device; auto-selected when omitted. + include_expensive: Re-run the simulator-backed metrics too. Off by + default, so a full board can be checked for fabricated *cheap* + metrics without an optimizer budget. + + Returns: + A `VerificationResult` whose `row` carries the runner's own numbers. + """ + key = {column: row.get(column) for column in ROW_KEY} + algo_id = str(row.get("algo_id")) + generator_cls = BUILTIN_GENERATORS.get(algo_id) + if generator_cls is None: + return VerificationResult( + key=key, + status="unknown_generator", + detail=( + f"No registered adapter named {algo_id!r}. A model has to be in `engiopt/generators/` " + "for anyone to be able to rebuild it; submit the adapter before submitting the scores." + ), + ) + + try: + generator = _load_at_recorded_revision(row, generator_cls, device=device) + except Exception as exc: # noqa: BLE001 - any failure to fetch is the same verdict to the caller + return VerificationResult(key=key, status="unresolvable", detail=str(exc)) + + recorded_hash = row.get("checkpoint_hash") + if recorded_hash and generator.checkpoint_hash != recorded_hash: + return VerificationResult( + key=key, + status="weights_moved", + detail=( + f"The package at {row.get('checkpoint_repo')}/{row.get('checkpoint_path')} now hashes to " + f"{generator.checkpoint_hash}, but the row was scored on {recorded_hash}. The weights behind " + "this score no longer exist at that address; re-submit the current ones under their own row." + ), + ) + + evaluator = evaluator or Evaluator.for_problem( + str(row.get("problem_id")), spec=f"{row.get('problem_id')}/{row.get('spec_version')}", device=device + ) + rescored = evaluator.score(generator, include_expensive=include_expensive) + rescored.update( + { + "verified": True, + "verified_by": verifier, + "verified_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"), + } + ) + uncorroborated = _uncorroborated(row, rescored) + rescored["flags"] = ",".join(integrity_flags(rescored, evaluator.spec)) + return VerificationResult( + key=key, + status="verified", + detail=_corroboration_note(uncorroborated), + row=rescored, + uncorroborated=uncorroborated, + ) + + +def _load_at_recorded_revision( + row: dict[str, Any], generator_cls: type[Generator], *, device: th.device | None +) -> Generator: + """Rebuild the model from the exact package a row names. + + Deliberately not `from_pretrained`: that resolves by problem/algo/seed + convention and would silently follow the *current* canonical checkpoint. A + verification has to read the address written down at scoring time -- + including the revision, so that a later upload to the same path cannot + launder itself into an old row's score. + + Raises: + ValueError: If the row carries no usable checkpoint address. + """ + repo = row.get("checkpoint_repo") + path = row.get("checkpoint_path") + if not repo or not path: + raise ValueError("Row carries no checkpoint address, so there is nothing to fetch.") + + resolved = resolve_checkpoint_reference( + model_source="hf", + model_ref=f"hf://{repo}/{path}", + required_files=list(generator_cls.checkpoint_files), + revision=row.get("checkpoint_revision") or None, + ) + from engibench.utils.all_problems import BUILTIN_PROBLEMS + + problem_id = str(row.get("problem_id")) + problem = BUILTIN_PROBLEMS[problem_id]() + device = device or pick_device() + return generator_cls.build( + resolved, + problem=problem, + device=device, + problem_id=problem_id, + seed=int(row["seed"]), + run_config=resolved.run_config, + condition_keys=condition_keys_for(problem, resolved), + condition_stats=condition_stats_for(resolved), + checkpoint_revision=resolved.revision, + checkpoint_hash=resolved.content_hash, + checkpoint_repo=resolved.repo_id, + checkpoint_path=resolved.package_path, + ) + + +def _uncorroborated(submitted: dict[str, Any], recomputed: dict[str, Any]) -> dict[str, tuple[float, float]]: + """Metrics whose submitted value the re-run did not reproduce. + + Compared relative to the recomputed magnitude, so a gap means the same thing + for `mmd` near 0.01 as for `cog` in the thousands. Values the runner did not + recompute -- expensive metrics on a cheap pass -- are skipped rather than + counted as disagreement. + """ + differences: dict[str, tuple[float, float]] = {} + for name in METRICS.columns(): + if name not in submitted or name not in recomputed: + continue + claimed, actual = _as_float(submitted[name]), _as_float(recomputed[name]) + if claimed is None or actual is None: + continue + scale = max(abs(actual), abs(claimed), 1e-12) + if abs(claimed - actual) / scale > CORROBORATION_TOLERANCE: + differences[name] = (claimed, actual) + return differences + + +def _corroboration_note(uncorroborated: dict[str, tuple[float, float]]) -> str: + """One line saying whether the submitted numbers survived the re-run.""" + if not uncorroborated: + return "re-scored; submitted values corroborated" + reported = ", ".join( + f"{name} claimed {claimed:.4g} vs {actual:.4g}" for name, (claimed, actual) in sorted(uncorroborated.items()) + ) + return f"re-scored; submitted values NOT reproduced ({reported})" + + +def _as_float(value: Any) -> float | None: + """Coerce a cell to a comparable float, or None when it is not one.""" + try: + number = float(value) + except (TypeError, ValueError): + return None + return None if math.isnan(number) else number + + +def verify_board( + board: pd.DataFrame, + *, + verifier: str, + problem_id: str | None = None, + only_unverified: bool = True, + include_expensive: bool = False, + device: th.device | None = None, +) -> tuple[pd.DataFrame, list[VerificationResult]]: + """Re-run every eligible row of a board, returning the verified replacements. + + Evaluators are built once per `(problem_id, spec_version)` and reused, since + resolving a spec draws its conditions and reference designs. + + Args: + board: The published leaderboard. + verifier: Recorded in `verified_by` on every row this stamps. + problem_id: Restrict to one problem. + only_unverified: Skip rows already carrying a verification stamp. + include_expensive: Re-run simulator-backed metrics as well. + device: Torch device; auto-selected when omitted. + + Returns: + The re-scored rows (ready to merge into the board), and one + `VerificationResult` per row attempted, including the failures. + """ + if board.empty: + return board, [] + candidates = board + if problem_id is not None: + candidates = candidates[candidates["problem_id"] == problem_id] + if only_unverified and "verified" in candidates.columns: + candidates = candidates[~candidates["verified"].fillna(value=False).astype(bool)] + + evaluators: dict[tuple[str, str], Evaluator] = {} + results: list[VerificationResult] = [] + verified_rows: list[dict[str, Any]] = [] + for record in candidates.to_dict("records"): + cache_key = (str(record.get("problem_id")), str(record.get("spec_version"))) + if cache_key not in evaluators: + try: + evaluators[cache_key] = Evaluator.for_problem( + cache_key[0], spec=f"{cache_key[0]}/{cache_key[1]}", device=device + ) + # A spec that will not resolve fails every row under it, not the run. + except Exception as exc: # noqa: BLE001 + results.append( + VerificationResult( + key={column: record.get(column) for column in ROW_KEY}, + status="unresolvable", + detail=f"could not load spec {cache_key[0]}/{cache_key[1]}: {exc}", + ) + ) + continue + result = verify_row( + record, + verifier=verifier, + evaluator=evaluators[cache_key], + device=device, + include_expensive=include_expensive, + ) + results.append(result) + if result.row is not None: + verified_rows.append(result.row) + return pd.DataFrame(verified_rows), results diff --git a/engiopt/specs/beams2d/v1.json b/engiopt/specs/beams2d/v1.json index deb48742..1e3e5a43 100644 --- a/engiopt/specs/beams2d/v1.json +++ b/engiopt/specs/beams2d/v1.json @@ -1,12 +1,17 @@ { "condition_digest": "f65e2c02551e7e08", "condition_seed": 1, + "copy_corpus_size": 512, + "copy_tol": 0.01, "dataset_id": "IDEALLab/beams_2d_50_100_v0", "dataset_revision": "ccb64d4f09cf66a5ed9ba2081438d695b9e2438d", "engibench_version": "0.2.0+0a028c03d02b", + "max_copy_rate": 0.5, "metrics": [ "mmd", "dpp", + "novelty", + "cond_sens", "viol", "iog", "cog", @@ -23,6 +28,11 @@ "overhang_constraint" ], "problem_id": "beams2d", + "required_seeds": [ + 1, + 2, + 3 + ], "sigma": 10.0, "version": "v1", "volfrac_tol": 0.01, diff --git a/engiopt/specs/heatconduction2d/v1.json b/engiopt/specs/heatconduction2d/v1.json index 56875d8b..b2186a46 100644 --- a/engiopt/specs/heatconduction2d/v1.json +++ b/engiopt/specs/heatconduction2d/v1.json @@ -1,12 +1,17 @@ { "condition_digest": "40030d1b2482e003", "condition_seed": 1, + "copy_corpus_size": 512, + "copy_tol": 0.01, "dataset_id": "IDEALLab/heat_conduction_2d_v0", "dataset_revision": "9f07c1e70d244f783e34a71c96db621173871dc1", "engibench_version": "0.2.0+0a028c03d02b", + "max_copy_rate": 0.5, "metrics": [ "mmd", "dpp", + "novelty", + "cond_sens", "viol", "iog", "cog", @@ -21,6 +26,11 @@ "length" ], "problem_id": "heatconduction2d", + "required_seeds": [ + 1, + 2, + 3 + ], "sigma": 10.0, "version": "v1", "volfrac_tol": 0.01, diff --git a/engiopt/specs/photonics2d/v1.json b/engiopt/specs/photonics2d/v1.json index 62a8f73b..e00b6341 100644 --- a/engiopt/specs/photonics2d/v1.json +++ b/engiopt/specs/photonics2d/v1.json @@ -1,12 +1,17 @@ { "condition_digest": "6b8330bee5a8718c", "condition_seed": 1, + "copy_corpus_size": 512, + "copy_tol": 0.01, "dataset_id": "IDEALLab/photonics_2d_120_120_v1", "dataset_revision": "be811ff69d67cd8c79148e29f2a68d293ec00955", "engibench_version": "0.2.0+0a028c03d02b", + "max_copy_rate": 0.5, "metrics": [ "mmd", "dpp", + "novelty", + "cond_sens", "viol", "iog", "cog", @@ -25,6 +30,11 @@ "num_optimization_steps" ], "problem_id": "photonics2d", + "required_seeds": [ + 1, + 2, + 3 + ], "sigma": 10.0, "version": "v1", "volfrac_tol": 0.01, diff --git a/engiopt/specs/thermoelastic2d/v1.json b/engiopt/specs/thermoelastic2d/v1.json index 4dee1208..38c4302f 100644 --- a/engiopt/specs/thermoelastic2d/v1.json +++ b/engiopt/specs/thermoelastic2d/v1.json @@ -1,12 +1,17 @@ { "condition_digest": "818d93c2a2d5ac25", "condition_seed": 1, + "copy_corpus_size": 512, + "copy_tol": 0.01, "dataset_id": "IDEALLab/thermoelastic_2d_v1", "dataset_revision": "7d5b4a694bd1a8f3f25fd745ac9f669764e29ea5", "engibench_version": "0.2.0+0a028c03d02b", + "max_copy_rate": 0.5, "metrics": [ "mmd", "dpp", + "novelty", + "cond_sens", "viol", "iog", "cog", @@ -26,6 +31,11 @@ "weight" ], "problem_id": "thermoelastic2d", + "required_seeds": [ + 1, + 2, + 3 + ], "sigma": 10.0, "version": "v1", "volfrac_tol": 0.01, diff --git a/engiopt/verify.py b/engiopt/verify.py new file mode 100644 index 00000000..3f9816ba --- /dev/null +++ b/engiopt/verify.py @@ -0,0 +1,132 @@ +"""Re-run published leaderboard rows from their checkpoint addresses. + +Verification is what a public board has instead of trust. Every row names the +repo, path, revision, and content hash of the weights it was scored on, so +anyone can fetch them and score them again:: + + # check what the public board claims, without publishing anything + python -m engiopt.verify --board IDEALLab/engiopt-leaderboard --problem-id beams2d + + # the official runner: re-score and stamp the rows it reproduced + python -m engiopt.verify --board IDEALLab/engiopt-leaderboard --verifier IDEALLab-ci --publish + + # include the simulator-backed metrics, which the cheap pass leaves alone + python -m engiopt.verify --board IDEALLab/engiopt-leaderboard --include-expensive --publish + +Only `--publish` writes, and it writes the *runner's own* numbers. Running this +without it is a read-only audit that anybody can perform, which is the property +that keeps the official runner honest too. +""" + +from __future__ import annotations + +import dataclasses +import sys + +import pandas as pd +import tyro + +from engiopt.evaluation.leaderboard import load_from_hub +from engiopt.evaluation.leaderboard import push_to_hub +from engiopt.evaluation.verify import VerificationResult +from engiopt.evaluation.verify import verify_board + + +@dataclasses.dataclass +class Args: + """Command-line arguments for verifying a published leaderboard.""" + + board: str = "IDEALLab/engiopt-leaderboard" + """HuggingFace dataset repo holding the leaderboard, or a local CSV path.""" + verifier: str = "" + """Who is running this, recorded in `verified_by`. + + Required with `--publish`: an unsigned verification says only that some + machine somewhere agreed, which is not a thing a reader can weigh.""" + problem_id: str | None = None + """Restrict to one problem.""" + include_expensive: bool = False + """Re-run the simulator-backed metrics too. Slow, and usually a separate pass.""" + recheck_verified: bool = False + """Also re-run rows that already carry a stamp, e.g. after a metric changed.""" + publish: bool = False + """Write the re-scored rows back to the board. Off by default: auditing is + the common case, and it should never need write access.""" + output_csv: str | None = None + """Also write the re-scored rows here.""" + + +def _load_board(reference: str) -> pd.DataFrame: + """Read the board from a HuggingFace dataset repo or a local CSV.""" + if reference.endswith(".csv"): + return pd.read_csv(reference) + return load_from_hub(reference) + + +def _report(results: list[VerificationResult]) -> None: + """Print one line per row attempted, grouped so the failures are readable.""" + by_status: dict[str, list[VerificationResult]] = {} + for result in results: + by_status.setdefault(result.status, []).append(result) + + for status in sorted(by_status): + rows = by_status[status] + print(f"\n{status} ({len(rows)}):") + for result in rows: + label = f"{result.key.get('algo_id')} cfg {result.key.get('config_fingerprint')} seed {result.key.get('seed')}" + print(f" {label}: {result.detail}") + + +def main(args: Args) -> int: + """Verify a board, optionally publishing the re-scored rows. + + Returns: + A process exit status. Non-zero when a row could not be reproduced, so + an unattended runner surfaces a moved checkpoint or a missing adapter + rather than reporting a clean sweep over rows it silently skipped. + """ + if args.publish and not args.verifier: + print("--publish needs --verifier: a verification records who performed it.") + return 2 + + board = _load_board(args.board) + if board.empty: + print(f"{args.board} holds no rows; nothing to verify.") + return 0 + + verified, results = verify_board( + board, + verifier=args.verifier or "unsigned-audit", + problem_id=args.problem_id, + only_unverified=not args.recheck_verified, + include_expensive=args.include_expensive, + ) + if not results: + print("No rows were eligible for verification.") + return 0 + + _report(results) + reproduced = sum(1 for result in results if result.ok) + uncorroborated = [result for result in results if result.ok and result.uncorroborated] + print(f"\n{reproduced}/{len(results)} row(s) reproduced.") + if uncorroborated: + print( + f"{len(uncorroborated)} reproduced row(s) did not match the submitted numbers. " + "The board now carries the re-scored values." + ) + + if args.output_csv and not verified.empty: + verified.to_csv(args.output_csv, index=False) + print(f"Wrote {len(verified)} re-scored row(s) to {args.output_csv}") + + if args.publish and not verified.empty: + # The verification runner is the one publisher whose `verified=True` is + # not a self-assertion, so it is also the one that bypasses admission. + merged = push_to_hub(verified, args.board, check_admission=False) + print(f"Published {len(verified)} verified row(s); board now holds {len(merged)} rows.") + + return 0 if reproduced == len(results) else 1 + + +if __name__ == "__main__": + sys.exit(main(tyro.cli(Args))) diff --git a/tests/conftest.py b/tests/conftest.py index 91127ea3..9afbcb6e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -59,9 +59,16 @@ def __init__( @property def dataset(self) -> dict[str, Any]: - """One split whose columns are the declared conditions, all scalar.""" + """One split whose columns are the declared conditions, plus optimal designs. + + `optimal_design` is here because the memorization check draws its corpus + from the training split -- the designs a model could have been fitted on + are exactly the ones it could be reproducing. + """ rows = 3 - return {"train": FakeDataset({key: [0.5] * rows for key in self.conditions_keys})} + columns: dict[str, list[Any]] = {key: [0.5] * rows for key in self.conditions_keys} + columns["optimal_design"] = [np.full(self.design_space.shape, value) for value in (0.1, 0.2, 0.3)] + return {"train": FakeDataset(columns)} def check_constraints(self, design: Any, config: dict[str, Any]) -> FakeViolations: """Mirror `Problem.check_constraints`: declared constraints plus design-space membership. @@ -115,6 +122,15 @@ def __getitem__(self, item: Any) -> Any: def __len__(self) -> int: return len(next(iter(self._columns.values()))) + def select(self, indices: Any) -> FakeDataset: + """Reorder or subset rows, as `datasets.Dataset.select` does. + + Needed by the condition-sensitivity check, which re-draws from a + generator under a permutation of the same conditions. + """ + order = list(indices) + return FakeDataset({name: [values[i] for i in order] for name, values in self._columns.items()}) + @pytest.fixture def mixed_condition_dataset() -> FakeDataset: diff --git a/tests/test_checkpoint_integrity.py b/tests/test_checkpoint_integrity.py index 27b8fdd4..903de337 100644 --- a/tests/test_checkpoint_integrity.py +++ b/tests/test_checkpoint_integrity.py @@ -214,3 +214,128 @@ def fake_upload(*, metadata: dict[str, Any], **_: Any) -> str: ) assert all(entry[checkpoint_store.PACKAGE_COMPLETE_FIELD] is True for entry in written) + + +# ---------------------------------------------------------------------- +# The hash must cover everything that changes what the model generates +# ---------------------------------------------------------------------- + + +def test_changing_fitted_preprocessing_changes_the_hash(tmp_path: Path) -> None: + """Normalizer bounds are part of the model even though they never reach a state dict. + + Several 1D models fit min/max bounds on the training split and denormalize + through them on the way out, so a package with identical weights and + different bounds *decodes differently*. If the hash missed that, a + behaviourally different checkpoint would inherit the previous one's + leaderboard row via `--skip-existing`. + """ + package = _write_package( + tmp_path / "pkg", + metadata={"algo": "demo", "design_normalizer": {"min": [0.0], "max": [1.0]}}, + ) + before = checkpoint_store.hash_package_contents(str(package)) + + (package / "metadata.json").write_text(json.dumps({"algo": "demo", "design_normalizer": {"min": [0.0], "max": [2.0]}})) + + assert checkpoint_store.hash_package_contents(str(package)) != before + + +def test_changing_condition_statistics_changes_the_hash(tmp_path: Path) -> None: + """The same argument for the models that rescale their conditions instead.""" + package = _write_package(tmp_path / "pkg", metadata={"condition_stats": {"mean": [0.5], "std": [0.1]}}) + before = checkpoint_store.hash_package_contents(str(package)) + + (package / "metadata.json").write_text(json.dumps({"condition_stats": {"mean": [0.9], "std": [0.1]}})) + + assert checkpoint_store.hash_package_contents(str(package)) != before + + +def test_the_same_weights_at_two_addresses_hash_alike(tmp_path: Path) -> None: + """One run writes its package to both a config path and the canonical path. + + Those differ in the metadata recording where they sit, and they must not + therefore look like two different models -- the hash answers "which + weights", not "which location". + """ + shared = {"algo": "demo", "seed": 1, "condition_keys": ["volfrac"]} + config_package = _write_package( + tmp_path / "cfg", metadata={**shared, "hf_package_path": "beams2d/cfg_abc12345/seed_1", "hf_revision": "aaa"} + ) + canonical_package = _write_package( + tmp_path / "canonical", metadata={**shared, "hf_package_path": "beams2d/seed_1", "hf_revision": "bbb"} + ) + + assert checkpoint_store.hash_package_contents(str(config_package)) == checkpoint_store.hash_package_contents( + str(canonical_package) + ) + + +def test_a_promotion_does_not_change_what_the_weights_hash_to(tmp_path: Path) -> None: + """Promotion records where a package came from; it does not alter the model.""" + package = _write_package(tmp_path / "pkg", metadata={"algo": "demo"}) + before = checkpoint_store.hash_package_contents(str(package)) + + (package / "metadata.json").write_text(json.dumps({"algo": "demo", "promoted_from": "beams2d/cfg_abc12345/seed_1"})) + + assert checkpoint_store.hash_package_contents(str(package)) == before + + +# ---------------------------------------------------------------------- +# An incomplete package must never load, however complete it looks +# ---------------------------------------------------------------------- + + +def test_an_incomplete_package_is_refused_even_with_every_file_present(tmp_path: Path) -> None: + """Completeness is checked before file discovery, not while explaining a missing file. + + A stage that happens to have written everything the *next* stage's loader + asks for would otherwise load successfully, and a model from a run that + never finished would reach a leaderboard row with nothing to indicate it. + """ + package = _write_package(tmp_path / "pkg", metadata={"stage": "vqgan", checkpoint_store.PACKAGE_COMPLETE_FIELD: False}) + + with pytest.raises(FileNotFoundError) as excinfo: + checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth"], source="local" + ) + + assert "incomplete" in str(excinfo.value) + assert "did not finish" in str(excinfo.value) + + +def test_a_complete_package_still_loads(tmp_path: Path) -> None: + """The guard must not cost ordinary packages their ability to load.""" + package = _write_package(tmp_path / "pkg", metadata={checkpoint_store.PACKAGE_COMPLETE_FIELD: True}) + + resolved = checkpoint_store._load_package_from_directory( + root_dir=str(package), required_files=["generator.pth"], source="local" + ) + + assert "generator.pth" in resolved.files + + +# ---------------------------------------------------------------------- +# An upload must not blend two runs at one path +# ---------------------------------------------------------------------- + + +def test_a_complete_upload_clears_whatever_was_at_that_path() -> None: + """`upload_folder` overwrites but does not delete, so a re-train leaves strays. + + Re-training a configuration whose previous run wrote *more* files leaves the + extras in place, and discovery then serves a package mixing two runs' + weights -- which loads, because every filename a loader looks for is there. + """ + assert checkpoint_store._stale_file_patterns({checkpoint_store.PACKAGE_COMPLETE_FIELD: True}) == ["*"] + # Absent marker means an ordinary single-stage model, which is complete. + assert checkpoint_store._stale_file_patterns({}) == ["*"] + + +def test_a_staged_upload_leaves_the_earlier_stages_alone() -> None: + """VQGAN's second stage uploads `vqgan.pth` without re-uploading `cvqgan.pth`. + + Deleting there would destroy the first stage's work -- the very thing staged + uploads exist to protect. + """ + assert checkpoint_store._stale_file_patterns({checkpoint_store.PACKAGE_COMPLETE_FIELD: False}) is None diff --git a/tests/test_engibench_provenance.py b/tests/test_engibench_provenance.py index 1ad294f3..684df7aa 100644 --- a/tests/test_engibench_provenance.py +++ b/tests/test_engibench_provenance.py @@ -38,7 +38,7 @@ def read_text(name: str) -> str | None: monkeypatch.setattr(importlib_metadata, "distribution", lambda _name: _Distribution()) - assert spec_mod._installed_vcs_commit() == "0a028c03d02b" + assert spec_mod._installed_vcs_commit("engibench") == "0a028c03d02b" def test_a_wheel_without_vcs_info_reports_nothing(monkeypatch: pytest.MonkeyPatch) -> None: @@ -53,7 +53,7 @@ def read_text(_name: str) -> str | None: monkeypatch.setattr(importlib_metadata, "distribution", lambda _name: _Distribution()) - assert spec_mod._installed_vcs_commit() is None + assert spec_mod._installed_vcs_commit("engibench") is None def test_an_installed_package_is_never_read_as_a_checkout(tmp_path: Path) -> None: @@ -66,7 +66,7 @@ def test_an_installed_package_is_never_read_as_a_checkout(tmp_path: Path) -> Non site_packages = tmp_path / ".venv" / "lib" / "python3.12" / "site-packages" site_packages.mkdir(parents=True) - assert spec_mod._source_checkout_commit(site_packages) is None + assert spec_mod._checkout_commit(site_packages) is None def test_a_directory_inside_an_unrelated_repo_is_not_that_repo(tmp_path: Path) -> None: @@ -79,7 +79,7 @@ def test_a_directory_inside_an_unrelated_repo_is_not_that_repo(tmp_path: Path) - # `nested` sits inside a git repo but is not its root, so it is not a checkout # of the package -- exactly the shape a vendored install takes. - assert spec_mod._source_checkout_commit(repo / "nested") is None + assert spec_mod._checkout_commit(repo / "nested") is None def test_freeze_and_the_row_report_the_same_kind_of_string() -> None: @@ -102,15 +102,15 @@ def test_the_runtime_engibench_is_recorded_on_every_row() -> None: def test_the_version_carries_a_sha_when_one_is_identifiable(monkeypatch: pytest.MonkeyPatch) -> None: """The composed value is `+`, which is what the specs record.""" - monkeypatch.setattr(spec_mod, "_installed_vcs_commit", lambda: "abcdef123456") + monkeypatch.setattr(spec_mod, "_installed_vcs_commit", lambda _dist: "abcdef123456") assert spec_mod.engibench_version().endswith("+abcdef123456") def test_the_version_degrades_to_the_bare_release(monkeypatch: pytest.MonkeyPatch) -> None: """With nothing identifying the source, the release version is all there is.""" - monkeypatch.setattr(spec_mod, "_installed_vcs_commit", lambda: None) - monkeypatch.setattr(spec_mod, "_source_checkout_commit", lambda _root: None) + monkeypatch.setattr(spec_mod, "_installed_vcs_commit", lambda _dist: None) + monkeypatch.setattr(spec_mod, "_checkout_commit", lambda _root: None) assert "+" not in spec_mod.engibench_version() diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py index 809b7f91..7cfe9aae 100644 --- a/tests/test_evaluation.py +++ b/tests/test_evaluation.py @@ -87,8 +87,13 @@ def test_builtin_metrics_declare_their_cost() -> None: `viol` is cheap: feasibility describes the design as generated, so it is judged by a constraint check rather than by running the optimizer. + + The two integrity metrics are cheap too, and that matters more than it + sounds: `novelty` and `cond_sens` decide whether a row can be ranked at all, + so a board that could only afford the cheap pass would otherwise have to + rank models it had never checked for memorization. """ - assert {spec.name for spec in METRICS.select(cost="cheap")} == {"mmd", "dpp", "viol"} + assert {spec.name for spec in METRICS.select(cost="cheap")} == {"mmd", "dpp", "viol", "novelty", "cond_sens"} assert {spec.name for spec in METRICS.select(cost="expensive")} == {"iog", "cog", "fog"} @@ -157,6 +162,7 @@ def _board() -> pd.DataFrame: "problem_id": "p", "algo_id": "a", "config_fingerprint": "cfg_a", + "checkpoint_repo": "someone/engiopt-a", "seed": 1, "spec_version": "v1", "mmd": 0.1, @@ -166,6 +172,7 @@ def _board() -> pd.DataFrame: "problem_id": "p", "algo_id": "a", "config_fingerprint": "cfg_a", + "checkpoint_repo": "someone/engiopt-a", "seed": 2, "spec_version": "v1", "mmd": 0.3, @@ -175,6 +182,7 @@ def _board() -> pd.DataFrame: "problem_id": "p", "algo_id": "b", "config_fingerprint": "cfg_b", + "checkpoint_repo": "someone/engiopt-b", "seed": 1, "spec_version": "v1", "mmd": 0.05, @@ -200,8 +208,8 @@ def test_rank_aggregates_seeds_with_the_median_by_default() -> None: def test_disagreement_surfaces_conflicting_rankings() -> None: """The two metrics crown different winners; the view must show both.""" table = disagreement(_board(), ["mmd", "dpp"]) - assert table.loc[("p", "b", "cfg_b", "v1"), "mmd"] == 1 - assert table.loc[("p", "a", "cfg_a", "v1"), "dpp"] == 1 + assert table.loc[("p", "b", "cfg_b", "someone/engiopt-b", "v1"), "mmd"] == 1 + assert table.loc[("p", "a", "cfg_a", "someone/engiopt-a", "v1"), "dpp"] == 1 def _two_configs_board() -> pd.DataFrame: @@ -256,7 +264,7 @@ def test_append_rows_supersedes_a_rerun_instead_of_duplicating(tmp_path: Any) -> stored = pd.read_csv(destination) assert len(stored) == 3 - match = stored.set_index(ROW_KEY).loc[("p", "a", "cfg_a", 1, "v1")] + match = stored.set_index(ROW_KEY).loc[("p", "a", "cfg_a", "someone/engiopt-a", 1, "v1")] assert match["mmd"] == pytest.approx(0.999) diff --git a/tests/test_leaderboard_integrity.py b/tests/test_leaderboard_integrity.py new file mode 100644 index 00000000..8f649dee --- /dev/null +++ b/tests/test_leaderboard_integrity.py @@ -0,0 +1,558 @@ +"""Tests for the things a *public* leaderboard needs that an internal one does not. + +An internal board can assume its rows were produced in good faith by someone who +had no reason to fabricate them. A public board cannot assume either, and the +two failure modes are independent: + +- Someone reports numbers no model produced. Defended by re-execution, not by + schema -- see `engiopt.evaluation.verify`. +- Someone reports honest numbers from a model that games the metric rather than + solving the problem. Defended by measuring the gaming, since the scoring + protocol is public and cannot be un-published. + +The headline case for the second is a retrieval system. The spec names which +conditions are scored and the dataset supplies the optimal design for each, so +returning those designs tops `mmd`, `iog`, and `fog` by construction. The tests +below assert that such a model is *detected*, not that it scores badly: it +should score wonderfully and be visibly a copier. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pandas as pd +import pytest + +from engiopt.evaluation.context import EvaluationContext +from engiopt.evaluation.leaderboard import eligible_rows +from engiopt.evaluation.leaderboard import rank +from engiopt.evaluation.registry import METRICS +from engiopt.evaluation.spec import EvalSpec +from engiopt.evaluation.submission import admission_problems +from engiopt.evaluation.submission import FLAG_IGNORES_CONDITIONS +from engiopt.evaluation.submission import FLAG_MEMORIZED +from engiopt.evaluation.submission import FLAG_UNVERIFIED +from engiopt.evaluation.submission import integrity_flags +from engiopt.evaluation.submission import prepare_submission +from engiopt.evaluation.submission import SubmissionRejectedError + +# Importing the metrics package registers the built-ins. +import engiopt.evaluation.metrics # noqa: F401 # isort: skip + + +def _context(problem: Any, gen: np.ndarray, ref: np.ndarray, **kwargs: Any) -> EvaluationContext: + return EvaluationContext(problem=problem, problem_id="fake", gen_designs=gen, ref_designs=ref, **kwargs) + + +# ---------------------------------------------------------------------- +# Memorization: a retrieval system must be visible as one +# ---------------------------------------------------------------------- + + +def test_a_generator_returning_the_reference_designs_is_caught(fake_problem: Any) -> None: + """The attack this metric exists for. + + A lookup table keyed on the condition vector returns the dataset-optimal + design for each scored condition. That is the *definition* of a perfect + `mmd` and a near-zero `iog`, so no amount of care in those metrics can + distinguish it from a model that learned the problem. This one can. + """ + rng = np.random.default_rng(0) + ref = rng.random((10, *fake_problem.design_space.shape)) + ctx = _context(fake_problem, ref.copy(), ref) + + scores = METRICS["novelty"].fn(ctx) + + assert scores["copy_rate"] == 1.0 + assert scores["novelty"] == pytest.approx(0.0, abs=1e-12) + # And it does indeed post a perfect distribution score, which is the point. + assert METRICS["mmd"].fn(ctx) == pytest.approx(0.0, abs=1e-9) + + +def test_a_generator_producing_its_own_designs_is_not_flagged(fake_problem: Any) -> None: + """The check must not fire on a model that merely resembles the data.""" + rng = np.random.default_rng(1) + ref = rng.random((10, *fake_problem.design_space.shape)) + ctx = _context(fake_problem, rng.random((10, *fake_problem.design_space.shape)), ref) + + scores = METRICS["novelty"].fn(ctx) + + assert scores["copy_rate"] == 0.0 + assert scores["novelty"] > 0 + + +def test_copying_the_training_split_is_caught_too(fake_problem: Any) -> None: + """Not only the scored references are copyable -- the whole public dataset is. + + A model that memorized the training split and emits rows of it under + whatever conditions it is given never touches a reference design, so a check + that looked only at those would report it as perfectly novel. + """ + rng = np.random.default_rng(2) + ref = rng.random((6, *fake_problem.design_space.shape)) + train = rng.random((20, *fake_problem.design_space.shape)) + ctx = _context(fake_problem, train[:6].copy(), ref, copy_corpus_fn=lambda: train) + + assert METRICS["novelty"].fn(ctx)["copy_rate"] == 1.0 + + +def test_near_copies_count_as_copies(fake_problem: Any) -> None: + """Adding imperceptible noise to a retrieved design must not launder it. + + Otherwise the defence is defeated by one line, and the tolerance is what + decides how much perturbation counts as having generated something. + """ + rng = np.random.default_rng(3) + ref = rng.random((8, *fake_problem.design_space.shape)) + barely_perturbed = ref + rng.normal(scale=1e-4, size=ref.shape) + + scores = METRICS["novelty"].fn(_context(fake_problem, barely_perturbed, ref, copy_tol=0.01)) + + assert scores["copy_rate"] == 1.0 + + +def test_novelty_is_diagnostic_and_cannot_be_ranked_on() -> None: + """Ranking on novelty would put pure noise in first place. + + The metric has no good direction -- zero means retrieval, large means only + "unlike the data", which is what a broken model also achieves. Registering a + direction would have created a new thing to game while closing an old one. + """ + assert METRICS["novelty"].higher_is_better is None + with pytest.raises(ValueError, match="diagnostic"): + rank(pd.DataFrame([{"problem_id": "p", "algo_id": "a", "novelty": 0.5}]), "novelty") + + +# ---------------------------------------------------------------------- +# Conditioning: a model must be shown to read its brief +# ---------------------------------------------------------------------- + + +def test_a_model_ignoring_its_conditions_scores_zero_sensitivity(fake_problem: Any) -> None: + """Same output whatever the brief: the signature of conditions never reaching the network.""" + rng = np.random.default_rng(4) + designs = rng.random((8, *fake_problem.design_space.shape)) + ctx = _context(fake_problem, designs, designs.copy(), resample_permuted=lambda _order: designs) + + assert METRICS["cond_sens"].fn(ctx) == pytest.approx(0.0, abs=1e-12) + + +def test_a_model_reading_its_conditions_scores_above_zero(fake_problem: Any) -> None: + """A generator whose output follows the condition order must register a response.""" + rng = np.random.default_rng(5) + designs = rng.random((8, *fake_problem.design_space.shape)) + ctx = _context(fake_problem, designs, designs.copy(), resample_permuted=lambda order: designs[order]) + + assert METRICS["cond_sens"].fn(ctx) > 0 + + +def test_sensitivity_is_unavailable_rather_than_zero_without_a_comparison(fake_problem: Any) -> None: + """An unconditional *problem* must not be reported as a model ignoring conditions. + + Zero would be a claim about the model. NaN is the honest answer: there was + nothing to shuffle, so nothing was measured. + """ + rng = np.random.default_rng(6) + designs = rng.random((8, *fake_problem.design_space.shape)) + + assert np.isnan(METRICS["cond_sens"].fn(_context(fake_problem, designs, designs.copy()))) + + +def test_an_honestly_unconditional_model_is_not_flagged() -> None: + """Being unconditional is allowed by the contract; claiming otherwise is not. + + The flag fires on the *inconsistency* between what a model declares and what + it does, so a model registered as unconditional keeps a clean row. + """ + row = {"algo_id": "gan_cnn_2d", "cond_sens": 0.0, "verified": True} + + from engiopt.utils.all_generators import BUILTIN_GENERATORS + + assert not BUILTIN_GENERATORS["gan_cnn_2d"].conditional + assert FLAG_IGNORES_CONDITIONS not in integrity_flags(row) + + +def test_a_conditional_model_that_ignores_conditions_is_flagged() -> None: + """The inconsistent case: declares conditional, output does not move.""" + row = {"algo_id": "cgan_cnn_2d", "cond_sens": 0.0, "verified": True} + + from engiopt.utils.all_generators import BUILTIN_GENERATORS + + assert BUILTIN_GENERATORS["cgan_cnn_2d"].conditional + assert FLAG_IGNORES_CONDITIONS in integrity_flags(row) + + +# ---------------------------------------------------------------------- +# Admission: what a public board refuses outright +# ---------------------------------------------------------------------- + + +def _publishable(**overrides: Any) -> dict[str, Any]: + row = { + "problem_id": "beams2d", + "algo_id": "cgan_cnn_2d", + "config_fingerprint": "023dd1fb", + "seed": 1, + "spec_version": "v1", + "checkpoint_repo": "someone/engiopt-cgan-cnn-2d", + "checkpoint_path": "beams2d/seed_1", + "checkpoint_revision": "cafe1234", + "checkpoint_hash": "0011223344556677", + "mmd": 0.04, + } + row.update(overrides) + return row + + +def test_a_row_with_no_checkpoint_address_is_refused() -> None: + """Nobody can re-run it, so it is not a result -- it is a claim with no referent.""" + problems = admission_problems(_publishable(checkpoint_repo=None, checkpoint_path=None)) + + assert any("no fetchable checkpoint" in problem for problem in problems) + + +def test_a_row_evaluated_from_a_local_directory_is_refused() -> None: + """`--model-source local` leaves no repo or revision, which is the common accident. + + The message has to name the fix, because the submitter's evaluation + succeeded and they have no reason to suspect the checkpoint was the problem. + """ + problems = admission_problems(_publishable(checkpoint_repo=None, checkpoint_path=None, checkpoint_revision=None)) + + assert any("--checkpoint-backend hf" in problem for problem in problems) + + +def test_a_submitter_cannot_stamp_their_own_row_as_verified() -> None: + """Verification means somebody else re-ran it. Self-service defeats the entire mechanism.""" + problems = admission_problems(_publishable(verified=True)) + + assert any("cannot be self-asserted" in problem for problem in problems) + + +def test_preparing_a_submission_clears_any_verification_stamp() -> None: + """Even absent malice, a stale True from a re-published row must not survive.""" + prepared = prepare_submission(pd.DataFrame([_publishable()])) + + assert not prepared["verified"].iloc[0] + assert FLAG_UNVERIFIED in prepared["flags"].iloc[0] + + +def test_every_rejected_row_is_reported_at_once() -> None: + """Fixing a batch one error per round-trip is the difference between usable and not.""" + frame = pd.DataFrame([_publishable(checkpoint_repo=None), _publishable(seed=2, algo_id=None)]) + + with pytest.raises(SubmissionRejectedError) as excinfo: + prepare_submission(frame) + + assert len(excinfo.value.problems_by_row) == 2 + + +def test_a_memorizing_row_is_published_but_flagged() -> None: + """Published, because hiding it throws away what the board just learned. + + Flagged, because calling it first place would be the board endorsing the + thing it exists to detect. + """ + prepared = prepare_submission(pd.DataFrame([_publishable(copy_rate=0.9)]), EvalSpec(problem_id="beams2d")) + + assert len(prepared) == 1 + assert FLAG_MEMORIZED in prepared["flags"].iloc[0] + + +# ---------------------------------------------------------------------- +# Eligibility: what gets ranked, as distinct from what gets published +# ---------------------------------------------------------------------- + + +def _scored(**overrides: Any) -> dict[str, Any]: + row = _publishable(verified=True, flags="") + row.update(overrides) + return row + + +def test_unverified_rows_are_published_but_never_ranked() -> None: + """The whole trust model in one assertion.""" + frame = pd.DataFrame([_scored(verified=False, flags=FLAG_UNVERIFIED)]) + + assert eligible_rows(frame).empty + + +def test_flagged_rows_stay_on_the_board_and_out_of_the_ranking() -> None: + """A retrieval system belongs in the table and not in the ordering.""" + frame = pd.DataFrame([_scored(algo_id="honest"), _scored(algo_id="copier", flags=FLAG_MEMORIZED)]) + + eligible = eligible_rows(frame) + + assert list(eligible["algo_id"]) == ["honest"] + assert len(frame) == 2 + + +def test_an_entry_missing_a_required_seed_is_not_ranked() -> None: + """The cherry-picking defence. + + Twenty seeds run, the best three published, and a median over a maximum + looks exactly like an honest median in every individual row. Requiring + *named* seeds removes the choice; requiring a count would not. + """ + spec = EvalSpec(problem_id="beams2d", required_seeds=(1, 2, 3)) + frame = pd.DataFrame([_scored(seed=1), _scored(seed=2)]) + + assert eligible_rows(frame, spec).empty + + +def test_an_entry_covering_every_required_seed_is_ranked() -> None: + """And the rule must not lock out the people who followed it.""" + spec = EvalSpec(problem_id="beams2d", required_seeds=(1, 2, 3)) + frame = pd.DataFrame([_scored(seed=1), _scored(seed=2), _scored(seed=3)]) + + ranked = rank(frame, "mmd", eval_spec=spec) + + assert len(ranked) == 1 + assert ranked["n_seeds"].iloc[0] == 3 + + +def test_seed_coverage_is_judged_per_entry_not_across_the_board() -> None: + """One complete entry must not lend its seeds to an incomplete one.""" + spec = EvalSpec(problem_id="beams2d", required_seeds=(1, 2, 3)) + frame = pd.DataFrame( + [ + _scored(algo_id="complete", seed=1), + _scored(algo_id="complete", seed=2), + _scored(algo_id="complete", seed=3), + _scored(algo_id="partial", seed=1), + ] + ) + + assert set(eligible_rows(frame, spec)["algo_id"]) == {"complete"} + + +def test_ranking_can_be_asked_for_the_raw_table() -> None: + """A local run has nothing verified yet, and still wants to see its own ordering.""" + frame = pd.DataFrame([_scored(algo_id="a", verified=False, mmd=0.1), _scored(algo_id="b", verified=False, mmd=0.2)]) + + ranked = rank(frame, "mmd", eligible_only=False) + + assert list(ranked["algo_id"]) == ["a", "b"] + + +def test_ranking_an_empty_eligible_set_does_not_explode() -> None: + """A board where nothing is verified yet is the normal state on day one.""" + frame = pd.DataFrame([_scored(verified=False)]) + + assert rank(frame, "mmd").empty + + +# ---------------------------------------------------------------------- +# End to end: the evaluator has to actually feed the integrity metrics +# ---------------------------------------------------------------------- + + +def _evaluator(problem: Any, spec: EvalSpec, ref: np.ndarray, conditions: Any) -> Any: + """An `Evaluator` wired to a fake problem, bypassing spec resolution. + + `Evaluator.for_problem` would load a real EngiBench problem and draw a real + dataset. What is under test here is the wiring between the evaluator and the + integrity metrics -- the permuted re-sample and the copy corpus -- which is + the part that no direct-context test can reach. + """ + import torch as th + + from engiopt.evaluation.evaluator import Evaluator + from engiopt.evaluation.spec import ResolvedSpec + + resolved = ResolvedSpec( + spec=spec, + conditions_tensor=th.tensor([[float(i), 0.5] for i in range(len(ref))]), + conditions=conditions, + ref_designs=ref, + indices=np.arange(len(ref)), + condition_keys=("volfrac", "rmin"), + ) + return Evaluator(problem=problem, problem_id="fake", resolved=resolved, device=th.device("cpu")) + + +def _generator(problem: Any, produce: Any, *, algo_id: str = "demo", conditional: bool = True) -> Any: + """A minimal real `Generator` whose `_sample` is supplied by the caller.""" + from engiopt.core import Generator + + class _Demo(Generator): + design_kinds = ("2d",) + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + + @classmethod + def build(cls, resolved: Any, problem: Any, device: Any, **base: Any) -> Any: + raise NotImplementedError + + def _sample(self, conditions: Any, n: int) -> Any: + return produce(conditions, n) + + _Demo.algo_id = algo_id + _Demo.conditional = conditional + import torch as th + + return _Demo(problem=problem, problem_id="fake", seed=1, device=th.device("cpu")) + + +def test_the_evaluator_detects_a_lookup_table_end_to_end(fake_problem: Any) -> None: + """The full path: sample, build the corpus, score, and report `copy_rate`. + + This is the claim the whole memorization defence rests on, so it is asserted + through the evaluator rather than against a hand-built context -- a metric + that works but is never fed would pass every other test in this file. + """ + spec = EvalSpec(problem_id="fake", n_samples=4, metrics=("mmd", "novelty")) + ref = np.stack([np.full(fake_problem.design_space.shape, v) for v in (0.4, 0.5, 0.6, 0.7)]) + conditions = fake_problem.dataset["train"].select([0, 1, 2, 0]) + + evaluator = _evaluator(fake_problem, spec, ref, conditions) + lookup_table = _generator(fake_problem, lambda _conditions, _n: ref) + row = evaluator.score(lookup_table) + + assert row["copy_rate"] == 1.0 + assert row["mmd"] == pytest.approx(0.0, abs=1e-9) + assert FLAG_MEMORIZED in integrity_flags(row, spec) + + +def test_the_evaluator_leaves_an_honest_model_unflagged(fake_problem: Any) -> None: + """The same path, for a model that generates rather than retrieves.""" + spec = EvalSpec(problem_id="fake", n_samples=4, metrics=("mmd", "novelty", "cond_sens")) + ref = np.stack([np.full(fake_problem.design_space.shape, v) for v in (0.4, 0.5, 0.6, 0.7)]) + conditions = fake_problem.dataset["train"].select([0, 1, 2, 0]) + rng = np.random.default_rng(7) + + evaluator = _evaluator(fake_problem, spec, ref, conditions) + # Output tracks the first condition column, so it is genuinely conditional. + honest = _generator( + fake_problem, + lambda conditions, n: ( + np.stack([np.full(fake_problem.design_space.shape, 0.05 * float(c[0])) for c in conditions.tensor]) + + rng.normal(scale=0.05, size=(n, *fake_problem.design_space.shape)) + ), + ) + row = evaluator.score(honest) + + assert row["copy_rate"] == 0.0 + assert row["cond_sens"] > 0 + assert not [flag for flag in integrity_flags(row, spec) if flag != FLAG_UNVERIFIED] + + +def test_the_evaluator_catches_a_conditional_model_that_ignores_conditions(fake_problem: Any) -> None: + """Declares itself conditional, returns the identical batch for any brief.""" + spec = EvalSpec(problem_id="fake", n_samples=4, metrics=("cond_sens",)) + ref = np.stack([np.full(fake_problem.design_space.shape, v) for v in (0.4, 0.5, 0.6, 0.7)]) + conditions = fake_problem.dataset["train"].select([0, 1, 2, 0]) + fixed = np.full((4, *fake_problem.design_space.shape), 0.25) + + evaluator = _evaluator(fake_problem, spec, ref, conditions) + row = evaluator.score(_generator(fake_problem, lambda _conditions, _n: fixed, conditional=True)) + + assert row["cond_sens"] == pytest.approx(0.0, abs=1e-12) + + +def test_the_permuted_re_sample_holds_the_latent_draw_fixed(fake_problem: Any) -> None: + """The comparison is only meaningful if the *only* thing that changed is the brief. + + Both draws are seeded identically, so a model whose randomness is the same + and whose conditions differ isolates the conditional response. If the seed + were not replayed, every model would look condition-sensitive because every + model would have re-rolled its noise. + """ + import torch as th + + spec = EvalSpec(problem_id="fake", n_samples=4, metrics=("cond_sens",)) + ref = np.stack([np.full(fake_problem.design_space.shape, v) for v in (0.4, 0.5, 0.6, 0.7)]) + conditions = fake_problem.dataset["train"].select([0, 1, 2, 0]) + + def _pure_noise(_conditions: Any, n: int) -> Any: + return th.rand(n, *fake_problem.design_space.shape) + + evaluator = _evaluator(fake_problem, spec, ref, conditions) + row = evaluator.score(_generator(fake_problem, _pure_noise)) + + # Noise that ignores its conditions must read as zero response, not as a + # large one produced by re-rolling the random draw. + assert row["cond_sens"] == pytest.approx(0.0, abs=1e-12) + + +def test_the_copy_corpus_is_drawn_once_for_a_whole_sweep(fake_problem: Any) -> None: + """Every model in a sweep is checked against the same corpus, fetched once. + + Both halves matter: a per-model corpus would be slow, and a corpus that + varied between models would make their `copy_rate` values incomparable. + """ + spec = EvalSpec(problem_id="fake", n_samples=4, metrics=("novelty",)) + ref = np.stack([np.full(fake_problem.design_space.shape, v) for v in (0.4, 0.5, 0.6, 0.7)]) + conditions = fake_problem.dataset["train"].select([0, 1, 2, 0]) + evaluator = _evaluator(fake_problem, spec, ref, conditions) + + draws = 0 + original = evaluator._draw_copy_corpus + + def _counting_draw() -> Any: + nonlocal draws + draws += 1 + return original() + + evaluator._draw_copy_corpus = _counting_draw # type: ignore[method-assign] + for _ in range(3): + evaluator.score(_generator(fake_problem, lambda _c, n: np.zeros((n, *fake_problem.design_space.shape)))) + + assert draws == 1 + + +# ---------------------------------------------------------------------- +# Two people, one model name +# ---------------------------------------------------------------------- + + +def test_two_contributors_of_the_same_model_do_not_share_a_row() -> None: + """On a public board the model name plus config plus seed is not unique. + + Two people who both train `cgan_cnn_2d` with this repository's default + hyperparameters on seed 1 produce identical identity columns and genuinely + different weights. Merging them would let either one's push replace the + other's verified row, silently and with no malice required. + """ + from engiopt.evaluation.leaderboard import merge_rows + + mine = pd.DataFrame([_scored(checkpoint_repo="me/engiopt-cgan-cnn-2d", mmd=0.10)]) + theirs = pd.DataFrame([_scored(checkpoint_repo="you/engiopt-cgan-cnn-2d", mmd=0.20)]) + + merged = merge_rows(mine, theirs) + + assert len(merged) == 2 + assert set(merged["mmd"]) == {0.10, 0.20} + + +def test_re_evaluating_your_own_checkpoint_still_supersedes_your_row() -> None: + """The de-duplication that `ROW_KEY` exists for must keep working.""" + from engiopt.evaluation.leaderboard import merge_rows + + first = pd.DataFrame([_scored(mmd=0.10)]) + second = pd.DataFrame([_scored(mmd=0.20)]) + + merged = merge_rows(first, second) + + assert len(merged) == 1 + assert merged["mmd"].iloc[0] == 0.20 + + +def test_two_contributors_are_ranked_as_separate_entries() -> None: + """And their seeds must not be pooled into one median over a mixture.""" + spec = EvalSpec(problem_id="beams2d", required_seeds=(1,)) + frame = pd.DataFrame( + [ + _scored(checkpoint_repo="me/engiopt-cgan-cnn-2d", seed=1, mmd=0.10), + _scored(checkpoint_repo="you/engiopt-cgan-cnn-2d", seed=1, mmd=0.20), + ] + ) + + ranked = rank(frame, "mmd", eval_spec=spec) + + assert len(ranked) == 2 + assert list(ranked["rank"]) == [1, 2] + assert ranked["checkpoint_repo"].iloc[0] == "me/engiopt-cgan-cnn-2d" diff --git a/tests/test_leaderboard_publishing.py b/tests/test_leaderboard_publishing.py index 668d84a8..cf9324d1 100644 --- a/tests/test_leaderboard_publishing.py +++ b/tests/test_leaderboard_publishing.py @@ -41,12 +41,22 @@ def hub_error(error_cls: type[Exception], message: str, status: int = 500) -> Ex def _row(algo: str, mmd: float) -> dict[str, Any]: + """A publishable row: identity, a fetchable checkpoint address, and a score. + + The address is not decoration. `push_to_hub` refuses rows without one, + because a row nobody can re-run cannot be verified and so can never be + ranked -- see `engiopt.evaluation.submission`. + """ return { "problem_id": "p", "algo_id": algo, "config_fingerprint": "default", "seed": 1, "spec_version": "v1", + "checkpoint_repo": f"someone/engiopt-{algo}", + "checkpoint_path": "p/seed_1", + "checkpoint_revision": "cafe1234", + "checkpoint_hash": f"hash-{algo}", "mmd": mmd, } diff --git a/tests/test_verification.py b/tests/test_verification.py new file mode 100644 index 00000000..8bb21685 --- /dev/null +++ b/tests/test_verification.py @@ -0,0 +1,254 @@ +"""Tests for re-running a published row from its checkpoint address. + +Verification is the only defence a public board has against a number that no +model produced, so the properties here are about what it refuses to accept: +weights that have moved since they were scored, addresses that lead nowhere, and +models nobody can rebuild. The successful path matters less than those, because +a verifier that stamps rows it did not actually reproduce is worse than none -- +it launders the claim it was supposed to check. +""" + +from __future__ import annotations + +from typing import Any + +import pandas as pd +import pytest + +from engiopt.evaluation import verify as verify_mod +from engiopt.evaluation.spec import EvalSpec +from engiopt.evaluation.submission import FLAG_MEMORIZED +from engiopt.evaluation.verify import verify_board +from engiopt.evaluation.verify import verify_row + + +class _FakeGenerator: + """Stands in for a rebuilt model; only its checkpoint hash is consulted.""" + + algo_id = "cgan_cnn_2d" + conditional = True + checkpoint_files = ("generator.pth",) + + def __init__(self, content_hash: str = "0011223344556677"): + self.checkpoint_hash = content_hash + + +class _FakeEvaluator: + """Returns fixed scores, standing in for a real scoring pass.""" + + def __init__(self, scores: dict[str, Any] | None = None): + self.spec = EvalSpec(problem_id="beams2d") + self.scores = scores or {"mmd": 0.04, "copy_rate": 0.0} + self.calls = 0 + + def score(self, generator: Any, *, include_expensive: bool = False) -> dict[str, Any]: + del generator, include_expensive + self.calls += 1 + return dict(self.scores) + + +def _row(**overrides: Any) -> dict[str, Any]: + row = { + "problem_id": "beams2d", + "algo_id": "cgan_cnn_2d", + "config_fingerprint": "023dd1fb", + "seed": 1, + "spec_version": "v1", + "checkpoint_repo": "someone/engiopt-cgan-cnn-2d", + "checkpoint_path": "beams2d/seed_1", + "checkpoint_revision": "cafe1234", + "checkpoint_hash": "0011223344556677", + "mmd": 0.04, + "verified": False, + } + row.update(overrides) + return row + + +@pytest.fixture +def loads_fake_generator(monkeypatch: pytest.MonkeyPatch) -> None: + """Rebuild any row into a `_FakeGenerator`, so no network or weights are needed.""" + monkeypatch.setattr(verify_mod, "_load_at_recorded_revision", lambda row, cls, device=None: _FakeGenerator()) + + +# ---------------------------------------------------------------------- +# What verification refuses +# ---------------------------------------------------------------------- + + +def test_a_model_with_no_registered_adapter_cannot_be_verified() -> None: + """Scores for a model nobody can rebuild are unfalsifiable by construction.""" + result = verify_row(_row(algo_id="a_model_nobody_shipped"), verifier="ci") + + assert result.status == "unknown_generator" + assert not result.ok + assert result.row is None + + +def test_weights_that_moved_since_scoring_are_not_re_scored_into_place(monkeypatch: pytest.MonkeyPatch) -> None: + """The swap attack: score a good checkpoint, then overwrite the path. + + Re-scoring whatever is at the address now would quietly attach a fresh, + honest score to a row whose provenance is a lie. The row has to fail + instead, because what it claims to describe no longer exists. + """ + monkeypatch.setattr( + verify_mod, "_load_at_recorded_revision", lambda row, cls, device=None: _FakeGenerator("ffffffffffffffff") + ) + evaluator = _FakeEvaluator() + + result = verify_row(_row(), verifier="ci", evaluator=evaluator) + + assert result.status == "weights_moved" + assert evaluator.calls == 0, "a moved checkpoint must not be scored at all" + assert "no longer exist at that address" in result.detail + + +def test_an_unreachable_checkpoint_fails_rather_than_being_skipped(monkeypatch: pytest.MonkeyPatch) -> None: + """A row that cannot be fetched is a failed verification, not an absent one.""" + + def _explode(row: Any, cls: Any, device: Any = None) -> None: + raise FileNotFoundError("HF package path not found") + + monkeypatch.setattr(verify_mod, "_load_at_recorded_revision", _explode) + + result = verify_row(_row(), verifier="ci", evaluator=_FakeEvaluator()) + + assert result.status == "unresolvable" + assert "not found" in result.detail + + +def test_a_row_without_an_address_has_nothing_to_fetch() -> None: + """Admission should have caught this, but verification must not assume it did.""" + result = verify_row(_row(checkpoint_repo=None, checkpoint_path=None), verifier="ci") + + assert result.status == "unresolvable" + assert "no checkpoint address" in result.detail + + +# ---------------------------------------------------------------------- +# What verification publishes +# ---------------------------------------------------------------------- + + +def test_the_verified_row_carries_the_runners_numbers(loads_fake_generator: None) -> None: + """Not a pass/fail on the submitted ones. + + Sampling from one seed on different hardware genuinely produces different + designs, so a comparison would need a tolerance loose enough to admit real + fudging. Re-scoring sidesteps that: the runner's number is the number, and + the submitted one becomes a claim reported as corroborated or not. + """ + evaluator = _FakeEvaluator({"mmd": 0.42, "copy_rate": 0.0}) + + result = verify_row(_row(mmd=0.001), verifier="ideallab-ci", evaluator=evaluator) + + assert result.ok + assert result.row is not None + assert result.row["mmd"] == 0.42 + assert result.row["verified"] is True + assert result.row["verified_by"] == "ideallab-ci" + assert result.row["verified_at"] + + +def test_an_unreproduced_claim_is_reported_without_failing_the_row(loads_fake_generator: None) -> None: + """The row still verifies -- the weights are real and were scored. + + What is not true is the submitted number, and saying so is more useful than + refusing the row, because the honest explanation (different hardware) and + the dishonest one produce the same signal. + """ + result = verify_row(_row(mmd=0.001), verifier="ci", evaluator=_FakeEvaluator({"mmd": 0.42})) + + assert result.ok + assert "mmd" in result.uncorroborated + assert result.uncorroborated["mmd"] == (0.001, 0.42) + assert "NOT reproduced" in result.detail + + +def test_a_reproduced_claim_says_so(loads_fake_generator: None) -> None: + """Small differences are noise, not evidence, and must not read as an accusation.""" + result = verify_row(_row(mmd=0.0400), verifier="ci", evaluator=_FakeEvaluator({"mmd": 0.0401})) + + assert not result.uncorroborated + assert "corroborated" in result.detail + + +def test_verification_re_derives_the_flags_rather_than_trusting_them(loads_fake_generator: None) -> None: + """A submitter who cleared their own `memorized` flag must not keep it cleared.""" + evaluator = _FakeEvaluator({"mmd": 0.04, "copy_rate": 0.95}) + + result = verify_row(_row(flags=""), verifier="ci", evaluator=evaluator) + + assert result.row is not None + assert FLAG_MEMORIZED in result.row["flags"] + + +# ---------------------------------------------------------------------- +# Fetching the package the row actually names +# ---------------------------------------------------------------------- + + +def test_the_recorded_revision_is_what_gets_fetched(monkeypatch: pytest.MonkeyPatch) -> None: + """Not the current head. + + Pinning the revision is what stops a later upload to the same path from + laundering itself into an old row's score. Without it, "verified" would mean + "some checkpoint at this path scores well today". + """ + captured: dict[str, Any] = {} + + def _fake_resolve(**kwargs: Any) -> Any: + captured.update(kwargs) + raise FileNotFoundError("stop here; the call is what is under test") + + monkeypatch.setattr(verify_mod, "resolve_checkpoint_reference", _fake_resolve) + + verify_row(_row(checkpoint_revision="deadbeef"), verifier="ci", evaluator=_FakeEvaluator()) + + assert captured["revision"] == "deadbeef" + assert captured["model_ref"] == "hf://someone/engiopt-cgan-cnn-2d/beams2d/seed_1" + + +# ---------------------------------------------------------------------- +# Sweeping a whole board +# ---------------------------------------------------------------------- + + +def test_already_verified_rows_are_left_alone_by_default( + loads_fake_generator: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """Re-running everything on every pass would make the runner cost grow with the board.""" + monkeypatch.setattr(verify_mod.Evaluator, "for_problem", classmethod(lambda cls, *a, **k: _FakeEvaluator())) + board = pd.DataFrame([_row(seed=1, verified=True), _row(seed=2, verified=False)]) + + _, results = verify_board(board, verifier="ci") + + assert len(results) == 1 + assert results[0].key["seed"] == 2 + + +def test_one_unverifiable_row_does_not_stop_the_sweep(monkeypatch: pytest.MonkeyPatch) -> None: + """An unattended runner has to get through the board and report at the end.""" + monkeypatch.setattr(verify_mod.Evaluator, "for_problem", classmethod(lambda cls, *a, **k: _FakeEvaluator())) + monkeypatch.setattr(verify_mod, "_load_at_recorded_revision", lambda row, cls, device=None: _FakeGenerator()) + board = pd.DataFrame([_row(seed=1, algo_id="not_a_real_model"), _row(seed=2)]) + + verified, results = verify_board(board, verifier="ci") + + assert {result.status for result in results} == {"unknown_generator", "verified"} + assert len(verified) == 1 + + +def test_an_unloadable_spec_fails_its_rows_and_not_the_run(monkeypatch: pytest.MonkeyPatch) -> None: + """A row naming a spec version that was never committed is a bad row, not a crash.""" + + def _no_such_spec(cls: Any, *args: Any, **kwargs: Any) -> None: + raise FileNotFoundError("No eval spec at beams2d/v9.json") + + monkeypatch.setattr(verify_mod.Evaluator, "for_problem", classmethod(_no_such_spec)) + + _, results = verify_board(pd.DataFrame([_row(spec_version="v9")]), verifier="ci") + + assert results[0].status == "unresolvable" + assert "could not load spec" in results[0].detail From b87913f815e5f0235373363c819f88408440492c Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 11 Aug 2026 14:51:07 +0200 Subject: [PATCH 16/24] Stop a model falling silently out of a run, and format the docs as CI does Two things the live run turned up. `--generators gan_cnn_2d vqgan --config-fingerprints gan_cnn_2d:6293adb3` scored gan_cnn_2d and never mentioned vqgan. Scoping the only entry to another algorithm left vqgan with an empty fingerprint list, so the loop body never ran: no row, no load attempt, no error. A leaderboard quietly missing an entrant is worse than one reporting a load failure, so a model no entry is scoped to now falls back to its canonical checkpoint. Ruff 0.16 formats Python blocks inside Markdown and CI runs it repo-wide, while the local check here was scoped to engiopt/ and tests/. Reformatted, and the lesson is the general one: run the command CI runs, not a subset of it. Co-Authored-By: Claude Opus 5 (1M context) --- CONTRIBUTING_A_MODEL.md | 2 +- LEADERBOARD.md | 2 +- engiopt/evaluate.py | 9 ++++++++- tests/test_leaderboard_integrity.py | 23 +++++++++++++++++++++++ 4 files changed, 33 insertions(+), 3 deletions(-) diff --git a/CONTRIBUTING_A_MODEL.md b/CONTRIBUTING_A_MODEL.md index 35bf7254..35c6b7ce 100644 --- a/CONTRIBUTING_A_MODEL.md +++ b/CONTRIBUTING_A_MODEL.md @@ -152,7 +152,7 @@ fitted during training travels with the weights instead, and is replayed on load ```python save_checkpoint_package( ..., - condition_stats={"mean": [...], "std": [...]}, # if you rescaled conditions + condition_stats={"mean": [...], "std": [...]}, # if you rescaled conditions condition_normalizer=normalizer_state(cond_norm), # if you used a min/max Normalizer design_normalizer=normalizer_state(design_norm), ) diff --git a/LEADERBOARD.md b/LEADERBOARD.md index 612b4e7f..34ab5cb9 100644 --- a/LEADERBOARD.md +++ b/LEADERBOARD.md @@ -154,7 +154,7 @@ from engiopt.evaluation.spec import EvalSpec board = load_from_hub("IDEALLab/engiopt-leaderboard") spec = EvalSpec.load("beams2d/v1") -rank(board, "fog", eval_spec=spec) # the public ordering +rank(board, "fog", eval_spec=spec) # the public ordering rank(board, "fog", eval_spec=spec, eligible_only=False) # everything, including claims disagreement(board, ["mmd", "dpp", "fog"], eval_spec=spec) ``` diff --git a/engiopt/evaluate.py b/engiopt/evaluate.py index 88e8c0b4..9ef50144 100644 --- a/engiopt/evaluate.py +++ b/engiopt/evaluate.py @@ -202,6 +202,13 @@ def _fingerprints_for(requested: tuple[str, ...], algo: str) -> tuple[str | None scopes an entry to its owner; a bare fingerprint still applies to every algorithm, which is what a single-model run wants. + A model that no entry is scoped to falls back to its canonical checkpoint + rather than to nothing. Returning an empty tuple would drop it before the + load is even attempted -- no row, no error, no message -- so + `--generators gan_cnn_2d vqgan --config-fingerprints gan_cnn_2d:6293adb3` + would quietly score one model and never mention the other. A leaderboard + silently missing an entrant is worse than one reporting a load failure. + Args: requested: Raw `--config-fingerprints` values. algo: The algorithm being loaded. @@ -216,7 +223,7 @@ def _fingerprints_for(requested: tuple[str, ...], algo: str) -> tuple[str | None for entry in requested if ":" not in entry or entry.split(":", 1)[0] == algo ] - return tuple(selected) + return tuple(selected) or (None,) def _load_generators(args: Args, evaluator: Evaluator) -> list[Generator]: diff --git a/tests/test_leaderboard_integrity.py b/tests/test_leaderboard_integrity.py index 8f649dee..7765ccb3 100644 --- a/tests/test_leaderboard_integrity.py +++ b/tests/test_leaderboard_integrity.py @@ -556,3 +556,26 @@ def test_two_contributors_are_ranked_as_separate_entries() -> None: assert len(ranked) == 2 assert list(ranked["rank"]) == [1, 2] assert ranked["checkpoint_repo"].iloc[0] == "me/engiopt-cgan-cnn-2d" + + +# ---------------------------------------------------------------------- +# A model must never fall off a run without saying so +# ---------------------------------------------------------------------- + + +def test_a_model_no_fingerprint_is_scoped_to_falls_back_to_canonical() -> None: + """Otherwise it is dropped before the load is attempted -- silently. + + `--generators gan_cnn_2d vqgan --config-fingerprints gan_cnn_2d:6293adb3` + scopes the only entry to `gan_cnn_2d`, leaving `vqgan` with an empty + fingerprint list. An empty list means the loop body never runs, so vqgan + produces no row, no error, and no message. A leaderboard quietly missing an + entrant is worse than one reporting a load failure. + """ + from engiopt.evaluate import _fingerprints_for + + assert _fingerprints_for(("gan_cnn_2d:6293adb3",), "vqgan") == (None,) + # The cases that already worked must keep working. + assert _fingerprints_for(("gan_cnn_2d:6293adb3",), "gan_cnn_2d") == ("6293adb3",) + assert _fingerprints_for(("abc123",), "vqgan") == ("abc123",) + assert _fingerprints_for((), "vqgan") == (None,) From 8680e06bd4ea02fcf6f6058e0cb9c010cab7668a Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 11 Aug 2026 14:54:40 +0200 Subject: [PATCH 17/24] Say that the conditioning check costs a second sampling pass `cheap` in this registry means "cannot reach a simulator", which cond_sens honours -- but it draws the batch twice, so selecting it doubles generation cost. Free for a GAN, noticeable for diffusion, and unpleasant for a model that samples a pixel at a time. Documented alongside the --metrics escape hatch, which was already supported and unmentioned. Co-Authored-By: Claude Opus 5 (1M context) --- CONTRIBUTING_A_MODEL.md | 5 +++++ engiopt/evaluation/metrics/builtin.py | 6 ++++++ 2 files changed, 11 insertions(+) diff --git a/CONTRIBUTING_A_MODEL.md b/CONTRIBUTING_A_MODEL.md index 35c6b7ce..9469f49c 100644 --- a/CONTRIBUTING_A_MODEL.md +++ b/CONTRIBUTING_A_MODEL.md @@ -189,6 +189,11 @@ Watch two columns while you develop: training designs rather than generating, and the board will publish it without ranking it. +`cond_sens` draws the batch a second time, under shuffled conditions, so it +doubles sampling cost. That is nothing for a GAN and noticeable for a diffusion +model — pass `--metrics mmd dpp viol` while iterating if it slows you down, then +run the full list before publishing. + ## 5. Publish it See [LEADERBOARD.md](LEADERBOARD.md) for the full submission path. The short diff --git a/engiopt/evaluation/metrics/builtin.py b/engiopt/evaluation/metrics/builtin.py index a7c089b8..b3de299b 100644 --- a/engiopt/evaluation/metrics/builtin.py +++ b/engiopt/evaluation/metrics/builtin.py @@ -130,6 +130,12 @@ def cond_sens(ctx: EvaluationContext) -> float: Its job is to stop an unconditional model quietly collecting a conditional model's `mmd` score, not to be maximized. + `cheap` here means what it means everywhere in this registry: it cannot + reach a simulator or optimizer. It is not free, though -- it draws the batch + a second time, so selecting it doubles generation cost. Immaterial for a + GAN's tenth of a second; noticeable for a diffusion model, and worth knowing + before pointing it at one that samples a pixel at a time. + Returns NaN when there is nothing to compare: an unconditional problem, or fewer than two samples. """ From 6c3b8c97fff8eda5c348f2e92321c0721ee7f0e6 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 11 Aug 2026 16:23:02 +0200 Subject: [PATCH 18/24] Drop a convenience property nothing calls `ResolvedCheckpoint.reference` was added alongside repo_id/package_path on the assumption verification would want a formatted hf:// string; it builds that inline instead. An unused accessor on a core dataclass is a maintenance claim with no payer. Co-Authored-By: Claude Opus 5 (1M context) --- engiopt/checkpoint_store.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/engiopt/checkpoint_store.py b/engiopt/checkpoint_store.py index 5a397f58..6d952173 100644 --- a/engiopt/checkpoint_store.py +++ b/engiopt/checkpoint_store.py @@ -154,13 +154,6 @@ class ResolvedCheckpoint: package_path: str | None = None """Path of the package inside `repo_id`, e.g. `beams2d/cfg_023dd1fb/seed_1`.""" - @property - def reference(self) -> str | None: - """A `hf://entity/repo/path` reference that round-trips through `resolve_checkpoint_reference`.""" - if self.repo_id is None or self.package_path is None: - return None - return f"hf://{self.repo_id}/{self.package_path}" - def build_hf_repo_id(hf_entity: str, hf_repo_prefix: str, algo: str) -> str: """Return the canonical HF repo id for an EngiOpt model family.""" From 9329620ebdca6318625129402d0d52e76d422e55 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 11 Aug 2026 16:26:19 +0200 Subject: [PATCH 19/24] Delete an unused flag and de-duplicate the metric-cell reader `FLAG_INCOMPLETE_SEEDS` was declared and never raised: seed coverage is a property of the entry, so `eligible_rows` drops those rows during ranking rather than flagging them individually. A flag constant nobody emits reads like a feature that exists. `_as_float` had been written twice, once per module, for the same job with slightly different blank handling. Promoted to `as_metric_value` in submission, where the NaN rule belongs -- a metric that did not run must be neither flagged as a zero nor reported as a claim the verifier failed to reproduce. Co-Authored-By: Claude Opus 5 (1M context) --- engiopt/evaluation/submission.py | 19 +++++++++---------- engiopt/evaluation/verify.py | 13 ++----------- 2 files changed, 11 insertions(+), 21 deletions(-) diff --git a/engiopt/evaluation/submission.py b/engiopt/evaluation/submission.py index 32b5003d..4b6b05ca 100644 --- a/engiopt/evaluation/submission.py +++ b/engiopt/evaluation/submission.py @@ -44,9 +44,6 @@ FLAG_UNVERIFIED = "unverified" """No runner has re-fetched these weights and reproduced these numbers.""" -FLAG_INCOMPLETE_SEEDS = "incomplete_seeds" -"""The entry does not cover every seed the spec requires; see `EvalSpec.required_seeds`.""" - FLAG_IGNORES_CONDITIONS = "ignores_conditions" """Output did not move when the conditions did, on a problem where they vary.""" @@ -96,7 +93,7 @@ def integrity_flags(row: dict[str, Any], spec: EvalSpec | None = None) -> list[s ranking instead -- see `leaderboard.rank`. """ flags: list[str] = [] - copy_rate = _as_float(row.get("copy_rate")) + copy_rate = as_metric_value(row.get("copy_rate")) max_copy_rate = spec.max_copy_rate if spec is not None else 0.5 if copy_rate is not None and copy_rate > max_copy_rate: flags.append(FLAG_MEMORIZED) @@ -118,7 +115,7 @@ def _declares_conditioning_it_does_not_use(row: dict[str, Any]) -> bool: An honestly unconditional model is not flagged. The contract permits those, and they are handed conditions only so that this can be measured. """ - sensitivity = _as_float(row.get("cond_sens")) + sensitivity = as_metric_value(row.get("cond_sens")) if sensitivity is None or sensitivity > 0: return False from engiopt.utils.all_generators import BUILTIN_GENERATORS @@ -174,12 +171,14 @@ def __init__(self, problems_by_row: dict[int, list[str]]) -> None: super().__init__("\n".join(lines)) -def _as_float(value: Any) -> float | None: - """A cell as a usable number, or None when it carries no measurement. +def as_metric_value(value: Any) -> float | None: + """A leaderboard cell as a usable number, or None when it carries no measurement. - NaN counts as no measurement, which matters: `cond_sens` is NaN on an - unconditional problem, and treating that as zero would flag every model on - it for ignoring conditions that were never supplied. + NaN counts as no measurement, which matters in both places this is used: a + metric that did not run must not be flagged as a zero, and it must not be + reported as a claim the verifier failed to reproduce. `cond_sens` is NaN on + an unconditional problem, and reading that as 0.0 would accuse every model + on it of ignoring conditions it was never given. """ if _is_blank(value): return None diff --git a/engiopt/evaluation/verify.py b/engiopt/evaluation/verify.py index 0ea6c6e1..3ff0d296 100644 --- a/engiopt/evaluation/verify.py +++ b/engiopt/evaluation/verify.py @@ -27,7 +27,6 @@ from dataclasses import dataclass from dataclasses import field import datetime as dt -import math from typing import Any, Literal, TYPE_CHECKING import pandas as pd @@ -39,6 +38,7 @@ from engiopt.evaluation.evaluator import Evaluator from engiopt.evaluation.leaderboard import ROW_KEY from engiopt.evaluation.registry import METRICS +from engiopt.evaluation.submission import as_metric_value from engiopt.evaluation.submission import integrity_flags from engiopt.utils.all_generators import BUILTIN_GENERATORS @@ -223,7 +223,7 @@ def _uncorroborated(submitted: dict[str, Any], recomputed: dict[str, Any]) -> di for name in METRICS.columns(): if name not in submitted or name not in recomputed: continue - claimed, actual = _as_float(submitted[name]), _as_float(recomputed[name]) + claimed, actual = as_metric_value(submitted[name]), as_metric_value(recomputed[name]) if claimed is None or actual is None: continue scale = max(abs(actual), abs(claimed), 1e-12) @@ -242,15 +242,6 @@ def _corroboration_note(uncorroborated: dict[str, tuple[float, float]]) -> str: return f"re-scored; submitted values NOT reproduced ({reported})" -def _as_float(value: Any) -> float | None: - """Coerce a cell to a comparable float, or None when it is not one.""" - try: - number = float(value) - except (TypeError, ValueError): - return None - return None if math.isnan(number) else number - - def verify_board( board: pd.DataFrame, *, From af6b018b083516ba1a064ba96cb7bf1eb085f8c1 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 11 Aug 2026 17:26:38 +0200 Subject: [PATCH 20/24] Finish the viol wording, and cover freeze_spec Two gaps in my own response to review, found while writing the replies. The CLI's --include-expensive help still listed feasibility as simulator-backed. That was half of Soheyl's point and I had only fixed the other half in CONTRIBUTING; viol has been a constraint check since the earlier round. `freeze_spec` gained parameters for every contract field but no test, and he asked for one. Two now: that the values reach the frozen spec, and that the CLI's defaults are the dataclass's rather than a second copy that can drift without either side failing. Co-Authored-By: Claude Opus 5 (1M context) --- engiopt/evaluate.py | 8 ++++- tests/test_specs.py | 72 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 1 deletion(-) diff --git a/engiopt/evaluate.py b/engiopt/evaluate.py index 9ef50144..5d82c047 100644 --- a/engiopt/evaluate.py +++ b/engiopt/evaluate.py @@ -69,7 +69,13 @@ class Args: metrics: tuple[str, ...] = () """Metric names; defaults to the spec's list.""" include_expensive: bool = False - """Run simulator-backed metrics (COG/IOG/FOG, feasibility). Slow.""" + """Also run the optimality gaps (IOG/COG/FOG), which invoke the optimizer. Slow. + + Feasibility is *not* here. `viol` describes the design as generated, so it + is a constraint check rather than a solver run, and it still reports when + the optimizer refuses to start from an invalid design -- the case where the + answer matters most. The cheap pass is `mmd`, `dpp`, `viol`, `novelty` and + `cond_sens`.""" output_csv: str = "leaderboard_{problem_id}.csv" """Where to append results locally; may include `{problem_id}`.""" push_to: str | None = None diff --git a/tests/test_specs.py b/tests/test_specs.py index e2b7f30b..4279d10f 100644 --- a/tests/test_specs.py +++ b/tests/test_specs.py @@ -192,3 +192,75 @@ class _Problem: dataset_id="IDEALLab/beams_2d_50_100_v0", ) spec.check_problem_definition(_Problem()) + + +# ---------------------------------------------------------------------- +# freeze_spec must not silently reset fields it forgot to expose +# ---------------------------------------------------------------------- + + +class _FreezeProbe: + """A problem stand-in; freeze_spec only resets it and reads its conditions.""" + + conditions_keys: ClassVar[tuple[str, ...]] = ("volfrac",) + + def reset(self, seed: int | None = None) -> None: + """Match `Problem.reset`, which freeze_spec calls before sampling.""" + + +def test_freeze_spec_carries_every_contract_field(monkeypatch: pytest.MonkeyPatch, tmp_path: Any) -> None: + """A parameter this entry point omits is a field it resets to the dataclass default. + + `volume_condition` and `volfrac_tol` decide `viol`, and three of the four + committed specs set a non-null volume condition. Freezing a v2 through the + documented CLI without them would hand the new version a different + feasibility contract than the one it succeeds -- silently, since the spec + would still write and still load. + """ + from engibench.utils import all_problems + + from engiopt.evaluation import spec as spec_mod + + captured: dict[str, Any] = {} + + def _capture(self: EvalSpec, problem: Any) -> EvalSpec: + captured.update(dataclasses.asdict(self)) + return self + + monkeypatch.setattr(EvalSpec, "freeze", _capture) + monkeypatch.setattr(EvalSpec, "save", lambda self: tmp_path / "v2.json") + monkeypatch.setitem(all_problems.BUILTIN_PROBLEMS, "freeze_probe", _FreezeProbe) + + spec_mod.freeze_spec( + "freeze_probe", + version="v2", + volume_condition="volfrac", + volfrac_tol=0.05, + required_seeds=(1, 2, 3, 4), + copy_tol=0.02, + max_copy_rate=0.25, + copy_corpus_size=64, + ) + + assert captured["volume_condition"] == "volfrac" + assert captured["volfrac_tol"] == 0.05 + assert captured["required_seeds"] == (1, 2, 3, 4) + assert captured["copy_tol"] == 0.02 + assert captured["max_copy_rate"] == 0.25 + assert captured["copy_corpus_size"] == 64 + + +def test_freeze_spec_defaults_track_the_dataclass() -> None: + """The CLI's defaults must not become a second, drifting copy of the contract's. + + Two hand-maintained copies diverge invisibly: both sides still typecheck and + the spec still writes, but a freshly frozen version quietly disagrees with + the one it succeeds. + """ + import inspect + + from engiopt.evaluation import spec as spec_mod + + parameters = inspect.signature(spec_mod.freeze_spec).parameters + for field_name in ("metrics", "volfrac_tol", "required_seeds", "copy_tol", "max_copy_rate", "copy_corpus_size"): + assert parameters[field_name].default == getattr(EvalSpec, field_name), field_name From 963fc0261f852825abdd5f2c77af295e2a4e8652 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Wed, 2 Sep 2026 13:03:00 +0200 Subject: [PATCH 21/24] Separate a generator that cannot fit from one nobody trained `--list-generators --check-availability` queried every registered generator regardless of design kind, so on beams2d it reported ten missing checkpoints and told the reader "they can be trained and evaluated". Seven of those ten cannot serve a 2D problem at all -- `evaluator.py` refuses the pairing -- so the count invented seven migrations nobody owes and the sentence was false for each of them. The real number is three. `generators_for` already answered this and the listing path never called it. It does now, which also drops seven pointless Hub requests, and the third state prints as "does not fit beams2d" rather than being folded into the missing set. Records the decision behind the number: W&B-era checkpoints are not being migrated, because a migrated checkpoint arrives without the fingerprint and content hash that make a package canonically addressable and a row verifiable. Retraining costs less than a provenance story nobody can check. Also drops the README sentence promising two scripts per algorithm, which this branch stopped being true when it deleted the thirteen evaluation scripts. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 17 ++++++- docs/checkpoint_layout.md | 17 +++++++ engiopt/evaluate.py | 68 ++++++++++++++++---------- tests/test_evaluate_listing.py | 87 ++++++++++++++++++++++++++++++++++ 4 files changed, 164 insertions(+), 25 deletions(-) create mode 100644 tests/test_evaluate_listing.py diff --git a/README.md b/README.md index b93e9275..2d6f8484 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,21 @@ As much as we can, we follow the [CleanRL](https://github.com/vwxyzjn/cleanrl) p [vqgan](engiopt/generators/vqgan) | Inverse Design | 2D | ✅ | VQVAE + Transformer [pixel_cnn_pp_2d](engiopt/generators/pixel_cnn_pp_2d) | Inverse Design | 2D | ✅ | PixelCNN++ Autoregressive Model +Every algorithm above is registered, meaning this repository holds an adapter that +can rebuild it. Being *evaluable* additionally needs published weights, and the two +are not the same: `cgan_cnn_2d`, `diffusion_2d_cond`, `gan_cnn_2d`, and `vqgan` have +checkpoints on HuggingFace, and the rest must be trained before they can be scored. +The `Dimensions` column is the other half of the answer — a 1D or 3D generator has no +`beams2d` checkpoint because it cannot serve a 2D problem at all, not because one is +missing. For the live answer rather than this snapshot: + +``` +python -m engiopt.evaluate --problem-id beams2d --list-generators --check-availability +``` + +Historical W&B-era checkpoints are **not** being migrated; see +[docs/checkpoint_layout.md](docs/checkpoint_layout.md#historical-checkpoints). + ## Dashboards HuggingFace hosts everything that has to be reloaded or compared -- model weights, run configs, evaluation metrics, and the leaderboard. WandB hosts what you only look at: loss curves and sample images. Nothing in the evaluation path requires WandB, so training with `--track false` produces exactly the same checkpoints and scores. You can access some of our runs at https://wandb.ai/engibench/engiopt. WandB dashboards @@ -77,7 +92,7 @@ huggingface-cli login ``` ### Inverse design -Usually, we provide two scripts per algorithm: one to train the model, and one to evaluate it. +Each generator provides its own training script and adapter; evaluation is handled through the shared `python -m engiopt.evaluate` command. To train a model, you can run (for example): diff --git a/docs/checkpoint_layout.md b/docs/checkpoint_layout.md index 37b30db2..5458538f 100644 --- a/docs/checkpoint_layout.md +++ b/docs/checkpoint_layout.md @@ -37,6 +37,23 @@ Each package contains the model files plus: | `metadata.json` | problem, algo, seed, fingerprint, primary files, W&B run URL | | `metrics.json` | evaluation scores, when `--attach-metrics` was used | +## Historical checkpoints + +Checkpoints trained before this layout existed lived in W&B artifacts, and they +are **not being migrated**. W&B is no longer a checkpoint source, so those +weights are reachable only by retraining. + +The decision is deliberate rather than deferred. A migrated checkpoint would +arrive without the two things this layout exists to provide -- a config +fingerprint identifying the hyperparameters that produced it, and a content hash +tying a leaderboard row to exact bytes -- so it could be stored but never +canonically addressed or verified. Retraining a model costs less than a +provenance story nobody can check. + +What this means in practice: a generator with no published package is trained, +not recovered. `--list-generators --check-availability` reports which those are +for a given problem, and counts only generators whose design kind can serve it. + ## Writing one Training scripts call `save_checkpoint_package` and splat diff --git a/engiopt/evaluate.py b/engiopt/evaluate.py index 5d82c047..bda57745 100644 --- a/engiopt/evaluate.py +++ b/engiopt/evaluate.py @@ -121,40 +121,55 @@ class Args: def _print_generators(args: Args) -> None: """Print every registered generator, and say which ones can actually be loaded. - Registered and available are different things, and the gap is wide enough to - mislead: an adapter is code in this repository, while a usable model also - needs published weights. Listing all fourteen as though they were - interchangeable sends people to `--generators all` and a wall of load - failures that look like bugs. + Three states, not two, and collapsing the last two misleads in opposite + directions. An adapter is code in this repository; a model that *fits the + problem* additionally has a matching design kind; a model you can evaluate + additionally has published weights. Reporting a 1D generator as missing a + `beams2d` checkpoint invents a migration nobody owes, and reporting all + fourteen as interchangeable sends people to `--generators all` and a wall of + load failures that look like bugs. """ - print(f"{len(BUILTIN_GENERATORS)} generators registered:\n") - published = _published_packages(args) + fitting = _fitting_generator_names(args.problem_id) + print(f"{len(BUILTIN_GENERATORS)} generators registered, {len(fitting)} of them fit {args.problem_id}:\n") + published = _published_packages(args, fitting) for name, generator in sorted(BUILTIN_GENERATORS.items()): kinds = "/".join(generator.design_kinds) conditioning = "conditional" if generator.conditional else "unconditional" - availability = "" if published is None else f" {_availability_label(published.get(name))}" - print(f" {name:<20} {kinds:<10} {conditioning:<15}{availability}") - if published is not None: - missing = sorted(name for name, count in published.items() if not count) - if missing: - print( - f"\n{len(missing)} generator(s) have no published {args.problem_id} checkpoint under " - f"{args.hf_entity}: {', '.join(missing)}.\n" - "They can be trained and evaluated, but `--generators all` will report them as load " - "failures until someone publishes weights. W&B is not a checkpoint source." - ) - else: + status = _status(args.problem_id, name, fitting, published) + print(f" {name:<20} {kinds:<10} {conditioning:<15} {status}".rstrip()) + if published is None: print( f"\nPass --check-availability to also query {args.hf_entity} for which of these have " f"published {args.problem_id} checkpoints." ) + return + missing = sorted(name for name in fitting if not published.get(name)) + if missing: + print( + f"\n{len(missing)} of the {len(fitting)} generators that fit {args.problem_id} have no published " + f"checkpoint under {args.hf_entity}: {', '.join(missing)}.\n" + "They can be trained and evaluated, but `--generators all` will report them as load " + "failures until someone publishes weights. W&B is not a checkpoint source." + ) -def _published_packages(args: Args) -> dict[str, int] | None: - """How many packages each generator has for this problem, or None if not asked. +def _status(problem_id: str, name: str, fitting: set[str], published: dict[str, int] | None) -> str: + """One column saying why a generator is or is not usable on this problem.""" + if name not in fitting: + return f"does not fit {problem_id}" + if published is None: + return "" + return _availability_label(published.get(name)) + + +def _published_packages(args: Args, fitting: set[str]) -> dict[str, int] | None: + """How many packages each fitting generator has, or None if not asked. Behind a flag because it is one Hub request per generator, and `-h`-adjacent - commands should not depend on the network. + commands should not depend on the network. Generators that cannot serve the + problem are not queried: their repository holds nothing for it by + construction, so the answer would be zero for a reason unrelated to + publishing. """ if not args.check_availability: return None @@ -162,7 +177,7 @@ def _published_packages(args: Args) -> dict[str, int] | None: from engiopt.checkpoint_store import list_packages counts: dict[str, int] = {} - for name in BUILTIN_GENERATORS: + for name in sorted(fitting): repo = build_hf_repo_id(args.hf_entity, args.hf_repo_prefix, name) try: counts[name] = len(list_packages(repo, args.problem_id)) @@ -193,9 +208,14 @@ def _resolve_generator_names(requested: tuple[str, ...], problem_id: str) -> lis """Expand `all` to every registered generator compatible with the problem.""" if "all" not in requested: return list(requested) + return sorted(_fitting_generator_names(problem_id)) + + +def _fitting_generator_names(problem_id: str) -> set[str]: + """Names of the generators whose design kind matches this problem's.""" from engibench.utils.all_problems import BUILTIN_PROBLEMS - return sorted(generators_for(BUILTIN_PROBLEMS[problem_id]())) + return set(generators_for(BUILTIN_PROBLEMS[problem_id]())) def _fingerprints_for(requested: tuple[str, ...], algo: str) -> tuple[str | None, ...]: diff --git a/tests/test_evaluate_listing.py b/tests/test_evaluate_listing.py new file mode 100644 index 00000000..21826de0 --- /dev/null +++ b/tests/test_evaluate_listing.py @@ -0,0 +1,87 @@ +"""Tests for what `--list-generators` claims about a problem. + +Registered, fitting, and available are three states, and the listing is the only +place a user sees the difference. Collapsing the last two invents migrations that +nobody owes: a 1D generator has no `beams2d` checkpoint because it cannot serve a +2D problem, not because someone failed to publish one. +""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from engiopt import evaluate +from engiopt.utils.all_generators import BUILTIN_GENERATORS + +PUBLISHED = {"cgan_cnn_2d", "diffusion_2d_cond", "gan_cnn_2d", "vqgan"} + + +@pytest.fixture +def fake_hub(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Answer package queries from a fixed set, and record who was asked.""" + from engiopt import checkpoint_store + + asked: list[str] = [] + + def _list_packages(repo_id: str, problem_id: str | None = None) -> list[str]: + algo = repo_id.split("engiopt-", 1)[1].replace("-", "_") + asked.append(algo) + return ["beams2d/seed_1"] if algo in PUBLISHED else [] + + monkeypatch.setattr(checkpoint_store, "list_packages", _list_packages) + return asked + + +def _listing(**overrides: object) -> evaluate.Args: + return dataclasses.replace(evaluate.Args(), list_generators=True, **overrides) # type: ignore[arg-type] + + +def test_only_generators_that_fit_the_problem_are_queried(fake_hub: list[str], capsys: pytest.CaptureFixture) -> None: + """A repository holds nothing for a problem its model cannot serve, so do not ask.""" + evaluate._print_generators(_listing(check_availability=True)) + capsys.readouterr() + + fitting = evaluate._fitting_generator_names("beams2d") + assert set(fake_hub) == fitting + assert fitting < set(BUILTIN_GENERATORS), "beams2d must not fit every registered generator" + + +def test_an_incompatible_generator_is_not_reported_as_missing( + fake_hub: list[str], + capsys: pytest.CaptureFixture, +) -> None: + """1D and 3D models must read as inapplicable, never as unpublished.""" + evaluate._print_generators(_listing(check_availability=True)) + out = capsys.readouterr().out + + assert "cgan_1d" in out, "every registered generator is still listed" + incompatible = set(BUILTIN_GENERATORS) - evaluate._fitting_generator_names("beams2d") + for name in incompatible: + row = next(line for line in out.splitlines() if line.strip().startswith(name)) + assert "does not fit beams2d" in row + assert "no published checkpoints" not in row + + +def test_the_missing_count_covers_only_generators_that_fit( + fake_hub: list[str], + capsys: pytest.CaptureFixture, +) -> None: + """The number a reader takes away is the number of migrations actually owed.""" + evaluate._print_generators(_listing(check_availability=True)) + out = capsys.readouterr().out + + fitting = evaluate._fitting_generator_names("beams2d") + owed = len(fitting - PUBLISHED) + assert f"{owed} of the {len(fitting)} generators that fit beams2d have no published" in out + assert owed < len(set(BUILTIN_GENERATORS) - PUBLISHED), "the whole point is that this is smaller" + + +def test_without_the_flag_nothing_touches_the_network(fake_hub: list[str], capsys: pytest.CaptureFixture) -> None: + """`--list-generators` alone must stay usable offline.""" + evaluate._print_generators(_listing()) + out = capsys.readouterr().out + + assert fake_hub == [] + assert "--check-availability" in out From b087001ff0c0e40d5a42d28783556be987714d16 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Wed, 2 Sep 2026 13:03:16 +0200 Subject: [PATCH 22/24] Load checkpoints with the restricted unpickler A `.pth` file is pickled Python, and unpickling executes instructions carried in the file, so `torch.load` on a checkpoint from someone else's account runs their code. The recorded revision and content hash say which bytes were loaded, not that loading them was safe -- and this branch is what makes a remote package the only way to reach a checkpoint at all. All sixteen load sites now pass `weights_only=True`. Thirteen simply never set it; VQGAN opted into the unrestricted form at three sites and did not need to. Verified against the Hub rather than assumed: all nine weight files across the four published packages -- cgan_cnn_2d, diffusion_2d_cond, gan_cnn_2d, vqgan -- load unchanged under the restriction, holding nothing but state dicts, tensors and scalars. A test pins this across `core.py` and every registered adapter, so a new model cannot quietly reintroduce it. This is not enough for unattended verification and LEADERBOARD.md now says so instead of implying otherwise: safetensors, validated JSON for reconstruction and preprocessing state, header checks before allocation, and an isolated runner are all still missing. They are deliberately not here, because nothing outside IDEALLab can submit yet -- so there is no package that would need converting, and a submission format built before the runner that consumes it is a guess. Corrects the claim that submission needs no IDEALLab access, which was true of the checkpoints and false of `--push-to`. Co-Authored-By: Claude Opus 5 (1M context) --- LEADERBOARD.md | 41 ++++++++++++++++++- engiopt/core.py | 4 +- engiopt/generators/_template/adapter.py | 2 +- engiopt/generators/cgan_1d/adapter.py | 4 +- engiopt/generators/cgan_2d/adapter.py | 4 +- engiopt/generators/cgan_bezier/adapter.py | 4 +- engiopt/generators/cgan_cnn_2d/adapter.py | 2 +- engiopt/generators/cgan_cnn_3d/adapter.py | 4 +- engiopt/generators/cgan_vae/adapter.py | 2 +- engiopt/generators/diffusion_1d/adapter.py | 4 +- .../generators/diffusion_2d_cond/adapter.py | 2 +- engiopt/generators/gan_1d/adapter.py | 4 +- engiopt/generators/gan_2d/adapter.py | 4 +- engiopt/generators/gan_bezier/adapter.py | 4 +- engiopt/generators/gan_cnn_2d/adapter.py | 2 +- engiopt/generators/pixel_cnn_pp_2d/adapter.py | 4 +- engiopt/generators/vqgan/adapter.py | 6 +-- tests/test_generator_registry.py | 19 +++++++++ 18 files changed, 96 insertions(+), 20 deletions(-) diff --git a/LEADERBOARD.md b/LEADERBOARD.md index 34ab5cb9..f6929fb9 100644 --- a/LEADERBOARD.md +++ b/LEADERBOARD.md @@ -13,8 +13,8 @@ solving the problem. Neither defence helps with the other. ## Submitting -Nothing here needs write access to `IDEALLab`. Publish your checkpoints to your -own HuggingFace account: +Publish your checkpoints to your own HuggingFace account -- nothing about the +weights needs `IDEALLab`: ```bash python engiopt/generators/my_model/my_model.py \ @@ -30,6 +30,14 @@ python -m engiopt.evaluate \ Your rows land immediately, marked `verified=false`. They are visible on the board and absent from the ranking until a runner re-scores them. +**The last step needs write access to the board repository.** `--push-to` uploads +the merged CSV directly, so submission is currently open to people who already +have it. That is a real limit on how public this board is, and it is the one +piece of the submission path that is not yet self-service; a no-access route +(a pending manifest opened as a HuggingFace Community PR) is tracked alongside +the safe-loading work below. Everything else above -- the checkpoints, the +address a row records, the audit command -- already works from any account. + Three things are worth knowing before you run it. **Your checkpoints must be on the Hub.** Every row records the repo, path, @@ -81,6 +89,35 @@ The runner is not privileged. It is the first auditor, and because every row carries a full address, the same command run by anyone else produces the same answer — which is what keeps the runner honest too. +### Loading someone else's weights + +A `.pth` file is not data. It is pickled Python, and unpickling executes +instructions carried inside the file, so `torch.load` on a stranger's checkpoint +runs a stranger's code. A recorded revision and content hash say *which* bytes +were loaded; they say nothing about whether loading them was safe. + +Every checkpoint load in this repository therefore passes `weights_only=True`, +which swaps in an unpickler that can rebuild tensors and refuses arbitrary +imports and calls. All nine weight files across the four published packages load +under it unchanged, so this costs nothing today and closes the obvious hole. + +**It is not sufficient for unattended verification, and this document should not +be read as claiming otherwise.** A restricted unpickler narrows the attack +surface; it does not make an arbitrary tensor file safe to feed to an adapter, +and it does nothing about resource exhaustion. Before a runner re-scores +submissions from the public without a human in the loop, the submission format +needs to carry weights in a format with no code path at all (safetensors), +validated JSON for reconstruction settings and preprocessing state, and +adapter-side checks on tensor names, shapes and dtypes read from the header +before any allocation — and the runner itself needs isolation, no credentials, +and time and memory limits. + +None of that is built here, because nothing outside `IDEALLab` can submit yet: +`--push-to` still requires write access to the board. The two are the same piece +of work and should land together, since building a submission format before the +thing that consumes it means guessing at what it needs. Until then, the runner +is run by hand against checkpoints whose origin is known. + ## What gets ranked A row is published if it can be checked, and ranked if checking it went well. diff --git a/engiopt/core.py b/engiopt/core.py index 68f7c0c7..9da38e1b 100644 --- a/engiopt/core.py +++ b/engiopt/core.py @@ -349,7 +349,9 @@ def build( config = resolved.run_config net = MyNet(latent_dim=config["latent_dim"], design_shape=problem.design_space.shape).to(device) - net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.load_state_dict( + th.load(resolved.files["generator.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) net.eval() return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) diff --git a/engiopt/generators/_template/adapter.py b/engiopt/generators/_template/adapter.py index 24085703..c6c706b4 100644 --- a/engiopt/generators/_template/adapter.py +++ b/engiopt/generators/_template/adapter.py @@ -78,7 +78,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device # n_conds=len(condition_keys_for(problem, resolved)), # design_shape=design_shape_of(problem), # ).to(device) - # state = th.load(resolved.files["generator.pth"], map_location=device) + # state = th.load(resolved.files["generator.pth"], map_location=device, weights_only=True) # net.load_state_dict(state[cls.primary_state_key]) # net.eval() net = cls._build_network(config, problem, device, resolved) diff --git a/engiopt/generators/cgan_1d/adapter.py b/engiopt/generators/cgan_1d/adapter.py index a106c776..1922da3f 100644 --- a/engiopt/generators/cgan_1d/adapter.py +++ b/engiopt/generators/cgan_1d/adapter.py @@ -56,7 +56,9 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device design_normalizer=design_normalizer, conds_normalizer=conds_normalizer, ).to(device) - net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.load_state_dict( + th.load(resolved.files["generator.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) net.eval() return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) diff --git a/engiopt/generators/cgan_2d/adapter.py b/engiopt/generators/cgan_2d/adapter.py index e3b8fb37..58749fa4 100644 --- a/engiopt/generators/cgan_2d/adapter.py +++ b/engiopt/generators/cgan_2d/adapter.py @@ -42,7 +42,9 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device design_shape=problem.design_space.shape, generator_output_activation=config.get("generator_output_activation", "tanh"), ).to(device) - net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.load_state_dict( + th.load(resolved.files["generator.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) net.eval() return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) diff --git a/engiopt/generators/cgan_bezier/adapter.py b/engiopt/generators/cgan_bezier/adapter.py index 2fb1d839..4bdf6e04 100644 --- a/engiopt/generators/cgan_bezier/adapter.py +++ b/engiopt/generators/cgan_bezier/adapter.py @@ -65,7 +65,9 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device eps=_EPS, scalar_features=1, ).to(device) - net.load_state_dict(th.load(resolved.files["bezier_generator.pth"], map_location=device)[cls.primary_state_key]) + net.load_state_dict( + th.load(resolved.files["bezier_generator.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) net.eval() return cls( net=net, latent_dim=config["latent_dim"], noise_dim=config["noise_dim"], problem=problem, device=device, **base diff --git a/engiopt/generators/cgan_cnn_2d/adapter.py b/engiopt/generators/cgan_cnn_2d/adapter.py index 5120de20..3f3231e4 100644 --- a/engiopt/generators/cgan_cnn_2d/adapter.py +++ b/engiopt/generators/cgan_cnn_2d/adapter.py @@ -43,7 +43,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device design_shape=problem.design_space.shape, generator_output_activation=config.get("generator_output_activation", "tanh"), ) - checkpoint = th.load(resolved.files["generator.pth"], map_location=device) + checkpoint = th.load(resolved.files["generator.pth"], map_location=device, weights_only=True) net.load_state_dict(checkpoint[cls.primary_state_key]) net.eval().to(device) return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) diff --git a/engiopt/generators/cgan_cnn_3d/adapter.py b/engiopt/generators/cgan_cnn_3d/adapter.py index 1cf3ed90..9417f26e 100644 --- a/engiopt/generators/cgan_cnn_3d/adapter.py +++ b/engiopt/generators/cgan_cnn_3d/adapter.py @@ -55,7 +55,9 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device n_conds=len(condition_keys_for(problem, resolved)), design_shape=problem.design_space.shape, ) - net.load_state_dict(th.load(resolved.files["generator_3d.pth"], map_location=device)[cls.primary_state_key]) + net.load_state_dict( + th.load(resolved.files["generator_3d.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) net.eval().to(device) return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) diff --git a/engiopt/generators/cgan_vae/adapter.py b/engiopt/generators/cgan_vae/adapter.py index 8323e900..7290fd24 100644 --- a/engiopt/generators/cgan_vae/adapter.py +++ b/engiopt/generators/cgan_vae/adapter.py @@ -42,7 +42,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device n_conds=len(condition_keys_for(problem, resolved)), design_shape=problem.design_space.shape, ) - checkpoint = th.load(resolved.files["multiview_3d_vaegan.pth"], map_location=device) + checkpoint = th.load(resolved.files["multiview_3d_vaegan.pth"], map_location=device, weights_only=True) net.load_state_dict(checkpoint[cls.primary_state_key]) net.eval().to(device) return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) diff --git a/engiopt/generators/diffusion_1d/adapter.py b/engiopt/generators/diffusion_1d/adapter.py index d6eed741..2770c003 100644 --- a/engiopt/generators/diffusion_1d/adapter.py +++ b/engiopt/generators/diffusion_1d/adapter.py @@ -60,7 +60,9 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device seq_length=int(np.prod(padded_shape)), auto_normalize=config.get("auto_norm", True), ).to(device) - diffusion.load_state_dict(th.load(resolved.files["model.pth"], map_location=device)[cls.primary_state_key]) + diffusion.load_state_dict( + th.load(resolved.files["model.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) diffusion.eval() return cls( net=diffusion, diff --git a/engiopt/generators/diffusion_2d_cond/adapter.py b/engiopt/generators/diffusion_2d_cond/adapter.py index ac005c0e..b4d6b88b 100644 --- a/engiopt/generators/diffusion_2d_cond/adapter.py +++ b/engiopt/generators/diffusion_2d_cond/adapter.py @@ -52,7 +52,7 @@ def __init__( def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device, **base: Any) -> Diffusion2DCond: """Rebuild the UNet and its noise schedule from the checkpoint package.""" config = resolved.run_config - checkpoint = th.load(resolved.files["model.pth"], map_location=device) + checkpoint = th.load(resolved.files["model.pth"], map_location=device, weights_only=True) net = UNet2DConditionModel( sample_size=problem.design_space.shape, in_channels=1, diff --git a/engiopt/generators/gan_1d/adapter.py b/engiopt/generators/gan_1d/adapter.py index 8aebf264..7e5f6535 100644 --- a/engiopt/generators/gan_1d/adapter.py +++ b/engiopt/generators/gan_1d/adapter.py @@ -45,7 +45,9 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device design_shape=design_shape_of(problem), design_normalizer=design_normalizer, ).to(device) - net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.load_state_dict( + th.load(resolved.files["generator.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) net.eval() return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) diff --git a/engiopt/generators/gan_2d/adapter.py b/engiopt/generators/gan_2d/adapter.py index e791e40a..edc89f65 100644 --- a/engiopt/generators/gan_2d/adapter.py +++ b/engiopt/generators/gan_2d/adapter.py @@ -36,7 +36,9 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device """Load a trained GAN-2D generator from its checkpoint package.""" config = resolved.run_config net = GAN2DNet(latent_dim=config["latent_dim"], design_shape=problem.design_space.shape).to(device) - net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + net.load_state_dict( + th.load(resolved.files["generator.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) net.eval() return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) diff --git a/engiopt/generators/gan_bezier/adapter.py b/engiopt/generators/gan_bezier/adapter.py index 062d7ae2..010ebefc 100644 --- a/engiopt/generators/gan_bezier/adapter.py +++ b/engiopt/generators/gan_bezier/adapter.py @@ -66,7 +66,9 @@ def build( eps=_EPS, scalar_features=1, ).to(device) - net.load_state_dict(th.load(resolved.files["bezier_generator.pth"], map_location=device)[cls.primary_state_key]) + net.load_state_dict( + th.load(resolved.files["bezier_generator.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) net.eval() return cls( net=net, latent_dim=config["latent_dim"], noise_dim=config["noise_dim"], problem=problem, device=device, **base diff --git a/engiopt/generators/gan_cnn_2d/adapter.py b/engiopt/generators/gan_cnn_2d/adapter.py index fe284e18..889e40ba 100644 --- a/engiopt/generators/gan_cnn_2d/adapter.py +++ b/engiopt/generators/gan_cnn_2d/adapter.py @@ -41,7 +41,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device """Load a trained GAN-CNN generator from its checkpoint package.""" config = resolved.run_config net = GANCNN2DNet(latent_dim=config["latent_dim"], design_shape=problem.design_space.shape) - checkpoint = th.load(resolved.files["generator.pth"], map_location=device) + checkpoint = th.load(resolved.files["generator.pth"], map_location=device, weights_only=True) net.load_state_dict(checkpoint[cls.primary_state_key]) net.eval().to(device) return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) diff --git a/engiopt/generators/pixel_cnn_pp_2d/adapter.py b/engiopt/generators/pixel_cnn_pp_2d/adapter.py index 16b4fa22..6b969d8c 100644 --- a/engiopt/generators/pixel_cnn_pp_2d/adapter.py +++ b/engiopt/generators/pixel_cnn_pp_2d/adapter.py @@ -55,7 +55,9 @@ def build( input_channels=1, nr_conditions=len(condition_keys_for(problem, resolved)), ) - net.load_state_dict(th.load(resolved.files["model.pth"], map_location=device)[cls.primary_state_key]) + net.load_state_dict( + th.load(resolved.files["model.pth"], map_location=device, weights_only=True)[cls.primary_state_key] + ) net.eval().to(device) return cls( net=net, diff --git a/engiopt/generators/vqgan/adapter.py b/engiopt/generators/vqgan/adapter.py index 67fbd730..31ee19bd 100644 --- a/engiopt/generators/vqgan/adapter.py +++ b/engiopt/generators/vqgan/adapter.py @@ -62,7 +62,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device latent_dim=config["latent_dim"], num_codebook_vectors=config["num_codebook_vectors"], ) - vqgan.load_state_dict(th.load(resolved.files["vqgan.pth"], map_location=device, weights_only=False)["vqgan"]) + vqgan.load_state_dict(th.load(resolved.files["vqgan.pth"], map_location=device, weights_only=True)["vqgan"]) vqgan.eval().to(device) cvqgan = VQGAN( device=device, @@ -91,7 +91,7 @@ def build(cls, resolved: ResolvedCheckpoint, problem: Problem, device: th.device dropout=config["dropout"], ) net.load_state_dict( - th.load(resolved.files["transformer.pth"], map_location=device, weights_only=False)[cls.primary_state_key] + th.load(resolved.files["transformer.pth"], map_location=device, weights_only=True)[cls.primary_state_key] ) net.eval().to(device) return cls(net=net, latent_size=config["latent_size"], problem=problem, device=device, **base) @@ -109,7 +109,7 @@ def _load_cvqgan(cls, package_root: str, *, config: dict[str, Any], device: th.d if config["conditional"]: raise FileNotFoundError(f"Conditional VQGAN needs cvqgan.pth, missing from {package_root}") return None - return th.load(path, map_location=device, weights_only=False)["cvqgan"] + return th.load(path, map_location=device, weights_only=True)["cvqgan"] def _sample(self, conditions: ConditionBatch, n: int) -> th.Tensor: """Sample a full grid of latent codes autoregressively, then decode it.""" diff --git a/tests/test_generator_registry.py b/tests/test_generator_registry.py index 69a73966..58ab0301 100644 --- a/tests/test_generator_registry.py +++ b/tests/test_generator_registry.py @@ -113,3 +113,22 @@ def test_registry_is_a_plain_mapping_like_builtin_problems() -> None: """`BUILTIN_GENERATORS` mirrors `BUILTIN_PROBLEMS`: a dict of name -> class.""" assert isinstance(BUILTIN_GENERATORS, dict) assert all(isinstance(g, type) and issubclass(g, Generator) for g in BUILTIN_GENERATORS.values()) + + +def test_every_checkpoint_load_uses_the_restricted_unpickler() -> None: + """A `.pth` file is pickled Python, so loading one runs whatever it contains. + + `weights_only=True` swaps in an unpickler that can rebuild tensors and + nothing else. Every published checkpoint loads under it, so the unrestricted + form buys nothing and costs arbitrary code execution the moment a checkpoint + comes from someone else's HuggingFace account. + """ + adapters = [Path("engiopt/generators", name, "adapter.py") for name in BUILTIN_GENERATORS] + sources = [Path("engiopt/core.py"), Path("engiopt/generators/_template/adapter.py"), *adapters] + + offenders = [] + for path in sources: + for number, line in enumerate(path.read_text().splitlines(), start=1): + if "th.load(" in line and "weights_only=True" not in line: + offenders.append(f"{path}:{number}") + assert not offenders, "checkpoint loads without weights_only=True: " + ", ".join(offenders) From dfeb373e7f1cbfc23e1bfa04e859ac659191a5ab Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 8 Sep 2026 16:16:53 +0200 Subject: [PATCH 23/24] Say who can publish a row, and why history is not migrated Two documentation corrections from review. README claimed "Anyone may publish", contradicting LEADERBOARD.md, which says in two places that `--push-to` needs write access to the board. The checkpoints need no access; the row does. README now says which, and names #78 as the self-service route. The non-migration rationale was wrong on the facts: a migrated package can be hashed, and a fingerprint can be derived when the run config is recoverable. The real reason is the manual work of recovering and validating that metadata across historical W&B runs, and the incomplete provenance it would still leave. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 2 +- docs/checkpoint_layout.md | 13 +++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 2d6f8484..3abf5d2a 100644 --- a/README.md +++ b/README.md @@ -157,7 +157,7 @@ Publishing downloads the existing board, merges on the row key, and uploads the Every row records what produced it: **which repo, path, revision, and content hash** the weights came from, the EngiOpt version, and the EngiBench version that ran the evaluation. The spec records what it was frozen against, including the pinned dataset revision, so a change to either side is visible rather than silently shifting every number. -**Anyone may publish; nothing is ranked until it is re-run.** Rows land `verified=false`, and a runner re-fetches the checkpoint at its recorded revision and scores it itself before they enter the ranking: +**Contributors with write access to the leaderboard repository may publish provisional rows; external self-service submission is tracked in #78. Nothing is ranked until it is re-run.** Rows land `verified=false`, and a runner re-fetches the checkpoint at its recorded revision and scores it itself before they enter the ranking: ``` python -m engiopt.verify --board IDEALLab/engiopt-leaderboard # audit; writes nothing diff --git a/docs/checkpoint_layout.md b/docs/checkpoint_layout.md index 5458538f..5293afcd 100644 --- a/docs/checkpoint_layout.md +++ b/docs/checkpoint_layout.md @@ -43,12 +43,13 @@ Checkpoints trained before this layout existed lived in W&B artifacts, and they are **not being migrated**. W&B is no longer a checkpoint source, so those weights are reachable only by retraining. -The decision is deliberate rather than deferred. A migrated checkpoint would -arrive without the two things this layout exists to provide -- a config -fingerprint identifying the hyperparameters that produced it, and a content hash -tying a leaderboard row to exact bytes -- so it could be stored but never -canonically addressed or verified. Retraining a model costs less than a -provenance story nobody can check. +The decision is deliberate rather than deferred. A migrated package could be +hashed, and a config fingerprint could be derived when the complete run +configuration is recoverable. Recovering and validating that metadata +consistently across historical W&B runs would take substantial manual work and +may still leave incomplete provenance. We therefore require missing packages to +be retrained and published in the current format rather than migrating legacy +artifacts. What this means in practice: a generator with no published package is trained, not recovered. `--list-generators --check-availability` reports which those are From 32670ba297855589c6c9261e3e47c21689859d17 Mon Sep 17 00:00:00 2001 From: Matthew Keeler Date: Tue, 8 Sep 2026 16:41:35 +0200 Subject: [PATCH 24/24] Make the docs agree with the code they describe A pass over every document this branch ships, checking each claim against the code and against the other documents. Contradictions with the code: - CONTRIBUTING's build() example loaded a checkpoint without weights_only=True, which every adapter now passes and a test pins. - The same example used checkpoint_identity without importing it. - The PR template still asked for evaluate_.py and for restoring models from W&B artifacts; both are gone. It now asks for the adapter and the HuggingFace package. - cgan_cnn_3d's README listed --lr_enc. That model has no encoder. - "a metric is added once rather than thirteen times": there are fourteen. Merge artifacts: - LEADERBOARD.md said "the mechanism this PR adds". - README cited #78 bare, which does not link from a README; LEADERBOARD.md described the self-service route without naming the issue at all. - Two README sentences read as migration notes ("now use", "no longer") rather than as documentation. Spelling: American throughout, including docstrings and test comments this branch added (defence, behaviour, honour, modelling). Housekeeping: one dash style per document; engiopt/.DS_Store untracked and ignored; the EngiBench typo and the broken \" escape in the citation; trailing slashes on two table links. Co-Authored-By: Claude Fable 5.1 --- .github/PULL_REQUEST_TEMPLATE.md | 8 ++++---- .gitignore | 3 +++ CONTRIBUTING_A_MODEL.md | 12 +++++++----- LEADERBOARD.md | 13 +++++++------ README.md | 21 +++++++++------------ docs/checkpoint_layout.md | 2 +- engiopt/.DS_Store | Bin 8196 -> 0 bytes engiopt/checkpoint_store.py | 2 +- engiopt/core.py | 2 +- engiopt/evaluation/context.py | 2 +- engiopt/evaluation/metrics/builtin.py | 4 ++-- engiopt/evaluation/verify.py | 2 +- engiopt/generators/cgan_cnn_3d/README.md | 1 - engiopt/generators/cgan_vae/README.md | 2 +- tests/test_evaluation.py | 2 +- tests/test_leaderboard_integrity.py | 6 +++--- tests/test_verification.py | 2 +- 17 files changed, 43 insertions(+), 41 deletions(-) delete mode 100644 engiopt/.DS_Store diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 043060f0..0dc300d5 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -44,10 +44,10 @@ To upload images to a PR -- simply drag and drop an image while in edit mode and - [ ] The code is reproducible: random seeds are set, PyTorch determinism is enabled - [ ] Hyperparameters are configurable via command-line arguments using `tyro` - [ ] WandB logging is integrated with `--track` flag support -- [ ] The model can be saved and restored via WandB artifacts (`--save-model` flag) +- [ ] The model is saved as a HuggingFace checkpoint package via `--save-model` (see `docs/checkpoint_layout.md`) ### Algorithm Completeness (for new algorithms) -- [ ] Both training script (`algorithm.py`) and evaluation script (`evaluate_algorithm.py`) are provided +- [ ] Both the training script (`algorithm.py`) and the `Generator` adapter (`adapter.py`) are provided; evaluation runs through `python -m engiopt.evaluate` - [ ] The algorithm works with EngiBench's `Problem` interface - [ ] The algorithm is added to the README table with correct metadata @@ -71,7 +71,7 @@ For example, - [ ] The code follows the CleanRL single-file philosophy where appropriate. - [ ] The code is understandable and commented. No large code blocks are left unexplained. Can I read and understand the code easily? - [ ] There is no merge conflict. -- [ ] For new algorithms: both training and evaluation scripts are provided. -- [ ] For new algorithms: WandB integration is complete (logging, artifacts, reproducibility). +- [ ] For new algorithms: the training script and `Generator` adapter are provided, and `python -m engiopt.evaluate` scores the model. +- [ ] For new algorithms: `--track` logging works and the checkpoint lands on HuggingFace with `--save-model`. - [ ] For bugfixes: it is a robust fix and not a hacky workaround. - [ ] The changes do not break existing trained models or evaluation workflows without good reason. diff --git a/.gitignore b/.gitignore index 463236ba..495e39e1 100644 --- a/.gitignore +++ b/.gitignore @@ -171,3 +171,6 @@ logs/* .ruff_cache/ engibench_studies/* + +# macOS +.DS_Store diff --git a/CONTRIBUTING_A_MODEL.md b/CONTRIBUTING_A_MODEL.md index 9469f49c..87c822fa 100644 --- a/CONTRIBUTING_A_MODEL.md +++ b/CONTRIBUTING_A_MODEL.md @@ -29,6 +29,7 @@ model lands where the loader expects: ```python from engiopt.checkpoint_store import save_checkpoint_package +from engiopt.core import checkpoint_identity save_checkpoint_package( checkpoint_backend="hf", @@ -49,7 +50,7 @@ save_checkpoint_package( `checkpoint_identity(args)` is what lets many hyperparameter settings coexist. Each run is stored at `{problem_id}/cfg_{fingerprint}/seed_{seed}`, and a run using your script's *default* hyperparameters additionally claims the canonical -`{problem_id}/seed_{seed}` -- which is what the bare model name resolves to. A +`{problem_id}/seed_{seed}` — which is what the bare model name resolves to. A sweep therefore can never redefine what your model means. `run_config` matters: it is what `from_pretrained` reads to reconstruct your @@ -81,7 +82,7 @@ n_conds = len(condition_keys_for(problem, resolved)) ``` That is what keeps a checkpoint loadable after a problem gains a condition, and -what makes the evaluator refuse -- with a clear message -- to feed your model +what makes the evaluator refuse — with a clear message — to feed your model conditions it was not trained on. Also reuse the shared helpers rather than rewriting them: @@ -121,7 +122,8 @@ Then implement the two methods: def build(cls, resolved, problem, device, **base): config = resolved.run_config net = MyNet(latent_dim=config["latent_dim"], design_shape=problem.design_space.shape).to(device) - net.load_state_dict(th.load(resolved.files["generator.pth"], map_location=device)[cls.primary_state_key]) + state = th.load(resolved.files["generator.pth"], map_location=device, weights_only=True) + net.load_state_dict(state[cls.primary_state_key]) net.eval() return cls(net=net, latent_dim=config["latent_dim"], problem=problem, device=device, **base) ``` @@ -184,7 +186,7 @@ Watch two columns while you develop: - **`cond_sens`** should be greater than zero if you declared `conditional = True`. Exactly zero means your conditions are not reaching the network, which - is a wiring bug far more often than a modelling choice. + is a wiring bug far more often than a modeling choice. - **`copy_rate`** should be near zero. High means your model is reproducing training designs rather than generating, and the board will publish it without ranking it. @@ -234,7 +236,7 @@ in a notebook cell and it will appear in the next leaderboard you build. | Loss curves, sample images | Weights & Biases, when `--track` is on | HuggingFace holds everything that must be reloaded or compared; W&B holds what -you only look at. Nothing in the evaluation path requires W&B -- training with +you only look at. Nothing in the evaluation path requires W&B — training with `--track false` produces exactly the same checkpoints and scores. To attach a model's scores to its own checkpoint: diff --git a/LEADERBOARD.md b/LEADERBOARD.md index f6929fb9..4d076cbf 100644 --- a/LEADERBOARD.md +++ b/LEADERBOARD.md @@ -3,7 +3,7 @@ A public board has two problems an internal one does not, and they are independent. Someone can report numbers no model produced, and someone can report perfectly honest numbers from a model that games the metric instead of -solving the problem. Neither defence helps with the other. +solving the problem. Neither defense helps with the other. - Against **fabrication**: re-execution. Nothing is ranked until a runner has fetched the weights itself and reproduced the score. @@ -13,7 +13,7 @@ solving the problem. Neither defence helps with the other. ## Submitting -Publish your checkpoints to your own HuggingFace account -- nothing about the +Publish your checkpoints to your own HuggingFace account — nothing about the weights needs `IDEALLab`: ```bash @@ -34,9 +34,10 @@ board and absent from the ranking until a runner re-scores them. the merged CSV directly, so submission is currently open to people who already have it. That is a real limit on how public this board is, and it is the one piece of the submission path that is not yet self-service; a no-access route -(a pending manifest opened as a HuggingFace Community PR) is tracked alongside -the safe-loading work below. Everything else above -- the checkpoints, the -address a row records, the audit command -- already works from any account. +(a pending manifest opened as a HuggingFace Community PR) is tracked in +[#78](https://github.com/IDEALLab/EngiOpt/issues/78), alongside the +safe-loading work below. Everything else above — the checkpoints, the address a +row records, the audit command — already works from any account. Three things are worth knowing before you run it. @@ -178,7 +179,7 @@ designs. That is the right benchmark for a generative design model, and it is not built here: freezing such a spec means an optimizer run per condition. The mechanism -this PR adds — flags, eligibility, and a spec version that carries its own +already in place — flags, eligibility, and a spec version that carries its own thresholds — is what a v2 would plug into. ## Reading the board diff --git a/README.md b/README.md index 3abf5d2a..55db08e8 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ This repository contains the code for optimization and machine learning algorith ## Coding Philosophy As much as we can, we follow the [CleanRL](https://github.com/vwxyzjn/cleanrl) philosophy: single-file, high-quality implementations with research-friendly features: -* Single-file implementation: every training detail is in one file, so you can easily understand and modify the code. Evaluation is not per-model — every model is scored by one shared evaluator through the `Generator` contract in `adapter.py`, so a metric is added once rather than thirteen times. +* Single-file implementation: every training detail is in one file, so you can easily understand and modify the code. Evaluation is not per-model — every model is scored by one shared evaluator through the `Generator` contract in `adapter.py`, so a metric is added once rather than once per model. * High-quality: we use type hints, docstrings, and comments to make the code easy to understand. We also rely on linters for formatting and checking our code. * Logging: we use experiment tracking tools like [Weights & Biases](https://wandb.ai/site) to log the results of our experiments. All our "official" runs are logged in the [EngiOpt project](https://wandb.ai/engibench/engiopt). * Reproducibility: we seed all the random number generators, make PyTorch deterministic, report the hyperparameters and code in WandB. @@ -35,8 +35,8 @@ As much as we can, we follow the [CleanRL](https://github.com/vwxyzjn/cleanrl) p [gan_bezier](engiopt/generators/gan_bezier/) | Inverse Design | 1D | ❌ | GAN + Bezier layer [gan_cnn_2d](engiopt/generators/gan_cnn_2d/) | Inverse Design | 2D | ❌ | GAN + CNN [surrogate_model](engiopt/surrogate_model/) | Surrogate Model | 1D | ❌ | MLP -[vqgan](engiopt/generators/vqgan) | Inverse Design | 2D | ✅ | VQVAE + Transformer -[pixel_cnn_pp_2d](engiopt/generators/pixel_cnn_pp_2d) | Inverse Design | 2D | ✅ | PixelCNN++ Autoregressive Model +[vqgan](engiopt/generators/vqgan/) | Inverse Design | 2D | ✅ | VQVAE + Transformer +[pixel_cnn_pp_2d](engiopt/generators/pixel_cnn_pp_2d/) | Inverse Design | 2D | ✅ | PixelCNN++ Autoregressive Model Every algorithm above is registered, meaning this repository holds an adapter that can rebuild it. Being *evaluable* additionally needs published weights, and the two @@ -54,7 +54,7 @@ Historical W&B-era checkpoints are **not** being migrated; see [docs/checkpoint_layout.md](docs/checkpoint_layout.md#historical-checkpoints). ## Dashboards -HuggingFace hosts everything that has to be reloaded or compared -- model weights, run configs, evaluation metrics, and the leaderboard. WandB hosts what you only look at: loss curves and sample images. Nothing in the evaluation path requires WandB, so training with `--track false` produces exactly the same checkpoints and scores. You can access some of our runs at https://wandb.ai/engibench/engiopt. +HuggingFace hosts everything that has to be reloaded or compared — model weights, run configs, evaluation metrics, and the leaderboard. WandB hosts what you only look at: loss curves and sample images. Nothing in the evaluation path requires WandB, so training with `--track false` produces exactly the same checkpoints and scores. You can access some of our runs at https://wandb.ai/engibench/engiopt. WandB dashboards @@ -157,7 +157,7 @@ Publishing downloads the existing board, merges on the row key, and uploads the Every row records what produced it: **which repo, path, revision, and content hash** the weights came from, the EngiOpt version, and the EngiBench version that ran the evaluation. The spec records what it was frozen against, including the pinned dataset revision, so a change to either side is visible rather than silently shifting every number. -**Contributors with write access to the leaderboard repository may publish provisional rows; external self-service submission is tracked in #78. Nothing is ranked until it is re-run.** Rows land `verified=false`, and a runner re-fetches the checkpoint at its recorded revision and scores it itself before they enter the ranking: +**Contributors with write access to the leaderboard repository may publish provisional rows; external self-service submission is tracked in [#78](https://github.com/IDEALLab/EngiOpt/issues/78). Nothing is ranked until it is re-run.** Rows land `verified=false`, and a runner re-fetches the checkpoint at its recorded revision and scores it itself before they enter the ranking: ``` python -m engiopt.verify --board IDEALLab/engiopt-leaderboard # audit; writes nothing @@ -180,14 +180,11 @@ The current surrogate model comprises several steps: See this [notebook](https://github.com/IDEALLab/EngiOpt/blob/main/engiopt/surrogate_model/case_study_pe_notebook.ipynb) for an example. -Surrogate-model optimization paths now use the same checkpoint abstraction. For example, the power-electronics optimizer can consume: +Surrogate-model optimization uses the same checkpoint abstraction. For example, the power-electronics optimizer can consume: * HF package refs such as `hf://IDEALLab/engiopt-mlp-tabular-only/power_electronics/DcGain/seed_42` * local checkpoint package directories -HuggingFace is the only checkpoint backend. W&B artifacts are no longer a model source anywhere in the codebase; see [docs/checkpoint_layout.md](docs/checkpoint_layout.md) for the package layout. - - - +HuggingFace is the only checkpoint backend; see [docs/checkpoint_layout.md](docs/checkpoint_layout.md) for the package layout. ## Colab notebooks We have some colab notebooks that show how to use some of the EngiBench/EngiOpt features. @@ -197,7 +194,7 @@ We have some colab notebooks that show how to use some of the EngiBench/EngiOpt ## Citing -If you use EngiBenc/EngiOpt in your research, please cite the following paper: +If you use EngiBench/EngiOpt in your research, please cite the following paper: ```bibtex @misc{felten_engibench_2025, @@ -206,7 +203,7 @@ If you use EngiBenc/EngiOpt in your research, please cite the following paper: doi = {10.48550/arXiv.2508.00831}, urldate = {2025-08-07}, publisher = {arXiv}, - author = {Felten, Florian and Apaza, Gabriel and B\¨aunlich, Gerhard and Diniz, Cashen and Dong, Xuliang and Drake, Arthur and Habibi, Milad and Hoffman, Nathaniel J. and Keeler, Matthew and Massoudi, Soheyl and VanGessel, Francis G. and Fuge, Mark}, + author = {Felten, Florian and Apaza, Gabriel and B{\"a}unlich, Gerhard and Diniz, Cashen and Dong, Xuliang and Drake, Arthur and Habibi, Milad and Hoffman, Nathaniel J. and Keeler, Matthew and Massoudi, Soheyl and VanGessel, Francis G. and Fuge, Mark}, month = jun, year = {2025}, } diff --git a/docs/checkpoint_layout.md b/docs/checkpoint_layout.md index 5293afcd..c8157258 100644 --- a/docs/checkpoint_layout.md +++ b/docs/checkpoint_layout.md @@ -79,7 +79,7 @@ save_checkpoint_package( ## Reading one ```python -# the default-hyperparameter checkpoint -- what the bare model name means +# the default-hyperparameter checkpoint — what the bare model name means gen = BUILTIN_GENERATORS["cgan_cnn_2d"].from_pretrained(problem, problem_id="beams2d", seed=1) # one specific configuration from a sweep diff --git a/engiopt/.DS_Store b/engiopt/.DS_Store deleted file mode 100644 index 716392b4aa3bbc2c6144b68e316ff6bb5341334f..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 8196 zcmeHM&2G~`5T5O?LnsspP=!FU#I-6-f+`S~5XyloL2v*Rk~nEC96PcdqP8l^9e4#^ zfjbYvyKsVUX5G|wVh{%?Qg^G}8Qb&CtiRc{GcFOS#>lAXFs1BL;^!2iGi-r3x&g7?0@YDU9=Vc=Xcz|RL8 zon<|d6D{S@fk8?D$R(VX1|4;PVC+QJ6FJdRpzx`y2bDloQer3#$90>{A?t~pXsO|( zG@MitStS*UlIXzMVos{Ar5OzahJjfIxOOi~&F+$)soxKtctN-01sz!Mrh16J^+Ou8 z6V#9ns7r11ed@8F9O_FBwFcXe_9%kw2ww|a@}Tqgf?PaxD~Buzk`aB07gLmtr9R{k ztC-auY<W3F|Qtl*%-pv=L~13GwNbqJhPct#GIv3 zH_5?W#OM)X@E4W{`T(4H!g-ehcyc6ruIR;#ru|+NOdouX;87#zF~z*5aaZQza37di zGM6k5in&gESSiknFg@RPmMN@Bbt%Q~g{j}-oPgZ|OkqZNE#@!-$wUa&x&`YOtYt78 z$jney`~av!$y^_MAwVWu;4YovdHVqzX8c`7+!65XfmeAf9tSvbpn2I0Fspxp?u?^h zviHAshMpV6+uPsN!s62M%Br<$m95v!{b1aT+wmxBw8NKt^ehOx&WYacJa>EJ*80^) zfgiVBzZVL?b$T#)^}_X?VB82sz7q+ysV`V1tJGS*adcFxRqt)um0D%zc+)<*U)|oc z@78vXk4x6I+jky3ZN2pde(({i3QZ+H3203pPsdNte#6qcX6So?&%2wwtbxT4T>9X) z4<G4(`Fw8R;-fBqqW StCZ=lB=i24B$(mcG4K th.Tensor | npt.NDArray """ # ------------------------------------------------------------------ - # Shared behaviour + # Shared behavior # ------------------------------------------------------------------ def sample( diff --git a/engiopt/evaluation/context.py b/engiopt/evaluation/context.py index 578e17bb..0525602d 100644 --- a/engiopt/evaluation/context.py +++ b/engiopt/evaluation/context.py @@ -352,7 +352,7 @@ def is_infeasible(self, design: Any, conditions: dict[str, Any] | None) -> bool: This is what a new problem gets for free. 2. The volume-fraction budget named by the spec's `volume_condition`, when the problem has one. Missing that target is a design failing to - honour its brief rather than an invalid design, and no EngiBench + honor its brief rather than an invalid design, and no EngiBench constraint covers it. A problem with neither -- photonics2d has no volume budget -- is scored diff --git a/engiopt/evaluation/metrics/builtin.py b/engiopt/evaluation/metrics/builtin.py index b3de299b..731452e8 100644 --- a/engiopt/evaluation/metrics/builtin.py +++ b/engiopt/evaluation/metrics/builtin.py @@ -2,7 +2,7 @@ These wrap the numerics in `engiopt.metrics` and declare, for each one, the question it answers, whether it needs a simulator, and which direction is -better. Behaviour is identical to the per-model `evaluate_*.py` scripts these +better. Behavior is identical to the per-model `evaluate_*.py` scripts these replaced, so numbers remain comparable. Adding a metric means adding one decorated function here (or in your own module @@ -77,7 +77,7 @@ def novelty(ctx: EvaluationContext) -> dict[str, float]: committed spec. A lookup table keyed on the condition vector therefore tops `mmd`, `iog`, and `fog` -- not by cheating the implementation, but because those metrics are *defined* as closeness to exactly the designs it returns. - No amount of care in the evaluator changes that; the only defence available + No amount of care in the evaluator changes that; the only defense available to a public board is to measure retrieval and say so. Two columns, because they answer different questions: diff --git a/engiopt/evaluation/verify.py b/engiopt/evaluation/verify.py index 3ff0d296..b531a2c8 100644 --- a/engiopt/evaluation/verify.py +++ b/engiopt/evaluation/verify.py @@ -1,7 +1,7 @@ """Re-running a published row from its checkpoint address. The board cannot trust the numbers it is given. Whoever computed a row also -chose what to report, and no amount of schema design changes that -- the defence +chose what to report, and no amount of schema design changes that -- the defense against a fabricated score is not a stricter column, it is running the model again. diff --git a/engiopt/generators/cgan_cnn_3d/README.md b/engiopt/generators/cgan_cnn_3d/README.md index ec208d3d..8958feb0 100644 --- a/engiopt/generators/cgan_cnn_3d/README.md +++ b/engiopt/generators/cgan_cnn_3d/README.md @@ -45,7 +45,6 @@ python -m engiopt.evaluate --problem-id heatconduction3d --generators cgan_cnn_3 - `--batch_size 8`: Batch size - `--lr_gen 0.0025`: Generator learning rate - `--lr_disc 10e-5`: Discriminator learning rate -- `--lr_enc 0.001`: Encoder learning rate - `--latent_dim 64`: Dimensionality of latent space ### Training Dynamics diff --git a/engiopt/generators/cgan_vae/README.md b/engiopt/generators/cgan_vae/README.md index 84037d38..150b8926 100644 --- a/engiopt/generators/cgan_vae/README.md +++ b/engiopt/generators/cgan_vae/README.md @@ -60,7 +60,7 @@ python -m engiopt.evaluate --problem-id heatconduction3d --generators cgan_vae - -## References: +## References Original VAE-GAN: https://arxiv.org/pdf/1512.09300 Multiview approach: https://github.com/bryonkucharski/Multiview-3D-VAE-GAN diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py index 7cfe9aae..e6e80d68 100644 --- a/tests/test_evaluation.py +++ b/tests/test_evaluation.py @@ -111,7 +111,7 @@ def test_cheap_metrics_never_touch_the_solver(fake_problem: Any) -> None: # ---------------------------------------------------------------------- -# Registry behaviour +# Registry behavior # ---------------------------------------------------------------------- diff --git a/tests/test_leaderboard_integrity.py b/tests/test_leaderboard_integrity.py index 7765ccb3..26b73da0 100644 --- a/tests/test_leaderboard_integrity.py +++ b/tests/test_leaderboard_integrity.py @@ -101,7 +101,7 @@ def test_copying_the_training_split_is_caught_too(fake_problem: Any) -> None: def test_near_copies_count_as_copies(fake_problem: Any) -> None: """Adding imperceptible noise to a retrieved design must not launder it. - Otherwise the defence is defeated by one line, and the tolerance is what + Otherwise the defense is defeated by one line, and the tolerance is what decides how much perturbation counts as having generated something. """ rng = np.random.default_rng(3) @@ -290,7 +290,7 @@ def test_flagged_rows_stay_on_the_board_and_out_of_the_ranking() -> None: def test_an_entry_missing_a_required_seed_is_not_ranked() -> None: - """The cherry-picking defence. + """The cherry-picking defense. Twenty seeds run, the best three published, and a median over a maximum looks exactly like an honest median in every individual row. Requiring @@ -400,7 +400,7 @@ def _sample(self, conditions: Any, n: int) -> Any: def test_the_evaluator_detects_a_lookup_table_end_to_end(fake_problem: Any) -> None: """The full path: sample, build the corpus, score, and report `copy_rate`. - This is the claim the whole memorization defence rests on, so it is asserted + This is the claim the whole memorization defense rests on, so it is asserted through the evaluator rather than against a hand-built context -- a metric that works but is never fed would pass every other test in this file. """ diff --git a/tests/test_verification.py b/tests/test_verification.py index 8bb21685..8359e233 100644 --- a/tests/test_verification.py +++ b/tests/test_verification.py @@ -1,6 +1,6 @@ """Tests for re-running a published row from its checkpoint address. -Verification is the only defence a public board has against a number that no +Verification is the only defense a public board has against a number that no model produced, so the properties here are about what it refuses to accept: weights that have moved since they were scored, addresses that lead nowhere, and models nobody can rebuild. The successful path matters less than those, because