From 3fa3c24452d20e9721bdaffe258785963868039e Mon Sep 17 00:00:00 2001
From: Mitchell Victoriano <mitchell.victoriano@gmail.com>
Date: Tue, 16 Jun 2026 16:57:05 -0700
Subject: [PATCH 1/4] Changes to support badger integration

---
 xopt/generators/bayesian/bax/algorithms.py    | 60 ++++++++++--------
 xopt/generators/bayesian/bax_generator.py     | 63 ++++++++++++-------
 .../generators/bayesian/bayesian_generator.py | 51 ++++++++-------
 xopt/generators/bayesian/utils.py             | 15 +++--
 4 files changed, 111 insertions(+), 78 deletions(-)

diff --git a/xopt/generators/bayesian/bax/algorithms.py b/xopt/generators/bayesian/bax/algorithms.py
index bace5ca5b..9b8c105d7 100644
--- a/xopt/generators/bayesian/bax/algorithms.py
+++ b/xopt/generators/bayesian/bax/algorithms.py
@@ -1,21 +1,30 @@
 from abc import ABC, abstractmethod
-from typing import ClassVar, Dict, List, Tuple
+from typing import Any, TypedDict
 
+import numpy as np
 import torch
 from botorch.models.model import Model, ModelList
+from numpy.typing import NDArray
 from pydantic import Field, PositiveInt, computed_field
 from torch import Tensor
-
 from xopt.pydantic import XoptBaseModel
 
 
+class ExecutionPathsResult(TypedDict):
+    test_points: Tensor
+    posterior_samples: Tensor
+    execution_paths: Tensor
+    solution_center: NDArray[np.floating]
+    solution_entropy: float
+
+
 class Algorithm(XoptBaseModel, ABC):
     """
     Base class for algorithms used in BAX.
 
     Attributes
     ----------
-    name : ClassVar[str]
+    name : str
         The name of the algorithm.
     n_samples : PositiveInt
         Number of execution paths to generate.
@@ -28,20 +37,20 @@ class Algorithm(XoptBaseModel, ABC):
         Evaluate the virtual objective 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
     @property
+    @computed_field
     def class_path(self) -> str:
         return f"{self.__class__.__module__}.{self.__class__.__name__}"
 
     @abstractmethod
     def get_execution_paths(
         self, model: Model, bounds: Tensor
-    ) -> Tuple[Tensor, Tensor, Dict]:
+    ) -> tuple[Tensor, Tensor, ExecutionPathsResult]:
         """
         Get execution paths for the algorithm.
 
