Merge pull request #386 from JaxGaussianProcesses/tchristie/bo

Add decision making functionality

gpjax/__init__.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from gpjax import integrators
+from gpjax import (
+    decision_making,
+    integrators,
+)
 from gpjax.base import (
     Module,
     param_field,
@@ -81,6 +84,7 @@
     "Module",
     "param_field",
     "cite",
+    "decision_making",
     "kernels",
     "fit",
     "Prior",

gpjax/citation.py

@@ -10,6 +10,11 @@
 )
 from jaxlib.xla_extension import PjitFunction
 
+from gpjax.decision_making.test_functions import (
+    Forrester,
+    LogarithmicGoldsteinPrice,
+)
+from gpjax.decision_making.utility_functions import ThompsonSampling
 from gpjax.kernels import (
     RFF,
     ArcCosine,
@@ -196,3 +201,41 @@ def _(tree) -> PaperCitation:
         booktitle="Uncertainty in Artificial Intelligence",
         citation_type="article",
     )
+
+
+####################
+# Decision making citations
+####################
+@cite.register(ThompsonSampling)
+def _(tree) -> PaperCitation:
+    return PaperCitation(
+        citation_key="wilson2020efficiently",
+        title="Efficiently sampling functions from Gaussian process posteriors",
+        authors="Wilson, James and Borovitskiy, Viacheslav and Terenin, Alexander and Mostowsky, Peter and Deisenroth, Marc",
+        year="2020",
+        booktitle="International Conference on Machine Learning",
+        citation_type="article",
+    )
+
+
+@cite.register(Forrester)
+def _(tree) -> BookCitation:
+    return BookCitation(
+        citation_key="forrester2008engineering",
+        authors="Forrester, Alexander and Sobester, Andras and Keane, Andy",
+        title="Engineering design via surrogate modelling: a practical guide",
+        year="2008",
+        publisher="John Wiley & Sons",
+    )
+
+
+@cite.register(LogarithmicGoldsteinPrice)
+def _(tree) -> PaperCitation:
+    return PaperCitation(
+        citation_key="picheny2013benchmark",
+        authors="Picheny, Victor and Wagner, Tobias and Ginsbourger, David",
+        title="A benchmark of kriging-based infill criteria for noisy optimization",
+        year="2013",
+        booktitle="Structural and multidisciplinary optimization",
+        citation_type="article",
+    )

gpjax/decision_making/__init__.py

@@ -0,0 +1,63 @@
+# Copyright 2023 The JaxGaussianProcesses Contributors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from gpjax.decision_making.decision_maker import (
+    AbstractDecisionMaker,
+    UtilityDrivenDecisionMaker,
+)
+from gpjax.decision_making.posterior_handler import PosteriorHandler
+from gpjax.decision_making.search_space import (
+    AbstractSearchSpace,
+    ContinuousSearchSpace,
+)
+from gpjax.decision_making.test_functions import (
+    AbstractContinuousTestFunction,
+    Forrester,
+    LogarithmicGoldsteinPrice,
+    Quadratic,
+)
+from gpjax.decision_making.utility_functions import (
+    AbstractSinglePointUtilityFunctionBuilder,
+    AbstractUtilityFunctionBuilder,
+    SinglePointUtilityFunction,
+    ThompsonSampling,
+    UtilityFunction,
+)
+from gpjax.decision_making.utility_maximizer import (
+    AbstractSinglePointUtilityMaximizer,
+    AbstractUtilityMaximizer,
+    ContinuousSinglePointUtilityMaximizer,
+)
+from gpjax.decision_making.utils import build_function_evaluator
+
+__all__ = [
+    "AbstractUtilityFunctionBuilder",
+    "AbstractUtilityMaximizer",
+    "AbstractDecisionMaker",
+    "AbstractSearchSpace",
+    "AbstractSinglePointUtilityFunctionBuilder",
+    "AbstractSinglePointUtilityMaximizer",
+    "UtilityFunction",
+    "build_function_evaluator",
+    "ContinuousSinglePointUtilityMaximizer",
+    "ContinuousSearchSpace",
+    "UtilityDrivenDecisionMaker",
+    "AbstractContinuousTestFunction",
+    "Forrester",
+    "LogarithmicGoldsteinPrice",
+    "PosteriorHandler",
+    "Quadratic",
+    "SinglePointUtilityFunction",
+    "ThompsonSampling",
+]