From 369a81cc4b4e6f2697aef228925b6845cd51476e Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Wed, 30 Apr 2025 13:49:56 -0400
Subject: [PATCH 01/20] removed closed curve constraint

---
 engibench/problems/airfoil/v0.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 90da86e9..438280e9 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -36,11 +36,6 @@
 
 
 @constraint(categories=IMPL)
-def is_closed(design: DesignType) -> None:
-    """Check if a curve is closed."""
-    curve = design["coords"]
-    assert curve[0] == curve[-1], "design: Curve is not closed"
-
 
 def self_intersect(curve: npt.NDArray[np.float64]) -> tuple[int, npt.NDArray[np.float64], npt.NDArray[np.float64]] | None:
     """Determines if two segments a and b intersect."""

From 4e3e71a4bb5c18ae4ee3a98b426cbfdc0bad3d4e Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Wed, 30 Apr 2025 17:57:41 -0400
Subject: [PATCH 02/20] Latest remote

---
 engibench/problems/airfoil/v0.py | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 438280e9..fa785164 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -1,6 +1,18 @@
-"""Airfoil 2D problem.
+r"""Airfoil problem.
 
 Filename convention is that folder paths do not end with /. For example, /path/to/folder is correct, but /path/to/folder/ is not.
+
+             .:-===+=+==-:
+     .==.                        .:-++=:....
+ .-:                                           .:--:::.
+-            Airfoil v.0                        :====--:-===
+:-                                    .:==:.
+   .-::.                     ::::-:.
+          ..::::----::::..
+
++-+-+-+-+-+-+-+-+-+
+|E|n|g|i|B|e|n|c|h|
++-+-+-+-+-+-+-+-+-+
 """
 
 from __future__ import annotations
@@ -142,7 +154,7 @@ class Airfoil(Problem[DesignType]):
             "angle_of_attack": spaces.Box(low=0.0, high=10.0, shape=(1,), dtype=np.float32),
         }
     )
-    design_constraints = (is_closed, does_not_self_intersect)
+    design_constraints = (does_not_self_intersect)
     dataset_id = "IDEALLab/airfoil_v0"
     container_id = "mdolab/public:u22-gcc-ompi-stable"
     __local_study_dir: str
@@ -390,13 +402,11 @@ def simulate(self, design: DesignType, config: dict[str, Any] | None = None, mpi
             "use_altitude": False,
             "output_dir": "'" + self.__docker_study_dir + "/output/'",
             "mesh_fname": "'" + self.__docker_study_dir + "/design.cgns'",
-            "task": "'analysis'",  # TODO(cashend): We can add the option to perform a polar analysis.
-            # https://github.com/IDEALLab/EngiBench/issues/15
+            "task": "'analysis'",
             **dict(self.conditions),
             **(config or {}),
         }
         self.__design_to_simulator_input(design, base_config)
-
         replace_template_values(
             self.__local_study_dir + "/airfoil_analysis.py",
             base_config,
@@ -561,6 +571,8 @@ def random_design(self) -> tuple[dict[str, Any], int]:
 
     print(problem.simulate(design, config=config, mpicores=8))
 
+    problem.reset(seed=1, cleanup=True)
+
     # Get design and conditions from the dataset, render design
     opt_design, optisteps_history = problem.optimize(design, config=config, mpicores=8)
     print(optisteps_history)

From 24509ae41579a96a5b38a9377aa9f3050b01ac5f Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 05:55:09 -0400
Subject: [PATCH 03/20] minor changes + optisteps viewer

---
 engibench/problems/airfoil/v0.py | 67 ++++++++++++++++++++++++++------
 1 file changed, 56 insertions(+), 11 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index fa785164..7cc4da53 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -527,12 +527,13 @@ def optimize(
 
         return {"coords": opt_coords, "angle_of_attack": starting_point["angle_of_attack"]}, optisteps_history
 
-    def render(self, design: DesignType, *, open_window: bool = False) -> Any:
+    def render(self, design: DesignType, *, open_window: bool = False, save_render: bool = False) -> Any:
         """Renders the design in a human-readable format.
 
         Args:
             design (dict): The design to render.
             open_window (bool): If True, opens a window with the rendered design.
+            save_render (bool): If True, saves the rendered design to a file in the study directory.
 
         Returns:
             Any: The rendered design.
@@ -541,21 +542,59 @@ def render(self, design: DesignType, *, open_window: bool = False) -> Any:
 
         fig, ax = plt.subplots()
         coords = design["coords"]
-
+        alpha = design["angle_of_attack"]
         ax.scatter(coords[0], coords[1], s=10, alpha=0.7)
-        plt.ylim(-0.15, 0.15)
+        ax.set_title(r"$\alpha$="+str(np.round(alpha,2)) + r"$^\circ$")
+        ax.axis("equal")
+        ax.axis("off")
+        ax.set_xlim((-0.005, 1.005))
+        ax.set_ylim((-0.125,0.125))
         if open_window:
             plt.show()
