gpjax 0.8.1__tar.gz → 0.9.0__tar.gz

{gpjax-0.8.1 → gpjax-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpjax
-Version: 0.8.1
+Version: 0.9.0
 Summary: Gaussian processes in JAX.
 Home-page: https://github.com/JaxGaussianProcesses/GPJax
 License: Apache-2.0
@@ -12,16 +12,17 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: beartype (>=0.16.2,<0.17.0)
-Requires-Dist: cola-ml (>=0.0.5,<0.0.6)
-Requires-Dist: jax (>=0.4.16)
-Requires-Dist: jaxlib (>=0.4.16)
-Requires-Dist: jaxtyping (>=0.2.15,<0.3.0)
-Requires-Dist: optax (>=0.1.4,<0.2.0)
-Requires-Dist: orbax-checkpoint (>=0.2.3)
-Requires-Dist: simple-pytree (>=0.1.7,<0.2.0)
-Requires-Dist: tensorflow-probability (>=0.22.0,<0.23.0)
-Requires-Dist: tqdm (>=4.65.0,<5.0.0)
+Requires-Dist: beartype (>=0.16.1,<0.17.0)
+Requires-Dist: cola-ml (==0.0.5)
+Requires-Dist: flax (>=0.8.4,<0.9.0)
+Requires-Dist: jax (<0.4.28)
+Requires-Dist: jaxlib (<0.4.28)
+Requires-Dist: jaxopt (>=0.8.3,<0.9.0)
+Requires-Dist: jaxtyping (>=0.2.10,<0.3.0)
+Requires-Dist: numpy (<2.0.0)
+Requires-Dist: optax (>=0.2.1,<0.3.0)
+Requires-Dist: tensorflow-probability (>=0.24.0,<0.25.0)
+Requires-Dist: tqdm (>=4.66.2,<5.0.0)
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Repository, https://github.com/JaxGaussianProcesses/GPJax
 Description-Content-Type: text/markdown
@@ -100,10 +101,9 @@ helped to shape GPJax into the package it is today.
 ## Notebook examples
 
 > - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/examples/regression/)
-> - [**Classification with MCMC**](https://docs.jaxgaussianprocesses.com/examples/classification/)
+> - [**Classification**](https://docs.jaxgaussianprocesses.com/examples/classification/)
 > - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/collapsed_vi/)
 > - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/uncollapsed_vi/)
