Refactor GA Output Directory Code into Reusable Base 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,8 +201,8 @@
     "\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",
+    " - `vocs.txt`: The VOCS object so that the objectives, constraints, and decision variables are retained alongside the data\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",
     "\n",
@@ -288,20 +287,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 +342,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "xopt-dev",
    "language": "python",
    "name": "python3"
   },
@@ -371,7 +356,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.10"
+   "version": "3.12.0"
   }
  },
  "nbformat": 4,

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/base.py

@@ -0,0 +1,186 @@
+from pydantic import Field, field_validator
+import logging
+import os
+import pandas as pd
+import time
+
+from ..checkpoints import CheckpointMixin
+from ..deduplicated import DeduplicatedGeneratorBase
+
+POPULATION_METADATA_COLUMNS = [
+    "xopt_generation",
+    "xopt_candidate_idx",
+    "xopt_runtime",
+    "xopt_error",
+]
+
+
+class GAGeneratorBase(CheckpointMixin, DeduplicatedGeneratorBase):
+    """
+    Base class for genetic algorithm generators which write output and checkpoints.
+
+    Handles the output directory, log file, and periodic checkpointing on behalf of
+    subclasses. Subclasses call `end_generation` once each time a generation is
+    completed and everything else is taken care of.
+
+    Nothing is written to disk until the generator is used, so building or
+    deserializing one never touches the filesystem.
+
+    Parameters
+    ----------
+    output_dir : str or os.PathLike, optional
+        Directory to save algorithm state and population history, or None to write
+        nothing. Stored as a string, unexpanded; environment variables and "~" are
+        expanded when the path is used. If the directory already contains data, a
+        number is appended to avoid overwriting it.
+    checkpoint_freq : int, default=1
+        Frequency (in generations) at which checkpoints are saved. Set to -1 to
+        disable checkpointing.
+    log_level : int
+        Level of log messages written to "log.txt".
+
+    Attributes
+    ----------
+    expanded_output_dir : str or None
+        `output_dir` with environment variables and "~" expanded. All file writes go
+        here.
+    """
+
+    output_dir: str | None = None
+    checkpoint_freq: int = Field(
+        1,
+        description="How often (in generations) to save checkpoints (set to -1 to disable)",
+    )
+    log_level: int = Field(
+        logging.INFO, description="Log message level output to log.txt"
+    )
+    _output_prepared: bool = (
+        False  # Whether the output directory has been resolved and created
+    )
+
+    @field_validator("output_dir", mode="before")
+    @classmethod
+    def validate_output_dir(cls, value):
+        """Accept any os.PathLike, storing it as a string."""
+        if isinstance(value, os.PathLike):
+            return os.fspath(value)
+        return value
+
+    @property
+    def expanded_output_dir(self) -> str | None:
+        """Output directory with environment variables and "~" expanded."""
+        if self.output_dir is None:
+            return None
+        return os.path.expanduser(os.path.expandvars(self.output_dir))
+
+    def model_post_init(self, context):
+        # Get a unique logger per object. Naming it after the concrete class keeps
+        # records propagating through that class's module logger.
+        self._logger = logging.getLogger(
+            f"{type(self).__module__}.{type(self).__name__}.{id(self)}"
+        )
+        self._logger.setLevel(self.log_level)
+
+    def _prepare_output(self) -> None:
+        """
+        Resolve and create the output directory and begin logging to file.
+
+        Repeated calls do nothing. If the requested directory already holds data, a
+        number is appended and `output_dir` is updated to the path actually used.
+        """
+        if (self.output_dir is None) or self._output_prepared:
+            return
+
+        # Check if directory exists and do collision avoidance. Resolve into a local
+        # so the field is only assigned once, since assignment revalidates the model.
+        # Suffixes are applied to the unexpanded path, but tested against the expanded one.
+        requested = self.output_dir
+        counter = 2
+        output_dir = requested
+        expanded = self.expanded_output_dir
+        while os.path.exists(expanded) and os.listdir(expanded):
+            output_dir = f"{requested}_{counter}"
+            expanded = os.path.expanduser(os.path.expandvars(output_dir))
+            counter += 1
+        if output_dir != requested:
+            self._logger.info(
+                f'detected existing output_dir "{requested}" and corrected '
+                f'to "{output_dir}" to avoid overwriting'
+            )
+        self.output_dir = output_dir
+
+        # We are now setup
+        os.makedirs(self.expanded_output_dir, exist_ok=True)
+        self._output_prepared = True
+
+        # Set up file logging
+        log_file_path = os.path.join(self.expanded_output_dir, "log.txt")
+        file_handler = logging.FileHandler(log_file_path, mode="w")
+        file_handler.setLevel(self.log_level)
+        file_handler.setFormatter(
+            logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+        )
+        self._logger.addHandler(file_handler)
+        self._logger.info(f"routing log output to file: {log_file_path}")
+
+        # Record the problem definition alongside the data
+        # Note: this is necessary to include in output for users running analysis on the results
+        # ie to plot Pareto front, you need to know the names and direction of the objectives
+        with open(os.path.join(self.expanded_output_dir, "vocs.txt"), "w") as f:
+            f.write(self.vocs.model_dump_json())
+
+    def end_generation(self, generation_index: int, population: list[dict]) -> None:
+        """
+        Record a completed generation, writing output and checkpoints as configured.
+
+        Parameters
+        ----------
+        generation_index : int
+            Index of the generation which was just completed.
+        population : list of dict
+            The individuals making up the completed population.
+        """
+        self._prepare_output()
+        if self.output_dir is None:
+            return
+        output_dir = self.expanded_output_dir
+        save_start_t = time.perf_counter()
+
+        # Save all Xopt data
+        self.data.to_csv(os.path.join(output_dir, "data.csv"), index=False)
+
+        # Construct the DataFrame for this population
+        pop_df = pd.DataFrame(population)
+        pop_df["xopt_generation"] = generation_index
+
+        # Normalize the columns in the DataFrame
+        # Avoid schema changing part way through optimization so we can write CSV in append mode
+        pop_df = pop_df.reindex(
+            columns=self.vocs.all_names + POPULATION_METADATA_COLUMNS
+        )
+
+        # Write population DataFrame to file
+        csv_path = os.path.join(output_dir, "populations.csv")
+        pop_df.to_csv(
+            csv_path, index=False, mode="a", header=not os.path.isfile(csv_path)
+        )
+        self._logger.debug(
+            f'saved optimization data to "{output_dir}" '
+            f"in {1000 * (time.perf_counter() - save_start_t):.2f}ms"
+        )
+
+        # Save a checkpoint if one is due
+        if self.checkpoint_freq > 0 and (generation_index % self.checkpoint_freq == 0):
+            checkpoint_path = self._save_checkpoint(
+                os.path.join(output_dir, "checkpoints")
+            )
+            self._logger.debug(f'saved checkpoint file "{checkpoint_path}"')
+
+    def close_log_file(self):
+        """
+        Closes out the log file (if used)
+        """
+        for handler in list(self._logger.handlers):
+            if isinstance(handler, logging.FileHandler):
+                handler.close()
+            self._logger.removeHandler(handler)