Generalize hartmann experimenter.

PiperOrigin-RevId: 628444635

vizier/_src/benchmarks/experimenters/synthetic/hartmann.py

@@ -14,108 +14,87 @@
 
 from __future__ import annotations
 
-"""HartMann 6-Dimensional function."""
+"""HartMann function."""
 
+import functools
 from typing import Sequence
 import numpy as np
 from vizier import pyvizier as vz
 from vizier._src.benchmarks.experimenters import experimenter
 from vizier._src.benchmarks.experimenters import numpy_experimenter
 
 
-def _hartmann6d(x: np.ndarray) -> float:
-  """Hartmann function.
-
-  Args:
-    x: Shape (6,) array
-
-  Returns:
-    Shape (,) array.
-  """
-  # For math notations:
-  # pylint: disable=invalid-name
-  alpha = np.asarray([1.0, 1.2, 3.0, 3.2])
-  A = np.asarray([
-      [10, 3, 17, 3.5, 1.7, 8],
-      [0.05, 10, 17, 0.1, 8, 14],
-      [3, 3.5, 1.7, 10, 17, 8],
-      [17, 8, 0.05, 10, 0.1, 14],
-  ])
-  P = 1e-4 * np.asarray([
-      [1312, 1696, 5569, 124, 8283, 5886],
-      [2329, 4135, 8307, 3736, 1004, 9991],
-      [2348, 1451, 3522, 2883, 3047, 6650],
-      [4047, 8828, 8732, 5743, 1091, 381],
-  ])
-  y = -alpha @ np.exp(-np.sum(A * (x - P) ** 2, axis=1))
-  return y
-
-
-class Hartmann6DExperimenter(experimenter.Experimenter):
-  """6D minimization function. See https://www.sfu.ca/~ssurjano/hart6.html."""
+def _hartmann(
+    x: np.ndarray, alpha: np.ndarray, A: np.ndarray, P: np.ndarray
+) -> float:
+  # pylint:disable=invalid-name
+  return -alpha @ np.exp(-np.sum(A * (x - P) ** 2, axis=1))
+
+
+class HartmannExperimenter(experimenter.Experimenter):
+  """General Hartmann minimization function."""
+
+  # pylint:disable=invalid-name
+  def __init__(self, alpha: np.ndarray, A: np.ndarray, P: np.ndarray):
+    self._dimension = A.shape[-1]
 
-  def __init__(self):
     self._impl = numpy_experimenter.NumpyExperimenter(
-        _hartmann6d, self.problem_statement()
+        functools.partial(_hartmann, alpha=alpha, A=A, P=P),
+        self.problem_statement(),
     )
 
   def evaluate(self, suggestions: Sequence[vz.Trial]) -> None:
     self._impl.evaluate(suggestions)
 
   def problem_statement(self) -> vz.ProblemStatement:
     problem = vz.ProblemStatement()
-    for i in range(1, 7):
+    for i in range(1, self._dimension + 1):
       problem.search_space.root.add_float_param(f"x{i}", 0, 1)
     problem.metric_information.append(
         vz.MetricInformation(name="value", goal=vz.ObjectiveMetricGoal.MINIMIZE)
     )
     return problem
 
 
-def _hartmann3d(x: np.ndarray) -> float:
-  """Hartmann function.
-
-  Args:
-    x: Shape (6,) array
-
-  Returns:
-    Shape (,) array.
-  """
-  # For math notations:
-  # pylint: disable=invalid-name
-  alpha = np.asarray([1.0, 1.2, 3.0, 3.2])
-  A = np.asarray([
-      [3, 10, 30],
-      [0.1, 10, 35],
-      [3, 10, 30],
-      [0.1, 10, 35],
-  ])
-  P = 1e-4 * np.asarray([
-      [3689, 1170, 2673],
-      [4699, 4387, 7470],
-      [1091, 8732, 5547],
-      [381, 5743, 8828],
-  ])
-  y = -alpha @ np.exp(-np.sum(A * (x - P) ** 2, axis=1))
-  return y
-
-
-class Hartmann3DExperimenter(experimenter.Experimenter):
-  """3D minimization function. See https://www.sfu.ca/~ssurjano/hart3.html."""
+class Hartmann3DExperimenter(HartmannExperimenter):
+  """See https://www.sfu.ca/~ssurjano/hart3.html."""
 
   def __init__(self):
