Minor changes

engibench/problems/airfoil/v0.py

@@ -1,6 +1,18 @@
-"""Airfoil 2D 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.
+
+             .:-===+=+==-:
+     .==.                        .:-++=:....
+ .-:                                           .:--:::.
+-            Airfoil v.0                        :====--:-===
+:-                                    .:==:.
+   .-::.                     ::::-:.
+          ..::::----::::..
+
++-+-+-+-+-+-+-+-+-+
+|E|n|g|i|B|e|n|c|h|
++-+-+-+-+-+-+-+-+-+
 """
 
 from __future__ import annotations
@@ -35,13 +47,6 @@
 DesignType = dict[str, Any]
 
 
-@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."""
     # intersection: find t such that (p + t dp - q) x dq = 0 with 0 <= t <= 1
@@ -147,7 +152,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
@@ -395,13 +400,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,
@@ -522,12 +525,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: 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 (bool): If True, saves the rendered design to a file in the study directory.
 
         Returns:
             Any: The rendered design.
@@ -536,37 +540,85 @@ 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))
+
+        if open_window:
+            plt.show()
+        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: 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 (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:
+            plt.savefig(self.__local_study_dir + "/optisteps.png", dpi=300, bbox_inches="tight")
+        plt.close(fig)
         return fig, ax
 
-    def random_design(self) -> 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_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["train"]["initial_design"]), dtype=int)
-        initial_design = self.dataset["train"]["initial_design"][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
 
 
 if __name__ == "__main__":
+    # Initialize the problem
     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))
 
+    # 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
     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=True)

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_split: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
+        Args:
+            dataset_split (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_split]), dtype=int)
+        return np.array(self.dataset[dataset_split][design_key][rnd]), rnd
 
 
 if __name__ == "__main__":

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_split: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
+        Args:
+            dataset_split (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_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.

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_split: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random design.
 
+        Args:
+            dataset_split (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_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.

engibench/problems/photonics2d/v0.py

@@ -665,7 +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) -> tuple[npt.NDArray, int]:
+    def random_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.
 
         Can return a design with small random variations or a uniform design, or can pull
@@ -674,6 +676,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_split (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 +689,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_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."""

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_split: str = "train", design_key: str = "initial_design") -> tuple[npt.NDArray, int]:
         """Samples a valid random initial design.
 
+        Args:
+            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["train"]["initial_design"]), dtype=int)
+        rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split][design_key]), dtype=int)
 
-        return np.array(self.dataset["train"]["initial_design"][rnd]), rnd
+        return np.array(self.dataset[dataset_split][design_key][rnd]), rnd
 
     def reset(self, seed: int | None = None) -> None:
         """Reset the problem.