Refactor Generator Checkpointing into Reusable Class

docs/examples/ga/nsga2/nsga2_python.ipynb

@@ -14,7 +14,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import json\n",
     "import logging\n",
     "import matplotlib.pyplot as plt\n",
     "import os\n",
@@ -28,7 +27,7 @@
     "    SimulatedBinaryCrossover,\n",
     ")\n",
     "from xopt.resources.test_functions.zdt import construct_zdt\n",
-    "from xopt import Xopt, Evaluator, VOCS"
+    "from xopt import Xopt, Evaluator"
    ]
   },
   {
@@ -202,7 +201,6 @@
     "\n",
     "The output files are the following.\n",
     " - `data.csv`: All data evaluated during the optimization\n",
-    " - `vocs.txt`: The VOCS object so that the objectives, constraints, decision variables are retained alongside the data\n",
     " - `populations.csv`: Each population is written here with a column `xopt_generation` to distinguish which generation the row belongs to\n",
     " - `checkpoints`: This generator periodically saves its full state to timestamped files in this directory\n",
     " - `log.txt`: Log output from the generator is recorded to this file\n",
@@ -288,20 +286,6 @@
     "df.head()"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Read the VOCS object back in. This can be used for data analysis / restarting optimizations\n",
-    "with open(os.path.join(my_xopt.generator.output_dir, \"vocs.txt\")) as f:\n",
-    "    vocs_from_file = VOCS(**json.load(f))\n",
-    "\n",
-    "# Show the objectives\n",
-    "vocs_from_file.objectives"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -357,7 +341,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "xopt-dev",
    "language": "python",
    "name": "python3"
   },
@@ -371,7 +355,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.10"
+   "version": "3.12.0"
   }
  },
  "nbformat": 4,

docs/examples/ga/nsga2/yaml_interface/index.md

@@ -255,7 +255,6 @@ Navigate to the output directory and observe the files there.
 - `populations.csv`: Each completed population is recorded to this file
 - `data.csv`: Contains all evaluated inviduals
 - `log.txt`: A record of all log messages the genreator emitted during its run
-- `vocs.txt`: A copy of the variable, objectives, and constraints (VOCs) definitions
 - `checkpoints/`:  This directory contains checkpoint files which are used with the `checkpoint_file` key of the generator to restart an optimization.
 
 

docs/examples/ga/nsga2/yaml_interface/nsga2_yaml.ipynb

@@ -52,7 +52,6 @@
     "- `populations.csv`: Each completed population is recorded to this file\n",
     "- `data.csv`: Contains all evaluated inviduals\n",
     "- `log.txt`: A record of all log messages the genreator emitted during its run\n",
-    "- `vocs.txt`: A copy of the variable, objectives, and constraints (VOCs) definitions\n",
     "- `checkpoints/`: This directory contains checkpoint files which are used with the `checkpoint_file` key of the generator to restart an optimization."
    ]
   },

xopt/generators/checkpoints.py

