gpjax 0.8.2__py3-none-any.whl → 0.9.1__py3-none-any.whl

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.2"
+__version__ = "0.9.1"
 
 __all__ = [
     "base",
@@ -49,6 +51,7 @@ __all__ = [
     "likelihoods",
     "mean_functions",
     "objectives",
+    "parameters",
     "variational_families",
     "Dataset",
     "cite",

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/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/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/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/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/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(
@@ -127,12 +139,7 @@ class LogarithmicGoldsteinPrice(AbstractContinuousTestFunction):
         x1 = 4.0 * x[:, 0] - 2.0
         x2 = 4.0 * x[:, 1] - 2.0
         a = 1.0 + (x1 + x2 + 1.0) ** 2 * (
-            19.0
-            - 14.0 * x1
-            + 3.0 * (x1**2)
-            - 14.0 * x2
-            + 6.0 * x1 * x2
-            + 3.0 * (x2**2)
+            19.0 - 14.0 * x1 + 3.0 * (x1**2) - 14.0 * x2 + 6.0 * x1 * x2 + 3.0 * (x2**2)
         )
         b = 30.0 + (2.0 * x1 - 3.0 * x2) ** 2 * (
             18.0
@@ -148,7 +155,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/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/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/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
 

gpjax/decision_making/utility_functions/expected_improvement.py

@@ -0,0 +1,112 @@
+# Copyright 2024 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 dataclasses import dataclass
+from functools import partial
+
+from beartype.typing import Mapping
+import jax.numpy as jnp
+import tensorflow_probability.substrates.jax as tfp
+
+from gpjax.dataset import Dataset
+from gpjax.decision_making.utility_functions.base import (
+    AbstractSinglePointUtilityFunctionBuilder,
+    SinglePointUtilityFunction,
+)
+from gpjax.decision_making.utils import (
+    OBJECTIVE,
+    get_best_latent_observation_val,
+)
+from gpjax.gps import ConjugatePosterior
+from gpjax.typing import (
+    Array,
+    Float,
+    KeyArray,
+)
+
+
+@dataclass
+class ExpectedImprovement(AbstractSinglePointUtilityFunctionBuilder):
+    """
+    Expected Improvement acquisition function as introduced by [Močkus,
+    1974](https://link.springer.com/chapter/10.1007/3-540-07165-2_55). The "best"
+    incumbent value is defined as the lowest posterior mean value evaluated at the the
+    previously observed points. This enables the acquisition function to be utilised with noisy observations.
+    """
+
+    def build_utility_function(
+        self,
+        posteriors: Mapping[str, ConjugatePosterior],
+        datasets: Mapping[str, Dataset],
+        key: KeyArray,
+    ) -> SinglePointUtilityFunction:
+        r"""
+        Build the Expected Improvement acquisition function. This computes the expected
+        improvement over the "best" of the previously observed points, utilising the
+        posterior distribution of the surrogate model. For posterior distribution
+        $`f(\cdot)`$, and best incumbent value $`\eta`$, this is defined
+        as:
+        ```math
+        \alpha_{\text{EI}}(\mathbf{x}) = \mathbb{E}\left[\max(0, \eta - f(\mathbf{x}))\right]
+        ```
+
+        Args:
+            posteriors (Mapping[str, ConjugatePosterior]): Dictionary of posteriors to
+            used to form the utility function. One posteriors must correspond to the
+            `OBJECTIVE` key, as we utilise the objective posterior to form the utility
+            function.
+            datasets (Mapping[str, Dataset]): Dictionary of datasets used to form the
+            utility function. Keys in `datasets` should correspond to keys in
+            `posteriors`. One of the datasets must correspond to the `OBJECTIVE` key.
+            key (KeyArray): JAX PRNG key used for random number generation.
+
+        Returns:
+            SinglePointUtilityFunction: The Expected Improvement acquisition function to
+            to be *maximised* in order to decide which point to query next.
+        """
+        self.check_objective_present(posteriors, datasets)
+        objective_posterior = posteriors[OBJECTIVE]
+        objective_dataset = datasets[OBJECTIVE]
+
+        if not isinstance(objective_posterior, ConjugatePosterior):
+            raise ValueError(
+                "Objective posterior must be a ConjugatePosterior to compute the Expected Improvement."
+            )
+
+        if (
+            objective_dataset.X is None
+            or objective_dataset.n == 0
+            or objective_dataset.y is None
+        ):
+            raise ValueError("Objective dataset must contain at least one item")
+
+        eta = get_best_latent_observation_val(objective_posterior, objective_dataset)
+        return partial(
+            _expected_improvement, objective_posterior, objective_dataset, eta
+        )
+
+
+def _expected_improvement(
+    objective_posterior: ConjugatePosterior,
+    objective_dataset: Dataset,
+    eta: Float[Array, ""],
+    x: Float[Array, "N D"],
+) -> Float[Array, "N 1"]:
+    latent_dist = objective_posterior(x, objective_dataset)
+    mean = latent_dist.mean()
+    var = latent_dist.variance()
+    normal = tfp.distributions.Normal(mean, jnp.sqrt(var))
+    return jnp.expand_dims(
+        ((eta - mean) * normal.cdf(eta) + var * normal.prob(eta)), -1
+    )