Merge pull request #255 from optimas-org/enable_new_ax

Enable new ax

Author: RemiLehe

.github/workflows/unix-openmpi.yml

@@ -28,7 +28,7 @@ jobs:
     - shell: bash -l {0}
       name: Install dependencies
       run: |
-        conda install numpy=1 pandas pytorch cpuonly -c pytorch
+        conda install numpy pandas pytorch cpuonly -c pytorch
         conda install -c conda-forge mpi4py openmpi=5.*
         pip install .[test]
     - shell: bash -l {0}

.github/workflows/unix.yml

@@ -28,7 +28,7 @@ jobs:
     - shell: bash -l {0}
       name: Install dependencies
       run: |
-        conda install numpy=1 pandas pytorch cpuonly -c pytorch
+        conda install numpy pandas pytorch cpuonly -c pytorch
         conda install -c conda-forge mpi4py mpich
         pip install .[test]
     - shell: bash -l {0}

doc/environment.yaml

@@ -5,7 +5,7 @@ dependencies:
   - pip
   - pip:
     - -e ..
-    - ax-platform == 0.4.0
+    - ax-platform >= 0.5.0
     - autodoc_pydantic >= 2.0.1
     - ipykernel
     - matplotlib

optimas/generators/ax/developer/multitask.py

@@ -3,12 +3,11 @@
 import os
 from copy import deepcopy
 from typing import List, Dict, Tuple, Optional, Union
+from pyre_extensions import assert_is_instance
 
 import numpy as np
 import torch
-from packaging import version
 
-from ax.version import version as ax_version
 from ax.core.arm import Arm
 from ax.core.batch_trial import BatchTrial
 from ax.core.multi_type_experiment import MultiTypeExperiment
@@ -22,8 +21,40 @@
 from ax.core.observation import ObservationFeatures
 from ax.core.generator_run import GeneratorRun
 from ax.storage.json_store.save import save_experiment
-from ax.storage.metric_registry import register_metric
-from ax.modelbridge.factory import get_MTGP_LEGACY as get_MTGP
+from ax.storage.metric_registry import register_metrics
+
+from ax.modelbridge.registry import Models, ST_MTGP_trans
+
+try:
+    # For Ax >= 0.5.0
+    from ax.modelbridge.transforms.derelativize import Derelativize
+    from ax.modelbridge.transforms.convert_metric_names import (
+        ConvertMetricNames,
+    )
+    from ax.modelbridge.transforms.trial_as_task import TrialAsTask
+    from ax.modelbridge.transforms.stratified_standardize_y import (
+        StratifiedStandardizeY,
+    )
+    from ax.modelbridge.transforms.task_encode import TaskChoiceToIntTaskChoice
+    from ax.modelbridge.registry import MBM_X_trans
+
+    MT_MTGP_trans = MBM_X_trans + [
+        Derelativize,
+        ConvertMetricNames,
+        TrialAsTask,
+        StratifiedStandardizeY,
+        TaskChoiceToIntTaskChoice,
+    ]
+
+except ImportError:
+    # For Ax < 0.5.0
+    from ax.modelbridge.registry import MT_MTGP_trans
+
+from ax.core.experiment import Experiment
+from ax.core.data import Data
+from ax.modelbridge.transforms.convert_metric_names import (
+    tconfig_from_mt_experiment,
+)
 
 from optimas.generators.ax.base import AxGenerator
 from optimas.core import (
@@ -37,13 +68,80 @@
 )
 from .ax_metric import AxMetric
 
-
 # Define generator states.
 NOT_STARTED = "not_started"
 LOFI_RETURNED = "lofi_returned"
 HIFI_RETURNED = "hifi_returned"
 
 