@@ -0,0 +1,123 @@
+from datetime import datetime
+from pydantic import BaseModel, Field, model_validator
+import json
+import os
+
+from ..vocs import VOCS
+
+
+class CheckpointMixin(BaseModel):
+    """
+    Mix-in class adding checkpoint saving and loading to a generator.
+
+    Checkpoints are written to a caller-supplied directory. The VOCS object is
+    serialized into the checkpoint itself. Legacy checkpoints which predate this
+    instead carry the VOCS object in a "vocs.txt" file one level above the
+    directory holding the checkpoints and are still supported when loading.
+
+    The host class must provide a ``vocs`` attribute and a ``to_json`` method.
+    Writing of the checkpoint file is the responsibility of the concrete class
+    as it will be implementation dependent.
+
+    Parameters
+    ----------
+    checkpoint_file : str, optional
+        Path to checkpoint file to load from. If provided, the generator will be
+        initialized from the checkpoint state. User-specified parameters will
+        override checkpoint values.
+    """
+
+    checkpoint_file: str | None = Field(
+        None, description="Path to checkpoint file to load from", exclude=True
+    )
+
+    @staticmethod
+    def _load_checkpoint_data(fname: str) -> dict:
+        """
+        Internal function to load generator data from checkpoint file as well as VOCS object.
+
+        Parameters
+        ----------
+        fname : str
+            Path to the checkpoint file
+
+        Returns
+        -------
+        dict
+            Dictionary containing VOCS and checkpoint data
+        """
+        # Load the checkpoint
+        with open(fname) as f:
+            checkpoint_data = json.load(f)
+
+        if "vocs" in checkpoint_data:
+            return checkpoint_data
+
+        # Legacy checkpoints w/o VOCS
+        vocs_fname = os.path.join(os.path.dirname(fname), "../vocs.txt")
+        if not os.path.exists(vocs_fname):
+            raise ValueError(
+                f'Checkpoint "{fname}" does not contain a VOCS object and no '
+                f'VOCS file was found at "{vocs_fname}".'
+            )
+
+        with open(vocs_fname) as f:
+            vocs = VOCS(**json.load(f))
+
+        return {"vocs": vocs, **checkpoint_data}
+
+    @model_validator(mode="before")
+    @classmethod
+    def load_from_checkpoint(cls, values):
+        """
+        Load from checkpoint file if checkpoint_file is provided.
+        """
+        # Case when a checkpoint file has been supplied
+        if isinstance(values, dict) and "checkpoint_file" in values:
+            checkpoint_file = values.pop("checkpoint_file")
+            if checkpoint_file is not None:
+                # Load checkpoint data
+                checkpoint_data = cls._load_checkpoint_data(checkpoint_file)
+
+                # Merge with user data precedence
+                merged_data = {**checkpoint_data, **values}
+                return merged_data
+
+        # No checkpoint
+        return values
+
+    def _save_checkpoint(self, path: str | os.PathLike) -> str:
+        """
+        Write a checkpoint of the generator state to disk.
+
+        Parameters
+        ----------
+        path : str or os.PathLike
+            Directory into which the checkpoint file is written. Created if it
+            does not already exist.
+
+        Returns
+        -------
+        str
+            Path to the checkpoint file which was written.
+        """
+        # Set up the output directory
+        os.makedirs(path, exist_ok=True)
+
+        # Create a base filename
+        base_checkpoint_filename = datetime.now().strftime("%Y%m%d_%H%M%S")
+        checkpoint_path = os.path.join(path, f"{base_checkpoint_filename}_1.txt")
+
+        # Check if file exists and increment counter until we find a free filename
+        counter = 2
+        while os.path.exists(checkpoint_path):
+            checkpoint_path = os.path.join(
+                path, f"{base_checkpoint_filename}_{counter}.txt"
+            )
+            counter += 1
+
+        # Now we have a unique filename
+        with open(checkpoint_path, "w") as f:
+            f.write(self.to_json())
+
+        return checkpoint_path

xopt/generators/ga/nsga2.py

@@ -1,8 +1,6 @@
-from datetime import datetime
 from itertools import chain
 from pydantic import Field, Discriminator, model_validator
 from typing import Annotated
-import json
 import logging
 import numpy as np
 import os
@@ -14,6 +12,7 @@
 from ...errors import DataError
 from ...generator import StateOwner
 from ...vocs import VOCS
