From 3cbc262c89862b56e4acddf8e4bd8147dc07797f Mon Sep 17 00:00:00 2001
From: "Christopher M. Pierce" <contact@chris-pierce.com>
Date: Thu, 6 Aug 2026 18:33:01 -0700
Subject: [PATCH 1/5] setup mixin class to handle checkpointing features

---
 xopt/generators/checkpoints.py            | 123 ++++++++++++++++++++++
 xopt/tests/generators/test_checkpoints.py |  84 +++++++++++++++
 2 files changed, 207 insertions(+)
 create mode 100644 xopt/generators/checkpoints.py
 create mode 100644 xopt/tests/generators/test_checkpoints.py

diff --git a/xopt/generators/checkpoints.py b/xopt/generators/checkpoints.py
new file mode 100644
index 00000000..d234fcb7
--- /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 "checkpoints" subdirectory of a caller-supplied
+    directory, with the VOCS object written alongside it as "vocs.txt".
+
+    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 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 generator '
+                "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
+
+    def _save_checkpoint(self, path: str | os.PathLike) -> str:
+        """
+        Write the VOCS object and a checkpoint of the generator state to disk.
+
+        Parameters
+        ----------
+        path : str or os.PathLike
+            Directory into which "vocs.txt" and the "checkpoints" subdirectory
+            containing the checkpoint file are written.
+
+        Returns
+        -------
+        str
+            Path to the checkpoint file which was written.
+        """
+        # Set up the output directory and write the VOCS object needed to reload
+        checkpoint_dir = os.path.join(path, "checkpoints")
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        with open(os.path.join(path, "vocs.txt"), "w") as f:
+            json.dump(self.vocs.model_dump(), f)
+
+        # Create a base filename
+        base_checkpoint_filename = datetime.now().strftime("%Y%m%d_%H%M%S")
+        checkpoint_path = os.path.join(
+            checkpoint_dir, 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(
+                checkpoint_dir, 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/tests/generators/test_checkpoints.py b/xopt/tests/generators/test_checkpoints.py
new file mode 100644
index 00000000..3f0726cb
--- /dev/null
+++ b/xopt/tests/generators/test_checkpoints.py
@@ -0,0 +1,84 @@
+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 test_save_checkpoint_layout(tmp_path):
+    generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE))
+    checkpoint_path = generator._save_checkpoint(tmp_path)
+
+    # VOCS is written to the parent directory and the checkpoint into "checkpoints"
+    vocs_path = tmp_path / "vocs.txt"
+    assert vocs_path.is_file()
+    assert os.path.dirname(checkpoint_path) == str(tmp_path / "checkpoints")
+    assert os.listdir(tmp_path / "checkpoints") == [os.path.basename(checkpoint_path)]
+
+    # Filename follows the timestamp plus deduplication index scheme
+    _, index = parse_checkpoint_filename(os.path.basename(checkpoint_path))
+    assert index == 1
+
+    # Both files hold valid JSON and the VOCS object round trips
+    with open(checkpoint_path) as f:
+        assert "counter" in json.load(f)
+    with open(vocs_path) as f:
+        assert VOCS(**json.load(f)) == 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)
+
+    # VOCS comes from the checkpoint output directory, 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)
+
+    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)
+    second = generator._save_checkpoint(tmp_path)
+
+    assert first != second
+    assert len(os.listdir(tmp_path / "checkpoints")) == 2
+
+
+def test_load_checkpoint_missing_vocs(tmp_path):
+    generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE))
+    checkpoint_path = generator._save_checkpoint(tmp_path)
+    os.remove(tmp_path / "vocs.txt")
+
+    with pytest.raises(ValueError, match="Could not load VOCS file"):
+        CheckpointingRandomGenerator(checkpoint_file=checkpoint_path)

From d59f14c0a5b016021829d9dbdf4acc05d5645d61 Mon Sep 17 00:00:00 2001
From: "Christopher M. Pierce" <contact@chris-pierce.com>
Date: Thu, 6 Aug 2026 18:37:06 -0700
Subject: [PATCH 2/5] refactor `NSGA2Generator` to use new `CheckpointMixin`

---
 xopt/generators/ga/nsga2.py | 94 ++-----------------------------------
 1 file changed, 4 insertions(+), 90 deletions(-)

diff --git a/xopt/generators/ga/nsga2.py b/xopt/generators/ga/nsga2.py
index ba716ccc..da8aa13b 100644
--- a/xopt/generators/ga/nsga2.py
+++ b/xopt/generators/ga/nsga2.py
@@ -1,4 +1,3 @@
-from datetime import datetime
 from itertools import chain
 from pydantic import Field, Discriminator, model_validator
 from typing import Annotated
