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/cov.py b/pymc/gp/cov.py index 905fdef930..56aac7d39f 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_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. + `J` is a vector of `np.arange(m)`. + + 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) + 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/pymc/gp/hsgp_approx.py b/pymc/gp/hsgp_approx.py index 0e13c77e25..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 @@ -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: 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 * np.pi) / period # angular frequency defining the periodicity + 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) + return phi_cos, phi_sin + + class HSGP(Base): R""" Hilbert Space Gaussian process approximation. @@ -76,12 +95,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. + 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 @@ -96,13 +117,16 @@ class HSGP(Base): 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. 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. - cov_func: None, 2D array, or instance of Covariance - The covariance function. Defaults to zero. + 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: 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. @@ -154,13 +178,13 @@ 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, ): 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." ) @@ -180,11 +204,17 @@ 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 is not None: + 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 self._m = m @@ -193,6 +223,7 @@ def __init__( if L is not None: self._L = pt.as_tensor(L) self._c = c + self._parameterization = parameterization super().__init__(mean_func=mean_func, cov_func=cov_func) @@ -250,7 +281,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) @@ -267,7 +298,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)) @@ -292,7 +323,9 @@ def prior_linearized(self, Xs: TensorLike): # 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)) + self.L = pt.as_tensor(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) @@ -317,8 +350,8 @@ 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) + 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) @@ -346,6 +379,7 @@ 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) @@ -372,3 +406,247 @@ def conditional(self, name: str, Xnew: TensorLike, dims: Optional[str] = None): """ 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: 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. + + 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 + 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, scale=scale, 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: int, + scale: Optional[Union[float, TensorLike]] = 1.0, + *, + mean_func: Mean = Zero(), + cov_func: Periodic, + ): + 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, 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." + ) + + 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 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 + 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: + scale = pm.HalfNormal(10) + cov_func = pm.gp.cov.Periodic(1, period=1, ls=ell) + + # 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 + # 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 = 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 + 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 * beta[:m]) + phi_sin[..., 1:] @ (psd[1:] * beta[m:]) + ) + ... + + + # 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, 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 + 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) + + 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[:m]) # type: ignore + + phi_sin[..., 1:] @ (psd[1:] * self._beta[m:]) # 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) + 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[: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 + 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) diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py index 750d5fe6a1..637dac29ec 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,24 @@ 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, 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_approx(J) + .flatten() + .eval() + ) + npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) + class TestLinear: def test_1d(self): diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py index 96465e9437..7a1c70de9b 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,15 +105,8 @@ 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 +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() @@ -125,7 +118,13 @@ 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="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 @@ -142,6 +141,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]) + pm.gp.HSGP(m=[50, 50], L=[12, 12], parameterization="wrong", 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) @@ -150,6 +155,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,7 +171,13 @@ 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("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"), + ], + ) 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 @@ -186,9 +198,15 @@ 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." - - @pytest.mark.parametrize("parameterization", ["centered", "noncentered"]) + 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"), + ], + ) 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 @@ -208,3 +226,84 @@ 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): + 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, 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`.\ + 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, 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, scale=eta, cov_func=cov_func) + f1 = hsgp.prior("f1", X=X1) + + 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) + + 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 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) + 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."