+from ..checkpoints import CheckpointMixin
 from ..deduplicated import DeduplicatedGeneratorBase
 from ..utils import fast_dominated_argsort
 from .operators import (
@@ -312,7 +311,7 @@ def generate_candidates_from_population(
 ########################################################################################################################
 
 
-class NSGA2Generator(DeduplicatedGeneratorBase, StateOwner):
+class NSGA2Generator(CheckpointMixin, DeduplicatedGeneratorBase, StateOwner):
     """
     Non-dominated Sorting Genetic Algorithm II (NSGA-II) generator.  Implements the NSGA-II algorithm
     for multi-objective optimization as described in [1]. This generator accomdates user selected mutation
@@ -373,11 +372,6 @@ class NSGA2Generator(DeduplicatedGeneratorBase, StateOwner):
     supports_constraints: bool = True
     supports_single_objective: bool = True
 
-    # Checkpoint loading
-    checkpoint_file: str | None = Field(
-        None, description="Path to checkpoint file to load from", exclude=True
-    )
-
     population_size: int = Field(50, description="Population size")
     crossover_operator: Annotated[
         (
@@ -436,58 +430,6 @@ def model_post_init(self, context):
         self._logger = logging.getLogger(f"{__name__}.NSGA2Generator.{id(self)}")
         self._logger.setLevel(self.log_level)
 
-    @staticmethod
-    def _load_checkpoint_data(fname: str) -> dict:
-        """
-        Internal function to load generator data from checkpoint file as well as VOCS object.
-
-        Parameters
-        ----------
-        fname : str
-            Path to the checkpoint file
-
-        Returns
-        -------
-        dict
-            Dictionary containing VOCS and checkpoint data
-        """
-        # Load the VOCS object
-        vocs_fname = os.path.join(os.path.dirname(fname), "../vocs.txt")
-        if not os.path.exists(vocs_fname):
-            raise ValueError(
-                f'Could not load VOCS file at "{vocs_fname}". Complete NSGA2Generator '
-                "output directory is required for loading from checkpoint."
-            )
-
-        with open(vocs_fname) as f:
-            vocs = VOCS(**json.load(f))
-
-        # Load the checkpoint
-        with open(fname) as f:
-            checkpoint_data = json.load(f)
-
-        return {"vocs": vocs, **checkpoint_data}
-
-    @model_validator(mode="before")
-    @classmethod
-    def load_from_checkpoint(cls, values):
-        """
-        Load from checkpoint file if checkpoint_file is provided.
-        """
-        # Case when a checkpoint file has been supplied
-        if isinstance(values, dict) and "checkpoint_file" in values:
-            checkpoint_file = values.pop("checkpoint_file")
-            if checkpoint_file is not None:
-                # Load checkpoint data
-                checkpoint_data = cls._load_checkpoint_data(checkpoint_file)
-
-                # Merge with user data precedence
-                merged_data = {**checkpoint_data, **values}
-                return merged_data
-
-        # No checkpoint
-        return values
-
     @model_validator(mode="after")
     def vocs_compatible(self):
         """
@@ -650,8 +592,6 @@ def add_data(self, new_data: pd.DataFrame):
 
                 # Save all Xopt data
                 self.data.to_csv(os.path.join(self.output_dir, "data.csv"), index=False)
-                with open(os.path.join(self.output_dir, "vocs.txt"), "w") as f:
-                    json.dump(self.vocs.dict(), f)
 
                 # Construct the DataFrame for this population
                 pop_df = pd.DataFrame(self.pop)
@@ -682,37 +622,10 @@ def add_data(self, new_data: pd.DataFrame):
                 if self.checkpoint_freq > 0 and (
                     self.n_generations % self.checkpoint_freq == 0
                 ):
-                    self._save_checkpoint()
-
-    def _save_checkpoint(self):
-        # Confirm we are ready to save checkpoint
-        if self.output_dir is None:
-            raise ValueError("Cannot save checkpoint without an output directory")
-        self.ensure_output_dir_setup()
-
-        # Create a base filename
-        os.makedirs(os.path.join(self.output_dir, "checkpoints"), exist_ok=True)
-        base_checkpoint_filename = datetime.now().strftime("%Y%m%d_%H%M%S")
-        checkpoint_path = os.path.join(
-            self.output_dir,
-            "checkpoints",
-            f"{base_checkpoint_filename}_1.txt",
-        )
-
-        # Check if file exists and increment counter until we find a free filename
-        counter = 2
-        while os.path.exists(checkpoint_path):
-            checkpoint_path = os.path.join(
-                self.output_dir,
-                "checkpoints",
-                f"{base_checkpoint_filename}_{counter}.txt",
-            )
-            counter += 1
-
-        # Now we have a unique filename
-        with open(checkpoint_path, "w") as f:
-            f.write(self.to_json())
-        self._logger.debug(f'saved checkpoint file "{checkpoint_path}"')
+                    checkpoint_path = self._save_checkpoint(
+                        os.path.join(self.output_dir, "checkpoints")
+                    )
+                    self._logger.debug(f'saved checkpoint file "{checkpoint_path}"')
 
     def set_data(self, data):
         self.data = data

xopt/tests/generators/ga/test_nsga2.py

@@ -76,7 +76,6 @@ def test_nsga2_output_data():
         # Verify that the data files are created
         assert os.path.exists(os.path.join(output_dir, "data.csv"))
         assert os.path.exists(os.path.join(output_dir, "populations.csv"))
-        assert os.path.exists(os.path.join(output_dir, "vocs.txt"))
         assert os.path.exists(os.path.join(output_dir, "log.txt"))
 
         # Read the data file and check its contents
@@ -110,11 +109,6 @@ def test_nsga2_output_data():
         # Check that the populations file contains the expected columns
         assert "xopt_generation" in pop_df.columns
 
-        # Check that the VOCS file contains valid JSON
-        with open(os.path.join(output_dir, "vocs.txt"), "r") as f:
-            vocs_dict = json.load(f)
-            VOCS(**vocs_dict)
-
         # Verify that the log file exists and has content
         with open(os.path.join(output_dir, "log.txt"), "r") as f:
             log_content = f.read()