-> - [**BlackJax Integration**](https://docs.jaxgaussianprocesses.com/examples/classification/#mcmc-inference)
 > - [**Laplace Approximation**](https://docs.jaxgaussianprocesses.com/examples/classification/#laplace-approximation)
 > - [**Inference on Non-Euclidean Spaces**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
 > - [**Inference on Graphs**](https://docs.jaxgaussianprocesses.com/examples/graph_kernels/)
@@ -174,13 +174,10 @@ posterior = prior * likelihood
 # Define an optimiser
 optimiser = ox.adam(learning_rate=1e-2)
 
-# Define the marginal log-likelihood
-negative_mll = jit(gpx.objectives.ConjugateMLL(negative=True))
-
 # Obtain Type 2 MLEs of the hyperparameters
 opt_posterior, history = gpx.fit(
     model=posterior,
-    objective=negative_mll,
+    objective=gpx.objectives.conjugate_mll,
     train_data=D,
     optim=optimiser,
     num_iters=500,

{gpjax-0.8.1 → gpjax-0.9.0}/README.md

@@ -72,10 +72,9 @@ helped to shape GPJax into the package it is today.
 ## Notebook examples
 
 > - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/examples/regression/)
-> - [**Classification with MCMC**](https://docs.jaxgaussianprocesses.com/examples/classification/)
+> - [**Classification**](https://docs.jaxgaussianprocesses.com/examples/classification/)
 > - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/collapsed_vi/)
 > - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/uncollapsed_vi/)
-> - [**BlackJax Integration**](https://docs.jaxgaussianprocesses.com/examples/classification/#mcmc-inference)
 > - [**Laplace Approximation**](https://docs.jaxgaussianprocesses.com/examples/classification/#laplace-approximation)
 > - [**Inference on Non-Euclidean Spaces**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
 > - [**Inference on Graphs**](https://docs.jaxgaussianprocesses.com/examples/graph_kernels/)
@@ -146,13 +145,10 @@ posterior = prior * likelihood
 # Define an optimiser
 optimiser = ox.adam(learning_rate=1e-2)
 
-# Define the marginal log-likelihood
-negative_mll = jit(gpx.objectives.ConjugateMLL(negative=True))
-
 # Obtain Type 2 MLEs of the hyperparameters
 opt_posterior, history = gpx.fit(
     model=posterior,
-    objective=negative_mll,
+    objective=gpx.objectives.conjugate_mll,
     train_data=D,
     optim=optimiser,
     num_iters=500,

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/__init__.py

@@ -12,8 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from warnings import filterwarnings
+
+from beartype.roar import BeartypeDecorHintPep585DeprecationWarning
+
+filterwarnings("ignore", category=BeartypeDecorHintPep585DeprecationWarning)
+
 from gpjax import (
-    base,
     decision_making,
     gps,
     integrators,
@@ -21,12 +26,9 @@ from gpjax import (
     likelihoods,
     mean_functions,
     objectives,
+    parameters,
     variational_families,
 )
-from gpjax.base import (
-    Module,
-    param_field,
-)
 from gpjax.citation import cite
 from gpjax.dataset import Dataset
 from gpjax.fit import (
@@ -38,7 +40,7 @@ __license__ = "MIT"
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.8.1"
+__version__ = "0.9.0"
 
 __all__ = [
     "base",
@@ -49,6 +51,7 @@ __all__ = [
     "likelihoods",
     "mean_functions",
     "objectives",
+    "parameters",
     "variational_families",
     "Dataset",
     "cite",

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/citation.py

@@ -23,13 +23,6 @@ from gpjax.kernels import (
     Matern32,
     Matern52,
 )
-from gpjax.objectives import (
-    ELBO,
-    CollapsedELBO,
-    ConjugateMLL,
-    LogPosteriorDensity,
-    NonConjugateMLL,
-)
 
 CitationType = Union[None, str, Dict[str, str]]
 
@@ -158,46 +151,6 @@ def _(tree) -> PaperCitation:
     )
 
 
-####################
-# Objective citations
-####################
-@cite.register(ConjugateMLL)
-@cite.register(NonConjugateMLL)
-@cite.register(LogPosteriorDensity)
-def _(tree) -> BookCitation:
-    return BookCitation(
-        citation_key="rasmussen2006gaussian",
-        title="Gaussian Processes for Machine Learning",
-        authors="Rasmussen, Carl Edward and Williams, Christopher K",
-        year="2006",
-        publisher="MIT press Cambridge, MA",
-        volume="2",
-    )
-
-
-@cite.register(CollapsedELBO)
-def _(tree) -> PaperCitation:
-    return PaperCitation(
-        citation_key="titsias2009variational",
-        title="Variational learning of inducing variables in sparse Gaussian processes",
-        authors="Titsias, Michalis",
-        year="2009",
-        booktitle="International Conference on Artificial Intelligence and Statistics",
-    )
-
-
-@cite.register(ELBO)
-def _(tree) -> PaperCitation:
-    return PaperCitation(
-        citation_key="hensman2013gaussian",
-        title="Gaussian Processes for Big Data",
-        authors="Hensman, James and Fusi, Nicolo and Lawrence, Neil D",
-        year="2013",
-        booktitle="Uncertainty in Artificial Intelligence",
-        citation_type="article",
-    )
-
-
 ####################
 # Decision making citations
 ####################

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/dataset.py

@@ -12,40 +12,39 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
 from dataclasses import dataclass
 import warnings
 
 from beartype.typing import Optional
+import jax
 import jax.numpy as jnp
 from jaxtyping import Num
-from simple_pytree import Pytree
 
 from gpjax.typing import Array
 
 
 @dataclass
-class Dataset(Pytree):
+@jax.tree_util.register_pytree_node_class
+class Dataset:
     r"""Base class for datasets.
 
-    Attributes
-    ----------
-        X (Optional[Num[Array, "N D"]]): input data.
-        y (Optional[Num[Array, "N Q"]]): output data.
+    Args:
+        X: input data.
+        y: output data.
     """
 
     X: Optional[Num[Array, "N D"]] = None
     y: Optional[Num[Array, "N Q"]] = None
 
     def __post_init__(self) -> None:
-        r"""Checks that the shapes of $`X`$ and $`y`$ are compatible,
-        and provides warnings regarding the precision of $`X`$ and $`y`$."""
+        r"""Checks that the shapes of $X$ and $y$ are compatible,
+        and provides warnings regarding the precision of $X$ and $y$."""
         _check_shape(self.X, self.y)
         _check_precision(self.X, self.y)
 
     def __repr__(self) -> str:
         r"""Returns a string representation of the dataset."""
-        repr = f"- Number of observations: {self.n}\n- Input dimension: {self.in_dim}"
+        repr = f"Dataset(Number of observations: {self.n:=} - Input dimension: {self.in_dim})"
         return repr
 
     def is_supervised(self) -> bool:
@@ -76,14 +75,21 @@ class Dataset(Pytree):
 
     @property
     def in_dim(self) -> int:
-        r"""Dimension of the inputs, $`X`$."""
+        r"""Dimension of the inputs, $X$."""
         return self.X.shape[1]
 
+    def tree_flatten(self):
+        return (self.X, self.y), None
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        return cls(*children)
+
 
 def _check_shape(
     X: Optional[Num[Array, "..."]], y: Optional[Num[Array, "..."]]
 ) -> None:
-    r"""Checks that the shapes of $`X`$ and $`y`$ are compatible."""
+    r"""Checks that the shapes of $X$ and $y$ are compatible."""
     if X is not None and y is not None and X.shape[0] != y.shape[0]:
         raise ValueError(
             "Inputs, X, and outputs, y, must have the same number of rows."
@@ -104,7 +110,7 @@ def _check_shape(
 def _check_precision(
     X: Optional[Num[Array, "..."]], y: Optional[Num[Array, "..."]]
 ) -> None:
-    r"""Checks the precision of $`X`$ and $`y`."""
+    r"""Checks the precision of $X$ and $y`."""
     if X is not None and X.dtype != jnp.float64:
         warnings.warn(
             "X is not of type float64. "

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/decision_making/decision_maker.py

@@ -58,26 +58,23 @@ class AbstractDecisionMaker(ABC):
     the black-box function of interest at this point.
 
     Attributes:
-        search_space (AbstractSearchSpace): Search space over which we can evaluate the
-        function(s) of interest.
-        posterior_handlers (Dict[str, PosteriorHandler]): Dictionary of posterior
-            handlers, which are used to update posteriors throughout the decision making
-            loop. Note that the word `posteriors` is used for consistency with GPJax, but these
-            objects are typically referred to as `models` in the model-based decision
-            making literature. Tags are used to distinguish between posteriors. In a typical
-            Bayesian optimisation setup one of the tags will be `OBJECTIVE`, defined in
-            decision_making.utils.
-        datasets (Dict[str, Dataset]): Dictionary of datasets, which are augmented with
+        search_space: Search space over which we can evaluate the function(s) of interest.
+        posterior_handlers: dictionary of posterior handlers, which are used to update
+            posteriors throughout the decision making loop. Note that the word `posteriors`
+            is used for consistency with GPJax, but these objects are typically referred to
+            as `models` in the model-based decision making literature. Tags are used to
+            distinguish between posteriors. In a typical Bayesian optimisation setup one of
+            the tags will be `OBJECTIVE`, defined in `decision_making.utils`.
+        datasets: dictionary of datasets, which are augmented with
             observations throughout the decision making loop. In a typical setup they are
             also used to update the posteriors, using the `posterior_handlers`. Tags are used
             to distinguish datasets, and correspond to tags in `posterior_handlers`.
-        key (KeyArray): JAX random key, used to generate random numbers.
-        batch_size (int): Number of points to query at each step of the decision making
+        key: JAX random key, used to generate random numbers.
+        batch_size: Number of points to query at each step of the decision making
             loop. Note that `SinglePointUtilityFunction`s are only capable of generating
             one point to be queried at each iteration of the decision making loop.
-        post_ask (List[Callable]): List of functions to be executed after each ask step.
-        post_tell (List[Callable]): List of functions to be executed after each tell
-            step.
+        post_ask: List of functions to be executed after each ask step.
+        post_tell: List of functions to be executed after each tell step.
     """
 
     search_space: AbstractSearchSpace
@@ -140,10 +137,10 @@ class AbstractDecisionMaker(ABC):
         Add newly observed data to datasets and update the corresponding posteriors.
 
         Args:
-            observation_datasets (Mapping[str, Dataset]): Dictionary of datasets
-            containing new observations. Tags are used to distinguish datasets, and
-            correspond to tags in `posterior_handlers` and `self.datasets`.
-            key (KeyArray): JAX PRNG key for controlling random state.
+            observation_datasets: dictionary of datasets containing new observations.
+                Tags are used to distinguish datasets, and correspond to tags in
+                `posterior_handlers` and `self.datasets`.
+            key: JAX PRNG key for controlling random state.
         """
         if observation_datasets.keys() != self.datasets.keys():
             raise ValueError(

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/decision_making/posterior_handler.py

@@ -27,7 +27,7 @@ from gpjax.gps import (
     AbstractPosterior,
     AbstractPrior,
 )
-from gpjax.objectives import AbstractObjective
+from gpjax.objectives import Objective
 from gpjax.typing import KeyArray
 
 LikelihoodBuilder = Callable[[int], AbstractLikelihood]
@@ -42,21 +42,18 @@ class PosteriorHandler:
     observed.
 
     Attributes:
-        prior (AbstractPrior): Prior to use when forming the posterior.
-        likelihood_builder (LikelihoodBuilder): Function which takes the number of
-        datapoints as input and returns a likelihood object initialised with the given
-        number of datapoints.
-        optimization_objective (AbstractObjective): Objective to use for optimizing the
-        posterior hyperparameters.
-        optimizer (ox.GradientTransformation): Optax optimizer to use for optimizing the
-        posterior hyperparameters.
-        num_optimization_iterations (int): Number of iterations to optimize
-        the posterior hyperparameters for.
+        prior: prior to use when forming the posterior.
+        likelihood_builder: function which takes the number of datapoints as input and
+            returns a likelihood object initialised with the given number of datapoints.
+        optimization_objective: objective to use for optimizing the posterior hyperparameters.
+        optimizer: an optax optimizer to use for optimizing the posterior hyperparameters.
+        num_optimization_iterations: the number of iterations to optimize
+            the posterior hyperparameters for.
     """
 
     prior: AbstractPrior
     likelihood_builder: LikelihoodBuilder
-    optimization_objective: AbstractObjective
+    optimization_objective: Objective
     optimizer: ox.GradientTransformation
     num_optimization_iters: int
 
@@ -71,10 +68,10 @@ class PosteriorHandler:
         Initialise (and optionally optimize) a posterior using the given dataset.
 
         Args:
-            dataset (Dataset): Dataset to get posterior for.
-            optimize (bool): Whether to optimize the posterior hyperparameters.
-            key (Optional[KeyArray]): A JAX PRNG key which is used for optimizing the posterior
-            hyperparameters.
+            dataset: dataset to get posterior for.
+            optimize: whether to optimize the posterior hyperparameters.
+            key: a JAX PRNG key which is used for optimizing the posterior
+                hyperparameters.
 
         Returns:
             Posterior for the given dataset.
@@ -108,14 +105,14 @@ class PosteriorHandler:
         set as in the `likelihood_builder` function.
 
         Args:
-            dataset: Dataset to get posterior for.
-            previous_posterior: Posterior being updated. This is supplied as one may
-            wish to simply increase the number of datapoints in the likelihood, without
-            optimizing the posterior hyperparameters, in which case the previous
-            posterior can be used to obtain the previously set prior hyperparameters.
-            optimize: Whether to optimize the posterior hyperparameters.
+            dataset: dataset to get posterior for.
+            previous_posterior: posterior being updated. This is supplied as one may
+                wish to simply increase the number of datapoints in the likelihood, without
+                optimizing the posterior hyperparameters, in which case the previous
+                posterior can be used to obtain the previously set prior hyperparameters.
+            optimize: whether to optimize the posterior hyperparameters.
             key: A JAX PRNG key which is used for optimizing the posterior
-            hyperparameters.
+                hyperparameters.
         """
         posterior = previous_posterior.prior * self.likelihood_builder(dataset.n)
 

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/decision_making/search_space.py

@@ -56,7 +56,7 @@ class AbstractSearchSpace(ABC):
 
 @dataclass
 class ContinuousSearchSpace(AbstractSearchSpace):
-    """The `ContinuousSearchSpace` class is used to bound the domain of continuous real functions of dimension $`D`$."""
+    """The `ContinuousSearchSpace` class is used to bound the domain of continuous real functions of dimension $D$."""
 
     lower_bounds: Float[Array, " D"]
     upper_bounds: Float[Array, " D"]

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/decision_making/test_functions/continuous_functions.py

@@ -12,24 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from abc import (
-    ABC,
-    abstractmethod,
-)
+from abc import abstractmethod
 from dataclasses import dataclass
 
 import jax.numpy as jnp
 from jaxtyping import (
     Array,
     Float,
+    Num,
 )
+import tensorflow_probability.substrates.jax as tfp
 
 from gpjax.dataset import Dataset
 from gpjax.decision_making.search_space import ContinuousSearchSpace
+from gpjax.gps import AbstractMeanFunction
 from gpjax.typing import KeyArray
 
 
-class AbstractContinuousTestFunction(ABC):
+class AbstractContinuousTestFunction(AbstractMeanFunction):
     """
     Abstract base class for continuous test functions.
 
@@ -43,19 +43,28 @@ class AbstractContinuousTestFunction(ABC):
     minimizer: Float[Array, "1 D"]
     minimum: Float[Array, "1 1"]
 
-    def generate_dataset(self, num_points: int, key: KeyArray) -> Dataset:
+    def generate_dataset(
+        self, num_points: int, key: KeyArray, obs_stddev: float = 0.0
+    ) -> Dataset:
         """
         Generate a toy dataset from the test function.
 
         Args:
             num_points (int): Number of points to sample.
             key (KeyArray): JAX PRNG key.
+            obs_stddev (float): (Optional) standard deviation of Gaussian distributed
+            noise added to observations.
 
         Returns:
             Dataset: Dataset of points sampled from the test function.
         """
         X = self.search_space.sample(num_points=num_points, key=key)
-        y = self.evaluate(X)
+        gaussian_noise = tfp.distributions.Normal(
+            jnp.zeros(num_points), obs_stddev * jnp.ones(num_points)
+        )
+        y = self.evaluate(X) + jnp.transpose(
+            gaussian_noise.sample(sample_shape=[1], seed=key)
+        )
         return Dataset(X=X, y=y)
 
     def generate_test_points(
@@ -73,6 +82,9 @@ class AbstractContinuousTestFunction(ABC):
         """
         return self.search_space.sample(num_points=num_points, key=key)
 
+    def __call__(self, x: Num[Array, "N D"]) -> Float[Array, "N 1"]:
+        return self.evaluate(x)
+
     @abstractmethod
     def evaluate(self, x: Float[Array, "N D"]) -> Float[Array, "N 1"]:
         """
@@ -92,7 +104,7 @@ class Forrester(AbstractContinuousTestFunction):
     """
     Forrester function introduced in 'Engineering design via surrogate modelling: a
     practical guide' (Forrester et al. 2008), rescaled to have zero mean and unit
-    variance over $`[0, 1]`$.
+    variance over $[0, 1]$.
     """
 
     search_space = ContinuousSearchSpace(
@@ -113,7 +125,7 @@ class LogarithmicGoldsteinPrice(AbstractContinuousTestFunction):
     """
     Logarithmic Goldstein-Price function introduced in 'A benchmark of kriging-based
     infill criteria for noisy optimization' (Picheny et al. 2013), which has zero mean
-    and unit variance over $`[0, 1]^2`$.
+    and unit variance over $[0, 1]^2$.
     """
 
     search_space = ContinuousSearchSpace(
@@ -148,7 +160,7 @@ class LogarithmicGoldsteinPrice(AbstractContinuousTestFunction):
 @dataclass
 class Quadratic(AbstractContinuousTestFunction):
     """
-    Toy quadratic function defined over $`[0, 1]`$.
+    Toy quadratic function defined over $[0, 1]$.
     """
 
     search_space = ContinuousSearchSpace(

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/decision_making/test_functions/non_conjugate_functions.py

@@ -17,15 +17,15 @@ from dataclasses import dataclass
 
 import jax.numpy as jnp
 import jax.random as jr
-from jaxtyping import (
-    Array,
-    Float,
-    Integer,
-)
 
 from gpjax.dataset import Dataset
 from gpjax.decision_making.search_space import ContinuousSearchSpace
-from gpjax.typing import KeyArray
+from gpjax.typing import (
+    Array,
+    Float,
+    Int,
+    KeyArray,
+)
 
 
 @dataclass
@@ -74,7 +74,7 @@ class PoissonTestFunction:
         return self.search_space.sample(num_points=num_points, key=key)
 
     @abstractmethod
-    def evaluate(self, x: Float[Array, "N 1"]) -> Integer[Array, "N 1"]:
+    def evaluate(self, x: Float[Array, "N 1"]) -> Int[Array, "N 1"]:
         """
         Evaluate the test function at a set of points. Function taken from
         https://docs.jaxgaussianprocesses.com/examples/poisson/#dataset.
@@ -83,7 +83,7 @@ class PoissonTestFunction:
             x (Float[Array, 'N D']): Points to evaluate the test function at.
 
         Returns:
-            Integer[Array, 'N 1']: Values of the test function at the points.
+            Float[Array, 'N 1']: Values of the test function at the points.
         """
         key = jr.key(42)
         f = lambda x: 2.0 * jnp.sin(3 * x) + 0.5 * x

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/decision_making/utility_functions/__init__.py

@@ -18,12 +18,20 @@ from gpjax.decision_making.utility_functions.base import (
     SinglePointUtilityFunction,
     UtilityFunction,
 )
+from gpjax.decision_making.utility_functions.expected_improvement import (
+    ExpectedImprovement,
+)
+from gpjax.decision_making.utility_functions.probability_of_improvement import (
+    ProbabilityOfImprovement,
+)
 from gpjax.decision_making.utility_functions.thompson_sampling import ThompsonSampling
 
 __all__ = [
     "UtilityFunction",
     "AbstractUtilityFunctionBuilder",
     "AbstractSinglePointUtilityFunctionBuilder",
+    "ExpectedImprovement",
     "SinglePointUtilityFunction",
     "ThompsonSampling",
+    "ProbabilityOfImprovement",
 ]

{gpjax-0.8.1 → gpjax-0.9.0}/gpjax/decision_making/utility_functions/base.py

@@ -35,8 +35,8 @@ from gpjax.typing import (
 SinglePointUtilityFunction = Callable[[Float[Array, "N D"]], Float[Array, "N 1"]]
 """
 Type alias for utility functions which don't support batching, and instead characterise
-the utility of querying a single point, rather than a batch of points. They take an array of points of shape $`[N, D]`$
-and return the value of the utility function at each point in an array of shape $`[N, 1]`$.
+the utility of querying a single point, rather than a batch of points. They take an array of points of shape $[N, D]$
+and return the value of the utility function at each point in an array of shape $[N, 1]$.
 """
 
 
@@ -65,14 +65,12 @@ class AbstractSinglePointUtilityFunctionBuilder(ABC):
         datasets.
 
         Args:
-            posteriors (Mapping[str, AbstractPosterior]): Dictionary of posteriors to be
-            used to form the utility function.
-            datasets (Mapping[str, Dataset]): Dictionary of datasets which may be used
-            to form the utility function.
+            posteriors: dictionary of posteriors to be used to form the utility function.
+            datasets: dictionary of datasets which may be used to form the utility function.
 
         Raises:
             ValueError: If the objective posterior or dataset are not present in the
-            posteriors or datasets.
+                posteriors or datasets.
         """
         if OBJECTIVE not in posteriors.keys():
             raise ValueError("Objective posterior not found in posteriors")
@@ -90,15 +88,13 @@ class AbstractSinglePointUtilityFunctionBuilder(ABC):
         Build a `UtilityFunction` from a set of posteriors and datasets.
 
         Args:
-            posteriors (Mapping[str, AbstractPosterior]): Dictionary of posteriors to be
-            used to form the utility function.
-            datasets (Mapping[str, Dataset]): Dictionary of datasets which may be used
-            to form the utility function.
-            key (KeyArray): JAX PRNG key used for random number generation.
+            posteriors: dictionary of posteriors to be used to form the utility function.
+            datasets: dictionary of datasets which may be used to form the utility function.
+            key: JAX PRNG key used for random number generation.
 
         Returns:
             SinglePointUtilityFunction: Utility function to be *maximised* in order to
-            decide which point to query next.
+                decide which point to query next.
         """
         raise NotImplementedError