+        if save_render:
+            plt.savefig(self.__local_study_dir + "/airfoil.png", dpi=300, bbox_inches="tight")
+        plt.close(fig)
         return fig, ax
 
-    def random_design(self) -> tuple[dict[str, Any], int]:
+    def render_optisteps(self, optisteps_history: list[OptiStep], open_window: bool = False, save_render: bool = False) -> Any:
+        """Renders the optimization step history.
+
+        Args:
+            optisteps_history (list[OptiStep]): The optimization steps to render.
+            open_window (bool): If True, opens a window with the rendered design.
+            save_render (bool): If True, saves the rendered design to a file in the study directory.
+
+        Returns:
+            Any: Rendered optimization step history.
+        """
+        import matplotlib.pyplot as plt
+
+        fig, ax = plt.subplots()
+        steps = np.array([step.step for step in optisteps_history])
+        objectives = np.array([step.obj_values[0][0] for step in optisteps_history])
+        ax.plot(steps, objectives, label="Drag Coefficient")
+        ax.set_title("Optimization Steps")
+        ax.set_xlabel("Iteration")
+        ax.set_ylabel("Drag counts")
+        if open_window:
+            plt.show()
+        if save_render:
+            plt.savefig(self.__local_study_dir + "/optisteps.png", dpi=300, bbox_inches="tight")
+        plt.close(fig)
+        return fig, ax
+
+    def random_design(self, design_key="initial_design") -> tuple[dict[str, Any], int]:
         """Samples a valid random initial design.
 
+        Args:
+            design_key (str): The key to use for the design in the dataset.
+                Defaults to "initial_design".
+
         Returns:
             tuple[dict[str, Any], int]: The valid random design and the index of the design in the dataset.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset["train"]["initial_design"]), dtype=int)
-        initial_design = self.dataset["train"]["initial_design"][rnd]
+        rnd = self.np_random.integers(low=0, high=len(self.dataset["train"][design_key]), dtype=int)
+        initial_design = self.dataset["train"][design_key][rnd]
         return {"coords": np.array(initial_design["coords"]), "angle_of_attack": initial_design["angle_of_attack"]}, rnd
 
 
@@ -563,17 +602,23 @@ def random_design(self) -> tuple[dict[str, Any], int]:
     problem = Airfoil()
     problem.reset(seed=0, cleanup=True)
 
+    # Retrieve the dataset
     dataset = problem.dataset
-    # Get design and conditions from the dataset
-    # Print Dataset object keys
+
+    # Get random initial design and optimized conditions from the dataset + the index
     design, idx = problem.random_design()
+
+    # Get the config conditions from the dataset
     config = dataset["train"].select_columns(problem.conditions_keys)[idx]
 
+    # Simulate the design
     print(problem.simulate(design, config=config, mpicores=8))
 
-    problem.reset(seed=1, cleanup=True)
+    # Cleanup the study directory
+    problem.reset(seed=0, cleanup=False)
 
     # Get design and conditions from the dataset, render design
     opt_design, optisteps_history = problem.optimize(design, config=config, mpicores=8)
-    print(optisteps_history)
-    problem.render(opt_design, open_window=True)
+
+    # Render the final optimized design
+    problem.render(opt_design, open_window=False, save_render=True)

From 2c87523932d8028062db5c064246fa5d5ef9d056 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 05:59:23 -0400
Subject: [PATCH 04/20] naming

---
 engibench/problems/airfoil/v0.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 7cc4da53..79b4354e 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -527,13 +527,13 @@ def optimize(
 
         return {"coords": opt_coords, "angle_of_attack": starting_point["angle_of_attack"]}, optisteps_history
 
-    def render(self, design: DesignType, *, open_window: bool = False, save_render: bool = False) -> Any:
+    def render(self, design: DesignType, *, open_window: bool = False, save: bool = False) -> Any:
         """Renders the design in a human-readable format.
 
         Args:
             design (dict): The design to render.
             open_window (bool): If True, opens a window with the rendered design.
-            save_render (bool): If True, saves the rendered design to a file in the study directory.
+            save (bool): If True, saves the rendered design to a file in the study directory.
 
         Returns:
             Any: The rendered design.
