copybara-service · Nov 14, 2023 · Nov 14, 2023
diff --git a/vizier/_src/benchmarks/analyzers/convergence_curve.py b/vizier/_src/benchmarks/analyzers/convergence_curve.py
 
 
 @attr.define
-class ConvergenceCurveConverter:
+class StatefulCurveConverter(abc.ABC):
+  """Converter that updates its state upon each curve conversion.
+
+  This is to ensure that distributive property holds:
+  converter.convert(trials1 + trials2) = converter.convert(trials1) +
+  converter.convert(trials2)
+  """
+
+  @abc.abstractmethod
+  def convert(self, trials: Sequence[pyvizier.Trial]) -> ConvergenceCurve:
+    """Returns a ConvergenceCurve corresponding to trials and updates state."""
+
+
+@attr.define
+class ConvergenceCurveConverter(StatefulCurveConverter):
   """Converter for Trial sequence to ConvergenceCurve.
 
   Attributes:
@@ -242,6 +256,8 @@ class ConvergenceCurveConverter:
       measurements_type: ['final', 'intermediate', 'all']
       batch_size: Number of trials in each batch. In each batch, the order of
         trials is ignored and the best trial is used.
+      _best_yval: Best y-value seen. Updated in convert.
+      _cumulative_cost: Cumulative cost of trials. Updated in convert.
   """
 
   metric_information: pyvizier.MetricInformation
@@ -252,13 +268,17 @@ class ConvergenceCurveConverter:
   measurements_type: str = attr.field(default='final', kw_only=True)
   batch_size: int = attr.field(default=1, kw_only=True)
 
+  # Private attributes for stateful converts.
+  _best_yval: float = attr.field(default=np.nan, kw_only=True)
+  _cumulative_cost: float = attr.field(default=0.0, kw_only=True)
+
   def convert(self, trials: Sequence[pyvizier.Trial]) -> ConvergenceCurve:
     """Returns ConvergenceCurve with a single curve."""
     if not trials:
       raise ValueError(f'No trials provided {trials}')
 
-    yvals = [np.nan]
-    xvals = [0]
+    yvals = [self._best_yval]
+    xvals = [self._cumulative_cost]
     candidates = []
 
     for i in range(0, len(trials), self.batch_size):
@@ -286,6 +306,8 @@ def convert(self, trials: Sequence[pyvizier.Trial]) -> ConvergenceCurve:
       candidates = []
 
     yvals = np.asarray(yvals[1:])
+    self._best_yval = self.comparator(yvals)
+    self._cumulative_cost = xvals[-1]
     if self.metric_information.goal == pyvizier.ObjectiveMetricGoal.MAXIMIZE:
       trend = ConvergenceCurve.YTrend.INCREASING
       flipped = False
@@ -310,7 +332,7 @@ def comparator(self):
         == pyvizier.ObjectiveMetricGoal.MAXIMIZE) else np.nanmin
 
 
-class HypervolumeCurveConverter:
+class HypervolumeCurveConverter(StatefulCurveConverter):
   """Converts Trials to cumulative hypervolume curve for multiobjective."""
 
   def __init__(
@@ -351,6 +373,9 @@ def create_metric_converter(mc):
         ],
     )
     self._origin_value = reference_value
+    # TODO: Speed this up with hypervolume vector tracking.
+    self._min_trial_idx = 1
+    self._pareto_frontier = np.empty(shape=(0, len(metric_informations)))
 
   def convert(self, trials: Sequence[pyvizier.Trial]) -> ConvergenceCurve:
     """Returns ConvergenceCurve with a curve of shape 1 x len(trials)."""
@@ -381,15 +406,32 @@ def convert(self, trials: Sequence[pyvizier.Trial]) -> ConvergenceCurve:
           )
         origin = self._origin_value
 
+    # Calculate cumulative hypervolume with the Pareto frontier.
+    all_metrics = np.vstack(
+        [self._pareto_frontier, metrics]
+    )  # shape is [num_pareto_points + num_points, feature dimension]
     front = multimetric.ParetoFrontier(
-        points=metrics,
+        points=all_metrics,
         origin=origin,
         num_vectors=self._num_vectors,
         cum_hypervolume_base=xla_pareto.jax_cum_hypervolume_origin,
     )