-    self._impl = numpy_experimenter.NumpyExperimenter(
-        _hartmann3d, self.problem_statement()
+    super().__init__(
+        alpha=np.array([1.0, 1.2, 3.0, 3.2]),
+        A=np.array([
+            [3, 10, 30],
+            [0.1, 10, 35],
+            [3, 10, 30],
+            [0.1, 10, 35],
+        ]),
+        P=1e-4
+        * np.array([
+            [3689, 1170, 2673],
+            [4699, 4387, 7470],
+            [1091, 8732, 5547],
+            [381, 5743, 8828],
+        ]),
     )
 
-  def evaluate(self, suggestions: Sequence[vz.Trial]) -> None:
-    self._impl.evaluate(suggestions)
 
-  def problem_statement(self) -> vz.ProblemStatement:
-    problem = vz.ProblemStatement()
-    for i in range(1, 4):
-      problem.search_space.root.add_float_param(f"x{i}", 0, 1)
-    problem.metric_information.append(
-        vz.MetricInformation(name="value", goal=vz.ObjectiveMetricGoal.MINIMIZE)
+class Hartmann6DExperimenter(HartmannExperimenter):
+  """See https://www.sfu.ca/~ssurjano/hart6.html."""
+
+  def __init__(self):
+    super().__init__(
+        alpha=np.array([1.0, 1.2, 3.0, 3.2]),
+        A=np.array([
+            [10, 3, 17, 3.5, 1.7, 8],
+            [0.05, 10, 17, 0.1, 8, 14],
+            [3, 3.5, 1.7, 10, 17, 8],
+            [17, 8, 0.05, 10, 0.1, 14],
+        ]),
+        P=1e-4
+        * np.array([
+            [1312, 1696, 5569, 124, 8283, 5886],
+            [2329, 4135, 8307, 3736, 1004, 9991],
+            [2348, 1451, 3522, 2883, 3047, 6650],
+            [4047, 8828, 8732, 5743, 1091, 381],
+        ]),
     )
-    return problem

vizier/_src/benchmarks/experimenters/synthetic/hartmann_test.py

@@ -23,66 +23,48 @@
 from absl.testing import absltest
 
 
-class Hartmann6DExperimenterTest(absltest.TestCase):
-
-  def test_numpy_fn(self):
-    np.testing.assert_allclose(
-        hartmann._hartmann6d(
-            np.asarray(
-                [0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573]
-            )
-        ),
-        -3.32237,
-        atol=1e-5,
-    )
+class Hartmann3DExperimenterTest(absltest.TestCase):
 
   def test_experimenter_argmin(self):
     trial = vz.Trial(
         parameters={
-            f'x{i+1}': x
-            for i, x in enumerate(
-                [0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573]
-            )
+            f'x{i+1}': x for i, x in enumerate([0.114614, 0.555649, 0.852547])
         }
     )
-    hartmann.Hartmann6DExperimenter().evaluate([trial])
+    hartmann.Hartmann3DExperimenter().evaluate([trial])
     self.assertAlmostEqual(
         trial.final_measurement_or_die.metrics.get_value('value', np.nan),
-        -3.32237,
+        -3.86278,
         places=5,
     )
 
   def test_experimenter(self):
     experimenter_testing.assert_evaluates_random_suggestions(
-        self, hartmann.Hartmann6DExperimenter()
+        self, hartmann.Hartmann3DExperimenter()
     )
 
 
-class Hartmann3DExperimenterTest(absltest.TestCase):
-
-  def test_numpy_fn(self):
-    np.testing.assert_allclose(
-        hartmann._hartmann3d(np.asarray([0.114614, 0.555649, 0.852547])),
-        -3.86278,
-        atol=1e-5,
-    )
+class Hartmann6DExperimenterTest(absltest.TestCase):
 
   def test_experimenter_argmin(self):
     trial = vz.Trial(
         parameters={
-            f'x{i+1}': x for i, x in enumerate([0.114614, 0.555649, 0.852547])
+            f'x{i+1}': x
+            for i, x in enumerate(
+                [0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573]
+            )
         }
     )
-    hartmann.Hartmann3DExperimenter().evaluate([trial])
+    hartmann.Hartmann6DExperimenter().evaluate([trial])
     self.assertAlmostEqual(
         trial.final_measurement_or_die.metrics.get_value('value', np.nan),
-        -3.86278,
+        -3.32237,
         places=5,
     )
 
   def test_experimenter(self):
     experimenter_testing.assert_evaluates_random_suggestions(
-        self, hartmann.Hartmann3DExperimenter()
+        self, hartmann.Hartmann6DExperimenter()
     )
 
 

vizier/benchmarks/experimenters/__init__.py

@@ -42,4 +42,5 @@
 from vizier._src.benchmarks.experimenters.synthetic.branin import Branin2DExperimenter
 from vizier._src.benchmarks.experimenters.synthetic.hartmann import Hartmann3DExperimenter
 from vizier._src.benchmarks.experimenters.synthetic.hartmann import Hartmann6DExperimenter
+from vizier._src.benchmarks.experimenters.synthetic.hartmann import HartmannExperimenter
 from vizier._src.benchmarks.experimenters.synthetic.simplekd import SimpleKDExperimenter