@@ -551,18 +551,18 @@ def render(self, design: DesignType, *, open_window: bool = False, save_render:
         ax.set_ylim((-0.125,0.125))
         if open_window:
             plt.show()
-        if save_render:
+        if save:
             plt.savefig(self.__local_study_dir + "/airfoil.png", dpi=300, bbox_inches="tight")
         plt.close(fig)
         return fig, ax
 
-    def render_optisteps(self, optisteps_history: list[OptiStep], open_window: bool = False, save_render: bool = False) -> Any:
+    def render_optisteps(self, optisteps_history: list[OptiStep], open_window: bool = False, save: bool = False) -> Any:
         """Renders the optimization step history.
 
         Args:
             optisteps_history (list[OptiStep]): The optimization steps to render.
             open_window (bool): If True, opens a window with the rendered design.
-            save_render (bool): If True, saves the rendered design to a file in the study directory.
+            save (bool): If True, saves the rendered design to a file in the study directory.
 
         Returns:
             Any: Rendered optimization step history.
@@ -578,7 +578,7 @@ def render_optisteps(self, optisteps_history: list[OptiStep], open_window: bool
         ax.set_ylabel("Drag counts")
         if open_window:
             plt.show()
-        if save_render:
+        if save:
             plt.savefig(self.__local_study_dir + "/optisteps.png", dpi=300, bbox_inches="tight")
         plt.close(fig)
         return fig, ax
@@ -621,4 +621,4 @@ def random_design(self, design_key="initial_design") -> tuple[dict[str, Any], in
     opt_design, optisteps_history = problem.optimize(design, config=config, mpicores=8)
 
     # Render the final optimized design
-    problem.render(opt_design, open_window=False, save_render=True)
+    problem.render(opt_design, open_window=False, save=True)

From 323d8b06f63dc2f45baaac576b488567c2382085 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 06:01:49 -0400
Subject: [PATCH 05/20] comments

---
 engibench/problems/airfoil/v0.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 79b4354e..01b2b468 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -583,22 +583,24 @@ def render_optisteps(self, optisteps_history: list[OptiStep], open_window: bool
         plt.close(fig)
         return fig, ax
 
-    def random_design(self, design_key="initial_design") -> tuple[dict[str, Any], int]:
+    def random_design(self, dataset_key="train", design_key="initial_design") -> tuple[dict[str, Any], int]:
         """Samples a valid random initial design.
 
         Args:
+            dataset_key (str): The key to use for the dataset. Defaults to "train".
             design_key (str): The key to use for the design in the dataset.
                 Defaults to "initial_design".
 
         Returns:
             tuple[dict[str, Any], int]: The valid random design and the index of the design in the dataset.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset["train"][design_key]), dtype=int)
-        initial_design = self.dataset["train"][design_key][rnd]
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key][design_key]), dtype=int)
+        initial_design = self.dataset[dataset_key][design_key][rnd]
         return {"coords": np.array(initial_design["coords"]), "angle_of_attack": initial_design["angle_of_attack"]}, rnd
 
 
 if __name__ == "__main__":
+    # Initialize the problem
     problem = Airfoil()
     problem.reset(seed=0, cleanup=True)
 
@@ -614,7 +616,7 @@ def random_design(self, design_key="initial_design") -> tuple[dict[str, Any], in
     # Simulate the design
     print(problem.simulate(design, config=config, mpicores=8))
 
-    # Cleanup the study directory
+    # Cleanup the study directory; will delete the previous contents from simulate in this case
     problem.reset(seed=0, cleanup=False)
 
     # Get design and conditions from the dataset, render design

From 1be3352910d5ef9c1d6b3de579fa2cdcbefc527b Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 06:05:23 -0400
Subject: [PATCH 06/20] beams2d random_design keys

---
 engibench/problems/beams2d/v0.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/engibench/problems/beams2d/v0.py b/engibench/problems/beams2d/v0.py
index c85078d8..5fead245 100644
--- a/engibench/problems/beams2d/v0.py
+++ b/engibench/problems/beams2d/v0.py
@@ -300,16 +300,20 @@ def render(self, design: np.ndarray, *, open_window: bool = False) -> Any:
             plt.show()
         return fig, ax
 
-    def random_design(self) -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_key="train", design_key="initial_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
+        Args:
+            dataset_key (str): The key for the dataset to sample from.
+            design_key (str): The key for the design to sample from.
+
         Returns:
             Tuple of:
                 np.ndarray: The valid random design.
                 int: The random index selected.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset["train"]), dtype=int)
-        return np.array(self.dataset["train"]["optimal_design"][rnd]), rnd
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key]), dtype=int)
+        return np.array(self.dataset[dataset_key][design_key][rnd]), rnd
 
 
 if __name__ == "__main__":

From 9dc796737b8fb6494cb321cb59cf42d1e6c5f625 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 06:07:29 -0400
Subject: [PATCH 07/20] random_design types

