Tests: Limit max iterations in optimize

engibench/problems/heatconduction2d/templates/optimize_heat_conduction_2d.py

@@ -50,7 +50,7 @@
 # NN: Grid size
 # vol_f: Volume fraction
 # width: Adiabatic boundary width
-NN, vol_f, width, output_path = cast_argv(int, float, float, str)
+NN, vol_f, width, max_iter, output_path = cast_argv(int, float, float, int, str)
 # Load Initial Design Data
 image = np_array_from_stdin()
 
@@ -199,7 +199,12 @@ def length(self):
 # Define filename for IPOPT log
 log_filename = os.path.join(output_dir, f"solution_V={vol_f}_w={width}.txt")
 # Set optimization solver parameters
-solver_params = {"acceptable_tol": 1.0e-3, "maximum_iterations": 100, "file_print_level": 5, "output_file": log_filename}
+solver_params = {
+    "acceptable_tol": 1.0e-3,
+    "maximum_iterations": max_iter,
+    "file_print_level": 5,
+    "output_file": log_filename,
+}
 solver = IPOPTSolver(problem, parameters=solver_params)
 # -------------------------------
 # Store and Save Results
@@ -239,6 +244,7 @@ def length(self):
 xdmf_filename.write(a_opt)
 print("v={vol_f}")
 print("w={width}")
-np.savez(output_path, design=RES_OPTults, OptiStep=np.array(objective_values))
+# `[:, None]` to make the output array 2D:
+np.savez(output_path, design=RES_OPTults, OptiStep=np.array(objective_values)[:, None])
 for f in glob.glob("/home/fenics/shared/templates/RES_OPT/TEMP*"):
     os.remove(f)

engibench/problems/heatconduction2d/v0.py

@@ -72,6 +72,8 @@ class Config(Conditions):
             int, bounded(lower=1).category(THEORY), bounded(lower=10, upper=1000).warning().category(IMPL)
         ] = 101
         """Resolution of the design space for the initialization"""
+        max_iter: int = 100
+        """Maximum number of iterations for the solver in `optimize()`."""
 
     config: Config
 
