Coverage for src/chebpy/gpr.py: 100%
179 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-22 21:50 +0000
1"""Gaussian process regression with Chebfun representations.
3Implements Gaussian process regression (GPR) following the algorithm described
4in Rasmussen & Williams, *Gaussian Processes for Machine Learning*, MIT Press,
52006, and the MATLAB Chebfun ``gpr.m`` by The University of Oxford and The
6Chebfun Developers.
8The posterior mean, variance, and (optionally) random samples from the
9posterior are all returned as Chebfun / Quasimatrix objects so that they can
10be manipulated with the full ChebPy toolkit (differentiation, integration,
11rootfinding, etc.).
13Reference:
14 C. E. Rasmussen & C. K. I. Williams, "Gaussian Processes for Machine
15 Learning", MIT Press, 2006.
16"""
18from __future__ import annotations
20from collections.abc import Callable
21from dataclasses import dataclass, field
22from typing import cast
24import numpy as np
25from numpy.typing import ArrayLike
27from .algorithms import chebpts2
28from .chebfun import Chebfun
29from .quasimatrix import Quasimatrix
30from .settings import _preferences as prefs
33# ---------------------------------------------------------------------------
34# Options container
35# ---------------------------------------------------------------------------
36@dataclass
37class _GPROptions:
38 """Parsed options for a GPR call."""
40 sigma: float = 1.0
41 sigma_given: bool = False
42 length_scale: float = 0.0
43 noise: float = 0.0
44 domain: np.ndarray = field(default_factory=lambda: np.array([-1.0, 1.0]))
45 trig: bool = False
46 n_samples: int = 0
49# ---------------------------------------------------------------------------
50# Kernel helpers
51# ---------------------------------------------------------------------------
54def _kernel_matrix(
55 x1: np.ndarray,
56 x2: np.ndarray,
57 opts: _GPROptions,
58) -> np.ndarray:
59 """Evaluate the covariance kernel k(x1_i, x2_j) for all pairs."""
60 r = x1[:, None] - x2[None, :]
61 if opts.trig:
62 period = opts.domain[1] - opts.domain[0]
63 return cast(np.ndarray, opts.sigma**2 * np.exp(-2.0 / opts.length_scale**2 * np.sin(np.pi / period * r) ** 2))
64 return cast(np.ndarray, opts.sigma**2 * np.exp(-0.5 / opts.length_scale**2 * r**2))
67def _log_marginal_likelihood(
68 length_scale: float | np.ndarray,
69 x: np.ndarray,
70 y: np.ndarray,
71 opts: _GPROptions,
72) -> float | np.ndarray:
73 """Negative log marginal likelihood (eq. 2.30 in Rasmussen & Williams).
75 Accepts scalar or array *length_scale* so that it can be wrapped as a
76 Chebfun for optimisation.
77 """
78 scalar_input = np.ndim(length_scale) == 0
79 ls = np.atleast_1d(np.asarray(length_scale, dtype=float))
80 n = len(x)
81 rx = x[:, None] - x[None, :]
82 result = np.empty_like(ls)
84 for idx in np.ndindex(ls.shape):
85 l_val = ls[idx]
86 if opts.trig:
87 period = opts.domain[1] - opts.domain[0]
88 cov_mat = opts.sigma**2 * np.exp(-2.0 / l_val**2 * np.sin(np.pi / period * rx) ** 2)
89 else:
90 cov_mat = opts.sigma**2 * np.exp(-0.5 / l_val**2 * rx**2)
92 if opts.noise != 0:
93 cov_mat += opts.noise**2 * np.eye(n)
94 else:
95 cov_mat += 1e-15 * n * opts.sigma**2 * np.eye(n)
97 chol_l = np.linalg.cholesky(cov_mat)
98 alpha = np.linalg.solve(chol_l.T, np.linalg.solve(chol_l, y))
99 lml = -0.5 * y @ alpha - np.sum(np.log(np.diag(chol_l))) - 0.5 * n * np.log(2 * np.pi)
100 result[idx] = lml
102 return float(result.item()) if scalar_input else result.ravel()
105# ---------------------------------------------------------------------------
106# Length-scale selection via max log marginal likelihood
107# ---------------------------------------------------------------------------
110def _select_length_scale(x: np.ndarray, y: np.ndarray, opts: _GPROptions) -> float:
111 """Choose the length-scale that maximises the log marginal likelihood."""
112 n = len(x)
113 dom_size = opts.domain[1] - opts.domain[0]
115 if opts.trig:
116 lo, hi = 1.0 / (2 * n), 10.0
117 else:
118 lo, hi = dom_size / (2 * np.pi * n), 10.0 / np.pi * dom_size
120 # Heuristic: shrink the right end of the search domain if the lml is
121 # monotonically decreasing (mirrors the MATLAB implementation).
122 f1 = float(_log_marginal_likelihood(lo, x, y, opts))
123 f2 = float(_log_marginal_likelihood(hi, x, y, opts))
124 while f1 > f2 and hi / lo > 1 + 1e-4:
125 new_bound = lo + (hi - lo) / 10.0
126 f_new = float(_log_marginal_likelihood(new_bound, x, y, opts))
127 if f_new > f1:
128 break
129 hi = new_bound
130 f2 = f_new
132 # Maximise using golden-section search (negated to find the max).
133 return _golden_section_max(lambda ls: float(_log_marginal_likelihood(ls, x, y, opts)), lo, hi)
136def _golden_section_max(f: Callable[[float], float], a: float, b: float, tol: float = 1e-6) -> float:
137 """Golden-section search for the scalar argmax of *f* on [a, b]."""
138 gr = (np.sqrt(5.0) + 1.0) / 2.0
139 c = b - (b - a) / gr
140 d = a + (b - a) / gr
141 while abs(b - a) > tol * (abs(a) + abs(b)):
142 if f(c) > f(d):
143 b = d
144 else:
145 a = c
146 c = b - (b - a) / gr
147 d = a + (b - a) / gr
148 return 0.5 * (a + b)
151# ---------------------------------------------------------------------------
152# Public API
153# ---------------------------------------------------------------------------
156def _parse_inputs(
157 x: ArrayLike,
158 y: ArrayLike,
159 *,
160 sigma: float | None,
161 noise: float,
162 trig: bool,
163 n_samples: int,
164) -> tuple[np.ndarray, np.ndarray, _GPROptions, float]:
165 """Validate inputs and build the initial options container.
167 Returns ``(x_arr, y_arr, opts, scaling_factor)``.
168 """
169 x_arr = np.asarray(x, dtype=float).ravel()
170 y_arr = np.asarray(y, dtype=float).ravel()
171 if x_arr.shape != y_arr.shape:
172 msg = "x and y must have the same length."
173 raise ValueError(msg)
175 opts = _GPROptions(trig=trig, noise=noise, n_samples=n_samples)
177 scaling_factor = 1.0
178 if sigma is not None:
179 opts.sigma = sigma
180 opts.sigma_given = True
181 else:
182 if len(y_arr) > 0:
183 scaling_factor = float(np.max(np.abs(y_arr)))
184 opts.sigma_given = False
185 opts.sigma = scaling_factor
187 return x_arr, y_arr, opts, scaling_factor
190def _infer_domain(
191 x_arr: np.ndarray,
192 opts: _GPROptions,
193 domain: tuple[float, float] | list[float] | np.ndarray | None,
194) -> None:
195 """Set ``opts.domain`` from *domain* or from the observation locations."""
196 if domain is not None:
197 opts.domain = np.asarray(domain, dtype=float)
198 elif len(x_arr) == 0:
199 opts.domain = np.array([-1.0, 1.0])
200 elif len(x_arr) == 1:
201 opts.domain = np.array([x_arr[0] - 1, x_arr[0] + 1])
202 elif opts.trig:
203 span = float(np.max(x_arr) - np.min(x_arr))
204 opts.domain = np.array([float(np.min(x_arr)), float(np.max(x_arr)) + 0.1 * span])
205 else:
206 opts.domain = np.array([float(np.min(x_arr)), float(np.max(x_arr))])
209def _infer_length_scale(
210 x_arr: np.ndarray,
211 y_arr: np.ndarray,
212 opts: _GPROptions,
213 scaling_factor: float,
214 length_scale: float | None,
215) -> None:
216 """Set ``opts.length_scale`` — user-supplied or auto-selected."""
217 if length_scale is not None:
218 opts.length_scale = length_scale
219 return
221 if len(x_arr) == 0:
222 opts.length_scale = 1.0
223 return
225 y_n = y_arr / scaling_factor if scaling_factor != 0 else y_arr
227 if not opts.sigma_given:
228 tmp = _GPROptions(
229 sigma=1.0,
230 sigma_given=True,
231 noise=opts.noise / scaling_factor if scaling_factor != 0 else opts.noise,
232 domain=opts.domain,
233 trig=opts.trig,
234 )
235 y_opt = y_n
236 else:
237 tmp = _GPROptions(
238 sigma=opts.sigma,
239 sigma_given=True,
240 noise=opts.noise,
241 domain=opts.domain,
242 trig=opts.trig,
243 )
244 y_opt = y_arr
246 opts.length_scale = _select_length_scale(x_arr, y_opt, tmp)
249def _posterior_chebfuns(
250 x_arr: np.ndarray,
251 y_arr: np.ndarray,
252 opts: _GPROptions,
253 scaling_factor: float,
254 n_samples: int,
255) -> tuple[Chebfun, Chebfun] | tuple[Chebfun, Chebfun, Quasimatrix]:
256 """Compute posterior mean, variance, and optional samples as Chebfuns."""
257 n = len(x_arr)
258 cov_mat = _kernel_matrix(x_arr, x_arr, opts)
259 if opts.noise == 0:
260 cov_mat += 1e-15 * scaling_factor**2 * n * np.eye(n)
261 else:
262 cov_mat += opts.noise**2 * np.eye(n)
264 chol_l = np.linalg.cholesky(cov_mat)
265 alpha = np.linalg.solve(chol_l.T, np.linalg.solve(chol_l, y_arr))
267 # Sample grid: Chebyshev points for the default tech, equispaced points
268 # for the periodic (Trigtech) case. Using the right grid here is critical
269 # because the constructed Chebfun pieces are then built via
270 # ``Chebfun.initfun_fixedlen`` whose underlying tech evaluates at exactly
271 # this grid.
272 sample_size = min(20 * n, 2000)
273 if opts.trig:
274 # n equispaced points on [a, b) — matches Trigtech._trigpts mapped to domain
275 x_sample = opts.domain[0] + (opts.domain[1] - opts.domain[0]) * np.arange(sample_size) / sample_size
276 else:
277 t = chebpts2(sample_size)
278 x_sample = 0.5 * (opts.domain[1] - opts.domain[0]) * t + 0.5 * (opts.domain[0] + opts.domain[1])
280 in_x = np.isin(x_sample, x_arr)
282 k_star = _kernel_matrix(x_sample, x_arr, opts)
283 if opts.noise:
284 k_star += opts.noise**2 * (np.abs(x_sample[:, None] - x_arr[None, :]) == 0)
286 # Posterior mean
287 mean_vals = k_star @ alpha
289 # Posterior variance
290 k_ss = _kernel_matrix(x_sample, x_sample, opts)
291 if opts.noise:
292 k_ss += opts.noise**2 * np.diag(in_x.astype(float))
294 v = np.linalg.solve(chol_l, k_star.T)
295 var_diag = np.diag(k_ss) - np.sum(v**2, axis=0)
296 var_diag = np.maximum(var_diag, 0.0)
298 # Build Chebfuns under the appropriate tech. For ``trig=True`` this
299 # produces Trigtech-backed pieces so that downstream calculus is performed
300 # in Fourier space.
301 tech_name = "Trigtech" if opts.trig else prefs.tech
302 with prefs:
303 prefs.tech = tech_name
304 f_mean = Chebfun.initfun_fixedlen(lambda _z: mean_vals, sample_size, opts.domain)
305 f_var = Chebfun.initfun_fixedlen(lambda _z: var_diag, sample_size, opts.domain)
307 if n_samples <= 0:
308 return f_mean, f_var
310 # Posterior samples
311 cov_post = k_ss - v.T @ v
312 cov_post = 0.5 * (cov_post + cov_post.T)
313 cov_post += 1e-12 * scaling_factor**2 * n * np.eye(sample_size)
314 chol_s = np.linalg.cholesky(cov_post)
316 draws = mean_vals[:, None] + chol_s @ np.random.randn(sample_size, n_samples)
317 cols: list[Chebfun] = []
318 for j in range(n_samples):
319 cols.append(
320 Chebfun.initfun_fixedlen(
321 lambda _z, _j=j: draws[:, _j],
322 sample_size,
323 opts.domain,
324 )
325 )
326 return f_mean, f_var, Quasimatrix(cols)
329def gpr(
330 x: ArrayLike,
331 y: ArrayLike,
332 *,
333 domain: tuple[float, float] | list[float] | np.ndarray | None = None,
334 sigma: float | None = None,
335 length_scale: float | None = None,
336 noise: float = 0.0,
337 trig: bool = False,
338 n_samples: int = 0,
339) -> tuple[Chebfun, Chebfun] | tuple[Chebfun, Chebfun, Quasimatrix]:
340 """Gaussian process regression returning Chebfun objects.
342 Given observations ``(x, y)`` of a latent function, compute the posterior
343 mean and variance of a Gaussian process with zero prior mean and a squared
344 exponential kernel::
346 k(x, x') = sigma**2 * exp(-0.5 / L**2 * (x - x')**2)
348 When ``trig=True`` a periodic variant is used instead::
350 k(x, x') = sigma**2 * exp(-2 / L**2 * sin(pi * (x - x') / P)**2)
352 where *P* is the period (length of the domain).
354 Args:
355 x: Observation locations (1-D array-like).
356 y: Observation values (same length as *x*).
357 domain: Domain ``[a, b]`` for the output Chebfuns. Defaults to
358 ``[min(x), max(x)]`` (or slightly extended for ``trig``).
359 sigma: Signal variance of the kernel. Defaults to ``max(|y|)``.
360 length_scale: Length-scale *L* of the kernel. If ``None``, it is
361 chosen to maximise the log marginal likelihood.
362 noise: Standard deviation of i.i.d. Gaussian observation noise.
363 The kernel diagonal is augmented by ``noise**2``.
364 trig: If ``True``, use a periodic squared-exponential kernel.
365 n_samples: Number of independent posterior samples to draw. When
366 positive, a :class:`Quasimatrix` with *n_samples* columns is
367 returned as the third element of the output tuple.
369 Returns:
370 ``(f_mean, f_var)`` — posterior mean and variance as Chebfun objects.
371 If ``n_samples > 0``, returns ``(f_mean, f_var, samples)`` where
372 *samples* is a Quasimatrix whose columns are independent draws from
373 the posterior.
375 Raises:
376 ValueError: If *x* and *y* have different lengths or are empty.
378 Examples:
379 >>> import numpy as np
380 >>> from chebpy.gpr import gpr
381 >>> rng = np.random.default_rng(1)
382 >>> x = -2 + 4 * rng.random(10)
383 >>> y = np.sin(np.exp(x))
384 >>> f_mean, f_var = gpr(x, y, domain=[-2, 2])
386 Reference:
387 C. E. Rasmussen & C. K. I. Williams, "Gaussian Processes for Machine
388 Learning", MIT Press, 2006.
389 """
390 x_arr, y_arr, opts, scaling_factor = _parse_inputs(
391 x,
392 y,
393 sigma=sigma,
394 noise=noise,
395 trig=trig,
396 n_samples=n_samples,
397 )
398 _infer_domain(x_arr, opts, domain)
399 _infer_length_scale(x_arr, y_arr, opts, scaling_factor, length_scale)
401 # No data → return prior
402 if len(x_arr) == 0:
403 f_mean = Chebfun.initconst(0.0, opts.domain)
404 f_var = Chebfun.initconst(opts.sigma**2, opts.domain)
405 if n_samples > 0:
406 return f_mean, f_var, _prior_samples(opts, scaling_factor, n_samples)
407 return f_mean, f_var
409 return _posterior_chebfuns(x_arr, y_arr, opts, scaling_factor, n_samples)
412def _prior_samples(
413 opts: _GPROptions,
414 scaling_factor: float,
415 n_samples: int,
416) -> Quasimatrix:
417 """Draw samples from the GP prior (no observations)."""
418 sample_size = 1000
419 if opts.trig:
420 x_sample = np.linspace(opts.domain[0], opts.domain[1], sample_size)
421 else:
422 t = chebpts2(sample_size)
423 x_sample = 0.5 * (opts.domain[1] - opts.domain[0]) * t + 0.5 * (opts.domain[0] + opts.domain[1])
425 k_ss = _kernel_matrix(x_sample, x_sample, opts)
426 k_ss += 1e-12 * scaling_factor**2 * np.eye(sample_size)
427 chol_s = np.linalg.cholesky(k_ss)
429 f_mean_vals = np.zeros(sample_size)
430 draws = f_mean_vals[:, None] + chol_s @ np.random.randn(sample_size, n_samples)
432 cols: list[Chebfun] = []
433 for j in range(n_samples):
434 cols.append(
435 Chebfun.initfun_fixedlen(
436 lambda _z, _j=j: draws[:, _j],
437 sample_size,
438 opts.domain,
439 )
440 )
441 return Quasimatrix(cols)