Repository: cornellius-gp/gpytorch
Pull request #577 (2019-03-16 22:01:38)

[WIP] implement cylindrical kernel
Implement the Cylindrical Kernel as defined in http://proceedings.mlr.press/v80/oh18a.html. Addresses issue #189.

- [x] Implement the Kernel
- [x] Sanity check on toy regression problem
- [ ] Write unit tests


Patch:
diff --git a/docs/source/kernels.rst b/docs/source/kernels.rst
index f9acc0c0..36d30564 100644
--- a/docs/source/kernels.rst
+++ b/docs/source/kernels.rst
@@ -27,6 +27,14 @@ Standard Kernels
 .. autoclass:: CosineKernel
    :members:
 
+
+:hidden:`CylindricalKernel`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: CylindricalKernel
+   :members:
+
+
 :hidden:`LinearKernel`
 ~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/gpytorch/kernels/__init__.py b/gpytorch/kernels/__init__.py
index 7ee10c48..447e2357 100644
--- a/gpytorch/kernels/__init__.py
+++ b/gpytorch/kernels/__init__.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 from .additive_structure_kernel import AdditiveStructureKernel
+from .cylindrical_kernel import CylindricalKernel
 from .cosine_kernel import CosineKernel
 from .multi_device_kernel import MultiDeviceKernel
 from .grid_interpolation_kernel import GridInterpolationKernel
