Update Augmented Rosenbrock function

Summary: Fixed a typo and extended it to more dimensions in `X`.

Reviewed By: bkarrer

Differential Revision: D16090823

fbshipit-source-id: 86f7b941d3eaeff8d531877e15f099ee468916f8

Author: VilockLi

botorch/test_functions/aug_rosenbrock.py

@@ -5,29 +5,39 @@
 from torch import Tensor
 
 
-# This function has infinitely many global maximizers.
+GLOBAL_MAXIMIZER = [1.0]
 GLOBAL_MAXIMUM = 0.0
 
 
 def neg_aug_rosenbrock(X: Tensor):
-    r"""Augmented Rosenbrock test function.
+    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:
+    d-dimensional function (usually evaluated on `[-5, 10]^(d-2) * [0, 1]^2`),
+    the last two dimensions are the fidelity parameters:
 
-        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(x) = sum_{i=1}^{d-1} (100 (x_{i+1} - x_i^2 + 0.1 * (1-x_{d-1}))^2
+                + (x_i - 1 + 0.1 * (1 - x_d)^2)^2)'
 
-    f has infinitely many minimizers, with `f_min = 0.0`
+    f has one minimizer for its global minimum at
+
+        `z_1 = (1, 1, ..., 1)`
+
+    with `f(z_i) = 0.0`
 
     Args:
-        X: A Tensor of size `4` or `k x 4` (`k` batch evaluations).
+        X: A Tensor of size `d` or `k x d` (`k` batch evaluations).
+
     Returns:
-        `-f(X)`, the negative value of the augmented Rosenbrock function.
+        `-f(X)`, the negative value of the Augmented-Rosenbrock function.
     """
     batch = X.ndimension() > 1
     X = X if batch else X.unsqueeze(0)
+    X_curr = X[..., :-3]
+    X_next = X[..., 1:-2]
     result = (
-        -100 * (X[:, 1] - X[:, 0] ** 2 + 0.1 * (1 - X[:, 2])) ** 2
-        + (X[:, 0] - 1 + 0.1 * (1 - X[:, 3]) ** 2) ** 2
-    )
+        -(
+            100 * (X_next - X_curr ** 2 + 0.1 * (1 - X[..., -2].unsqueeze(-1))) ** 2
+            + (X_curr - 1 + 0.1 * (1 - X[..., -1].unsqueeze(-1)) ** 2) ** 2
+        )
+    ).sum(dim=-1)
     return result if batch else result.squeeze(0)

test/test_functions/test_aug_rosenbroc.py …st/test_functions/test_aug_rosenbrock.pytest/test_functions/test_aug_rosenbroc.py renamed to test/test_functions/test_aug_rosenbrock.py test/test_functions/test_aug_rosenbroc.py renamed to test/test_functions/test_aug_rosenbrock.py

@@ -5,14 +5,21 @@
 import unittest
 
 import torch
-from botorch.test_functions.aug_rosenbrock import GLOBAL_MAXIMUM, neg_aug_rosenbrock
+from botorch.test_functions.aug_rosenbrock import (
+    GLOBAL_MAXIMIZER,
+    GLOBAL_MAXIMUM,
+    neg_aug_rosenbrock,
+)
+
+
+DIMENSION = 5
 
 
 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)
+            X = torch.zeros(5, device=device, dtype=dtype)
             res = neg_aug_rosenbrock(X)
             self.assertEqual(res.dtype, dtype)
             self.assertEqual(res.device.type, device.type)
@@ -25,7 +32,7 @@ def test_single_eval_neg_aug_rosenbrock_cuda(self):
     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)
+            X = torch.zeros(2, 5, device=device, dtype=dtype)
             res = neg_aug_rosenbrock(X)
             self.assertEqual(res.dtype, dtype)
             self.assertEqual(res.device.type, device.type)
@@ -39,14 +46,13 @@ 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]],
+                GLOBAL_MAXIMIZER * DIMENSION,
                 device=device,
                 dtype=dtype,
                 requires_grad=True,
             )
             res = neg_aug_rosenbrock(X)
-            for r in res:
-                self.assertAlmostEqual(r.item(), GLOBAL_MAXIMUM, places=4)
+            self.assertAlmostEqual(res.item(), GLOBAL_MAXIMUM, places=4)
             grad = torch.autograd.grad(res.sum(), X)[0]
             self.assertLess(grad.abs().max().item(), 1e-4)