diff --git a/docs/examples/gp_model_creation/model_creation.ipynb b/docs/examples/gp_model_creation/model_creation.ipynb index 9ab3ee3e4..a4180f151 100644 --- a/docs/examples/gp_model_creation/model_creation.ipynb +++ b/docs/examples/gp_model_creation/model_creation.ipynb @@ -221,7 +221,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "xopt-dev", "language": "python", "name": "python3" }, @@ -235,7 +235,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.11" + "version": "3.14.2" } }, "nbformat": 4, diff --git a/docs/examples/gp_model_creation/saas.ipynb b/docs/examples/gp_model_creation/saas.ipynb new file mode 100644 index 000000000..fe4e5cbc8 --- /dev/null +++ b/docs/examples/gp_model_creation/saas.ipynb @@ -0,0 +1,286 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + } + }, + "source": [ + "## Demonstrating SAAS Functionality in Xopt GP Models\n", + "This example compares two models trained on the same sparse problem:\n", + "1. A standard GP model.\n", + "2. A GP model with SAAS priors enabled via `saas_outputs`.\n", + "\n", + "The synthetic objective depends only on `x0` and `x1`, while the remaining dimensions are irrelevant. This makes it easy to see how SAAS encourages sparsity in learned lengthscales." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# set values if testing\n", + "import os\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", + "\n", + "from xopt import Xopt, Evaluator\n", + "from xopt.generators import RandomGenerator\n", + "from xopt.resources.test_functions.rosenbrock import make_rosenbrock_vocs\n", + "from xopt.generators.bayesian.models import StandardModelConstructor\n", + "from xopt.generators.bayesian.visualize import visualize_model\n", + "\n", + "SMOKE_TEST = os.environ.get(\"SMOKE_TEST\")\n", + "\n", + "# Use a higher-dimensional input space to make sparsity visible.\n", + "n_dim = 20\n", + "n_initial = 15\n", + "vocs = make_rosenbrock_vocs(n_dim)\n", + "vocs.observables = [\"z\"]\n", + "\n", + "\n", + "def evaluate_func(X):\n", + " return {\n", + " \"y\": X[\"x0\"] ** 2 + X[\"x1\"] ** 2 + X[\"x10\"] ** 2,\n", + " \"z\": X[\"x0\"] ** 2,\n", + " }" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + } + }, + "source": [ + "## Generate Training Data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "evaluator = Evaluator(function=evaluate_func)\n", + "generator = RandomGenerator(vocs=vocs)\n", + "X = Xopt(generator=generator, evaluator=evaluator)\n", + "X.random_evaluate(n_initial)\n", + "\n", + "data = X.data\n", + "data[vocs.variable_names + [\"y\", \"z\"]].head()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Build Standard and SAAS Models\n", + "Both models are trained on the same data so any difference in behavior comes from the SAAS prior configuration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "standard_constructor = StandardModelConstructor()\n", + "saas_constructor = StandardModelConstructor(saas_outputs=[\"y\", \"z\"])\n", + "\n", + "# Build two models from identical data.\n", + "standard_model = standard_constructor.build_model(\n", + " input_names=vocs.variable_names,\n", + " outcome_names=[\"y\", \"z\"],\n", + " data=data,\n", + ")\n", + "\n", + "saas_model = saas_constructor.build_model(\n", + " input_names=vocs.variable_names,\n", + " outcome_names=[\"y\", \"z\"],\n", + " data=data,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def get_lengthscale_vector(gp_model):\n", + " for name, value in gp_model.named_parameters():\n", + " if \"lengthscale\" in name:\n", + " return value.detach().cpu().numpy().ravel()\n", + " raise RuntimeError(\"No lengthscale parameter found\")\n", + "\n", + "\n", + "standard_y_model = standard_model.models[vocs.output_names.index(\"y\")]\n", + "saas_y_model = saas_model.models[vocs.output_names.index(\"y\")]\n", + "\n", + "comparison = pd.DataFrame(\n", + " {\n", + " \"variable\": vocs.variable_names,\n", + " \"standard_lengthscale\": get_lengthscale_vector(standard_y_model),\n", + " \"saas_lengthscale\": get_lengthscale_vector(saas_y_model),\n", + " }\n", + ")\n", + "\n", + "comparison.sort_values(\"saas_lengthscale\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + } + }, + "source": [ + "## Compare Learned Lengthscales\n", + "Smaller lengthscales correspond to stronger local sensitivity for that variable.\n", + "Because this objective is sparse, SAAS should concentrate sensitivity on `x0` and `x1` \n", + "while the standard model does not until there is a significant amount of data." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_data = comparison.set_index(\"variable\")[\n", + " [\"standard_lengthscale\", \"saas_lengthscale\"]\n", + "]\n", + "\n", + "ax = plot_data.plot(kind=\"bar\", figsize=(10, 4), rot=0)\n", + "ax.set_ylabel(\"Lengthscale\")\n", + "ax.set_title(\"Standard GP vs SAAS GP lengthscales for output y\")\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "reference_point = data.iloc[-1][vocs.variable_names].to_dict()\n", + "\n", + "fig, ax = visualize_model(\n", + " standard_model,\n", + " vocs,\n", + " data,\n", + " variable_names=[\"x0\", \"x1\"],\n", + " reference_point=reference_point,\n", + ")\n", + "fig.suptitle(\"Standard GP\", y=1.02)\n", + "\n", + "fig, ax = visualize_model(\n", + " saas_model,\n", + " vocs,\n", + " data,\n", + " variable_names=[\"x0\", \"x1\"],\n", + " reference_point=reference_point,\n", + ")\n", + "fig.suptitle(\"SAAS GP\", y=1.02)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Add more data and rebuild models to see how lengthscales evolve." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "X.random_evaluate(n_initial)\n", + "\n", + "# rebuild models with more data to see how lengthscales evolve\n", + "standard_model = standard_constructor.build_model(\n", + " input_names=vocs.variable_names,\n", + " outcome_names=[\"y\", \"z\"],\n", + " data=X.data,\n", + ")\n", + "\n", + "saas_model = saas_constructor.build_model(\n", + " input_names=vocs.variable_names,\n", + " outcome_names=[\"y\", \"z\"],\n", + " data=X.data,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "standard_y_model = standard_model.models[vocs.output_names.index(\"y\")]\n", + "saas_y_model = saas_model.models[vocs.output_names.index(\"y\")]\n", + "\n", + "comparison = pd.DataFrame(\n", + " {\n", + " \"variable\": vocs.variable_names,\n", + " \"standard_lengthscale\": get_lengthscale_vector(standard_y_model),\n", + " \"saas_lengthscale\": get_lengthscale_vector(saas_y_model),\n", + " }\n", + ")\n", + "\n", + "comparison.sort_values(\"saas_lengthscale\")\n", + "\n", + "plot_data = comparison.set_index(\"variable\")[\n", + " [\"standard_lengthscale\", \"saas_lengthscale\"]\n", + "]\n", + "\n", + "ax = plot_data.plot(kind=\"bar\", figsize=(10, 4), rot=0)\n", + "ax.set_ylabel(\"Lengthscale\")\n", + "ax.set_title(\"Standard GP vs SAAS GP lengthscales for output y\")\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Takeaway\n", + "SAAS is enabled by setting `saas_outputs` in `StandardModelConstructor`.\n", + "In sparse problems, SAAS generally produces more selective lengthscales and can improve robustness when many dimensions are weakly relevant." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "xopt-dev", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.2" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/xopt/generators/bayesian/base_model.py b/xopt/generators/bayesian/base_model.py index f722bf51f..4ea4f9235 100644 --- a/xopt/generators/bayesian/base_model.py +++ b/xopt/generators/bayesian/base_model.py @@ -12,6 +12,7 @@ from gpytorch.mlls import VariationalELBO from pydantic import ConfigDict from torch import Tensor +from botorch.models.map_saas import get_map_saas_model from xopt.generators.bayesian.custom_botorch.heteroskedastic import ( XoptHeteroskedasticSingleTaskGP, @@ -247,3 +248,53 @@ def build_approximate_gp( fit_gpytorch_mll(mll) return model + + @staticmethod + def build_map_saas_gp( + X: Tensor, Y: Tensor, Yvar: Tensor, train: bool = True, **kwargs + ): + """ + Utility method for creating and training a MAP SAAS SingleTaskGP model. + + Parameters + ---------- + X : Tensor + Training data for input variables. + Y : Tensor + Training data for outcome variables. + Yvar : Tensor + Training data for outcome variable variances. + train : bool, True + Flag to specify if hyperparameter training should take place + **kwargs + Additional keyword arguments for model configuration. + + Returns + ------- + Model + The trained MAP SAAS SingleTaskGP model. + + Notes + ----- + MAP SAAS modeling can be unstable when the number of dimensions is high and the amount of data is low. + Your results may vary and keep an eye on warnings. + + """ + + if X.shape[0] == 0 or Y.shape[0] == 0: + raise ValueError("no data found to train model!") + + model = get_map_saas_model(X, Y, Yvar, **kwargs) + + if train: + try: + with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + mll = ExactMarginalLogLikelihood(model.likelihood, model) + fit_gpytorch_mll(mll) + except ModelFittingError: + warnings.warn( + "Model fitting failed for MAP SAAS GP. Returning untrained model." + ) + + return model diff --git a/xopt/generators/bayesian/models/standard.py b/xopt/generators/bayesian/models/standard.py index 7eedea46e..52f8445da 100644 --- a/xopt/generators/bayesian/models/standard.py +++ b/xopt/generators/bayesian/models/standard.py @@ -130,6 +130,9 @@ class StandardModelConstructor(ModelConstructor): description="specify if inputs should be transformed inside the gp " "model, can optionally specify a dict of specifications", ) + saas_outputs: List[str] = Field( + [], description="list of output names to apply SAAS priors to" + ) custom_noise_prior: Optional[Prior] = Field( None, description="specify custom noise prior for the GP likelihood, " @@ -335,6 +338,31 @@ def build_model( "mean_module": mean_module, } + # handle saas model construction + if outcome_name in self.saas_outputs: + if kwargs.pop("covar_module", None) is not None: + warnings.warn( + f"Covariance module specified for output {outcome_name} will be overwritten by SAAS model construction." + ) + if kwargs.pop("mean_module", None) is not None: + warnings.warn( + f"Mean module specified for output {outcome_name} will be overwritten by SAAS model construction." + ) + models.append( + self.build_map_saas_gp( + train_X.to(**tkwargs), + train_Y.to(**tkwargs), + train_Yvar.to(**tkwargs) if train_Yvar is not None else None, + train=False, + **kwargs, + ) + ) + + # add in likelihood if needed + # if models[-1].likelihood is None: + # models[-1].likelihood = self.get_likelihood() + continue + if train_Yvar is None: # train basic single-task-gp model models.append( diff --git a/xopt/tests/generators/bayesian/test_model_constructor.py b/xopt/tests/generators/bayesian/test_model_constructor.py index a9bfd95ee..ddd04a574 100644 --- a/xopt/tests/generators/bayesian/test_model_constructor.py +++ b/xopt/tests/generators/bayesian/test_model_constructor.py @@ -2,6 +2,7 @@ import os import time from copy import deepcopy +from unittest.mock import patch import numpy as np import pandas as pd @@ -17,7 +18,7 @@ from gpytorch import ExactMarginalLogLikelihood from gpytorch.kernels import PeriodicKernel, PolynomialKernel, ScaleKernel from gpytorch.likelihoods import GaussianLikelihood -from gpytorch.means import ConstantMean +from gpytorch.means import ConstantMean, LinearMean from gpytorch.priors import GammaPrior from pydantic import ValidationError @@ -832,6 +833,30 @@ def test_build_models_with_training(self): ) assert isinstance(model_approx, SingleTaskVariationalGP) + def test_build_map_saas_gp(self): + torch.manual_seed(0) + X = torch.randn(10, 2) + Y = torch.randn(10, 1) + Yvar = torch.ones(10, 1) * 0.1 + + with pytest.raises(ValueError, match="no data found to train model!"): + StandardModelConstructor.build_map_saas_gp( + torch.empty(0, 2), torch.empty(0, 1), torch.empty(0, 1) + ) + + with patch( + "xopt.generators.bayesian.base_model.fit_gpytorch_mll", + side_effect=ModelFittingError("fitting failed"), + ): + with pytest.warns( + UserWarning, + match="Model fitting failed for MAP SAAS GP. Returning untrained model.", + ): + model = StandardModelConstructor.build_map_saas_gp( + X, Y, Yvar, train=True + ) + assert isinstance(model, SingleTaskGP) + def test_approximate_gp(self): test_vocs = deepcopy(TEST_VOCS) test_data = deepcopy(TEST_DATA) @@ -1265,3 +1290,49 @@ def test_batched_empty_data(self): constructor = BatchedModelConstructor() with pytest.raises(ValueError, match="no data found"): constructor.build_model_from_vocs(test_vocs, empty_data) + + def test_saas_model_constructor(self): + test_vocs = deepcopy(TEST_VOCS) + test_data = deepcopy(TEST_DATA) + + constructor = StandardModelConstructor( + saas_outputs=["y1"], + covar_modules={"y1": ScaleKernel(PeriodicKernel())}, + mean_modules={"y1": LinearMean(TEST_VOCS.n_variables)}, + ) + with ( + patch.object( + StandardModelConstructor, + "build_map_saas_gp", + wraps=StandardModelConstructor.build_map_saas_gp, + ) as mock_build_map_saas_gp, + pytest.warns( + UserWarning, match="will be overwritten by SAAS model construction" + ) as warning_records, + ): + model = constructor.build_model_from_vocs(test_vocs, test_data) + + warning_messages = {str(record.message) for record in warning_records} + expected_warning_messages = { + "Covariance module specified for output y1 will be overwritten by SAAS model construction.", + "Mean module specified for output y1 will be overwritten by SAAS model construction.", + } + assert expected_warning_messages.issubset(warning_messages) + assert mock_build_map_saas_gp.call_count == 1 + assert isinstance(model.models[0], SingleTaskGP) + assert isinstance(model.models[0].covar_module, ScaleKernel) + assert isinstance(model.models[0].mean_module, ConstantMean) + assert hasattr(model.models[0].covar_module.base_kernel, "lengthscale") + assert not isinstance(model.models[0].covar_module.base_kernel, PeriodicKernel) + assert not isinstance(model.models[0].mean_module, LinearMean) + assert isinstance(model, ModelListGP) + + constructor_multi = StandardModelConstructor(saas_outputs=["y1", "c1"]) + with patch.object( + StandardModelConstructor, + "build_map_saas_gp", + wraps=StandardModelConstructor.build_map_saas_gp, + ) as mock_build_map_saas_gp_multi: + model_multi = constructor_multi.build_model_from_vocs(test_vocs, test_data) + assert mock_build_map_saas_gp_multi.call_count == 2 + assert isinstance(model_multi, ModelListGP)