fix v0.py warm-start and remove redundant calcs in optimize

engibench/problems/airfoil/v0.py

@@ -6,7 +6,7 @@
 import dataclasses
 from dataclasses import dataclass
 from dataclasses import field
-import importlib
+from importlib.util import find_spec
 import json
 import os
 import shutil
@@ -37,7 +37,7 @@
 from engibench.utils.files import clone_dir
 
 # Allow loading pyoptsparse histories even if pyoptsparse is not installed:
-if importlib.util.find_spec("pyoptsparse") is None:
+if find_spec("pyoptsparse") is None:
     from engibench.problems.airfoil import fake_pyoptsparse
 
     sys.modules["pyoptsparse"] = fake_pyoptsparse

engibench/problems/beams2d/__init__.py

@@ -1,5 +1,5 @@
 """Beams2D problem module."""
 
-from engibench.problems.beams2d.v0 import Beams2D
+from engibench.problems.beams2d.v1 import Beams2D
 
 __all__ = ["Beams2D"]

engibench/problems/beams2d/v0.py

@@ -212,7 +212,7 @@ def rmin_bound(rmin: float, nelx: int, nely: int) -> None:
 
     design_constraints = (volume_fraction_bound,)
     design_space = spaces.Box(low=0.0, high=1.0, shape=(Config.nely, Config.nelx), dtype=np.float64)
-    dataset_id = f"IDEALLab/beams_2d_{Config.nely}_{Config.nelx}_v{version}"
+    dataset_id = f"IDEALLab/beams_2d_{Config.nely}_{Config.nelx}_v0"
     container_id = None
 
     def __init__(self, seed: int = 0, config: dict[str, Any] | None = None):
@@ -235,7 +235,7 @@ def __init__(self, seed: int = 0, config: dict[str, Any] | None = None):
         self.nelx = self.config.nelx
         self.nely = self.config.nely
         self.design_space = spaces.Box(low=0.0, high=1.0, shape=(self.nely, self.nelx), dtype=np.float64)