---
 engibench/problems/airfoil/v0.py | 2 +-
 engibench/problems/beams2d/v0.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 01b2b468..59ce62d5 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -583,7 +583,7 @@ def render_optisteps(self, optisteps_history: list[OptiStep], open_window: bool
         plt.close(fig)
         return fig, ax
 
-    def random_design(self, dataset_key="train", design_key="initial_design") -> tuple[dict[str, Any], int]:
+    def random_design(self, dataset_key: str ="train", design_key: str ="initial_design") -> tuple[dict[str, Any], int]:
         """Samples a valid random initial design.
 
         Args:
diff --git a/engibench/problems/beams2d/v0.py b/engibench/problems/beams2d/v0.py
index 5fead245..1164a26a 100644
--- a/engibench/problems/beams2d/v0.py
+++ b/engibench/problems/beams2d/v0.py
@@ -300,7 +300,7 @@ def render(self, design: np.ndarray, *, open_window: bool = False) -> Any:
             plt.show()
         return fig, ax
 
-    def random_design(self, dataset_key="train", design_key="initial_design") -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_key: str = "train", design_key: str = "initial_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
         Args:

From 864d944ab6e939df865572b30ca54cf8fd4b9808 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 13:46:48 -0400
Subject: [PATCH 08/20] minor fixes

---
 engibench/problems/airfoil/v0.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 59ce62d5..96cf2126 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -583,7 +583,7 @@ def render_optisteps(self, optisteps_history: list[OptiStep], open_window: bool
         plt.close(fig)
         return fig, ax
 
-    def random_design(self, dataset_key: str ="train", design_key: str ="initial_design") -> tuple[dict[str, Any], int]:
+    def random_design(self, dataset_key: str = "train", design_key: str = "initial_design") -> tuple[dict[str, Any], int]:
         """Samples a valid random initial design.
 
         Args:

From cd333b3d5163a00395c8844b1059e7feb25b04ef Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 13:47:13 -0400
Subject: [PATCH 09/20] heatconduction random_design

---
 engibench/problems/heatconduction2d/v0.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/engibench/problems/heatconduction2d/v0.py b/engibench/problems/heatconduction2d/v0.py
index 1d612bfb..b54b2daf 100644
--- a/engibench/problems/heatconduction2d/v0.py
+++ b/engibench/problems/heatconduction2d/v0.py
@@ -231,16 +231,20 @@ def initialize_design(self, volume: float | None = None, resolution: int | None
 
         return np.load(design_file)
 
-    def random_design(self) -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_key: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
+        Args:
+            dataset_key (str): The key for the dataset to sample from.
+            design_key (str): The key for the design to sample from.
+
         Returns:
             Tuple of:
                 np.ndarray: The valid random design.
                 int: The random index selected.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset["train"]["optimal_design"]))
-        return np.array(self.dataset["train"]["optimal_design"][rnd]), int(rnd)
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key][design_key]))
+        return np.array(self.dataset[dataset_key][design_key][rnd]), int(rnd)
 
     def render(self, design: npt.NDArray, *, open_window: bool = False) -> Any:
         """Renders the design in a human-readable format.

From af62fd9c19b3f6b3092e07a5b30a6a88e6cca507 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 13:55:48 -0400
Subject: [PATCH 10/20] photonics random_design

---
 engibench/problems/photonics2d/v0.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/engibench/problems/photonics2d/v0.py b/engibench/problems/photonics2d/v0.py
index 7f424b82..fb15095b 100644
--- a/engibench/problems/photonics2d/v0.py
+++ b/engibench/problems/photonics2d/v0.py
@@ -665,7 +665,8 @@ def _randomized_noise_field_design(self, noise: float = 0.001, blur: int = 0) ->
 
         return rho_start.astype(np.float32)
 
-    def random_design(self, noise: float | None = None, blur: int = 0) -> tuple[npt.NDArray, int]:
+
+    def random_design(self, noise: float | None = None, blur: int = 0, dataset_key: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Generates a random initial design.
 
         Can return a design with small random variations or a uniform design, or can pull
@@ -674,6 +675,8 @@ def random_design(self, noise: float | None = None, blur: int = 0) -> tuple[npt.
         Args:
             noise (float|None): If None, pull from dataset. If float, use that as the noise level.
             blur (int): The amount of pixel blurring to apply to random field. Only active if noise is used.
+            dataset_key (str): The key for the dataset to sample from.
+            design_key (str): The key for the design to sample from.
 
         Returns:
             tuple[npt.NDArray, int]: The starting design array (rho) and an integer (0).
@@ -685,8 +688,8 @@ def random_design(self, noise: float | None = None, blur: int = 0) -> tuple[npt.
         if noise is not None:
             rho_start = self._randomized_noise_field_design(noise=noise, blur=blur)
             return rho_start, 0
-        rnd = self.np_random.integers(low=0, high=len(self.dataset["train"]), dtype=int)
-        return np.array(self.dataset["train"]["optimal_design"][rnd]), rnd
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key]), dtype=int)
+        return np.array(self.dataset[dataset_key][design_key][rnd]), rnd
 
     def reset(self, seed: int | None = None, **kwargs) -> None:
         """Resets the problem, which in this case, is just the random seed."""

From e34f7083d3930fbacc1ac92bcc6d3ed9b092630e Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 13:58:44 -0400
Subject: [PATCH 11/20] power electronics random_design

---
 engibench/problems/power_electronics/v0.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/engibench/problems/power_electronics/v0.py b/engibench/problems/power_electronics/v0.py
index 2bed3dc8..2345be2f 100644
--- a/engibench/problems/power_electronics/v0.py
+++ b/engibench/problems/power_electronics/v0.py
@@ -174,15 +174,19 @@ def render(self, design: npt.NDArray, *, open_window: bool = False) -> None:  #
         nx.draw(G, pos, with_labels=True, node_color=node_colors, node_size=200, font_size=10)
         plt.show()
 
-    def random_design(self) -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_key: str = "train", design_key: str = "initial_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random initial design.
 
+        Args:
+            dataset_key (str): The key for the dataset to sample from.
+            design_key (str): The key for the design to sample from.
+
         Returns:
             DesignType: The valid random design.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset["train"]["initial_design"]), dtype=int)
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key][design_key]), dtype=int)
 
-        return np.array(self.dataset["train"]["initial_design"][rnd]), rnd
+        return np.array(self.dataset[dataset_key][design_key][rnd]), rnd
 
     def reset(self, seed: int | None = None) -> None:
         """Reset the problem.

