Update three augmented test functions for multi-fidelity bayesian optimization; updated _init_.py and branin.py (#168)

Summary:
Pull Request resolved: #168

We wrote three augmented test functions for multi-fidelity bayesian optimization from section 5.3 in https://arxiv.org/abs/1903.04703; _init_.py and branin.py were updated as well.

Reviewed By: bkarrer

Differential Revision: D15614824

fbshipit-source-id: eecc2938a17ce0def811d223c94ab98845cc9334

Author: VilockLi

botorch/test_functions/__init__.py

@@ -2,6 +2,9 @@
 
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
+from .aug_branin import neg_aug_branin
+from .aug_hartmann6 import neg_aug_hartmann6
+from .aug_rosenbrock import neg_aug_rosenbrock
 from .branin import neg_branin
 from .eggholder import neg_eggholder
 from .hartmann6 import neg_hartmann6
@@ -17,4 +20,7 @@
     "neg_holder_table",
     "neg_michalewicz",
     "neg_styblinski_tang",
+    "neg_aug_branin",
+    "neg_aug_hartmann6",
+    "neg_aug_rosenbrock",
 ]

botorch/test_functions/aug_branin.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import math
+
+import torch
+from torch import Tensor
+
+
+# This function has infinitely many global maximizers
+GLOBAL_MAXIMUM = -0.397887
+
+
+def neg_aug_branin(X: Tensor) -> Tensor:
+    r"""Negative augmented-Branin test function.
+
+    The last dimension of X is the fidelity.
+    3-dimensional function with domain `[-5, 10] x [0, 15] * [0,1]`:
+
+        B(x) = (x_2 - (b - 0.1 * (1 - x_3))x_1^2 + c x_1 - r)^2 +
+            10 (1-t) cos(x_1) + 10
+
+    Here `b`, `c`, `r` and `t` are constants where `b = 5.1 / (4 * math.pi ** 2)`
+    `c = 5 / math.pi`, `r = 6`, `t = 1 / (8 * math.pi)`.
+    B has infinitely many minimizers with `x_1 = -pi, pi, 3pi`
+
+    and `B_min = 0.397887`
+
+    Args:
+        X: A Tensor of size `3` or `k x 3` (`k` batch evaluations).
+
+    Returns:
+        `-B(X)`, the negative value of the augmented-Branin function.
+    """
+    batch = X.ndimension() > 1
+    X = X if batch else X.unsqueeze(0)
+    t1 = (
+        X[:, 1]
+        - (5.1 / (4 * math.pi ** 2) - 0.1 * (1 - X[:, 2])) * X[:, 0] ** 2
+        + 5 / math.pi * X[:, 0]
+        - 6
+    )
+    t2 = 10 * (1 - 1 / (8 * math.pi)) * torch.cos(X[:, 0])
+    B = t1 ** 2 + t2 + 10
+    result = -B
+    return result if batch else result.squeeze(0)

botorch/test_functions/aug_hartmann6.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import torch
+from torch import Tensor
+
+from .hartmann6 import ALPHA, A, P
+
+
+GLOBAL_MAXIMIZER = [0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573, 1]
+
+
+def neg_aug_hartmann6(X: Tensor) -> Tensor:
+    r"""Negative augmented Hartmann6 test function.
+
+    The last dimension of X is the fidelity parameter.
+    7-dimensional function (typically evaluated on `[0, 1]^7`):
+
+        H(x) = -(ALPHA_1 - 0.1 * (1-x_7)) * exp(- sum_{j=1}^6 A_1j (x_j - P_1j) ** 2) -
+            sum_{i=2}^4 ALPHA_i exp( - sum_{j=1}^6 A_ij (x_j - P_ij) ** 2)
+
+    H has unique global minimizer
+    x = [0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573, 1]
+
+    with `H_min = -3.32237`
+
+    Args:
+        X: A Tensor of size `7` or `k x 7` (k batch evaluations).
+
+    Returns:
+        `-H(X)`, the negative value of the augmented Hartmann6 function.
+    """
+    batch = X.ndimension() > 1
+    X = X if batch else X.unsqueeze(0)
+    inner_sum = torch.sum(
+        X.new(A) * (X[:, :6].unsqueeze(1) - 0.0001 * X.new(P)) ** 2, dim=2
+    )
+    alpha1 = ALPHA[0] - 0.1 * (1 - X[:, 6])
+    H = (
+        -torch.sum(X.new(ALPHA)[1:] * torch.exp(-inner_sum)[:, 1:], dim=1)
+        - alpha1 * torch.exp(-inner_sum)[:, 0]
+    )
+    result = -H
+    return result if batch else result.squeeze(0)

