optimeo.bo

   1# Copyright (c) 2025 Colin BOUSIGE
   2# Contact: colin.bousige@cnrs.fr
   3#
   4# This program is free software: you can redistribute it and/or modify
   5# it under the terms of the MIT License as published by
   6# the Free Software Foundation, either version 3 of the License, or
   7# any later version. 
   8
   9"""
  10This module provides a class for optimizing experiments using Bayesian Optimization (BO) with the [Ax platform](https://ax.dev/).
  11It includes methods for initializing the experiment, suggesting trials, predicting outcomes, and plotting results.
  12
  13You can see an example notebook [here](../examples/bo.ipynb).
  14
  15"""
  16
  17import warnings
  18warnings.simplefilter(action='ignore', category=FutureWarning)
  19warnings.simplefilter(action='ignore', category=DeprecationWarning)
  20warnings.simplefilter(action='ignore', category=UserWarning)
  21warnings.simplefilter(action='ignore', category=RuntimeError)
  22
  23import numpy as np
  24import pandas as pd
  25from janitor import clean_names
  26from typing import Any, Dict, List, Optional, Union
  27
  28from ax.core.observation import ObservationFeatures, TrialStatus
  29from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
  30from ax.modelbridge.modelbridge_utils import get_pending_observation_features
  31from ax.modelbridge.registry import Models
  32from ax.plot.contour import interact_contour, plot_contour
  33from ax.plot.pareto_frontier import plot_pareto_frontier
  34from ax.plot.pareto_utils import compute_posterior_pareto_frontier
  35from ax.plot.slice import plot_slice
  36from ax.service.ax_client import AxClient, ObjectiveProperties
  37from botorch.acquisition.analytic import *
  38import plotly.graph_objects as go
  39
  40# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
  41
  42class BOExperiment:
  43    """
  44    BOExperiment is a class designed to facilitate Bayesian Optimization experiments using the [Ax platform](https://ax.dev/).
  45    It encapsulates the experiment setup, including features, outcomes, constraints, and optimization methods.
  46    
  47    Parameters
  48    ----------
  49    features: Dict[str, Dict[str, Any]]
  50        A dictionary defining the features of the experiment, including their types and ranges.
  51        Each feature is represented as a dictionary with keys 'type', 'data', and 'range'.
  52        - 'type': The type of the feature (e.g., 'int', 'float', 'text').
  53        - 'data': The observed data for the feature.
  54        - 'range': The range of values for the feature.
  55    outcomes: Dict[str, Dict[str, Any]]
  56        A dictionary defining the outcomes of the experiment, including their types and observed data.
  57        Each outcome is represented as a dictionary with keys 'type' and 'data'.
  58        - 'type': The type of the outcome (e.g., 'int', 'float').
  59        - 'data': The observed data for the outcome.
  60    ranges: Optional[Dict[str, Dict[str, Any]]]
  61        A dictionary defining the ranges of the features. Default is `None`.
  62        If not provided, the ranges will be inferred from the features data.
  63        The ranges should be in the format `{'feature_name': [minvalue,maxvalue]}`.
  64    N: int
  65        The number of trials to suggest in each optimization step. Must be a positive integer.
  66    maximize: Union[bool, Dict[str, bool]]
  67        A boolean or dict indicating whether to maximize the outcomes in the form `{'outcome1':True, 'outcome2':False}`.
  68        If a single boolean is provided, it is applied to all outcomes. Default is `True`.
  69    fixed_features: Optional[Dict[str, Any]]
  70        A dictionary defining fixed features with their values. Default is `None`.
  71        If provided, the fixed features will be treated as fixed parameters in the generation process.
  72        The fixed features should be in the format `{'feature_name': value}`.
  73        The values should be the fixed values for the respective features.
  74    outcome_constraints: Optional[List[str]]
  75        Constraints on the outcomes, specified as a list of strings. Default is `None`.
  76        The constraints should be in the format `{'outcome_name': [minvalue,maxvalue]}`.
  77    feature_constraints: Optional[List[str]]
  78        Constraints on the features, specified as a list of strings. Default is `None`.
  79        The constraints should be in the format `{'feature_name': [minvalue,maxvalue]}`.
  80    optim: str
  81        The optimization method to use, either 'bo' for Bayesian Optimization or 'sobol' for Sobol sequence. Default is 'bo'.
  82    acq_func: Optional[Dict[str, Any]]
  83        The acquisition function to use for the optimization process. It must be a dict with 2 keys:
  84        - `acqf`: the acquisition function class to use (e.g., `UpperConfidenceBound`),
  85        - `acqf_kwargs`: a dict of the kwargs to pass to the acquisition function class. (e.g. `{'beta': 0.1}`).
  86        
  87        If not provided, the default acquisition function is used (`LogExpectedImprovement` or `qLogExpectedImprovement` if N>1).
  88    
  89    Attributes
  90    ----------
  91    
  92    features: Dict[str, Dict[str, Any]]
  93        A dictionary defining the features of the experiment, including their types and ranges.
  94    outcomes: Dict[str, Dict[str, Any]]
  95        A dictionary defining the outcomes of the experiment, including their types and observed data.
  96    N: int
  97        The number of trials to suggest in each optimization step. Must be a positive integer.
  98    maximize: Union[bool, List[bool]]
  99        A boolean or list of booleans indicating whether to maximize the outcomes.
 100        If a single boolean is provided, it is applied to all outcomes.
 101    outcome_constraints: Optional[Dict[str, Dict[str, float]]]
 102        Constraints on the outcomes, specified as a dictionary or list of dictionaries.
 103    feature_constraints: Optional[List[Dict[str, Any]]]
 104        Constraints on the features, specified as a list of dictionaries.
 105    optim: str
 106        The optimization method to use, either 'bo' for Bayesian Optimization or 'sobol' for Sobol sequence.
 107    data: pd.DataFrame
 108        A DataFrame representing the current data in the experiment, including features and outcomes.
 109    acq_func: dict
 110        The acquisition function to use for the optimization process. 
 111    generator_run:
 112        The generator run for the experiment, used to generate new candidates.
 113    model:
 114        The model used for predictions in the experiment.
 115    ax_client:
 116        The AxClient for the experiment, used to manage trials and data.
 117    gs:
 118        The generation strategy for the experiment, used to generate new candidates.
 119    parameters:
 120        The parameters for the experiment, including their types and ranges.
 121    names:
 122        The names of the features in the experiment.
 123    fixed_features:
 124        The fixed features for the experiment, used to generate new candidates.
 125    candidate:
 126        The candidate(s) suggested by the optimization process.
 127        
 128
 129    Methods
 130    -------
 131    
 132    - <b>initialize_ax_client()</b>:
 133        Initializes the AxClient with the experiment's parameters, objectives, and constraints.
 134    - <b>suggest_next_trials()</b>:
 135        Suggests the next set of trials based on the current model and optimization strategy.
 136        Returns a DataFrame containing the suggested trials and their predicted outcomes.
 137    - <b>predict(params: List[Dict[str, Any]]) -> List[Dict[str, float]]</b>:
 138        Predicts the outcomes for a given set of parameters using the current model.
 139        Returns a list of predicted outcomes for the given parameters.
 140    - <b>update_experiment(params: Dict[str, Any], outcomes: Dict[str, Any])</b>:
 141        Updates the experiment with new parameters and outcomes, and reinitializes the AxClient.
 142    - <b>plot_model(metricname: Optional[str] = None, slice_values: Optional[Dict[str, Any]]  None, linear: bool = False)`</b>:
 143        Plots the model's predictions for the experiment's parameters and outcomes.
 144        If metricname is None, the first outcome metric is used.
 145        If slice_values is provided, it slices the plot at those values.
 146        If linear is True, it plots a linear slice plot.
 147        If the experiment has only one feature, it plots a slice plot.
 148        If the experiment has multiple features, it plots a contour plot.
 149        Returns a Plotly figure of the model's predictions.
 150    - <b>plot_optimization_trace(optimum: Optional[float] = None)</b>:
 151        Plots the optimization trace, showing the progress of the optimization over trials.
 152        If the experiment has multiple outcomes, it raises a warning and returns None.
 153        Returns a Plotly figure of the optimization trace.
 154    - <b>plot_pareto_frontier()</b>:
 155        Plots the Pareto frontier for multi-objective optimization experiments.
 156        If the experiment has only one outcome, it raises a warning and returns None.
 157        Returns a Plotly figure of the Pareto frontier.
 158    - <b>get_best_parameters() -> pd.DataFrame</b>:
 159        Returns the best parameters found by the optimization process.
 160        If the experiment has multiple outcomes, it returns a DataFrame of the Pareto optimal parameters.
 161        If the experiment has only one outcome, it returns a DataFrame of the best parameters and their outcomes.
 162        The DataFrame contains the best parameters and their corresponding outcomes.
 163    - <b>clear_trials()</b>:
 164        Clears all trials in the experiment.
 165        This is useful for resetting the experiment before suggesting new trials.
 166    - <b>set_model()</b>:
 167        Sets the model to be used for predictions.
 168        This method is called after initializing the AxClient.
 169    - <b>set_gs()</b>:
 170        Sets the generation strategy for the experiment.
 171        This method is called after initializing the AxClient.
 172
 173
 174    Example
 175    -------
 176    ```python
 177    features, outcomes = read_experimental_data('data.csv', out_pos=[-2, -1])
 178    experiment = BOExperiment(features, 
 179                              outcomes, 
 180                              N=5, 
 181                              maximize={'out1':True, 'out2':False}
 182                              )
 183    experiment.suggest_next_trials()
 184    experiment.plot_model(metricname='outcome1')
 185    experiment.plot_model(metricname='outcome2', linear=True)
 186    experiment.plot_model(metricname='outcome1', slice_values={'feature1': 5})
 187    experiment.plot_optimization_trace()
 188    experiment.plot_pareto_frontier()
 189    experiment.get_best_parameters()
 190    experiment.update_experiment({'feature1': [4]}, {'outcome1': [0.4]})
 191    experiment.plot_model()
 192    experiment.plot_optimization_trace()
 193    experiment.plot_pareto_frontier()
 194    experiment.get_best_parameters()
 195    ```
 196    """
 197
 198    def __init__(self,
 199                 features: Dict[str, Dict[str, Any]],
 200                 outcomes: Dict[str, Dict[str, Any]],
 201                 ranges: Optional[Dict[str, Dict[str, Any]]] = None,
 202                 N=1,
 203                 maximize: Union[bool, Dict[str, bool]] = True,
 204                 fixed_features: Optional[Dict[str, Any]] = None,
 205                 outcome_constraints: Optional[List[str]] = None,
 206                 feature_constraints: Optional[List[str]] = None,
 207                 optim='bo',
 208                 acq_func=None) -> None:
 209        self._first_initialization_done = False
 210        self.ranges              = ranges
 211        self.features            = features
 212        self.names               = list(self._features.keys())
 213        self.fixed_features      = fixed_features
 214        self.outcomes            = outcomes
 215        self.N                   = N
 216        self.maximize            = maximize
 217        self.outcome_constraints = outcome_constraints
 218        self.feature_constraints = feature_constraints
 219        self.optim               = optim
 220        self.acq_func            = acq_func
 221        self.candidate = None
 222        """The candidate(s) suggested by the optimization process."""
 223        self.ax_client = None
 224        """Ax's client for the experiment."""
 225        self.model = None
 226        """Ax's Gaussian Process model."""
 227        self.parameters = None
 228        """Ax's parameters for the experiment."""
 229        self.generator_run = None
 230        """Ax's generator run for the experiment."""
 231        self.gs = None
 232        """Ax's generation strategy for the experiment."""
 233        self.initialize_ax_client()
 234        self.Nmetrics = len(self.ax_client.objective_names)
 235        """The number of metrics in the experiment."""
 236        self._first_initialization_done = True
 237        """To indicate that the first initialization is done so that we don't call `initialize_ax_client()` again."""
 238        self.pareto_frontier = None
 239        """The Pareto frontier for multi-objective optimization experiments."""
 240
 241    @property
 242    def features(self):
 243        """
 244        A dictionary defining the features of the experiment, including their types and ranges.
 245        
 246        Example
 247        -------
 248        ```python
 249        features = {
 250            'feature1': {'type': 'int', 
 251                         'data': [1, 2, 3], 
 252                         'range': [1, 3]},
 253            'feature2': {'type': 'float', 
 254                         'data': [0.1, 0.2, 0.3], 
 255                         'range': [0.1, 0.3]},
 256            'feature3': {'type': 'text', 
 257                         'data': ['A', 'B', 'C'], 
 258                         'range': ['A', 'B', 'C']}
 259            }
 260        ```
 261        """
 262        return self._features
 263
 264    @features.setter
 265    def features(self, value):
 266        """
 267        Set the features of the experiment with validation.
 268        """
 269        if not isinstance(value, dict):
 270            raise ValueError("features must be a dictionary")
 271        self._features = value
 272        for name in self._features.keys():
 273            if self.ranges and name in self.ranges.keys():
 274                self._features[name]['range'] = self.ranges[name]
 275            else:
 276                if self._features[name]['type'] == 'text':
 277                    self._features[name]['range'] = list(set(self._features[name]['data']))
 278                elif self._features[name]['type'] == 'int':
 279                    self._features[name]['range'] = [int(np.min(self._features[name]['data'])),
 280                                                     int(np.max(self._features[name]['data']))]
 281                elif self._features[name]['type'] == 'float':
 282                    self._features[name]['range'] = [float(np.min(self._features[name]['data'])),
 283                                                     float(np.max(self._features[name]['data']))]
 284        if self._first_initialization_done:
 285            self.initialize_ax_client()
 286    
 287    @property
 288    def ranges(self):
 289        """
 290        A dictionary defining the ranges of the features. Default is `None`.
 291        
 292        If not provided, the ranges will be inferred from the features data.
 293        The ranges should be in the format `{'feature_name': [minvalue,maxvalue]}`.
 294        """
 295        return self._ranges
 296
 297    @ranges.setter
 298    def ranges(self, value):
 299        """
 300        Set the ranges of the features with validation.
 301        """
 302        if value is not None:
 303            if not isinstance(value, dict):
 304                raise ValueError("ranges must be a dictionary")
 305        self._ranges = value
 306    
 307    @property
 308    def names(self):
 309        """
 310        The names of the features.
 311        """
 312        return self._names
 313    
 314    @names.setter
 315    def names(self, value):
 316        """
 317        Set the names of the features.
 318        """
 319        if not isinstance(value, list):
 320            raise ValueError("names must be a list")
 321        self._names = value
 322
 323    @property
 324    def outcomes(self):
 325        """
 326        A dictionary defining the outcomes of the experiment, including their types and observed data.
 327        
 328        Example
 329        -------
 330        ```python
 331        outcomes = {
 332            'outcome1': {'type': 'float', 
 333                         'data': [0.1, 0.2, 0.3]},
 334            'outcome2': {'type': 'float', 
 335                         'data': [1.0, 2.0, 3.0]}
 336            }
 337        ```
 338        """
 339        return self._outcomes
 340
 341    @outcomes.setter
 342    def outcomes(self, value):
 343        """
 344        Set the outcomes of the experiment with validation.
 345        """
 346        if not isinstance(value, dict):
 347            raise ValueError("outcomes must be a dictionary")
 348        self._outcomes = value
 349        self.out_names = list(value.keys())
 350        if self._first_initialization_done:
 351            self.initialize_ax_client()
 352    
 353    @property
 354    def fixed_features(self):
 355        """
 356        A dictionary defining fixed features with their values. Default is `None`.
 357        If provided, the fixed features will be treated as fixed parameters in the generation process.
 358        The fixed features should be in the format `{'feature_name': value}`.
 359        The values should be the fixed values for the respective features.
 360        """
 361        return self._fixed_features
 362
 363    @fixed_features.setter
 364    def fixed_features(self, value):
 365        """
 366        Set the fixed features of the experiment.
 367        """
 368        self._fixed_features = None
 369        if value is not None:
 370            if not isinstance(value, dict):
 371                raise ValueError("fixed_features must be a dictionary")
 372            for name in value.keys():
 373                if name not in self.names:
 374                    raise ValueError(f"Fixed feature '{name}' not found in features")
 375            # fixed_features should be an ObservationFeatures object
 376            self._fixed_features = ObservationFeatures(parameters=value)
 377        if self._first_initialization_done:
 378            self.set_gs()
 379
 380    @property
 381    def N(self):
 382        """
 383        The number of trials to suggest in each optimization step. Must be a positive integer. Default is `1`.
 384        """
 385        return self._N
 386
 387    @N.setter
 388    def N(self, value):
 389        """
 390        Set the number of trials to suggest in each optimization step with validation.
 391        """
 392        if not isinstance(value, int) or value <= 0:
 393            raise ValueError("N must be a positive integer")
 394        self._N = value
 395        if self._first_initialization_done:
 396            self.set_gs()
 397
 398    @property
 399    def maximize(self):
 400        """
 401        A boolean or dict indicating whether to maximize the outcomes in the form `{'outcome1':True, 'outcome2':False}`.
 402        If a single boolean is provided, it is applied to all outcomes. Default is `True`.
 403        """
 404        return self._maximize
 405
 406    @maximize.setter
 407    def maximize(self, value):
 408        """
 409        Set the maximization setting for the outcomes with validation.
 410        """
 411        if isinstance(value, bool):
 412            self._maximize = {out: value for out in self.out_names}
 413        elif isinstance(value, dict) and len(value) == len(self._outcomes):
 414            self._maximize = {k:v for k,v in value.items() if 
 415                              (k in self.out_names and isinstance(v, bool))}
 416        else:
 417            raise ValueError("maximize must be a boolean or a list of booleans with the same length as outcomes")
 418        if self._first_initialization_done:
 419            self.initialize_ax_client()
 420
 421    @property
 422    def outcome_constraints(self):
 423        """
 424        Constraints on the outcomes, specified as a list of strings. Default is `None`.
 425        """
 426        return self._outcome_constraints
 427
 428    @outcome_constraints.setter
 429    def outcome_constraints(self, value):
 430        """
 431        Set the outcome constraints of the experiment with validation.
 432        """
 433        if isinstance(value, str):
 434            self._outcome_constraints = [value]
 435        elif isinstance(value, list):
 436            self._outcome_constraints = value
 437        else:
 438            self._outcome_constraints = None
 439        if self._first_initialization_done:
 440            self.initialize_ax_client()
 441
 442    @property
 443    def feature_constraints(self):
 444        """
 445        Constraints on the features, specified as a list of strings. Default is `None`.
 446        
 447        Example
 448        -------
 449        ```python
 450        feature_constraints = [
 451            'feature1 <= 10.0',
 452            'feature1 + 2*feature2 >= 3.0'
 453        ]
 454        ```
 455        """
 456        return self._feature_constraints
 457
 458    @feature_constraints.setter
 459    def feature_constraints(self, value):
 460        """
 461        Set the feature constraints of the experiment with validation.
 462        """
 463        if isinstance(value, dict):
 464            self._feature_constraints = [value]
 465        elif isinstance(value, list):
 466            self._feature_constraints = value
 467        elif isinstance(value, str):
 468            self._feature_constraints = [value]
 469        else:
 470            self._feature_constraints = None
 471        if self._first_initialization_done:
 472            self.initialize_ax_client()
 473
 474    @property
 475    def optim(self):
 476        """
 477        The optimization method to use, either `'bo'` for Bayesian Optimization or `'sobol'` for Sobol sequence. Default is `'bo'`.
 478        """
 479        return self._optim
 480
 481    @optim.setter
 482    def optim(self, value):
 483        """
 484        Set the optimization method with validation.
 485        """
 486        value = value.lower()
 487        if value not in ['bo', 'sobol']:
 488            raise ValueError("Optimization method must be either 'bo' or 'sobol'")
 489        self._optim = value
 490        if self._first_initialization_done:
 491            self.set_gs()
 492
 493    @property
 494    def data(self) -> pd.DataFrame:
 495        """
 496        Returns a DataFrame of the current data in the experiment, including features and outcomes.
 497        """
 498        feature_data = {name: info['data'] for name, info in self._features.items()}
 499        outcome_data = {name: info['data'] for name, info in self._outcomes.items()}
 500        data_dict = {**feature_data, **outcome_data}
 501        return pd.DataFrame(data_dict)
 502
 503    @data.setter
 504    def data(self, value: pd.DataFrame):
 505        """
 506        Sets the features and outcomes data from a given DataFrame.
 507        """
 508        if not isinstance(value, pd.DataFrame):
 509            raise ValueError("Data must be a pandas DataFrame")
 510
 511        feature_columns = [col for col in value.columns if col in self._features]
 512        outcome_columns = [col for col in value.columns if col in self._outcomes]
 513
 514        for col in feature_columns:
 515            self._features[col]['data'] = value[col].tolist()
 516
 517        for col in outcome_columns:
 518            self._outcomes[col]['data'] = value[col].tolist()
 519
 520        if self._first_initialization_done:
 521            self.initialize_ax_client()
 522
 523    @property
 524    def pareto_frontier(self):
 525        """
 526        The Pareto frontier for multi-objective optimization experiments.
 527        """
 528        return self._pareto_frontier
 529    
 530    @pareto_frontier.setter
 531    def pareto_frontier(self, value):
 532        """
 533        Set the Pareto frontier of the experiment.
 534        """
 535        self._pareto_frontier = value
 536        
 537    
 538    @property
 539    def acq_func(self):
 540        """
 541        The acquisition function to use for the optimization process. It must be a dict with 2 keys:
 542        - `acqf`: the acquisition function class to use (e.g., `UpperConfidenceBound`),
 543        - `acqf_kwargs`: a dict of the kwargs to pass to the acquisition function class. (e.g. `{'beta': 0.1}`).
 544        
 545        If not provided, the default acquisition function is used (`LogExpectedImprovement` or `qLogExpectedImprovement` if N>1).
 546        
 547        Example
 548        -------
 549        ```python
 550        acq_func = {
 551            'acqf': UpperConfidenceBound,
 552            'acqf_kwargs': {'beta': 0.1} # lower value = exploitation, higher value = exploration
 553        }
 554        ```
 555        """
 556        return self._acq_func
 557    
 558    @acq_func.setter
 559    def acq_func(self, value):
 560        """
 561        Set the acquisition function with validation.
 562        """
 563        self._acq_func = value
 564        if self._first_initialization_done:
 565            self.set_gs()
 566
 567    def __repr__(self):
 568        return self.__str__()
 569
 570    def __str__(self):
 571        """
 572        Return a string representation of the BOExperiment instance.
 573        """
 574        return f"""
 575BOExperiment(
 576    N={self.N},
 577    maximize={self.maximize},
 578    outcome_constraints={self.outcome_constraints},
 579    feature_constraints={self.feature_constraints},
 580    optim={self.optim}
 581)
 582
 583Input data:
 584
 585{self.data}
 586        """
 587
 588    def initialize_ax_client(self):
 589        """
 590        Initialize the AxClient with the experiment's parameters, objectives, and constraints.
 591        """
 592        print('\n========   INITIALIZING MODEL   ========\n')
 593        self.ax_client = AxClient(verbose_logging=False, 
 594                                  suppress_storage_errors=True)
 595        self.parameters = []
 596        for name, info in self._features.items():
 597            if info['type'] == 'text':
 598                self.parameters.append({
 599                    "name": name,
 600                    "type": "choice",
 601                    "values": [str(val) for val in info['range']],
 602                    "value_type": "str"})
 603            elif info['type'] == 'int':
 604                self.parameters.append({
 605                    "name": name,
 606                    "type": "range",
 607                    "bounds": [int(np.min(info['range'])),
 608                               int(np.max(info['range']))],
 609                    "value_type": "int"})
 610            elif info['type'] == 'float':
 611                self.parameters.append({
 612                    "name": name,
 613                    "type": "range",
 614                    "bounds": [float(np.min(info['range'])),
 615                               float(np.max(info['range']))],
 616                    "value_type": "float"})
 617        
 618        self.ax_client.create_experiment(
 619            name="bayesian_optimization",
 620            parameters=self.parameters,
 621            objectives={k: ObjectiveProperties(minimize=not v) 
 622                        for k,v in self._maximize.items() 
 623                        if isinstance(v, bool) and k in self._outcomes.keys()},
 624            parameter_constraints=self._feature_constraints,
 625            outcome_constraints=self._outcome_constraints,
 626            overwrite_existing_experiment=True
 627        )
 628
 629        if len(next(iter(self._outcomes.values()))['data']) > 0:
 630            for i in range(len(next(iter(self._outcomes.values()))['data'])):
 631                params = {name: info['data'][i] for name, info in self._features.items()}
 632                outcomes = {name: info['data'][i] for name, info in self._outcomes.items()}
 633                self.ax_client.attach_trial(params)
 634                self.ax_client.complete_trial(trial_index=i, raw_data=outcomes)
 635
 636        self.set_model()
 637        self.set_gs()
 638
 639    def set_model(self):
 640        """
 641        Set the model to be used for predictions.
 642        This method is called after initializing the AxClient.
 643        """
 644        self.model = Models.BOTORCH_MODULAR(
 645                experiment=self.ax_client.experiment,
 646                data=self.ax_client.experiment.fetch_data()
 647                )
 648    
 649    def set_gs(self):
 650        """
 651        Set the generation strategy for the experiment.
 652        This method is called after initializing the AxClient.
 653        """
 654        self.clear_trials()
 655        if self._optim == 'bo':
 656            if not self.model:
 657                self.set_model()
 658            if self.acq_func is None:
 659                self.gs = GenerationStrategy(
 660                    steps=[GenerationStep(
 661                                model=Models.BOTORCH_MODULAR,
 662                                num_trials=-1,  # No limitation on how many trials should be produced from this step
 663                                max_parallelism=3,  # Parallelism limit for this step, often lower than for Sobol
 664                            )
 665                        ]
 666                    )
 667            else:
 668                self.gs = GenerationStrategy(
 669                    steps=[GenerationStep(
 670                                model=Models.BOTORCH_MODULAR,
 671                                num_trials=-1,  # No limitation on how many trials should be produced from this step
 672                                max_parallelism=3,  # Parallelism limit for this step, often lower than for Sobol
 673                                model_configs={"botorch_model_class": self.acq_func['acqf']},
 674                                model_gen_options={"acquisition_options": self.acq_func['acqf_kwargs']}
 675                            )
 676                        ]
 677                    )
 678        elif self._optim == 'sobol':
 679            self.gs = GenerationStrategy(
 680                steps=[GenerationStep(
 681                            model=Models.SOBOL,
 682                            num_trials=-1,  # How many trials should be produced from this generation step
 683                            should_deduplicate=True,  # Deduplicate the trials
 684                            # model_kwargs={"seed": 165478},  # Any kwargs you want passed into the model
 685                            model_gen_kwargs={},  # Any kwargs you want passed to `modelbridge.gen`
 686                        )
 687                    ]
 688                )
 689        self.generator_run = self.gs.gen(
 690                experiment=self.ax_client.experiment,  # Ax `Experiment`, for which to generate new candidates
 691                data=None,  # Ax `Data` to use for model training, optional.
 692                n=self._N,  # Number of candidate arms to produce
 693                fixed_features=self._fixed_features, 
 694                pending_observations=get_pending_observation_features(
 695                    self.ax_client.experiment
 696                ),  # Points that should not be re-generated
 697            )
 698    
 699    def clear_trials(self):
 700        """
 701        Clear all trials in the experiment.
 702        """
 703        # Get all pending trial indices
 704        pending_trials = [k for k,i in self.ax_client.experiment.trials.items() 
 705                            if i.status==TrialStatus.CANDIDATE]
 706        for i in pending_trials:
 707            self.ax_client.experiment.trials[i].mark_abandoned()
 708    
 709    def suggest_next_trials(self, with_predicted=True):
 710        """
 711        Suggest the next set of trials based on the current model and optimization strategy.
 712
 713        Returns
 714        -------
 715
 716        pd.DataFrame: 
 717            DataFrame containing the suggested trials and their predicted outcomes.
 718        """
 719        self.clear_trials()
 720        if self.ax_client is None:
 721            self.initialize_ax_client()
 722        if self._N == 1:
 723            self.candidate = self.ax_client.experiment.new_trial(self.generator_run)
 724        else:
 725            self.candidate = self.ax_client.experiment.new_batch_trial(self.generator_run)
 726        trials = self.ax_client.get_trials_data_frame()
 727        trials = trials[trials['trial_status'] == 'CANDIDATE']
 728        trials = trials[[name for name in self.names]]
 729        if with_predicted:
 730            topred = [trials.iloc[i].to_dict() for i in range(len(trials))]
 731            preds = pd.DataFrame(self.predict(topred))
 732            # add 'predicted_' to the names of the pred dataframe
 733            preds.columns = [f'Predicted_{col}' for col in preds.columns]
 734            preds = preds.reset_index(drop=True)
 735            trials = trials.reset_index(drop=True)
 736            return pd.concat([trials, preds], axis=1)
 737        else:
 738            return trials
 739
 740    def predict(self, params):
 741        """
 742        Predict the outcomes for a given set of parameters using the current model.
 743
 744        Parameters
 745        ----------
 746
 747        params : List[Dict[str, Any]]
 748            List of parameter dictionaries for which to predict outcomes.
 749
 750        Returns
 751        -------
 752
 753        List[Dict[str, float]]: 
 754            List of predicted outcomes for the given parameters.
 755        """
 756        if self.ax_client is None:
 757            self.initialize_ax_client()
 758        obs_feats = [ObservationFeatures(parameters=p) for p in params]
 759        f, _ = self.model.predict(obs_feats)
 760        return f
 761
 762    def update_experiment(self, params, outcomes):
 763        """
 764        Update the experiment with new parameters and outcomes, and reinitialize the AxClient.
 765
 766        Parameters
 767        ----------
 768
 769        params : Dict[str, Any]
 770            Dictionary of new parameters to update the experiment with.
 771
 772        outcomes : Dict[str, Any]
 773            Dictionary of new outcomes to update the experiment with.
 774        """
 775        # append new data to the features and outcomes dictionaries
 776        for k, v in zip(params.keys(), params.values()):
 777            if k not in self._features:
 778                raise ValueError(f"Parameter '{k}' not found in features")
 779            if isinstance(v, np.ndarray):
 780                v = v.tolist()
 781            if not isinstance(v, list):
 782                v = [v]
 783            self._features[k]['data'] += v
 784        for k, v in zip(outcomes.keys(), outcomes.values()):
 785            if k not in self._outcomes:
 786                raise ValueError(f"Outcome '{k}' not found in outcomes")
 787            if isinstance(v, np.ndarray):
 788                v = v.tolist()
 789            if not isinstance(v, list):
 790                v = [v]
 791            self._outcomes[k]['data'] += v
 792        self.initialize_ax_client()
 793
 794    def plot_model(self, metricname=None, slice_values={}, linear=False):
 795        """
 796        Plot the model's predictions for the experiment's parameters and outcomes.
 797
 798        Parameters
 799        ----------
 800
 801        metricname : Optional[str]
 802            The name of the metric to plot. If None, the first outcome metric is used.
 803
 804        slice_values : Optional[Dict[str, Any]]
 805            Dictionary of slice values for plotting.
 806
 807        linear : bool
 808            Whether to plot a linear slice plot. Default is False.
 809
 810        Returns
 811        -------
 812
 813        plotly.graph_objects.Figure: 
 814            Plotly figure of the model's predictions.
 815        """
 816        if self.ax_client is None:
 817            self.initialize_ax_client()
 818            self.suggest_next_trials()
 819
 820        cand_name = 'Candidate' if self._N == 1 else 'Candidates'
 821        mname = self.ax_client.objective_names[0] if metricname is None else metricname
 822        param_name = [name for name in self.names if name not in slice_values.keys()]
 823        par_numeric = [name for name in param_name if self._features[name]['type'] in ['int', 'float']]
 824        if len(par_numeric)==1:
 825            fig = plot_slice(
 826                    model=self.model,
 827                    metric_name=mname,
 828                    density=100,
 829                    param_name=par_numeric[0],
 830                    generator_runs_dict={cand_name: self.generator_run},
 831                    slice_values=slice_values
 832                    )
 833        elif len(par_numeric)==2:
 834            fig = plot_contour(
 835                    model=self.model,
 836                    metric_name=mname,
 837                    param_x=par_numeric[0],
 838                    param_y=par_numeric[1],
 839                    generator_runs_dict={cand_name: self.generator_run},
 840                    slice_values=slice_values
 841                    )
 842        else:
 843            fig = interact_contour(
 844                    model=self.model,
 845                    generator_runs_dict={cand_name: self.generator_run},
 846                    metric_name=mname,
 847                    slice_values=slice_values,
 848                )
 849
 850        # Turn the figure into a plotly figure
 851        plotly_fig = go.Figure(fig.data)
 852
 853        # Modify only the "In-sample" markers
 854        trials = self.ax_client.get_trials_data_frame()
 855        trials = trials[trials['trial_status'] == 'CANDIDATE']
 856        trials = trials[[name for name in self.names]]
 857        for trace in plotly_fig.data:
 858            if trace.type == "contour":  # Check if it's a contour plot
 859                trace.colorscale = "viridis"  # Apply Viridis colormap
 860            if 'marker' in trace:  # Modify only the "In-sample" markers
 861                trace.marker.color = "white"  # Change marker color
 862                trace.marker.symbol = "circle"  # Change marker style
 863                trace.marker.size = 10
 864                trace.marker.line.width = 2
 865                trace.marker.line.color = 'black'
 866                if trace.text is not None:
 867                    trace.text = [t.replace('Arm', '<b>Sample').replace("_0","</b>") for t in trace.text]
 868            if trace.legendgroup == cand_name:  # Modify only the "Candidate" markers
 869                trace.marker.color = "red"  # Change marker color
 870                trace.name = cand_name
 871                trace.marker.symbol = "x"
 872                trace.marker.size = 12
 873                trace.marker.opacity = 1
 874                # Add hover info
 875                trace.hoverinfo = "text"  # Enable custom text for hover
 876                trace.hoverlabel = dict(bgcolor="#f8d5cd", font_color='black')
 877                if trace.text is not None:
 878                    trace.text = [t.replace("<i>","").replace("</i>","") for t in trace.text]
 879                trace.text = [
 880                    f"<b>Candidate {i+1}</b><br>{'<br>'.join([f'{col}: {val}' for col, val in trials.iloc[i].items()])}"
 881                    for t in trace.text
 882                    for i in range(len(trials))
 883                ]
 884        plotly_fig.update_layout(
 885            plot_bgcolor="white",  # White background
 886            legend=dict(bgcolor='rgba(0,0,0,0)'),
 887            margin=dict(l=10, r=10, t=50, b=50),
 888            xaxis=dict(
 889                showgrid=True,  # Enable grid
 890                gridcolor="lightgray",  # Light gray grid lines
 891                zeroline=False,
 892                zerolinecolor="black",  # Black zero line
 893                showline=True,
 894                linewidth=1,
 895                linecolor="black",  # Black border
 896                mirror=True
 897            ),
 898            yaxis=dict(
 899                showgrid=True,  # Enable grid
 900                gridcolor="lightgray",  # Light gray grid lines
 901                zeroline=False,
 902                zerolinecolor="black",  # Black zero line
 903                showline=True,
 904                linewidth=1,
 905                linecolor="black",  # Black border
 906                mirror=True
 907            ),
 908            xaxis2=dict(
 909                showgrid=True,  # Enable grid
 910                gridcolor="lightgray",  # Light gray grid lines
 911                zeroline=False,
 912                zerolinecolor="black",  # Black zero line
 913                showline=True,
 914                linewidth=1,
 915                linecolor="black",  # Black border
 916                mirror=True
 917            ),
 918            yaxis2=dict(
 919                showgrid=True,  # Enable grid
 920                gridcolor="lightgray",  # Light gray grid lines
 921                zeroline=False,
 922                zerolinecolor="black",  # Black zero line
 923                showline=True,
 924                linewidth=1,
 925                linecolor="black",  # Black border
 926                mirror=True
 927            ),
 928        )
 929        return plotly_fig
 930
 931    def plot_optimization_trace(self, optimum=None):
 932        """
 933        Plot the optimization trace, showing the progress of the optimization over trials.
 934
 935        Parameters
 936        ----------
 937
 938        optimum : Optional[float]
 939            The optimal value to plot on the optimization trace.
 940
 941        Returns
 942        -------
 943
 944        plotly.graph_objects.Figure: 
 945            Plotly figure of the optimization trace.
 946        """
 947        if self.ax_client is None:
 948            self.initialize_ax_client()
 949        if len(self._outcomes) > 1:
 950            print("Optimization trace is not available for multi-objective optimization.")
 951            return None
 952        fig = self.ax_client.get_optimization_trace(objective_optimum=optimum)
 953        fig = go.Figure(fig.data)
 954        for trace in fig.data:
 955            # add hover info
 956            trace.hoverinfo = "x+y"
 957        fig.update_layout(
 958            plot_bgcolor="white",  # White background
 959            legend=dict(bgcolor='rgba(0,0,0,0)'),
 960            margin=dict(l=50, r=10, t=50, b=50),
 961            xaxis=dict(
 962                showgrid=True,  # Enable grid
 963                gridcolor="lightgray",  # Light gray grid lines
 964                zeroline=False,
 965                zerolinecolor="black",  # Black zero line
 966                showline=True,
 967                linewidth=1,
 968                linecolor="black",  # Black border
 969                mirror=True
 970            ),
 971            yaxis=dict(
 972                showgrid=True,  # Enable grid
 973                gridcolor="lightgray",  # Light gray grid lines
 974                zeroline=False,
 975                zerolinecolor="black",  # Black zero line
 976                showline=True,
 977                linewidth=1,
 978                linecolor="black",  # Black border
 979                mirror=True
 980            ),
 981        )
 982        return fig
 983
 984    def compute_pareto_frontier(self):
 985        """
 986        Compute the Pareto frontier for multi-objective optimization experiments.
 987
 988        Returns
 989        -------
 990        The Pareto frontier.
 991        """
 992        if self.ax_client is None:
 993            self.initialize_ax_client()
 994        if len(self._outcomes) < 2:
 995            print("Pareto frontier is not available for single-objective optimization.")
 996            return None
 997        
 998        objectives = self.ax_client.experiment.optimization_config.objective.objectives
 999        self.pareto_frontier = compute_posterior_pareto_frontier(
1000            experiment=self.ax_client.experiment,
1001            data=self.ax_client.experiment.fetch_data(),
1002            primary_objective=objectives[1].metric,
1003            secondary_objective=objectives[0].metric,
1004            absolute_metrics=[o.metric_names[0] for o in objectives],
1005            num_points=20,
1006        )
1007        return self.pareto_frontier
1008    
1009    def plot_pareto_frontier(self, show_error_bars=True):
1010        """
1011        Plot the Pareto frontier for multi-objective optimization experiments.
1012
1013        Parameters
1014        ----------
1015        show_error_bars : bool, optional
1016            Whether to show error bars on the plot. Default is True.
1017
1018        Returns
1019        -------
1020        plotly.graph_objects.Figure: 
1021            Plotly figure of the Pareto frontier.
1022        """
1023        if self.pareto_frontier is None:
1024            return None
1025        
1026        fig = plot_pareto_frontier(self.pareto_frontier)
1027        fig = go.Figure(fig.data)
1028        
1029        # Modify traces to show/hide error bars
1030        if not show_error_bars:
1031            for trace in fig.data:
1032                # Remove error bars by setting them to None
1033                if hasattr(trace, 'error_x') and trace.error_x is not None:
1034                    trace.error_x = None
1035                if hasattr(trace, 'error_y') and trace.error_y is not None:
1036                    trace.error_y = None
1037        
1038        fig.update_layout(
1039            plot_bgcolor="white",  # White background
1040            legend=dict(bgcolor='rgba(0,0,0,0)'),
1041            margin=dict(l=50, r=10, t=50, b=50),
1042            xaxis=dict(
1043                showgrid=True,  # Enable grid
1044                gridcolor="lightgray",  # Light gray grid lines
1045                zeroline=False,
1046                zerolinecolor="black",  # Black zero line
1047                showline=True,
1048                linewidth=1,
1049                linecolor="black",  # Black border
1050                mirror=True
1051            ),
1052            yaxis=dict(
1053                showgrid=True,  # Enable grid
1054                gridcolor="lightgray",  # Light gray grid lines
1055                zeroline=False,
1056                zerolinecolor="black",  # Black zero line
1057                showline=True,
1058                linewidth=1,
1059                linecolor="black",  # Black border
1060                mirror=True
1061            ),
1062        )
1063        return fig
1064
1065    def get_best_parameters(self):
1066        """
1067        Return the best parameters found by the optimization process.
1068
1069        Returns
1070        -------
1071
1072        pd.DataFrame: 
1073            DataFrame containing the best parameters and their outcomes.
1074        """
1075        if self.ax_client is None:
1076            self.initialize_ax_client()
1077        if self.Nmetrics == 1:
1078            best_parameters = self.ax_client.get_best_parameters()[0]
1079            best_outcomes = self.ax_client.get_best_parameters()[1]
1080            best_parameters.update(best_outcomes[0])
1081            best = pd.DataFrame(best_parameters, index=[0])
1082        else:
1083            best_parameters = self.ax_client.get_pareto_optimal_parameters()
1084            best = ordered_dict_to_dataframe(best_parameters)
1085        return best
1086
1087# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
1088
1089def flatten_dict(d, parent_key="", sep="_"):
1090    """
1091    Flatten a nested dictionary.
1092    """
1093    items = []
1094    for k, v in d.items():
1095        new_key = f"{parent_key}{sep}{k}" if parent_key else k
1096        if isinstance(v, dict):
1097            items.extend(flatten_dict(v, new_key, sep=sep).items())
1098        else:
1099            items.append((new_key, v))
1100    return dict(items)
1101
1102# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
1103
1104def ordered_dict_to_dataframe(data):
1105    """
1106    Convert an OrderedDict with arbitrary nesting to a DataFrame.
1107    """
1108    dflat = flatten_dict(data)
1109    out = []
1110
1111    for key, value in dflat.items():
1112        main_dict = value[0]
1113        sub_dict = value[1][0]
1114        out.append([value for value in main_dict.values()] +
1115                   [value for value in sub_dict.values()])
1116
1117    df = pd.DataFrame(out, columns=[key for key in main_dict.keys()] +
1118                                   [key for key in sub_dict.keys()])
1119    return df
1120
1121# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
1122
1123def read_experimental_data(file_path: str, out_pos=[-1]) -> (Dict[str, Dict[str, Any]], Dict[str, Dict[str, Any]]):
1124    """
1125    Read experimental data from a CSV file and format it into features and outcomes dictionaries.
1126
1127    Parameters
1128    ----------
1129    file_path (str) 
1130        Path to the CSV file containing experimental data.
1131    out_pos (list of int)
1132        Column indices of the outcome variables. Default is the last column.
1133
1134    Returns
1135    -------
1136    Tuple[Dict[str, Dict[str, Any]], Dict[str, Dict[str, Any]]]
1137        Formatted features and outcomes dictionaries.
1138    """
1139    data = pd.read_csv(file_path)
1140    data = clean_names(data, remove_special=True, case_type='preserve')
1141    outcome_column_name = data.columns[out_pos]
1142    features = data.loc[:, ~data.columns.isin(outcome_column_name)].copy()
1143    outcomes = data[outcome_column_name].copy()
1144
1145    feature_definitions = {}
1146    for column in features.columns:
1147        if features[column].dtype == 'object':
1148            unique_values = features[column].unique()
1149            feature_definitions[column] = {'type': 'text',
1150                                           'range': unique_values.tolist()}
1151        elif features[column].dtype in ['int64', 'float64']:
1152            min_val = features[column].min()
1153            max_val = features[column].max()
1154            feature_type = 'int' if features[column].dtype == 'int64' else 'float'
1155            feature_definitions[column] = {'type': feature_type,
1156                                           'range': [min_val, max_val]}
1157
1158    formatted_features = {name: {'type': info['type'],
1159                                 'data': features[name].tolist(),
1160                                 'range': info['range']}
1161                          for name, info in feature_definitions.items()}
1162    # same for outcomes with just type and data
1163    outcome_definitions = {}
1164    for column in outcomes.columns:
1165        if outcomes[column].dtype == 'object':
1166            unique_values = outcomes[column].unique()
1167            outcome_definitions[column] = {'type': 'text',
1168                                           'data': unique_values.tolist()}
1169        elif outcomes[column].dtype in ['int64', 'float64']:
1170            min_val = outcomes[column].min()
1171            max_val = outcomes[column].max()
1172            outcome_type = 'int' if outcomes[column].dtype == 'int64' else 'float'
1173            outcome_definitions[column] = {'type': outcome_type,
1174                                           'data': outcomes[column].tolist()}
1175    formatted_outcomes = {name: {'type': info['type'],
1176                                 'data': outcomes[name].tolist()}
1177                           for name, info in outcome_definitions.items()}
1178    return formatted_features, formatted_outcomes