Feature: Adds probability of improvement as an acquisition function #458

miguelgondu · 2024-07-01T12:38:53Z

Type of changes

New feature
Documentation / docstrings
Tests

Checklist

I've formatted the new code by running poetry run pre-commit run --all-files --show-diff-on-failure before committing.
I've added tests for new code.
I've added docstrings for the new code.

Checklist for this PR in particular

Adding tests particular to the PI acquisition function.
Decide on whether we want to call these UtilityFunctions or AcquisitionFunctions. (We decided to keep them as utility_functions)
Decide on whether we want to write a whole new tutorial for these, or just expand the tutorial on BO. (We decided to add to an existing tutorial)
Add docstrings.

Description

This PR expands the decision_making module by adding Probability of Improvement as a potential acquisition function. It refactors the tests that are joint to utility functions, and provides a tutorial for how to use Probability of Improvement in practice (pattern-matching from the current tutorial on Thompson Sampling).

Issue Number: N/A

…unctions

github-actions

Thank you for opening your first PR into GPJax!

If you have not heard from us in a while, please feel free to ping @gpjax/developers or anyone who has commented on the PR. Most of our reviewers are volunteers and sometimes things fall through the cracks.

You can also join us on Slack for real-time discussion.

For details on testing, writing docs, and our review process, please see the developer guide

We strive to be a welcoming and open project. Please follow our Code of Conduct.

miguelgondu · 2024-07-01T12:49:33Z

Hi @gpjax/developers,

I'm close to done with adding Probability of Improvement (PI) as an acquisition function. Before I finish this PR I would like to ask two questions:

Should I create an entire new tutorial for BayesOpt beyond Thompson Sampling? Should I just update the current tutorial with a brief example of how to use PI? I vote for the latter. Currently I'm writing a new tutorial, but there's plenty of overlap with the current one on BO. It's worth noting that the current tutorial doesn't use the built-in functionalities for optimizing the acquisition function.
I'm currently adding PI under the same folder as Thompson Sampling (i.e. utility_functions). I feel a better name for them might be acquisition_functions, but I'm happy either way. Let me know if I should be providing it under a different model than the abstract Utility Function.

Thomas-Christie

Nice job - have just left a few minor comments.

gpjax/decision_making/utility_functions/probability_of_improvement.py

Thomas-Christie · 2024-07-02T21:03:31Z

gpjax/decision_making/utility_functions/probability_of_improvement.py

+                "Objective posterior must be a ConjugatePosterior to draw an approximate sample."
+            )
+
+        objective_dataset = datasets[OBJECTIVE]


Probably best to have something along the lines of

if objective_dataset.X is None or objective_dataset.n == 0: raise ValueError("Objective dataset must contain at least one item")

given that we use the objective dataset to find best_y.

Thomas-Christie · 2024-07-02T21:05:18Z

gpjax/decision_making/utility_functions/probability_of_improvement.py

+            predictive_dist = objective_posterior.predict(x_test, objective_dataset)
+
+            # Assuming that the goal is to minimize the objective function
+            best_y = objective_dataset.y.min()


It could be useful to define best_y as the minimum posterior mean value at any of the observed points, which would also handle the case where observations could be noisy.

Sounds good. I've updated how we compute best_y. It is the first time I see it being computed like that, though. If I'm not mistaken, other frameworks use the min of the dataset.

Yeah - for instance this approach is mentioned in "A benchmark of kriging-based infill criteria for noisy
optimization", section 3.2.

Thomas-Christie · 2024-07-02T21:10:03Z

gpjax/decision_making/utility_functions/probability_of_improvement.py

