Normalize 2D diffusion designs to DDPM sample scale

[SoheylM]

Summary

I added a follow-up fix for the 2D conditional diffusion sampler/training scale after checking the qualitative Beams2D outputs from the previous diffusion fixes.

What changed:

• Normalize 2D training designs from their observed training-data bounds to the DDPM sample range [-1, 1] instead of assuming the model should train directly on [0, 1].
• During sampling, reconstruct the predicted clean sample x0, clip that prediction in normalized [-1, 1], and use it to compute the posterior mean.
• Save design_min and design_max in the checkpoint so evaluation can inverse-transform generated samples back to the original problem scale.
• Keep evaluation backward-compatible with older checkpoints that do not have saved design bounds.
• Add focused tests for arbitrary-bound normalization and the clipped-x0 posterior step.

Why:

• The previous fixes made the schedule/noise path saner, but local Beams2D sampling showed the custom sampler could still let generated values drift far outside the training range.
• This makes the 2D diffusion path problem-agnostic: problems are no longer assumed to already live in [0, 1], and outputs are mapped back to the original bounds for metrics.

Validation:

• conda run -n engibench python -m pytest -q
• conda run -n engibench ruff check engiopt/diffusion_2d_cond/diffusion_2d_cond.py engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py tests/test_diffusion_fixes.py
• conda run -n engibench ruff format --check engiopt/diffusion_2d_cond/diffusion_2d_cond.py engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py tests/test_diffusion_fixes.py
• git diff --check
• conda run -n engibench python -m compileall -q engiopt/diffusion_2d_cond tests
• One-epoch Beams2D smoke training run with the new scaling path completed successfully.

[SoheylM]

I wanted to add a bit more context on why I changed the 2D diffusion scaling/sampling path here.

The main issue I found after the previous diffusion fixes was that the 2D sampler could still drift outside the scale of the design problem during reverse diffusion. For Beams2D, the old sampling path could produce raw values far outside the valid design range, which then showed up as washed backgrounds, ghost structures, or collapsed samples. This PR makes the scale contract explicit:

• Use the EngiBench problem.design_space.low / problem.design_space.high bounds as the source of truth for design scaling.
• Normalize designs to the DDPM sample range [-1, 1] for training.
• During sampling, reconstruct predicted clean x0, clip that clean prediction in [-1, 1], and use it to compute the posterior mean.
• Save the design bounds in the checkpoint so evaluation can map generated samples back to the original problem scale.
• Keep a training-data min/max fallback only if the problem bounds are not finite/usable.

I also ran a short Beams2D sanity training run with the new path. This was not meant to be a final paper-quality model, just a check that the implementation behaves sensibly:

• 50 training epochs on Beams2D
• checkpoint epoch: 49
• final batch loss: 0.001645
• generated raw range after inverse transform: 0.000083 to 0.999917
• fraction of generated pixels clipped back to bounds: 0.000000
• mean absolute difference vs the sampled test references: 0.113108

Qualitatively, the 50-epoch samples are not perfect, but they are much better behaved than the previous unstable sampler path: the generated designs stayed in range before the defensive final evaluation-time bounds clip was needed, and several Beams2D structures are recognizable under test-set conditions. I see this as a better implementation because it uses the EngiBench problem definition as the scale source of truth and keeps the DDPM sampler internally consistent, rather than relying on accidental dataset extrema or a final-output clip to rescue out-of-range samples.

Validation I ran locally:

• conda run -n engibench python -m pytest -q
• conda run -n engibench ruff check engiopt/diffusion_1d/diffusion_1d.py engiopt/diffusion_1d/evaluate_diffusion_1d.py engiopt/diffusion_2d_cond/diffusion_2d_cond.py engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py tests
• conda run -n engibench ruff format --check engiopt/diffusion_1d/diffusion_1d.py engiopt/diffusion_1d/evaluate_diffusion_1d.py engiopt/diffusion_2d_cond/diffusion_2d_cond.py engiopt/diffusion_2d_cond/evaluate_diffusion_2d_cond.py tests
• git diff --check

One note: this still uses the local EngiOpt DDPM sampler rather than switching fully to Hugging Face DDPMScheduler. I think that is acceptable here because the sampler is now small, explicit, and covered by focused tests for the sensitive scale/posterior behavior. A full migration to Diffusers scheduler could still be a good future cleanup, but this PR keeps the change conservative and aligned with the existing code structure.

[SoheylM]

Adding the qualitative sanity-check figure from the short Beams2D run so the numerical results above have visual context.

What this shows:

• Top row: reference optimal designs from the Beams2D test split for sampled test conditions.
• Middle row: generated designs from the 50-epoch diffusion checkpoint trained with the PR Normalize 2D diffusion designs to DDPM sample scale #68 scaling/sampling path.
• Bottom row: absolute difference between generated and reference designs.

This is not intended as a final paper-quality model. It is a quick qualitative sanity check that the implementation is better behaved: generated designs remain inside the valid design bounds after inverse transform, several structures are recognizable, and the sampler no longer depends on the final evaluation-time bounds clip to rescue out-of-range samples. There is still a defensive final clip in evaluation, but in this sanity run the measured clipped fraction was 0.0; the point is that the reverse process itself stayed inside the intended normalized scale. The remaining washed backgrounds / imperfect members are model-quality issues to keep investigating, but the scale contract in this PR is much cleaner and more problem-consistent.

[g-braeunlich]

LGTM