ENH: add exponentiation of a covariance function with a scalar (pymc-devs#3852)

* ENH: add exponentiation with a scalar

* fix the scalar condition

* add examples in notebook

Author: tirthasheshpatel

RELEASE-NOTES.md

@@ -30,6 +30,7 @@
 - Distinguish between `Data` and `Deterministic` variables when graphing models with graphviz. PR [#3491](https://github.com/pymc-devs/pymc3/pull/3491).
 - Sequential Monte Carlo - Approximate Bayesian Computation step method is now available. The implementation is in an experimental stage and will be further improved.
 - Added `Matern12` covariance function for Gaussian processes. This is the Matern kernel with nu=1/2.
+- Added exponentiation of a covariance function with a scalar. See PR[#3852](https://github.com/pymc-devs/pymc3/pull/3852)
 - Progressbar reports number of divergences in real time, when available [#3547](https://github.com/pymc-devs/pymc3/pull/3547).
 - Sampling from variational approximation now allows for alternative trace backends [#3550].
 - Infix `@` operator now works with random variables and deterministics [#3619](https://github.com/pymc-devs/pymc3/pull/3619).

docs/source/notebooks/GP-MeansAndCovs.ipynb

@@ -13,9 +13,11 @@
 #   limitations under the License.
 
 import numpy as np
+import theano
 import theano.tensor as tt
 from functools import reduce
 from operator import mul, add
+from numbers import Number
 
 __all__ = [
     "Constant",
@@ -100,6 +102,22 @@ def __radd__(self, other):
     def __rmul__(self, other):
         return self.__mul__(other)
 
+    def __pow__(self, other):
+        if(
+            isinstance(other, theano.compile.SharedVariable) and
+            other.get_value().squeeze().shape == ()
+        ):
+            other = tt.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")
+        
+
     def __array_wrap__(self, result):
         """
         Required to allow radd/rmul by numpy arrays.
@@ -172,6 +190,19 @@ def __call__(self, X, Xs=None, diag=False):
         return reduce(mul, self.merge_factors(X, Xs, diag))
 
 
+class Exponentiated(Covariance):
+    def __init__(self, kernel, power):
+        self.kernel = kernel
+        self.power = power
+        super().__init__(
+            input_dim=self.kernel.input_dim,
+            active_dims=self.kernel.active_dims
+        )
+
+    def __call__(self, X, Xs=None, diag=False):
+        return self.kernel(X, Xs, diag=diag) ** self.power
+
+
 class Kron(Covariance):
     r"""Form a covariance object that is the kronecker product of other covariances.
 

pymc3/gp/cov.py

@@ -237,6 +237,61 @@ def test_multiops(self):
         npt.assert_allclose(np.diag(K1), K2d, atol=1e-5)
         npt.assert_allclose(np.diag(K2), K1d, atol=1e-5)
 
+class TestCovExponentiation:
+    def test_symexp_cov(self):
+        X = np.linspace(0, 1, 10)[:, None]
+        with pm.Model() as model:
+            cov1 = pm.gp.cov.ExpQuad(1, 0.1)
+            cov = cov1 ** 2
+        K = theano.function([], cov(X))()
+        npt.assert_allclose(K[0, 1], 0.53940 ** 2, atol=1e-3)
+        # check diagonal
+        Kd = theano.function([], cov(X, diag=True))()
+        npt.assert_allclose(np.diag(K), Kd, atol=1e-5)
+
+    def test_covexp_numpy(self):
+        X = np.linspace(0, 1, 10)[:, None]
+        with pm.Model() as model:
+            a = np.array([[2]])
+            cov = pm.gp.cov.ExpQuad(1, 0.1) ** a
+        K = theano.function([], cov(X))()
+        npt.assert_allclose(K[0, 1], 0.53940 ** 2, atol=1e-3)
+        # check diagonal
+        Kd = theano.function([], cov(X, diag=True))()
+        npt.assert_allclose(np.diag(K), Kd, atol=1e-5)
+
+    def test_covexp_theano(self):
+        X = np.linspace(0, 1, 10)[:, None]
+        with pm.Model() as model:
+            a = tt.alloc(2.0, 1, 1)
+            cov = pm.gp.cov.ExpQuad(1, 0.1) ** a
+        K = theano.function([], cov(X))()
+        npt.assert_allclose(K[0, 1], 0.53940 ** 2, atol=1e-3)
+        # check diagonal
+        Kd = theano.function([], cov(X, diag=True))()
+        npt.assert_allclose(np.diag(K), Kd, atol=1e-5)
+
+    def test_covexp_shared(self):
+        X = np.linspace(0, 1, 10)[:, None]
+        with pm.Model() as model:
+            a = theano.shared(2.0)
+            cov = pm.gp.cov.ExpQuad(1, 0.1) ** a
+        K = theano.function([], cov(X))()
+        npt.assert_allclose(K[0, 1], 0.53940 ** 2, atol=1e-3)
+        # check diagonal
+        Kd = theano.function([], cov(X, diag=True))()
+        npt.assert_allclose(np.diag(K), Kd, atol=1e-5)
+
+    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 pm.Model() as model:
+                a = np.array([[1.0, 2.0]])
+                cov = pm.gp.cov.ExpQuad(1, 0.1) ** a
+
 
 class TestCovKron:
     def test_symprod_cov(self):