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.html).
  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
 239    @property
 240    def features(self):
 241        """
 242        A dictionary defining the features of the experiment, including their types and ranges.
 243        
 244        Example
 245        -------
 246        ```python
 247        features = {
 248            'feature1': {'type': 'int', 
 249                         'data': [1, 2, 3], 
 250                         'range': [1, 3]},
 251            'feature2': {'type': 'float', 
 252                         'data': [0.1, 0.2, 0.3], 
 253                         'range': [0.1, 0.3]},
 254            'feature3': {'type': 'text', 
 255                         'data': ['A', 'B', 'C'], 
 256                         'range': ['A', 'B', 'C']}
 257            }
 258        ```
 259        """
 260        return self._features
 261
 262    @features.setter
 263    def features(self, value):
 264        """
 265        Set the features of the experiment with validation.
 266        """
 267        if not isinstance(value, dict):
 268            raise ValueError("features must be a dictionary")
 269        self._features = value
 270        for name in self._features.keys():
 271            if self.ranges and name in self.ranges.keys():
 272                self._features[name]['range'] = self.ranges[name]
 273            else:
 274                if self._features[name]['type'] == 'text':
 275                    self._features[name]['range'] = list(set(self._features[name]['data']))
 276                elif self._features[name]['type'] == 'int':
 277                    self._features[name]['range'] = [int(np.min(self._features[name]['data'])),
 278                                                     int(np.max(self._features[name]['data']))]
 279                elif self._features[name]['type'] == 'float':
 280                    self._features[name]['range'] = [float(np.min(self._features[name]['data'])),
 281                                                     float(np.max(self._features[name]['data']))]
 282        if self._first_initialization_done:
 283            self.initialize_ax_client()
 284    
 285    @property
 286    def ranges(self):
 287        """
 288        A dictionary defining the ranges of the features. Default is `None`.
 289        
 290        If not provided, the ranges will be inferred from the features data.
 291        The ranges should be in the format `{'feature_name': [minvalue,maxvalue]}`.
 292        """
 293        return self._ranges
 294
 295    @ranges.setter
 296    def ranges(self, value):
 297        """
 298        Set the ranges of the features with validation.
 299        """
 300        if value is not None:
 301            if not isinstance(value, dict):
 302                raise ValueError("ranges must be a dictionary")
 303        self._ranges = value
 304    
 305    @property
 306    def names(self):
 307        """
 308        The names of the features.
 309        """
 310        return self._names
 311    
 312    @names.setter
 313    def names(self, value):
 314        """
 315        Set the names of the features.
 316        """
 317        if not isinstance(value, list):
 318            raise ValueError("names must be a list")
 319        self._names = value
 320
 321    @property
 322    def outcomes(self):
 323        """
 324        A dictionary defining the outcomes of the experiment, including their types and observed data.
 325        
 326        Example
 327        -------
 328        ```python
 329        outcomes = {
 330            'outcome1': {'type': 'float', 
 331                         'data': [0.1, 0.2, 0.3]},
 332            'outcome2': {'type': 'float', 
 333                         'data': [1.0, 2.0, 3.0]}
 334            }
 335        ```
 336        """
 337        return self._outcomes
 338
 339    @outcomes.setter
 340    def outcomes(self, value):
 341        """
 342        Set the outcomes of the experiment with validation.
 343        """
 344        if not isinstance(value, dict):
 345            raise ValueError("outcomes must be a dictionary")
 346        self._outcomes = value
 347        self.out_names = list(value.keys())
 348        if self._first_initialization_done:
 349            self.initialize_ax_client()
 350    
 351    @property
 352    def fixed_features(self):
 353        """
 354        A dictionary defining fixed features with their values. Default is `None`.
 355        If provided, the fixed features will be treated as fixed parameters in the generation process.
 356        The fixed features should be in the format `{'feature_name': value}`.
 357        The values should be the fixed values for the respective features.
 358        """
 359        return self._fixed_features
 360
 361    @fixed_features.setter
 362    def fixed_features(self, value):
 363        """
 364        Set the fixed features of the experiment.
 365        """
 366        self._fixed_features = None
 367        if value is not None:
 368            if not isinstance(value, dict):
 369                raise ValueError("fixed_features must be a dictionary")
 370            for name in value.keys():
 371                if name not in self.names:
 372                    raise ValueError(f"Fixed feature '{name}' not found in features")
 373            # fixed_features should be an ObservationFeatures object
 374            self._fixed_features = ObservationFeatures(parameters=value)
 375        if self._first_initialization_done:
 376            self.set_gs()
 377
 378    @property
 379    def N(self):
 380        """
 381        The number of trials to suggest in each optimization step. Must be a positive integer. Default is `1`.
 382        """
 383        return self._N
 384
 385    @N.setter
 386    def N(self, value):
 387        """
 388        Set the number of trials to suggest in each optimization step with validation.
 389        """
 390        if not isinstance(value, int) or value <= 0:
 391            raise ValueError("N must be a positive integer")
 392        self._N = value
 393        if self._first_initialization_done:
 394            self.set_gs()
 395
 396    @property
 397    def maximize(self):
 398        """
 399        A boolean or dict indicating whether to maximize the outcomes in the form `{'outcome1':True, 'outcome2':False}`.
 400        If a single boolean is provided, it is applied to all outcomes. Default is `True`.
 401        """
 402        return self._maximize
 403
 404    @maximize.setter
 405    def maximize(self, value):
 406        """
 407        Set the maximization setting for the outcomes with validation.
 408        """
 409        if isinstance(value, bool):
 410            self._maximize = {out: value for out in self.out_names}
 411        elif isinstance(value, dict) and len(value) == len(self._outcomes):
 412            self._maximize = {k:v for k,v in value.items() if 
 413                              (k in self.out_names and isinstance(v, bool))}
 414        else:
 415            raise ValueError("maximize must be a boolean or a list of booleans with the same length as outcomes")
 416        if self._first_initialization_done:
 417            self.initialize_ax_client()
 418
 419    @property
 420    def outcome_constraints(self):
 421        """
 422        Constraints on the outcomes, specified as a list of strings. Default is `None`.
 423        """
 424        return self._outcome_constraints
 425
 426    @outcome_constraints.setter
 427    def outcome_constraints(self, value):
 428        """
 429        Set the outcome constraints of the experiment with validation.
 430        """
 431        if isinstance(value, str):
 432            self._outcome_constraints = [value]
 433        elif isinstance(value, list):
 434            self._outcome_constraints = value
 435        else:
 436            self._outcome_constraints = None
 437        if self._first_initialization_done:
 438            self.initialize_ax_client()
 439
 440    @property
 441    def feature_constraints(self):
 442        """
 443        Constraints on the features, specified as a list of strings. Default is `None`.
 444        
 445        Example
 446        -------
 447        ```python
 448        feature_constraints = [
 449            'feature1 <= 10.0',
 450            'feature1 + 2*feature2 >= 3.0'
 451        ]
 452        ```
 453        """
 454        return self._feature_constraints
 455
 456    @feature_constraints.setter
 457    def feature_constraints(self, value):
 458        """
 459        Set the feature constraints of the experiment with validation.
 460        """
 461        if isinstance(value, dict):
 462            self._feature_constraints = [value]
 463        elif isinstance(value, list):
 464            self._feature_constraints = value
 465        elif isinstance(value, str):
 466            self._feature_constraints = [value]
 467        else:
 468            self._feature_constraints = None
 469        if self._first_initialization_done:
 470            self.initialize_ax_client()
 471
 472    @property
 473    def optim(self):
 474        """
 475        The optimization method to use, either `'bo'` for Bayesian Optimization or `'sobol'` for Sobol sequence. Default is `'bo'`.
 476        """
 477        return self._optim
 478
 479    @optim.setter
 480    def optim(self, value):
 481        """
 482        Set the optimization method with validation.
 483        """
 484        value = value.lower()
 485        if value not in ['bo', 'sobol']:
 486            raise ValueError("Optimization method must be either 'bo' or 'sobol'")
 487        self._optim = value
 488        if self._first_initialization_done:
 489            self.set_gs()
 490
 491    @property
 492    def data(self) -> pd.DataFrame:
 493        """
 494        Returns a DataFrame of the current data in the experiment, including features and outcomes.
 495        """
 496        feature_data = {name: info['data'] for name, info in self._features.items()}
 497        outcome_data = {name: info['data'] for name, info in self._outcomes.items()}
 498        data_dict = {**feature_data, **outcome_data}
 499        return pd.DataFrame(data_dict)
 500
 501    @data.setter
 502    def data(self, value: pd.DataFrame):
 503        """
 504        Sets the features and outcomes data from a given DataFrame.
 505        """
 506        if not isinstance(value, pd.DataFrame):
 507            raise ValueError("Data must be a pandas DataFrame")
 508
 509        feature_columns = [col for col in value.columns if col in self._features]
 510        outcome_columns = [col for col in value.columns if col in self._outcomes]
 511
 512        for col in feature_columns:
 513            self._features[col]['data'] = value[col].tolist()
 514
 515        for col in outcome_columns:
 516            self._outcomes[col]['data'] = value[col].tolist()
 517
 518        if self._first_initialization_done:
 519            self.initialize_ax_client()
 520    
 521    @property
 522    def acq_func(self):
 523        """
 524        The acquisition function to use for the optimization process. It must be a dict with 2 keys:
 525        - `acqf`: the acquisition function class to use (e.g., `UpperConfidenceBound`),
 526        - `acqf_kwargs`: a dict of the kwargs to pass to the acquisition function class. (e.g. `{'beta': 0.1}`).
 527        
 528        If not provided, the default acquisition function is used (`LogExpectedImprovement` or `qLogExpectedImprovement` if N>1).
 529        
 530        Example
 531        -------
 532        ```python
 533        acq_func = {
 534            'acqf': UpperConfidenceBound,
 535            'acqf_kwargs': {'beta': 0.1} # lower value = exploitation, higher value = exploration
 536        }
 537        ```
 538        """
 539        return self._acq_func
 540    
 541    @acq_func.setter
 542    def acq_func(self, value):
 543        """
 544        Set the acquisition function with validation.
 545        """
 546        self._acq_func = value
 547        if self._first_initialization_done:
 548            self.set_gs()
 549
 550    def __repr__(self):
 551        return self.__str__()
 552
 553    def __str__(self):
 554        """
 555        Return a string representation of the BOExperiment instance.
 556        """
 557        return f"""
 558BOExperiment(
 559    N={self.N},
 560    maximize={self.maximize},
 561    outcome_constraints={self.outcome_constraints},
 562    feature_constraints={self.feature_constraints},
 563    optim={self.optim}
 564)
 565
 566Input data:
 567
 568{self.data}
 569        """
 570
 571    def initialize_ax_client(self):
 572        """
 573        Initialize the AxClient with the experiment's parameters, objectives, and constraints.
 574        """
 575        print('\n========   INITIALIZING MODEL   ========\n')
 576        self.ax_client = AxClient(verbose_logging=False, 
 577                                  suppress_storage_errors=True)
 578        self.parameters = []
 579        for name, info in self._features.items():
 580            if info['type'] == 'text':
 581                self.parameters.append({
 582                    "name": name,
 583                    "type": "choice",
 584                    "values": [str(val) for val in info['range']],
 585                    "value_type": "str"})
 586            elif info['type'] == 'int':
 587                self.parameters.append({
 588                    "name": name,
 589                    "type": "range",
 590                    "bounds": [int(np.min(info['range'])),
 591                               int(np.max(info['range']))],
 592                    "value_type": "int"})
 593            elif info['type'] == 'float':
 594                self.parameters.append({
 595                    "name": name,
 596                    "type": "range",
 597                    "bounds": [float(np.min(info['range'])),
 598                               float(np.max(info['range']))],
 599                    "value_type": "float"})
 600        
 601        self.ax_client.create_experiment(
 602            name="bayesian_optimization",
 603            parameters=self.parameters,
 604            objectives={k: ObjectiveProperties(minimize=not v) 
 605                        for k,v in self._maximize.items() 
 606                        if isinstance(v, bool) and k in self._outcomes.keys()},
 607            parameter_constraints=self._feature_constraints,
 608            outcome_constraints=self._outcome_constraints,
 609            overwrite_existing_experiment=True
 610        )
 611
 612        if len(next(iter(self._outcomes.values()))['data']) > 0:
 613            for i in range(len(next(iter(self._outcomes.values()))['data'])):
 614                params = {name: info['data'][i] for name, info in self._features.items()}
 615                outcomes = {name: info['data'][i] for name, info in self._outcomes.items()}
 616                self.ax_client.attach_trial(params)
 617                self.ax_client.complete_trial(trial_index=i, raw_data=outcomes)
 618
 619        self.set_model()
 620        self.set_gs()
 621
 622    def set_model(self):
 623        """
 624        Set the model to be used for predictions.
 625        This method is called after initializing the AxClient.
 626        """
 627        self.model = Models.BOTORCH_MODULAR(
 628                experiment=self.ax_client.experiment,
 629                data=self.ax_client.experiment.fetch_data()
 630                )
 631    
 632    def set_gs(self):
 633        """
 634        Set the generation strategy for the experiment.
 635        This method is called after initializing the AxClient.
 636        """
 637        self.clear_trials()
 638        if self._optim == 'bo':
 639            if not self.model:
 640                self.set_model()
 641            if self.acq_func is None:
 642                self.gs = GenerationStrategy(
 643                    steps=[GenerationStep(
 644                                model=Models.BOTORCH_MODULAR,
 645                                num_trials=-1,  # No limitation on how many trials should be produced from this step
 646                                max_parallelism=3,  # Parallelism limit for this step, often lower than for Sobol
 647                            )
 648                        ]
 649                    )
 650            else:
 651                self.gs = GenerationStrategy(
 652                    steps=[GenerationStep(
 653                                model=Models.BOTORCH_MODULAR,
 654                                num_trials=-1,  # No limitation on how many trials should be produced from this step
 655                                max_parallelism=3,  # Parallelism limit for this step, often lower than for Sobol
 656                                model_kwargs={"botorch_acqf_class": self.acq_func['acqf'],
 657                                              "acquisition_options": self.acq_func['acqf_kwargs']}
 658                            )
 659                        ]
 660                    )
 661        elif self._optim == 'sobol':
 662            self.gs = GenerationStrategy(
 663                steps=[GenerationStep(
 664                            model=Models.SOBOL,
 665                            num_trials=-1,  # How many trials should be produced from this generation step
 666                            should_deduplicate=True,  # Deduplicate the trials
 667                            # model_kwargs={"seed": 165478},  # Any kwargs you want passed into the model
 668                            model_gen_kwargs={},  # Any kwargs you want passed to `modelbridge.gen`
 669                        )
 670                    ]
 671                )
 672        self.generator_run = self.gs.gen(
 673                experiment=self.ax_client.experiment,  # Ax `Experiment`, for which to generate new candidates
 674                data=None,  # Ax `Data` to use for model training, optional.
 675                n=self._N,  # Number of candidate arms to produce
 676                fixed_features=self._fixed_features, 
 677                pending_observations=get_pending_observation_features(
 678                    self.ax_client.experiment
 679                ),  # Points that should not be re-generated
 680            )
 681    
 682    def clear_trials(self):
 683        """
 684        Clear all trials in the experiment.
 685        """
 686        # Get all pending trial indices
 687        pending_trials = [k for k,i in self.ax_client.experiment.trials.items() 
 688                            if i.status==TrialStatus.CANDIDATE]
 689        for i in pending_trials:
 690            self.ax_client.experiment.trials[i].mark_abandoned()
 691    
 692    def suggest_next_trials(self, with_predicted=True):
 693        """
 694        Suggest the next set of trials based on the current model and optimization strategy.
 695
 696        Returns
 697        -------
 698
 699        pd.DataFrame: 
 700            DataFrame containing the suggested trials and their predicted outcomes.
 701        """
 702        self.clear_trials()
 703        if self.ax_client is None:
 704            self.initialize_ax_client()
 705        if self._N == 1:
 706            self.candidate = self.ax_client.experiment.new_trial(self.generator_run)
 707        else:
 708            self.candidate = self.ax_client.experiment.new_batch_trial(self.generator_run)
 709        trials = self.ax_client.get_trials_data_frame()
 710        trials = trials[trials['trial_status'] == 'CANDIDATE']
 711        trials = trials[[name for name in self.names]]
 712        if with_predicted:
 713            topred = [trials.iloc[i].to_dict() for i in range(len(trials))]
 714            preds = pd.DataFrame(self.predict(topred))
 715            # add 'predicted_' to the names of the pred dataframe
 716            preds.columns = [f'Predicted_{col}' for col in preds.columns]
 717            preds = preds.reset_index(drop=True)
 718            trials = trials.reset_index(drop=True)
 719            return pd.concat([trials, preds], axis=1)
 720        else:
 721            return trials
 722
 723    def predict(self, params):
 724        """
 725        Predict the outcomes for a given set of parameters using the current model.
 726
 727        Parameters
 728        ----------
 729
 730        params : List[Dict[str, Any]]
 731            List of parameter dictionaries for which to predict outcomes.
 732
 733        Returns
 734        -------
 735
 736        List[Dict[str, float]]: 
 737            List of predicted outcomes for the given parameters.
 738        """
 739        if self.ax_client is None:
 740            self.initialize_ax_client()
 741        obs_feats = [ObservationFeatures(parameters=p) for p in params]
 742        f, _ = self.model.predict(obs_feats)
 743        return f
 744
 745    def update_experiment(self, params, outcomes):
 746        """
 747        Update the experiment with new parameters and outcomes, and reinitialize the AxClient.
 748
 749        Parameters
 750        ----------
 751
 752        params : Dict[str, Any]
 753            Dictionary of new parameters to update the experiment with.
 754
 755        outcomes : Dict[str, Any]
 756            Dictionary of new outcomes to update the experiment with.
 757        """
 758        # append new data to the features and outcomes dictionaries
 759        for k, v in zip(params.keys(), params.values()):
 760            if k not in self._features:
 761                raise ValueError(f"Parameter '{k}' not found in features")
 762            if isinstance(v, np.ndarray):
 763                v = v.tolist()
 764            if not isinstance(v, list):
 765                v = [v]
 766            self._features[k]['data'] += v
 767        for k, v in zip(outcomes.keys(), outcomes.values()):
 768            if k not in self._outcomes:
 769                raise ValueError(f"Outcome '{k}' not found in outcomes")
 770            if isinstance(v, np.ndarray):
 771                v = v.tolist()
 772            if not isinstance(v, list):
 773                v = [v]
 774            self._outcomes[k]['data'] += v
 775        self.initialize_ax_client()
 776
 777    def plot_model(self, metricname=None, slice_values={}, linear=False):
 778        """
 779        Plot the model's predictions for the experiment's parameters and outcomes.
 780
 781        Parameters
 782        ----------
 783
 784        metricname : Optional[str]
 785            The name of the metric to plot. If None, the first outcome metric is used.
 786
 787        slice_values : Optional[Dict[str, Any]]
 788            Dictionary of slice values for plotting.
 789
 790        linear : bool
 791            Whether to plot a linear slice plot. Default is False.
 792
 793        Returns
 794        -------
 795
 796        plotly.graph_objects.Figure: 
 797            Plotly figure of the model's predictions.
 798        """
 799        if self.ax_client is None:
 800            self.initialize_ax_client()
 801            self.suggest_next_trials()
 802
 803        cand_name = 'Candidate' if self._N == 1 else 'Candidates'
 804        mname = self.ax_client.objective_names[0] if metricname is None else metricname
 805        param_name = [name for name in self.names if name not in slice_values.keys()]
 806        par_numeric = [name for name in param_name if self._features[name]['type'] in ['int', 'float']]
 807        if len(par_numeric)==1:
 808            fig = plot_slice(
 809                    model=self.model,
 810                    metric_name=mname,
 811                    density=100,
 812                    param_name=par_numeric[0],
 813                    generator_runs_dict={cand_name: self.generator_run},
 814                    slice_values=slice_values
 815                    )
 816        elif len(par_numeric)==2:
 817            fig = plot_contour(
 818                    model=self.model,
 819                    metric_name=mname,
 820                    param_x=par_numeric[0],
 821                    param_y=par_numeric[1],
 822                    generator_runs_dict={cand_name: self.generator_run},
 823                    slice_values=slice_values
 824                    )
 825        else:
 826            fig = interact_contour(
 827                    model=self.model,
 828                    generator_runs_dict={cand_name: self.generator_run},
 829                    metric_name=mname,
 830                    slice_values=slice_values,
 831                )
 832
 833        # Turn the figure into a plotly figure
 834        plotly_fig = go.Figure(fig.data)
 835
 836        # Modify only the "In-sample" markers
 837        trials = self.ax_client.get_trials_data_frame()
 838        trials = trials[trials['trial_status'] == 'CANDIDATE']
 839        trials = trials[[name for name in self.names]]
 840        for trace in plotly_fig.data:
 841            if trace.type == "contour":  # Check if it's a contour plot
 842                trace.colorscale = "viridis"  # Apply Viridis colormap
 843            if 'marker' in trace:  # Modify only the "In-sample" markers
 844                trace.marker.color = "white"  # Change marker color
 845                trace.marker.symbol = "circle"  # Change marker style
 846                trace.marker.size = 10
 847                trace.marker.line.width = 2
 848                trace.marker.line.color = 'black'
 849                if trace.text is not None:
 850                    trace.text = [t.replace('Arm', '<b>Sample').replace("_0","</b>") for t in trace.text]
 851            if trace.legendgroup == cand_name:  # Modify only the "Candidate" markers
 852                trace.marker.color = "red"  # Change marker color
 853                trace.name = cand_name
 854                trace.marker.symbol = "x"
 855                trace.marker.size = 12
 856                trace.marker.opacity = 1
 857                # Add hover info
 858                trace.hoverinfo = "text"  # Enable custom text for hover
 859                trace.hoverlabel = dict(bgcolor="#f8d5cd", font_color='black')
 860                if trace.text is not None:
 861                    trace.text = [t.replace("<i>","").replace("</i>","") for t in trace.text]
 862                trace.text = [
 863                    f"<b>Candidate {i+1}</b><br>{'<br>'.join([f'{col}: {val}' for col, val in trials.iloc[i].items()])}"
 864                    for t in trace.text
 865                    for i in range(len(trials))
 866                ]
 867        plotly_fig.update_layout(
 868            plot_bgcolor="white",  # White background
 869            legend=dict(bgcolor='rgba(0,0,0,0)'),
 870            margin=dict(l=10, r=10, t=50, b=50),
 871            xaxis=dict(
 872                showgrid=True,  # Enable grid
 873                gridcolor="lightgray",  # Light gray grid lines
 874                zeroline=False,
 875                zerolinecolor="black",  # Black zero line
 876                showline=True,
 877                linewidth=1,
 878                linecolor="black",  # Black border
 879                mirror=True
 880            ),
 881            yaxis=dict(
 882                showgrid=True,  # Enable grid
 883                gridcolor="lightgray",  # Light gray grid lines
 884                zeroline=False,
 885                zerolinecolor="black",  # Black zero line
 886                showline=True,
 887                linewidth=1,
 888                linecolor="black",  # Black border
 889                mirror=True
 890            ),
 891            xaxis2=dict(
 892                showgrid=True,  # Enable grid
 893                gridcolor="lightgray",  # Light gray grid lines
 894                zeroline=False,
 895                zerolinecolor="black",  # Black zero line
 896                showline=True,
 897                linewidth=1,
 898                linecolor="black",  # Black border
 899                mirror=True
 900            ),
 901            yaxis2=dict(
 902                showgrid=True,  # Enable grid
 903                gridcolor="lightgray",  # Light gray grid lines
 904                zeroline=False,
 905                zerolinecolor="black",  # Black zero line
 906                showline=True,
 907                linewidth=1,
 908                linecolor="black",  # Black border
 909                mirror=True
 910            ),
 911        )
 912        return plotly_fig
 913
 914    def plot_optimization_trace(self, optimum=None):
 915        """
 916        Plot the optimization trace, showing the progress of the optimization over trials.
 917
 918        Parameters
 919        ----------
 920
 921        optimum : Optional[float]
 922            The optimal value to plot on the optimization trace.
 923
 924        Returns
 925        -------
 926
 927        plotly.graph_objects.Figure: 
 928            Plotly figure of the optimization trace.
 929        """
 930        if self.ax_client is None:
 931            self.initialize_ax_client()
 932        if len(self._outcomes) > 1:
 933            print("Optimization trace is not available for multi-objective optimization.")
 934            return None
 935        fig = self.ax_client.get_optimization_trace(objective_optimum=optimum)
 936        fig = go.Figure(fig.data)
 937        for trace in fig.data:
 938            # add hover info
 939            trace.hoverinfo = "x+y"
 940        fig.update_layout(
 941            plot_bgcolor="white",  # White background
 942            legend=dict(bgcolor='rgba(0,0,0,0)'),
 943            margin=dict(l=50, r=10, t=50, b=50),
 944            xaxis=dict(
 945                showgrid=True,  # Enable grid
 946                gridcolor="lightgray",  # Light gray grid lines
 947                zeroline=False,
 948                zerolinecolor="black",  # Black zero line
 949                showline=True,
 950                linewidth=1,
 951                linecolor="black",  # Black border
 952                mirror=True
 953            ),
 954            yaxis=dict(
 955                showgrid=True,  # Enable grid
 956                gridcolor="lightgray",  # Light gray grid lines
 957                zeroline=False,
 958                zerolinecolor="black",  # Black zero line
 959                showline=True,
 960                linewidth=1,
 961                linecolor="black",  # Black border
 962                mirror=True
 963            ),
 964        )
 965        return fig
 966
 967    def plot_pareto_frontier(self):
 968        """
 969        Plot the Pareto frontier for multi-objective optimization experiments.
 970
 971        Returns
 972        -------
 973
 974        plotly.graph_objects.Figure: 
 975            Plotly figure of the Pareto frontier.
 976        """
 977        if self.ax_client is None:
 978            self.initialize_ax_client()
 979        if len(self._outcomes) < 2:
 980            print("Pareto frontier is not available for single-objective optimization.")
 981            return None
 982        objectives = self.ax_client.experiment.optimization_config.objective.objectives
 983        frontier = compute_posterior_pareto_frontier(
 984            experiment=self.ax_client.experiment,
 985            data=self.ax_client.experiment.fetch_data(),
 986            primary_objective=objectives[1].metric,
 987            secondary_objective=objectives[0].metric,
 988            absolute_metrics=self.out_names,
 989            num_points=20,
 990        )
 991        fig = plot_pareto_frontier(frontier)
 992        fig = go.Figure(fig.data)
 993        fig.update_layout(
 994            plot_bgcolor="white",  # White background
 995            legend=dict(bgcolor='rgba(0,0,0,0)'),
 996            margin=dict(l=50, r=10, t=50, b=50),
 997            xaxis=dict(
 998                showgrid=True,  # Enable grid
 999                gridcolor="lightgray",  # Light gray grid lines
1000                zeroline=False,
1001                zerolinecolor="black",  # Black zero line
1002                showline=True,
1003                linewidth=1,
1004                linecolor="black",  # Black border
1005                mirror=True
1006            ),
1007            yaxis=dict(
1008                showgrid=True,  # Enable grid
1009                gridcolor="lightgray",  # Light gray grid lines
1010                zeroline=False,
1011                zerolinecolor="black",  # Black zero line
1012                showline=True,
1013                linewidth=1,
1014                linecolor="black",  # Black border
1015                mirror=True
1016            ),
1017        )
1018        return fig
1019
1020    def get_best_parameters(self):
1021        """
1022        Return the best parameters found by the optimization process.
1023
1024        Returns
1025        -------
1026
1027        pd.DataFrame: 
1028            DataFrame containing the best parameters and their outcomes.
1029        """
1030        if self.ax_client is None:
1031            self.initialize_ax_client()
1032        if self.Nmetrics == 1:
1033            best_parameters = self.ax_client.get_best_parameters()[0]
1034            best_outcomes = self.ax_client.get_best_parameters()[1]
1035            best_parameters.update(best_outcomes[0])
1036            best = pd.DataFrame(best_parameters, index=[0])
1037        else:
1038            best_parameters = self.ax_client.get_pareto_optimal_parameters()
1039            best = ordered_dict_to_dataframe(best_parameters)
1040        return best
1041
1042# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
1043
1044def flatten_dict(d, parent_key="", sep="_"):
1045    """
1046    Flatten a nested dictionary.
1047    """
1048    items = []
1049    for k, v in d.items():
1050        new_key = f"{parent_key}{sep}{k}" if parent_key else k
1051        if isinstance(v, dict):
1052            items.extend(flatten_dict(v, new_key, sep=sep).items())
1053        else:
1054            items.append((new_key, v))
1055    return dict(items)
1056
1057# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
1058
1059def ordered_dict_to_dataframe(data):
1060    """
1061    Convert an OrderedDict with arbitrary nesting to a DataFrame.
1062    """
1063    dflat = flatten_dict(data)
1064    out = []
1065
1066    for key, value in dflat.items():
1067        main_dict = value[0]
1068        sub_dict = value[1][0]
1069        out.append([value for value in main_dict.values()] +
1070                   [value for value in sub_dict.values()])
1071
1072    df = pd.DataFrame(out, columns=[key for key in main_dict.keys()] +
1073                                   [key for key in sub_dict.keys()])
1074    return df
1075
1076# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
1077
1078def read_experimental_data(file_path: str, out_pos=[-1]) -> (Dict[str, Dict[str, Any]], Dict[str, Dict[str, Any]]):
1079    """
1080    Read experimental data from a CSV file and format it into features and outcomes dictionaries.
1081
1082    Parameters
1083    ----------
1084    file_path (str) 
1085        Path to the CSV file containing experimental data.
1086    out_pos (list of int)
1087        Column indices of the outcome variables. Default is the last column.
1088
1089    Returns
1090    -------
1091    Tuple[Dict[str, Dict[str, Any]], Dict[str, Dict[str, Any]]]
1092        Formatted features and outcomes dictionaries.
1093    """
1094    data = pd.read_csv(file_path)
1095    data = clean_names(data, remove_special=True, case_type='preserve')
1096    outcome_column_name = data.columns[out_pos]
1097    features = data.loc[:, ~data.columns.isin(outcome_column_name)].copy()
1098    outcomes = data[outcome_column_name].copy()
1099
1100    feature_definitions = {}
1101    for column in features.columns:
1102        if features[column].dtype == 'object':
1103            unique_values = features[column].unique()
1104            feature_definitions[column] = {'type': 'text',
1105                                           'range': unique_values.tolist()}
1106        elif features[column].dtype in ['int64', 'float64']:
1107            min_val = features[column].min()
1108            max_val = features[column].max()
1109            feature_type = 'int' if features[column].dtype == 'int64' else 'float'
1110            feature_definitions[column] = {'type': feature_type,
1111                                           'range': [min_val, max_val]}
1112
1113    formatted_features = {name: {'type': info['type'],
1114                                 'data': features[name].tolist(),
1115                                 'range': info['range']}
1116                          for name, info in feature_definitions.items()}
1117    # same for outcomes with just type and data
1118    outcome_definitions = {}
1119    for column in outcomes.columns:
1120        if outcomes[column].dtype == 'object':
1121            unique_values = outcomes[column].unique()
1122            outcome_definitions[column] = {'type': 'text',
1123                                           'data': unique_values.tolist()}
1124        elif outcomes[column].dtype in ['int64', 'float64']:
1125            min_val = outcomes[column].min()
1126            max_val = outcomes[column].max()
1127            outcome_type = 'int' if outcomes[column].dtype == 'int64' else 'float'
1128            outcome_definitions[column] = {'type': outcome_type,
1129                                           'data': outcomes[column].tolist()}
1130    formatted_outcomes = {name: {'type': info['type'],
1131                                 'data': outcomes[name].tolist()}
1132                           for name, info in outcome_definitions.items()}
1133    return formatted_features, formatted_outcomes