diff --git a/xopt/generators/bayesian/bax/algorithms.py b/xopt/generators/bayesian/bax/algorithms.py index f89fa5cdf..3dea6e016 100644 --- a/xopt/generators/bayesian/bax/algorithms.py +++ b/xopt/generators/bayesian/bax/algorithms.py @@ -1,11 +1,10 @@ from abc import ABC, abstractmethod -from typing import ClassVar, List +from typing import Any import torch from botorch.models.model import Model, ModelList from pydantic import BaseModel, ConfigDict, Field, PositiveInt, computed_field from torch import Tensor - from xopt.pydantic import XoptBaseModel @@ -56,7 +55,7 @@ class Algorithm(XoptBaseModel, ABC): Attributes ---------- - name : ClassVar[str] + name : str The name of the algorithm. n_samples : PositiveInt Number of execution paths to generate. @@ -69,12 +68,12 @@ class Algorithm(XoptBaseModel, ABC): Perform the virtual measurement and calculate objective values at the given inputs. """ - name: ClassVar[str] = "base_algorithm" + name: str = Field(default="base_algorithm", frozen=True) n_samples: PositiveInt = Field( default=20, description="number of execution paths to generate" ) - @computed_field + @computed_field # type: ignore[prop-decorator] @property def class_path(self) -> str: return f"{self.__class__.__module__}.{self.__class__.__name__}" @@ -105,7 +104,7 @@ def perform_virtual_measurement( x: Tensor, bounds: Tensor, n_samples: int, - tkwargs: dict = None, + tkwargs: dict[str, Any] | None = None, ) -> VirtualMeasurementResult: """ Evaluate the virtual objective at the given inputs. @@ -120,7 +119,7 @@ def perform_virtual_measurement( The bounds for the optimization. n_samples : int The number of samples to generate. - tkwargs : dict, optional + tkwargs : dict[str, Any] | None, optional Additional keyword arguments for the evaluation. Returns @@ -148,7 +147,7 @@ class GridScanAlgorithm(Algorithm, ABC): Create a mesh for evaluating posteriors on. """ - name = "grid_scan_algorithm" + name: str = Field(default="grid_scan", frozen=True) n_mesh_points: PositiveInt = Field( default=10, description="number of mesh points along each axis" ) @@ -200,13 +199,14 @@ class GridOptimize(GridScanAlgorithm): Methods ------- - get_execution_paths(self, model: Model, bounds: Tensor) -> Tuple[Tensor, Tensor, Dict] + get_execution_paths(self, model: Model, bounds: Tensor) -> tuple[Tensor, Tensor, ExecutionPathsResult] Get execution paths that minimize the objective function. perform_virtual_measurement(self, model: Model, x: Tensor, bounds: Tensor, n_samples: int, tkwargs: dict = None) -> VirtualMeasurementResult Evaluate the virtual measurement and calculate objective values (samples). """ - observable_names_ordered: List[str] = Field( + name: str = Field(default="grid_optimize", frozen=True) + observable_names_ordered: list[str] = Field( description="names of observable/objective models used in this algorithm", ) minimize: bool = True @@ -228,7 +228,7 @@ def execute(self, model: Model, bounds: Tensor) -> GridOptimizeResult: Contains best_inputs, best_objective, input_execution_paths, output_execution_paths, and additional results. """ # build evaluation mesh - test_points = self.create_mesh(bounds) + test_points: Tensor = self.create_mesh(bounds) if isinstance(model, ModelList): test_points = test_points.to(model.models[0].train_targets) else: @@ -273,7 +273,7 @@ def perform_virtual_measurement( x: Tensor, bounds: Tensor, n_samples: int, - tkwargs: dict = None, + tkwargs: dict[str, Any] | None = None, ) -> VirtualMeasurementResult: """ Perform the virtual measurement (samples). @@ -318,6 +318,7 @@ class CurvatureGridOptimize(GridOptimize): Perform the virtual measurement (samples) with curvature. """ + name: str = Field(default="curvature_grid_optimize", frozen=True) use_mean: bool = False def perform_virtual_measurement( @@ -326,7 +327,7 @@ def perform_virtual_measurement( x: Tensor, bounds: Tensor, n_samples: int, - tkwargs: dict = None, + tkwargs: dict[str, Any] | None = None, ) -> VirtualMeasurementResult: """ Evaluate the virtual objective (samples) with curvature. diff --git a/xopt/generators/bayesian/bax_generator.py b/xopt/generators/bayesian/bax_generator.py index ee0603493..eafbbf9a9 100644 --- a/xopt/generators/bayesian/bax_generator.py +++ b/xopt/generators/bayesian/bax_generator.py @@ -1,10 +1,11 @@ -from copy import deepcopy import importlib import logging import pickle -from typing import Dict, List, Optional +from copy import deepcopy +from typing import Any, Hashable, Optional, cast from botorch.models import ModelListGP, SingleTaskGP +from gpytorch import Module from pydantic import ( Field, SerializeAsAny, @@ -12,12 +13,18 @@ field_validator, model_validator, ) - +from pydantic.fields import ModelPrivateAttr, PrivateAttr +from xopt.errors import VOCSError from xopt.generators.bayesian.bax.acquisition import ModelListExpectedInformationGain from xopt.generators.bayesian.bax.algorithms import Algorithm, GridOptimize from xopt.generators.bayesian.bayesian_generator import BayesianGenerator -from xopt.generators.bayesian.turbo import EntropyTurboController, SafetyTurboController +from xopt.generators.bayesian.turbo import ( + EntropyTurboController, + SafetyTurboController, + TurboController, +) from xopt.generators.bayesian.utils import validate_turbo_controller_center +from xopt.vocs import VOCS logger = logging.getLogger() @@ -56,48 +63,89 @@ class BaxGenerator(BayesianGenerator): supports_no_objective: bool = True supports_discrete_variables: bool = False algorithm: SerializeAsAny[Algorithm] = Field( - description="algorithm evaluated in the BAX process" + default=GridOptimize(observable_names_ordered=[]), + description="algorithm evaluated in the BAX process", ) - algorithm_results: Optional[Dict] = Field( + algorithm_results: Optional[dict] = Field( None, description="dictionary results from algorithm", exclude=True ) algorithm_results_file: Optional[str] = Field( None, description="file name to save algorithm results at every step" ) _n_calls: int = 0 - _compatible_turbo_controllers = [EntropyTurboController, SafetyTurboController] + _compatible_turbo_controllers: list[type[TurboController]] = PrivateAttr( + default=[EntropyTurboController, SafetyTurboController] + ) # NOTE: this is meant for use in Badger, TODO: add it to Xopt - _compatible_algorithms = [GridOptimize] + _compatible_algorithms: list[type[Algorithm]] = PrivateAttr(default=[GridOptimize]) @model_validator(mode="after") - def validate_model_after(self): + def validate_model_after(self) -> "BaxGenerator": # validate turbo controller center if it exists validate_turbo_controller_center(self) return self + @field_validator("vocs", mode="after") + @classmethod + def validate_vocs(cls, v: VOCS, info: ValidationInfo) -> VOCS: + # Preserve inherited Bayesian VOCS validation behavior. + # v = super().validate_vocs(v, info) + + # assert that the generator had no objectives + if not v.n_objectives == 0: + raise VOCSError("BAX generator only supports problems with no objectives") + + return v + @field_validator("algorithm", mode="before") - def validate_algorithm(cls, v, info: ValidationInfo): + @classmethod + def validate_algorithm(cls, v: Any, info: ValidationInfo) -> Any: if isinstance(v, dict): - try: + if "class_path" in v: class_path = v.pop("class_path") module_name, class_name = class_path.rsplit(".", 1) - except KeyError: - raise ValueError("Algorithm dictionary must contain 'class_path' key") - - try: - algorithm_class = getattr( - importlib.import_module(module_name), class_name + try: + algorithm_class = getattr( + importlib.import_module(module_name), class_name + ) + except ModuleNotFoundError: + raise ValueError(f"Cannot import '{module_name}.{class_name}'") + elif "name" in v: + name = v["name"] + algorithm_class = next( + ( + c + for c in cls._compatible_algorithms.default + if c.model_fields["name"].default == name + ), + None, + ) + if algorithm_class is None: + raise ValueError( + f"Unknown algorithm name '{name}'. " + f"Provide one of {[c.model_fields['name'].default for c in cls._compatible_algorithms.default]} " + f"or supply 'class_path'." + ) + else: + raise ValueError( + "Algorithm dictionary must contain 'class_path' or 'name' key" ) - except ModuleNotFoundError: - raise ValueError(f"Cannot import '{module_name}.{class_name}'") v = algorithm_class.model_validate(v) return v - def generate(self, n_candidates: int) -> List[Dict]: + @classmethod + def get_compatible_algorithms(cls) -> list[type[Algorithm]]: + compatible = cls._compatible_algorithms + compatible_list: list[type[Algorithm]] = [] + if isinstance(compatible, ModelPrivateAttr): + compatible_list = cast(list[type[Algorithm]], compatible.get_default()) + return compatible_list + + def generate(self, n_candidates: int) -> list[dict[Hashable, Any]]: """ Generate a specified number of candidate samples. @@ -108,19 +156,19 @@ def generate(self, n_candidates: int) -> List[Dict]: Returns ------- - List[Dict] + list[dict[Hashable, Any]] A list of dictionaries containing the generated samples. """ self._n_calls += 1 return super().generate(n_candidates) - def _get_acquisition(self, model) -> ModelListExpectedInformationGain: + def _get_acquisition(self, model: Module) -> ModelListExpectedInformationGain: """ Get the acquisition function. Parameters ---------- - model : Model + model : Module The model to use for the acquisition function. Returns diff --git a/xopt/generators/bayesian/bayesian_generator.py b/xopt/generators/bayesian/bayesian_generator.py index f9ced1f42..cc1e637fd 100644 --- a/xopt/generators/bayesian/bayesian_generator.py +++ b/xopt/generators/bayesian/bayesian_generator.py @@ -1,39 +1,37 @@ import logging -from math import prod import os import time import warnings -from copy import deepcopy from abc import ABC, abstractmethod +from copy import deepcopy from itertools import islice, product -from typing import Any, Dict, List, Optional, Union, cast +from math import prod +from typing import Any, Dict, Hashable, List, Optional, Union, cast import numpy as np import pandas as pd import torch from botorch.acquisition import ( - FixedFeatureAcquisitionFunction, - qUpperConfidenceBound, AcquisitionFunction, + FixedFeatureAcquisitionFunction, MCAcquisitionObjective, + qUpperConfidenceBound, ) from botorch.models.model import Model from botorch.sampling.get_sampler import get_sampler +from gest_api.vocs import VOCS, DiscreteVariable, MaximizeObjective, MinimizeObjective from gpytorch import Module from pydantic import ( Field, - field_validator, PositiveInt, SerializeAsAny, + field_validator, model_validator, ) -from pydantic.fields import PrivateAttr, ModelPrivateAttr +from pydantic.fields import ModelPrivateAttr, PrivateAttr from pydantic_core.core_schema import ValidationInfo from torch import Tensor - -from gest_api.vocs import DiscreteVariable, MinimizeObjective, MaximizeObjective - -from xopt.errors import VOCSError, XoptError, FeasibilityError +from xopt.errors import FeasibilityError, VOCSError, XoptError from xopt.generator import Generator from xopt.generators.bayesian.base_model import ModelConstructor from xopt.generators.bayesian.custom_botorch.constrained_acquisition import ( @@ -49,19 +47,19 @@ ) from xopt.generators.bayesian.models.time_dependent import TimeDependentModelConstructor from xopt.generators.bayesian.objectives import ( + CustomXoptObjective, create_constraint_callables, create_mc_objective, - CustomXoptObjective, ) from xopt.generators.bayesian.turbo import ( TurboController, ) from xopt.generators.bayesian.utils import ( + compute_hypervolume_and_pf, interpolate_points, rectilinear_domain_union, set_botorch_weights, validate_turbo_controller_base, - compute_hypervolume_and_pf, validate_turbo_controller_center, ) from xopt.generators.bayesian.visualize import visualize_generator_model @@ -76,7 +74,6 @@ has_discrete_variables, ) - logger = logging.getLogger() # It seems pydantic v2 does not auto-register models anymore @@ -203,6 +200,29 @@ class BayesianGenerator(Generator, ABC): default=[LBFGSOptimizer, GridOptimizer] ) + @field_validator("vocs", mode="after") + @classmethod + def validate_vocs(cls, v: VOCS, info: ValidationInfo) -> VOCS: + if v.n_constraints > 0 and not info.data["supports_constraints"]: + raise VOCSError("this generator does not support constraints") + + if has_discrete_variables(v) and not info.data["supports_discrete_variables"]: + raise VOCSError("this generator does not support discrete variables") + + # assertion that at least one objective exists is done in model_validator below + + if v.n_objectives == 1: + if not info.data["supports_single_objective"]: + raise VOCSError( + "this generator does not support single objective optimization" + ) + elif v.n_objectives > 1 and not info.data["supports_multi_objective"]: + raise VOCSError( + "this generator does not support multi-objective optimization" + ) + + return v + @classmethod def get_compatible_turbo_controllers(cls) -> list[type[TurboController] | None]: compatible = cls._compatible_turbo_controllers @@ -226,7 +246,7 @@ def get_compatible_numerical_optimizers( @field_validator("model", mode="before") @classmethod - def validate_torch_modules(cls, value: Any): + def validate_torch_modules(cls, value: Any) -> Any: if isinstance(value, str): if value.startswith("base64:"): value = decode_torch_module(value) @@ -238,7 +258,7 @@ def validate_torch_modules(cls, value: Any): @field_validator("gp_constructor", mode="before") @classmethod - def validate_gp_constructor(cls, value: Any): + def validate_gp_constructor(cls, value: Any) -> Any: constructor_dict = { "standard": StandardModelConstructor, "batched": BatchedModelConstructor, @@ -266,7 +286,7 @@ def validate_gp_constructor(cls, value: Any): @field_validator("numerical_optimizer", mode="before") @classmethod - def validate_numerical_optimizer(cls, value: Any): + def validate_numerical_optimizer(cls, value: Any) -> Any: optimizer_dict: dict[str, type[NumericalOptimizer]] = { "grid": GridOptimizer, "LBFGS": LBFGSOptimizer, @@ -293,7 +313,7 @@ def validate_numerical_optimizer(cls, value: Any): @field_validator("turbo_controller", mode="before") @classmethod - def validate_turbo_controller(cls, value: Any, info: ValidationInfo): + def validate_turbo_controller(cls, value: Any, info: ValidationInfo) -> Any: """note default behavior is no use of turbo""" if value is None: return value @@ -313,7 +333,7 @@ def validate_turbo_controller(cls, value: Any, info: ValidationInfo): @field_validator("computation_time", mode="before") @classmethod - def validate_computation_time(cls, value: Any): + def validate_computation_time(cls, value: Any) -> Any: if value is None: return value elif isinstance(value, pd.DataFrame): @@ -328,7 +348,7 @@ def validate_computation_time(cls, value: Any): return value @model_validator(mode="after") - def validate_model_after(self): + def validate_model_after(self) -> "BayesianGenerator": # validate turbo controller center if it exists validate_turbo_controller_center(self) @@ -346,7 +366,7 @@ def validate_model_after(self): return self - def add_data(self, new_data: pd.DataFrame): + def add_data(self, new_data: pd.DataFrame) -> None: """ Add new data to the generator for Bayesian Optimization. @@ -361,7 +381,7 @@ def add_data(self, new_data: pd.DataFrame): """ self.data = pd.concat([self.data, new_data], axis=0, ignore_index=True) - def generate(self, n_candidates: int): + def generate(self, n_candidates: int) -> list[dict[Hashable, Any]]: """ Generate candidates using Bayesian Optimization. @@ -372,7 +392,7 @@ def generate(self, n_candidates: int): Returns ------- - List[Dict] + list[dict[Hashable, Any]] A list of dictionaries containing the generated candidates. Raises diff --git a/xopt/generators/bayesian/utils.py b/xopt/generators/bayesian/utils.py index 286571ecc..38ce6520c 100644 --- a/xopt/generators/bayesian/utils.py +++ b/xopt/generators/bayesian/utils.py @@ -5,23 +5,22 @@ import gpytorch import numpy as np import pandas as pd -from pydantic import ValidationInfo import torch from botorch.acquisition import AcquisitionFunction from botorch.models import ModelListGP from botorch.models.model import Model from botorch.models.utils import multioutput_to_batch_mode_transform -from botorch.utils.multi_objective import is_non_dominated, Hypervolume - -from gest_api.vocs import MinimizeObjective, MaximizeObjective, ExploreObjective - +from botorch.utils.multi_objective import Hypervolume, is_non_dominated +from gest_api.vocs import ExploreObjective, MaximizeObjective, MinimizeObjective +from pydantic import ValidationInfo +from xopt.generator import Generator from xopt.generators.bayesian.turbo import TurboController from xopt.vocs import VOCS, random_inputs def get_training_data( input_names: List[str], outcome_name: str, data: pd.DataFrame -) -> (torch.Tensor, torch.Tensor): +) -> tuple[torch.Tensor, torch.Tensor]: """ Creates training data from input data frame. @@ -215,7 +214,7 @@ def rectilinear_domain_union(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: return out_bounds -def interpolate_points(df, num_points=10): +def interpolate_points(df: pd.DataFrame, num_points: int = 10) -> pd.DataFrame: """ Generates interpolated points between two points specified by a pandas DataFrame. @@ -309,7 +308,7 @@ def validate_turbo_controller_base( ) -def validate_turbo_controller_center(generator): +def validate_turbo_controller_center(generator: Generator) -> None: if generator.turbo_controller is not None: # Check that values for center_x are within trust region bounds trust_region = generator.turbo_controller.get_trust_region(generator) diff --git a/xopt/pydantic.py b/xopt/pydantic.py index ae1501c0a..b1a7544b3 100644 --- a/xopt/pydantic.py +++ b/xopt/pydantic.py @@ -29,8 +29,8 @@ from pydantic import ( BaseModel, ConfigDict, - create_model, Field, + create_model, field_serializer, field_validator, model_serializer, @@ -253,7 +253,7 @@ class XoptBaseModel(BaseModel): @model_validator(mode="before") @classmethod - def validate_files(cls, data: Any): + def validate_files(cls, data: Any) -> Any: if not isinstance(data, dict): return data for key, value in data.items(): @@ -273,13 +273,13 @@ def validate_files(cls, data: Any): def serialize_json(self, sinfo: SerializationInfo) -> dict: return orjson_dumps_except_root(self) - def to_json(self, **kwargs) -> str: + def to_json(self, **kwargs: Any) -> str: return orjson_dumps(self, **kwargs) - def json(self, **kwargs): + def json(self, **kwargs: Any) -> str: return self.to_json(**kwargs) - def yaml(self, **kwargs): + def yaml(self, **kwargs: Any) -> str: """serialize first then dump to yaml string""" output = json.loads( self.to_json( @@ -289,7 +289,7 @@ def yaml(self, **kwargs): return yaml.dump(output) @classmethod - def from_file(cls, filename: str): + def from_file(cls, filename: str) -> "XoptBaseModel": if not os.path.exists(filename): raise OSError(f"file {filename} is not found") @@ -297,11 +297,11 @@ def from_file(cls, filename: str): return cls.from_yaml(file) @classmethod - def from_yaml(cls, yaml_obj: [str, TextIO]): + def from_yaml(cls, yaml_obj: str | TextIO) -> "XoptBaseModel": return cls.model_validate(remove_none_values(yaml.safe_load(yaml_obj))) @classmethod - def from_dict(cls, config: dict): + def from_dict(cls, config: dict) -> "XoptBaseModel": return cls.model_validate(remove_none_values(config))