Repository: cornellius-gp/gpytorch
Pull request #719 (2019-06-05 02:05:46)

Add horseshoe prior
Unlike spearmint or cornell-moe, which use `log(log(1 + 3.0 * (self.scale / X)) ** 2))`,  this implementation takes the average of the lower and upper bound on the pdf provided in the following reference:

C. M. Carvalho, N. G. Polson, and J. G. Scott. The horseshoe estimator for sparse signals. Biometrika, 2010.


Patch:
diff --git a/gpytorch/priors/__init__.py b/gpytorch/priors/__init__.py
index 59ccac77..0ece1e34 100644
--- a/gpytorch/priors/__init__.py
+++ b/gpytorch/priors/__init__.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 from .prior import Prior
+from .horseshoe_prior import HorseshoePrior
 from .lkj_prior import LKJCholeskyFactorPrior, LKJCovariancePrior, LKJPrior
 from .smoothed_box_prior import SmoothedBoxPrior
 from .torch_priors import GammaPrior, MultivariateNormalPrior, NormalPrior
@@ -12,6 +13,7 @@ from .torch_priors import GammaPrior, MultivariateNormalPrior, NormalPrior
 __all__ = [
     "Prior",
     "GammaPrior",
+    "HorseshoePrior",
     "LKJPrior",
     "LKJCholeskyFactorPrior",
     "LKJCovariancePrior",
diff --git a/gpytorch/priors/horseshoe_prior.py b/gpytorch/priors/horseshoe_prior.py
new file mode 100644
index 00000000..dbc263a7
--- /dev/null
+++ b/gpytorch/priors/horseshoe_prior.py
@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+
+import math
+from numbers import Number
+
+import torch
+from gpytorch.priors.prior import Prior
+from torch.distributions import constraints
+from torch.nn import Module as TModule
+
+
+class HorseshoePrior(Prior):
+    """Horseshoe prior.
+
+    There is no analytical form for the horeshoe prior's pdf, but it
+    satisfies a tight bound of the form `lb(x) <= pdf(x) <= ub(x)`, where
+
+        lb(x) = K/2 * log(1 + 4 * (scale / x) ** 2)
+        ub(x) = K * log(1 + 2 * (scale / x) ** 2)
+
+    with `K = 1 / sqrt(2 pi^3)`. Here, we simply use
+
+        pdf(x) ~ (lb(x) + ub(x)) / 2
+
+    Reference: C. M. Carvalho, N. G. Polson, and J. G. Scott.
+        The horseshoe estimator for sparse signals. Biometrika, 2010.
+    """
+
+    arg_constraints = {"scale": constraints.positive}
+    support = constraints.real
+    _validate_args = True
+
+    def __init__(self, scale, validate_args=False, transform=None):
+        TModule.__init__(self)
+        if isinstance(scale, Number):
+            scale = torch.tensor(float(scale))
+        self.K = 1 / math.sqrt(2 * math.pi ** 3)
+        self.scale = scale
+        super().__init__(scale.shape, validate_args=validate_args)
+        # now need to delete to be able to register buffer
+        del self.scale
+        self.register_buffer("scale", scale)
+        self._transform = transform
+
+    def log_prob(self, X):
+        A = (self.scale / self.transform(X)) ** 2
+        lb = self.K / 2 * torch.log(1 + 4 * A)
+        ub = self.K * torch.log(1 + 2 * A)
+        return torch.log((lb + ub) / 2)


Test patch:
diff --git a/test/priors/test_horseshoe_prior.py b/test/priors/test_horseshoe_prior.py
new file mode 100644
index 00000000..2c3f1514
--- /dev/null
+++ b/test/priors/test_horseshoe_prior.py
@@ -0,0 +1,71 @@
+#!/usr/bin/env python3
+
+import unittest
+
+import torch
+from gpytorch.priors import HorseshoePrior
+
+
+class TestHorseshoePrior(unittest.TestCase):
+    def test_horseshoe_prior_to_gpu(self):
+        if torch.cuda.is_available():
+            prior = HorseshoePrior(0.1).cuda()
+            self.assertEqual(prior.scale.device.type, "cuda")
+
+    def test_horseshoe_prior_validate_args(self):
+        with self.assertRaises(ValueError):
+            HorseshoePrior(-0.1, validate_args=True)
+
+    def test_horseshoe_prior_log_prob(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        scale = torch.tensor(0.1, device=device)
+        prior = HorseshoePrior(scale)
+        t = torch.rand(3, device=device)
+        prior.log_prob(t)
+
+    def test_horseshoe_prior_log_prob_cuda(self):
+        if torch.cuda.is_available():
+            return self.test_horseshoe_prior_log_prob(cuda=True)
+
+    def test_horseshoe_prior_log_prob_log_transform(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        scale = torch.tensor(0.1, device=device)
+        prior = HorseshoePrior(scale)
+        prior_tf = HorseshoePrior(scale, transform=torch.exp)
+        t = torch.tensor(0.5, device=device)
+        self.assertTrue(torch.equal(prior_tf.log_prob(t), prior.log_prob(t.exp())))
+
+    def test_horseshoe_prior_log_prob_log_transform_cuda(self):
+        if torch.cuda.is_available():
+            return self.test_horseshoe_prior_log_prob_log_transform(cuda=True)
+
+    def test_horseshoe_prior_batch_log_prob(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+
+        scale = torch.tensor([0.1, 0.25], device=device)
+        prior = HorseshoePrior(scale)
+        t = torch.ones(2, device=device)
+        prior.log_prob(t)
+        t = torch.ones(2, 2, device=device)
+        prior.log_prob(t)
+        with self.assertRaises(RuntimeError):
+            prior.log_prob(torch.ones(3, device=device))
+
+        scale = torch.tensor([[0.1, 0.25], [0.3, 1.0]], device=device)
+        prior = HorseshoePrior(scale)
+        t = torch.ones(2, device=device)
+        prior.log_prob(t)
+        t = torch.ones(2, 2, device=device)
+        prior.log_prob(t)
+        with self.assertRaises(RuntimeError):
+            prior.log_prob(torch.ones(3, device=device))
+        with self.assertRaises(RuntimeError):
+            prior.log_prob(torch.ones(2, 3, device=device))
+
+    def test_horseshoe_prior_batch_log_prob_cuda(self):
+        if torch.cuda.is_available():
+            return self.test_horseshoe_prior_batch_log_prob(cuda=True)
+
+
+if __name__ == "__main__":
+    unittest.main()