Skip to content

cds2.bayesopt

Gaussian Process and Bayesian optimization.

Bayesian optimization with Gaussian-process surrogate and acquisition functions.

GaussianProcess

Gaussian-process regressor with RBF kernel.

Uses the squared-exponential kernel k(x, x') = sigma_f^2 * exp(-0.5 * ||x-x'||^2 / length_scale^2) with a small observation noise added to the diagonal.

Source code in src/cds2/bayesopt.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
class GaussianProcess:
    """Gaussian-process regressor with RBF kernel.

    Uses the squared-exponential kernel
    ``k(x, x') = sigma_f^2 * exp(-0.5 * ||x-x'||^2 / length_scale^2)``
    with a small observation noise added to the diagonal.
    """

    def __init__(
        self,
        length_scale: float = 1.0,
        sigma_f: float = 1.0,
        noise: float = 1e-8,
    ) -> None:
        if length_scale <= 0 or sigma_f <= 0 or noise < 0:
            msg = "length_scale and sigma_f must be >0, noise >=0"  # pragma: no cover
            raise ValueError(msg)  # pragma: no cover
        self.length_scale: float = float(length_scale)
        self.sigma_f: float = float(sigma_f)
        self.noise: float = float(noise)
        self.X_train_: FloatArray | None = None
        self.y_train_: FloatArray | None = None
        self.L_: FloatArray | None = None
        self.alpha_: FloatArray | None = None

    def fit(self, x_train: object, y_train: object) -> GaussianProcess:
        """Fit the GP to training data."""
        x_arr = _as_2d(x_train)
        if x_arr.ndim != 2:
            msg = "x_train must be 2-D"  # pragma: no cover
            raise ValueError(msg)  # pragma: no cover
        y_arr = np.asarray(y_train, dtype=float).ravel()
        if y_arr.size != x_arr.shape[0]:
            msg = "x_train and y_train must have matching first dimension"  # pragma: no cover
            raise ValueError(msg)  # pragma: no cover
        if x_arr.shape[0] == 0:
            msg = "training data must be non-empty"  # pragma: no cover
            raise ValueError(msg)  # pragma: no cover
        n = x_arr.shape[0]
        k_mat = _rbf_kernel(x_arr, x_arr, self.length_scale, self.sigma_f)
        # leverage cds2.linalg for Cholesky; add jitter for stability
        jitter = self.noise**2
        k_mat = k_mat + np.eye(n, dtype=float) * (jitter + 1e-10)
        try:
            lower = cds2_linalg.cholesky(k_mat)
        except np.linalg.LinAlgError:  # pragma: no cover
            # add larger jitter and retry  # pragma: no cover
            k_mat = k_mat + np.eye(n, dtype=float) * 1e-6  # pragma: no cover
            lower = cds2_linalg.cholesky(k_mat)  # pragma: no cover
        # solve for alpha = K^{-1} y via Cholesky
        # Use triangular solves via scipy for stability
        # Solve L * v = y  -> v
        v_vec = sla.solve_triangular(lower, y_arr, lower=True, check_finite=False)
        # Solve L.T * alpha = v
        alpha = sla.solve_triangular(lower.T, v_vec, lower=False, check_finite=False)
        self.X_train_ = np.asarray(x_arr, dtype=float)
        self.y_train_ = np.asarray(y_arr, dtype=float)
        self.L_ = np.asarray(lower, dtype=float)
        self.alpha_ = np.asarray(alpha, dtype=float)
        return self

    def predict(self, x_star: object) -> tuple[FloatArray, FloatArray]:
        """Predict mean and standard deviation at new points."""
        if self.X_train_ is None or self.y_train_ is None or self.L_ is None or self.alpha_ is None:
            msg = "model is not fitted"
            raise RuntimeError(msg)
        x_arr = _as_2d(x_star, dim=self.X_train_.shape[1])
        if x_arr.shape[1] != self.X_train_.shape[1]:
            msg = "x_star dimension must match training data"  # pragma: no cover
            raise ValueError(msg)  # pragma: no cover
        k_star = _rbf_kernel(x_arr, self.X_train_, self.length_scale, self.sigma_f)
        # mean
        mu = k_star @ self.alpha_
        # variance via triangular solve: v = L^{-1} K_star^T
        v_mat = sla.solve_triangular(self.L_, k_star.T, lower=True, check_finite=False)
        k_star_star = float(self.sigma_f**2)
        # var = k(x*, x*) - v^T v for each x*
        var = k_star_star - np.sum(v_mat * v_mat, axis=0)
        var = np.maximum(var, 0.0)
        std = np.sqrt(var)
        return np.asarray(mu, dtype=float), np.asarray(std, dtype=float)