-        self.dataset_id = f"IDEALLab/beams_2d_{self.nely}_{self.nelx}_v{self.version}"
+        self.dataset_id = f"IDEALLab/beams_2d_{self.nely}_{self.nelx}_v0"
 
     def simulate(
         self, design: npt.NDArray, config: dict[str, Any] | None = None, *, ce: npt.NDArray | None = None
@@ -385,18 +385,20 @@ def random_design(self, dataset_split: str = "train", design_key: str = "optimal
         return np.array(self.dataset[dataset_split][design_key][rnd]), rnd
 
 
-if __name__ == "__main__":
-    # Provides a way to instantiate the problem without having to pass configs to optimize or simulate later.
-    # Possible sets of nely and nelx: (25, 50), (50, 100), and (100, 200)
-    # If a new nely and nelx are not passed in, uses the default conditions.
+def main(problem_type: type[Problem]) -> None:
+    """Provides a way to instantiate the problem without having to pass configs to optimize or simulate later.
 
-    problem = Beams2D(seed=0)
+    Possible sets of nely and nelx: (25, 50), (50, 100), and (100, 200)
+    If a new nely and nelx are not passed in, uses the default conditions.
+    """
+    problem = problem_type(seed=0)
 
+    assert hasattr(problem, "nelx")
+    assert hasattr(problem, "nely")
     print(f"Loading dataset for nely={problem.nely}, nelx={problem.nelx}.")
     dataset = problem.dataset
 
     # Example of getting the training set
-    optimal_train = dataset["train"]["optimal_design"]
     c_train = dataset["train"]["c"]
     condition_keys = [f.name for f in dataclasses.fields(problem.Conditions)]
     params_train = dataset["train"].select_columns(condition_keys)
@@ -406,7 +408,7 @@ def random_design(self, dataset_split: str = "train", design_key: str = "optimal
     design, idx = problem.random_design()
     config = params_train[idx]
     compliance = c_train[idx]
-    fig, ax = problem.render(design, open_window=True)
+    problem.render(design, open_window=True)
 
     print(f"Verifying compliance via simulation. Reference value: {compliance:.4f}")
 
@@ -421,8 +423,12 @@ def random_design(self, dataset_split: str = "train", design_key: str = "optimal
     problem.reset(seed=1)
 
     # NOTE: optimal_design and optisteps_history[-1].stored_design are interchangeable.
-    optimal_design, optisteps_history = problem.optimize(config=config)
+    optimal_design, optisteps_history = problem.optimize(config=config, starting_point=None)
     print(f"Final compliance: {optisteps_history[-1].obj_values[0]:.4f}")
     print(f"Final design volume fraction: {optimal_design.sum() / (np.prod(optimal_design.shape)):.4f}")
 
-    fig, ax = problem.render(optimal_design, open_window=True)
+    problem.render(optimal_design, open_window=True)
+
+
+if __name__ == "__main__":
+    main(Beams2D)

engibench/problems/beams2d/v1.py

@@ -0,0 +1,112 @@
+# ruff: noqa: N806
+# Disabled variable name conventions
+
+"""Beams 2D problem."""
+
+from copy import deepcopy
+import dataclasses
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+from engibench.problems.beams2d.backend import calc_sensitivity
+from engibench.problems.beams2d.backend import design_to_image
+from engibench.problems.beams2d.backend import image_to_design
+from engibench.problems.beams2d.backend import inner_opt
+from engibench.problems.beams2d.backend import overhang_filter_d
+from engibench.problems.beams2d.backend import overhang_filter_x
+from engibench.problems.beams2d.backend import State
+from engibench.problems.beams2d.v0 import Beams2D as Beams2D_v0
+from engibench.problems.beams2d.v0 import ExtendedOptiStep
+from engibench.problems.beams2d.v0 import main
+from engibench.utils.upcast import upcast
+
+
+class Beams2D(Beams2D_v0):
+    r"""Beam 2D topology optimization problem - Version 1 (v1).
+
+    ### v1
+    This version augments v0 by fixing a minor detail in the v0 warm-start optimization process.
+    Specifically, when warm-starting from a provided design, a small epsilon value is added to
+    avoid zero-density values that could lead to gradient issues. The datasets themselves remain unchanged.
+
+    All other behavior is identical to v0.
+    See v0.py for full baseline documentation.
+    """
+
+    version = 1
+
+    def optimize(
+        self, starting_point: npt.NDArray | None = None, config: dict[str, Any] | None = None
+    ) -> tuple[np.ndarray, list[ExtendedOptiStep]]:
+        """Optimizes the design of a beam.
+
+        Args:
+            starting_point (npt.NDArray or None): The design to begin warm-start optimization from (optional).
+            config (dict): A dictionary with configuration (e.g., boundary conditions) for the optimization.
+
+        Returns:
+            Tuple[np.ndarray, dict]: The optimized design and its performance.
+        """
+        base_config = self.Config(**{**dataclasses.asdict(self.simulate_config), **(config or {})})
+
+        self.__st = State.new(base_config.nelx, base_config.nely, base_config.rmin, base_config.forcedist)
+
+        # Returns the full history of the optimization instead of just the last step
+        optisteps_history = []
+
+        if starting_point is None:
+            xPhys = base_config.volfrac * np.ones((base_config.nelx, base_config.nely), dtype=float)
+            x = xPhys.ravel()
+        else:
+            starting_point = image_to_design(starting_point)
+            assert starting_point is not None
+            eps = 1e-4
+            x = (
+                (1 - eps) * starting_point + eps * base_config.volfrac
+            )  # add tiny non-zero values to avoid warm-start gradient issues for zero values
+            xPhys = x.reshape((base_config.nelx, base_config.nely))
+
+        xPrint = overhang_filter_x(xPhys) if base_config.overhang_constraint else xPhys.ravel()
+        loop, change = (0, 1.0)
+
+        while change > self.__st.min_change and loop < base_config.max_iter:
+            ce = calc_sensitivity(xPrint, st=self.__st, cfg=dataclasses.asdict(base_config))
+            simulate_config = upcast(base_config)
+            self.reset_called = True  # override for multiple reset calls in optimize
+            c = self.simulate(xPrint, ce=ce, config=dataclasses.asdict(simulate_config))
+
+            # Record the current state in optisteps_history
+            current_step = ExtendedOptiStep(obj_values=np.array(c), step=loop)
+            current_step.design = np.array(xPrint)
+            optisteps_history.append(current_step)
+
+            loop += 1
+
+            dc = (-base_config.penal * xPrint ** (base_config.penal - 1) * (self.__st.Emax - self.__st.Emin)) * ce
+            dv = np.ones(base_config.nely * base_config.nelx)
+            # MATLAB implementation:
+            if base_config.overhang_constraint:
+                xPrint, dc, dv = overhang_filter_d(xPhys, dc, dv)
+            else:
+                xPrint = xPhys.ravel()
+
+            dc = np.asarray(self.__st.H * (dc[np.newaxis].T / self.__st.Hs))[:, 0]
+            dv = np.asarray(self.__st.H * (dv[np.newaxis].T / self.__st.Hs))[:, 0]
+            # Ensure dc remains nonpositive
+            dc = np.clip(dc, None, 0.0)
+
+            xnew, xPhys, xPrint = inner_opt(x, self.__st, dc, dv, dataclasses.asdict(base_config))
+            # Compute the change by the inf. norm
+            change = np.linalg.norm(  # type: ignore[assignment]
+                xnew.reshape(base_config.nelx * base_config.nely, 1) - x.reshape(base_config.nelx * base_config.nely, 1),
+                np.inf,
+            )
+            x = deepcopy(xnew)
+
+        return design_to_image(xPrint, base_config.nelx, base_config.nely), optisteps_history
+
+
+if __name__ == "__main__":
+    main(Beams2D)

engibench/problems/heatconduction2d/templates/optimize_heat_conduction_2d.py

@@ -7,7 +7,7 @@
 """
 
 import glob
-import importlib
+from importlib.util import find_spec
 import os
 import re
 
@@ -40,7 +40,7 @@
 from engibench.utils.cli import np_array_from_stdin
 
 # Ensure IPOPT is available
-if importlib.util.find_spec("pyadjoint.ipopt") is None:
+if find_spec("pyadjoint.ipopt") is None:
     raise ImportError("""This example depends on IPOPT and Python ipopt bindings. \
     When compiling IPOPT, make sure to link against HSL, as it \
     is a necessity for practical problems.""")

engibench/problems/heatconduction3d/templates/optimize_heat_conduction_3d.py

@@ -7,7 +7,7 @@
 """
 
 import glob
-import importlib
+from importlib.util import find_spec
 import os
 import re
 
@@ -41,7 +41,7 @@
 from engibench.utils.cli import np_array_from_stdin
 
 # Ensure IPOPT is available
-if importlib.util.find_spec("pyadjoint.ipopt") is None:
+if find_spec("pyadjoint.ipopt") is None:
     raise ImportError("""This example depends on IPOPT and Python ipopt bindings. \
     When compiling IPOPT, make sure to link against HSL, as it \
     is a necessity for practical problems.""")

engibench/problems/heatconduction3d/v0.py

@@ -253,7 +253,7 @@ def render(self, design: npt.NDArray, *, open_window: bool = False) -> Figure:
         size = len(design) + 1
 
         fig = plt.figure()
-        ax = fig.add_subplot(111, projection="3d")
+        ax: Any = fig.add_subplot(111, projection="3d")
         x, y, z = np.indices((size + 1, size + 1, size + 1)) / size  # Normalize to [0,1]
         # Define which voxels to plot
         threshold = 0.7