Add GP Wrapped Periodic Kernel

pymc/gp/cov.py

@@ -41,6 +41,7 @@
     "Cosine",
     "Periodic",
     "WarpedInput",
+    "WrappedPeriodic",
     "Gibbs",
     "Coregion",
     "ScaledCov",
@@ -502,12 +503,20 @@
 
     def euclidean_dist(self, X, Xs):
         r2 = self.square_dist(X, Xs)
+        return self._sqrt(r2)
+
+    def _sqrt(self, r2):
         return pt.sqrt(r2 + 1e-12)
 
     def diag(self, X):
         return pt.alloc(1.0, X.shape[0])
 
     def full(self, X, Xs=None):
+        X, Xs = self._slice(X, Xs)
+        r2 = self.square_dist(X, Xs)
+        return self.full_from_distance(r2, squared=True)
+
+    def full_from_distance(self, dist, squared=False):
         raise NotImplementedError
 
     def power_spectral_density(self, omega):
@@ -544,8 +553,14 @@
         f1 = X.dimshuffle(0, "x", 1)
         f2 = Xs.dimshuffle("x", 0, 1)
         r = np.pi * (f1 - f2) / self.period
-        r = pt.sum(pt.square(pt.sin(r) / self.ls), 2)
-        return pt.exp(-0.5 * r)
+        r2 = pt.sum(pt.square(pt.sin(r) / self.ls), 2)
+        return self.full_from_distance(r2, squared=True)
+
+    def full_from_distance(self, dist, squared=False):
+        # NOTE: This is the same as the ExpQuad as we assume the periodicity
+        # has already been accounted for in the distance
+        r2 = dist if squared else dist**2
+        return pt.exp(-0.5 * r2)
 
 
 class ExpQuad(Stationary):
@@ -559,9 +574,9 @@
 
     """
 
-    def full(self, X, Xs=None):
-        X, Xs = self._slice(X, Xs)
-        return pt.exp(-0.5 * self.square_dist(X, Xs))
+    def full_from_distance(self, dist, squared=False):
+        r2 = dist if squared else dist**2
+        return pt.exp(-0.5 * r2)
 
     def power_spectral_density(self, omega):
         r"""
@@ -592,10 +607,10 @@
         super().__init__(input_dim, ls, ls_inv, active_dims)
         self.alpha = alpha
 
-    def full(self, X, Xs=None):
-        X, Xs = self._slice(X, Xs)
+    def full_from_distance(self, dist, squared=False):
+        r2 = dist if squared else dist**2
         return pt.power(
-            (1.0 + 0.5 * self.square_dist(X, Xs) * (1.0 / self.alpha)),
+            (1.0 + 0.5 * r2 * (1.0 / self.alpha)),
             -1.0 * self.alpha,
         )
 
