Badger integration

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",

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

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,

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.
         """

bax_algorithms/visualize.py

@@ -91,15 +91,17 @@ 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)
 
     # 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 = {}