diff --git a/bax_algorithms/emittance.py b/bax_algorithms/emittance.py index ef241c3..c142746 100644 --- a/bax_algorithms/emittance.py +++ b/bax_algorithms/emittance.py @@ -1,5 +1,6 @@ -from pydantic import Field, PositiveInt -from typing import Optional +from pydantic import Field, PositiveInt, field_validator, field_serializer +from typing import Optional, Any +import ast import torch from torch import Tensor @@ -461,38 +462,42 @@ class VirtualEmittanceMeasurementResult(VirtualMeasurementResult): class EmittanceAlgorithm(Algorithm): - x_key: str = Field( + name: str = Field("minimize_emittance", frozen=True) + x_key: str | None = Field( None, description="key designating the beamsize squared output in x from evaluate function", ) - y_key: str = Field( + + y_key: str | None = Field( None, description="key designating the beamsize squared output in y from evaluate function", ) energy: float = Field(1.0, description="Beam energy in [eV]") q_len: float = Field( - description="the longitudinal thickness of the measurement quadrupole" + 0.08, description="the longitudinal thickness of the measurement quadrupole" ) - rmat_x: Tensor = Field( - None, description="tensor shape 2x2 containing downstream rmat for x dimension" + rmat_x: Tensor | None = Field( + Tensor([[1.0, 1.0], [0.0, 1.0]]), + description="2x2 Tensor containing downstream rmat for x dimension", ) - rmat_y: Tensor = Field( - None, description="tensor shape 2x2 containing downstream rmat for y dimension" + rmat_y: Tensor | None = Field( + Tensor([[1.0, 1.0], [0.0, 1.0]]), + description="2x2 Tensor containing downstream rmat for y dimension", ) - twiss0_x: Tensor = Field( - None, - description="1d tensor length 2 containing design x-twiss: [beta0_x, alpha0_x] (for bmag)", + twiss0_x: Tensor | None = Field( + Tensor([1.0, 0.0]), + description="List length 2 containing design x-twiss: [beta0_x, alpha0_x] (for bmag)", ) - twiss0_y: Tensor = Field( - None, - description="1d tensor length 2 containing design y-twiss: [beta0_y, alpha0_y] (for bmag)", + twiss0_y: Tensor | None = Field( + Tensor([1.0, 0.0]), + description="List length 2 containing design y-twiss: [beta0_y, alpha0_y] (for bmag)", ) meas_dim: int = Field( - None, + 0, description="index identifying the measurement quad dimension in the model", ) n_steps_measurement_param: int = Field( - 3, description="number of steps to use in the virtual measurement scans" + 5, description="number of steps to use in the virtual measurement scans" ) thin_lens: bool = Field( False, @@ -502,9 +507,6 @@ class EmittanceAlgorithm(Algorithm): True, description="Whether to multiply the emit by the bmag to get virtual objective.", ) - results: dict = Field( - {}, description="Dictionary to store results from emittance calculcation" - ) maxiter_fit: int = Field( 20, description="Maximum number of iterations in nonlinear emittance fitting." ) @@ -513,6 +515,41 @@ class EmittanceAlgorithm(Algorithm): description="Whether to retain beamsize values only around the minimum from each scan.", ) + @field_validator("rmat_x", "rmat_y", "twiss0_x", "twiss0_y", mode="before") + @classmethod + def validate_tensors(cls, v: Any) -> Tensor: + """Accept tensors, (possibly nested) lists/tuples, or their string + representations (e.g. "[[1.0, 1.0], [0.0, 1.0]]" or "1.0, 0.0") and + convert them into a double-precision tensor.""" + if isinstance(v, Tensor): + return v + if isinstance(v, str): + stripped = v.strip() + if stripped.startswith("[") or stripped.startswith("("): + v = ast.literal_eval(stripped) + else: + v = [item for item in stripped.split(",")] + if isinstance(v, (list, tuple)): + float_list = cls._to_nested_floats(v) + return torch.tensor(float_list, dtype=torch.double) + raise ValueError(f"Cannot convert {v} to a Tensor.") + + @classmethod + def _to_nested_floats(cls, v: Any) -> Any: + """Recursively convert a (possibly nested) list/tuple to floats, + preserving the nesting structure.""" + if isinstance(v, (list, tuple)): + return [cls._to_nested_floats(item) for item in v] + return float(v) + + @field_serializer("rmat_x", "rmat_y", "twiss0_x", "twiss0_y") + def serialize_tensor(self, v: Any) -> Any: + """Serialize tensor fields to nested lists so they round-trip through + JSON/YAML (pydantic otherwise dumps unknown types as ``'torch.Tensor'``).""" + if isinstance(v, Tensor): + return v.tolist() + return v + @property def x_idx(self) -> int: """ @@ -790,6 +827,7 @@ def _crop_quad_scans( class PathwiseMinimizeEmittance(EmittanceAlgorithm, PathwiseOptimization): + name: str = Field("pathwise_minimize_emittance", frozen=True) n_batch: PositiveInt = Field( 1, description="Number of sample batches to optimize, with each batch containing self.n_samples", diff --git a/bax_algorithms/pathwise/base.py b/bax_algorithms/pathwise/base.py index 8020e3c..3f693b5 100644 --- a/bax_algorithms/pathwise/base.py +++ b/bax_algorithms/pathwise/base.py @@ -1,11 +1,12 @@ # to be added to basic algorithms in Xopt -from abc import abstractmethod -from bax_algorithms.pathwise.optimize import VirtualOptimizer, DifferentialEvolution +from bax_algorithms.pathwise.optimize import DifferentialEvolution from botorch.models.model import Model, ModelList from botorch.sampling.pathwise.posterior_samplers import draw_matheron_paths from pydantic import Field -from xopt.generators.bayesian.bax.algorithms import Algorithm +from xopt.generators.bayesian.bax.algorithms import ( + Algorithm, +) from torch import Tensor import torch from typing import List @@ -39,49 +40,30 @@ class PathwiseOptimization(Algorithm): Get the bounds for virtual optimization. """ - name = "pathwise_optimization" - optimizer: VirtualOptimizer = Field( + name: str = Field("pathwise_optimization", frozen=True) + optimizer: DifferentialEvolution = Field( DifferentialEvolution(), description="Optimizer for virtual objective." ) - results: dict = Field( - default=None, - description="dictionary containing algorithm results", - ) observable_names_ordered: List[str] = Field( - default=None, description="names of observable models used in this algorithm", ) - @abstractmethod - def perform_virtual_measurement( - self, - model: Model, - x: Tensor, - bounds: Tensor, - n_samples: int = None, - tkwargs: dict = None, - ) -> dict: - """ - Defines how the measurement of the virtual objective should be performed. - Stores results in a dictionary. - Returned dictionary must contain key 'objective' containing the virtual objective results. - """ - return {"objective": None} - def evaluate_virtual_objective( self, model: Model, x: Tensor, bounds: Tensor, - n_samples: int = None, - tkwargs: dict = None, + n_samples: int | None = None, ) -> Tensor: """ Performs virtual measurement and extracts virtual objective value from resultant dictionary. """ measurement_result = self.perform_virtual_measurement( - model, x, bounds, n_samples, tkwargs + model, + x, + bounds, + n_samples, ) return measurement_result.objective diff --git a/bax_algorithms/pathwise/optimize.py b/bax_algorithms/pathwise/optimize.py index cd311c1..e7b5686 100644 --- a/bax_algorithms/pathwise/optimize.py +++ b/bax_algorithms/pathwise/optimize.py @@ -31,6 +31,7 @@ def optimize( """ Minimizes virtual objective sample functions and returns optimal inputs. """ + raise NotImplementedError("This method should be implemented in subclasses.") @abstractmethod def _wrap_virtual_objective( @@ -44,6 +45,7 @@ def _wrap_virtual_objective( """ Wraps virtual objective function so inputs/outputs are suitable for optimization method. """ + raise NotImplementedError("This method should be implemented in subclasses.") @abstractmethod def _get_virtual_optimization_bounds( @@ -52,6 +54,7 @@ def _get_virtual_optimization_bounds( """ Get bounds for virtual optimization (may not be the same as bounds passed to optimizer). """ + raise NotImplementedError("This method should be implemented in subclasses.") def _get_target_function( self, diff --git a/bax_algorithms/solenoid_alignment.py b/bax_algorithms/solenoid_alignment.py index 347392c..42b88dc 100644 --- a/bax_algorithms/solenoid_alignment.py +++ b/bax_algorithms/solenoid_alignment.py @@ -19,21 +19,21 @@ class VirtualAlignmentMeasurementResult(VirtualMeasurementResult): class PathwiseSolenoidAlignment(PathwiseOptimization): - name: str = Field("PathwiseSolenoidAlignment", frozen=True) - x_key: str = Field( + name: str = Field("pathwise_solenoid_alignment", frozen=True) + x_key: str | None = Field( None, description="key designating the centroid position in x from evaluate function", ) - y_key: str = Field( + y_key: str | None = Field( None, description="key designating the centroid poisition in y from evaluate function", ) - meas_dim: int = Field( + meas_dim: int | None = Field( None, description="index identifying the measurement quad dimension in the model", ) n_steps_measurement_param: int = Field( - 3, description="number of steps to use in the virtual measurement scans" + 5, description="number of steps to use in the virtual measurement scans" ) n_batch: PositiveInt = Field( 1, @@ -57,8 +57,8 @@ def y_idx(self) -> int: return self.observable_names_ordered.index(self.y_key) def perform_virtual_measurement( - self, model, x, bounds, tkwargs: dict = None, n_samples: int = None - ): + self, model: Model, x: Tensor, bounds: Tensor, n_samples: int | None = None + ) -> VirtualAlignmentMeasurementResult: """ inputs: model: a botorch ModelListGP @@ -80,7 +80,6 @@ def perform_virtual_measurement( model, x_tuning, bounds, - tkwargs, n_samples, ) @@ -96,9 +95,7 @@ def perform_virtual_measurement( return virtual_alignment_result - def get_meas_scan_inputs( - self, x_tuning: Tensor, bounds: Tensor, tkwargs: dict = None - ): + def get_meas_scan_inputs(self, x_tuning: Tensor, bounds: Tensor) -> Tensor: """ A function that generates the inputs for virtual emittance measurement scans at the tuning configurations specified by x_tuning. @@ -119,10 +116,9 @@ def get_meas_scan_inputs( # expand the x tensor to represent quad measurement scans # at the locations in tuning parameter space specified by X - tkwargs = tkwargs if tkwargs else {"dtype": torch.double, "device": "cpu"} x_meas = torch.linspace( - *bounds.T[self.meas_dim], self.n_steps_measurement_param, **tkwargs + *bounds.T[self.meas_dim], self.n_steps_measurement_param ) # prepare column of measurement scans coordinates @@ -147,7 +143,11 @@ def get_meas_scan_inputs( return x def evaluate_posterior_misalignment( - self, model, x_tuning, bounds, tkwargs: dict = None, n_samples: int = None + self, + model: Model, + x_tuning: Tensor, + bounds: Tensor, + n_samples: int | None = None, ): """ inputs: @@ -158,10 +158,9 @@ def evaluate_posterior_misalignment( """ assert len(x_tuning.shape) in [2, 3] # x_tuning must be shape (n_tuning_configs, n_tuning_dims) or (n_samples, n_tuning_configs, ndim) - tkwargs = tkwargs if tkwargs else {"dtype": torch.double, "device": "cpu"} x = self.get_meas_scan_inputs( - x_tuning, bounds, tkwargs + x_tuning, bounds ) # result shape n_tuning_configs*n_steps x ndim centroid_position = self.evaluate_virtual_observables(model, x, n_samples) @@ -188,7 +187,7 @@ def evaluate_posterior_misalignment( return misalignment - def execute(self, model: Model, bounds: Tensor) -> Tensor: + def execute(self, model: Model, bounds: Tensor) -> OptimizationAlgorithmResult: best_tuning_inputs_list = [] best_objective_list = [] best_scan_inputs_list = [] @@ -234,7 +233,7 @@ def execute(self, model: Model, bounds: Tensor) -> Tensor: return algorithm_result - def _get_optimization_indeces(self, bounds) -> Tensor: + def _get_optimization_indeces(self, bounds: Tensor) -> Tensor: """ Get indeces specifying parameters for virtual objective optimization. """ diff --git a/bax_algorithms/visualize.py b/bax_algorithms/visualize.py index a1d9de1..5a20a00 100644 --- a/bax_algorithms/visualize.py +++ b/bax_algorithms/visualize.py @@ -91,7 +91,9 @@ def visualize_virtual_measurement_result( f"which are not in generator.vocs.variable_names." ) tkwargs = generator.tkwargs - x = _generate_input_mesh(vocs, variable_names, reference_point, n_grid, tkwargs) + x = _generate_input_mesh( + vocs, variable_names, data, reference_point, n_grid, tkwargs + ) # get bax observable models and bounds bax_model, bounds = get_bax_model_and_bounds(generator) @@ -99,7 +101,7 @@ def visualize_virtual_measurement_result( # get virtual measurement (sample) results kwargs = kwargs if kwargs else {} measurement_result = generator.algorithm.perform_virtual_measurement( - bax_model, x, bounds, tkwargs=tkwargs, n_samples=n_samples, **kwargs + bax_model, x, bounds, n_samples=n_samples, **kwargs ).model_dump() # create figure and subplots @@ -219,7 +221,7 @@ def plot_bax_objective_convergence( if file_name.startswith(file_prefix) ] file_names = sorted(file_names, key=lambda x: int(x[prefix_len:-ext_len])) - file_paths = [os.path.abspath(file_name) for file_name in file_names] + file_paths = [os.path.join(directory, file_name) for file_name in file_names] results_dicts = [] aggregated_results_dict = {} @@ -278,7 +280,7 @@ def plot_bax_input_convergence( if file_name.startswith(file_prefix) ] file_names = sorted(file_names, key=lambda x: int(x[prefix_len:-ext_len])) - file_paths = [os.path.abspath(file_name) for file_name in file_names] + file_paths = [os.path.join(directory, file_name) for file_name in file_names] results_dicts = [] aggregated_results_dict = {}