Refactor acquisition function minimally by nabenabe0928 · Pull Request #6166 · optuna/optuna

nabenabe0928 · 2025-06-18T05:28:50Z

Motivation

This PR introduces the abstraction of the acquisition function for GPSampler.

Description of the changes

Eliminate data class
Make the acquisition function abstract

nabenabe0928 · 2025-06-18T05:31:16Z

@sawa3030 @contramundum53
Could you review this PR?

Copilot

Pull Request Overview

A refactor to introduce a class-based abstraction for acquisition functions in GPSampler and remove the old dataclass/API.

Replace the AcquisitionFunctionParams and free-function API with BaseAcquisitionFunc subclasses (LogEI, LogPI, UCB, LCB, ConstrainedLogEI, LogEHVI).
Update all samplers, terminators, tests, and optimization helpers to use the new class-based API.
Extend GPRegressor with internal caching (_cache_matrix) and simplify its kernel/posterior interfaces.

Reviewed Changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/samplers_tests/test_gp.py	Updated import and use of `LogEI` class.
tests/gp_tests/test_gp.py	Adjusted `GPRegressor` instantiation to new signature.
tests/gp_tests/test_acqf.py	Switched tests to the new `BaseAcquisitionFunc` API.
optuna/terminator/improvement/evaluator.py	Replaced `create_acqf_params` with class-based calls.
optuna/terminator/improvement/emmr.py	Removed old acqf_params code; use `gpr.posterior`.
optuna/samplers/_gp/sampler.py	Refactored sampler methods to accept `BaseAcquisitionFunc`.
optuna/_gp/optim_sample.py	Deleted obsolete `optimize_acqf_sample`.
optuna/_gp/optim_mixed.py	Updated helpers to use `BaseAcquisitionFunc`.
optuna/_gp/gp.py	Added caching, revised `kernel` & `posterior` API.
optuna/_gp/acqf.py	Implemented `BaseAcquisitionFunc` and subclasses.

Comments suppressed due to low confidence (2)

optuna/samplers/_gp/sampler.py:232

Newly constructed GPRegressor instances in constraints_gprs lack a call to _cache_matrix(), so their posterior will assert on missing cached matrices. Call gpr._cache_matrix() right after instantiation before appending.

optuna/_gp/acqf.py:231

The stabilizing_noise annotation appears in the function body instead of the __init__ signature, causing a syntax error. Move it into the parameter list and remove the misplaced closing parenthesis.

        stabilizing_noise: float = 1e-12,

contramundum53

LGTM.

nabenabe0928 · 2025-06-20T02:16:33Z

import optuna


def objective(trial: optuna.Trial) -> float:
    x = trial.suggest_float("x", -5, 5)
    y = trial.suggest_float("y", -5, 5)
    return x**2 + y**2


def multi_objective(trial: optuna.Trial) -> tuple[float, float]:
    x = trial.suggest_float("x", -5, 5)
    y = trial.suggest_float("y", -5, 5)
    return x**2 + y**2, (x - 2)**2 + (y - 2)**2


def constraints(trial: optuna.trial.FrozenTrial) -> tuple[float, float]:
    x = trial.params["x"]
    y = trial.params["y"]
    return (x - 2, y - 2)


mode = ["single", "multi", "constr"][2]
if mode == "single":
    study = optuna.create_study(sampler=optuna.samplers.GPSampler(seed=0))
    obj_func = objective
elif mode == "multi":
    study = optuna.create_study(sampler=optuna.samplers.GPSampler(seed=0), directions=["minimize"]*2)
    obj_func = multi_objective
elif mode == "constr":
    study = optuna.create_study(sampler=optuna.samplers.GPSampler(seed=0, constraints_func=constraints))
    obj_func = objective

study.optimize(obj_func, n_trials=20)
trials = study.trials
print((trials[-1].datetime_complete - trials[0].datetime_start).total_seconds())

The benchmarking code to check the reproducibility.

sawa3030 · 2025-06-20T02:26:43Z