@@ -25,6 +26,7 @@ __all__ = [
     "Kernel",
     "AdditiveKernel",
     "AdditiveStructureKernel",
+    "CylindricalKernel",
     "MultiDeviceKernel",
     "CosineKernel",
     "GridKernel",
diff --git a/gpytorch/kernels/cylindrical_kernel.py b/gpytorch/kernels/cylindrical_kernel.py
new file mode 100644
index 00000000..7c87e535
--- /dev/null
+++ b/gpytorch/kernels/cylindrical_kernel.py
@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+
+import torch
+from typing import Optional
+from .kernel import Kernel
+from ..constraints import Positive, Interval
+from ..priors import Prior
+from .. import settings
+
+
+class CylindricalKernel(Kernel):
+    r"""
+    Computes a covariance matrix based on the Cylindrical Kernel between
+    inputs :math:`mathbf{x_1}` and :math:`mathbf{x_2}`.
+    It was proposed in `BOCK: Bayesian Optimization with Cylindrical Kernels`.
+    See http://proceedings.mlr.press/v80/oh18a.html for more details
+
+    .. note::
+        The data must lie completely within the unit ball.
+
+    Args:
+        :attr:`num_angular_weights` (int):
+            The number of components in the angular kernel
+        :attr:`radial_base_kernel` (gpytorch.kernel):
+            The base kernel for computing the radial kernel
+        :attr:`batch_size` (int, optional):
+            Set this if the data is batch of input data.
+            It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `1`
+        :attr:`eps` (float):
+            Small floating point number used to improve numerical stability
+            in kernel computations. Default: `1e-6`
+        :attr:`param_transform` (function, optional):
+            Set this if you want to use something other than softplus to ensure positiveness of parameters.
+        :attr:`inv_param_transform` (function, optional):
+            Set this to allow setting parameters directly in transformed space and sampling from priors.
+            Automatically inferred for common transformations such as torch.exp or torch.nn.functional.softplus.
+
+    """
+    def __init__(
+        self,
+        num_angular_weights: int,
+        radial_base_kernel: Kernel,
+        eps: Optional[int] = 1e-6,
+        angular_weights_prior: Optional[Prior] = None,
+        angular_weights_constraint: Optional[Interval] = None,
+        alpha_prior: Optional[Prior] = None,
+        alpha_constraint: Optional[Interval] = None,
+        beta_prior: Optional[Prior] = None,
+        beta_constraint: Optional[Interval] = None,
+        **kwargs
+    ):
+        if angular_weights_constraint is None:
+            angular_weights_constraint = Positive()
+
+        if alpha_constraint is None:
+            alpha_constraint = Positive()
+
+        if beta_constraint is None:
+            beta_constraint = Positive()
+
+        super().__init__(has_lengthscale=False, **kwargs)
+        self.num_angular_weights = num_angular_weights
+        self.radial_base_kernel = radial_base_kernel
+        self.eps = eps
+
+        self.register_parameter(name="raw_angular_weights",
+                                parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, num_angular_weights)))
+        self.register_constraint("raw_angular_weights", angular_weights_constraint)
+        self.register_parameter(name="raw_alpha", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)))
+        self.register_constraint("raw_alpha", alpha_constraint)
+        self.register_parameter(name="raw_beta", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)))
+        self.register_constraint("raw_beta", beta_constraint)
+
+        if angular_weights_prior is not None:
+            self.register_prior(
+                "angular_weights_prior", angular_weights_prior,
+                lambda: self.angular_weights, lambda v: self._set_angular_weights(v)
+            )
+        if alpha_prior is not None:
+            self.register_prior(
+                "alpha_prior", alpha_prior, lambda: self.alpha, lambda v: self._set_alpha(v)
+            )
+        if beta_prior is not None:
+            self.register_prior(
+                "beta_prior", beta_prior, lambda: self.beta, lambda v: self._set_beta(v)
+            )
+
+    @property
+    def angular_weights(self) -> torch.Tensor:
+        return self.raw_angular_weights_constraint.transform(self.raw_angular_weights)
+
+    @angular_weights.setter
+    def angular_weights(self, value: torch.Tensor) -> None:
+        if not torch.is_tensor(value):
+            value = torch.tensor(value)
+
+        self.initialize(raw_angular_weights=self.raw_angular_weights_constraint.inverse_transform(value))
+
+    @property
+    def alpha(self) -> torch.Tensor:
+        return self.raw_alpha_constraint.transform(self.raw_alpha)
+
+    @alpha.setter
+    def alpha(self, value: torch.Tensor) -> None:
+        if not torch.is_tensor(value):
+            value = torch.tensor(value)
+
+        self.initialize(raw_alpha=self.raw_alpha_constraint.inverse_transform(value))
+
+    @property
+    def beta(self) -> torch.Tensor:
+        return self.raw_beta_constraint.transform(self.raw_beta)
+
+    @beta.setter
+    def beta(self, value: torch.Tensor) -> None:
+        if not torch.is_tensor(value):
+            value = torch.tensor(value)
+
+        self.initialize(raw_beta=self.raw_beta_constraint.inverse_transform(value))
+
+    def forward(
+        self,
+        x1: torch.Tensor,
+        x2: torch.Tensor,
+        diag: Optional[bool] = False,
+        **params
+    ) -> torch.Tensor:
+
+        x1_, x2_ = x1.clone(), x2.clone()
+        # Jitter datapoints that are exactly 0
+        x1_[x1_ == 0], x2_[x2_ == 0] = x1_[x1_ == 0] + self.eps, x2_[x2_ == 0] + self.eps
+        r1, r2 = x1_.norm(dim=-1, keepdim=True), x2_.norm(dim=-1, keepdim=True)
+
+        if torch.any(r1 > 1.) or torch.any(r2 > 1.):
+            raise RuntimeError("Cylindrical kernel not defined for data points with radius > 1. Scale your data!")
+
+        a1, a2 = x1.div(r1), x2.div(r2)
+        if not diag:
+            gram_mat = a1.matmul(a2.transpose(-2, -1))
+            for p in range(self.num_angular_weights):
+                if p == 0:
+                    angular_kernel = self.angular_weights[..., 0, None, None]
+                else:
+                    angular_kernel = angular_kernel + self.angular_weights[..., p, None, None].mul(gram_mat.pow(p))
+        else:
+            gram_mat = a1.mul(a2).sum(-1)
+            for p in range(self.num_angular_weights):
+                if p == 0:
+                    angular_kernel = self.angular_weights[..., 0, None]
+                else:
+                    angular_kernel = angular_kernel + self.angular_weights[..., p, None].mul(gram_mat.pow(p))
+
+        with settings.lazily_evaluate_kernels(False):
+            radial_kernel = self.radial_base_kernel(
+                self.kuma(r1),
+                self.kuma(r2),
+                diag=diag,
+                **params
+            )
+        return radial_kernel.mul(angular_kernel)
+
+    def kuma(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.alpha.view(*self.batch_shape, 1, 1)
+        beta = self.beta.view(*self.batch_shape, 1, 1)
+
+        res = 1 - (1 - x.pow(alpha) + self.eps).pow(beta)
+        return res
+
+    def num_outputs_per_input(self, x1: torch.Tensor, x2: torch.Tensor) -> int:
+        return self.radial_base_kernel.num_outputs_per_input(x1, x2)


Test patch:
diff --git a/test/kernels/_base_kernel_test_case.py b/test/kernels/_base_kernel_test_case.py
index 0df9e725..77d2ea1b 100644
--- a/test/kernels/_base_kernel_test_case.py
+++ b/test/kernels/_base_kernel_test_case.py
@@ -13,9 +13,18 @@ class BaseKernelTestCase(object):
     def create_kernel_ard(self, num_dims, **kwargs):
         raise NotImplementedError()
 
+    def create_data_no_batch(self):
+        return torch.randn(50, 10)
+
+    def create_data_single_batch(self):
+        return torch.randn(2, 50, 2)
+
+    def create_data_double_batch(self):
+        return torch.randn(3, 2, 50, 2)
+
     def test_active_dims_list(self):
         kernel = self.create_kernel_no_ard(active_dims=[0, 2, 4, 6])
-        x = torch.randn(50, 10)
+        x = self.create_data_no_batch()
         covar_mat = kernel(x).evaluate_kernel().evaluate()
         kernel_basic = self.create_kernel_no_ard()
         covar_mat_actual = kernel_basic(x[:, [0, 2, 4, 6]]).evaluate_kernel().evaluate()
@@ -25,7 +34,7 @@ class BaseKernelTestCase(object):
     def test_active_dims_range(self):
         active_dims = list(range(3, 9))
         kernel = self.create_kernel_no_ard(active_dims=active_dims)
-        x = torch.randn(50, 10)
+        x = self.create_data_no_batch()
         covar_mat = kernel(x).evaluate_kernel().evaluate()
         kernel_basic = self.create_kernel_no_ard()
         covar_mat_actual = kernel_basic(x[:, active_dims]).evaluate_kernel().evaluate()
@@ -34,7 +43,7 @@ class BaseKernelTestCase(object):
 
     def test_no_batch_kernel_single_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard()
-        x = torch.randn(2, 50, 2)
+        x = self.create_data_single_batch()
         batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
 
         actual_mat_1 = kernel(x[0]).evaluate_kernel().evaluate()
@@ -45,7 +54,7 @@ class BaseKernelTestCase(object):
 
     def test_single_batch_kernel_single_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
-        x = torch.randn(2, 50, 2)
+        x = self.create_data_single_batch()
         batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
 
         actual_mat_1 = kernel(x[0]).evaluate_kernel().evaluate()
@@ -56,7 +65,7 @@ class BaseKernelTestCase(object):
 
     def test_no_batch_kernel_double_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
-        x = torch.randn(3, 2, 50, 2)
+        x = self.create_data_double_batch()
         batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
 
         ij_actual_covars = []
@@ -76,7 +85,7 @@ class BaseKernelTestCase(object):
         except NotImplementedError:
             return
 
-        x = torch.randn(3, 2, 50, 2)
+        x = self.create_data_double_batch()
         batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
 
         ij_actual_covars = []
@@ -92,7 +101,7 @@ class BaseKernelTestCase(object):
 
     def test_smoke_double_batch_kernel_double_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([3, 2]))
