feat(beams2d): record the optimization trajectory on each OptiStep

engibench/problems/beams2d/v0.py

@@ -226,12 +226,9 @@ def optimize(
             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
+            # The design this step was evaluated at, taken before the overhang
+            # filter below rebinds xPrint to the next one.
+            design = np.array(xPrint)
 
             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)
@@ -252,6 +249,40 @@ def optimize(
                 xnew.reshape(base_config.nelx * base_config.nely, 1) - x.reshape(base_config.nelx * base_config.nely, 1),
                 np.inf,
             )
+
+            # Record the current state in optisteps_history, now that the move
+            # this design led to is known.
+            #
+            # x_sensitivities is the filtered objective sensitivity alone, one
+            # value per design variable, so that it has the same shape as the
+            # design it belongs to. The volume sensitivity dv is deliberately not
+            # stacked alongside it: consumers flatten this field, so an extra
+            # channel doubles its length with nothing recording that it did, and
+            # photonics2d -- the other 2D problem that reports sensitivities --
+            # reports the objective gradient by itself.
+            #
+            # The move is measured in printed density, the space the recorded
+            # design is in, so that design + update is the next step's design;
+            # inner_opt also returns the raw density field, and differencing that
+            # instead would mix the two spaces. The objective delta needs the
+            # *next* step's objective, so it is filled in on the following pass,
+            # and the last step keeps None rather than costing an extra solve.
+            obj_values = np.array(c)
+            if optisteps_history:
+                previous = optisteps_history[-1]
+                previous.obj_values_update = obj_values - previous.obj_values
+            current_step = ExtendedOptiStep(
+                obj_values=obj_values,
+                step=loop,
+                x=design,
+                x_sensitivities=dc.copy(),
+                x_update=xPrint - design,
+            )
+            current_step.design = design
+            optisteps_history.append(current_step)
+
+            loop += 1
+
             x = deepcopy(xnew)
 
         return design_to_image(xPrint, base_config.nelx, base_config.nely), optisteps_history

engibench/problems/beams2d/v1.py

@@ -76,12 +76,9 @@ def optimize(
             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
+            # The design this step was evaluated at, taken before the overhang
+            # filter below rebinds xPrint to the next one.
+            design = np.array(xPrint)
 
             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)
@@ -102,6 +99,40 @@ def optimize(
                 xnew.reshape(base_config.nelx * base_config.nely, 1) - x.reshape(base_config.nelx * base_config.nely, 1),
                 np.inf,
             )
+
+            # Record the current state in optisteps_history, now that the move
+            # this design led to is known.
+            #
+            # x_sensitivities is the filtered objective sensitivity alone, one
+            # value per design variable, so that it has the same shape as the
+            # design it belongs to. The volume sensitivity dv is deliberately not
+            # stacked alongside it: consumers flatten this field, so an extra
+            # channel doubles its length with nothing recording that it did, and
+            # photonics2d -- the other 2D problem that reports sensitivities --
+            # reports the objective gradient by itself.
+            #
+            # The move is measured in printed density, the space the recorded
+            # design is in, so that design + update is the next step's design;
+            # inner_opt also returns the raw density field, and differencing that
+            # instead would mix the two spaces. The objective delta needs the
+            # *next* step's objective, so it is filled in on the following pass,
+            # and the last step keeps None rather than costing an extra solve.
+            obj_values = np.array(c)
+            if optisteps_history:
+                previous = optisteps_history[-1]
+                previous.obj_values_update = obj_values - previous.obj_values
+            current_step = ExtendedOptiStep(
+                obj_values=obj_values,
+                step=loop,
+                x=design,
+                x_sensitivities=dc.copy(),
+                x_update=xPrint - design,
+            )
+            current_step.design = design
+            optisteps_history.append(current_step)
+
+            loop += 1
+
             x = deepcopy(xnew)
 
         return design_to_image(xPrint, base_config.nelx, base_config.nely), optisteps_history

tests/test_beams2d.py

@@ -0,0 +1,79 @@
+"""Tests for the Beams2D problem.
+
+`optimize` records, for every step, the design it was evaluated at, the filtered
+sensitivities there, and the move the optimizer made from it. These tests pin each
+field to the step it belongs to.
+
+A small grid and few iterations keep this fast.
+"""
+
+from itertools import pairwise
+
+import numpy as np
+import pytest
+
+from engibench.problems.beams2d import Beams2D
+
+NELX = 20
+NELY = 10
+MAX_ITER = 4
+VOLFRAC = 0.5
+N_ELEMS = NELX * NELY
+
+
+@pytest.fixture(scope="module")
+def problem() -> Beams2D:
+    return Beams2D(seed=0, config={"nelx": NELX, "nely": NELY, "volfrac": VOLFRAC})
+
+
+@pytest.fixture(scope="module")
+def optimized(problem: Beams2D) -> tuple[np.ndarray, list]:
+    """Run optimize once and reuse across tests (the expensive step)."""
+    problem.reset(seed=0)
+    start = np.full(problem.design_space.shape, VOLFRAC, dtype=np.float64)
+    # max_iter goes through optimize rather than the constructor: it lives on
+    # Config, not SimulateConfig, so optimize rebuilds it from the default.
+    return problem.optimize(start, config={"max_iter": MAX_ITER})
+
+
+def test_history_is_numbered_from_zero(optimized: tuple) -> None:
+    _opt, history = optimized
+    assert 1 <= len(history) <= MAX_ITER
+    assert [step.step for step in history] == list(range(len(history)))
+
+
+def test_every_step_records_design_sensitivities_and_update(optimized: tuple) -> None:
+    _opt, history = optimized
+    for step in history:
+        assert step.design.shape == (N_ELEMS,)
+        # `x` is the core field every generic consumer reads; `design` is kept
+        # for the callers that already use it.
+        assert step.x is step.design
+        # One sensitivity per design variable, so it is shaped like the design.
+        assert step.x_sensitivities.shape == step.x.shape
+        assert step.x_update.shape == (N_ELEMS,)
+
+
+def test_sensitivities_are_nonpositive(optimized: tuple) -> None:
+    """Compliance falls as material is added, and the optimizer clips dc at zero."""
+    _opt, history = optimized
+    for step in history:
+        assert np.all(step.x_sensitivities <= 0.0)
+        assert np.all(np.isfinite(step.x_sensitivities))
+
+
+def test_update_is_the_move_to_the_next_design(optimized: tuple) -> None:
+    """x_update is the step the optimizer took, so it lands on the next design."""
+    _opt, history = optimized
+    for step, following in pairwise(history):
+        np.testing.assert_allclose(step.x + step.x_update, following.x, atol=1e-12)
+
+
+def test_objective_delta_is_filled_in_except_on_the_last_step(optimized: tuple) -> None:
+    """A step's delta needs the next step's objective, so the last one has none."""
+    _opt, history = optimized
+    if len(history) < 2:  # noqa: PLR2004 - a single step has no delta to check
+        pytest.skip("optimization converged in one step")
+    for step, following in pairwise(history):
+        np.testing.assert_allclose(step.obj_values_update, following.obj_values - step.obj_values)
+    assert history[-1].obj_values_update is None