@@ -14,6 +13,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 +312,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 +373,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 +431,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):
         """
@@ -682,37 +625,8 @@ 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(self.output_dir)
+                    self._logger.debug(f'saved checkpoint file "{checkpoint_path}"')
 
     def set_data(self, data):
         self.data = data

From 36b38ea907a9bbedc072152e2e686e10e24cbf5c Mon Sep 17 00:00:00 2001
From: "Christopher M. Pierce" <contact@chris-pierce.com>
Date: Thu, 6 Aug 2026 18:42:43 -0700
Subject: [PATCH 3/5] don't write `vocs.txt` from main generator

---
 docs/examples/ga/nsga2/nsga2_python.ipynb              | 6 +++---
 docs/examples/ga/nsga2/yaml_interface/index.md         | 2 +-
 docs/examples/ga/nsga2/yaml_interface/nsga2_yaml.ipynb | 2 +-
 xopt/generators/ga/nsga2.py                            | 3 ---
 4 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/docs/examples/ga/nsga2/nsga2_python.ipynb b/docs/examples/ga/nsga2/nsga2_python.ipynb
index 08d5ee2a..96bc4282 100644
--- a/docs/examples/ga/nsga2/nsga2_python.ipynb
+++ b/docs/examples/ga/nsga2/nsga2_python.ipynb
@@ -202,7 +202,7 @@
     "\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",
+    " - `vocs.txt`: The VOCS object so that the objectives, constraints, decision variables are retained alongside the data. Written as each checkopint is emitted\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",
@@ -357,7 +357,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "xopt-dev",
    "language": "python",
    "name": "python3"
   },
@@ -371,7 +371,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 2e293e6d..9c0592c4 100644
--- a/docs/examples/ga/nsga2/yaml_interface/index.md
+++ b/docs/examples/ga/nsga2/yaml_interface/index.md
@@ -255,7 +255,7 @@ 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
+- `vocs.txt`: A copy of the variable, objectives, and constraints (VOCs) definitions, written as each checkpoint is emitted.
 - `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 1f0635d3..063be280 100644
--- a/docs/examples/ga/nsga2/yaml_interface/nsga2_yaml.ipynb
+++ b/docs/examples/ga/nsga2/yaml_interface/nsga2_yaml.ipynb
@@ -52,7 +52,7 @@
     "- `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",
+    "- `vocs.txt`: A copy of the variable, objectives, and constraints (VOCs) definitions, written as each checkpoint is emitted\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/ga/nsga2.py b/xopt/generators/ga/nsga2.py
index da8aa13b..402fd46a 100644
--- a/xopt/generators/ga/nsga2.py
+++ b/xopt/generators/ga/nsga2.py
@@ -1,7 +1,6 @@
 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
@@ -593,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)

From 15c163a8680e937cdd16128dd8a1c32da5141da0 Mon Sep 17 00:00:00 2001
From: "Christopher M. Pierce" <contact@chris-pierce.com>
Date: Thu, 6 Aug 2026 19:34:02 -0700
Subject: [PATCH 4/5] we don't need to actually save vocs.txt anymore

---
 docs/examples/ga/nsga2/nsga2_python.ipynb     | 18 +-----
 .../examples/ga/nsga2/yaml_interface/index.md |  1 -
 .../ga/nsga2/yaml_interface/nsga2_yaml.ipynb  |  1 -
 xopt/generators/checkpoints.py                | 31 +++++-----
 xopt/tests/generators/ga/test_nsga2.py        |  6 --
 xopt/tests/generators/test_checkpoints.py     | 56 ++++++++++++++++---
 6 files changed, 65 insertions(+), 48 deletions(-)

diff --git a/docs/examples/ga/nsga2/nsga2_python.ipynb b/docs/examples/ga/nsga2/nsga2_python.ipynb
index 96bc4282..89576db4 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. Written as each checkopint is emitted\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,
diff --git a/docs/examples/ga/nsga2/yaml_interface/index.md b/docs/examples/ga/nsga2/yaml_interface/index.md
index 9c0592c4..91d4ff82 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, written as each checkpoint is emitted.
 - `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 063be280..63da9e81 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, written as each checkpoint is emitted\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
index d234fcb7..1d4d519c 100644
--- a/xopt/generators/checkpoints.py
+++ b/xopt/generators/checkpoints.py
@@ -11,7 +11,9 @@ class CheckpointMixin(BaseModel):
     Mix-in class adding checkpoint saving and loading to a generator.
 
     Checkpoints are written to a "checkpoints" subdirectory of a caller-supplied
-    directory, with the VOCS object written alongside it as "vocs.txt".
+    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 beside the checkpoint directory 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
@@ -44,21 +46,24 @@ def _load_checkpoint_data(fname: str) -> dict:
         dict
             Dictionary containing VOCS and checkpoint data
         """
