From e3f598acd6e332fa1cc245b27405d3c203ef51eb Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Thu, 4 Aug 2022 23:55:42 -0700 Subject: [PATCH 01/17] add power spectral density methods to ExpQuad, Matern32, Matern52. allow additivity --- pymc/gp/cov.py | 134 +++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 118 insertions(+), 16 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 41919a0dfb..f1524d3829 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -46,6 +46,20 @@ ] +def _maybe_squeeze(value): + if ( + isinstance(value, aesara.compile.SharedVariable) + and value.get_value().squeeze().shape == () + ): + return at.squeeze(value) + elif np.asarray(value).squeeze().shape == (): + return np.squeeze(value) + elif isinstance(value, Number): + return value + + raise ValueError("A scalar value is required.") + + class Covariance: r""" Base class for all kernels/covariance functions. @@ -67,6 +81,14 @@ def __init__(self, input_dim, active_dims=None): else: self.active_dims = np.asarray(active_dims, int) + @property + def D(self): + """ The dimensionality of the input, as taken from the + `active_dims`. + """ + # Evaluate lazily in-case this changes. + return len(self.active_dims) + def __call__(self, X, Xs=None, diag=False): r""" Evaluate the kernel/covariance function. @@ -121,19 +143,7 @@ def __rmul__(self, other): return self.__mul__(other) def __pow__(self, other): - if ( - isinstance(other, aesara.compile.SharedVariable) - and other.get_value().squeeze().shape == () - ): - other = at.squeeze(other) - return Exponentiated(self, other) - elif isinstance(other, Number): - return Exponentiated(self, other) - elif np.asarray(other).squeeze().shape == (): - other = np.squeeze(other) - return Exponentiated(self, other) - - raise ValueError("A covariance function can only be exponentiated by a scalar value") + return Exponentiated(self, self._maybe_squeeze(other)) def __array_wrap__(self, result): """ @@ -159,6 +169,9 @@ def __array_wrap__(self, result): raise TypeError( f"Unknown Covariance combination type {result[0][0]}. Known types are `Add` or `Prod`." ) + + def psd(self): + raise NotImplementedError("Only covariance functions of type `Stationary` may implement a `.psd`.") class Combination(Covariance): @@ -174,7 +187,10 @@ def __init__(self, factor_list): else: self.factor_list.append(factor) - def merge_factors(self, X, Xs=None, diag=False): + def merge_factors_cov(self, X, Xs=None, diag=False): + """ Called to evaluate either all the sums or all the + products of kernels that are possible to evaluate. + """ factor_list = [] for factor in self.factor_list: # make sure diag=True is handled properly @@ -201,15 +217,51 @@ def merge_factors(self, X, Xs=None, diag=False): factor_list.append(factor) return factor_list + def merge_factors_psd(self, omega): + """ Called to evaluatate spectral densities of combination + kernels when possible. Implements a more restricted set + of rules than `merge_factors_cov` -- just additivity of + stationary covariances with defined power spectral densities + and mutliplication by scalars. + """ + factor_list = [] + for i, factor in enumerate(self.factor_list): + if isinstance(factor, Covariance): + # If it's a covariance, calculate the psd + factor_list.append(factor.psd(omega)) + elif isinstance( + factor, + ( + TensorConstant, + TensorVariable, + TensorSharedVariable, + ), + ): + # If it's an aesara variable, make sure it's a scalar + ## TODO this check doesnt do the right thing, arrays pass through + factor = _maybe_squeeze(factor) + factor_list.append(factor) + + else: + # Otherwise try to keep going + factor_list.append(factor) + return factor_list + class Add(Combination): def __call__(self, X, Xs=None, diag=False): - return reduce(add, self.merge_factors(X, Xs, diag)) + return reduce(add, self.merge_factors_cov(X, Xs, diag)) + def psd(self, omega): + return reduce(add, self.merge_factors_psd(omega)) + class Prod(Combination): def __call__(self, X, Xs=None, diag=False): - return reduce(mul, self.merge_factors(X, Xs, diag)) + return reduce(mul, self.merge_factors_cov(X, Xs, diag)) + + def psd(self, omega): + return reduce(mul, self.merge_factors_psd(omega)) class Exponentiated(Covariance): @@ -406,6 +458,9 @@ def diag(self, X): def full(self, X, Xs=None): raise NotImplementedError + + def psd(self, omega): + raise NotImplementedError class Periodic(Stationary): @@ -450,11 +505,24 @@ class ExpQuad(Stationary): .. math:: k(x, x') = \mathrm{exp}\left[ -\frac{(x - x')^2}{2 \ell^2} \right] + + The power spectral density is: + + .. math:: + + S(\boldsymbol\omega) = + (\sqrt(2 \pi)^D \prod_{i}^{D}\ell_i + \exp\left( -\frac{1}{2} \sum_{i}^{D}\ell_i^2 \omega_i^{2} \right) """ def full(self, X, Xs=None): X, Xs = self._slice(X, Xs) return at.exp(-0.5 * self.square_dist(X, Xs)) + + def psd(self, omega): + ls = at.ones(self.D) * self.ls + c = at.power(at.sqrt(2.0 * np.pi), self.D) * at.prod(ls) + return c * at.exp(-0.5 * at.dot(at.square(omega), at.square(ls))) class RatQuad(Stationary): @@ -487,12 +555,29 @@ class Matern52(Stationary): k(x, x') = \left(1 + \frac{\sqrt{5(x - x')^2}}{\ell} + \frac{5(x-x')^2}{3\ell^2}\right) \mathrm{exp}\left[ - \frac{\sqrt{5(x - x')^2}}{\ell} \right] + + The power spectral density is: + + .. math:: + + S(\boldsymbol\omega) = + \frac{2^D \pi^{\frac{D}{2}} \Gamma(\frac{D+5}{2}) 5^{5/2}} + {\frac{3}{4}\sqrt{\pi}} + \prod_{i=1}^{D}\ell_{i} + \left(5 + \sum_{i=1}^{D}\ell_{i}^2 \boldsymbol\omega_{i}^{2}\right)^{-\frac{D+5}{2}} """ def full(self, X, Xs=None): X, Xs = self._slice(X, Xs) r = self.euclidean_dist(X, Xs) return (1.0 + np.sqrt(5.0) * r + 5.0 / 3.0 * at.square(r)) * at.exp(-1.0 * np.sqrt(5.0) * r) + + def psd(self, omega): + ls = at.ones(self.D) * self.ls + D52 = (self.D + 5) / 2 + num = at.power(2, self.D) * at.power(np.pi, self.D / 2) * at.gamma(D52) * at.power(5, 5 / 2) + den = 0.75 * at.sqrt(np.pi) + return (num / den) * at.prod(ls) * at.power(5.0 + at.dot(at.square(omega), at.square(ls)), -1 * D52) class Matern32(Stationary): @@ -503,12 +588,29 @@ class Matern32(Stationary): k(x, x') = \left(1 + \frac{\sqrt{3(x - x')^2}}{\ell}\right) \mathrm{exp}\left[ - \frac{\sqrt{3(x - x')^2}}{\ell} \right] + + The power spectral density is: + + .. math:: + + S(\boldsymbol\omega) = + \frac{2^D \pi^{D/2} \Gamma\left(\frac{D+3}{2}\right) 3^{3/2}} + {\frac{1}{2}\sqrt{\pi}} + \prod_{i=1}^{D}\ell_{i} + \left(3 + \sum_{i=1}^{D}\ell_{i}^2 \boldsymbol\omega_{i}^{2}\right)^{-\frac{D+3}{2}} """ def full(self, X, Xs=None): X, Xs = self._slice(X, Xs) r = self.euclidean_dist(X, Xs) return (1.0 + np.sqrt(3.0) * r) * at.exp(-np.sqrt(3.0) * r) + + def psd(self, omega): + ls = at.ones(self.D) * self.ls + D32 = (self.D + 3) / 2 + num = at.power(2, self.D) * at.power(np.pi, self.D / 2) * at.gamma(D32) * at.power(3, 3 / 2) + den = 0.5 * at.sqrt(np.pi) + return (num / den) * at.prod(ls) * at.power(3.0 + at.dot(at.square(omega), at.square(ls)), -1 * D32) class Matern12(Stationary): From 7ddd11e31054468595d2a6e2d6db530850c420d9 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Sat, 6 Aug 2022 18:21:48 -0700 Subject: [PATCH 02/17] composite kernels for psd, start tests --- pymc/gp/cov.py | 89 +++++++++++++++++++++++++++++-------------- pymc/tests/test_gp.py | 83 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 142 insertions(+), 30 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index f1524d3829..bb1fb2a5b3 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -14,6 +14,7 @@ import warnings +from collections import Counter from functools import reduce from numbers import Number from operator import add, mul @@ -46,7 +47,7 @@ ] -def _maybe_squeeze(value): +def _verify_scalar(value): if ( isinstance(value, aesara.compile.SharedVariable) and value.get_value().squeeze().shape == () @@ -131,6 +132,14 @@ def _slice(self, X, Xs): return X, Xs def __add__(self, other): + if isinstance(other, Number): + # If it's a scalar, cast as Constant covariance. This + # allows validation for power spectral density calc. + try: + other = Constant(c=_verify_scalar(other)) + except ValueError: + pass + return Add([self, other]) def __mul__(self, other): @@ -143,7 +152,7 @@ def __rmul__(self, other): return self.__mul__(other) def __pow__(self, other): - return Exponentiated(self, self._maybe_squeeze(other)) + return Exponentiated(self, _verify_scalar(other)) def __array_wrap__(self, result): """ @@ -160,19 +169,21 @@ def __array_wrap__(self, result): A = np.zeros((r, c)) for i in range(r): for j in range(c): - A[i, j] = result[i, j].factor_list[1] + r = result[i, j].factor_list[1] + if isinstance(r, Constant): + # Counteract the elemwise Add edgecase + r = r.c + A[i, j] = r if isinstance(result[0][0], Add): return result[0][0].factor_list[0] + A elif isinstance(result[0][0], Prod): return result[0][0].factor_list[0] * A else: raise TypeError( - f"Unknown Covariance combination type {result[0][0]}. Known types are `Add` or `Prod`." + f"Unknown Covariance combination type {result[0][0]}. " + "Known types are `Add` or `Prod`." ) - def psd(self): - raise NotImplementedError("Only covariance functions of type `Stationary` may implement a `.psd`.") - class Combination(Covariance): def __init__(self, factor_list): @@ -222,29 +233,34 @@ def merge_factors_psd(self, omega): kernels when possible. Implements a more restricted set of rules than `merge_factors_cov` -- just additivity of stationary covariances with defined power spectral densities - and mutliplication by scalars. + and multiplication by scalars. """ factor_list = [] - for i, factor in enumerate(self.factor_list): + for factor in self.factor_list: if isinstance(factor, Covariance): - # If it's a covariance, calculate the psd - factor_list.append(factor.psd(omega)) - elif isinstance( - factor, - ( - TensorConstant, - TensorVariable, - TensorSharedVariable, - ), - ): - # If it's an aesara variable, make sure it's a scalar - ## TODO this check doesnt do the right thing, arrays pass through - factor = _maybe_squeeze(factor) - factor_list.append(factor) - + # If it's a covariance try to calculate the psd + try: + factor_list.append(factor.psd(omega)) + + except (AttributeError, NotImplementedError) as e: + + if isinstance(factor, Stationary): + raise NotImplementedError( + "No power spectral density method has been " + f"implemented for {factor}." + ) from e + + else: + raise ValueError( + "Power specral densities, `.psd(omega)`, can only " + "be calculated for `Stationary` covariance " + f"functions. {factor} is non-stationary." + ) from e + else: - # Otherwise try to keep going + # Otherwise defer the reduction to later factor_list.append(factor) + return factor_list @@ -261,6 +277,18 @@ def __call__(self, X, Xs=None, diag=False): return reduce(mul, self.merge_factors_cov(X, Xs, diag)) def psd(self, omega): + check = Counter( + [ + isinstance(factor, Covariance) + for factor in self.factor_list + ] + ) + if check.get(True) >= 2: + raise NotImplementedError( + "The power spectral density of products of covariance " + "functions is not implemented." + ) + return reduce(mul, self.merge_factors_psd(omega)) @@ -521,8 +549,9 @@ def full(self, X, Xs=None): def psd(self, omega): ls = at.ones(self.D) * self.ls - c = at.power(at.sqrt(2.0 * np.pi), self.D) * at.prod(ls) - return c * at.exp(-0.5 * at.dot(at.square(omega), at.square(ls))) + c = at.power(at.sqrt(2.0 * np.pi), self.D) + exp = at.exp(-0.5 * at.dot(at.square(omega), at.square(ls))) + return c * at.prod(ls) * exp class RatQuad(Stationary): @@ -577,7 +606,8 @@ def psd(self, omega): D52 = (self.D + 5) / 2 num = at.power(2, self.D) * at.power(np.pi, self.D / 2) * at.gamma(D52) * at.power(5, 5 / 2) den = 0.75 * at.sqrt(np.pi) - return (num / den) * at.prod(ls) * at.power(5.0 + at.dot(at.square(omega), at.square(ls)), -1 * D52) + pow = at.power(5.0 + at.dot(at.square(omega), at.square(ls)), -1 * D52) + return (num / den) * at.prod(ls) * pow class Matern32(Stationary): @@ -610,7 +640,8 @@ def psd(self, omega): D32 = (self.D + 3) / 2 num = at.power(2, self.D) * at.power(np.pi, self.D / 2) * at.gamma(D32) * at.power(3, 3 / 2) den = 0.5 * at.sqrt(np.pi) - return (num / den) * at.prod(ls) * at.power(3.0 + at.dot(at.square(omega), at.square(ls)), -1 * D32) + pow = at.power(3.0 + at.dot(at.square(omega), at.square(ls)), -1 * D32) + return (num / den) * at.prod(ls) * pow class Matern12(Stationary): diff --git a/pymc/tests/test_gp.py b/pymc/tests/test_gp.py index 6e398bd004..9b4fb82778 100644 --- a/pymc/tests/test_gp.py +++ b/pymc/tests/test_gp.py @@ -259,6 +259,62 @@ def test_inv_rightprod(self): cov = M + pm.gp.cov.ExpQuad(1, 1.0) +class TestCovPSD: + def test_covpsd_add(self): + L = 10.0 + omega = np.pi * np.arange(1, 101) / (2 * L) + with pm.Model() as model: + cov1 = 2 * pm.gp.cov.ExpQuad(1, 0.1) + cov2 = 5 * pm.gp.cov.ExpQuad(1, 1.0) + cov = cov1 + cov2 + psd1 = cov1.psd(omega[:, None]).eval() + psd2 = cov2.psd(omega[:, None]).eval() + psd = cov.psd(omega[:, None]).eval() + npt.assert_allclose(psd, psd1 + psd2) + + def test_copsd_multiply(self): + # This could be implemented via convolution + L = 10.0 + omega = np.pi * np.arange(1, 101) / (2 * L) + with pm.Model() as model: + cov1 = 2 * pm.gp.cov.ExpQuad(1, ls=1) + cov2 = pm.gp.cov.ExpQuad(1, ls=1) + + with pytest.raises(NotImplementedError): + psd = (cov1 * cov2).psd(omega[:, None]).eval() + + def test_covpsd_nonstationary1(self): + L = 10.0 + omega = np.pi * np.arange(1, 101) / (2 * L) + with pm.Model() as model: + cov = 2 * pm.gp.cov.Linear(1, c=5) + + with pytest.raises(ValueError): + psd = cov.psd(omega[:, None]).eval() + + def test_covpsd_nonstationary2(self): + L = 10.0 + omega = np.pi * np.arange(1, 101) / (2 * L) + with pm.Model() as model: + cov = 2 * pm.gp.cov.ExpQuad(1, ls=1) + 10.0 + + with pytest.raises(ValueError): + psd = cov.psd(omega[:, None]).eval() + + def test_covpsd_notimplemented(self): + + class NewStationaryCov(pm.gp.cov.Stationary): + pass + + L = 10.0 + omega = np.pi * np.arange(1, 101) / (2 * L) + with pm.Model() as model: + cov = 2 * NewStationaryCov(1, ls=1) + + with pytest.raises(NotImplementedError): + psd = cov.psd(omega[:, None]).eval() + + class TestCovExponentiation: def test_symexp_cov(self): X = np.linspace(0, 1, 10)[:, None] @@ -306,7 +362,7 @@ def test_covexp_shared(self): def test_invalid_covexp(self): X = np.linspace(0, 1, 10)[:, None] - with pytest.raises(ValueError, match=r"can only be exponentiated by a scalar value"): + with pytest.raises(ValueError, match=r"scalar value is required"): with pm.Model() as model: a = np.array([[1.0, 2.0]]) cov = pm.gp.cov.ExpQuad(1, 0.1) ** a @@ -406,6 +462,22 @@ def test_stable(self): assert not np.any(dists < 0) +def test_psd_1d_matches_2d(cov_func): + L = 10.0 + M = 100 + cov_1d = cov_func(1, ls=0.1) + omega_1d = np.pi * np.arange(1, M + 1) / (2 * L) + psd_1d = cov_1d.psd(omega_1d[:, None]).eval() + + cov_2d = cov_func(2, ls=0.1) + S = np.meshgrid(*[np.arange(1, 1 + M) for _ in range(cov_2d.D)]) + S = np.vstack([s.flatten() for s in S]).T + omega_2d = np.pi * S / (2 * L) + psd_2d = cov_2d.psd(omega_2d).eval().reshape(M, M) + + npt.assert_allclose(psd_1d, np.sqrt(np.diag(psd_2d)), atol=1e-3) + + class TestExpQuad: def test_1d(self): X = np.linspace(0, 1, 10)[:, None] @@ -418,6 +490,9 @@ def test_1d(self): # check diagonal Kd = cov(X, diag=True).eval() npt.assert_allclose(np.diag(K), Kd, atol=1e-5) + + def test_psd_1d_matches_2d(self): + test_psd_1d_matches_2d(pm.gp.cov.ExpQuad) def test_2d(self): X = np.linspace(0, 1, 10).reshape(5, 2) @@ -526,6 +601,9 @@ def test_1d(self): # check diagonal Kd = cov(X, diag=True).eval() npt.assert_allclose(np.diag(K), Kd, atol=1e-5) + + def test_psd_1d_matches_2d(self): + test_psd_1d_matches_2d(pm.gp.cov.Matern52) class TestMatern32: @@ -540,6 +618,9 @@ def test_1d(self): # check diagonal Kd = cov(X, diag=True).eval() npt.assert_allclose(np.diag(K), Kd, atol=1e-5) + + def test_psd_1d_matches_2d(self): + test_psd_1d_matches_2d(pm.gp.cov.Matern32) class TestMatern12: From 667109699bc8d47e69f6504f21307721a6be7600 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Sat, 6 Aug 2022 18:51:07 -0700 Subject: [PATCH 03/17] remove test_matches_1d_2d for now --- pymc/tests/test_gp.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/pymc/tests/test_gp.py b/pymc/tests/test_gp.py index 9b4fb82778..54282faffc 100644 --- a/pymc/tests/test_gp.py +++ b/pymc/tests/test_gp.py @@ -460,22 +460,6 @@ def test_stable(self): cov = pm.gp.cov.ExpQuad(2, 0.1) dists = cov.square_dist(X, X).eval() assert not np.any(dists < 0) - - -def test_psd_1d_matches_2d(cov_func): - L = 10.0 - M = 100 - cov_1d = cov_func(1, ls=0.1) - omega_1d = np.pi * np.arange(1, M + 1) / (2 * L) - psd_1d = cov_1d.psd(omega_1d[:, None]).eval() - - cov_2d = cov_func(2, ls=0.1) - S = np.meshgrid(*[np.arange(1, 1 + M) for _ in range(cov_2d.D)]) - S = np.vstack([s.flatten() for s in S]).T - omega_2d = np.pi * S / (2 * L) - psd_2d = cov_2d.psd(omega_2d).eval().reshape(M, M) - - npt.assert_allclose(psd_1d, np.sqrt(np.diag(psd_2d)), atol=1e-3) class TestExpQuad: @@ -491,8 +475,6 @@ def test_1d(self): Kd = cov(X, diag=True).eval() npt.assert_allclose(np.diag(K), Kd, atol=1e-5) - def test_psd_1d_matches_2d(self): - test_psd_1d_matches_2d(pm.gp.cov.ExpQuad) def test_2d(self): X = np.linspace(0, 1, 10).reshape(5, 2) @@ -601,9 +583,6 @@ def test_1d(self): # check diagonal Kd = cov(X, diag=True).eval() npt.assert_allclose(np.diag(K), Kd, atol=1e-5) - - def test_psd_1d_matches_2d(self): - test_psd_1d_matches_2d(pm.gp.cov.Matern52) class TestMatern32: @@ -618,9 +597,6 @@ def test_1d(self): # check diagonal Kd = cov(X, diag=True).eval() npt.assert_allclose(np.diag(K), Kd, atol=1e-5) - - def test_psd_1d_matches_2d(self): - test_psd_1d_matches_2d(pm.gp.cov.Matern32) class TestMatern12: From ee44cbb7fede7e6a3df04d0d4b9268515f838d9b Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Sat, 6 Aug 2022 18:54:04 -0700 Subject: [PATCH 04/17] mark merge_factors with leading underscore --- pymc/gp/cov.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index bb1fb2a5b3..200393f0e6 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -198,7 +198,7 @@ def __init__(self, factor_list): else: self.factor_list.append(factor) - def merge_factors_cov(self, X, Xs=None, diag=False): + def _merge_factors_cov(self, X, Xs=None, diag=False): """ Called to evaluate either all the sums or all the products of kernels that are possible to evaluate. """ @@ -228,10 +228,10 @@ def merge_factors_cov(self, X, Xs=None, diag=False): factor_list.append(factor) return factor_list - def merge_factors_psd(self, omega): + def _merge_factors_psd(self, omega): """ Called to evaluatate spectral densities of combination kernels when possible. Implements a more restricted set - of rules than `merge_factors_cov` -- just additivity of + of rules than `_merge_factors_cov` -- just additivity of stationary covariances with defined power spectral densities and multiplication by scalars. """ @@ -266,15 +266,15 @@ def merge_factors_psd(self, omega): class Add(Combination): def __call__(self, X, Xs=None, diag=False): - return reduce(add, self.merge_factors_cov(X, Xs, diag)) + return reduce(add, self._merge_factors_cov(X, Xs, diag)) def psd(self, omega): - return reduce(add, self.merge_factors_psd(omega)) + return reduce(add, self._merge_factors_psd(omega)) class Prod(Combination): def __call__(self, X, Xs=None, diag=False): - return reduce(mul, self.merge_factors_cov(X, Xs, diag)) + return reduce(mul, self._merge_factors_cov(X, Xs, diag)) def psd(self, omega): check = Counter( @@ -289,7 +289,7 @@ def psd(self, omega): "functions is not implemented." ) - return reduce(mul, self.merge_factors_psd(omega)) + return reduce(mul, self._merge_factors_psd(omega)) class Exponentiated(Covariance): From a2ebd05aa0a91e176aa2f06ce418f52090f772d8 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Sat, 6 Aug 2022 19:09:52 -0700 Subject: [PATCH 05/17] hopefully make precommit happy --- pymc/gp/cov.py | 76 +++++++++++++++++++------------------------ pymc/tests/test_gp.py | 28 ++++++++-------- 2 files changed, 47 insertions(+), 57 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 200393f0e6..e77ed23aaf 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -48,10 +48,7 @@ def _verify_scalar(value): - if ( - isinstance(value, aesara.compile.SharedVariable) - and value.get_value().squeeze().shape == () - ): + if isinstance(value, aesara.compile.SharedVariable) and value.get_value().squeeze().shape == (): return at.squeeze(value) elif np.asarray(value).squeeze().shape == (): return np.squeeze(value) @@ -84,12 +81,12 @@ def __init__(self, input_dim, active_dims=None): @property def D(self): - """ The dimensionality of the input, as taken from the - `active_dims`. + """The dimensionality of the input, as taken from the + `active_dims`. """ # Evaluate lazily in-case this changes. return len(self.active_dims) - + def __call__(self, X, Xs=None, diag=False): r""" Evaluate the kernel/covariance function. @@ -139,7 +136,7 @@ def __add__(self, other): other = Constant(c=_verify_scalar(other)) except ValueError: pass - + return Add([self, other]) def __mul__(self, other): @@ -183,7 +180,7 @@ def __array_wrap__(self, result): f"Unknown Covariance combination type {result[0][0]}. " "Known types are `Add` or `Prod`." ) - + class Combination(Covariance): def __init__(self, factor_list): @@ -199,7 +196,7 @@ def __init__(self, factor_list): self.factor_list.append(factor) def _merge_factors_cov(self, X, Xs=None, diag=False): - """ Called to evaluate either all the sums or all the + """Called to evaluate either all the sums or all the products of kernels that are possible to evaluate. """ factor_list = [] @@ -229,7 +226,7 @@ def _merge_factors_cov(self, X, Xs=None, diag=False): return factor_list def _merge_factors_psd(self, omega): - """ Called to evaluatate spectral densities of combination + """Called to evaluatate spectral densities of combination kernels when possible. Implements a more restricted set of rules than `_merge_factors_cov` -- just additivity of stationary covariances with defined power spectral densities @@ -241,28 +238,28 @@ def _merge_factors_psd(self, omega): # If it's a covariance try to calculate the psd try: factor_list.append(factor.psd(omega)) - + except (AttributeError, NotImplementedError) as e: - + if isinstance(factor, Stationary): raise NotImplementedError( "No power spectral density method has been " f"implemented for {factor}." ) from e - + else: raise ValueError( "Power specral densities, `.psd(omega)`, can only " "be calculated for `Stationary` covariance " f"functions. {factor} is non-stationary." ) from e - + else: # Otherwise defer the reduction to later factor_list.append(factor) - + return factor_list - + class Add(Combination): def __call__(self, X, Xs=None, diag=False): @@ -270,25 +267,20 @@ def __call__(self, X, Xs=None, diag=False): def psd(self, omega): return reduce(add, self._merge_factors_psd(omega)) - + class Prod(Combination): def __call__(self, X, Xs=None, diag=False): return reduce(mul, self._merge_factors_cov(X, Xs, diag)) def psd(self, omega): - check = Counter( - [ - isinstance(factor, Covariance) - for factor in self.factor_list - ] - ) + check = Counter([isinstance(factor, Covariance) for factor in self.factor_list]) if check.get(True) >= 2: raise NotImplementedError( "The power spectral density of products of covariance " "functions is not implemented." ) - + return reduce(mul, self._merge_factors_psd(omega)) @@ -486,7 +478,7 @@ def diag(self, X): def full(self, X, Xs=None): raise NotImplementedError - + def psd(self, omega): raise NotImplementedError @@ -533,20 +525,20 @@ class ExpQuad(Stationary): .. math:: k(x, x') = \mathrm{exp}\left[ -\frac{(x - x')^2}{2 \ell^2} \right] - + The power spectral density is: - + .. math:: - - S(\boldsymbol\omega) = - (\sqrt(2 \pi)^D \prod_{i}^{D}\ell_i + + S(\boldsymbol\omega) = + (\sqrt(2 \pi)^D \prod_{i}^{D}\ell_i \exp\left( -\frac{1}{2} \sum_{i}^{D}\ell_i^2 \omega_i^{2} \right) """ def full(self, X, Xs=None): X, Xs = self._slice(X, Xs) return at.exp(-0.5 * self.square_dist(X, Xs)) - + def psd(self, omega): ls = at.ones(self.D) * self.ls c = at.power(at.sqrt(2.0 * np.pi), self.D) @@ -584,12 +576,12 @@ class Matern52(Stationary): k(x, x') = \left(1 + \frac{\sqrt{5(x - x')^2}}{\ell} + \frac{5(x-x')^2}{3\ell^2}\right) \mathrm{exp}\left[ - \frac{\sqrt{5(x - x')^2}}{\ell} \right] - + The power spectral density is: - + .. math:: - - S(\boldsymbol\omega) = + + S(\boldsymbol\omega) = \frac{2^D \pi^{\frac{D}{2}} \Gamma(\frac{D+5}{2}) 5^{5/2}} {\frac{3}{4}\sqrt{\pi}} \prod_{i=1}^{D}\ell_{i} @@ -600,7 +592,7 @@ def full(self, X, Xs=None): X, Xs = self._slice(X, Xs) r = self.euclidean_dist(X, Xs) return (1.0 + np.sqrt(5.0) * r + 5.0 / 3.0 * at.square(r)) * at.exp(-1.0 * np.sqrt(5.0) * r) - + def psd(self, omega): ls = at.ones(self.D) * self.ls D52 = (self.D + 5) / 2 @@ -618,12 +610,12 @@ class Matern32(Stationary): k(x, x') = \left(1 + \frac{\sqrt{3(x - x')^2}}{\ell}\right) \mathrm{exp}\left[ - \frac{\sqrt{3(x - x')^2}}{\ell} \right] - + The power spectral density is: - + .. math:: - - S(\boldsymbol\omega) = + + S(\boldsymbol\omega) = \frac{2^D \pi^{D/2} \Gamma\left(\frac{D+3}{2}\right) 3^{3/2}} {\frac{1}{2}\sqrt{\pi}} \prod_{i=1}^{D}\ell_{i} @@ -634,7 +626,7 @@ def full(self, X, Xs=None): X, Xs = self._slice(X, Xs) r = self.euclidean_dist(X, Xs) return (1.0 + np.sqrt(3.0) * r) * at.exp(-np.sqrt(3.0) * r) - + def psd(self, omega): ls = at.ones(self.D) * self.ls D32 = (self.D + 3) / 2 diff --git a/pymc/tests/test_gp.py b/pymc/tests/test_gp.py index 54282faffc..8bd4086de8 100644 --- a/pymc/tests/test_gp.py +++ b/pymc/tests/test_gp.py @@ -266,12 +266,12 @@ def test_covpsd_add(self): with pm.Model() as model: cov1 = 2 * pm.gp.cov.ExpQuad(1, 0.1) cov2 = 5 * pm.gp.cov.ExpQuad(1, 1.0) - cov = cov1 + cov2 + cov = cov1 + cov2 psd1 = cov1.psd(omega[:, None]).eval() psd2 = cov2.psd(omega[:, None]).eval() psd = cov.psd(omega[:, None]).eval() npt.assert_allclose(psd, psd1 + psd2) - + def test_copsd_multiply(self): # This could be implemented via convolution L = 10.0 @@ -279,41 +279,40 @@ def test_copsd_multiply(self): with pm.Model() as model: cov1 = 2 * pm.gp.cov.ExpQuad(1, ls=1) cov2 = pm.gp.cov.ExpQuad(1, ls=1) - + with pytest.raises(NotImplementedError): psd = (cov1 * cov2).psd(omega[:, None]).eval() - + def test_covpsd_nonstationary1(self): L = 10.0 omega = np.pi * np.arange(1, 101) / (2 * L) with pm.Model() as model: cov = 2 * pm.gp.cov.Linear(1, c=5) - + with pytest.raises(ValueError): psd = cov.psd(omega[:, None]).eval() - + def test_covpsd_nonstationary2(self): L = 10.0 omega = np.pi * np.arange(1, 101) / (2 * L) with pm.Model() as model: cov = 2 * pm.gp.cov.ExpQuad(1, ls=1) + 10.0 - + with pytest.raises(ValueError): psd = cov.psd(omega[:, None]).eval() - + def test_covpsd_notimplemented(self): - class NewStationaryCov(pm.gp.cov.Stationary): pass - + L = 10.0 omega = np.pi * np.arange(1, 101) / (2 * L) with pm.Model() as model: cov = 2 * NewStationaryCov(1, ls=1) - + with pytest.raises(NotImplementedError): psd = cov.psd(omega[:, None]).eval() - + class TestCovExponentiation: def test_symexp_cov(self): @@ -460,8 +459,8 @@ def test_stable(self): cov = pm.gp.cov.ExpQuad(2, 0.1) dists = cov.square_dist(X, X).eval() assert not np.any(dists < 0) - - + + class TestExpQuad: def test_1d(self): X = np.linspace(0, 1, 10)[:, None] @@ -474,7 +473,6 @@ def test_1d(self): # check diagonal Kd = cov(X, diag=True).eval() npt.assert_allclose(np.diag(K), Kd, atol=1e-5) - def test_2d(self): X = np.linspace(0, 1, 10).reshape(5, 2) From f4595571382918bfad77221a58adc220bfe1ba0f Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Fri, 12 Aug 2022 13:07:14 -0700 Subject: [PATCH 06/17] make factor_list hidden, check correct input_dims, active_dims for psd combinations --- pymc/gp/cov.py | 93 +++++++++++++++++++++++++++++++++----------------- 1 file changed, 62 insertions(+), 31 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index e77ed23aaf..4394f26d17 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -78,6 +78,9 @@ def __init__(self, input_dim, active_dims=None): self.active_dims = np.arange(input_dim) else: self.active_dims = np.asarray(active_dims, int) + + if max(self.active_dims) > self.input_dim: + raise ValueError("Values in `active_dims` can't be larger than `input_dim`.") @property def D(self): @@ -108,10 +111,10 @@ def __call__(self, X, Xs=None, diag=False): def diag(self, X): raise NotImplementedError - def full(self, X, Xs): + def full(self, X, Xs=None): raise NotImplementedError - def _slice(self, X, Xs): + def _slice(self, X, Xs=None): xdims = X.shape[-1] if isinstance(xdims, Variable): xdims = xdims.eval() @@ -127,7 +130,7 @@ def _slice(self, X, Xs): if Xs is not None: Xs = at.as_tensor_variable(Xs[:, self.active_dims]) return X, Xs - + def __add__(self, other): if isinstance(other, Number): # If it's a scalar, cast as Constant covariance. This @@ -166,15 +169,15 @@ def __array_wrap__(self, result): A = np.zeros((r, c)) for i in range(r): for j in range(c): - r = result[i, j].factor_list[1] + r = result[i, j]._factor_list[1] if isinstance(r, Constant): # Counteract the elemwise Add edgecase r = r.c A[i, j] = r if isinstance(result[0][0], Add): - return result[0][0].factor_list[0] + A + return result[0][0]._factor_list[0] + A elif isinstance(result[0][0], Prod): - return result[0][0].factor_list[0] * A + return result[0][0]._factor_list[0] * A else: raise TypeError( f"Unknown Covariance combination type {result[0][0]}. " @@ -184,23 +187,46 @@ def __array_wrap__(self, result): class Combination(Covariance): def __init__(self, factor_list): - input_dim = max( - factor.input_dim for factor in factor_list if isinstance(factor, Covariance) + """Use constituent factors to get input_dim and active_dims for the Combination covariance. + """ + # Check if all input_dim are the same in factor_list + input_dims = set( + factor.input_dim + for factor in factor_list + if isinstance(factor, Covariance) ) - super().__init__(input_dim=input_dim) - self.factor_list = [] + if len(input_dims) != 1: + raise ValueError("All covariances must have the same `input_dim`.") + else: + input_dim = input_dims.pop() + + # Union all active_dims sets in factor_list for the combination covariance + active_dims = np.sort( + np.asarray(list(set.union( + *[ + set(factor.active_dims) + for factor in factor_list + if isinstance(factor, Covariance) + ] + )), dtype=int) + ) + + super().__init__(input_dim=input_dim, active_dims=active_dims) + + # Set up combination kernel, flatten out factor_list so that + self._factor_list = [] for factor in factor_list: if isinstance(factor, self.__class__): - self.factor_list.extend(factor.factor_list) + self._factor_list.extend(factor._factor_list) else: - self.factor_list.append(factor) + self._factor_list.append(factor) def _merge_factors_cov(self, X, Xs=None, diag=False): """Called to evaluate either all the sums or all the products of kernels that are possible to evaluate. """ factor_list = [] - for factor in self.factor_list: + for factor in self._factor_list: # make sure diag=True is handled properly if isinstance(factor, Covariance): factor_list.append(factor(X, Xs, diag)) @@ -226,15 +252,22 @@ def _merge_factors_cov(self, X, Xs=None, diag=False): return factor_list def _merge_factors_psd(self, omega): - """Called to evaluatate spectral densities of combination - kernels when possible. Implements a more restricted set - of rules than `_merge_factors_cov` -- just additivity of - stationary covariances with defined power spectral densities - and multiplication by scalars. + """Called to evaluatate spectral densities of combination kernels when possible. Implements + a more restricted set of rules than `_merge_factors_cov` -- just additivity of stationary + covariances with defined power spectral densities and multiplication by scalars. Also, the + active_dims for all covariances in the sum must be the same. """ factor_list = [] - for factor in self.factor_list: + for factor in self._factor_list: if isinstance(factor, Covariance): + + # Allow merging covariances for psd only if active_dims are the same + if set(self.active_dims) != set(factor.active_dims): + raise ValueError( + "For power spectral density calculations `active_dims` must be the same " + "for all covariances in the sum." + ) + # If it's a covariance try to calculate the psd try: factor_list.append(factor.psd(omega)) @@ -243,15 +276,13 @@ def _merge_factors_psd(self, omega): if isinstance(factor, Stationary): raise NotImplementedError( - "No power spectral density method has been " - f"implemented for {factor}." + f"No power spectral density method has been implemented for {factor}." ) from e else: raise ValueError( - "Power specral densities, `.psd(omega)`, can only " - "be calculated for `Stationary` covariance " - f"functions. {factor} is non-stationary." + "Power spectral densities, `.psd(omega)`, can only be calculated for " + f"`Stationary` covariance functions. {factor} is non-stationary." ) from e else: @@ -274,7 +305,7 @@ def __call__(self, X, Xs=None, diag=False): return reduce(mul, self._merge_factors_cov(X, Xs, diag)) def psd(self, omega): - check = Counter([isinstance(factor, Covariance) for factor in self.factor_list]) + check = Counter([isinstance(factor, Covariance) for factor in self._factor_list]) if check.get(True) >= 2: raise NotImplementedError( "The power spectral density of products of covariance " @@ -314,7 +345,7 @@ def __init__(self, factor_list): self.input_dims = [factor.input_dim for factor in factor_list] input_dim = sum(self.input_dims) super().__init__(input_dim=input_dim) - self.factor_list = factor_list + self._factor_list = factor_list def _split(self, X, Xs): indices = np.cumsum(self.input_dims) @@ -327,11 +358,12 @@ def _split(self, X, Xs): def __call__(self, X, Xs=None, diag=False): X_split, Xs_split = self._split(X, Xs) - covs = [cov(x, xs, diag) for cov, x, xs in zip(self.factor_list, X_split, Xs_split)] + covs = [cov(x, xs, diag) for cov, x, xs in zip(self._factor_list, X_split, Xs_split)] return reduce(mul, covs) -class Constant(Covariance): +class Constant: + # Don't subclass `Covariance` so input_dim and active_dims aren't required. r""" Constant valued covariance function. @@ -341,7 +373,6 @@ class Constant(Covariance): """ def __init__(self, c): - super().__init__(1, None) self.c = c def diag(self, X): @@ -363,8 +394,8 @@ class WhiteNoise(Covariance): k(x, x') = \sigma^2 \mathrm{I} """ - def __init__(self, sigma): - super().__init__(1, None) + def __init__(self, input_dim, sigma, active_dims=None): + super().__init__(input_dim, active_dims) self.sigma = sigma def diag(self, X): From 99223a0115fdc021e1ea3e7e7e97e0e2d9431877 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Fri, 12 Aug 2022 13:33:38 -0700 Subject: [PATCH 07/17] run precommit --- pymc/gp/cov.py | 48 ++++++++++++++++++++++++------------------------ 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 4394f26d17..4b6f2ab8ca 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -78,7 +78,7 @@ def __init__(self, input_dim, active_dims=None): self.active_dims = np.arange(input_dim) else: self.active_dims = np.asarray(active_dims, int) - + if max(self.active_dims) > self.input_dim: raise ValueError("Values in `active_dims` can't be larger than `input_dim`.") @@ -130,7 +130,7 @@ def _slice(self, X, Xs=None): if Xs is not None: Xs = at.as_tensor_variable(Xs[:, self.active_dims]) return X, Xs - + def __add__(self, other): if isinstance(other, Number): # If it's a scalar, cast as Constant covariance. This @@ -187,33 +187,33 @@ def __array_wrap__(self, result): class Combination(Covariance): def __init__(self, factor_list): - """Use constituent factors to get input_dim and active_dims for the Combination covariance. - """ + """Use constituent factors to get input_dim and active_dims for the Combination covariance.""" # Check if all input_dim are the same in factor_list - input_dims = set( - factor.input_dim - for factor in factor_list - if isinstance(factor, Covariance) - ) + input_dims = {factor.input_dim for factor in factor_list if isinstance(factor, Covariance)} if len(input_dims) != 1: raise ValueError("All covariances must have the same `input_dim`.") else: input_dim = input_dims.pop() - + # Union all active_dims sets in factor_list for the combination covariance active_dims = np.sort( - np.asarray(list(set.union( - *[ - set(factor.active_dims) - for factor in factor_list - if isinstance(factor, Covariance) - ] - )), dtype=int) + np.asarray( + list( + set.union( + *[ + set(factor.active_dims) + for factor in factor_list + if isinstance(factor, Covariance) + ] + ) + ), + dtype=int, + ) ) - + super().__init__(input_dim=input_dim, active_dims=active_dims) - - # Set up combination kernel, flatten out factor_list so that + + # Set up combination kernel, flatten out factor_list so that self._factor_list = [] for factor in factor_list: if isinstance(factor, self.__class__): @@ -253,21 +253,21 @@ def _merge_factors_cov(self, X, Xs=None, diag=False): def _merge_factors_psd(self, omega): """Called to evaluatate spectral densities of combination kernels when possible. Implements - a more restricted set of rules than `_merge_factors_cov` -- just additivity of stationary - covariances with defined power spectral densities and multiplication by scalars. Also, the + a more restricted set of rules than `_merge_factors_cov` -- just additivity of stationary + covariances with defined power spectral densities and multiplication by scalars. Also, the active_dims for all covariances in the sum must be the same. """ factor_list = [] for factor in self._factor_list: if isinstance(factor, Covariance): - + # Allow merging covariances for psd only if active_dims are the same if set(self.active_dims) != set(factor.active_dims): raise ValueError( "For power spectral density calculations `active_dims` must be the same " "for all covariances in the sum." ) - + # If it's a covariance try to calculate the psd try: factor_list.append(factor.psd(omega)) From 576bd708b7e2aa6e819a752365992d87308e2179 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Wed, 31 Aug 2022 11:22:47 -0700 Subject: [PATCH 08/17] split Covariance into Covariance and BaseCovariance, where BaseCovariance doesnt need to know input_dims and active_dims and Covariance does. This allows Constant and WhiteNoise to work without specified input_dim and active_dims --- pymc/gp/cov.py | 124 +++++++++++++++++++++++++++---------------------- 1 file changed, 68 insertions(+), 56 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index a41b97c552..ec1c6f3868 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -58,37 +58,10 @@ def _verify_scalar(value): raise ValueError("A scalar value is required.") -class Covariance: - r""" - Base class for all kernels/covariance functions. - - Parameters - ---------- - input_dim: integer - The number of input dimensions, or columns of X (or Xs) - the kernel will operate on. - active_dims: List of integers - Indicate which dimension or column of X the covariance - function operates on. +class BaseCovariance: + """ + Base class for kernels/covariance functions. """ - - def __init__(self, input_dim, active_dims=None): - self.input_dim = input_dim - if active_dims is None: - self.active_dims = np.arange(input_dim) - else: - self.active_dims = np.asarray(active_dims, int) - - if max(self.active_dims) > self.input_dim: - raise ValueError("Values in `active_dims` can't be larger than `input_dim`.") - - @property - def D(self): - """The dimensionality of the input, as taken from the - `active_dims`. - """ - # Evaluate lazily in-case this changes. - return len(self.active_dims) def __call__(self, X, Xs=None, diag=False): r""" @@ -114,27 +87,10 @@ def diag(self, X): def full(self, X, Xs=None): raise NotImplementedError - def _slice(self, X, Xs=None): - xdims = X.shape[-1] - if isinstance(xdims, Variable): - xdims = xdims.eval() - if self.input_dim != xdims: - warnings.warn( - f"Only {self.input_dim} column(s) out of {xdims} are" - " being used to compute the covariance function. If this" - " is not intended, increase 'input_dim' parameter to" - " the number of columns to use. Ignore otherwise.", - UserWarning, - ) - X = at.as_tensor_variable(X[:, self.active_dims]) - if Xs is not None: - Xs = at.as_tensor_variable(Xs[:, self.active_dims]) - return X, Xs - def __add__(self, other): + # If it's a scalar, cast as Constant covariance. This allows validation for power spectral + # density calc. if isinstance(other, Number): - # If it's a scalar, cast as Constant covariance. This - # allows validation for power spectral density calc. try: other = Constant(c=_verify_scalar(other)) except ValueError: @@ -185,15 +141,67 @@ def __array_wrap__(self, result): ) +class Covariance(BaseCovariance): + """ + Base class for kernels/covariance functions with input_dim and active_dims, which excludes + kernels like `Constant` and `WhiteNoise`. + + Parameters + ---------- + input_dim: integer + The number of input dimensions, or columns of X (or Xs) + the kernel will operate on. + active_dims: List of integers + Indicate which dimension or column of X the covariance + function operates on. + """ + + def __init__(self, input_dim, active_dims=None): + self.input_dim = input_dim + if active_dims is None: + self.active_dims = np.arange(input_dim) + else: + self.active_dims = np.asarray(active_dims, int) + + if max(self.active_dims) > self.input_dim: + raise ValueError("Values in `active_dims` can't be larger than `input_dim`.") + + @property + def D(self): + """The dimensionality of the input, as taken from the + `active_dims`. + """ + # Evaluate lazily in-case this changes. + return len(self.active_dims) + + def _slice(self, X, Xs=None): + xdims = X.shape[-1] + if isinstance(xdims, Variable): + xdims = xdims.eval() + if self.input_dim != xdims: + warnings.warn( + f"Only {self.input_dim} column(s) out of {xdims} are" + " being used to compute the covariance function. If this" + " is not intended, increase 'input_dim' parameter to" + " the number of columns to use. Ignore otherwise.", + UserWarning, + ) + X = at.as_tensor_variable(X[:, self.active_dims]) + if Xs is not None: + Xs = at.as_tensor_variable(Xs[:, self.active_dims]) + return X, Xs + + class Combination(Covariance): def __init__(self, factor_list): """Use constituent factors to get input_dim and active_dims for the Combination covariance.""" + # Check if all input_dim are the same in factor_list input_dims = {factor.input_dim for factor in factor_list if isinstance(factor, Covariance)} + if len(input_dims) != 1: raise ValueError("All covariances must have the same `input_dim`.") - else: - input_dim = input_dims.pop() + input_dim = input_dims.pop() # Union all active_dims sets in factor_list for the combination covariance active_dims = np.sort( @@ -227,14 +235,17 @@ def _merge_factors_cov(self, X, Xs=None, diag=False): """ factor_list = [] for factor in self._factor_list: + # make sure diag=True is handled properly - if isinstance(factor, Covariance): + if isinstance(factor, BaseCovariance): factor_list.append(factor(X, Xs, diag)) + elif isinstance(factor, np.ndarray): if np.ndim(factor) == 2 and diag: factor_list.append(np.diag(factor)) else: factor_list.append(factor) + elif isinstance( factor, ( @@ -243,12 +254,15 @@ def _merge_factors_cov(self, X, Xs=None, diag=False): TensorSharedVariable, ), ): + if factor.ndim == 2 and diag: factor_list.append(at.diag(factor)) else: factor_list.append(factor) + else: factor_list.append(factor) + return factor_list def _merge_factors_psd(self, omega): @@ -362,8 +376,7 @@ def __call__(self, X, Xs=None, diag=False): return reduce(mul, covs) -class Constant: - # Don't subclass `Covariance` so input_dim and active_dims aren't required. +class Constant(BaseCovariance): r""" Constant valued covariance function. @@ -394,8 +407,7 @@ class WhiteNoise(Covariance): k(x, x') = \sigma^2 \mathrm{I} """ - def __init__(self, input_dim, sigma, active_dims=None): - super().__init__(input_dim, active_dims) + def __init__(self, sigma): self.sigma = sigma def diag(self, X): From 1632c3a498bee0634003036fe485af74cc80e71a Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Wed, 31 Aug 2022 11:23:53 -0700 Subject: [PATCH 09/17] fix weird test with inconsitent input_dim --- pymc/tests/test_gp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/tests/test_gp.py b/pymc/tests/test_gp.py index 06325a8b78..1bee6d7952 100644 --- a/pymc/tests/test_gp.py +++ b/pymc/tests/test_gp.py @@ -395,7 +395,7 @@ def test_multiops(self): + pm.gp.cov.ExpQuad(1, 0.1) + pm.gp.cov.ExpQuad(1, 0.1) * pm.gp.cov.ExpQuad(1, 0.1) ) - cov2 = pm.gp.cov.ExpQuad(1, 0.1) * pm.gp.cov.ExpQuad(2, 0.1) + cov2 = pm.gp.cov.ExpQuad(2, 0.1) * pm.gp.cov.ExpQuad(2, 0.1) cov = pm.gp.cov.Kron([cov1, cov2]) K_true = kronecker(cov1(X1).eval(), cov2(X2).eval()).eval() K = cov(X).eval() From 40f1b62b39b78ec530f67e5feb76a63923a6b98b Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Wed, 31 Aug 2022 22:53:40 -0700 Subject: [PATCH 10/17] add tests for matern52 and matern32 --- pymc/gp/cov.py | 2 +- pymc/tests/test_gp.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index ec1c6f3868..fa6dfbe1fd 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -569,7 +569,7 @@ class ExpQuad(Stationary): k(x, x') = \mathrm{exp}\left[ -\frac{(x - x')^2}{2 \ell^2} \right] - The power spectral density is: + The power spectral density (`psd`) is: .. math:: diff --git a/pymc/tests/test_gp.py b/pymc/tests/test_gp.py index 1bee6d7952..fc31d54be9 100644 --- a/pymc/tests/test_gp.py +++ b/pymc/tests/test_gp.py @@ -506,6 +506,15 @@ def test_inv_lengthscale(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 + X = np.linspace(0, 1, 10)[:, None] + omega = np.linspace(0, 2, 50) + ell = 2.0 + true_1d_psd = np.sqrt(2 * np.pi * np.square(ell)) * np.exp(-0.5 * np.square(ell * omega)) + test_1d_psd = pm.gp.cov.ExpQuad(1, ls=ell).psd(omega[:, None]).flatten().eval() + npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) + class TestWhiteNoise: def test_1d(self): @@ -581,6 +590,16 @@ def test_1d(self): # check diagonal 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 + X = np.linspace(0, 1, 10)[:, None] + omega = np.linspace(0, 2, 50) + ell = 2.0 + lamda = np.sqrt(5) / ell + true_1d_psd = (16.0 / 3.0) * np.power(lamda, 5) * np.power(lamda**2 + omega**2, -3) + test_1d_psd = pm.gp.cov.Matern52(1, ls=ell).psd(omega[:, None]).flatten().eval() + npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) class TestMatern32: @@ -595,6 +614,16 @@ def test_1d(self): # check diagonal 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 + X = np.linspace(0, 1, 10)[:, None] + omega = np.linspace(0, 2, 50) + ell = 2.0 + lamda = np.sqrt(3) / ell + true_1d_psd = 4 * np.power(lamda, 3) * np.power(lamda**2 + omega**2, -2) + test_1d_psd = pm.gp.cov.Matern32(1, ls=ell).psd(omega[:, None]).flatten().eval() + npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) class TestMatern12: From 784ca32d58cfc57a8516154e8310ca515f27971f Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Wed, 31 Aug 2022 23:08:49 -0700 Subject: [PATCH 11/17] fix from precommit --- pymc/tests/test_gp.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pymc/tests/test_gp.py b/pymc/tests/test_gp.py index fc31d54be9..3f7de6a2f5 100644 --- a/pymc/tests/test_gp.py +++ b/pymc/tests/test_gp.py @@ -590,11 +590,11 @@ def test_1d(self): # check diagonal 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 X = np.linspace(0, 1, 10)[:, None] - omega = np.linspace(0, 2, 50) + omega = np.linspace(0, 2, 50) ell = 2.0 lamda = np.sqrt(5) / ell true_1d_psd = (16.0 / 3.0) * np.power(lamda, 5) * np.power(lamda**2 + omega**2, -3) @@ -614,11 +614,11 @@ def test_1d(self): # check diagonal 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 X = np.linspace(0, 1, 10)[:, None] - omega = np.linspace(0, 2, 50) + omega = np.linspace(0, 2, 50) ell = 2.0 lamda = np.sqrt(3) / ell true_1d_psd = 4 * np.power(lamda, 3) * np.power(lamda**2 + omega**2, -2) From 7df057bbaa9d7990a74839f3e0e01fad00a6d0e0 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Wed, 5 Oct 2022 21:46:57 -0700 Subject: [PATCH 12/17] add HSGP implementation from pymc-experimental --- pymc/gp/hsgp.py | 216 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 216 insertions(+) create mode 100644 pymc/gp/hsgp.py diff --git a/pymc/gp/hsgp.py b/pymc/gp/hsgp.py new file mode 100644 index 0000000000..623a9b3430 --- /dev/null +++ b/pymc/gp/hsgp.py @@ -0,0 +1,216 @@ +# Copyright 2022 The PyMC Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warnings + +from typing import Optional + +import aesara.tensor as at +import numpy as np + +import pymc as pm + +from pymc.gp.cov import Covariance +from pymc.gp.gp import Base +from pymc.gp.mean import Mean, Zero + + +class HSGP(Base): + R""" + Hilbert Space Gaussian process + + The `gp.HSGP` class is an implementation of the Hilbert Space Gaussian process. This + approximation is a linear model that uses a fixed set of basis vectors, whose coeficients are + random functions of a stationary covariance function's power spectral density. 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 additonally has an `approx_K` method which returns the approximate covariance + matrix. It supports a limited subset of additive covariances. + + For information on choosing appropriate `m`, `L`, and `c`, refer Ruitort-Mayol et. al. or to the + pymc examples documentation. + + Parameters + ---------- + m: list + The number of basis vectors to use for each active dimension (covariance parameter + `active_dim`). + 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`. + c: 1.5 + 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. + 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. + + Examples + -------- + .. code:: python + + # A three dimensional column vector of inputs. + X = np.random.randn(100, 3) + + with pm.Model() as model: + # Specify the covariance function. Three input dimensions, but we only want to use the + # last two. + cov_func = pm.gp.cov.ExpQuad(3, ls=0.1, active_dims=[1, 2]) + + # Specify the HSGP. Use 10 basis vectors across each active dimension, [1, 2] for a + # total of 10 * 10 = 100. The input X is normally distributed, so use a boundary + # condition that should easily contain all the points, from -6 to 6 in each dimension. + gp = pmx.gp.HSGP(n_basis=[10, 10], L=[6, 6], 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 + + - Solin, A., Sarkka, S. (2019) Hilbert Space Methods for Reduced-Rank Gaussian Process + Regression. + """ + + def __init__( + self, + m: int, + L: Optional = None, + c: float = 1.5, + *, + mean_func: Mean = Zero(), + cov_func: Optional[Covariance] = None, + ): + arg_err_msg = ( + "`m` and L, if provided, must be lists or tuples, with one element per active " + "dimension." + ) + try: + if len(m) != cov_func.D: + raise ValueError(arg_err_msg) + except TypeError as e: + raise ValueError(arg_err_msg) from e + + if L is not None and len(L) != cov_func.D: + raise ValueError(arg_err_msg) + + if L is None and c < 1.2: + warnings.warn( + "Most applications will require a `c >= 1.2` for accuracy at the boundaries of the " + "domain." + ) + + self.m = m + self.L = L + self.c = c + self.D = cov_func.D + + super().__init__(mean_func=mean_func, cov_func=cov_func) + + def __add__(self, other): + raise NotImplementedError("Additive HSGPs aren't supported ") + + def _set_boundary(self, X): + """Make L from X and c if L is not passed in.""" + if self.L is None: + # Define new L based on c and X range + La = at.abs(at.min(X, axis=0)) + Lb = at.abs(at.max(X, axis=0)) + self.L = self.c * at.max(at.stack((La, Lb)), axis=0) + else: + self.L = at.as_tensor_variable(self.L) + + @staticmethod + def _eigendecomposition(X, L, m, D): + """Construct the eigenvalues and eigenfunctions of the Laplace operator.""" + m_star = at.prod(m) + S = np.meshgrid(*[np.arange(1, 1 + m[d]) for d in range(D)]) + S = np.vstack([s.flatten() for s in S]).T + eigvals = at.square((np.pi * S) / (2 * L)) + phi = at.ones((X.shape[0], m_star)) + for d in range(D): + c = 1.0 / np.sqrt(L[d]) + phi *= c * at.sin(at.sqrt(eigvals[:, d]) * (at.tile(X[:, d][:, None], m_star) + L[d])) + omega = at.sqrt(eigvals) + return omega, phi, m_star + + def approx_K(self, X, L, m): + """A helper function which gives the approximate kernel or covariance matrix K. This can be + helpful when trying to see how well an approximation may work. + """ + X, _ = self.cov_func._slice(X) + omega, phi, _ = self._eigendecomposition(X, self.L, self.m, self.cov_func.D) + psd = self.cov_func.psd(omega) + return at.dot(phi * psd, at.transpose(phi)) + + def prior(self, name, X, dims=None): + R""" + Returns the (approximate) GP prior distribution evaluated over the input locations `X`. + + Parameters + ---------- + name: string + Name of the random variable + X: array-like + Function input values. + dims: None + Dimension name for the GP random variable. + """ + + X, _ = self.cov_func._slice(X) + self._set_boundary(X) + omega, phi, m_star = self._eigendecomposition(X, self.L, self.m, self.D) + psd = self.cov_func.psd(omega) + self.beta = pm.Normal(f"{name}_coeffs_", size=m_star) + self.f = pm.Deterministic( + name, self.mean_func(X) + at.squeeze(at.dot(phi, self.beta * psd)), dims + ) + return self.f + + def _build_conditional(self, name, Xnew): + Xnew, _ = self.cov_func._slice(Xnew) + omega, phi, _ = self._eigendecomposition(Xnew, self.L, self.m, self.D) + psd = self.cov_func.psd(omega) + return self.mean_func(Xnew) + at.squeeze(at.dot(phi, self.beta * psd)) + + def conditional(self, name, Xnew, dims=None): + R""" + Returns the (approximate) conditional distribution evaluated over new input locations + `Xnew`. + + Parameters + ---------- + name: string + Name of the random variable + Xnew: array-like + Function input values. + dims: None + Dimension name for the GP random variable. + """ + fnew = self._build_conditional(name, Xnew) + return pm.Deterministic(name, fnew, dims) From fcb96827491a4b087ccc59dbe67cb10e7aecb58f Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Mon, 10 Oct 2022 20:09:30 -0700 Subject: [PATCH 13/17] make hsgp importable --- pymc/gp/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pymc/gp/__init__.py b/pymc/gp/__init__.py index 9510f33de4..995a5f9b87 100644 --- a/pymc/gp/__init__.py +++ b/pymc/gp/__init__.py @@ -22,3 +22,4 @@ MarginalKron, MarginalSparse, ) +from pymc.gp.hsgp import HSGP From 6be8fdb73f0ef24a7385ec53040284cbe7170c5e Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Wed, 12 Oct 2022 16:39:58 -0700 Subject: [PATCH 14/17] make dims explicit keyword in Deterministic, fixes error --- pymc/gp/hsgp.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pymc/gp/hsgp.py b/pymc/gp/hsgp.py index 623a9b3430..cc665fc1e9 100644 --- a/pymc/gp/hsgp.py +++ b/pymc/gp/hsgp.py @@ -188,7 +188,7 @@ def prior(self, name, X, dims=None): psd = self.cov_func.psd(omega) self.beta = pm.Normal(f"{name}_coeffs_", size=m_star) self.f = pm.Deterministic( - name, self.mean_func(X) + at.squeeze(at.dot(phi, self.beta * psd)), dims + name, self.mean_func(X) + at.squeeze(at.dot(phi, self.beta * psd)), dims=dims ) return self.f @@ -213,4 +213,4 @@ def conditional(self, name, Xnew, dims=None): Dimension name for the GP random variable. """ fnew = self._build_conditional(name, Xnew) - return pm.Deterministic(name, fnew, dims) + return pm.Deterministic(name, fnew, dims=dims) From 17c4f12be31f72f3114213ab132f2ac82869e3ca Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Sat, 29 Oct 2022 14:50:57 -0700 Subject: [PATCH 15/17] fix docstring --- pymc/gp/hsgp.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pymc/gp/hsgp.py b/pymc/gp/hsgp.py index cc665fc1e9..63e8de8889 100644 --- a/pymc/gp/hsgp.py +++ b/pymc/gp/hsgp.py @@ -71,10 +71,10 @@ class HSGP(Base): # last two. cov_func = pm.gp.cov.ExpQuad(3, ls=0.1, active_dims=[1, 2]) - # Specify the HSGP. Use 10 basis vectors across each active dimension, [1, 2] for a - # total of 10 * 10 = 100. The input X is normally distributed, so use a boundary - # condition that should easily contain all the points, from -6 to 6 in each dimension. - gp = pmx.gp.HSGP(n_basis=[10, 10], L=[6, 6], cov_func=cov_func) + # Specify the HSGP. Use 50 basis vectors across each active dimension, [1, 2] for a + # total of 50 * 50 = 2500. The range of the data is inferred from X, and the boundary + # condition multiplier `c` uses 4 * half range. + gp = pm.gp.HSGP(m=[50, 50], c=4.0, cov_func=cov_func) # Place a GP prior over the function f. f = gp.prior("f", X=X) @@ -108,7 +108,7 @@ def __init__( ): arg_err_msg = ( "`m` and L, if provided, must be lists or tuples, with one element per active " - "dimension." + "dimension of the kernel or covariance function." ) try: if len(m) != cov_func.D: From 776ea464c14c95cd8a69009d7aa70ecb4c4cf657 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Fri, 4 Nov 2022 09:10:24 -0700 Subject: [PATCH 16/17] add drop_first and add docstring --- pymc/gp/hsgp.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/pymc/gp/hsgp.py b/pymc/gp/hsgp.py index 63e8de8889..52afbd07d9 100644 --- a/pymc/gp/hsgp.py +++ b/pymc/gp/hsgp.py @@ -54,6 +54,10 @@ class HSGP(Base): 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. + 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. mean_func: None, instance of Mean @@ -102,6 +106,7 @@ def __init__( m: int, L: Optional = None, c: float = 1.5, + drop_first=False, *, mean_func: Mean = Zero(), cov_func: Optional[Covariance] = None, @@ -125,6 +130,7 @@ def __init__( "domain." ) + self.drop_first = drop_first self.m = m self.L = L self.c = c @@ -186,10 +192,18 @@ def prior(self, name, X, dims=None): self._set_boundary(X) omega, phi, m_star = self._eigendecomposition(X, self.L, self.m, self.D) psd = self.cov_func.psd(omega) - self.beta = pm.Normal(f"{name}_coeffs_", size=m_star) - self.f = pm.Deterministic( - name, self.mean_func(X) + at.squeeze(at.dot(phi, self.beta * psd)), dims=dims - ) + + if self.drop_first: + print('dropping first') + self.beta = pm.Normal(f"{name}_coeffs_", size=m_star - 1) + self.f = pm.Deterministic( + name, self.mean_func(X) + at.squeeze(at.dot(phi[:, 1:], self.beta * psd[1:])), dims=dims + ) + else: + self.beta = pm.Normal(f"{name}_coeffs_", size=m_star) + self.f = pm.Deterministic( + name, self.mean_func(X) + at.squeeze(at.dot(phi, self.beta * psd)), dims=dims + ) return self.f def _build_conditional(self, name, Xnew): From c87520d2e2a404b90245a223bfb2158ee6468753 Mon Sep 17 00:00:00 2001 From: Bill Engels Date: Wed, 9 Nov 2022 11:51:14 -0800 Subject: [PATCH 17/17] fix missing sqrt!!! --- pymc/gp/hsgp.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pymc/gp/hsgp.py b/pymc/gp/hsgp.py index 52afbd07d9..ba7146f781 100644 --- a/pymc/gp/hsgp.py +++ b/pymc/gp/hsgp.py @@ -194,15 +194,14 @@ def prior(self, name, X, dims=None): psd = self.cov_func.psd(omega) if self.drop_first: - print('dropping first') self.beta = pm.Normal(f"{name}_coeffs_", size=m_star - 1) self.f = pm.Deterministic( - name, self.mean_func(X) + at.squeeze(at.dot(phi[:, 1:], self.beta * psd[1:])), dims=dims + name, self.mean_func(X) + at.squeeze(at.dot(phi[:, 1:], self.beta * at.sqrt(psd[1:]))), dims=dims ) else: self.beta = pm.Normal(f"{name}_coeffs_", size=m_star) self.f = pm.Deterministic( - name, self.mean_func(X) + at.squeeze(at.dot(phi, self.beta * psd)), dims=dims + name, self.mean_func(X) + at.squeeze(at.dot(phi, self.beta * at.sqrt(psd))), dims=dims ) return self.f