From 46a4b167c2fb0e5331a0e35ef059b15783c5eea5 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 14:04:40 -0400
Subject: [PATCH 12/20] double check beams2d

---
 engibench/problems/beams2d/v0.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/engibench/problems/beams2d/v0.py b/engibench/problems/beams2d/v0.py
index 1164a26a..393f26be 100644
--- a/engibench/problems/beams2d/v0.py
+++ b/engibench/problems/beams2d/v0.py
@@ -300,7 +300,7 @@ def render(self, design: np.ndarray, *, open_window: bool = False) -> Any:
             plt.show()
         return fig, ax
 
-    def random_design(self, dataset_key: str = "train", design_key: str = "initial_design") -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_key: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
         Args:

From 2c333119df77aab2e6061e520078d07055d83b9d Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 14:05:09 -0400
Subject: [PATCH 13/20] double check heatconduction3d

---
 engibench/problems/heatconduction3d/v0.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/engibench/problems/heatconduction3d/v0.py b/engibench/problems/heatconduction3d/v0.py
index 94f4cb29..500d6d19 100644
--- a/engibench/problems/heatconduction3d/v0.py
+++ b/engibench/problems/heatconduction3d/v0.py
@@ -233,16 +233,20 @@ def initialize_design(self, volume: float | None = None, resolution: int | None
 
         return np.load(design_file)
 
-    def random_design(self) -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_key: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
+        Args:
+            dataset_key (str): The key for the dataset to sample from.
+            design_key (str): The key for the design to sample from.
+
         Returns:
             Tuple of:
                 np.ndarray: The valid random design.
                 int: The random index selected.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset["train"]["optimal_design"]))
-        return np.array(self.dataset["train"]["optimal_design"][rnd]), int(rnd)
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key][design_key]), dtype=int)
+        return np.array(self.dataset[dataset_key][design_key][rnd]), rnd
 
     def render(self, design: npt.NDArray, *, open_window: bool = False) -> Any:
         """Renders the design in a human-readable format.

From 6a3f22bc676ff8e841ea573b5ba395711682b844 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 14:10:05 -0400
Subject: [PATCH 14/20] versioning consistency

---
 engibench/problems/airfoil/v0.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 96cf2126..d260b5bd 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -5,7 +5,7 @@
              .:-===+=+==-:
      .==.                        .:-++=:....
  .-:                                           .:--:::.
--            Airfoil v.0                        :====--:-===
+-            Airfoil v.0.0.1                    :====--:-===
 :-                                    .:==:.
    .-::.                     ::::-:.
           ..::::----::::..
@@ -119,7 +119,7 @@ class Airfoil(Problem[DesignType]):
     ## Dataset
     The dataset linked to this problem is hosted on the [Hugging Face Datasets Hub](https://huggingface.co/datasets/IDEALLab/airfoil_v0).
 
-    ### v0
+    ### v0.0.1
 
     #### Fields
     The dataset contains optimal design, conditions, objectives and these additional fields:

From 0db48cf223cb77f4b6bb3387754accb5258ae6bd Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 14:21:23 -0400
Subject: [PATCH 15/20] ruff

---
 engibench/problems/airfoil/v0.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index d260b5bd..59ab73e6 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -556,7 +556,7 @@ def render(self, design: DesignType, *, open_window: bool = False, save: bool =
         plt.close(fig)
         return fig, ax
 
-    def render_optisteps(self, optisteps_history: list[OptiStep], open_window: bool = False, save: bool = False) -> Any:
+    def render_optisteps(self, optisteps_history: list[OptiStep], *, open_window: bool = False, save: bool = False) -> Any:
         """Renders the optimization step history.
 
         Args:

From 450da70aed82ae621e893a25cab1dd0bf7615eb7 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Thu, 1 May 2025 14:23:03 -0400
Subject: [PATCH 16/20] ruff format

---
 engibench/problems/airfoil/v0.py     | 7 +++----
 engibench/problems/photonics2d/v0.py | 5 +++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 59ab73e6..53ca8356 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -48,7 +48,6 @@
 
 
 @constraint(categories=IMPL)