optuna/_gp/acqf.py

-def logpi(mean: torch.Tensor, var: torch.Tensor, f0: float) -> torch.Tensor:
-    # Return the integral of N(mean, var) from -inf to f0
-    # This is identical to the integral of N(0, 1) from -inf to (f0-mean)/sigma
-    # Return E_{y ~ N(mean, var)}[bool(y <= f0)]
-    sigma = torch.sqrt(var)
-    return torch.special.log_ndtr((f0 - mean) / sigma)
-
-
-def ucb(mean: torch.Tensor, var: torch.Tensor, beta: float) -> torch.Tensor:
-    return mean + torch.sqrt(beta * var)
-
-
-def lcb(mean: torch.Tensor, var: torch.Tensor, beta: float) -> torch.Tensor:
-    return mean - torch.sqrt(beta * var)


[note] Reviewed this part

sawa3030

These are my notes to help with understanding and to keep track of what to review.

sawa3030 · 2025-06-20T02:28:36Z

optuna/_gp/acqf.py

-# TODO(contramundum53): consider abstraction for acquisition functions.
-# NOTE: Acquisition function is not class on purpose to integrate numba in the future.
-class AcquisitionFunctionType(IntEnum):
-    LOG_EI = 0
-    UCB = 1
-    LCB = 2
-    LOG_PI = 3
-    LOG_EHVI = 4


[note]: Reviewed this part

sawa3030 · 2025-06-20T02:39:42Z

optuna/_gp/gp.py

+        assert (
+            self._cov_Y_Y_inv is not None and self._cov_Y_Y_inv_Y is not None
+        ), "Call cache_matrix before calling posterior."
+        cov_fx_fX = self.kernel(x[..., None, :], self._X_train)[..., 0, :]


[TODO]: Check if X equals to self._X_train

sawa3030 · 2025-06-20T02:41:36Z

optuna/_gp/gp.py

        # We apply the cholesky decomposition to efficiently compute log(|C|) and C^-1.

-        cov_fX_fX = self.kernel(is_categorical, X, X)
+        cov_fX_fX = self.kernel(self._X_train, self._X_train)


[TODO]: Check if X equals to self._X_train

sawa3030 · 2025-06-20T02:42:39Z

optuna/_gp/gp.py