fit(x_train, y_train)

Fit the GP to training data.

Source code in src/cds2/bayesopt.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def fit(self, x_train: object, y_train: object) -> GaussianProcess:
    """Fit the GP to training data."""
    x_arr = _as_2d(x_train)
    if x_arr.ndim != 2:
        msg = "x_train must be 2-D"  # pragma: no cover
        raise ValueError(msg)  # pragma: no cover
    y_arr = np.asarray(y_train, dtype=float).ravel()
    if y_arr.size != x_arr.shape[0]:
        msg = "x_train and y_train must have matching first dimension"  # pragma: no cover
        raise ValueError(msg)  # pragma: no cover
    if x_arr.shape[0] == 0:
        msg = "training data must be non-empty"  # pragma: no cover
        raise ValueError(msg)  # pragma: no cover
    n = x_arr.shape[0]
    k_mat = _rbf_kernel(x_arr, x_arr, self.length_scale, self.sigma_f)
    # leverage cds2.linalg for Cholesky; add jitter for stability
    jitter = self.noise**2
    k_mat = k_mat + np.eye(n, dtype=float) * (jitter + 1e-10)
    try:
        lower = cds2_linalg.cholesky(k_mat)
    except np.linalg.LinAlgError:  # pragma: no cover
        # add larger jitter and retry  # pragma: no cover
        k_mat = k_mat + np.eye(n, dtype=float) * 1e-6  # pragma: no cover
        lower = cds2_linalg.cholesky(k_mat)  # pragma: no cover
    # solve for alpha = K^{-1} y via Cholesky
    # Use triangular solves via scipy for stability
    # Solve L * v = y  -> v
    v_vec = sla.solve_triangular(lower, y_arr, lower=True, check_finite=False)
    # Solve L.T * alpha = v
    alpha = sla.solve_triangular(lower.T, v_vec, lower=False, check_finite=False)
    self.X_train_ = np.asarray(x_arr, dtype=float)
    self.y_train_ = np.asarray(y_arr, dtype=float)
    self.L_ = np.asarray(lower, dtype=float)
    self.alpha_ = np.asarray(alpha, dtype=float)
    return self

predict(x_star)

Predict mean and standard deviation at new points.

Source code in src/cds2/bayesopt.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def predict(self, x_star: object) -> tuple[FloatArray, FloatArray]:
    """Predict mean and standard deviation at new points."""
    if self.X_train_ is None or self.y_train_ is None or self.L_ is None or self.alpha_ is None:
        msg = "model is not fitted"
        raise RuntimeError(msg)
    x_arr = _as_2d(x_star, dim=self.X_train_.shape[1])
    if x_arr.shape[1] != self.X_train_.shape[1]:
        msg = "x_star dimension must match training data"  # pragma: no cover
        raise ValueError(msg)  # pragma: no cover
    k_star = _rbf_kernel(x_arr, self.X_train_, self.length_scale, self.sigma_f)
    # mean
    mu = k_star @ self.alpha_
    # variance via triangular solve: v = L^{-1} K_star^T
    v_mat = sla.solve_triangular(self.L_, k_star.T, lower=True, check_finite=False)
    k_star_star = float(self.sigma_f**2)
    # var = k(x*, x*) - v^T v for each x*
    var = k_star_star - np.sum(v_mat * v_mat, axis=0)
    var = np.maximum(var, 0.0)
    std = np.sqrt(var)
    return np.asarray(mu, dtype=float), np.asarray(std, dtype=float)

OptimizeResult dataclass

Outcome of Bayesian optimization.

Source code in src/cds2/bayesopt.py 
209
210
211
212
213
214
215
216
217
218
219
@dataclass(frozen=True)
class OptimizeResult:
    """Outcome of Bayesian optimization."""

    x: FloatArray
    fun: float
    success: bool
    message: str
    n_iterations: int
    xs: FloatArray
    ys: FloatArray

expected_improvement(mu, sigma, best_f, xi=0.01)

Expected Improvement for minimization.

Parameters:

Name Type Description Default
mu object

Predicted mean(s) at candidate points.

 required
sigma object

Predicted standard deviation(s).

 required
best_f float

Best observed objective value so far (minimum).

 required
xi float

Exploration-exploitation trade-off.

 0.01