@@ -137,6 +139,7 @@ def optimize(
         config = config or {}
         volume = config.get("volume", self.config.volume)
         length = config.get("length", self.config.length)
+        max_iter = config.get("max_iter", self.config.max_iter)
         resolution = config.get("resolution", self.config.resolution)
         if starting_point is None:
             starting_point = self.initialize_design(volume, resolution)
@@ -145,7 +148,7 @@ def optimize(
             run_container_script(
                 self.container_id,
                 Path(__file__).parent / "templates" / "optimize_heat_conduction_2d.py",
-                args=(resolution - 1, volume, length),
+                args=(resolution - 1, volume, length, max_iter),
                 stdin=np_array_to_bytes(starting_point),
                 output_path=f"RES_OPT/OUTPUT={volume}_w={length}.npz",
             )

engibench/problems/heatconduction3d/templates/optimize_heat_conduction_3d.py

@@ -51,7 +51,7 @@
 # NN: Grid size
 # vol_f: Volume fraction
 # width: Adiabatic boundary width
-NN, vol_f, width, output_path = cast_argv(int, float, float, str)
+NN, vol_f, width, max_iter, output_path = cast_argv(int, float, float, int, str)
 # Load Initial Design Data
 image = np_array_from_stdin()
 
@@ -224,7 +224,12 @@ def length(self):
 # Define filename for IPOPT log
 log_filename = os.path.join(output_dir, f"solution_V={vol_f}_w={width}.txt")
 # Set optimization solver parameters
-solver_params = {"acceptable_tol": 1.0e-100, "maximum_iterations": 100, "file_print_level": 5, "output_file": log_filename}
+solver_params = {
+    "acceptable_tol": 1.0e-100,
+    "maximum_iterations": max_iter,
+    "file_print_level": 5,
+    "output_file": log_filename,
+}
 solver = IPOPTSolver(problem, parameters=solver_params)
 # -------------------------------
 # Store and Save Results
@@ -265,6 +270,7 @@ def length(self):
 xdmf_filename.write(a_opt)
 print(f"v={vol_f}")
 print(f"w={width}")
-np.savez(output_path, design=RES_OPTults, OptiStep=np.array(objective_values))
+# `[:, None]` to make the output array 2D:
+np.savez(output_path, design=RES_OPTults, OptiStep=np.array(objective_values)[:, None])
 for f in glob.glob("/home/fenics/shared/templates/RES_OPT/TEMP*"):
     os.remove(f)

engibench/problems/heatconduction3d/v0.py

@@ -68,6 +68,8 @@ class Config(Conditions):
             int, bounded(lower=1).category(THEORY), bounded(lower=10, upper=1000).warning().category(IMPL)
         ] = 51
         """Resolution of the design space"""
+        max_iter: int = 100
+        """Maximum number of iterations for the solver in `optimize()`."""
 
     config: Config
 
@@ -134,14 +136,15 @@ def optimize(
         volume = config.get("volume", self.config.volume)
         area = config.get("area", self.config.area)
         resolution = config.get("resolution", self.config.resolution)
+        max_iter = config.get("max_iter", self.config.max_iter)
         if starting_point is None:
             starting_point = self.initialize_design(volume, resolution)
 
         output = np.load(
             run_container_script(
                 self.container_id,
                 Path(__file__).parent / "templates" / "optimize_heat_conduction_3d.py",
-                args=(resolution - 1, volume, area),
+                args=(resolution - 1, volume, area, max_iter),
                 stdin=cli.np_array_to_bytes(starting_point),
                 output_path=f"RES_OPT/OUTPUT={volume}_w={area}.npz",
             )

engibench/problems/photonics2d/v0.py

@@ -76,7 +76,6 @@ class Photonics2D(Problem[npt.NDArray]):
     _space_slice = 8  # Extra space for source/probe slices (pixels)
 
     # Defaults for the optimization parameters
-    _num_optimization_steps_default = 200  # Default number of optimization steps
     _step_size_default = 1e-1  # Default step size for Adam optimizer
     _eta_default = 0.5
     _num_projections_default = 1
@@ -127,6 +126,8 @@ class Config(Conditions):
             bounded(lower=110, upper=300).warning().category(IMPL),
         ] = 120
         """number of grid cells in y"""
+        num_optimization_steps: int = 200
+        """Maximum number of optimization steps."""
 
     design_space = spaces.Box(low=0.0, high=1.0, shape=(Config.num_elems_x, Config.num_elems_y), dtype=np.float64)
 
@@ -307,7 +308,7 @@ def optimize(  # noqa: PLR0915
         num_elems_y = self.num_elems_y
         # Pull out optimization parameters from conditions
         # Parameters specific to optimization
-        num_optimization_steps = conditions.get("num_optimization_steps", self._num_optimization_steps_default)
+        num_optimization_steps = conditions["num_optimization_steps"]
         step_size = conditions.get("step_size", self._step_size_default)
         penalty_weight = conditions.get("penalty_weight", self._penalty_weight_default)
         self._eta = conditions.get("eta", self._eta_default)

engibench/problems/power_electronics/v0.py

@@ -57,6 +57,7 @@ class Conditions:
     dataset_id = "IDEALLab/power_electronics_v0"
     container_id = None
     config: Config
+    Config = Config
 
     def __init__(
         self,

engibench/problems/thermoelastic2d/model/fea_model.py

@@ -27,15 +27,17 @@
 class FeaModel:
     """Finite Element Analysis (FEA) model for coupled 2D thermoelastic topology optimization."""
 
-    def __init__(self, *, plot: bool = False, eval_only: bool | None = False) -> None:
+    def __init__(self, *, plot: bool = False, eval_only: bool | None = False, max_iter: int = MAX_ITERATIONS) -> None:
         """Instantiates a new model for the thermoelastic 2D problem.
 
         Args:
             plot: (bool, optional): If True, the updated design will be plotted at each iteration.
             eval_only: (bool, optional): If True, the model will only evaluate the design and return the objective values.
+            max_iter: Maximal number of iterations for the `run` method.
         """
         self.plot = plot
         self.eval_only = eval_only
+        self.max_iter = max_iter
 
     def get_initial_design(self, volume_fraction: float, nelx: int, nely: int) -> np.ndarray:
         """Generates the initial design variable field for the optimization process.
@@ -340,7 +342,7 @@ def run(self, bcs: dict[str, Any], x_init: np.ndarray | None = None) -> dict[str
                 f" It.: {iterr:4d} Obj.: {f0val:10.4f} Vol.: {np.sum(x) / (nelx * nely):6.3f} ch.: {change:6.3f} || t_forward:{t_forward:6.3f} + t_sensitivity:{t_sensitivity:6.3f} + t_sens_calc:{t_sensitivity_calc:6.3f} + t_mma: {t_mma:6.3f} = {t_total:6.3f}"
             )
 
-            if iterr > MAX_ITERATIONS:
+            if iterr > self.max_iter:
                 break
 
         print("Optimization finished...")

engibench/problems/thermoelastic2d/v0.py

@@ -19,6 +19,7 @@
 from engibench.core import ObjectiveDirection
 from engibench.core import OptiStep
 from engibench.core import Problem
+from engibench.problems.thermoelastic2d.model import fea_model
 from engibench.problems.thermoelastic2d.model.fea_model import FeaModel
 from engibench.problems.thermoelastic2d.utils import get_res_bounds
 from engibench.problems.thermoelastic2d.utils import indices_to_binary_matrix
@@ -84,6 +85,8 @@ class Config(Conditions):
 
         nelx: ClassVar[Annotated[int, bounded(lower=1).category(THEORY)]] = NELX
         nely: ClassVar[Annotated[int, bounded(lower=1).category(THEORY)]] = NELX
+        max_iter: int = fea_model.MAX_ITERATIONS
+        """Maximal number of iterations for optimize."""
 
         @constraint
         @staticmethod
@@ -142,7 +145,8 @@ def optimize(
         """
         boundary_dict = dataclasses.asdict(self.conditions)
         boundary_dict.update({k: v for k, v in (config or {}).items() if k in boundary_dict})
-        results = FeaModel(plot=False, eval_only=False).run(boundary_dict, x_init=starting_point)
+        max_iter = (config or {}).get("max_iter", self.Config.max_iter)
+        results = FeaModel(plot=False, eval_only=False, max_iter=max_iter).run(boundary_dict, x_init=starting_point)
         design = np.array(results["design"]).astype(np.float32)
         opti_steps = results["opti_steps"]
         return design, opti_steps

engibench/problems/thermoelastic3d/model/fem_model.py

@@ -25,15 +25,17 @@
 class FeaModel3D:
     """Finite Element Analysis (FEA) model for coupled 3D thermoelastic topology optimization."""
 
-    def __init__(self, *, plot: bool = False, eval_only: bool = False) -> None:
+    def __init__(self, *, plot: bool = False, eval_only: bool = False, max_iter: int = MAX_ITERATIONS) -> None:
         """Instantiates a new 3D thermoelastic model.
 
         Args:
             plot: If True, you can hook in your own plotting / volume rendering each iteration.
             eval_only: If True, evaluate the given design once and return objective components only.
+            max_iter: Maximal number of iterations for the `run` method.
         """
         self.plot = plot
         self.eval_only = eval_only
+        self.max_iter = max_iter
 
     def get_initial_design(self, volume_fraction: float, nelx: int, nely: int, nelz: int) -> np.ndarray:
         """Generates the initial design variable field for the optimization process.
@@ -404,7 +406,7 @@ def node_id(ix: int, iy: int, iz: int) -> int:
                 f"|| t_forward:{t_forward:6.3f} + t_adj:{t_adjoints:6.3f} + t_sens:{t_sens:6.3f} + t_mma:{t_mma:6.3f} = {t_total:6.3f}"
             )
 
-            if iterr > MAX_ITERATIONS:
+            if iterr > self.max_iter:
                 break
 
         print("3D optimization finished.")

engibench/problems/thermoelastic3d/v0.py

@@ -17,6 +17,7 @@
 from engibench.core import ObjectiveDirection
 from engibench.core import OptiStep
 from engibench.core import Problem
+from engibench.problems.thermoelastic3d.model import fem_model
 from engibench.problems.thermoelastic3d.model.fem_model import FeaModel3D
 
 NELX = NELY = NELZ = 16
@@ -95,6 +96,8 @@ class Config(Conditions):
         nelx: Annotated[int, bounded(lower=1).category(THEORY)] = NELX
         nely: Annotated[int, bounded(lower=1).category(THEORY)] = NELY
         nelz: Annotated[int, bounded(lower=1).category(THEORY)] = NELZ
+        max_iter: int = fem_model.MAX_ITERATIONS
+        """Maximal number of iterations for optimize."""
 
         @constraint
         @staticmethod
@@ -164,7 +167,9 @@ def optimize(
         """
         boundary_dict = dataclasses.asdict(self.conditions)
         boundary_dict.update({k: v for k, v in (config or {}).items() if k in boundary_dict})
-        results = FeaModel3D(plot=False, eval_only=False).run(boundary_dict, x_init=starting_point)
+        results = FeaModel3D(
+            plot=False, eval_only=False, max_iter=(config or {}).get("max_iter", self.Config.max_iter)
+        ).run(boundary_dict, x_init=starting_point)
         design = np.array(results["design"]).astype(np.float32)
         opti_steps = results["opti_steps"]
         return design, opti_steps

tests/test_problem_implementations.py

@@ -2,8 +2,9 @@
 
 import dataclasses
 import inspect
+import os
 import sys
-from typing import get_args, get_origin
+from typing import Any, get_args, get_origin
 
 import gymnasium
 from gymnasium import spaces
@@ -13,9 +14,6 @@
 from engibench import Problem
 from engibench.utils.all_problems import BUILTIN_PROBLEMS
 
-PYTHON_PROBLEMS = [p for p in BUILTIN_PROBLEMS.values() if p.container_id is None]
-CONTAINER_PROBLEMS = [p for p in BUILTIN_PROBLEMS.values() if p.container_id is not None]
-
 
 @pytest.mark.parametrize("problem_class", BUILTIN_PROBLEMS.values())
 def test_problem_impl(problem_class: type[Problem]) -> None:
@@ -88,10 +86,12 @@ def test_problem_impl(problem_class: type[Problem]) -> None:
     print(f"Done testing {problem_class.__name__}.")
 
 
-@pytest.mark.parametrize(
-    "problem_class",
-    PYTHON_PROBLEMS + (CONTAINER_PROBLEMS if sys.platform.startswith("linux") else []),
-)
+def problem_id(problem_class: type[Problem]) -> str:
+    id_, _ = problem_class.__module__.removeprefix("engibench.problems.").split(".", 1)
+    return id_
+
+
+@pytest.mark.parametrize("problem_class", BUILTIN_PROBLEMS.values())
 def test_python_problem_impl(problem_class: type[Problem]) -> None:
     """Check that all problems defined in Python files respect the API.
 
@@ -100,6 +100,8 @@ def test_python_problem_impl(problem_class: type[Problem]) -> None:
     2. The optimization produces valid designs within the design space
     3. The optimization history contains valid objective values
     """
+    if problem_class.container_id is not None and not sys.platform.startswith("linux"):
+        pytest.skip(f"Skipping containerized problem {problem_class.__name__} on non-linux platform")
     print(f"Testing optimization and simulation for {problem_class.__name__}...")
     # Initialize problem and get a random design
     problem = problem_class(seed=1)
@@ -121,15 +123,20 @@ def test_python_problem_impl(problem_class: type[Problem]) -> None:
     # Test optimization outputs
     print(f"Optimizing {problem_class.__name__}...")
     # Skip optimization test for power electronics, airfoil, and heat conduction problems
-    if (
-        problem_class.__module__.startswith("engibench.problems.power_electronics")
-        or problem_class.__module__.startswith("engibench.problems.airfoil")
-        or problem_class.__module__.startswith("engibench.problems.heatconduction")
-    ):
+    if problem_id(problem_class) == "airfoil":
         print(f"Skipping optimization test for {problem_class.__name__}")
         return
     problem.reset(seed=1)
-    optimal_design, history = problem.optimize(starting_point=design)
+    default_max_iter = 10
+    max_iter = os.environ.get("ENGIBENCH_MAX_ITER", default_max_iter)
+    max_iter_config: dict[str, Any] = {
+        key: max_iter for key in ("max_iter", "num_optimization_steps") if hasattr(problem_class.Config, key)
+    }
+    try:
+        optimal_design, history = problem.optimize(starting_point=design, config=max_iter_config)
+    except NotImplementedError:
+        print("Problem class {problem_class.__name__} does not implement optimize - Skipping optimize")
+        return
     if isinstance(problem.design_space, spaces.Box):
         assert np.all(optimal_design >= problem.design_space.low), (
             f"Problem {problem_class.__name__}: The optimal design should be within the design space."