botorch/test_functions/aug_rosenbrock.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from torch import Tensor
+
+
+# This function has infinitely many global maximizers.
+GLOBAL_MAXIMUM = 0.0
+
+
+def neg_aug_rosenbrock(X: Tensor):
+    r"""Augmented Rosenbrock test function.
+
+    4-dimensional function, the last two dimensions are for fidelity parameters,
+    X is usually in [-5,10]^2 * [0,1]^2:
+
+        f(x) = 100 (x_2 - x_1^2 + 0.1 * (1 - x_3))^2 + (x_1 - 1 + 0.1 * (1 - x_4)^2)^2
+
+    f has infinitely many minimizers, with `f_min = 0.0`
+
+    Args:
+        X: A Tensor of size `4` or `k x 4` (`k` batch evaluations).
+    Returns:
+        `-f(X)`, the negative value of the augmented Rosenbrock function.
+    """
+    batch = X.ndimension() > 1
+    X = X if batch else X.unsqueeze(0)
+    result = (
+        -100 * (X[:, 1] - X[:, 0] ** 2 + 0.1 * (1 - X[:, 2])) ** 2
+        + (X[:, 0] - 1 + 0.1 * (1 - X[:, 3]) ** 2) ** 2
+    )
+    return result if batch else result.squeeze(0)

botorch/test_functions/branin.py

@@ -17,13 +17,15 @@ def neg_branin(X: Tensor) -> Tensor:
 
     Two-dimensional function (usually evaluated on `[-5, 10] x [0, 15]`):
 
-        `B(x) = (x2 - b x_1^2 + c x_1 - r)^2 + 10 (1-t) cos(x_1) + 10`
+        `B(x) = (x_2 - b x_1^2 + c x_1 - r)^2 + 10 (1-t) cos(x_1) + 10`
 
+    Here `b`, `c`, `r` and `t` are constants where `b = 5.1 / (4 * math.pi ** 2)`
+    `c = 5 / math.pi`, `r = 6`, `t = 1 / (8 * math.pi)`
     B has 3 minimizers for its global minimum at
 
         `z_1 = (-pi, 12.275), z_2 = (pi, 2.275), z_3 = (9.42478, 2.475)`
 
-    with `B(z_i) = -0.397887`
+    with `B(z_i) = 0.397887`
 
     Args:
         X: A Tensor of size `2` or `k x 2` (`k` batch evaluations).

test/test_functions/test_aug_branin.py

