sparse axis aligned kernel for GP modeling

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,

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
+}

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