Source code in src/cds2/bayesopt.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def expected_improvement(
    mu: object,
    sigma: object,
    best_f: float,
    xi: float = 0.01,
) -> FloatArray:
    """Expected Improvement for minimization.

    Parameters
    ----------
    mu:
        Predicted mean(s) at candidate points.
    sigma:
        Predicted standard deviation(s).
    best_f:
        Best observed objective value so far (minimum).
    xi:
        Exploration-exploitation trade-off.
    """
    mu_arr = np.asarray(mu, dtype=float)
    sigma_arr = np.asarray(sigma, dtype=float)
    best = float(best_f)
    xi_val = float(xi)
    # broadcast shapes: mu_arr and sigma_arr should be compatible
    # handle scalar vs array
    imp = best - mu_arr - xi_val
    with np.errstate(divide="ignore", invalid="ignore"):
        z = np.divide(imp, sigma_arr, out=np.zeros_like(imp, dtype=float), where=sigma_arr > 0)
        cdf = norm.cdf(z)
        pdf = norm.pdf(z)
        ei = imp * cdf + sigma_arr * pdf
        ei = np.where(sigma_arr > 0, ei, 0.0)
        ei = np.maximum(ei, 0.0)
    return np.asarray(ei, dtype=float)

upper_confidence_bound(mu, sigma, kappa=2.0)

Upper confidence bound acquisition.

ucb = mu + kappa * sigma. For maximization, maximise this quantity; for minimization the optimizer uses -mu + kappa*sigma (i.e. upper_confidence_bound(-mu, sigma)).

Source code in src/cds2/bayesopt.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
def upper_confidence_bound(
    mu: object,
    sigma: object,
    kappa: float = 2.0,
) -> FloatArray:
    """Upper confidence bound acquisition.

    ``ucb = mu + kappa * sigma``. For maximization, maximise this
    quantity; for minimization the optimizer uses ``-mu + kappa*sigma``
    (i.e. ``upper_confidence_bound(-mu, sigma)``).
    """
    mu_arr = np.asarray(mu, dtype=float)
    sigma_arr = np.asarray(sigma, dtype=float)
    k = float(kappa)
    return np.asarray(mu_arr + k * sigma_arr, dtype=float)

bayes_opt(objective, bounds, n_init=5, n_iter=20, acquisition='ei', seed=None)

Bayesian optimization over a box-constrained domain.

Parameters:

Name Type Description Default
objective Callable[[FloatArray], float]

Function to minimize; takes a 1-D array of shape (dim,).

 required
bounds Sequence[tuple[float, float]] | FloatArray

Sequence of (low, high) pairs per dimension.

 required
n_init int

Number of random initial evaluations.

 5
n_iter int

Number of Bayesian optimization steps.

 20
acquisition str

Acquisition function: "ei" (expected improvement) or "ucb" (upper confidence bound / lower confidence bound for minimization).

 'ei'
seed int | None

Random seed for reproducibility.

 None