+# get_MTGP is not part of the Ax codebase, as of Ax 0.4.1, due to this PR:
+# https://github.com/facebook/Ax/pull/2508
+# Here we use `get_MTGP` https://ax.dev/docs/tutorials/multi_task/
+def get_MTGP(
+    experiment: Experiment,
+    data: Data,
+    search_space: Optional[SearchSpace] = None,
+    trial_index: Optional[int] = None,
+    device: torch.device = torch.device("cpu"),
+    dtype: torch.dtype = torch.double,
+) -> TorchModelBridge:
+    """Instantiate a Multi-task Gaussian Process (MTGP) model.
+
+    Points are generated with EI (Expected Improvement).
+    If the input experiment is a MultiTypeExperiment then a
+    Multi-type Multi-task GP model will be instantiated.
+    Otherwise, the model will be a Single-type Multi-task GP.
+    """
+    if isinstance(experiment, MultiTypeExperiment):
+        trial_index_to_type = {
+            t.index: t.trial_type for t in experiment.trials.values()
+        }
+        transforms = MT_MTGP_trans
+        transform_configs = {
+            "TrialAsTask": {
+                "trial_level_map": {"trial_type": trial_index_to_type}
+            },
+            "ConvertMetricNames": tconfig_from_mt_experiment(experiment),
+        }
+    else:
+        # Set transforms for a Single-type MTGP model.
+        transforms = ST_MTGP_trans
+        transform_configs = None
+
+    # Choose the status quo features for the experiment from the selected
+    # trial. If trial_index is None, we will look for a status quo from the
+    # last experiment trial to use as a status quo for the experiment.
+    if trial_index is None:
+        trial_index = len(experiment.trials) - 1
+    elif trial_index >= len(experiment.trials):
+        raise ValueError(
+            "trial_index is bigger than the number of experiment trials"
+        )
+
+    status_quo = experiment.trials[trial_index].status_quo
+    if status_quo is None:
+        status_quo_features = None
+    else:
+        status_quo_features = ObservationFeatures(
+            parameters=status_quo.parameters,
+            trial_index=trial_index,  # pyre-ignore[6]
+        )
+
+    return assert_is_instance(
+        Models.ST_MTGP(
+            experiment=experiment,
+            search_space=search_space or experiment.search_space,
+            data=data,
+            transforms=transforms,
+            transform_configs=transform_configs,
+            torch_dtype=dtype,
+            torch_device=device,
+            status_quo_features=status_quo_features,
+        ),
+        TorchModelBridge,
+    )
+
+
 class AxMultitaskGenerator(AxGenerator):
     """Multitask Bayesian optimization using the Ax developer API.
 
@@ -307,7 +405,9 @@ def _create_experiment(self) -> MultiTypeExperiment:
         )
 
         # Register metric in order to be able to save experiment to json file.
-        _, encoder_registry, decoder_registry = register_metric(AxMetric)
+        _, encoder_registry, decoder_registry = register_metrics(
+            {AxMetric: None}
+        )
         self._encoder_registry = encoder_registry
         self._decoder_registry = decoder_registry
 

optimas/generators/ax/import_error_dummy_generator.py

@@ -12,5 +12,5 @@ def __init__(self, *args, **kwargs) -> None:
         raise RuntimeError(
             "You need to install ax-platform, in order "
             "to use Ax-based generators in optimas.\n"
-            "e.g. with `pip install ax-platform >= 0.4.0`"
+            "e.g. with `pip install ax-platform > 0.5.0`"
         )

optimas/generators/ax/service/base.py

@@ -205,7 +205,7 @@ def _tell(self, trials: List[Trial]) -> None:
                             # i.e., min trials).
                             if isinstance(tc, (MinTrials, MaxTrials)):
                                 tc.threshold -= 1
-                        generation_strategy._maybe_move_to_next_step()
+                        generation_strategy._maybe_transition_to_next_node()
             finally:
                 if trial.ignored:
                     continue

optimas/generators/ax/service/multi_fidelity.py

@@ -124,7 +124,13 @@ def _create_generation_steps(
             GenerationStep(
                 model=Models.BOTORCH_MODULAR,
                 num_trials=-1,
-                model_kwargs=bo_model_kwargs,
+                model_kwargs={
+                    **bo_model_kwargs,
+                    "acquisition_options": {
+                        "X_pending": None,
+                        "constraints": None,
+                    },
+                },
                 model_gen_kwargs={
                     "model_gen_options": {
                         Keys.ACQF_KWARGS: {

optimas/generators/ax/service/single_fidelity.py

@@ -134,23 +134,13 @@ def _create_generation_steps(
     ) -> List[GenerationStep]:
         """Create generation steps for single-fidelity optimization."""
         # Select BO model.
+        # Ax 0.5.0 detects if multi-objective.
         if self._fully_bayesian:
-            if len(self.objectives) > 1:
-                # Use a SAAS model with qNEHVI acquisition function.
-                MODEL_CLASS = Models.FULLYBAYESIANMOO
-            else:
-                # Use a SAAS model with qNEI acquisition function.
-                MODEL_CLASS = Models.FULLYBAYESIAN
-            # Disable additional logs from fully Bayesian model.
-            bo_model_kwargs["disable_progbar"] = True
-            bo_model_kwargs["verbose"] = False
+            # Use a SAAS model with qNEHVI/qNEI acquisition function
+            MODEL_CLASS = Models.SAASBO
         else:
-            if len(self.objectives) > 1:
-                # Use a model with qNEHVI acquisition function.
-                MODEL_CLASS = Models.MOO
-            else:
-                # Use a model with qNEI acquisition function.
-                MODEL_CLASS = Models.GPEI
+            # Use a model with qNEHVI/qNEI acquisition function
+            MODEL_CLASS = Models.BOTORCH_MODULAR
 
         # Make generation strategy.
         steps = []

optimas/utils/ax/ax_model_manager.py

@@ -130,7 +130,7 @@ def _build_ax_client_from_dataframe(
         # allow calling `model.predict`. Using MOO for multiobjective is
         # needed because otherwise calls to `get_pareto_optimal_parameters`
         # would fail.
-        model = Models.GPEI if len(objectives) == 1 else Models.MOO
+        model = Models.BOTORCH_MODULAR
         gs = GenerationStrategy([GenerationStep(model=model, num_trials=-1)])
         ax_client = AxClient(generation_strategy=gs, verbose_logging=False)
         ax_client.create_experiment(

pyproject.toml

@@ -34,11 +34,11 @@ test = [
     'flake8',
     'pytest',
     'pytest-mpi',
-    'ax-platform == 0.4.0',
+    'ax-platform >=0.5.0',
     'matplotlib',
 ]
 all = [
-    'ax-platform == 0.4.0',
+    'ax-platform >=0.5.0',
     'matplotlib'
 ]