-
 def self_intersect(curve: npt.NDArray[np.float64]) -> tuple[int, npt.NDArray[np.float64], npt.NDArray[np.float64]] | None:
     """Determines if two segments a and b intersect."""
     # intersection: find t such that (p + t dp - q) x dq = 0 with 0 <= t <= 1
@@ -154,7 +153,7 @@ class Airfoil(Problem[DesignType]):
             "angle_of_attack": spaces.Box(low=0.0, high=10.0, shape=(1,), dtype=np.float32),
         }
     )
-    design_constraints = (does_not_self_intersect)
+    design_constraints = does_not_self_intersect
     dataset_id = "IDEALLab/airfoil_v0"
     container_id = "mdolab/public:u22-gcc-ompi-stable"
     __local_study_dir: str
@@ -544,11 +543,11 @@ def render(self, design: DesignType, *, open_window: bool = False, save: bool =
         coords = design["coords"]
         alpha = design["angle_of_attack"]
         ax.scatter(coords[0], coords[1], s=10, alpha=0.7)
-        ax.set_title(r"$\alpha$="+str(np.round(alpha,2)) + r"$^\circ$")
+        ax.set_title(r"$\alpha$=" + str(np.round(alpha, 2)) + r"$^\circ$")
         ax.axis("equal")
         ax.axis("off")
         ax.set_xlim((-0.005, 1.005))
-        ax.set_ylim((-0.125,0.125))
+        ax.set_ylim((-0.125, 0.125))
         if open_window:
             plt.show()
         if save:
diff --git a/engibench/problems/photonics2d/v0.py b/engibench/problems/photonics2d/v0.py
index fb15095b..8c91b4ec 100644
--- a/engibench/problems/photonics2d/v0.py
+++ b/engibench/problems/photonics2d/v0.py
@@ -665,8 +665,9 @@ def _randomized_noise_field_design(self, noise: float = 0.001, blur: int = 0) ->
 
         return rho_start.astype(np.float32)
 
-
-    def random_design(self, noise: float | None = None, blur: int = 0, dataset_key: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
+    def random_design(
+        self, noise: float | None = None, blur: int = 0, dataset_key: str = "train", design_key: str = "optimal_design"
+    ) -> tuple[npt.NDArray, int]:
         """Generates a random initial design.
 
         Can return a design with small random variations or a uniform design, or can pull

From 5f1d0f4f6d88a09bf07889d48623d65e360a7d64 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Fri, 2 May 2025 09:55:58 -0400
Subject: [PATCH 17/20] dataset_split airfoil

---
 engibench/problems/airfoil/v0.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 53ca8356..83cbdf8f 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -5,7 +5,7 @@
              .:-===+=+==-:
      .==.                        .:-++=:....
  .-:                                           .:--:::.
--            Airfoil v.0.0.1                    :====--:-===
+-            Airfoil v.0                        :====--:-===
 :-                                    .:==:.
    .-::.                     ::::-:.
           ..::::----::::..
@@ -47,7 +47,7 @@
 DesignType = dict[str, Any]
 
 
-@constraint(categories=IMPL)
+
 def self_intersect(curve: npt.NDArray[np.float64]) -> tuple[int, npt.NDArray[np.float64], npt.NDArray[np.float64]] | None:
     """Determines if two segments a and b intersect."""
     # intersection: find t such that (p + t dp - q) x dq = 0 with 0 <= t <= 1
@@ -118,7 +118,7 @@ class Airfoil(Problem[DesignType]):
     ## Dataset
     The dataset linked to this problem is hosted on the [Hugging Face Datasets Hub](https://huggingface.co/datasets/IDEALLab/airfoil_v0).
 
-    ### v0.0.1
+    ### v0
 
     #### Fields
     The dataset contains optimal design, conditions, objectives and these additional fields:
@@ -547,7 +547,7 @@ def render(self, design: DesignType, *, open_window: bool = False, save: bool =
         ax.axis("equal")
         ax.axis("off")
         ax.set_xlim((-0.005, 1.005))
-        ax.set_ylim((-0.125, 0.125))
+
         if open_window:
             plt.show()
         if save:
@@ -582,19 +582,19 @@ def render_optisteps(self, optisteps_history: list[OptiStep], *, open_window: bo
         plt.close(fig)
         return fig, ax
 
-    def random_design(self, dataset_key: str = "train", design_key: str = "initial_design") -> tuple[dict[str, Any], int]:
+    def random_design(self, dataset_split: str = "train", design_key: str = "initial_design") -> tuple[dict[str, Any], int]:
         """Samples a valid random initial design.
 
         Args:
-            dataset_key (str): The key to use for the dataset. Defaults to "train".
+            dataset_split (str): The key to use for the dataset. Defaults to "train".
             design_key (str): The key to use for the design in the dataset.
                 Defaults to "initial_design".
 
         Returns:
             tuple[dict[str, Any], int]: The valid random design and the index of the design in the dataset.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key][design_key]), dtype=int)