Source code in src/cds2/bayesopt.py 
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
def bayes_opt(
    objective: Callable[[FloatArray], float],
    bounds: Sequence[tuple[float, float]] | FloatArray,
    n_init: int = 5,
    n_iter: int = 20,
    acquisition: str = "ei",
    seed: int | None = None,
) -> OptimizeResult:
    """Bayesian optimization over a box-constrained domain.

    Parameters
    ----------
    objective:
        Function to minimize; takes a 1-D array of shape ``(dim,)``.
    bounds:
        Sequence of ``(low, high)`` pairs per dimension.
    n_init:
        Number of random initial evaluations.
    n_iter:
        Number of Bayesian optimization steps.
    acquisition:
        Acquisition function: ``"ei"`` (expected improvement) or ``"ucb"``
        (upper confidence bound / lower confidence bound for minimization).
    seed:
        Random seed for reproducibility.
    """
    if n_init < 1:
        msg = "n_init must be >= 1"  # pragma: no cover
        raise ValueError(msg)  # pragma: no cover
    if n_iter < 0:
        msg = "n_iter must be >= 0"  # pragma: no cover
        raise ValueError(msg)  # pragma: no cover
    bounds_arr = _prepare_bounds(bounds)
    dim = int(bounds_arr.shape[0])
    low = np.asarray(bounds_arr[:, 0], dtype=float)
    high = np.asarray(bounds_arr[:, 1], dtype=float)

    acq = acquisition.lower().strip()
    if acq in {"expected_improvement", "expected-improvement"}:
        acq = "ei"  # pragma: no cover
    elif acq in {"upper_confidence_bound", "upper-confidence-bound"}:
        acq = "ucb"  # pragma: no cover
    if acq not in {"ei", "ucb"}:
        msg = "acquisition must be 'ei' or 'ucb'"  # pragma: no cover
        raise ValueError(msg)  # pragma: no cover

    rng = np.random.default_rng(seed)

    # initial design uniformly at random
    x_init = rng.uniform(low, high, size=(n_init, dim))
    y_init = np.empty(n_init, dtype=float)
    for i in range(n_init):
        y_init[i] = _eval_objective(objective, x_init[i])

    x_all = np.asarray(x_init, dtype=float)
    y_all = np.asarray(y_init, dtype=float)

    # adaptive length scale scaled to bounds
    mean_range = float(np.mean(high - low))
    base_length_scale = max(mean_range * 0.25, 1e-2)

    for _ in range(n_iter):
        y_mean = float(np.mean(y_all))
        y_std = float(np.std(y_all))
        if y_std < 1e-9:
            y_std = 1.0  # pragma: no cover
        y_norm = (y_all - y_mean) / y_std

        gp = GaussianProcess(length_scale=base_length_scale, sigma_f=1.0, noise=1e-6)
        gp.fit(x_all, y_norm)
        best_f_norm = float(np.min(y_norm))

        # acquisition optimization via random sampling + local refinement
        n_candidates = 5000 if dim <= 3 else 2000 * dim
        if n_candidates > 20000:
            n_candidates = 20000  # pragma: no cover
        candidates = rng.uniform(low, high, size=(n_candidates, dim))
        mu_cand, sigma_cand = gp.predict(candidates)

        if acq == "ei":
            acq_vals = expected_improvement(mu_cand, sigma_cand, best_f_norm, xi=0.01)
        else:
            # for minimization, maximize -mu + kappa*sigma
            # upper_confidence_bound(-mu, sigma) == -mu + kappa*sigma
            acq_vals = upper_confidence_bound(-mu_cand, sigma_cand, kappa=2.0)

        best_idx = int(np.argmax(acq_vals))
        best_cand = np.asarray(candidates[best_idx], dtype=float)
        best_acq_val = float(acq_vals[best_idx])

        # local refinement from top candidates
        bounds_list: list[tuple[float, float]] = [
            (float(low[d]), float(high[d])) for d in range(dim)
        ]

        def neg_acq(x_flat: FloatArray) -> float:
            x_2d = np.asarray(x_flat, dtype=float).reshape(1, -1)
            mu_arr, sigma_arr = gp.predict(x_2d)
            mu0 = float(mu_arr[0])
            sigma0 = float(sigma_arr[0])
            if acq == "ei":
                ei_val = expected_improvement(
                    np.array([mu0]), np.array([sigma0]), best_f_norm, xi=0.01
                )[0]
                return -float(ei_val)
            # ucb for minimization
            kappa = 2.0
            # neg_acq = mu - kappa*sigma  (minimize -> maximize -mu+kappa*sigma)
            return float(mu0 - kappa * sigma0)

        # try refining the top few candidates
        top_n = min(5, n_candidates)
        top_indices = np.argsort(acq_vals)[-top_n:]
        best_x_next = best_cand
        for idx in top_indices:
            x0 = np.asarray(candidates[int(idx)], dtype=float)
            try:
                res = spo.minimize(neg_acq, x0, method="L-BFGS-B", bounds=bounds_list)
            except Exception:  # pragma: no cover
                continue  # pragma: no cover
            if not res.success:
                continue
            # acquisition at refined point is -res.fun
            try:
                acq_at_res = -float(res.fun)
            except Exception:  # pragma: no cover
                continue  # pragma: no cover
            if acq_at_res > best_acq_val:
                best_acq_val = acq_at_res
                best_x_next = np.asarray(res.x, dtype=float)

        # ensure within bounds (clip due to numerical issues)
        best_x_next = np.clip(best_x_next, low, high)

        y_next = _eval_objective(objective, best_x_next)
        x_all = np.vstack([x_all, best_x_next.reshape(1, -1)])
        y_all = np.append(y_all, y_next)

    best_index = int(np.argmin(y_all))
    best_x = np.asarray(x_all[best_index], dtype=float)
    best_fun = float(y_all[best_index])
    return OptimizeResult(
        x=best_x,
        fun=best_fun,
        success=True,
        message="Optimization terminated successfully.",
        n_iterations=int(n_iter),
        xs=np.asarray(x_all, dtype=float),
        ys=np.asarray(y_all, dtype=float),
    )