Improved docstrings for optimizer and suggest

obsidian/optimizer/bayesian.py

@@ -4,7 +4,7 @@
 
 from obsidian.parameters import ParamSpace, Target, Task
 from obsidian.surrogates import SurrogateBoTorch, DNN
-from obsidian.acquisition import aq_class_dict, aq_hp_defaults, valid_aqs
+from obsidian.acquisition import aq_class_dict, aq_defaults, aq_hp_defaults, valid_aqs
 from obsidian.surrogates import model_class_dict
 from obsidian.objectives import Index_Objective, Objective_Sequence
 from obsidian.exceptions import IncompatibleObjectiveError, UnsupportedError, UnfitError, DataWarning
@@ -41,7 +41,23 @@ class BayesianOptimizer(Optimizer):
         X_space (ParamSpace): The parameter space defining the search space for the optimization.
         surrogate (str | dict | list[str] | list[dict], optional): The surrogate model(s) to use.
             It can be a string representing a single model type, a dictionary specifying multiple model types
-            with their hyperparameters, or a list of strings or dictionaries. Defaults to ['GP'].
+            with their hyperparameters, or a list of strings or dictionaries.
+            
+            Defaults to ``'GP'``. Options are as follows:
+            
+            - ``'GP'``: Gaussian Process with default settings (Matern Kernel, Gamma covariance priors)
+            - ``'MixedGP'``: GP with mixed parameter types (continuous, categorical). Will be re-selected
+              by default if 'GP' is selected and input space is mixed.
+            - ``'DKL'``: GP with a NN feature-extractor (deep kernel learning)
+            - ``'GPflat'``: GP without priors. May result in optimization instability, but removes bias
+              for special situations.
+            - ``'GPprior'``: GP with custom priors on the mean, likelihood, and covariance
+            - ``'MTGP'``: Multi-task GP for multi-output optimization. Will be re-selected by default
+              if 'GP' is selected and the input space contains Task parameters.
+            - ``'DNN'``: Dropout neural network. Uses MC sampling to mask neurons during training and
+              to estimate uncertainty.
+            
+            
         seed (int | None, optional): The random seed to use. Defaults to None.
         verbose (int, optional): The verbosity level. Defaults to 1.
 
@@ -60,7 +76,7 @@ class BayesianOptimizer(Optimizer):
     """
     def __init__(self,
                  X_space: ParamSpace,
-                 surrogate: str | dict | list[str] | list[dict] = ['GP'],
+                 surrogate: str | dict | list[str] | list[dict] = 'GP',
                  seed: int | None = None,
                  verbose: int = 1):
 
@@ -498,7 +514,7 @@ def _parse_aq_kwargs(self,
 
         # Improvement aqs based on inflation or deflation of best point
         if aq in ['EI', 'PI']:
-            o_max = o.max(dim=0).values * (1+hps['Xi_f'])
+            o_max = o.max(dim=0).values * (1+hps['inflate'])
             aq_kwargs.update({'best_f': o_max})
 
         # UCB based on: mu + sqrt(beta) * sqrt(variance) = mu + sqrt(beta) * sd
@@ -540,10 +556,10 @@ def _parse_aq_kwargs(self,
     def suggest(self,
                 m_batch: int = 1,
                 target: Target | list[Target] | None = None,
+                acquisition: list[str] | list[dict] = None,
                 optim_sequential: bool = True,
                 optim_samples: int = 512,
                 optim_restarts: int = 10,
-                acquisition: list[str] | list[dict] = None,
                 objective: MCAcquisitionObjective | None = None,
                 out_constraints: list[Callable] | None = None,
                 eq_constraints: tuple[Tensor, Tensor, float] | None = None,
@@ -558,18 +574,43 @@ def suggest(self,
         function calculated from the expectation of a surrogate model.
 
         Args:
-            m_batch (int, optional): The number of experiments to suggest at once. The default is `1`.
+            m_batch (int, optional): The number of experiments to suggest at once. The default is ``1``.
             target (Target or list of Target, optional): The response(s) to be used for optimization,
+            acquisition (list of str or list of dict, optional): Indicator for the desired acquisition function(s).
+                A list will propose experiments for each acquisition function based on ``optim_sequential``.
+                
+                The default is ``['NEI']`` for single-output and ``['NEHVI']`` for multi-output.
+                Options are as follows:
+                
+                - ``'EI'``: Expected Improvement (relative to best of ``y_train``). Accepts hyperparameter
+                  ``'inflate'``, a positive or negative float to inflate/deflate the best point for explore/exploit.
+                - ``'NEI'``: Noisy Expected Improvement. More robust than ``EI`` and uses all of ``y_train``,
+                  but accepts no hyperparameters.
+                - ``'PI'``: Probability of Improvement (relative to best of ``y_train``). Accepts hyperparameter
+                  ``'inflate'``, a positive or negative float to inflate/deflate the best point for explore/exploit.
+                - ``'UCB'``: Upper Confidence Bound. Accepts hyperparameter ``'beta'``, a positive float which sets
+                  the number of standard deviations above the mean.
+                - ``'SR'``: Simple Regret
+                - ``'RS'``: Random Sampling
+                - ``'Mean'``: Mean of the posterior distribution (pure exploitation/maximization of objective)
+                - ``'SF'``: Space Filling. Requests points that maximize the minimumd distance to ``X_train`` based
+                  on Euclidean distance.
+                - ``'NIPV'``: Negative Integrated Posterior Variance. Requests the point which most improves the prediction
+                  interval for a random selection of points in the design space. Used for active learning.
+                - ``'EHVI'``: Expected Hypervolume Improvement. Can accept a ``ref_point``, otherwise a point just
+                  below the minimum of ``y_train``.
+                - ``'NEHVI'``: Noisy Expected Hypervolume Improvement. Can accept a ``ref_point``, otherwise a point
+                  just below the minimum of ``y_train``.
+                - ``'NParEGO'``: Noisy Pareto Efficient Global Optimization. Can accept ``scalarization_weights``, a
+                  list of weights for each objective.
+                
             optim_sequential (bool, optional): Whether or not to optimize batch designs sequentially
                 (by fantasy) or simultaneously. Default is `True`.
             optim_samples (int, optional): The number of samples to use for quasi Monte Carlo sampling
                 of the acquisition function. Also used for initializing the acquisition optimizer.
                 The default value is `512`.
             optim_restarts (int, optional): The number of restarts to use in the global optimization
                 of the acquisition function. The default value is `10`.
-            acquisition (list of str or list of dict, optional): Indicator for the desired acquisition function(s).
-                The default is `['EI']`. A list will propose experiments for each acquisition function
-                in parallel and independent.
             objective (MCAcquisitionObjective, optional): The objective function to be used for optimization.
                 The default is `None`.
             out_constraints (list of Callable, optional): A list of constraints to be applied to the output space.
@@ -634,7 +675,7 @@ def suggest(self,
 
         # Default to noisy expected improvement if no aq method is provided
         if not acquisition:
-            acquisition = ['NEI'] if optim_type == 'single' else ['NEHVI']
+            acquisition = [aq_defaults[optim_type]]
 
         if not isinstance(acquisition, list):
             raise TypeError('acquisition must be a list of strings or dictionaries')