-        x = torch.randn(3, 2, 50, 2)
+        x = self.create_data_double_batch()
         batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
         return batch_covar_mat
 
@@ -102,6 +111,6 @@ class BaseKernelTestCase(object):
         except NotImplementedError:
             return
 
-        x = torch.randn(3, 2, 50, 2)
+        x = self.create_data_double_batch()
         batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
         return batch_covar_mat
diff --git a/test/kernels/test_cylindrical_kernel.py b/test/kernels/test_cylindrical_kernel.py
new file mode 100644
index 00000000..dfe5ae64
--- /dev/null
+++ b/test/kernels/test_cylindrical_kernel.py
@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+
+import math
+import torch
+import unittest
+from gpytorch.kernels import CylindricalKernel, MaternKernel
+from test.kernels._base_kernel_test_case import BaseKernelTestCase
+
+
+class TestCylindricalKernel(unittest.TestCase, BaseKernelTestCase):
+    def create_kernel_no_ard(self, **kwargs):
+        return CylindricalKernel(5, MaternKernel(nu=2.5), **kwargs)
+
+    def create_data_no_batch(self):
+        return torch.rand(50, 10) / math.sqrt(10)
+
+    def create_data_single_batch(self):
+        return torch.rand(2, 50, 2) / math.sqrt(2)
+
+    def create_data_double_batch(self):
+        return torch.rand(3, 2, 50, 2) / math.sqrt(2)
+
+
+if __name__ == "__main__":
+    unittest.main()