@@ -54,7 +63,7 @@ def get_execution_paths(
 
         Returns
         -------
-        Tuple[Tensor, Tensor, Dict]
+        tuple[Tensor, Tensor, ExecutionPathsResult]
             The execution paths, their corresponding values, and additional results.
         """
         pass
@@ -66,7 +75,7 @@ def evaluate_virtual_objective(
         x: Tensor,
         bounds: Tensor,
         n_samples: int,
-        tkwargs: dict = None,
+        tkwargs: dict[str, Any] | None = None,
     ) -> Tensor:
         """
         Evaluate the virtual objective at the given inputs.
@@ -81,7 +90,7 @@ def evaluate_virtual_objective(
             The bounds for the optimization.
         n_samples : int
             The number of samples to generate.
-        tkwargs : dict, optional
+        tkwargs : dict[str, Any], optional
             Additional keyword arguments for the evaluation.
 
         Returns
@@ -109,7 +118,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"
     )
@@ -161,13 +170,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.
-    evaluate_virtual_objective(self, model: Model, x: Tensor, bounds: Tensor, n_samples: int, tkwargs: dict = None) -> Tensor
+    evaluate_virtual_objective(self, model: Model, x: Tensor, bounds: Tensor, n_samples: int, tkwargs: dict[str, Any] | None = None) -> Tensor
         Evaluate the virtual objective (samples).
     """
 
-    observable_names_ordered: List[str] = Field(
+    name: str = Field(default="grid_optimize", frozen=True)
+    observable_names_ordered: list[str] = Field(
         default=["y1"],
         description="names of observable/objective models used in this algorithm",
     )
@@ -175,7 +185,7 @@ class GridOptimize(GridScanAlgorithm):
 
     def get_execution_paths(
         self, model: Model, bounds: Tensor
-    ) -> Tuple[Tensor, Tensor, Dict]:
+    ) -> tuple[Tensor, Tensor, ExecutionPathsResult]:
         """
         Get execution paths that minimize the objective function.
 
@@ -188,11 +198,11 @@ def get_execution_paths(
 
         Returns
         -------
-        Tuple[Tensor, Tensor, Dict]
+        tuple[Tensor, Tensor, ExecutionPathsResult]
             The execution paths, their corresponding values, 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:
@@ -218,13 +228,13 @@ def get_execution_paths(
         solution_entropy = float(torch.log(x_opt.std(dim=0) ** 2).sum())
 
         # collect secondary results in a dict
-        results_dict = {
-            "test_points": test_points,
-            "posterior_samples": posterior_samples,
-            "execution_paths": torch.hstack((x_opt, y_opt)),
-            "solution_center": solution_center,
-            "solution_entropy": solution_entropy,
-        }
+        results_dict = ExecutionPathsResult(
+            test_points=test_points,
+            posterior_samples=posterior_samples,
+            execution_paths=torch.hstack((x_opt, y_opt)),
+            solution_center=solution_center,
+            solution_entropy=solution_entropy,
+        )
 
         # return execution paths
         return x_opt.unsqueeze(-2), y_opt.unsqueeze(-2), results_dict
@@ -235,7 +245,7 @@ def evaluate_virtual_objective(
         x: Tensor,
         bounds: Tensor,
         n_samples: int,
-        tkwargs: dict = None,
+        tkwargs: dict[str, Any] | None = None,
     ) -> Tensor:
         """
         Evaluate the virtual objective (samples).
@@ -289,7 +299,7 @@ def evaluate_virtual_objective(
         x: Tensor,
         bounds: Tensor,
         n_samples: int,
-        tkwargs: dict = None,
+        tkwargs: dict[str, Any] | None = None,
     ) -> Tensor:
         """
         Evaluate the virtual objective (samples) with curvature.
diff --git a/xopt/generators/bayesian/bax_generator.py b/xopt/generators/bayesian/bax_generator.py
index b942d4bf3..c050f8134 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,13 +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()
 
@@ -54,26 +60,30 @@ class BaxGenerator(BayesianGenerator):
 
     name = "bax"
     supports_constraints: bool = True
+
     supports_discrete_variables: bool = False
     algorithm: SerializeAsAny[Algorithm] = Field(
-        description="algorithm evaluated in the BAX process"
+        default=GridOptimize(), 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])
 
     @field_validator("vocs", mode="after")
-    def validate_vocs(cls, v, info: ValidationInfo):
+    @classmethod
+    def validate_vocs(cls, v: VOCS, info: ValidationInfo) -> VOCS:
         # Preserve inherited Bayesian VOCS validation behavior.
-        v = super().validate_vocs(v, info)
+        # v = super().validate_vocs(v, info)
 
         # assert that the generator had no objectives
         if not v.n_objectives == 0:
@@ -81,15 +91,9 @@ def validate_vocs(cls, v, info: ValidationInfo):
 
         return v
 
-    @model_validator(mode="after")
-    def validate_model_after(self):
-        # validate turbo controller center if it exists
-        validate_turbo_controller_center(self)
-
-        return self
-
     @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:
                 class_path = v.pop("class_path")
@@ -108,7 +112,22 @@ def validate_algorithm(cls, v, info: ValidationInfo):
 
         return v
 
-    def generate(self, n_candidates: int) -> List[Dict]:
+    @model_validator(mode="after")
+    def validate_model_after(self) -> "BaxGenerator":
+        # validate turbo controller center if it exists
+        validate_turbo_controller_center(self)
+
+        return self
+
+    @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.
 
@@ -119,19 +138,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 3615a7135..e22e27d1a 100644
--- a/xopt/generators/bayesian/bayesian_generator.py
+++ b/xopt/generators/bayesian/bayesian_generator.py
@@ -1,38 +1,36 @@
 import logging
-from math import prod
 import os
 import time
 import warnings
 from abc import ABC, abstractmethod
 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 (
@@ -48,19 +46,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
@@ -74,7 +72,6 @@
     has_discrete_variables,
 )
 
-
 logger = logging.getLogger()
 
 # It seems pydantic v2 does not auto-register models anymore
@@ -198,7 +195,8 @@ class BayesianGenerator(Generator, ABC):
     )
 
     @field_validator("vocs", mode="after")
-    def validate_vocs(cls, v, info: ValidationInfo):
+    @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")
 
@@ -242,8 +240,15 @@ 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):
+            null_options = [
+                "",
+                "none",
+                "null",
+            ]
+            if value.lower() in null_options:
+                return None
             if value.startswith("base64:"):
                 value = decode_torch_module(value)
             elif os.path.exists(value):
@@ -254,7 +259,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,
@@ -282,7 +287,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,
@@ -309,7 +314,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
@@ -329,7 +334,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):
@@ -344,7 +349,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)
 
@@ -362,7 +367,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.
 