-        initial_design = self.dataset[dataset_key][design_key][rnd]
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split][design_key]), dtype=int)
+        initial_design = self.dataset[dataset_split][design_key][rnd]
         return {"coords": np.array(initial_design["coords"]), "angle_of_attack": initial_design["angle_of_attack"]}, rnd
 
 

From 3f40ce5e49abb36722038e3e3360fe854a076a47 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Fri, 2 May 2025 09:56:16 -0400
Subject: [PATCH 18/20] dataset_split other problems

---
 engibench/problems/beams2d/v0.py           | 8 ++++----
 engibench/problems/heatconduction2d/v0.py  | 8 ++++----
 engibench/problems/heatconduction3d/v0.py  | 8 ++++----
 engibench/problems/photonics2d/v0.py       | 8 ++++----
 engibench/problems/power_electronics/v0.py | 8 ++++----
 5 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/engibench/problems/beams2d/v0.py b/engibench/problems/beams2d/v0.py
index 393f26be..177bc5f1 100644
--- a/engibench/problems/beams2d/v0.py
+++ b/engibench/problems/beams2d/v0.py
@@ -300,11 +300,11 @@ def render(self, design: np.ndarray, *, open_window: bool = False) -> Any:
             plt.show()
         return fig, ax
 
-    def random_design(self, dataset_key: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_split: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
         Args:
-            dataset_key (str): The key for the dataset to sample from.
+            dataset_split (str): The key for the dataset to sample from.
             design_key (str): The key for the design to sample from.
 
         Returns:
@@ -312,8 +312,8 @@ def random_design(self, dataset_key: str = "train", design_key: str = "optimal_d
                 np.ndarray: The valid random design.
                 int: The random index selected.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key]), dtype=int)
-        return np.array(self.dataset[dataset_key][design_key][rnd]), rnd
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split]), dtype=int)
+        return np.array(self.dataset[dataset_split][design_key][rnd]), rnd
 
 
 if __name__ == "__main__":
diff --git a/engibench/problems/heatconduction2d/v0.py b/engibench/problems/heatconduction2d/v0.py
index b54b2daf..d63875af 100644
--- a/engibench/problems/heatconduction2d/v0.py
+++ b/engibench/problems/heatconduction2d/v0.py
@@ -231,11 +231,11 @@ def initialize_design(self, volume: float | None = None, resolution: int | None
 
         return np.load(design_file)
 
-    def random_design(self, dataset_key: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_split: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
         Args:
-            dataset_key (str): The key for the dataset to sample from.
+            dataset_split (str): The key for the dataset to sample from.
             design_key (str): The key for the design to sample from.
 
         Returns:
@@ -243,8 +243,8 @@ def random_design(self, dataset_key: str = "train", design_key: str = "optimal_d
                 np.ndarray: The valid random design.
                 int: The random index selected.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key][design_key]))
-        return np.array(self.dataset[dataset_key][design_key][rnd]), int(rnd)
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split][design_key]))
+        return np.array(self.dataset[dataset_split][design_key][rnd]), int(rnd)
 
     def render(self, design: npt.NDArray, *, open_window: bool = False) -> Any:
         """Renders the design in a human-readable format.
diff --git a/engibench/problems/heatconduction3d/v0.py b/engibench/problems/heatconduction3d/v0.py
index 500d6d19..23ffabfc 100644
--- a/engibench/problems/heatconduction3d/v0.py
+++ b/engibench/problems/heatconduction3d/v0.py
@@ -233,11 +233,11 @@ def initialize_design(self, volume: float | None = None, resolution: int | None
 
         return np.load(design_file)
 
-    def random_design(self, dataset_key: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_split: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
         Args:
-            dataset_key (str): The key for the dataset to sample from.
+            dataset_split (str): The key for the dataset to sample from.
             design_key (str): The key for the design to sample from.
 
         Returns:
@@ -245,8 +245,8 @@ def random_design(self, dataset_key: str = "train", design_key: str = "optimal_d
                 np.ndarray: The valid random design.
                 int: The random index selected.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key][design_key]), dtype=int)
-        return np.array(self.dataset[dataset_key][design_key][rnd]), rnd
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split][design_key]), dtype=int)
+        return np.array(self.dataset[dataset_split][design_key][rnd]), rnd
 
     def render(self, design: npt.NDArray, *, open_window: bool = False) -> Any:
         """Renders the design in a human-readable format.