@@ -611,9 +626,8 @@
                    \mathrm{exp}\left[ - \frac{\sqrt{5(x - x')^2}}{\ell} \right]
     """
 
-    def full(self, X, Xs=None):
-        X, Xs = self._slice(X, Xs)
-        r = self.euclidean_dist(X, Xs)
+    def full_from_distance(self, dist, squared=False):
+        r = self._sqrt(dist) if squared else dist
         return (1.0 + np.sqrt(5.0) * r + 5.0 / 3.0 * pt.square(r)) * pt.exp(-1.0 * np.sqrt(5.0) * r)
 
     def power_spectral_density(self, omega):
@@ -651,9 +665,8 @@
                   \mathrm{exp}\left[ - \frac{\sqrt{3(x - x')^2}}{\ell} \right]
     """
 
-    def full(self, X, Xs=None):
-        X, Xs = self._slice(X, Xs)
-        r = self.euclidean_dist(X, Xs)
+    def full_from_distance(self, dist, squared=False):
+        r = self._sqrt(dist) if squared else dist
         return (1.0 + np.sqrt(3.0) * r) * pt.exp(-np.sqrt(3.0) * r)
 
     def power_spectral_density(self, omega):
@@ -690,9 +703,8 @@
         k(x, x') = \mathrm{exp}\left[ -\frac{(x - x')^2}{\ell} \right]
     """
 
-    def full(self, X, Xs=None):
-        X, Xs = self._slice(X, Xs)
-        r = self.euclidean_dist(X, Xs)
+    def full_from_distance(self, dist, squared=False):
+        r = self._sqrt(dist) if squared else dist
         return pt.exp(-r)
 
 
@@ -705,9 +717,9 @@
        k(x, x') = \mathrm{exp}\left[ -\frac{||x - x'||}{2\ell} \right]
     """
 
-    def full(self, X, Xs=None):
-        X, Xs = self._slice(X, Xs)
-        return pt.exp(-0.5 * self.euclidean_dist(X, Xs))
+    def full_from_distance(self, dist, squared=False):
+        r = self._sqrt(dist) if squared else dist
+        return pt.exp(-0.5 * r)
 
 
 class Cosine(Stationary):
@@ -718,9 +730,9 @@
        k(x, x') = \mathrm{cos}\left( 2 \pi \frac{||x - x'||}{ \ell^2} \right)
     """
 
-    def full(self, X, Xs=None):
-        X, Xs = self._slice(X, Xs)
-        return pt.cos(2.0 * np.pi * self.euclidean_dist(X, Xs))
+    def full_from_distance(self, dist, squared=False):
+        r = self._sqrt(dist) if squared else dist
+        return pt.cos(2.0 * np.pi * r)
 
 
 class Linear(Covariance):
@@ -814,6 +826,79 @@
         return self.cov_func(self.w(X, self.args), diag=True)
 
 
+class WrappedPeriodic(Covariance):
+    r"""
+    The `WrappedPeriodic` kernel constructs periodic kernels from any `Stationary` kernel.
+
+    This is done by warping the input with the function
+
+    .. math::
+        \mathbf{u}(x) = \left(
+            \mathrm{sin} \left( \frac{2\pi x}{T} \right),
+            \mathrm{cos} \left( \frac{2\pi x}{T} \right)
+        \right)
+
+    The original derivation as applied to the squared exponential kernel can
+    be found in [1] or referenced in Chapter 4, page 92 of [2].
+
+    Notes
+    -----
+    Note that we drop the resulting scaling by 4 found in the original derivation
+    so that the interpretation of the length scales is consistent between the
+    underlying kernel and the periodic version.
+
+    Examples
+    --------
+    In order to construct a kernel equivalent to the `Periodic` kernel you
+    can do the following (though using `Periodic` will likely be a bit faster):
+
+    .. code:: python
+
+        exp_quad = pm.gp.cov.ExpQuad(1, ls=0.5)
+        cov = pm.gp.cov.WrappedPeriodic(exp_quad, period=5)
+
+    References
+    ----------
+    .. [1] MacKay, D. J. C. (1998). Introduction to Gaussian Processes.
+       In Bishop, C. M., editor, Neural Networks and Machine Learning. Springer-Verlag
+    .. [2] Rasmussen, C. E., & Williams, C. K. I. (2006). Gaussian processes for machine learning. MIT Press.
+       http://gaussianprocess.org/gpml/chapters/
+
+    Parameters
+    ----------
+    cov_func: Stationary
+        Base kernel or covariance function
+    period: Period
+    """
+
+    def __init__(self, cov_func: Stationary, period):
+        if not isinstance(cov_func, Stationary):
+            raise TypeError("Must inherit from the Stationary class")
+        self.cov_func = cov_func
+        self.period = period
+
+    @property
+    def input_dim(self):
+        return self.cov_func.input_dim
+
+    @property
+    def active_dims(self):
+        return self.cov_func.active_dims
+
+    def full(self, X, Xs=None):
+        X, Xs = self._slice(X, Xs)
+        if Xs is None:
+            Xs = X
+        f1 = pt.expand_dims(X, axis=(0,))
+        f2 = pt.expand_dims(Xs, axis=(1,))
+        r = np.pi * (f1 - f2) / self.period
+        r2 = pt.sum(pt.square(pt.sin(r) / self.cov_func.ls), 2)
+        return self.cov_func.full_from_distance(r2, squared=True)
+
+    def diag(self, X):
+        return pt.alloc(1.0, X.shape[0])
+
+
 class Gibbs(Covariance):
     r"""
     The Gibbs kernel.  Use an arbitrary lengthscale function defined

tests/gp/test_cov.py

@@ -671,6 +671,28 @@ def test_raises(self):
             pm.gp.cov.WarpedInput(1, "str is not Covariance object", lambda x: x)
 
 
+class TestWrappedPeriodic:
+    def test_1d(self):
+        X = np.linspace(0, 1, 10)[:, None]
+        with pm.Model():
+            cov1 = pm.gp.cov.Periodic(1, ls=0.2, period=1)
+            cov2 = pm.gp.cov.WrappedPeriodic(
+                cov_func=pm.gp.cov.ExpQuad(1, ls=0.2),
+                period=1,
+            )
+        K1 = cov1(X).eval()
+        K2 = cov2(X).eval()
+        npt.assert_allclose(K1, K2, atol=1e-3)
+        K1d = cov1(X, diag=True).eval()
+        K2d = cov2(X, diag=True).eval()
+        npt.assert_allclose(K1d, K2d, atol=1e-3)
+
+    def test_raises(self):
+        lin_cov = pm.gp.cov.Linear(1, c=1)
+        with pytest.raises(TypeError, match="Must inherit from the Stationary class"):
+            pm.gp.cov.WrappedPeriodic(lin_cov, period=1)
+
+
 class TestGibbs:
     def test_1d(self):
         X = np.linspace(0, 2, 10)[:, None]