-            :, 0
-        ]
+        cov_Y_Y_chol_inv_Y = torch.linalg.solve_triangular(
+            cov_Y_Y_chol, self._y_train[:, None], upper=False


[TODO]: Check if Y equals to self._y_train

sawa3030 · 2025-06-20T03:29:00Z

optuna/_gp/acqf.py

-            beta=acqf_params.beta,
-            acqf_stabilizing_noise=acqf_params.acqf_stabilizing_noise,
-            acqf_params_for_constraints=acqf_params_for_constraints,
+class BaseAcquisitionFunc(ABC):


[TODO]: Verify that each member function in every AcquisitionFunction subclass corresponds to the appropriate function in the previous implementation.

Checked LogEI, UCB, LCB, and LGEHVI

Checked ConstrainedLogEI and LogPI

sawa3030 · 2025-06-20T03:33:33Z

optuna/_gp/acqf.py

-        return f_val
-
-
-def eval_acqf_no_grad(acqf_params: AcquisitionFunctionParams, x: np.ndarray) -> np.ndarray:


[TODO]: Review the implementation of eval_acqf_no_grad in this PR.

sawa3030 · 2025-06-20T03:33:55Z

optuna/_gp/acqf.py

-        return eval_acqf(acqf_params, torch.from_numpy(x)).detach().numpy()
-
-
-def eval_acqf_with_grad(


[TODO]: Review the implementation of eval_acqf_with_grad in this PR.

sawa3030 · 2025-06-27T02:02:18Z

import optuna


def objective(trial: optuna.Trial) -> float:
    x = trial.suggest_float("x", -5, 5)
    y = trial.suggest_float("y", -5, 5)
    return x**2 + y**2


def multi_objective(trial: optuna.Trial) -> tuple[float, float]:
    x = trial.suggest_float("x", -5, 5)
    y = trial.suggest_float("y", -5, 5)
    return x**2 + y**2, (x - 2)**2 + (y - 2)**2


def constraints(trial: optuna.trial.FrozenTrial) -> tuple[float, float]:
    x = trial.params["x"]
    y = trial.params["y"]
    return (x - 2, y - 2)


mode = ["single", "multi", "constr"][2]
if mode == "single":
    study = optuna.create_study(sampler=optuna.samplers.GPSampler(seed=0))
    obj_func = objective
elif mode == "multi":
    study = optuna.create_study(sampler=optuna.samplers.GPSampler(seed=0), directions=["minimize"]*2)
    obj_func = multi_objective
elif mode == "constr":
    study = optuna.create_study(sampler=optuna.samplers.GPSampler(seed=0, constraints_func=constraints))
    obj_func = objective

study.optimize(obj_func, n_trials=20)
trials = study.trials
print((trials[-1].datetime_complete - trials[0].datetime_start).total_seconds())

The benchmarking code to check the reproducibility.

note: A slight difference in results between the master branch and this PR was observed during multi-objective optimization. This discrepancy stems from minor errors in numerical calculations between NumPy and PyTorch. This difference is negligible and can be safely ignored in this context.

sawa3030

Thank you for your support in reviewing this PR. LGTM

nabenabe0928 · 2025-06-27T02:26:55Z

@sawa3030 @contramundum53
Thank you for your reviews, I will merge this PR:)

Minimal Acqf Refactor

e58f297

nabenabe0928 marked this pull request as ready for review June 18, 2025 05:28

nabenabe0928 assigned contramundum53 and sawa3030 Jun 18, 2025

nabenabe0928 added the code-fix Change that does not change the behavior, such as code refactoring. label Jun 18, 2025

nabenabe0928 added this to the v4.5.0 milestone Jun 18, 2025

contramundum53 requested a review from Copilot June 18, 2025 05:30

Copilot AI reviewed Jun 18, 2025

View reviewed changes

nabenabe0928 added 2 commits June 18, 2025 07:39

Apply wang's comment

db39da1

Fix CI

587d34f

contramundum53 approved these changes Jun 18, 2025

View reviewed changes

nabenabe0928 unassigned contramundum53 Jun 19, 2025

sawa3030 reviewed Jun 20, 2025

View reviewed changes

nabenabe0928 mentioned this pull request Jun 25, 2025

Include scipy-stubs in the type-check dependencies #6174

Merged

sawa3030 approved these changes Jun 27, 2025

View reviewed changes

nabenabe0928 merged commit 684f10a into optuna:master Jun 27, 2025
14 checks passed

nabenabe0928 unassigned sawa3030 Jun 27, 2025

This was referenced Jun 27, 2025

Improve the notational consistency in GP-related docs #6183

Closed

Fix stale docs in GP-related modules #6184

Merged

nabenabe0928 mentioned this pull request Oct 16, 2025

Developments of Gaussian-Process Based Bayesian Optimization (GPSampler) nabenabe0928/my-skills#1

Open

		return f_val


		def eval_acqf_no_grad(acqf_params: AcquisitionFunctionParams, x: np.ndarray) -> np.ndarray:

		return eval_acqf(acqf_params, torch.from_numpy(x)).detach().numpy()


		def eval_acqf_with_grad(

Uh oh!

Conversation

nabenabe0928 commented Jun 18, 2025

Motivation

Description of the changes

Uh oh!

nabenabe0928 commented Jun 18, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

contramundum53 left a comment

Choose a reason for hiding this comment

Uh oh!

nabenabe0928 commented Jun 20, 2025

Uh oh!

sawa3030 Jun 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sawa3030 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sawa3030 commented Jun 27, 2025

Uh oh!

sawa3030 left a comment

Choose a reason for hiding this comment

Uh oh!

nabenabe0928 commented Jun 27, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

sawa3030 Jun 20, 2025 •

edited

Loading