[WIP] Implement VOCS interface

.github/workflows/unix-noax.yml

@@ -31,6 +31,7 @@ jobs:
         conda install numpy pandas pytorch cpuonly -c pytorch
         conda install -c conda-forge mpi4py mpich
         pip install .[test]
+        pip install git+https://github.com/campa-consortium/generator_standard.git@obs_type
         pip uninstall --yes ax-platform # Run without Ax
     - shell: bash -l {0}
       name: Run unit tests without Ax

.github/workflows/unix-openmpi.yml

@@ -31,6 +31,7 @@ jobs:
         conda install numpy pandas pytorch cpuonly -c pytorch
         conda install -c conda-forge mpi4py openmpi=5.*
         pip install .[test]
+        pip install git+https://github.com/campa-consortium/generator_standard.git@obs_type
     - shell: bash -l {0}
       name: Run unit tests with openMPI
       run: |

.github/workflows/unix.yml

@@ -31,6 +31,7 @@ jobs:
         conda install numpy pandas pytorch cpuonly -c pytorch
         conda install -c conda-forge mpi4py mpich
         pip install .[test]
+        pip install git+https://github.com/campa-consortium/generator_standard.git@obs_type
     - shell: bash -l {0}
       name: Run unit tests with MPICH
       run: |

optimas/diagnostics/exploration_diagnostics.py

@@ -19,6 +19,7 @@
 from optimas.explorations import Exploration
 from optimas.utils.other import get_df_with_selection
 from optimas.utils.ax import AxModelManager
+from generator_standard.vocs import VOCS
 
 
 class ExplorationDiagnostics:
@@ -102,12 +103,27 @@ def _create_exploration(
                 analyzed_parameters.append(p)
 
         # Create exploration using dummy generator and evaluator.
+        # Create exploration using dummy generator and evaluator.
+        variables = {}
+        for vp in varying_parameters:
+            variables[vp.name] = [vp.lower_bound, vp.upper_bound]
+
+        vocs_objectives = {}
+        for obj in objectives:
+            vocs_objectives[obj.name] = (
+                "MINIMIZE" if obj.minimize else "MAXIMIZE"
+            )
+
+        observables = [param.name for param in analyzed_parameters]
+
+        vocs = VOCS(
+            variables=variables,
+            objectives=vocs_objectives,
+            observables=observables,
+        )
+
         return Exploration(
-            generator=Generator(
-                varying_parameters=varying_parameters,
-                objectives=objectives,
-                analyzed_parameters=analyzed_parameters,
-            ),
+            generator=Generator(vocs=vocs),
             evaluator=Evaluator(sim_function=None),
             history=history_path,
             exploration_dir_path=exploration_dir_path,

optimas/generators/ax/base.py

@@ -1,12 +1,13 @@
 """Contains the definition of the base Ax generator."""
 
-from typing import List, Optional
+from typing import Optional
 import logging
 
 import torch
 
 from optimas.core import Objective, TrialParameter, VaryingParameter, Parameter
 from optimas.generators.base import Generator
+from generator_standard.vocs import VOCS
 
 
 # Disable Ax loggers to get cleaner output. In principle, setting
@@ -22,13 +23,8 @@ class AxGenerator(Generator):
 
     Parameters
     ----------
-    varying_parameters : list of VaryingParameter
-        List of input parameters to vary.
-    objectives : list of Objective
-        List of optimization objectives.
-    analyzed_parameters : list of Parameter, optional
-        List of parameters to analyze at each trial, but which are not
-        optimization objectives. By default ``None``.
+    vocs : VOCS
+        VOCS object defining variables, objectives, constraints, and observables.
     use_cuda : bool, optional
         Whether to allow the generator to run on a CUDA GPU. By default
         ``False``.
@@ -63,9 +59,7 @@ class AxGenerator(Generator):
 
     def __init__(
         self,
-        varying_parameters: List[VaryingParameter],
-        objectives: List[Objective],
-        analyzed_parameters: Optional[List[Parameter]] = None,
+        vocs: VOCS,
         use_cuda: Optional[bool] = False,
         gpu_id: Optional[int] = 0,
         dedicated_resources: Optional[bool] = False,
@@ -77,9 +71,7 @@ def __init__(
         allow_updating_parameters: Optional[bool] = False,
     ) -> None:
         super().__init__(
-            varying_parameters=varying_parameters,
-            objectives=objectives,
-            analyzed_parameters=analyzed_parameters,
+            vocs=vocs,
             use_cuda=use_cuda,
             gpu_id=gpu_id,
             dedicated_resources=dedicated_resources,

optimas/generators/ax/developer/multitask.py

@@ -67,6 +67,7 @@
     TrialStatus,
 )
 from .ax_metric import AxMetric
+from generator_standard.vocs import VOCS, DiscreteVariable
 
 # Define generator states.
 NOT_STARTED = "not_started"
@@ -152,10 +153,8 @@ class AxMultitaskGenerator(AxGenerator):
 
     Parameters
     ----------
-    varying_parameters : list of VaryingParameter
-        List of input parameters to vary. One them should be a fidelity.
-    objectives : list of Objective
-        List of optimization objectives. Only one objective is supported.
+    vocs : VOCS
+        VOCS object defining variables, objectives, constraints, and observables.
     lofi_task, hifi_task : Task
         The low- and high-fidelity tasks.
     analyzed_parameters : list of Parameter, optional
@@ -184,31 +183,25 @@ class AxMultitaskGenerator(AxGenerator):
 
     def __init__(
         self,
-        varying_parameters: List[VaryingParameter],
-        objectives: List[Objective],
+        vocs: VOCS,
         lofi_task: Task,
         hifi_task: Task,
-        analyzed_parameters: Optional[List[Parameter]] = None,
         use_cuda: Optional[bool] = False,
         gpu_id: Optional[int] = 0,
         dedicated_resources: Optional[bool] = False,
         save_model: Optional[bool] = True,
         model_save_period: Optional[int] = 5,
         model_history_dir: Optional[str] = "model_history",
     ) -> None:
-        # As trial parameters these get written to history array
-        # Ax trial_index and arm toegther locate a point
-        # Multiple points (Optimas trials) can share the same Ax trial_index
-        custom_trial_parameters = [
-            TrialParameter("arm_name", "ax_arm_name", dtype="U32"),
-            TrialParameter("trial_type", "ax_trial_type", dtype="U32"),
-            TrialParameter("ax_trial_id", "ax_trial_index", dtype=int),
-        ]
-        self._check_inputs(varying_parameters, objectives, lofi_task, hifi_task)
+        self._check_inputs(vocs, lofi_task, hifi_task)
+
+        # Convert discrete variables to trial parameters before calling super().__init__
+        custom_trial_parameters = (
+            self._convert_discrete_variables_to_trial_parameters(vocs)
+        )
+
         super().__init__(
-            varying_parameters=varying_parameters,
-            objectives=objectives,
-            analyzed_parameters=analyzed_parameters,
+            vocs=vocs,
             use_cuda=use_cuda,
             gpu_id=gpu_id,
             dedicated_resources=dedicated_resources,
@@ -231,6 +224,25 @@ def __init__(
         self.gr_lofi = None
         self._experiment = self._create_experiment()
 
+        # Internal mapping: _id -> (arm_name, ax_trial_id, trial_type)
+        self._id_mapping = {}
+        self._next_id = 0
+
+    def _convert_discrete_variables_to_trial_parameters(
+        self, vocs: VOCS
+    ) -> List[TrialParameter]:
+        """Convert discrete variables from VOCS to TrialParameter objects."""
+        trial_parameters = []
+        for var_name, var_spec in vocs.variables.items():
+            if isinstance(var_spec, DiscreteVariable):
+                # Convert discrete variable to trial parameter
+                max_len = max(len(str(val)) for val in var_spec.values)
+                trial_param = TrialParameter(
+                    var_name, var_name, dtype=f"U{max_len}"
+                )
+                trial_parameters.append(trial_param)
+        return trial_parameters
+
     def get_gen_specs(
         self, sim_workers: int, run_params: Dict, sim_max: int
     ) -> Dict:
@@ -242,20 +254,30 @@ def get_gen_specs(
         gen_specs["out"].append(("task", str, max_length))
         return gen_specs
 
+    def _validate_vocs(self, vocs: VOCS) -> None:
+        """Validate VOCS for multitask generator."""
+        super()._validate_vocs(vocs)
+        # Check that only one objective has been given.
+        n_objectives = len(vocs.objectives)
+        assert n_objectives == 1, (
+            "Multitask generator supports only a single objective. "
+            "Objectives given: {}.".format(n_objectives)
+        )
+        # Check that there is a discrete variable called 'trial_type'
+        assert (
+            "trial_type" in vocs.variables
+        ), "Multitask generator requires a discrete variable named 'trial_type'"
+        assert isinstance(
+            vocs.variables["trial_type"], DiscreteVariable
+        ), "Variable 'trial_type' must be a discrete variable"
+
     def _check_inputs(
         self,
-        varying_parameters: List[VaryingParameter],
-        objectives: List[Objective],
+        vocs: VOCS,
         lofi_task: Task,
         hifi_task: Task,
     ) -> None:
         """Check that the generator inputs are valid."""
-        # Check that only one objective has been given.
-        n_objectives = len(objectives)
-        assert n_objectives == 1, (
-            "Multitask generator supports only a single objective. "
-            "Objectives given: {}.".format(n_objectives)
-        )
         # Check that the number of low-fidelity trials per iteration is larger
         # than that of high-fidelity trials.
         assert lofi_task.n_opt >= hifi_task.n_opt, (
@@ -274,11 +296,21 @@ def suggest(self, num_points: Optional[int]) -> List[dict]:
                     var.name: arm.parameters.get(var.name)
                     for var in self._varying_parameters
                 }
-                # SH for VOCS standard these will need to be 'variables'
-                # For now much match the trial parameter names.
-                point["ax_trial_id"] = trial_index
-                point["arm_name"] = arm.name
-                point["trial_type"] = trial_type
+                # SH We can use a discrete var here in vocs (converted for now to trial parameters)
+                # But unlike varying parameters the name refers to a fixed generator concept.
+                for trial_param in self._custom_trial_parameters:
+                    if trial_param.name == "trial_type":
+                        point[trial_param.name] = trial_type
+
+                # Generate unique _id and store mapping
+                current_id = self._next_id
+                self._id_mapping[current_id] = {
+                    "arm_name": arm.name,
+                    "ax_trial_id": trial_index,
+                    "trial_type": trial_type,
+                }
+                point["_id"] = current_id
+                self._next_id += 1
                 points.append(point)
         return points
 
@@ -295,6 +327,15 @@ def ingest(self, results: List[dict]) -> None:
                 custom_parameters=self._custom_trial_parameters,
             )
             trials.append(trial)
+
+        # Apply _id mapping to all trials before processing
+        for trial in trials:
+            if trial.gen_id is not None and trial.gen_id in self._id_mapping:
+                mapping = self._id_mapping[trial.gen_id]
+                trial.arm_name = mapping["arm_name"]
+                trial.ax_trial_id = mapping["ax_trial_id"]
+                # trial_type should already be in trial from custom_parameters
+
         if self.gen_state == NOT_STARTED:
             self._incorporate_external_data(trials)
         else:

optimas/generators/ax/service/ax_client.py

@@ -6,6 +6,7 @@
 from ax.core.objective import MultiObjective
 
 from optimas.core import Objective, VaryingParameter, Parameter
+from generator_standard.vocs import VOCS
 from .base import AxServiceGenerator
 
 
@@ -71,18 +72,19 @@ def __init__(
         model_save_period: Optional[int] = 5,
         model_history_dir: Optional[str] = "model_history",
     ):
-        varying_parameters = self._get_varying_parameters(ax_client)
-        objectives = self._get_objectives(ax_client)
+        # Create VOCS object from AxClient data
+        vocs = self._create_vocs_from_ax_client(ax_client)
+
+        # Add constraints to analyzed parameters
         analyzed_parameters = self._add_constraints_to_analyzed_parameters(
             analyzed_parameters, ax_client
         )
+
         use_cuda = self._use_cuda(ax_client)
         self._ax_client = ax_client
+
         super().__init__(
-            varying_parameters=varying_parameters,
-            objectives=objectives,
-            analyzed_parameters=analyzed_parameters,
-            enforce_n_init=True,
+            vocs=vocs,
             abandon_failed_trials=abandon_failed_trials,
             use_cuda=use_cuda,
             gpu_id=gpu_id,
@@ -92,6 +94,38 @@ def __init__(
             model_history_dir=model_history_dir,
         )
 
+    def _create_vocs_from_ax_client(self, ax_client: AxClient) -> VOCS:
+        """Create a VOCS object from the AxClient data."""
+        # Extract variables from search space
+        variables = {}
+        for _, p in ax_client.experiment.search_space.parameters.items():
+            variables[p.name] = [p.lower, p.upper]
+
+        # Extract objectives from optimization config
+        objectives = {}
+        ax_objective = ax_client.experiment.optimization_config.objective
+        if isinstance(ax_objective, MultiObjective):
+            ax_objectives = ax_objective.objectives
+        else:
+            ax_objectives = [ax_objective]
+
+        for ax_obj in ax_objectives:
+            obj_type = "MINIMIZE" if ax_obj.minimize else "MAXIMIZE"
+            objectives[ax_obj.metric_names[0]] = obj_type
+
+        # Extract observables from outcome constraints (if any)
+        observables = set()
+        ax_config = ax_client.experiment.optimization_config
+        if ax_config.outcome_constraints:
+            for constraint in ax_config.outcome_constraints:
+                observables.add(constraint.metric.name)
+
+        return VOCS(
+            variables=variables,
+            objectives=objectives,
+            observables=observables,
+        )
+
     def _get_varying_parameters(self, ax_client: AxClient):
         """Obtain the list of varying parameters from the AxClient."""
         varying_parameters = []