+            # Assuming that the goal is to minimize the objective function
+            best_y = objective_dataset.y.min()
+
+            return gaussian_cdf(


Instead of implementing our own gaussian_cdf function, IMO it makes sense to use the existing normal distribution provided by Tensorflow probability i.e. normal = tfp.distributions.Normal(predicitve_dist.mean(), predictive_dist.stddev()) and use the cdf defined on this.

Another alternative would be to update GaussianDistribution, right? For now, I'll create a normal as you state.

Thomas-Christie · 2024-07-02T21:13:58Z

tests/test_decision_making/test_utility_functions/test_utility_functions.py

+)
+
+
+@pytest.mark.filterwarnings(


Nice, this will help make testing additional utility functions a lot easier.

Thomas-Christie · 2024-07-02T21:15:11Z

tests/test_decision_making/test_utility_functions/test_utility_functions.py

+    posterior = generate_dummy_conjugate_posterior(dataset)
+    posteriors = {OBJECTIVE: posterior}
+    datasets = {OBJECTIVE: dataset}
+    ts_utility_builder = utility_function_builder(**utility_function_kwargs)


nit: Some of the names of variables in this file e.g. ts_utility_builder could do with updating now that these tests are more "utility function agnostic".

…ge on PI

miguelgondu · 2024-07-03T08:59:06Z

Hi @Thomas-Christie,

Thanks for the review! I've addressed all the comments. Could I ask you to take a second look? Let me know if anything else comes to mind.

Cheers.

Thomas-Christie · 2024-07-03T20:32:40Z

gpjax/decision_making/utility_functions/probability_of_improvement.py

+    More precisely, given a predictive posterior distribution of the objective
+    function $`f`$, the probability of improvement at a test point $`x`$ is defined as:
+    $$`\text{PI}(x) = \text{Prob}[f(x) < f(x_{\text{best}})]`$$
+    where $`x_{\text{best}}`$ is the minimizer of $`f`$ in the dataset.


Could update to something like where $x_{\text{best}}$ is the minimiser of the posterior mean at previously observed values, to handle noisy observations.

Updated! Forgot to change this.

Thomas-Christie · 2024-07-03T20:35:29Z

tests/test_decision_making/test_utility_functions/test_probability_of_improvement.py

+    assert jnp.isclose(utility_values, expected_utility_values).all()
+
+
+if __name__ == "__main__":


nit: I think you forgot to remove this after debugging.

Absolutely. My bad.

miguelgondu · 2024-07-05T08:10:56Z

All green! 🎉

thomaspinder · 2024-07-06T06:37:41Z

Super work @miguelgondu ! Thanks for reviewing @Thomas-Christie

miguelgondu added 4 commits June 30, 2024 11:38

Adds probability of improvement, refactors common tests for utility f…

eadb7bb

…unctions

Finishes a first draft of the tutorial

7ec1f89

Runs pre-commit hooks

c67083f

Fixes a Ruff error on jaxtyping annotation

e0fc9d2

miguelgondu marked this pull request as draft July 1, 2024 12:39

github-actions bot reviewed Jul 1, 2024

View reviewed changes

Adds docstrings to probability of improvement

623181a

miguelgondu changed the title ~~[Work in progress] Feature: Adds probability of improvement as an acquisition function~~ [WIP] Feature: Adds probability of improvement as an acquisition function Jul 1, 2024

miguelgondu added 5 commits July 1, 2024 15:38

Adds a simple test for PI

5d0a84c

Changes the tutorial a bit before deleting it

3924fa9

Updates the tutorial on decision making by mentioning PI

ba5bbdc

Makes a better test for PI by manually computing the CDF

b455c9a

Lints according to Ruff

1f67bfb

miguelgondu changed the title ~~[WIP] Feature: Adds probability of improvement as an acquisition function~~ Feature: Adds probability of improvement as an acquisition function Jul 1, 2024

miguelgondu marked this pull request as ready for review July 1, 2024 19:15

Thomas-Christie suggested changes Jul 2, 2024

View reviewed changes

miguelgondu added 3 commits July 3, 2024 10:46

Removes the manual CDF and replaces it with tfp, improves error messa…

64aeb2c

…ge on PI

Removes unused imports in utils

b651e90

Removes test of manual CDF

ceef630

miguelgondu requested a review from Thomas-Christie July 3, 2024 08:59

Thomas-Christie suggested changes Jul 3, 2024

View reviewed changes

Updates the documentation to reflect changes, removes if main

58d7bd2

miguelgondu requested a review from Thomas-Christie July 3, 2024 20:39

Thomas-Christie approved these changes Jul 3, 2024

View reviewed changes

thomaspinder approved these changes Jul 4, 2024

View reviewed changes

thomaspinder merged commit fc45b9f into JaxGaussianProcesses:main Jul 6, 2024
11 checks passed

miguelgondu deleted the feature-adds-probability-of-improvement branch July 6, 2024 06:44

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feature: Adds probability of improvement as an acquisition function #458

Feature: Adds probability of improvement as an acquisition function #458

miguelgondu commented Jul 1, 2024 •

edited

Loading

github-actions bot left a comment

miguelgondu commented Jul 1, 2024 •

edited

Loading

Thomas-Christie left a comment

Thomas-Christie Jul 2, 2024

miguelgondu Jul 3, 2024

Thomas-Christie Jul 2, 2024

miguelgondu Jul 3, 2024

Thomas-Christie Jul 3, 2024

Thomas-Christie Jul 2, 2024

miguelgondu Jul 3, 2024

Thomas-Christie Jul 2, 2024

miguelgondu Jul 3, 2024

Thomas-Christie Jul 2, 2024

miguelgondu Jul 3, 2024

miguelgondu commented Jul 3, 2024

Thomas-Christie Jul 3, 2024

miguelgondu Jul 3, 2024

Thomas-Christie Jul 3, 2024

miguelgondu Jul 3, 2024

miguelgondu commented Jul 5, 2024 •

edited

Loading

thomaspinder commented Jul 6, 2024

		assert jnp.isclose(utility_values, expected_utility_values).all()


		if __name__ == "__main__":

		)


		@pytest.mark.filterwarnings(

Feature: Adds probability of improvement as an acquisition function #458

Feature: Adds probability of improvement as an acquisition function #458

Conversation

miguelgondu commented Jul 1, 2024 • edited Loading

Type of changes

Checklist

Checklist for this PR in particular

Description

github-actions bot left a comment

Choose a reason for hiding this comment

miguelgondu commented Jul 1, 2024 • edited Loading

Thomas-Christie left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

miguelgondu commented Jul 3, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

miguelgondu commented Jul 5, 2024 • edited Loading

thomaspinder commented Jul 6, 2024

miguelgondu commented Jul 1, 2024 •

edited

Loading

miguelgondu commented Jul 1, 2024 •

edited

Loading

miguelgondu commented Jul 5, 2024 •

edited

Loading