-        # Load the VOCS object
+        # 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'Could not load VOCS file at "{vocs_fname}". Complete generator '
-                "output directory is required for loading from checkpoint."
+                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))
 
-        # Load the checkpoint
-        with open(fname) as f:
-            checkpoint_data = json.load(f)
-
         return {"vocs": vocs, **checkpoint_data}
 
     @model_validator(mode="before")
@@ -83,24 +88,22 @@ def load_from_checkpoint(cls, values):
 
     def _save_checkpoint(self, path: str | os.PathLike) -> str:
         """
-        Write the VOCS object and a checkpoint of the generator state to disk.
+        Write a checkpoint of the generator state to disk.
 
         Parameters
         ----------
         path : str or os.PathLike
-            Directory into which "vocs.txt" and the "checkpoints" subdirectory
-            containing the checkpoint file are written.
+            Directory into which the "checkpoints" subdirectory containing the
+            checkpoint file is written.
 
         Returns
         -------
         str
             Path to the checkpoint file which was written.
         """
-        # Set up the output directory and write the VOCS object needed to reload
+        # Set up the output directory
         checkpoint_dir = os.path.join(path, "checkpoints")
         os.makedirs(checkpoint_dir, exist_ok=True)
-        with open(os.path.join(path, "vocs.txt"), "w") as f:
-            json.dump(self.vocs.model_dump(), f)
 
         # Create a base filename
         base_checkpoint_filename = datetime.now().strftime("%Y%m%d_%H%M%S")
diff --git a/xopt/tests/generators/ga/test_nsga2.py b/xopt/tests/generators/ga/test_nsga2.py
index 635d93dc..850b819d 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
index 3f0726cb..0b2a3091 100644
--- a/xopt/tests/generators/test_checkpoints.py
+++ b/xopt/tests/generators/test_checkpoints.py
@@ -23,13 +23,25 @@ def parse_checkpoint_filename(filename: str) -> tuple[datetime, int]:
     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)
 
-    # VOCS is written to the parent directory and the checkpoint into "checkpoints"
-    vocs_path = tmp_path / "vocs.txt"
-    assert vocs_path.is_file()
+    # Only the "checkpoints" subdirectory is created, no separate VOCS file
+    assert os.listdir(tmp_path) == ["checkpoints"]
     assert os.path.dirname(checkpoint_path) == str(tmp_path / "checkpoints")
     assert os.listdir(tmp_path / "checkpoints") == [os.path.basename(checkpoint_path)]
 
@@ -37,11 +49,11 @@ def test_save_checkpoint_layout(tmp_path):
     _, index = parse_checkpoint_filename(os.path.basename(checkpoint_path))
     assert index == 1
 
-    # Both files hold valid JSON and the VOCS object round trips
+    # The checkpoint holds valid JSON and carries the VOCS object itself
     with open(checkpoint_path) as f:
-        assert "counter" in json.load(f)
-    with open(vocs_path) as f:
-        assert VOCS(**json.load(f)) == generator.vocs
+        checkpoint_data = json.load(f)
+    assert "counter" in checkpoint_data
+    assert VOCS(**checkpoint_data["vocs"]) == generator.vocs
 
 
 def test_checkpoint_round_trip(tmp_path):
@@ -75,10 +87,36 @@ def test_save_checkpoint_avoids_overwriting(tmp_path):
     assert len(os.listdir(tmp_path / "checkpoints")) == 2
 
 
-def test_load_checkpoint_missing_vocs(tmp_path):
+def test_load_legacy_checkpoint(tmp_path):
+    generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE), counter=17)
+    checkpoint_path = generator._save_checkpoint(tmp_path)
+    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)
+    make_legacy_checkpoint(checkpoint_path)
     os.remove(tmp_path / "vocs.txt")
 
-    with pytest.raises(ValueError, match="Could not load VOCS file"):
+    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)
+
+    # 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

From 26b32a154cee428beb624aa421320ce1de3f5004 Mon Sep 17 00:00:00 2001
From: "Christopher M. Pierce" <contact@chris-pierce.com>
Date: Thu, 6 Aug 2026 19:38:22 -0700
Subject: [PATCH 5/5] have save_checkpoints point to actual checkpoints
 directory

---
 xopt/generators/checkpoints.py            | 21 +++++++++------------
 xopt/generators/ga/nsga2.py               |  4 +++-
 xopt/tests/generators/test_checkpoints.py | 22 +++++++++++-----------
 3 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/xopt/generators/checkpoints.py b/xopt/generators/checkpoints.py
