From fe795f900e09aad3b7f34b1abc3439ac1230707c Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 24 Aug 2023 14:20:56 +0100 Subject: [PATCH 01/25] implement periodic psd and pass test --- pymc/gp/cov.py | 22 ++++++++++++++++++++++ tests/gp/test_cov.py | 17 +++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 8a6ec534ca..add2156661 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -791,6 +791,28 @@ def full_from_distance(self, dist: TensorLike, squared: bool = False) -> TensorV r2 = dist if squared else dist**2 return pt.exp(-0.5 * r2) + def power_spectral_density(self, m: TensorLike) -> TensorVariable: + """ + Technically, this is not a spectral density but these are the first `m` coefficients of + the low rank approximation for the periodic kernel, which are used in the same way. + + The coefficients of the HSGP approximation for the Periodic kernel are given by: + + .. math:: + + \tilde{q}_j^2 = \frac{2 \text{I}_j (\\ell^{-2})}{\\exp(\\ell^{-2})} \\ + \tilde{q}_0^2 = \frac{\text{I}_0 (\\ell^{-2})}{\\exp(\\ell^{-2})} + + where $\text{I}_j$ is the modified Bessel function of the first kind. + """ + a = 1 / pt.square(self.ls) + J = pt.arange(0, m) + c = pt.where(J > 0, 2, 1) + + q2 = c * pt.iv(J, a) / pt.exp(a) + + return q2 + class Linear(Covariance): r""" diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py index 750d5fe6a1..20f480d87c 100644 --- a/tests/gp/test_cov.py +++ b/tests/gp/test_cov.py @@ -18,6 +18,8 @@ import pytensor.tensor as pt import pytest +from scipy.special import iv + import pymc as pm from pymc.math import cartesian, kronecker @@ -639,6 +641,21 @@ def test_1d(self): Kd = cov(X, diag=True).eval() npt.assert_allclose(np.diag(K), Kd, atol=1e-5) + def test_psd(self): + # compare to simple 1d formula + ell = 2.0 + P = 0.1 + m = 5 + + a = 1 / np.square(ell) + J = np.arange(0, m) + true_1d_psd = np.where(J > 0, 2, 1) * iv(J, a) / np.exp(a) + + test_1d_psd = ( + pm.gp.cov.Periodic(1, period=P, ls=ell).power_spectral_density(m).flatten().eval() + ) + npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) + class TestLinear: def test_1d(self): From 788faeebe066cf2da74202a2892b9f975ec07f92 Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 24 Aug 2023 14:25:19 +0100 Subject: [PATCH 02/25] simplify hsgp test, set up for mark other cov_func --- tests/gp/test_hsgp_approx.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index 96465e9437..2521912437 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -105,15 +105,6 @@ def X2(self, data): def model(self): return pm.Model() - @pytest.fixture - def cov_func(self): - return pm.gp.cov.ExpQuad(1, ls=1) - - @pytest.fixture - def gp(self, cov_func): - gp = pm.gp.Latent(cov_func=cov_func) - return gp - def test_set_boundaries_1d(self, X1): X1s = X1 - np.mean(X1, axis=0) L = pm.gp.hsgp_approx.set_boundary(X1s, c=2).eval() @@ -150,6 +141,7 @@ def test_parametrization(self): cov_func = pm.gp.cov.ExpQuad(2, ls=[1, 2]) pm.gp.HSGP(m=[50, 50], L=[12, 12], cov_func=cov_func) + @pytest.mark.parametrize("cov_func", [pm.gp.cov.ExpQuad(1, ls=1)]) @pytest.mark.parametrize("drop_first", [True, False]) def test_parametrization_drop_first(self, model, cov_func, X1, drop_first): n_basis = 100 @@ -165,6 +157,7 @@ def test_parametrization_drop_first(self, model, cov_func, X1, drop_first): else: assert n_coeffs == n_basis, "one was dropped when it shouldn't have been" + @pytest.mark.parametrize("cov_func", [pm.gp.cov.ExpQuad(1, ls=1)]) @pytest.mark.parametrize("parameterization", ["centered", "noncentered"]) def test_prior(self, model, cov_func, X1, parameterization, rng): """Compare HSGP prior to unapproximated GP prior, pm.gp.Latent. Draw samples from the @@ -188,6 +181,7 @@ def test_prior(self, model, cov_func, X1, parameterization, rng): ) assert not reject, "H0 was rejected, even though HSGP and GP priors should match." + @pytest.mark.parametrize("cov_func", [pm.gp.cov.ExpQuad(1, ls=1)]) @pytest.mark.parametrize("parameterization", ["centered", "noncentered"]) def test_conditional(self, model, cov_func, X1, parameterization): """Compare HSGP conditional to unapproximated GP prior, pm.gp.Latent. Draw samples from the From ee76703f677b7c7673391b5a7b418ea1d9c2459d Mon Sep 17 00:00:00 2001 From: theorashid Date: Fri, 25 Aug 2023 15:56:12 +0100 Subject: [PATCH 03/25] attempt at if-else logic in HSGP class API --- pymc/gp/hsgp_approx.py | 110 ++++++++++++++++++++++++++--------- tests/gp/test_hsgp_approx.py | 11 +++- 2 files changed, 91 insertions(+), 30 deletions(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index adca16f600..39b571060a 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -23,7 +23,7 @@ import pymc as pm -from pymc.gp.cov import Covariance +from pymc.gp.cov import Covariance, Periodic from pymc.gp.gp import Base from pymc.gp.mean import Mean, Zero @@ -53,7 +53,7 @@ def calc_eigenvectors( m: Sequence[int], tl: ModuleType = np, ): - """Calculate eigenvectors of the Laplacian. These are used as basis vectors in the HSGP + """Calculate eigenvectors of the Laplacian. These are used as basis vectors in the HSGP approximation. """ m_star = int(np.prod(m)) @@ -66,6 +66,25 @@ def calc_eigenvectors( return phi +def calc_basis_periodic( + Xs: TensorLike, + period: TensorLike, + m: Sequence[int], + tl: ModuleType = np, +): + """ + Calculate basis vectors for the cosine series expansion of the periodic covariance function. + These are derived from the Taylor series representation of the covariance. + """ + w0 = (2 * tl.pi) / period # angular frequency defining the periodicity + m1 = tl.tile(w0 * Xs[:, None], m) + m2 = tl.diag(tl.arange(m)) + mw0x = m1 @ m2 + phi_cos = tl.cos(mw0x) + phi_sin = tl.sin(mw0x) + return phi_cos, phi_sin + + class HSGP(Base): R""" Hilbert Space Gaussian process approximation. @@ -160,7 +179,7 @@ def __init__( cov_func: Covariance, ): arg_err_msg = ( - "`m` and L, if provided, must be sequences with one element per active " + "`m` and `L`, if provided, must be sequences with one element per active " "dimension of the kernel or covariance function." ) @@ -186,6 +205,11 @@ def __init__( else: self._parameterization = parameterization + if isinstance(cov_func, Periodic) and cov_func.n_dims > 1: + raise ValueError( + "HSGP approximation for periodic kernel only implemented for 1-dimensional case." + ) + self._drop_first = drop_first self._m = m self._m_star = int(np.prod(self._m)) @@ -287,20 +311,26 @@ def prior_linearized(self, Xs: TensorLike): # Index Xs using input_dim and active_dims of covariance function Xs, _ = self.cov_func._slice(Xs) - # If not provided, use Xs and c to set L - if self._L is None: - assert isinstance(self._c, (numbers.Real, np.ndarray, pt.TensorVariable)) - self.L = set_boundary(Xs, self._c) + if isinstance(self.cov_func, Periodic): + phi_cos, phi_sin = calc_basis_periodic(Xs, self.cov_func.period, self._m, tl=pt) + psd = self.cov_func.power_spectral_density(self._m) + return (phi_cos, phi_sin), psd + else: - self.L = self._L + # If not provided, use Xs and c to set L + if self._L is None: + assert isinstance(self._c, (numbers.Real, np.ndarray, pt.TensorVariable)) + self.L = set_boundary(Xs, self._c) + else: + self.L = self._L - eigvals = calc_eigenvalues(self.L, self._m, tl=pt) - phi = calc_eigenvectors(Xs, self.L, eigvals, self._m, tl=pt) - omega = pt.sqrt(eigvals) - psd = self.cov_func.power_spectral_density(omega) + eigvals = calc_eigenvalues(self.L, self._m, tl=pt) + phi = calc_eigenvectors(Xs, self.L, eigvals, self._m, tl=pt) + omega = pt.sqrt(eigvals) + psd = self.cov_func.power_spectral_density(omega) - i = int(self._drop_first == True) - return phi[:, i:], pt.sqrt(psd[i:]) + i = int(self._drop_first == True) + return phi[:, i:], pt.sqrt(psd[i:]) def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: ignore R""" @@ -319,16 +349,26 @@ def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: self._X_mean = pt.mean(X, axis=0) phi, sqrt_psd = self.prior_linearized(X - self._X_mean) - if self._parameterization == "noncentered": - self._beta = pm.Normal( - f"{name}_hsgp_coeffs_", size=self._m_star - int(self._drop_first) + if isinstance(self.cov_func, Periodic): + (phi_cos, phi_sin), psd = self.prior_linearized(X - self._X_mean) + self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(self._m * 2 - 1)) + # The first eigenfunction for the sine component is zero + # and so does not contribute to the approximation. + f = phi_cos @ (psd * self._beta[: self._m]) + phi_sin[..., 1:] @ ( + psd[1:] * self._beta[self._m + 1 :] ) - self._sqrt_psd = sqrt_psd - f = self.mean_func(X) + phi @ (self._beta * self._sqrt_psd) - elif self._parameterization == "centered": - self._beta = pm.Normal(f"{name}_hsgp_coeffs_", sigma=sqrt_psd) - f = self.mean_func(X) + phi @ self._beta + else: + if self._parameterization == "noncentered": + self._beta = pm.Normal( + f"{name}_hsgp_coeffs_", size=self._m_star - int(self._drop_first) + ) + self._sqrt_psd = sqrt_psd + f = self.mean_func(X) + phi @ (self._beta * self._sqrt_psd) + + elif self._parameterization == "centered": + self._beta = pm.Normal(f"{name}_hsgp_coeffs_", sigma=sqrt_psd) + f = self.mean_func(X) + phi @ self._beta self.f = pm.Deterministic(name, f, dims=dims) return self.f @@ -346,15 +386,27 @@ def _build_conditional(self, Xnew): ) Xnew, _ = self.cov_func._slice(Xnew) - eigvals = calc_eigenvalues(self.L, self._m, tl=pt) - phi = calc_eigenvectors(Xnew - X_mean, self.L, eigvals, self._m, tl=pt) - i = int(self._drop_first == True) - if self._parameterization == "noncentered": - return self.mean_func(Xnew) + phi[:, i:] @ (beta * sqrt_psd) + if isinstance(self.cov_func, Periodic): + phi_cos, phi_sin = calc_basis_periodic( + Xnew - X_mean, self.cov_func.period, self._m, tl=pt + ) + psd = self.cov_func.power_spectral_density(self._m) + + phi = phi_cos @ (psd * beta[: self._m]) + phi_sin[..., 1:] @ ( + psd[1:] * beta[self._m + 1 :] + ) + return self.mean_func(Xnew) + phi + else: + eigvals = calc_eigenvalues(self.L, self._m, tl=pt) + phi = calc_eigenvectors(Xnew - X_mean, self.L, eigvals, self._m, tl=pt) + i = int(self._drop_first == True) + + if self._parameterization == "noncentered": + return self.mean_func(Xnew) + phi[:, i:] @ (beta * sqrt_psd) - elif self._parameterization == "centered": - return self.mean_func(Xnew) + phi[:, i:] @ beta + elif self._parameterization == "centered": + return self.mean_func(Xnew) + phi[:, i:] @ beta def conditional(self, name: str, Xnew: TensorLike, dims: Optional[str] = None): # type: ignore R""" diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index 2521912437..3616659adf 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -116,7 +116,9 @@ def test_set_boundaries_3d(self, X2): assert np.all(L == 10) def test_parametrization(self): - err_msg = "`m` and L, if provided, must be sequences with one element per active dimension" + err_msg = ( + "`m` and `L`, if provided, must be sequences with one element per active dimension" + ) with pytest.raises(ValueError, match=err_msg): # m must be a list @@ -133,6 +135,13 @@ def test_parametrization(self): cov_func = pm.gp.cov.ExpQuad(1, ls=0.1) pm.gp.HSGP(m=[500], L=[12, 12], cov_func=cov_func) + with pytest.raises( + ValueError, + match="HSGP approximation for periodic kernel only implemented for 1-dimensional case.", + ): + cov_func = pm.gp.cov.Periodic(2, period=1, ls=[1, 2]) + pm.gp.HSGP(m=[500, 500], L=[12, 12], cov_func=cov_func) + # pass without error, cov_func has 2 active dimensions, c given as scalar cov_func = pm.gp.cov.ExpQuad(3, ls=[1, 2], active_dims=[0, 2]) pm.gp.HSGP(m=[50, 50], c=2, cov_func=cov_func) From 9134d37f7c206ad3cf525959cbae26d26d13007d Mon Sep 17 00:00:00 2001 From: theorashid Date: Fri, 15 Sep 2023 12:06:38 +0100 Subject: [PATCH 04/25] np.pi in w0 --- pymc/gp/hsgp_approx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 39b571060a..1665953af0 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -76,7 +76,7 @@ def calc_basis_periodic( Calculate basis vectors for the cosine series expansion of the periodic covariance function. These are derived from the Taylor series representation of the covariance. """ - w0 = (2 * tl.pi) / period # angular frequency defining the periodicity + w0 = (2 * np.pi) / period # angular frequency defining the periodicity m1 = tl.tile(w0 * Xs[:, None], m) m2 = tl.diag(tl.arange(m)) mw0x = m1 @ m2 From 00071bdb02cb96910dd20020c245b5add36f7dd3 Mon Sep 17 00:00:00 2001 From: theorashid Date: Sat, 7 Oct 2023 17:58:07 +0100 Subject: [PATCH 05/25] check arguments for period and nonperiod case --- pymc/gp/hsgp_approx.py | 48 ++++++++++++++++++++++++------------ tests/gp/test_hsgp_approx.py | 4 +-- 2 files changed, 34 insertions(+), 18 deletions(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 1665953af0..fd24554a15 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -119,7 +119,10 @@ class HSGP(Base): drop_first: bool Default `False`. Sometimes the first basis vector is quite "flat" and very similar to the intercept term. When there is an intercept in the model, ignoring the first basis - vector may improve sampling. + vector may improve sampling. This argument will be deprecated in future versions. + parameterization: str + Whether to use `centred` or `noncentered` parameterization when multiplying the + basis by the coefficients. cov_func: None, 2D array, or instance of Covariance The covariance function. Defaults to zero. mean_func: None, instance of Mean @@ -173,7 +176,7 @@ def __init__( L: Optional[Sequence[float]] = None, c: Optional[numbers.Real] = None, drop_first: bool = False, - parameterization="noncentered", + parameterization: Optional[str] = "noncentered", *, mean_func: Mean = Zero(), cov_func: Covariance, @@ -190,24 +193,37 @@ def __init__( raise ValueError(arg_err_msg) m = tuple(m) - if (L is None and c is None) or (L is not None and c is not None): - raise ValueError("Provide one of `c` or `L`") + if isinstance(cov_func, Periodic): + if cov_func.n_dims > 1: + raise ValueError( + "HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case." + ) + if (L or c or parameterization) is not None: + warnings.warn( + "Argument `L`, `c` or `parameterization` supplied but not used for `Periodic` kernel." + ) - if L is not None and (not isinstance(L, Sequence) or len(L) != cov_func.n_dims): - raise ValueError(arg_err_msg) + else: + if (L is None and c is None) or (L is not None and c is not None): + raise ValueError("Provide one of `c` or `L`") - if L is None and c is not None and c < 1.2: - warnings.warn("For an adequate approximation `c >= 1.2` is recommended.") + if L is not None and (not isinstance(L, Sequence) or len(L) != cov_func.n_dims): + raise ValueError(arg_err_msg) - parameterization = parameterization.lower().replace("-", "").replace("_", "") - if parameterization not in ["centered", "noncentered"]: - raise ValueError("`parameterization` must be either 'centered' or 'noncentered'.") - else: - self._parameterization = parameterization + if L is None and c is not None and c < 1.2: + warnings.warn("For an adequate approximation `c >= 1.2` is recommended.") - if isinstance(cov_func, Periodic) and cov_func.n_dims > 1: - raise ValueError( - "HSGP approximation for periodic kernel only implemented for 1-dimensional case." + parameterization = parameterization.lower().replace("-", "").replace("_", "") + if parameterization not in ["centered", "noncentered"]: + raise ValueError("`parameterization` must be either 'centered' or 'noncentered'.") + else: + self._parameterization = parameterization + + if drop_first: + warnings.warn( + "The drop_first argument will be deprecated in future versions." + " See https://github.com/pymc-devs/pymc/pull/6877", + DeprecationWarning, ) self._drop_first = drop_first diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index 3616659adf..c33ed76228 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -137,10 +137,10 @@ def test_parametrization(self): with pytest.raises( ValueError, - match="HSGP approximation for periodic kernel only implemented for 1-dimensional case.", + match="HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case.", ): cov_func = pm.gp.cov.Periodic(2, period=1, ls=[1, 2]) - pm.gp.HSGP(m=[500, 500], L=[12, 12], cov_func=cov_func) + pm.gp.HSGP(m=[500, 500], cov_func=cov_func) # pass without error, cov_func has 2 active dimensions, c given as scalar cov_func = pm.gp.cov.ExpQuad(3, ls=[1, 2], active_dims=[0, 2]) From 72038927a861874be51cc64e70829308cadfc276 Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 9 Oct 2023 15:34:30 +0100 Subject: [PATCH 06/25] tests run but failing --- pymc/gp/cov.py | 4 ++-- pymc/gp/hsgp_approx.py | 26 +++++++++++++++----------- tests/gp/test_hsgp_approx.py | 20 ++++++++++++++++---- 3 files changed, 33 insertions(+), 17 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index add2156661..9b87fd13bf 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -791,10 +791,11 @@ def full_from_distance(self, dist: TensorLike, squared: bool = False) -> TensorV r2 = dist if squared else dist**2 return pt.exp(-0.5 * r2) - def power_spectral_density(self, m: TensorLike) -> TensorVariable: + def power_spectral_density(self, J: TensorLike) -> TensorVariable: """ Technically, this is not a spectral density but these are the first `m` coefficients of the low rank approximation for the periodic kernel, which are used in the same way. + `J` is a vector of `np.arange(m)`. The coefficients of the HSGP approximation for the Periodic kernel are given by: @@ -806,7 +807,6 @@ def power_spectral_density(self, m: TensorLike) -> TensorVariable: where $\text{I}_j$ is the modified Bessel function of the first kind. """ a = 1 / pt.square(self.ls) - J = pt.arange(0, m) c = pt.where(J > 0, 2, 1) q2 = c * pt.iv(J, a) / pt.exp(a) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index fd24554a15..4e91a5b242 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -76,9 +76,12 @@ def calc_basis_periodic( Calculate basis vectors for the cosine series expansion of the periodic covariance function. These are derived from the Taylor series representation of the covariance. """ + if len(m) != 1: + raise ValueError("`Periodic` basis vectors only implemented for 1-dimensional case.") + m = m[0] # for compatibility with other kernels, m must be a sequence w0 = (2 * np.pi) / period # angular frequency defining the periodicity - m1 = tl.tile(w0 * Xs[:, None], m) - m2 = tl.diag(tl.arange(m)) + m1 = tl.tile(w0 * Xs, m) + m2 = tl.diag(tl.arange(0, m, 1)) mw0x = m1 @ m2 phi_cos = tl.cos(mw0x) phi_sin = tl.sin(mw0x) @@ -329,7 +332,8 @@ def prior_linearized(self, Xs: TensorLike): if isinstance(self.cov_func, Periodic): phi_cos, phi_sin = calc_basis_periodic(Xs, self.cov_func.period, self._m, tl=pt) - psd = self.cov_func.power_spectral_density(self._m) + J = pt.arange(0, self._m[0], 1) + psd = self.cov_func.power_spectral_density(J) return (phi_cos, phi_sin), psd else: @@ -367,12 +371,12 @@ def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: if isinstance(self.cov_func, Periodic): (phi_cos, phi_sin), psd = self.prior_linearized(X - self._X_mean) - self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(self._m * 2 - 1)) + + m0 = self._m[0] + self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(m0 * 2 - 1)) # The first eigenfunction for the sine component is zero # and so does not contribute to the approximation. - f = phi_cos @ (psd * self._beta[: self._m]) + phi_sin[..., 1:] @ ( - psd[1:] * self._beta[self._m + 1 :] - ) + f = phi_cos @ (psd * self._beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) else: if self._parameterization == "noncentered": @@ -407,11 +411,11 @@ def _build_conditional(self, Xnew): phi_cos, phi_sin = calc_basis_periodic( Xnew - X_mean, self.cov_func.period, self._m, tl=pt ) - psd = self.cov_func.power_spectral_density(self._m) + m0 = self._m[0] + J = pt.arange(0, m0, 1) + psd = self.cov_func.power_spectral_density(J) - phi = phi_cos @ (psd * beta[: self._m]) + phi_sin[..., 1:] @ ( - psd[1:] * beta[self._m + 1 :] - ) + phi = phi_cos @ (psd * beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * beta[m0:]) return self.mean_func(Xnew) + phi else: eigvals = calc_eigenvalues(self.L, self._m, tl=pt) diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index c33ed76228..a7548debee 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -166,8 +166,14 @@ def test_parametrization_drop_first(self, model, cov_func, X1, drop_first): else: assert n_coeffs == n_basis, "one was dropped when it shouldn't have been" - @pytest.mark.parametrize("cov_func", [pm.gp.cov.ExpQuad(1, ls=1)]) - @pytest.mark.parametrize("parameterization", ["centered", "noncentered"]) + @pytest.mark.parametrize( + "cov_func,parameterization", + [ + (pm.gp.cov.ExpQuad(1, ls=1), "centered"), + (pm.gp.cov.ExpQuad(1, ls=1), "noncentered"), + (pm.gp.cov.Periodic(1, period=1, ls=1), None), + ], + ) def test_prior(self, model, cov_func, X1, parameterization, rng): """Compare HSGP prior to unapproximated GP prior, pm.gp.Latent. Draw samples from the prior and compare them using MMD two sample test. Tests both centered and non-centered @@ -190,8 +196,14 @@ def test_prior(self, model, cov_func, X1, parameterization, rng): ) assert not reject, "H0 was rejected, even though HSGP and GP priors should match." - @pytest.mark.parametrize("cov_func", [pm.gp.cov.ExpQuad(1, ls=1)]) - @pytest.mark.parametrize("parameterization", ["centered", "noncentered"]) + @pytest.mark.parametrize( + "cov_func,parameterization", + [ + (pm.gp.cov.ExpQuad(1, ls=1), "centered"), + (pm.gp.cov.ExpQuad(1, ls=1), "noncentered"), + (pm.gp.cov.Periodic(1, period=1, ls=1), None), + ], + ) def test_conditional(self, model, cov_func, X1, parameterization): """Compare HSGP conditional to unapproximated GP prior, pm.gp.Latent. Draw samples from the prior and compare them using MMD two sample test. Tests both centered and non-centered From 1e87c67a1b085e62c55a97624f7c105d7f17b7ff Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 9 Oct 2023 16:02:54 +0100 Subject: [PATCH 07/25] fix one of the mypy errors --- pymc/gp/hsgp_approx.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 4e91a5b242..b22c37f9be 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -78,10 +78,10 @@ def calc_basis_periodic( """ if len(m) != 1: raise ValueError("`Periodic` basis vectors only implemented for 1-dimensional case.") - m = m[0] # for compatibility with other kernels, m must be a sequence + m0 = m[0] # for compatibility with other kernels, m must be a sequence w0 = (2 * np.pi) / period # angular frequency defining the periodicity - m1 = tl.tile(w0 * Xs, m) - m2 = tl.diag(tl.arange(0, m, 1)) + m1 = tl.tile(w0 * Xs, m0) + m2 = tl.diag(tl.arange(0, m0, 1)) mw0x = m1 @ m2 phi_cos = tl.cos(mw0x) phi_sin = tl.sin(mw0x) From 82d317d3b769c72cf66cdd4e857262a949547900 Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 9 Oct 2023 16:03:36 +0100 Subject: [PATCH 08/25] periodic cov pytest uses J not m because mpyp --- tests/gp/test_cov.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py index 20f480d87c..1de9cffea6 100644 --- a/tests/gp/test_cov.py +++ b/tests/gp/test_cov.py @@ -648,11 +648,11 @@ def test_psd(self): m = 5 a = 1 / np.square(ell) - J = np.arange(0, m) + J = np.arange(0, m, 1) true_1d_psd = np.where(J > 0, 2, 1) * iv(J, a) / np.exp(a) test_1d_psd = ( - pm.gp.cov.Periodic(1, period=P, ls=ell).power_spectral_density(m).flatten().eval() + pm.gp.cov.Periodic(1, period=P, ls=ell).power_spectral_density(J).flatten().eval() ) npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) From 2105976d16d60879e4105cd76069e8bd90f9879d Mon Sep 17 00:00:00 2001 From: theorashid Date: Tue, 10 Oct 2023 11:12:00 +0100 Subject: [PATCH 09/25] avoid AttributeError by init parameterization --- pymc/gp/hsgp_approx.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index b22c37f9be..c1bce8e660 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -205,6 +205,7 @@ def __init__( warnings.warn( "Argument `L`, `c` or `parameterization` supplied but not used for `Periodic` kernel." ) + self._parameterization = parameterization else: if (L is None and c is None) or (L is not None and c is not None): From 11d79e15a57eeeaf20064cd86113a997bfddfb1a Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 12 Oct 2023 18:35:41 +0100 Subject: [PATCH 10/25] get rid of mypy "None" error --- pymc/gp/hsgp_approx.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index c1bce8e660..94d4b6fc00 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -205,7 +205,6 @@ def __init__( warnings.warn( "Argument `L`, `c` or `parameterization` supplied but not used for `Periodic` kernel." ) - self._parameterization = parameterization else: if (L is None and c is None) or (L is not None and c is not None): @@ -217,11 +216,12 @@ def __init__( if L is None and c is not None and c < 1.2: warnings.warn("For an adequate approximation `c >= 1.2` is recommended.") - parameterization = parameterization.lower().replace("-", "").replace("_", "") - if parameterization not in ["centered", "noncentered"]: - raise ValueError("`parameterization` must be either 'centered' or 'noncentered'.") - else: - self._parameterization = parameterization + if parameterization is not None: + parameterization = parameterization.lower().replace("-", "").replace("_", "") + if parameterization not in ["centered", "noncentered"]: + raise ValueError( + "`parameterization` must be either 'centered' or 'noncentered'." + ) if drop_first: warnings.warn( @@ -235,6 +235,7 @@ def __init__( self._m_star = int(np.prod(self._m)) self._L = L self._c = c + self._parameterization = parameterization super().__init__(mean_func=mean_func, cov_func=cov_func) From c57d17b093eb0f3a78804f825df0191568729fb2 Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 12 Oct 2023 19:17:08 +0100 Subject: [PATCH 11/25] self.prior_linearized inside if avoids double call --- pymc/gp/hsgp_approx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 94d4b6fc00..3d173948b7 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -369,7 +369,6 @@ def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: Dimension name for the GP random variable. """ self._X_mean = pt.mean(X, axis=0) - phi, sqrt_psd = self.prior_linearized(X - self._X_mean) if isinstance(self.cov_func, Periodic): (phi_cos, phi_sin), psd = self.prior_linearized(X - self._X_mean) @@ -381,6 +380,7 @@ def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: f = phi_cos @ (psd * self._beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) else: + phi, sqrt_psd = self.prior_linearized(X - self._X_mean) if self._parameterization == "noncentered": self._beta = pm.Normal( f"{name}_hsgp_coeffs_", size=self._m_star - int(self._drop_first) From 5e1599c2b227d1b82d038c617fa9a912cfc174df Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 12 Oct 2023 19:19:05 +0100 Subject: [PATCH 12/25] need to add the mean function in prior --- pymc/gp/hsgp_approx.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 3d173948b7..d17c2c4376 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -377,7 +377,11 @@ def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(m0 * 2 - 1)) # The first eigenfunction for the sine component is zero # and so does not contribute to the approximation. - f = phi_cos @ (psd * self._beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) + f = ( + self.mean_func(X) + + phi_cos @ (psd * self._beta[:m0]) + + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) + ) else: phi, sqrt_psd = self.prior_linearized(X - self._X_mean) From abab45bc3f97c608cca6e9007798dbebb67a9fce Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 16 Oct 2023 08:26:28 +0100 Subject: [PATCH 13/25] run pre-commit locally --- pymc/gp/hsgp_approx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 49b80d692e..4d0b7bf745 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -333,7 +333,7 @@ def prior_linearized(self, Xs: TensorLike): # Index Xs using input_dim and active_dims of covariance function Xs, _ = self.cov_func._slice(Xs) - + if isinstance(self.cov_func, Periodic): phi_cos, phi_sin = calc_basis_periodic(Xs, self.cov_func.period, self._m, tl=pt) J = pt.arange(0, self._m[0], 1) From 51391b90e3e5cc4ca87708b3ae38499277cea6ab Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 16 Oct 2023 08:54:29 +0100 Subject: [PATCH 14/25] add docs to test about why test_prior fails --- tests/gp/test_hsgp_approx.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index a7548debee..1e1043c396 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -171,13 +171,18 @@ def test_parametrization_drop_first(self, model, cov_func, X1, drop_first): [ (pm.gp.cov.ExpQuad(1, ls=1), "centered"), (pm.gp.cov.ExpQuad(1, ls=1), "noncentered"), - (pm.gp.cov.Periodic(1, period=1, ls=1), None), + # (pm.gp.cov.Periodic(1, period=1, ls=1), None), ], ) def test_prior(self, model, cov_func, X1, parameterization, rng): """Compare HSGP prior to unapproximated GP prior, pm.gp.Latent. Draw samples from the prior and compare them using MMD two sample test. Tests both centered and non-centered parameterizations. + + Note: for `pm.gp.cov.Periodic`, this test does not pass and has been commented out. + The test passes more often when subtracting the mean from the mean from the samples. + It might be that the period is slightly off for the approximate power spectral density. + See https://github.com/pymc-devs/pymc/pull/6877/ for the full discussion. """ with model: hsgp = pm.gp.HSGP(m=[200], c=2.0, parameterization=parameterization, cov_func=cov_func) From c9db47be15f79421e78f4c7fa15c2f2ff7cdfee2 Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 16 Oct 2023 09:32:33 +0100 Subject: [PATCH 15/25] freshen up the docs --- pymc/gp/cov.py | 2 +- pymc/gp/hsgp_approx.py | 17 ++++++++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 98a68f6a1d..5a7a71af2d 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -797,7 +797,7 @@ def power_spectral_density(self, J: TensorLike) -> TensorVariable: the low rank approximation for the periodic kernel, which are used in the same way. `J` is a vector of `np.arange(m)`. - The coefficients of the HSGP approximation for the Periodic kernel are given by: + The coefficients of the HSGP approximation for the `Periodic` kernel are given by: .. math:: diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 4d0b7bf745..dbc5496e17 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -114,18 +114,19 @@ class HSGP(Base): L: list The boundary of the space for each `active_dim`. It is called the boundary condition. Choose L such that the domain `[-L, L]` contains all points in the column of X given by the - `active_dim`. + `active_dim`. This parameter is not used for `Periodic` covariance. c: float The proportion extension factor. Used to construct L from X. Defined as `S = max|X|` such that `X` is in `[-S, S]`. `L` is the calculated as `c * S`. One of `c` or `L` must be - provided. Further information can be found in Ruitort-Mayol et. al. + provided. Further information can be found in Ruitort-Mayol et al. This parameter is not + used for `Periodic` covariance. drop_first: bool Default `False`. Sometimes the first basis vector is quite "flat" and very similar to the intercept term. When there is an intercept in the model, ignoring the first basis vector may improve sampling. This argument will be deprecated in future versions. parameterization: str Whether to use `centred` or `noncentered` parameterization when multiplying the - basis by the coefficients. + basis by the coefficients. This parameter is not used for `Periodic` covariance. cov_func: None, 2D array, or instance of Covariance The covariance function. Defaults to zero. mean_func: None, instance of Mean @@ -269,6 +270,16 @@ def prior_linearized(self, Xs: TensorLike): the GP is constructed. If not, a RuntimeError is raised. Second, the `Xs` needs to be zero-centered, so it's mean must be subtracted. An example is given below. + Note, for the `Periodic` covariance, the approximation is not a spectral density, but a low + rank approximation based on expanding the periodic covariance function into a series of + stochastic resonators. However, these are used in the same way as the HSGP approximation. + Further information can be found in Appendix B of Ruitort-Mayol et al. + We no longer need to specify `L` or `c` because we are dealing with a periodic space and we + aren't using the Laplacian with boundary conditions. Rather than returning the Laplace + eigenfunctions and the square root of the power spectral density, in the periodic case we + return the cosine and sine terms of the periodic basis as a tuple, `(phi_cos, phi_sin)`, + and the coefficients of tbe expansion. + Parameters ---------- Xs: array-like From 6ffd9da49dfa05b93a6e28004a7f765ff727d48f Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 16 Oct 2023 13:28:55 +0100 Subject: [PATCH 16/25] improve coverage of args --- tests/gp/test_hsgp_approx.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index 1e1043c396..e42b38060a 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -120,6 +120,10 @@ def test_parametrization(self): "`m` and `L`, if provided, must be sequences with one element per active dimension" ) + with pytest.raises(ValueError, match="Provide one of `c` or `L`"): + cov_func = pm.gp.cov.ExpQuad(1, ls=0.1) + pm.gp.HSGP(m=[500], c=2, L=[12], cov_func=cov_func) + with pytest.raises(ValueError, match=err_msg): # m must be a list cov_func = pm.gp.cov.ExpQuad(1, ls=0.1) @@ -135,6 +139,12 @@ def test_parametrization(self): cov_func = pm.gp.cov.ExpQuad(1, ls=0.1) pm.gp.HSGP(m=[500], L=[12, 12], cov_func=cov_func) + with pytest.raises( + ValueError, match="`parameterization` must be either 'centered' or 'noncentered'." + ): + cov_func = pm.gp.cov.ExpQuad(2, ls=[1, 2], parameterization="wrong") + pm.gp.HSGP(m=[50, 50], L=[12, 12], cov_func=cov_func) + with pytest.raises( ValueError, match="HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case.", From 3bfccd32b1c054d8b9c03151dc2ebc29af17ab41 Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 16 Oct 2023 13:37:24 +0100 Subject: [PATCH 17/25] mypy ignore type --- pymc/gp/hsgp_approx.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index dbc5496e17..e69d820712 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -392,8 +392,8 @@ def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: # and so does not contribute to the approximation. f = ( self.mean_func(X) - + phi_cos @ (psd * self._beta[:m0]) - + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) + + phi_cos @ (psd * self._beta[:m0]) # type: ignore + + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) # type: ignore ) else: From 7aa87e2c54378fa59e13643d83e2e972ee548439 Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 16 Oct 2023 14:49:11 +0100 Subject: [PATCH 18/25] fix parameterization test --- tests/gp/test_hsgp_approx.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index e42b38060a..0b1b2e9de1 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -142,8 +142,8 @@ def test_parametrization(self): with pytest.raises( ValueError, match="`parameterization` must be either 'centered' or 'noncentered'." ): - cov_func = pm.gp.cov.ExpQuad(2, ls=[1, 2], parameterization="wrong") - pm.gp.HSGP(m=[50, 50], L=[12, 12], cov_func=cov_func) + cov_func = pm.gp.cov.ExpQuad(2, ls=[1, 2]) + pm.gp.HSGP(m=[50, 50], L=[12, 12], parameterization="wrong", cov_func=cov_func) with pytest.raises( ValueError, From 15c09bc4f29bfa0cda1c3ca051cdfffb70d980ee Mon Sep 17 00:00:00 2001 From: theorashid Date: Wed, 18 Oct 2023 12:23:13 +0100 Subject: [PATCH 19/25] use power_spectral_density_approx because cant sum --- pymc/gp/cov.py | 2 +- pymc/gp/hsgp_approx.py | 8 +++++--- tests/gp/test_cov.py | 5 ++++- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 5a7a71af2d..56aac7d39f 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -791,7 +791,7 @@ def full_from_distance(self, dist: TensorLike, squared: bool = False) -> TensorV r2 = dist if squared else dist**2 return pt.exp(-0.5 * r2) - def power_spectral_density(self, J: TensorLike) -> TensorVariable: + def power_spectral_density_approx(self, J: TensorLike) -> TensorVariable: """ Technically, this is not a spectral density but these are the first `m` coefficients of the low rank approximation for the periodic kernel, which are used in the same way. diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index e69d820712..59bc6d585a 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -98,7 +98,9 @@ class HSGP(Base): is largely similar to `gp.Latent`. Like `gp.Latent`, it does not assume a Gaussian noise model and can be used with any likelihood, or as a component anywhere within a model. Also like `gp.Latent`, it has `prior` and `conditional` methods. It supports any sum of covariance - functions that implement a `power_spectral_density` method. + functions that implement a `power_spectral_density` method (Note, this therefore excludes the + `Periodic` covariance function, which uses a different set of basis functions for the + approximation.). For information on choosing appropriate `m`, `L`, and `c`, refer Ruitort-Mayol et. al. or to the PyMC examples that use HSGP. @@ -348,7 +350,7 @@ def prior_linearized(self, Xs: TensorLike): if isinstance(self.cov_func, Periodic): phi_cos, phi_sin = calc_basis_periodic(Xs, self.cov_func.period, self._m, tl=pt) J = pt.arange(0, self._m[0], 1) - psd = self.cov_func.power_spectral_density(J) + psd = self.cov_func.power_spectral_density_approx(J) return (phi_cos, phi_sin), psd else: @@ -432,7 +434,7 @@ def _build_conditional(self, Xnew): ) m0 = self._m[0] J = pt.arange(0, m0, 1) - psd = self.cov_func.power_spectral_density(J) + psd = self.cov_func.power_spectral_density_approx(J) phi = phi_cos @ (psd * beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * beta[m0:]) return self.mean_func(Xnew) + phi diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py index 1de9cffea6..637dac29ec 100644 --- a/tests/gp/test_cov.py +++ b/tests/gp/test_cov.py @@ -652,7 +652,10 @@ def test_psd(self): true_1d_psd = np.where(J > 0, 2, 1) * iv(J, a) / np.exp(a) test_1d_psd = ( - pm.gp.cov.Periodic(1, period=P, ls=ell).power_spectral_density(J).flatten().eval() + pm.gp.cov.Periodic(1, period=P, ls=ell) + .power_spectral_density_approx(J) + .flatten() + .eval() ) npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) From 0681ca24f145f04e2d28edfbd6af2e68f4ab1f35 Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 19 Oct 2023 11:04:15 +0100 Subject: [PATCH 20/25] tbe -> the typo docs --- pymc/gp/hsgp_approx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 59bc6d585a..6ba6231467 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -280,7 +280,7 @@ def prior_linearized(self, Xs: TensorLike): aren't using the Laplacian with boundary conditions. Rather than returning the Laplace eigenfunctions and the square root of the power spectral density, in the periodic case we return the cosine and sine terms of the periodic basis as a tuple, `(phi_cos, phi_sin)`, - and the coefficients of tbe expansion. + and the coefficients of the expansion. Parameters ---------- From 36eced3b84de9415c784a7f99a649c42a454f9b9 Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 19 Oct 2023 11:09:59 +0100 Subject: [PATCH 21/25] expand warning and change the triggering clause --- pymc/gp/hsgp_approx.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 6ba6231467..981878716a 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -204,9 +204,11 @@ def __init__( raise ValueError( "HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case." ) - if (L or c or parameterization) is not None: + # only for `parameterization == "centered"` to avoid triggering for default `parameterization` case + if ((L or c) is not None) or (parameterization == "centered"): warnings.warn( "Argument `L`, `c` or `parameterization` supplied but not used for `Periodic` kernel." + "For this kernel, `m` is the only parameter controlling the approximation that must be supplied." ) else: From 56759ce304db6533ace4a90ccd87eab090baa1b1 Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 9 Nov 2023 15:03:39 +0000 Subject: [PATCH 22/25] separate into HSGP and HSGPPeriodic --- pymc/gp/__init__.py | 2 +- pymc/gp/hsgp_approx.py | 385 +++++++++++++++++++++++++---------- tests/gp/test_hsgp_approx.py | 79 +++++-- 3 files changed, 345 insertions(+), 121 deletions(-) diff --git a/pymc/gp/__init__.py b/pymc/gp/__init__.py index 99c8023398..6b9fd72204 100644 --- a/pymc/gp/__init__.py +++ b/pymc/gp/__init__.py @@ -22,4 +22,4 @@ MarginalKron, MarginalSparse, ) -from pymc.gp.hsgp_approx import HSGP +from pymc.gp.hsgp_approx import HSGP, HSGPPeriodic diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 981878716a..751b535854 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -23,7 +23,7 @@ import pymc as pm -from pymc.gp.cov import Covariance, Periodic +from pymc.gp.cov import Covariance from pymc.gp.gp import Base from pymc.gp.mean import Mean, Zero @@ -98,14 +98,14 @@ class HSGP(Base): is largely similar to `gp.Latent`. Like `gp.Latent`, it does not assume a Gaussian noise model and can be used with any likelihood, or as a component anywhere within a model. Also like `gp.Latent`, it has `prior` and `conditional` methods. It supports any sum of covariance - functions that implement a `power_spectral_density` method (Note, this therefore excludes the - `Periodic` covariance function, which uses a different set of basis functions for the - approximation.). + functions that implement a `power_spectral_density` method. (Note, this excludes the + `Periodic` covariance function, which uses a different set of basis functions for a + low rank approximation, as described in `HSGPPeriodic`.). - For information on choosing appropriate `m`, `L`, and `c`, refer Ruitort-Mayol et. al. or to + For information on choosing appropriate `m`, `L`, and `c`, refer Ruitort-Mayol et al. or to the PyMC examples that use HSGP. - To with with the HSGP in its "linearized" form, as a matrix of basis vectors and and vector of + To work with the HSGP in its "linearized" form, as a matrix of basis vectors and a vector of coefficients, see the method `prior_linearized`. Parameters @@ -116,20 +116,19 @@ class HSGP(Base): L: list The boundary of the space for each `active_dim`. It is called the boundary condition. Choose L such that the domain `[-L, L]` contains all points in the column of X given by the - `active_dim`. This parameter is not used for `Periodic` covariance. + `active_dim`. c: float The proportion extension factor. Used to construct L from X. Defined as `S = max|X|` such that `X` is in `[-S, S]`. `L` is the calculated as `c * S`. One of `c` or `L` must be - provided. Further information can be found in Ruitort-Mayol et al. This parameter is not - used for `Periodic` covariance. + provided. Further information can be found in Ruitort-Mayol et al. drop_first: bool Default `False`. Sometimes the first basis vector is quite "flat" and very similar to the intercept term. When there is an intercept in the model, ignoring the first basis vector may improve sampling. This argument will be deprecated in future versions. parameterization: str Whether to use `centred` or `noncentered` parameterization when multiplying the - basis by the coefficients. This parameter is not used for `Periodic` covariance. - cov_func: None, 2D array, or instance of Covariance + basis by the coefficients. + cov_func: None, 2D array, or instance of `Covariance` The covariance function. Defaults to zero. mean_func: None, instance of Mean The mean function. Defaults to zero. @@ -199,34 +198,19 @@ def __init__( raise ValueError(arg_err_msg) m = tuple(m) - if isinstance(cov_func, Periodic): - if cov_func.n_dims > 1: - raise ValueError( - "HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case." - ) - # only for `parameterization == "centered"` to avoid triggering for default `parameterization` case - if ((L or c) is not None) or (parameterization == "centered"): - warnings.warn( - "Argument `L`, `c` or `parameterization` supplied but not used for `Periodic` kernel." - "For this kernel, `m` is the only parameter controlling the approximation that must be supplied." - ) - - else: - if (L is None and c is None) or (L is not None and c is not None): - raise ValueError("Provide one of `c` or `L`") + if (L is None and c is None) or (L is not None and c is not None): + raise ValueError("Provide one of `c` or `L`") - if L is not None and (not isinstance(L, Sequence) or len(L) != cov_func.n_dims): - raise ValueError(arg_err_msg) + if L is not None and (not isinstance(L, Sequence) or len(L) != cov_func.n_dims): + raise ValueError(arg_err_msg) - if L is None and c is not None and c < 1.2: - warnings.warn("For an adequate approximation `c >= 1.2` is recommended.") + if L is None and c is not None and c < 1.2: + warnings.warn("For an adequate approximation `c >= 1.2` is recommended.") - if parameterization is not None: - parameterization = parameterization.lower().replace("-", "").replace("_", "") - if parameterization not in ["centered", "noncentered"]: - raise ValueError( - "`parameterization` must be either 'centered' or 'noncentered'." - ) + if parameterization is not None: + parameterization = parameterization.lower().replace("-", "").replace("_", "") + if parameterization not in ["centered", "noncentered"]: + raise ValueError("`parameterization` must be either 'centered' or 'noncentered'.") if drop_first: warnings.warn( @@ -274,16 +258,6 @@ def prior_linearized(self, Xs: TensorLike): the GP is constructed. If not, a RuntimeError is raised. Second, the `Xs` needs to be zero-centered, so it's mean must be subtracted. An example is given below. - Note, for the `Periodic` covariance, the approximation is not a spectral density, but a low - rank approximation based on expanding the periodic covariance function into a series of - stochastic resonators. However, these are used in the same way as the HSGP approximation. - Further information can be found in Appendix B of Ruitort-Mayol et al. - We no longer need to specify `L` or `c` because we are dealing with a periodic space and we - aren't using the Laplacian with boundary conditions. Rather than returning the Laplace - eigenfunctions and the square root of the power spectral density, in the periodic case we - return the cosine and sine terms of the periodic basis as a tuple, `(phi_cos, phi_sin)`, - and the coefficients of the expansion. - Parameters ---------- Xs: array-like @@ -310,7 +284,7 @@ def prior_linearized(self, Xs: TensorLike): ell = pm.InverseGamma("ell", mu=5.0, sigma=5.0) cov_func = eta**2 * pm.gp.cov.ExpQuad(1, ls=ell) - # m = [200] means 200 basis vectors for the first dimenison + # m = [200] means 200 basis vectors for the first dimension # L = [10] means the approximation is valid from Xs = [-10, 10] gp = pm.gp.HSGP(m=[200], L=[10], cov_func=cov_func) @@ -327,7 +301,7 @@ def prior_linearized(self, Xs: TensorLike): # Specify standard normal prior in the coefficients. The number of which # is given by the number of basis vectors, which is also saved in the GP object # as m_star. - beta = pm.Normal("beta", size=gp.m_star) + beta = pm.Normal("beta", size=gp._m_star) # The (non-centered) GP approximation is given by f = pm.Deterministic("f", phi @ (beta * sqrt_psd)) @@ -349,27 +323,20 @@ def prior_linearized(self, Xs: TensorLike): # Index Xs using input_dim and active_dims of covariance function Xs, _ = self.cov_func._slice(Xs) - if isinstance(self.cov_func, Periodic): - phi_cos, phi_sin = calc_basis_periodic(Xs, self.cov_func.period, self._m, tl=pt) - J = pt.arange(0, self._m[0], 1) - psd = self.cov_func.power_spectral_density_approx(J) - return (phi_cos, phi_sin), psd - + # If not provided, use Xs and c to set L + if self._L is None: + assert isinstance(self._c, (numbers.Real, np.ndarray, pt.TensorVariable)) + self.L = pt.as_tensor(set_boundary(Xs, self._c)) else: - # If not provided, use Xs and c to set L - if self._L is None: - assert isinstance(self._c, (numbers.Real, np.ndarray, pt.TensorVariable)) - self.L = pt.as_tensor(set_boundary(Xs, self._c)) - else: - self.L = self._L + self.L = self._L - eigvals = calc_eigenvalues(self.L, self._m, tl=pt) - phi = calc_eigenvectors(Xs, self.L, eigvals, self._m, tl=pt) - omega = pt.sqrt(eigvals) - psd = self.cov_func.power_spectral_density(omega) + eigvals = calc_eigenvalues(self.L, self._m, tl=pt) + phi = calc_eigenvectors(Xs, self.L, eigvals, self._m, tl=pt) + omega = pt.sqrt(eigvals) + psd = self.cov_func.power_spectral_density(omega) - i = int(self._drop_first == True) - return phi[:, i:], pt.sqrt(psd[i:]) + i = int(self._drop_first == True) + return phi[:, i:], pt.sqrt(psd[i:]) def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: ignore R""" @@ -387,31 +354,17 @@ def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: """ self._X_mean = pt.mean(X, axis=0) - if isinstance(self.cov_func, Periodic): - (phi_cos, phi_sin), psd = self.prior_linearized(X - self._X_mean) - - m0 = self._m[0] - self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(m0 * 2 - 1)) - # The first eigenfunction for the sine component is zero - # and so does not contribute to the approximation. - f = ( - self.mean_func(X) - + phi_cos @ (psd * self._beta[:m0]) # type: ignore - + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) # type: ignore + phi, sqrt_psd = self.prior_linearized(X - self._X_mean) + if self._parameterization == "noncentered": + self._beta = pm.Normal( + f"{name}_hsgp_coeffs_", size=self._m_star - int(self._drop_first) ) + self._sqrt_psd = sqrt_psd + f = self.mean_func(X) + phi @ (self._beta * self._sqrt_psd) - else: - phi, sqrt_psd = self.prior_linearized(X - self._X_mean) - if self._parameterization == "noncentered": - self._beta = pm.Normal( - f"{name}_hsgp_coeffs_", size=self._m_star - int(self._drop_first) - ) - self._sqrt_psd = sqrt_psd - f = self.mean_func(X) + phi @ (self._beta * self._sqrt_psd) - - elif self._parameterization == "centered": - self._beta = pm.Normal(f"{name}_hsgp_coeffs_", sigma=sqrt_psd) - f = self.mean_func(X) + phi @ self._beta + elif self._parameterization == "centered": + self._beta = pm.Normal(f"{name}_hsgp_coeffs_", sigma=sqrt_psd) + f = self.mean_func(X) + phi @ self._beta self.f = pm.Deterministic(name, f, dims=dims) return self.f @@ -430,26 +383,248 @@ def _build_conditional(self, Xnew): Xnew, _ = self.cov_func._slice(Xnew) - if isinstance(self.cov_func, Periodic): - phi_cos, phi_sin = calc_basis_periodic( - Xnew - X_mean, self.cov_func.period, self._m, tl=pt + eigvals = calc_eigenvalues(self.L, self._m, tl=pt) + phi = calc_eigenvectors(Xnew - X_mean, self.L, eigvals, self._m, tl=pt) + i = int(self._drop_first == True) + + if self._parameterization == "noncentered": + return self.mean_func(Xnew) + phi[:, i:] @ (beta * sqrt_psd) + + elif self._parameterization == "centered": + return self.mean_func(Xnew) + phi[:, i:] @ beta + + def conditional(self, name: str, Xnew: TensorLike, dims: Optional[str] = None): # type: ignore + R""" + Returns the (approximate) conditional distribution evaluated over new input locations + `Xnew`. + + Parameters + ---------- + name + Name of the random variable + Xnew : array-like + Function input values. + dims: None + Dimension name for the GP random variable. + """ + fnew = self._build_conditional(Xnew) + return pm.Deterministic(name, fnew, dims=dims) + + +class HSGPPeriodic(Base): + R""" + Hilbert Space Gaussian process approximation for the Periodic covariance function. + + Note, this is not actually a Hilbert space approximation, but it comes from the same + paper (Ruitort-Mayol et al., 2022. See Appendix B) and follows the same spirit: using a basis + approximation to a Gaussian process. In this case, the approximation is based on a series of + stochastic resonators. + + For these reasons, we have followed the same API as `gp.HSGP`, and can be used as a drop-in + replacement for `gp.Latent`. Like `gp.Latent`, it has `prior` and `conditional` methods. + + For information on choosing appropriate `m`, refer to Ruitort-Mayol et al.. Note, this approximation + is only implemented for the 1-D case. + + To work with the approximation in its "linearized" form, as a matrix of basis vectors and a + vector of coefficients, see the method `prior_linearized`. + + Parameters + ---------- + m: list + The number of basis vectors to use. Must be a list with one element. + cov_func: Instance of `Periodic` covariance + The covariance function. Defaults to zero. + mean_func: None, instance of Mean + The mean function. Defaults to zero. + + Examples + -------- + .. code:: python + + # A three dimensional column vector of inputs. + X = np.random.rand(100, 3) + + with pm.Model() as model: + # Specify the covariance function, only for the 1-D case + cov_func = pm.gp.cov.Periodic(1, period=1, ls=0.1) + + # Specify the approximation with 25 basis vectors + gp = pm.gp.HSGPPeriodic(m=[25], cov_func=cov_func) + + # Place a GP prior over the function f. + f = gp.prior("f", X=X) + + ... + + # After fitting or sampling, specify the distribution + # at new points with .conditional + Xnew = np.linspace(-1, 2, 50)[:, None] + + with model: + fcond = gp.conditional("fcond", Xnew=Xnew) + + References + ---------- + - Ruitort-Mayol, G., and Anderson, M., and Solin, A., and Vehtari, A. (2022). Practical + Hilbert Space Approximate Bayesian Gaussian Processes for Probabilistic Programming + """ + + def __init__( + self, + m: Sequence[int], + *, + mean_func: Mean = Zero(), + cov_func: Covariance, + ): + arg_err_msg = ( + "`m` and `L`, if provided, must be sequences with one element per active " + "dimension of the kernel or covariance function." + ) + + if not isinstance(m, Sequence): + raise ValueError(arg_err_msg) + + if cov_func.n_dims > 1: + raise ValueError( + "HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case." ) - m0 = self._m[0] - J = pt.arange(0, m0, 1) - psd = self.cov_func.power_spectral_density_approx(J) + m = tuple(m) - phi = phi_cos @ (psd * beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * beta[m0:]) - return self.mean_func(Xnew) + phi - else: - eigvals = calc_eigenvalues(self.L, self._m, tl=pt) - phi = calc_eigenvectors(Xnew - X_mean, self.L, eigvals, self._m, tl=pt) - i = int(self._drop_first == True) + self._m = m - if self._parameterization == "noncentered": - return self.mean_func(Xnew) + phi[:, i:] @ (beta * sqrt_psd) + super().__init__(mean_func=mean_func, cov_func=cov_func) + + def prior_linearized(self, Xs: TensorLike): + """Linearized version of the approximation. Returns the cosine and sine bases and coeffients + of the expansion needed to create the GP. + + This function allows the user to bypass the GP interface and work directly with the basis + and coefficients directly. This format allows the user to create predictions using + `pm.set_data` similarly to a linear model. It also enables computational speed ups in + multi-GP models since they may share the same basis. + + Correct results when using `prior_linearized` in tandem with `pm.set_data` and + `pm.MutableData` require that the `Xs` are zero-centered, so it's mean must be subtracted. + An example is given below. + + Parameters + ---------- + Xs: array-like + Function input values. Assumes they have been mean subtracted or centered at zero. + + Returns + ------- + (phi_cos, phi_sin): Tuple[array-like] + List of either Numpy or PyTensor 2D array of the cosine and sine fixed basis vectors. + There are n rows, one per row of `Xs` and `m` columns, one for each basis vector. + psd: array-like + Either a Numpy or PyTensor 1D array of the coefficients of the expansion. + + Examples + -------- + .. code:: python + + # A one dimensional column vector of inputs. + X = np.linspace(0, 10, 100)[:, None] + + with pm.Model() as model: + cov_func = pm.gp.cov.Periodic(1, period=1, ls=ell) + + # m = [200] means 200 basis vectors for the first dimension + gp = pm.gp.HSGPPeriodic(m=[200], cov_func=cov_func) + + # Order is important. First calculate the mean, then make X a shared variable, + # then subtract the mean. When X is mutated later, the correct mean will be + # subtracted. + X_mean = np.mean(X, axis=0) + X = pm.MutableData("X", X) + Xs = X - X_mean + + # Pass the zero-subtracted Xs in to the GP + (phi_cos, phi_sin), psd, sqrt_psd = gp.prior_linearized(Xs=Xs) + + # Specify standard normal prior in the coefficients. The number of which + # is twice the number of basis vectors minus one. + # This is so that each cosine term has a `beta` and all but one of the + # sine terms, as first eigenfunction for the sine component is zero + m0 = gp._m[0] + beta = pm.Normal("beta", size=(m0 * 2 - 1)) + + # The (non-centered) GP approximation is given by + f = pm.Deterministic( + "f", + phi_cos @ (psd * self._beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) + ) + ... + + + # Then it works just like a linear regression to predict on new data. + # First mutate the data X, + x_new = np.linspace(-10, 10, 100) + with model: + model.set_data("X", x_new[:, None]) + + # and then make predictions for the GP using posterior predictive sampling. + with model: + ppc = pm.sample_posterior_predictive(idata, var_names=["f"]) + """ + Xs, _ = self.cov_func._slice(Xs) + + phi_cos, phi_sin = calc_basis_periodic(Xs, self.cov_func.period, self._m, tl=pt) + J = pt.arange(0, self._m[0], 1) + psd = self.cov_func.power_spectral_density_approx(J) + return (phi_cos, phi_sin), psd + + def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): + R""" + Returns the (approximate) GP prior distribution evaluated over the input locations `X`. + For usage examples, refer to `pm.gp.Latent`. + + Parameters + ---------- + name: str + Name of the random variable + X: array-like + Function input values. + dims: None + Dimension name for the GP random variable. + """ + self._X_mean = pt.mean(X, axis=0) + + (phi_cos, phi_sin), psd = self.prior_linearized(X - self._X_mean) + + m0 = self._m[0] + self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(m0 * 2 - 1)) + # The first eigenfunction for the sine component is zero + # and so does not contribute to the approximation. + f = ( + self.mean_func(X) + + phi_cos @ (psd * self._beta[:m0]) # type: ignore + + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) # type: ignore + ) + + self.f = pm.Deterministic(name, f, dims=dims) + return self.f + + def _build_conditional(self, Xnew): + try: + beta, X_mean = self._beta, self._X_mean + + except AttributeError: + raise ValueError( + "Prior is not set, can't create a conditional. Call `.prior(name, X)` first." + ) + + Xnew, _ = self.cov_func._slice(Xnew) + + phi_cos, phi_sin = calc_basis_periodic(Xnew - X_mean, self.cov_func.period, self._m, tl=pt) + m0 = self._m[0] + J = pt.arange(0, m0, 1) + psd = self.cov_func.power_spectral_density_approx(J) - elif self._parameterization == "centered": - return self.mean_func(Xnew) + phi[:, i:] @ beta + phi = phi_cos @ (psd * beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * beta[m0:]) + return self.mean_func(Xnew) + phi def conditional(self, name: str, Xnew: TensorLike, dims: Optional[str] = None): # type: ignore R""" diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index 0b1b2e9de1..89fbbdeec1 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -76,7 +76,7 @@ def two_sample_test(sample1, sample2, n_sims=1000, alpha=0.05): return h0, mmd, critical_value, mmd > critical_value -class TestHSGP: +class _BaseFixtures: @pytest.fixture def rng(self): return np.random.RandomState(10) @@ -105,6 +105,8 @@ def X2(self, data): def model(self): return pm.Model() + +class TestHSGP(_BaseFixtures): def test_set_boundaries_1d(self, X1): X1s = X1 - np.mean(X1, axis=0) L = pm.gp.hsgp_approx.set_boundary(X1s, c=2).eval() @@ -145,13 +147,6 @@ def test_parametrization(self): cov_func = pm.gp.cov.ExpQuad(2, ls=[1, 2]) pm.gp.HSGP(m=[50, 50], L=[12, 12], parameterization="wrong", cov_func=cov_func) - with pytest.raises( - ValueError, - match="HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case.", - ): - cov_func = pm.gp.cov.Periodic(2, period=1, ls=[1, 2]) - pm.gp.HSGP(m=[500, 500], cov_func=cov_func) - # pass without error, cov_func has 2 active dimensions, c given as scalar cov_func = pm.gp.cov.ExpQuad(3, ls=[1, 2], active_dims=[0, 2]) pm.gp.HSGP(m=[50, 50], c=2, cov_func=cov_func) @@ -188,11 +183,6 @@ def test_prior(self, model, cov_func, X1, parameterization, rng): """Compare HSGP prior to unapproximated GP prior, pm.gp.Latent. Draw samples from the prior and compare them using MMD two sample test. Tests both centered and non-centered parameterizations. - - Note: for `pm.gp.cov.Periodic`, this test does not pass and has been commented out. - The test passes more often when subtracting the mean from the mean from the samples. - It might be that the period is slightly off for the approximate power spectral density. - See https://github.com/pymc-devs/pymc/pull/6877/ for the full discussion. """ with model: hsgp = pm.gp.HSGP(m=[200], c=2.0, parameterization=parameterization, cov_func=cov_func) @@ -209,14 +199,13 @@ def test_prior(self, model, cov_func, X1, parameterization, rng): h0, mmd, critical_value, reject = two_sample_test( samples1, samples2, n_sims=500, alpha=0.01 ) - assert not reject, "H0 was rejected, even though HSGP and GP priors should match." + assert not reject, f"H0 was rejected, even though HSGP and GP priors should match." @pytest.mark.parametrize( "cov_func,parameterization", [ (pm.gp.cov.ExpQuad(1, ls=1), "centered"), (pm.gp.cov.ExpQuad(1, ls=1), "noncentered"), - (pm.gp.cov.Periodic(1, period=1, ls=1), None), ], ) def test_conditional(self, model, cov_func, X1, parameterization): @@ -238,3 +227,63 @@ def test_conditional(self, model, cov_func, X1, parameterization): samples1, samples2, n_sims=500, alpha=0.01 ) assert not reject, "H0 was rejected, even though HSGP prior and conditional should match." + + +class TestHSGPPeriodic(_BaseFixtures): + def test_parametrization(self): + with pytest.raises( + ValueError, + match="HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case.", + ): + cov_func = pm.gp.cov.Periodic(2, period=1, ls=[1, 2]) + pm.gp.HSGPPeriodic(m=[500, 500], cov_func=cov_func) + + @pytest.mark.parametrize("cov_func", [pm.gp.cov.Periodic(1, period=1, ls=1)]) + @pytest.mark.xfail( + reason="For `pm.gp.cov.Periodic`, this test does not pass.\ + The mmd is around `0.0468`.\ + The test passes more often when subtracting the mean from the mean from the samples.\ + It might be that the period is slightly off for the approximate power spectral density.\ + See https://github.com/pymc-devs/pymc/pull/6877/ for the full discussion." + ) + def test_prior(self, model, cov_func, X1, rng): + """Compare HSGPPeriodic prior to unapproximated GP prior, pm.gp.Latent. Draw samples from the + prior and compare them using MMD two sample test. + """ + with model: + hsgp = pm.gp.HSGPPeriodic(m=[200], cov_func=cov_func) + f1 = hsgp.prior("f1", X=X1) + + gp = pm.gp.Latent(cov_func=cov_func) + f2 = gp.prior("f2", X=X1) + + idata = pm.sample_prior_predictive(samples=1000, random_seed=rng) + + samples1 = az.extract(idata.prior["f1"])["f1"].values.T + samples2 = az.extract(idata.prior["f2"])["f2"].values.T + + h0, mmd, critical_value, reject = two_sample_test( + samples1, samples2, n_sims=500, alpha=0.01 + ) + assert not reject, f"H0 was rejected, {mmd} even though HSGP and GP priors should match." + + @pytest.mark.parametrize("cov_func", [pm.gp.cov.Periodic(1, period=1, ls=1)]) + def test_conditional_periodic(self, model, cov_func, X1): + """Compare HSGPPeriodic conditional to unapproximated GP prior, pm.gp.Latent. Draw samples + from the prior and compare them using MMD two sample test. The conditional should match the + prior when no data is observed. + """ + with model: + hsgp = pm.gp.HSGPPeriodic(m=[100], cov_func=cov_func) + f = hsgp.prior("f", X=X1) + fc = hsgp.conditional("fc", Xnew=X1) + + idata = pm.sample_prior_predictive(samples=1000) + + samples1 = az.extract(idata.prior["f"])["f"].values.T + samples2 = az.extract(idata.prior["fc"])["fc"].values.T + + h0, mmd, critical_value, reject = two_sample_test( + samples1, samples2, n_sims=500, alpha=0.01 + ) + assert not reject, "H0 was rejected, even though HSGP prior and conditional should match." From 227584f6422da1b7bf170ed736c4a5bc337286e4 Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 9 Nov 2023 15:05:41 +0000 Subject: [PATCH 23/25] remove commented periodic in HSGP test --- tests/gp/test_hsgp_approx.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index 89fbbdeec1..7797772d22 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -176,7 +176,6 @@ def test_parametrization_drop_first(self, model, cov_func, X1, drop_first): [ (pm.gp.cov.ExpQuad(1, ls=1), "centered"), (pm.gp.cov.ExpQuad(1, ls=1), "noncentered"), - # (pm.gp.cov.Periodic(1, period=1, ls=1), None), ], ) def test_prior(self, model, cov_func, X1, parameterization, rng): From 32d9e8635ebab75f8e0802e5c0abb2bd7568a2bb Mon Sep 17 00:00:00 2001 From: theorashid Date: Thu, 9 Nov 2023 15:20:46 +0000 Subject: [PATCH 24/25] appease mypy --- pymc/gp/hsgp_approx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 751b535854..7e7992a28b 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -576,7 +576,7 @@ def prior_linearized(self, Xs: TensorLike): psd = self.cov_func.power_spectral_density_approx(J) return (phi_cos, phi_sin), psd - def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): + def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: ignore R""" Returns the (approximate) GP prior distribution evaluated over the input locations `X`. For usage examples, refer to `pm.gp.Latent`. From 56963e0214c24767211044a84abef09472ea1add Mon Sep 17 00:00:00 2001 From: theorashid Date: Wed, 22 Nov 2023 12:46:45 +0000 Subject: [PATCH 25/25] add scale parameter for HSGP --- pymc/gp/hsgp_approx.py | 86 ++++++++++++++++++++---------------- tests/gp/test_hsgp_approx.py | 33 +++++++++++--- 2 files changed, 74 insertions(+), 45 deletions(-) diff --git a/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 7e7992a28b..c0b8d41806 100644 --- a/pymc/gp/hsgp_approx.py +++ b/pymc/gp/hsgp_approx.py @@ -23,7 +23,7 @@ import pymc as pm -from pymc.gp.cov import Covariance +from pymc.gp.cov import Covariance, Periodic from pymc.gp.gp import Base from pymc.gp.mean import Mean, Zero @@ -69,19 +69,16 @@ def calc_eigenvectors( def calc_basis_periodic( Xs: TensorLike, period: TensorLike, - m: Sequence[int], + m: int, tl: ModuleType = np, ): """ Calculate basis vectors for the cosine series expansion of the periodic covariance function. These are derived from the Taylor series representation of the covariance. """ - if len(m) != 1: - raise ValueError("`Periodic` basis vectors only implemented for 1-dimensional case.") - m0 = m[0] # for compatibility with other kernels, m must be a sequence w0 = (2 * np.pi) / period # angular frequency defining the periodicity - m1 = tl.tile(w0 * Xs, m0) - m2 = tl.diag(tl.arange(0, m0, 1)) + m1 = tl.tile(w0 * Xs, m) + m2 = tl.diag(tl.arange(0, m, 1)) mw0x = m1 @ m2 phi_cos = tl.cos(mw0x) phi_sin = tl.sin(mw0x) @@ -128,8 +125,8 @@ class HSGP(Base): parameterization: str Whether to use `centred` or `noncentered` parameterization when multiplying the basis by the coefficients. - cov_func: None, 2D array, or instance of `Covariance` - The covariance function. Defaults to zero. + cov_func: Covariance function, must be an instance of `Stationary` and implement a + `power_spectral_density` method. mean_func: None, instance of Mean The mean function. Defaults to zero. @@ -431,10 +428,11 @@ class HSGPPeriodic(Base): Parameters ---------- - m: list - The number of basis vectors to use. Must be a list with one element. - cov_func: Instance of `Periodic` covariance - The covariance function. Defaults to zero. + m: int + The number of basis vectors to use. Must be a positive integer. + scale: TensorLike + The standard deviation (square root of the variance) of the GP effect. Defaults to 1.0. + cov_func: Must be an instance of instance of `Periodic` covariance mean_func: None, instance of Mean The mean function. Defaults to zero. @@ -447,10 +445,11 @@ class HSGPPeriodic(Base): with pm.Model() as model: # Specify the covariance function, only for the 1-D case + scale = pm.HalfNormal(10) cov_func = pm.gp.cov.Periodic(1, period=1, ls=0.1) # Specify the approximation with 25 basis vectors - gp = pm.gp.HSGPPeriodic(m=[25], cov_func=cov_func) + gp = pm.gp.HSGPPeriodic(m=25, scale=scale, cov_func=cov_func) # Place a GP prior over the function f. f = gp.prior("f", X=X) @@ -472,31 +471,37 @@ class HSGPPeriodic(Base): def __init__( self, - m: Sequence[int], + m: int, + scale: Optional[Union[float, TensorLike]] = 1.0, *, mean_func: Mean = Zero(), - cov_func: Covariance, + cov_func: Periodic, ): - arg_err_msg = ( - "`m` and `L`, if provided, must be sequences with one element per active " - "dimension of the kernel or covariance function." - ) + arg_err_msg = "`m` must be a positive integer as the `Periodic` kernel approximation is only implemented for 1-dimensional case." - if not isinstance(m, Sequence): + if not isinstance(m, int): + raise ValueError(arg_err_msg) + + if m <= 0: raise ValueError(arg_err_msg) + if not isinstance(cov_func, Periodic): + raise ValueError( + "`cov_func` must be an instance of a `Periodic` kernel only. Use the `scale` parameter to control the variance." + ) + if cov_func.n_dims > 1: raise ValueError( "HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case." ) - m = tuple(m) self._m = m + self.scale = scale super().__init__(mean_func=mean_func, cov_func=cov_func) def prior_linearized(self, Xs: TensorLike): - """Linearized version of the approximation. Returns the cosine and sine bases and coeffients + """Linearized version of the approximation. Returns the cosine and sine bases and coefficients of the expansion needed to create the GP. This function allows the user to bypass the GP interface and work directly with the basis @@ -529,10 +534,11 @@ def prior_linearized(self, Xs: TensorLike): X = np.linspace(0, 10, 100)[:, None] with pm.Model() as model: + scale = pm.HalfNormal(10) cov_func = pm.gp.cov.Periodic(1, period=1, ls=ell) - # m = [200] means 200 basis vectors for the first dimension - gp = pm.gp.HSGPPeriodic(m=[200], cov_func=cov_func) + # m=200 means 200 basis vectors + gp = pm.gp.HSGPPeriodic(m=200, scale=scale, cov_func=cov_func) # Order is important. First calculate the mean, then make X a shared variable, # then subtract the mean. When X is mutated later, the correct mean will be @@ -542,19 +548,19 @@ def prior_linearized(self, Xs: TensorLike): Xs = X - X_mean # Pass the zero-subtracted Xs in to the GP - (phi_cos, phi_sin), psd, sqrt_psd = gp.prior_linearized(Xs=Xs) + (phi_cos, phi_sin), psd = gp.prior_linearized(Xs=Xs) # Specify standard normal prior in the coefficients. The number of which # is twice the number of basis vectors minus one. # This is so that each cosine term has a `beta` and all but one of the # sine terms, as first eigenfunction for the sine component is zero - m0 = gp._m[0] - beta = pm.Normal("beta", size=(m0 * 2 - 1)) + m = gp._m + beta = pm.Normal("beta", size=(m * 2 - 1)) # The (non-centered) GP approximation is given by f = pm.Deterministic( "f", - phi_cos @ (psd * self._beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) + phi_cos @ (psd * beta[:m]) + phi_sin[..., 1:] @ (psd[1:] * beta[m:]) ) ... @@ -572,8 +578,9 @@ def prior_linearized(self, Xs: TensorLike): Xs, _ = self.cov_func._slice(Xs) phi_cos, phi_sin = calc_basis_periodic(Xs, self.cov_func.period, self._m, tl=pt) - J = pt.arange(0, self._m[0], 1) - psd = self.cov_func.power_spectral_density_approx(J) + J = pt.arange(0, self._m, 1) + # rescale basis coefficients by the sqrt variance term + psd = self.scale * self.cov_func.power_spectral_density_approx(J) return (phi_cos, phi_sin), psd def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: ignore @@ -594,14 +601,14 @@ def prior(self, name: str, X: TensorLike, dims: Optional[str] = None): # type: (phi_cos, phi_sin), psd = self.prior_linearized(X - self._X_mean) - m0 = self._m[0] - self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(m0 * 2 - 1)) + m = self._m + self._beta = pm.Normal(f"{name}_hsgp_coeffs_", size=(m * 2 - 1)) # The first eigenfunction for the sine component is zero # and so does not contribute to the approximation. f = ( self.mean_func(X) - + phi_cos @ (psd * self._beta[:m0]) # type: ignore - + phi_sin[..., 1:] @ (psd[1:] * self._beta[m0:]) # type: ignore + + phi_cos @ (psd * self._beta[:m]) # type: ignore + + phi_sin[..., 1:] @ (psd[1:] * self._beta[m:]) # type: ignore ) self.f = pm.Deterministic(name, f, dims=dims) @@ -619,11 +626,12 @@ def _build_conditional(self, Xnew): Xnew, _ = self.cov_func._slice(Xnew) phi_cos, phi_sin = calc_basis_periodic(Xnew - X_mean, self.cov_func.period, self._m, tl=pt) - m0 = self._m[0] - J = pt.arange(0, m0, 1) - psd = self.cov_func.power_spectral_density_approx(J) + m = self._m + J = pt.arange(0, m, 1) + # rescale basis coefficients by the sqrt variance term + psd = self.scale * self.cov_func.power_spectral_density_approx(J) - phi = phi_cos @ (psd * beta[:m0]) + phi_sin[..., 1:] @ (psd[1:] * beta[m0:]) + phi = phi_cos @ (psd * beta[:m]) + phi_sin[..., 1:] @ (psd[1:] * beta[m:]) return self.mean_func(Xnew) + phi def conditional(self, name: str, Xnew: TensorLike, dims: Optional[str] = None): # type: ignore diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index 7797772d22..7a1c70de9b 100644 --- a/tests/gp/test_hsgp_approx.py +++ b/tests/gp/test_hsgp_approx.py @@ -230,14 +230,35 @@ def test_conditional(self, model, cov_func, X1, parameterization): class TestHSGPPeriodic(_BaseFixtures): def test_parametrization(self): + err_msg = "`m` must be a positive integer as the `Periodic` kernel approximation is only implemented for 1-dimensional case." + + with pytest.raises(ValueError, match=err_msg): + # `m` must be a positive integer, not a list + cov_func = pm.gp.cov.Periodic(1, period=1, ls=0.1) + pm.gp.HSGPPeriodic(m=[500], cov_func=cov_func) + + with pytest.raises(ValueError, match=err_msg): + # `m`` must be a positive integer + cov_func = pm.gp.cov.Periodic(1, period=1, ls=0.1) + pm.gp.HSGPPeriodic(m=-1, cov_func=cov_func) + + with pytest.raises( + ValueError, + match="`cov_func` must be an instance of a `Periodic` kernel only. Use the `scale` parameter to control the variance.", + ): + # `cov_func` must be `Periodic` only + cov_func = 5.0 * pm.gp.cov.Periodic(1, period=1, ls=0.1) + pm.gp.HSGPPeriodic(m=500, cov_func=cov_func) + with pytest.raises( ValueError, match="HSGP approximation for `Periodic` kernel only implemented for 1-dimensional case.", ): cov_func = pm.gp.cov.Periodic(2, period=1, ls=[1, 2]) - pm.gp.HSGPPeriodic(m=[500, 500], cov_func=cov_func) + pm.gp.HSGPPeriodic(m=500, scale=0.5, cov_func=cov_func) @pytest.mark.parametrize("cov_func", [pm.gp.cov.Periodic(1, period=1, ls=1)]) + @pytest.mark.parametrize("eta", [100.0]) @pytest.mark.xfail( reason="For `pm.gp.cov.Periodic`, this test does not pass.\ The mmd is around `0.0468`.\ @@ -245,15 +266,15 @@ def test_parametrization(self): It might be that the period is slightly off for the approximate power spectral density.\ See https://github.com/pymc-devs/pymc/pull/6877/ for the full discussion." ) - def test_prior(self, model, cov_func, X1, rng): + def test_prior(self, model, cov_func, eta, X1, rng): """Compare HSGPPeriodic prior to unapproximated GP prior, pm.gp.Latent. Draw samples from the prior and compare them using MMD two sample test. """ with model: - hsgp = pm.gp.HSGPPeriodic(m=[200], cov_func=cov_func) + hsgp = pm.gp.HSGPPeriodic(m=200, scale=eta, cov_func=cov_func) f1 = hsgp.prior("f1", X=X1) - gp = pm.gp.Latent(cov_func=cov_func) + gp = pm.gp.Latent(cov_func=eta**2 * cov_func) f2 = gp.prior("f2", X=X1) idata = pm.sample_prior_predictive(samples=1000, random_seed=rng) @@ -268,12 +289,12 @@ def test_prior(self, model, cov_func, X1, rng): @pytest.mark.parametrize("cov_func", [pm.gp.cov.Periodic(1, period=1, ls=1)]) def test_conditional_periodic(self, model, cov_func, X1): - """Compare HSGPPeriodic conditional to unapproximated GP prior, pm.gp.Latent. Draw samples + """Compare HSGPPeriodic conditional to HSGPPeriodic prior. Draw samples from the prior and compare them using MMD two sample test. The conditional should match the prior when no data is observed. """ with model: - hsgp = pm.gp.HSGPPeriodic(m=[100], cov_func=cov_func) + hsgp = pm.gp.HSGPPeriodic(m=100, cov_func=cov_func) f = hsgp.prior("f", X=X1) fc = hsgp.conditional("fc", Xnew=X1)