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