index 1d4d519c..c189e836 100644
--- a/xopt/generators/checkpoints.py
+++ b/xopt/generators/checkpoints.py
@@ -10,10 +10,10 @@ class CheckpointMixin(BaseModel):
     """
     Mix-in class adding checkpoint saving and loading to a generator.
 
-    Checkpoints are written to a "checkpoints" subdirectory of 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 beside the checkpoint directory and are still supported when loading.
+    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
@@ -93,8 +93,8 @@ def _save_checkpoint(self, path: str | os.PathLike) -> str:
         Parameters
         ----------
         path : str or os.PathLike
-            Directory into which the "checkpoints" subdirectory containing the
-            checkpoint file is written.
+            Directory into which the checkpoint file is written. Created if it
+            does not already exist.
 
         Returns
         -------
@@ -102,20 +102,17 @@ def _save_checkpoint(self, path: str | os.PathLike) -> str:
             Path to the checkpoint file which was written.
         """
         # Set up the output directory
-        checkpoint_dir = os.path.join(path, "checkpoints")
-        os.makedirs(checkpoint_dir, exist_ok=True)
+        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(
-            checkpoint_dir, f"{base_checkpoint_filename}_1.txt"
-        )
+        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(
-                checkpoint_dir, f"{base_checkpoint_filename}_{counter}.txt"
+                path, f"{base_checkpoint_filename}_{counter}.txt"
             )
             counter += 1
 
diff --git a/xopt/generators/ga/nsga2.py b/xopt/generators/ga/nsga2.py
index 402fd46a..e38f3b60 100644
--- a/xopt/generators/ga/nsga2.py
+++ b/xopt/generators/ga/nsga2.py
@@ -622,7 +622,9 @@ def add_data(self, new_data: pd.DataFrame):
                 if self.checkpoint_freq > 0 and (
                     self.n_generations % self.checkpoint_freq == 0
                 ):
-                    checkpoint_path = self._save_checkpoint(self.output_dir)
+                    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):
diff --git a/xopt/tests/generators/test_checkpoints.py b/xopt/tests/generators/test_checkpoints.py
index 0b2a3091..38034320 100644
--- a/xopt/tests/generators/test_checkpoints.py
+++ b/xopt/tests/generators/test_checkpoints.py
@@ -38,12 +38,12 @@ def make_legacy_checkpoint(checkpoint_path: str) -> None:
 
 def test_save_checkpoint_layout(tmp_path):
     generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE))
-    checkpoint_path = generator._save_checkpoint(tmp_path)
+    checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints")
 
-    # Only the "checkpoints" subdirectory is created, no separate VOCS file
-    assert os.listdir(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))
@@ -58,9 +58,9 @@ def test_save_checkpoint_layout(tmp_path):
 
 def test_checkpoint_round_trip(tmp_path):
     generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE), counter=17)
-    checkpoint_path = generator._save_checkpoint(tmp_path)
+    checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints")
 
-    # VOCS comes from the checkpoint output directory, not from the user
+    # 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
@@ -72,7 +72,7 @@ def test_checkpoint_round_trip(tmp_path):
 
 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)
+    checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints")
 
     reloaded = CheckpointingRandomGenerator(checkpoint_file=checkpoint_path, counter=99)
     assert reloaded.counter == 99
@@ -80,8 +80,8 @@ def test_checkpoint_user_values_take_precedence(tmp_path):
 
 def test_save_checkpoint_avoids_overwriting(tmp_path):
     generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE))
-    first = generator._save_checkpoint(tmp_path)
-    second = generator._save_checkpoint(tmp_path)
+    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
@@ -89,7 +89,7 @@ def test_save_checkpoint_avoids_overwriting(tmp_path):
 
 def test_load_legacy_checkpoint(tmp_path):
     generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE), counter=17)
-    checkpoint_path = generator._save_checkpoint(tmp_path)
+    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
@@ -100,7 +100,7 @@ def test_load_legacy_checkpoint(tmp_path):
 
 def test_load_legacy_checkpoint_missing_vocs_file(tmp_path):
     generator = CheckpointingRandomGenerator(vocs=deepcopy(TEST_VOCS_BASE))
-    checkpoint_path = generator._save_checkpoint(tmp_path)
+    checkpoint_path = generator._save_checkpoint(tmp_path / "checkpoints")
     make_legacy_checkpoint(checkpoint_path)
     os.remove(tmp_path / "vocs.txt")
 
@@ -110,7 +110,7 @@ def test_load_legacy_checkpoint_missing_vocs_file(tmp_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)
+    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)