Rework serialization to support pydantic codecs

docs/examples/basic/xopt_serialization.ipynb

@@ -0,0 +1,159 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "aab681aa",
+   "metadata": {},
+   "source": [
+    "# Serialization Options\n",
+    "Xopt objects serialize to YAML/JSON for checkpointing and restarts. What goes\n",
+    "into a dump is controlled by a few writer options, passed to `yaml()`, `json()`\n",
+    "or `dump()`:\n",
+    "\n",
+    "- `module_mode`: `\"drop\"` (default), `\"file\"` (write torch modules as `.pt`\n",
+    "  sidecar files next to the dump file) or `\"inline\"` (embed as base64 strings)\n",
+    "- `array_mode`: `\"list\"` (default) or `\"b64\"` (binary numpy/torch payloads)\n",
+    "- `df_mode`: `\"dict\"` (default) or `\"b64\"` for dataframes\n",
+    "- `compress`: `None` (default), `\"gzip\"` or `\"zstd\"`, with an optional `level`\n",
+    "\n",
+    "This example compares dump sizes for a sizeable Xopt instance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "797033b0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "\n",
+    "from xopt import Xopt, Evaluator, VOCS\n",
+    "from xopt.generators.bayesian import UpperConfidenceBoundGenerator\n",
+    "from xopt.resources.test_functions.sinusoid_1d import evaluate_sinusoid\n",
+    "\n",
+    "vocs = VOCS(variables={\"x1\": [0, 1.75 * math.pi]}, objectives={\"y1\": \"MINIMIZE\"})\n",
+    "X = Xopt(\n",
+    "    generator=UpperConfidenceBoundGenerator(vocs=vocs),\n",
+    "    evaluator=Evaluator(function=evaluate_sinusoid),\n",
+    ")\n",
+    "\n",
+    "# evaluate 500 random points and train the GP model\n",
+    "X.random_evaluate(500)\n",
+    "X.generator.train_model()\n",
+    "X.generator.model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8292f847",
+   "metadata": {},
+   "source": [
+    "## Default dump\n",
+    "By default torch modules (like the trained GP model above) are dropped and\n",
+    "arrays/dataframes are written as plain lists and dicts."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6291e520",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(f\"default: {len(X.yaml()) / 1024:.0f} KiB\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "56e2f2c2",
+   "metadata": {},
+   "source": [
+    "## Inline torch modules\n",
+    "With `module_mode=\"inline\"` the trained model is embedded in the dump as a\n",
+    "base64 payload, so the file is fully self-contained. Compression shrinks the\n",
+    "binary payloads considerably; combining it with `b64` array and dataframe\n",
+    "modes compresses the evaluation data as well."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "944cd42d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "\n",
+    "b64 = dict(module_mode=\"inline\", array_mode=\"b64\", df_mode=\"b64\")\n",
+    "variants = {\n",
+    "    \"inline raw (lists)\": dict(module_mode=\"inline\"),\n",
+    "    \"inline b64\": b64,\n",
+    "    \"inline b64 + gzip 9\": dict(**b64, compress=\"gzip\", level=9),\n",
+    "    \"inline b64 + zstd 3\": dict(**b64, compress=\"zstd\", level=3),\n",
+    "    \"inline b64 + zstd 22 (max)\": dict(**b64, compress=\"zstd\", level=22),\n",
+    "}\n",
+    "for name, kwargs in variants.items():\n",
+    "    start = time.perf_counter()\n",
+    "    size = len(X.yaml(**kwargs))\n",
+    "    elapsed = time.perf_counter() - start\n",
+    "    print(f\"{name:28s} {size / 1024:4.0f} KiB  {elapsed * 1e3:6.1f} ms\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a89d5733",
+   "metadata": {},
+   "source": [
+    "When `level` is not given, gzip defaults to level 9 and zstd to the\n",
+    "zstandard default (level 3). Higher zstd levels trade dump time for\n",
+    "size."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2626cce4",
+   "metadata": {},
+   "source": [
+    "## Reloading\n",
+    "Every variant reloads with `from_yaml`/`from_file`, including the trained\n",
+    "model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "419226f4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "X2 = Xopt.from_yaml(X.yaml(module_mode=\"inline\", compress=\"zstd\"))\n",
+    "X2.generator.model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "620b96b0",
+   "metadata": {},
+   "source": [
+    "For long optimization runs prefer `module_mode=\"file\"` (the default used by\n",
+    "`X.dump()` when `serialize_torch=True`): modules are written as `.pt` files\n",
+    "next to the dump file, and `Xopt.from_file` finds them from any working\n",
+    "directory."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/examples/sequential/rcds.ipynb

@@ -129,7 +129,6 @@
     "    max_evaluations: 100\n",
     "generator:\n",
     "    name: rcds\n",
-    "    x0: null\n",
     "    init_mat: null\n",
     "    noise: 0.00001\n",
     "    step: 0.01\n",
@@ -199,7 +198,6 @@
     "    max_evaluations: 400\n",
     "generator:\n",
     "    name: rcds\n",
-    "    x0: null\n",
     "    init_mat: null\n",
     "    noise: 1e-8\n",
     "    step: 0.01\n",

docs/index.md

@@ -77,7 +77,7 @@ generator:
     output_path: .
 
 evaluator:
-    function: my_function
+    function: my_module.my_function
     function_kwargs:
       my_arguments: 42
 

environment.yml

@@ -6,14 +6,14 @@ dependencies:
   - python>=3.10
   - deap
   - numpy
-  - pydantic>=2.3
+  - pydantic>=2.12
   - pyyaml
   - botorch>=0.13.0
   - scipy>=1.10.1
   - pandas
   - ipywidgets
   - tqdm
-  - orjson
+  - zstandard
   - matplotlib
 # parallel
   - mpi4py

mkdocs.yml

@@ -16,6 +16,7 @@ nav:
             - Working with generators: examples/basic/xopt_generator.ipynb
             - Stopping conditions: examples/basic/xopt_stopping_condition.ipynb
             - Checkpointing and restarts: examples/basic/checkpointing_and_restarts.ipynb
+            - Serialization options: examples/basic/xopt_serialization.ipynb
         - Bayesian:
             - Gaussian Process Model Creation:
                 - Basic example: examples/gp_model_creation/model_creation.ipynb

pyproject.toml

@@ -29,7 +29,7 @@ dependencies = [
   "pandas",
   "ipywidgets",
   "tqdm",
-  "orjson",
+  "zstandard",
   "matplotlib",
   "gest-api>=0.2"
 ]