-    hv_curve = front.hypervolume(is_cumulative=True)
+    all_hv_curve = front.hypervolume(is_cumulative=True)
+
+    # Remove the Pareto frontier add-in and update state.
+    hv_curve = all_hv_curve[len(self._pareto_frontier) :]
+    xs = np.asarray(
+        range(self._min_trial_idx, len(hv_curve) + self._min_trial_idx)
+    )
+    self._min_trial_idx += len(hv_curve)
+    algo = multimetric.FastParetoOptimalAlgorithm(
+        xla_pareto.JaxParetoOptimalAlgorithm()
+    )
+    pareto_points = algo.is_pareto_optimal(points=all_metrics)
+    self._pareto_frontier = all_metrics[pareto_points]
+
     return ConvergenceCurve(
-        xs=np.asarray(range(1, len(hv_curve) + 1)),
+        xs=xs,
         ys=np.asarray(hv_curve).reshape([1, -1]),
         trend=ConvergenceCurve.YTrend.INCREASING,
         ylabel='hypervolume',

diff --git a/vizier/_src/benchmarks/analyzers/convergence_curve_test.py b/vizier/_src/benchmarks/analyzers/convergence_curve_test.py
@@ -211,6 +211,36 @@ def test_convert_flip_signs_inf(self, goal, expected):
     np.testing.assert_array_equal(curve.xs, [1, 2, 3])
     np.testing.assert_array_equal(curve.ys, expected)
 
+  @parameterized.parameters(
+      (pyvizier.ObjectiveMetricGoal.MAXIMIZE, [2, 1, 4, 5], 2),
+      (pyvizier.ObjectiveMetricGoal.MAXIMIZE, [2, 1, 4], 1),
+      (pyvizier.ObjectiveMetricGoal.MAXIMIZE, [2, 1, 4, 7, 9, 8, 10], 4),
+      (pyvizier.ObjectiveMetricGoal.MINIMIZE, [3, 7, 2, 3, 4], 3),
+  )
+  def test_convert_with_state_updates(self, goal, values, split_idx):
+    trials = _gen_trials(values)
+    generator = convergence.ConvergenceCurveConverter(
+        pyvizier.MetricInformation(name='', goal=goal),
+        flip_signs_for_min=False,
+    )
+    curve = generator.convert(trials)
+
+    generator = convergence.ConvergenceCurveConverter(
+        pyvizier.MetricInformation(name='', goal=goal),
+        flip_signs_for_min=False,
+    )
+    first_trials = trials[:split_idx]
+    second_trials = trials[split_idx:]
+    first_curve = generator.convert(first_trials)
+    second_curve = generator.convert(second_trials)
+
+    np.testing.assert_array_equal(
+        curve.xs, np.concatenate((first_curve.xs, second_curve.xs), axis=-1)
+    )
+    np.testing.assert_array_equal(
+        curve.ys, np.concatenate((first_curve.ys, second_curve.ys), axis=-1)
+    )
+
 
 class HypervolumeCurveConverterTest(parameterized.TestCase):
 
@@ -335,6 +365,55 @@ def test_convert_with_inf_none_reference(self):
     np.testing.assert_array_equal(curve.xs, [1, 2])
     np.testing.assert_array_equal(curve.ys, [[0.0, 0.0]])
 
+  def test_convert_with_state(self):
+    generator = convergence.HypervolumeCurveConverter(
+        [
+            pyvizier.MetricInformation(
+                name='max', goal=pyvizier.ObjectiveMetricGoal.MAXIMIZE
+            ),
+            pyvizier.MetricInformation(
+                name='min', goal=pyvizier.ObjectiveMetricGoal.MINIMIZE
+            ),
+        ],
+        reference_value=np.array([0.0]),
+    )
+    pytrials = []
+    pytrials.append(
+        pyvizier.Trial().complete(
+            pyvizier.Measurement(metrics={'max': 4.0, 'min': 2.0})
+        )
+    )
+    pytrials.append(
+        pyvizier.Trial().complete(
+            pyvizier.Measurement(metrics={'max': 5.0, 'min': -1.0})
+        )
+    )
+    pytrials.append(
+        pyvizier.Trial().complete(
+            pyvizier.Measurement(metrics={'max': 4.0, 'min': -2.0})
+        )
+    )
+
+    curve = generator.convert(pytrials)
+    np.testing.assert_array_equal(curve.xs, [1, 2, 3])
+    np.testing.assert_array_almost_equal(curve.ys, [[0.0, 5.0, 9.0]], decimal=1)
+
+    pytrials = []
+    pytrials.append(
+        pyvizier.Trial().complete(
+            pyvizier.Measurement(metrics={'max': 2.0, 'min': -0.1})
+        )
+    )
+    pytrials.append(
+        pyvizier.Trial().complete(
+            pyvizier.Measurement(metrics={'max': 5.0, 'min': -2.0})
+        )
+    )
+
+    curve = generator.convert(pytrials)
+    np.testing.assert_array_equal(curve.xs, [4, 5])
+    np.testing.assert_array_almost_equal(curve.ys, [[9.0, 10.0]], decimal=1)
+
 
 class MultiMetricCurveConverterTest(parameterized.TestCase):
 

diff --git a/vizier/_src/benchmarks/analyzers/state_analyzer.py b/vizier/_src/benchmarks/analyzers/state_analyzer.py
@@ -106,22 +106,22 @@ def to_curve(
       raise ValueError('Empty States.')
 
     problem_statement = states[0].experimenter.problem_statement()
-    if not problem_statement.is_single_objective:
-      raise ValueError('Multiobjective Conversion not supported yet.')
 
-    converter = convergence_curve.ConvergenceCurveConverter(
-        problem_statement.metric_information.item(),
-        flip_signs_for_min=flip_signs_for_min,
-    )
     curves = []
     for state in states:
       if problem_statement != state.experimenter.problem_statement():
         raise ValueError(
             f'States must have same problem {problem_statement}'
             f' and {state.experimenter.problem_statement()}'
         )
-
       state_trials = state.algorithm.supporter.GetTrials()
+
+      converter = (
+          convergence_curve.MultiMetricCurveConverter.from_metrics_config(
+              problem_statement.metric_information,
+              flip_signs_for_min=flip_signs_for_min,
+          )
+      )
       curve = converter.convert(state_trials)
       curves.append(curve)
     return convergence_curve.ConvergenceCurve.align_xs(curves)[0]