Handle discrete variables in gen base class

Author: shuds13

optimas/generators/ax/developer/multitask.py

@@ -205,11 +205,6 @@ def __init__(
         ]
         self._check_inputs(vocs, lofi_task, hifi_task)
 
-        # Convert discrete variables to trial parameters before calling super().__init__
-        custom_trial_parameters.extend(
-            self._convert_discrete_variables_to_trial_parameters(vocs)
-        )
-
         super().__init__(
             vocs=vocs,
             use_cuda=use_cuda,
@@ -234,21 +229,6 @@ def __init__(
         self.gr_lofi = None
         self._experiment = self._create_experiment()
 
-    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:

optimas/generators/base.py

@@ -105,6 +105,12 @@ def __init__(
         self._custom_trial_parameters = (
             [] if custom_trial_parameters is None else custom_trial_parameters
         )
+
+        # Automatically add discrete variables as trial parameters
+        discrete_trial_params = (
+            self._convert_vocs_discrete_variables_to_trial_parameters()
+        )
+        self._custom_trial_parameters.extend(discrete_trial_params)
         self._allow_fixed_parameters = allow_fixed_parameters
         self._allow_updating_parameters = allow_updating_parameters
         self._gen_function = persistent_generator
@@ -191,6 +197,31 @@ def _convert_vocs_observables_to_parameters(self) -> List[Parameter]:
             parameters.append(param)
         return parameters
 
+    def _convert_vocs_discrete_variables_to_trial_parameters(
+        self,
+    ) -> List[TrialParameter]:
+        """Convert discrete variables from VOCS to TrialParameter objects.
+
+        Only converts discrete variables that were NOT already converted to
+        VaryingParameters.
+        """
+        trial_parameters = []
+        # Get the names of variables that were already converted to
+        # VaryingParameters
+        varying_param_names = {vp.name for vp in self._varying_parameters}
+
+        for var_name, var_spec in self._vocs.variables.items():
+            if isinstance(var_spec, DiscreteVariable):
+                # Only convert if it wasn't already converted to a
+                # VaryingParameter
+                if var_name not in varying_param_names:
+                    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
+
     @property
     def vocs(self) -> VOCS:
         """Get the VOCS object."""