@@ -377,7 +382,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.
 
@@ -388,7 +393,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 bba64f707..273d093ae 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)

From 7a42ab79602396f76ff335ff24084a3ed6a779e9 Mon Sep 17 00:00:00 2001
From: Mitchell Victoriano <mitchell.victoriano@gmail.com>
Date: Tue, 16 Jun 2026 17:46:02 -0700
Subject: [PATCH 2/4] Fixed missing class_path for validation

---
 xopt/generators/bayesian/bax/algorithms.py    |  2 +-
 xopt/generators/bayesian/bax_generator.py     | 35 ++++++++++++++-----
 .../generators/bayesian/bayesian_generator.py |  7 ----
 xopt/pydantic.py                              | 16 ++++-----
 4 files changed, 35 insertions(+), 25 deletions(-)

diff --git a/xopt/generators/bayesian/bax/algorithms.py b/xopt/generators/bayesian/bax/algorithms.py
index 9b8c105d7..57a7b59b3 100644
--- a/xopt/generators/bayesian/bax/algorithms.py
+++ b/xopt/generators/bayesian/bax/algorithms.py
@@ -42,8 +42,8 @@ class Algorithm(XoptBaseModel, ABC):
         default=20, description="number of execution paths to generate"
     )
 
+    @computed_field  # type: ignore[prop-decorator]
     @property
-    @computed_field
     def class_path(self) -> str:
         return f"{self.__class__.__module__}.{self.__class__.__name__}"
 
diff --git a/xopt/generators/bayesian/bax_generator.py b/xopt/generators/bayesian/bax_generator.py
index c050f8134..5d8d3b547 100644
--- a/xopt/generators/bayesian/bax_generator.py
+++ b/xopt/generators/bayesian/bax_generator.py
@@ -95,18 +95,35 @@ def validate_vocs(cls, v: VOCS, info: ValidationInfo) -> VOCS:
     @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)
 
diff --git a/xopt/generators/bayesian/bayesian_generator.py b/xopt/generators/bayesian/bayesian_generator.py
index e22e27d1a..3a8750942 100644
--- a/xopt/generators/bayesian/bayesian_generator.py
+++ b/xopt/generators/bayesian/bayesian_generator.py
@@ -242,13 +242,6 @@ def get_compatible_numerical_optimizers(
     @classmethod
     def validate_torch_modules(cls, value: Any) -> Any:
         if isinstance(value, str):
-            null_options = [
-                "",
-                "none",
-                "null",
-            ]
-            if value.lower() in null_options:
-                return None
             if value.startswith("base64:"):
                 value = decode_torch_module(value)
             elif os.path.exists(value):
diff --git a/xopt/pydantic.py b/xopt/pydantic.py
index e584c30f8..13483dfa0 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,
@@ -202,7 +202,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():
@@ -222,13 +222,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(
@@ -238,7 +238,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")
 
@@ -246,11 +246,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))
 
 

From dbe50bf2b8e5f5f925231bc6d690dc42a8cd1cf9 Mon Sep 17 00:00:00 2001
From: Mitchell Victoriano <mitchell.victoriano@gmail.com>
Date: Tue, 23 Jun 2026 18:55:41 -0700
Subject: [PATCH 3/4] Added missing type to tkwargs

---
 xopt/generators/bayesian/bax/algorithms.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/xopt/generators/bayesian/bax/algorithms.py b/xopt/generators/bayesian/bax/algorithms.py
index cbf7da017..b99585bde 100644
--- a/xopt/generators/bayesian/bax/algorithms.py
+++ b/xopt/generators/bayesian/bax/algorithms.py
@@ -104,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.
@@ -119,7 +119,7 @@ def perform_virtual_measurement(
             The bounds for the optimization.
         n_samples : int
             The number of samples to generate.
-        tkwargs : dict[str, Any], optional
+        tkwargs : dict[str, Any] | None, optional
             Additional keyword arguments for the evaluation.
 
         Returns

From 1d194b5619d2a3133d520cc8490a5f6c847e9fd9 Mon Sep 17 00:00:00 2001
From: Mitchell Victoriano <mitchell.victoriano@gmail.com>
Date: Wed, 1 Jul 2026 16:32:16 -0700
Subject: [PATCH 4/4] Added missing name field for curvature grid optimize

---
 xopt/generators/bayesian/bax/algorithms.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/xopt/generators/bayesian/bax/algorithms.py b/xopt/generators/bayesian/bax/algorithms.py
index b99585bde..3dea6e016 100644
--- a/xopt/generators/bayesian/bax/algorithms.py
+++ b/xopt/generators/bayesian/bax/algorithms.py
@@ -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(