diff --git a/docs/examples/ga/nsga2/nsga2_python.ipynb b/docs/examples/ga/nsga2/nsga2_python.ipynb index 08d5ee2ac..89576db48 100644 --- a/docs/examples/ga/nsga2/nsga2_python.ipynb +++ b/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, diff --git a/docs/examples/ga/nsga2/yaml_interface/index.md b/docs/examples/ga/nsga2/yaml_interface/index.md index 2e293e6de..91d4ff825 100644 --- a/docs/examples/ga/nsga2/yaml_interface/index.md +++ b/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. diff --git a/docs/examples/ga/nsga2/yaml_interface/nsga2_yaml.ipynb b/docs/examples/ga/nsga2/yaml_interface/nsga2_yaml.ipynb index 1f0635d36..63da9e817 100644 --- a/docs/examples/ga/nsga2/yaml_interface/nsga2_yaml.ipynb +++ b/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." ] }, diff --git a/xopt/generators/checkpoints.py b/xopt/generators/checkpoints.py new file mode 100644 index 000000000..c189e8366 --- /dev/null +++ b/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 diff --git a/xopt/generators/ga/nsga2.py b/xopt/generators/ga/nsga2.py index ba716cccb..e38f3b605 100644 --- a/xopt/generators/ga/nsga2.py +++ b/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 diff --git a/xopt/tests/generators/ga/test_nsga2.py b/xopt/tests/generators/ga/test_nsga2.py index 635d93dc2..850b819d0 100644 --- a/xopt/tests/generators/ga/test_nsga2.py +++ b/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() diff --git a/xopt/tests/generators/test_checkpoints.py b/xopt/tests/generators/test_checkpoints.py new file mode 100644 index 000000000..380343207 --- /dev/null +++ b/xopt/tests/generators/test_checkpoints.py @@ -0,0 +1,122 @@ +from copy import deepcopy +from datetime import datetime +import json +import os + +import pytest + +from xopt.generators.checkpoints import CheckpointMixin +from xopt.generators.random import RandomGenerator +from xopt.resources.testing import TEST_VOCS_BASE +from xopt.vocs import VOCS + + +class CheckpointingRandomGenerator(CheckpointMixin, RandomGenerator): + """Minimal host for the checkpoint mixin with a field to round trip.""" + + counter: int = 0 + + +def parse_checkpoint_filename(filename: str) -> tuple[datetime, int]: + """Split a checkpoint filename into its timestamp and deduplication index.""" + base, index = filename.rsplit("_", 1) + return datetime.strptime(base, "%Y%m%d_%H%M%S"), int(index.split(".")[0]) + + +def make_legacy_checkpoint(checkpoint_path: str) -> None: + """Rewrite a checkpoint into the legacy layout with VOCS held in "vocs.txt".""" + with open(checkpoint_path) as f: + checkpoint_data = json.load(f) + vocs = checkpoint_data.pop("vocs") + + with open(checkpoint_path, "w") as f: + json.dump(checkpoint_data, f) + legacy_dir = os.path.dirname(os.path.dirname(checkpoint_path)) + with open(os.path.join(legacy_dir, "vocs.txt"), "w") as f: + json.dump(vocs, f) + + +def test_save_checkpoint_layout(tmp_path): + generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE)) + checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints") + + # The checkpoint is written directly into the supplied directory + assert os.path.dirname(checkpoint_path) == str(tmp_path / "checkpoints") + assert os.listdir(tmp_path / "checkpoints") == [os.path.basename(checkpoint_path)] + assert os.listdir(tmp_path) == ["checkpoints"] + + # Filename follows the timestamp plus deduplication index scheme + _, index = parse_checkpoint_filename(os.path.basename(checkpoint_path)) + assert index == 1 + + # The checkpoint holds valid JSON and carries the VOCS object itself + with open(checkpoint_path) as f: + checkpoint_data = json.load(f) + assert "counter" in checkpoint_data + assert VOCS(**checkpoint_data["vocs"]) == generator.vocs + + +def test_checkpoint_round_trip(tmp_path): + generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE), counter=17) + checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints") + + # VOCS comes from the checkpoint itself, not from the user + reloaded = CheckpointingRandomGenerator(checkpoint_file=checkpoint_path) + assert reloaded.counter == 17 + assert reloaded.vocs == generator.vocs + + # The path used to load is not carried into the reloaded generator's state + assert reloaded.checkpoint_file is None + assert "checkpoint_file" not in json.loads(reloaded.to_json()) + + +def test_checkpoint_user_values_take_precedence(tmp_path): + generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE), counter=17) + checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints") + + reloaded = CheckpointingRandomGenerator(checkpoint_file=checkpoint_path, counter=99) + assert reloaded.counter == 99 + + +def test_save_checkpoint_avoids_overwriting(tmp_path): + generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE)) + first = generator._save_checkpoint(tmp_path / "checkpoints") + second = generator._save_checkpoint(tmp_path / "checkpoints") + + assert first != second + assert len(os.listdir(tmp_path / "checkpoints")) == 2 + + +def test_load_legacy_checkpoint(tmp_path): + generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE), counter=17) + checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints") + make_legacy_checkpoint(checkpoint_path) + + # VOCS is recovered from "vocs.txt" since the checkpoint does not carry it + reloaded = CheckpointingRandomGenerator(checkpoint_file=checkpoint_path) + assert reloaded.counter == 17 + assert reloaded.vocs == generator.vocs + + +def test_load_legacy_checkpoint_missing_vocs_file(tmp_path): + generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE)) + checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints") + make_legacy_checkpoint(checkpoint_path) + os.remove(tmp_path / "vocs.txt") + + with pytest.raises(ValueError, match="does not contain a VOCS object"): + CheckpointingRandomGenerator(checkpoint_file=checkpoint_path) + + +def test_embedded_vocs_preferred_over_legacy_file(tmp_path): + generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE)) + checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints") + + # A stale legacy file beside a modern checkpoint must be ignored + stale_vocs = deepcopy(TEST_VOCS_BASE) + stale_vocs.variables.pop(next(iter(stale_vocs.variables))) + with open(tmp_path / "vocs.txt", "w") as f: + json.dump(stale_vocs.model_dump(), f) + + reloaded = CheckpointingRandomGenerator(checkpoint_file=checkpoint_path) + assert reloaded.vocs == generator.vocs