diff --git a/engibench/problems/photonics2d/v0.py b/engibench/problems/photonics2d/v0.py
index 8c91b4ec..b57d2c06 100644
--- a/engibench/problems/photonics2d/v0.py
+++ b/engibench/problems/photonics2d/v0.py
@@ -666,7 +666,7 @@ def _randomized_noise_field_design(self, noise: float = 0.001, blur: int = 0) ->
         return rho_start.astype(np.float32)
 
     def random_design(
-        self, noise: float | None = None, blur: int = 0, dataset_key: str = "train", design_key: str = "optimal_design"
+        self, noise: float | None = None, blur: int = 0, dataset_split: str = "train", design_key: str = "optimal_design"
     ) -> tuple[npt.NDArray, int]:
         """Generates a random initial design.
 
@@ -676,7 +676,7 @@ def random_design(
         Args:
             noise (float|None): If None, pull from dataset. If float, use that as the noise level.
             blur (int): The amount of pixel blurring to apply to random field. Only active if noise is used.
-            dataset_key (str): The key for the dataset to sample from.
+            dataset_split (str): The key for the dataset to sample from.
             design_key (str): The key for the design to sample from.
 
         Returns:
@@ -689,8 +689,8 @@ def random_design(
         if noise is not None:
             rho_start = self._randomized_noise_field_design(noise=noise, blur=blur)
             return rho_start, 0
-        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key]), dtype=int)
-        return np.array(self.dataset[dataset_key][design_key][rnd]), rnd
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split]), dtype=int)
+        return np.array(self.dataset[dataset_split][design_key][rnd]), rnd
 
     def reset(self, seed: int | None = None, **kwargs) -> None:
         """Resets the problem, which in this case, is just the random seed."""
diff --git a/engibench/problems/power_electronics/v0.py b/engibench/problems/power_electronics/v0.py
index 2345be2f..b95ea28f 100644
--- a/engibench/problems/power_electronics/v0.py
+++ b/engibench/problems/power_electronics/v0.py
@@ -174,19 +174,19 @@ def render(self, design: npt.NDArray, *, open_window: bool = False) -> None:  #
         nx.draw(G, pos, with_labels=True, node_color=node_colors, node_size=200, font_size=10)
         plt.show()
 
-    def random_design(self, dataset_key: str = "train", design_key: str = "initial_design") -> tuple[npt.NDArray, int]:
+    def random_design(self, dataset_split: str = "train", design_key: str = "initial_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random initial design.
 
         Args:
-            dataset_key (str): The key for the dataset to sample from.
+            dataset_split (str): The key for the dataset to sample from.
             design_key (str): The key for the design to sample from.
 
         Returns:
             DesignType: The valid random design.
         """
-        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_key][design_key]), dtype=int)
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split][design_key]), dtype=int)
 
-        return np.array(self.dataset[dataset_key][design_key][rnd]), rnd
+        return np.array(self.dataset[dataset_split][design_key][rnd]), rnd
 
     def reset(self, seed: int | None = None) -> None:
         """Reset the problem.

From a5d196b3edbe0c708adc229702ca3b8d6a929bab Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Fri, 2 May 2025 11:01:10 -0400
Subject: [PATCH 19/20] minor changes

---
 engibench/problems/airfoil/v0.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 83cbdf8f..29c0340c 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -46,8 +46,6 @@
 
 DesignType = dict[str, Any]
 
-
-
 def self_intersect(curve: npt.NDArray[np.float64]) -> tuple[int, npt.NDArray[np.float64], npt.NDArray[np.float64]] | None:
     """Determines if two segments a and b intersect."""
     # intersection: find t such that (p + t dp - q) x dq = 0 with 0 <= t <= 1
@@ -153,7 +151,7 @@ class Airfoil(Problem[DesignType]):
             "angle_of_attack": spaces.Box(low=0.0, high=10.0, shape=(1,), dtype=np.float32),
         }
     )
-    design_constraints = does_not_self_intersect
+    design_constraints = (does_not_self_intersect,)
     dataset_id = "IDEALLab/airfoil_v0"
     container_id = "mdolab/public:u22-gcc-ompi-stable"
     __local_study_dir: str

From e32ff01362faa5e5395d76ccd9b9dcaa872de035 Mon Sep 17 00:00:00 2001
From: cashend <cashen.diniz123@gmail.com>
Date: Fri, 2 May 2025 11:06:35 -0400
Subject: [PATCH 20/20] ruff

---
 engibench/problems/airfoil/v0.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/engibench/problems/airfoil/v0.py b/engibench/problems/airfoil/v0.py
index 29c0340c..bb0fef13 100644
--- a/engibench/problems/airfoil/v0.py
+++ b/engibench/problems/airfoil/v0.py
@@ -1,4 +1,4 @@
-r"""Airfoil problem.
+"""Airfoil problem.
 
 Filename convention is that folder paths do not end with /. For example, /path/to/folder is correct, but /path/to/folder/ is not.
 
@@ -46,6 +46,7 @@
 
 DesignType = dict[str, Any]
 
+
 def self_intersect(curve: npt.NDArray[np.float64]) -> tuple[int, npt.NDArray[np.float64], npt.NDArray[np.float64]] | None:
     """Determines if two segments a and b intersect."""
     # intersection: find t such that (p + t dp - q) x dq = 0 with 0 <= t <= 1