@@ -0,0 +1,61 @@
+#! /usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import math
+import unittest
+
+import torch
+from botorch.test_functions.aug_branin import GLOBAL_MAXIMUM, neg_aug_branin
+
+
+class TestNegAugBranin(unittest.TestCase):
+    def test_single_eval_neg_aug_branin(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            X = torch.zeros(3, device=device, dtype=dtype)
+            res = neg_aug_branin(X)
+            self.assertEqual(res.dtype, dtype)
+            self.assertEqual(res.device.type, device.type)
+            self.assertEqual(res.shape, torch.Size())
+
+    def test_single_eval_neg_aug_branin_cuda(self):
+        if torch.cuda.is_available():
+            self.test_single_eval_neg_aug_branin(cuda=True)
+
+    def test_batch_eval_neg_aug_branin(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            X = torch.zeros(2, 3, device=device, dtype=dtype)
+            res = neg_aug_branin(X)
+            self.assertEqual(res.dtype, dtype)
+            self.assertEqual(res.device.type, device.type)
+            self.assertEqual(res.shape, torch.Size([2]))
+
+    def test_batch_eval_neg_aug_branin_cuda(self):
+        if torch.cuda.is_available():
+            self.test_batch_eval_neg_aug_branin(cuda=True)
+
+    def test_neg_aug_branin_global_maxima(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            X = torch.tensor(
+                [
+                    [-math.pi, 12.275, 1],
+                    [math.pi, 1.3867356039019576, 0.1],
+                    [math.pi, 1.781519779945532, 0.5],
+                    [math.pi, 2.1763039559891064, 0.9],
+                ],
+                device=device,
+                dtype=dtype,
+                requires_grad=True,
+            )
+            res = neg_aug_branin(X)
+            for r in res:
+                self.assertAlmostEqual(r.item(), GLOBAL_MAXIMUM, places=4)
+            grad = torch.autograd.grad(res.sum(), X)[0]
+            self.assertLess(grad.abs().max().item(), 1e-4)
+
+    def test_neg_aug_branin_global_maxima_cuda(self):
+        if torch.cuda.is_available():
+            self.test_neg_aug_branin_global_maxima(cuda=True)

test/test_functions/test_aug_hartmann6.py

@@ -0,0 +1,48 @@
+#! /usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import unittest
+
+import torch
+from botorch.test_functions.aug_hartmann6 import GLOBAL_MAXIMIZER, neg_aug_hartmann6
+from botorch.test_functions.hartmann6 import GLOBAL_MAXIMUM
+
+
+class TestNegAugHartmann6(unittest.TestCase):
+    def test_single_eval_neg_aug_hartmann6(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            X = torch.zeros(7, device=device, dtype=dtype)
+            res = neg_aug_hartmann6(X)
+            self.assertEqual(res.dtype, dtype)
+            self.assertEqual(res.device.type, device.type)
+            self.assertEqual(res.shape, torch.Size())
+
+    def test_single_eval_neg_aug_hartmann6_cuda(self):
+        if torch.cuda.is_available():
+            self.test_single_eval_neg_aug_hartmann6(cuda=True)
+
+    def test_batch_eval_neg_aug_hartmann6(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            X = torch.zeros(2, 7, device=device, dtype=dtype)
+            res = neg_aug_hartmann6(X)
+            self.assertEqual(res.dtype, dtype)
+            self.assertEqual(res.device.type, device.type)
+            self.assertEqual(res.shape, torch.Size([2]))
+
+    def test_batch_eval_neg_aug_hartmann6_cuda(self):
+        if torch.cuda.is_available():
+            self.test_batch_eval_neg_aug_hartmann6(cuda=True)
+
+    def test_neg_aug_hartmann6_global_maximum(self, cuda=False):
+        device = torch.device("scuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            X = torch.tensor(GLOBAL_MAXIMIZER, device=device, dtype=dtype)
+            res = neg_aug_hartmann6(X)
+            self.assertAlmostEqual(res.item(), GLOBAL_MAXIMUM, places=4)
+
+    def test_neg_aug_hartmann6_global_maximum_cuda(self):
+        if torch.cuda.is_available():
+            self.test_neg_aug_hartmann6_global_maximum(cuda=False)

test/test_functions/test_aug_rosenbroc.py

@@ -0,0 +1,55 @@
+#! /usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import unittest
+
+import torch
+from botorch.test_functions.aug_rosenbrock import GLOBAL_MAXIMUM, neg_aug_rosenbrock
+
+
+class TestNegAugRosenbrock(unittest.TestCase):
+    def test_single_eval_neg_aug_rosenbrock(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            X = torch.zeros(4, device=device, dtype=dtype)
+            res = neg_aug_rosenbrock(X)
+            self.assertEqual(res.dtype, dtype)
+            self.assertEqual(res.device.type, device.type)
+            self.assertEqual(res.shape, torch.Size())
+
+    def test_single_eval_neg_aug_rosenbrock_cuda(self):
+        if torch.cuda.is_available():
+            self.test_single_eval_neg_aug_rosenbrock(cuda=True)
+
+    def test_batch_eval_neg_aug_rosenbrock(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            X = torch.zeros(2, 4, device=device, dtype=dtype)
+            res = neg_aug_rosenbrock(X)
+            self.assertEqual(res.dtype, dtype)
+            self.assertEqual(res.device.type, device.type)
+            self.assertEqual(res.shape, torch.Size([2]))
+
+    def test_batch_eval_neg_aug_rosenbrock_cuda(self):
+        if torch.cuda.is_available():
+            self.test_batch_eval_neg_aug_rosenbrock(cuda=True)
+
+    def test_neg_aug_rosenbrock_global_maximum(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            X = torch.tensor(
+                [[1, 1, 1, 1], [1, 0.95, 0.5, 1], [1, 0.91, 0.1, 1], [1, 0.99, 0.9, 1]],
+                device=device,
+                dtype=dtype,
+                requires_grad=True,
+            )
+            res = neg_aug_rosenbrock(X)
+            for r in res:
+                self.assertAlmostEqual(r.item(), GLOBAL_MAXIMUM, places=4)
+            grad = torch.autograd.grad(res.sum(), X)[0]
+            self.assertLess(grad.abs().max().item(), 1e-4)
+
+    def test_neg_aug_rosenbrock_global_maximum_cuda(self):
+        if torch.cuda.is_available():
+            self.test_neg_aug_rosenbrock_global_maximum(cuda=False)