From 2f541c6d22755cfc33d17ab3998a406d9385f24e Mon Sep 17 00:00:00 2001
From: Konstantin Pieper <pieperk@ornl.gov>
Date: Fri, 17 Jul 2026 14:51:17 -0400
Subject: [PATCH 1/8] allow dataset_y updates with multiple values

---
 src/dial_dataclass/dial_dataclass.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/dial_dataclass/dial_dataclass.py b/src/dial_dataclass/dial_dataclass.py
index 8f45ebc..8f3ef6c 100644
--- a/src/dial_dataclass/dial_dataclass.py
+++ b/src/dial_dataclass/dial_dataclass.py
@@ -304,9 +304,13 @@ class DialWorkflowDatasetUpdate(BaseModel):
         description='The next collection of X values you want to append to your overall data',
         min_length=1,
     )
-    """the next collection of X values you want to append"""
-    next_y: float = Field(description='The next Y value you want to append to your overall data')
-    """the next Y value you want to append"""
+    next_y: Annotated[
+        float | list[float],
+        Field(
+            description=('The next Y value you want to append to your overall data'),
+        ),
+    ]
+
     kernel_args: dict[str, float | int | bool | str | list[float] | tuple] | None = Field(
         default=None
     )
@@ -326,7 +330,7 @@ class DialWorkflowDatasetUpdates(BaseModel):
 
     workflow_id: ValidatedObjectId
     next_x_list: list[list[float]] = Field(min_length=1)
-    next_y_list: list[float] = Field(min_length=1)
+    next_y_list: list[float | list[float]] = Field(min_length=1)
     kernel_args: dict[str, float | int | bool | str | list[float] | tuple] | None = None
     backend_args: dict[str, float | int | bool | str | list[float] | tuple] | None = None
     extra_args: dict[str, float | int | bool | str | list[float] | tuple] | None = None

From d60218d1a03a2e8a3bdc85e8a80ecdebd9723e02 Mon Sep 17 00:00:00 2001
From: Konstantin Pieper <pieperk@ornl.gov>
Date: Fri, 17 Jul 2026 12:50:10 -0400
Subject: [PATCH 2/8] update examples and benchmarks

- bump the SABLE version to get the newest features
- use the hyperparameter prior_std instead of alpha
- transition the sable backend and clients to use that feature, and
  clean up specification of unnecessary hyperparameters
- fix and adapt benchmarks to the new vesion
---
 scripts/1d_sable_client.py                 |  23 +-
 scripts/1d_sinusoidal_growth_client.py     |  27 +-
 src/dial_service/backends/sable_backend.py |  17 +-
 tests/benchmarks/test_rosenbrock.py        |  43 +--
 tests/benchmarks/test_strainmap.py         | 274 ++++++++--------
 uv.lock                                    | 363 ++++++++++++++++++++-
 6 files changed, 559 insertions(+), 188 deletions(-)

diff --git a/scripts/1d_sable_client.py b/scripts/1d_sable_client.py
index 727365b..1e3bed0 100644
--- a/scripts/1d_sable_client.py
+++ b/scripts/1d_sable_client.py
@@ -89,8 +89,6 @@ def __init__(self, service_destination: str):
         ]
         self.meshgrids = np.meshgrid(*self.grid_points, indexing='ij')
         self.x_grid = np.hstack([mg.reshape(-1, 1) for mg in self.meshgrids])
-        # Mirror the demo's RNG sequence before generating the noisy initial observations.
-        truth_model(self.x_grid[:, 0], 0.0, self.rng)
         self.y_raw, _ = truth_model(self.x_raw[:, 0], self.noise_level, self.rng)
 
         self.dataset_x = self.x_raw.reshape(-1, 1).tolist()
@@ -100,16 +98,16 @@ def __init__(self, service_destination: str):
         self.test_points = self.x_test.reshape(-1, 1).tolist()
 
         self.kernel = 'rbf'
-        self.kernel_args = {'x_range': [0.0, 1.0], 'sigma_range': [2.5e-3, 0.5], 'gamma': 0.1}
+        self.kernel_args = {
+            'sigma_range': [2.5e-3, 0.5],
+            'gamma': 0.1,
+        }
         self.backend = 'sable'
         self.backend_args = {
-            'n_features': 5000,
-            'alpha': 0.05,
-            'p': 1.25,
-            'n_iter_irls': 100,
+            'prior_std': 0.8,
         }
-        self.strategy = 'upper_confidence_bound'
-        self.strategy_args = {'exploit': 0.0, 'explore': 1.0}
+        self.strategy = 'uncertainty'
+        self.strategy_args = {}
         self.niter = 0
         self.max_iter = 30
         self.at_grids = True
@@ -246,13 +244,6 @@ def handle_next_points(self, payload):
         self.dataset_x.append(self.x_next)
         self.dataset_y.append(y_scalar)
 
-        # In this example we are running pure exploration, no optimization:
-        # optpos = np.argmax(self.dataset_y)
-        # y_opt = self.dataset_y[optpos]
-        # optimal_coords = self.dataset_x[optpos]
-        # coord_str = ', '.join([f'{coord:.2f}' for coord in optimal_coords])
-        # print(f'Optimal simulated datapoint at ({coord_str}), y={y_opt:.3f}\n')
-
     def graph(self):
         plt.clf()
 
diff --git a/scripts/1d_sinusoidal_growth_client.py b/scripts/1d_sinusoidal_growth_client.py
index 756f4e8..dd5d6d2 100644
--- a/scripts/1d_sinusoidal_growth_client.py
+++ b/scripts/1d_sinusoidal_growth_client.py
@@ -76,13 +76,18 @@ def __init__(self, service_destination: str):
         # Assume that there is some small noise in the measurements to stabilize the fit
         self.statistics_y = Normal(loc='y', scale=1e-6)
 
-        self.backend = 'sklearn'
+        # Set the prior standard deviation of the surrogate model
+        self.prior_std = 1.0
+
+        self.backend = 'sklearn'  # 'sklearn' or 'sable'
         if self.backend == 'sklearn':
             # configure kernel_hyperparameters
             self.kernel = 'matern'  # 'rbf' or 'matern'
+            # set kernel hyperparameters
+            prior_variance = self.prior_std**2
             self.kernel_args = {
                 'length_scale': 0.1,
-                'constant_value': 1.0,
+                'constant_value': prior_variance,
             }
             self.optimize_lengthscale = False
             if self.optimize_lengthscale:
@@ -95,19 +100,18 @@ def __init__(self, service_destination: str):
         elif self.backend == 'sable':
             self.kernel = 'rbf'
             self.kernel_args = {
-                'x_range': [-2.0, 2.0],
                 'sigma_range': [1.0e-3, 1.0],
-                'gamma': 0.1,
             }
+            # increase the prior standard deviation to avoid overconfidence
+            self.prior_std *= 5.0
             self.backend_args = {
-                'n_features': 10000,
-                'alpha': 0.0005,
-                'p': 1.25,
-                'n_iter_irls': 20,
+                'prior_std': self.prior_std,
             }
 
-        self.strategy = 'upper_confidence_bound'
-        self.strategy_args = {'exploit': 0.4, 'explore': 1}
+        strat = ('upper_confidence_bound', {'exploit': 0.4, 'explore': 1})
+        # strat = ('expected_improvement', {})
+        self.strategy, self.strategy_args = strat
+
         self.niter = 0
         self.max_iter = 20
         self.at_grids = True
@@ -177,7 +181,6 @@ def callback_message(self, operation: str, **kwargs) -> IntersectClientCallback:
                 kernel_args=self.kernel_args,
                 backend=self.backend,
                 backend_args=self.backend_args,
-                preprocess_standardize=True,
                 y_is_good=True,
             )
 
@@ -265,7 +268,7 @@ def graph(self):
             self.mean_grid + 2 * self.stddev_grid,
             self.mean_grid - 2 * self.stddev_grid,
             alpha=0.5,
-            label='Confidence Interval',
+            label='$2\\sigma$ Confidence Interval',
         )
         axs[0].scatter(
             np.array(self.dataset_x)[:-1, 0],
diff --git a/src/dial_service/backends/sable_backend.py b/src/dial_service/backends/sable_backend.py
index ba637b5..5762919 100644
--- a/src/dial_service/backends/sable_backend.py
+++ b/src/dial_service/backends/sable_backend.py
@@ -20,12 +20,19 @@
 
 def _get_model_kwargs(data) -> dict:
     backend_args = {} if data.backend_args is None else data.backend_args
-    return {
-        'n_features': backend_args.get('n_features', 10000),
-        'alpha': backend_args.get('alpha', 0.3),
-        'p': backend_args.get('p', 1.25),
-        'n_iter_irls': backend_args.get('n_iter_irls', 100),
+    # set some default args
+    model_kwargs = {
+        # default to a low number of features in general
+        'n_features': backend_args.get('n_features', 5000),
+        # set the prior standard deviation to a default value of 1.
+        # this works well for data where np.std(data.Y_train) ~ 1 (after output normalization)
+        'prior_std': 1.0,
+        # default to moderate sparsity (p = 1.2) instead of full sparsity (p = 1.) for now
+        'p': backend_args.get('p', 1.2),
     }
+    # update and add any user-supplied args
+    model_kwargs.update(backend_args)
+    return model_kwargs
 
 
 def _get_observation_errors(data, n_observations: int) -> np.ndarray:
diff --git a/tests/benchmarks/test_rosenbrock.py b/tests/benchmarks/test_rosenbrock.py
index 690b556..3543dc9 100644
--- a/tests/benchmarks/test_rosenbrock.py
+++ b/tests/benchmarks/test_rosenbrock.py
@@ -12,6 +12,7 @@
 # from pytest_benchmark.fixture import BenchmarkFixture
 from dial_dataclass import (
     DialInputSingleOtherStrategy,
+    Normal,
 )
 from dial_service import (
     core as dial_core,
@@ -89,15 +90,12 @@ def run_simulation(
     client_state = DialWorkflowCreationParamsService(
         dataset_x=dataset_x,
         dataset_y=dataset_y,
+        statistics_y=Normal(loc='y', scale=NOISE_LEVEL),
         bounds=INITIAL_BOUNDS,
         kernel='rbf',
         kernel_args={
             'length_scale': LENGTH_SCALE,
-            'length_scale_bounds': 'fixed',
-            'noise_level': NOISE_LEVEL,
-            'noise_level_bounds': 'fixed',
             'constant_value': CONSTANT_VALUE,
-            'constant_value_bounds': 'fixed',
         },
         y_is_good=False,  # we wish to minimize y (the error)
         backend='sklearn',  # "sklearn" or "gpax"
@@ -252,6 +250,23 @@ def test_benchmark_rosenbrock_accuracy(
     )
 
 
+def _run_single(task: tuple) -> tuple:
+    """Worker: run one accuracy benchmark and return (strategy_name, iterations, target, guess, history)."""
+    import json
+
+    strategy, strategy_args, run_index, dataset_x = task
+    strategy_name = f'{strategy}' + (f' {json.dumps(strategy_args)}' if strategy_args else '')
+    iterations, target, guess, history = accuracy_benchmark(strategy, strategy_args, dataset_x)
+    logger.info(
+        '  [%s] Run %d: iterations=%d, target=%.4f',
+        strategy_name,
+        run_index + 1,
+        iterations,
+        target,
+    )
+    return (strategy_name, strategy, strategy_args, iterations, target, guess, history)
+
+
 if __name__ == '__main__':
     """Generate HTML benchmark report with plots comparing different strategies."""
     import argparse
@@ -303,31 +318,21 @@ def positive_int_type(arg):
         'Running benchmarks with %d iterations per strategy using multiprocessing...', NUM_RUNS
     )
 
-    def _run_single(task: tuple) -> tuple:
-        """Worker: run one accuracy benchmark and return (strategy_name, iterations, target, guess, history)."""
-        strategy, strategy_args, run_index, dataset_x = task
-        strategy_name = f'{strategy}' + (f' {json.dumps(strategy_args)}' if strategy_args else '')
-        iterations, target, guess, history = accuracy_benchmark(strategy, strategy_args, dataset_x)
-        logger.info(
-            '  [%s] Run %d: iterations=%d, target=%.4f',
-            strategy_name,
-            run_index + 1,
-            iterations,
-            target,
-        )
-        return (strategy_name, strategy, strategy_args, iterations, target, guess, history)
-
     # Use the same dataset for each parameter we have
     datasets = [generate_initial_dataset() for _ in range(NUM_RUNS)]
     tasks = [
-        (strategy, strategy_args, run_index, dataset)
+        (strategy, strategy_args, run_index, dataset.copy())
         for strategy, strategy_args in parameters
         for run_index, dataset in enumerate(datasets)
     ]
 
+    # run in parallel
     with Pool() as pool:
         run_outputs = pool.map(_run_single, tasks)
 
+    # run in serial
+    # run_outputs = [_run_single(task) for task in tasks]
+
     # Aggregate results preserving strategy order
     strategy_order = [
         f'{strategy}' + (f' {json.dumps(strategy_args)}' if strategy_args else '')
diff --git a/tests/benchmarks/test_strainmap.py b/tests/benchmarks/test_strainmap.py
index 2bca51c..e91fa3e 100644
--- a/tests/benchmarks/test_strainmap.py
+++ b/tests/benchmarks/test_strainmap.py
@@ -16,6 +16,7 @@
 from dial_dataclass import (
     DialInputPredictions,
     DialInputSingleOtherStrategy,
+    Normal,
 )
 from dial_service import (
     core as dial_core,
@@ -42,102 +43,111 @@
 )
 
 
-def normalize_data(in_data):
-    m_data = (in_data.max() + in_data.min()) * 0.5
-    d_data = (in_data.max() - in_data.min()) * 0.5
-    return (in_data - m_data) / d_data
-
+def read_and_prepare_data():
+    # normalization helper
+    def normalize_data(in_data):
+        m_data = (in_data.max() + in_data.min()) * 0.5
+        d_data = (in_data.max() - in_data.min()) * 0.5
+        return (in_data - m_data) / d_data
+
+    # define constants and read data
+    # wall_st_index = 0
+    # wall_st_end = 676
+    x_col_index = 0  # z follow next
+    e_col_index = 2  # e11, e22 and e33
+    nrows = 26  # Rows in sim strain map of the wall
+    ncols = 26  # Cols in sim strain map of the wall
+    data = pd.read_csv(low_f_data_name)
+
+    Y, Z, E11, E22, E33, R11, R22, R33 = np.array(
+        [
+            data.values[:, x_col_index],
+            data.values[:, x_col_index + 1],
+            data.values[:, e_col_index],
+            data.values[:, e_col_index + 1],
+            data.values[:, e_col_index + 2],
+            data.values[:, e_col_index + 3],
+            data.values[:, e_col_index + 4],
+            data.values[:, e_col_index + 5],
+        ]
+    )
 
-wall_st_index = 0
-wall_st_end = 676
-x_col_index = 0  # z follow next
-e_col_index = 2  # e11, e22 and e33
-nrows = 26  # Rows in sim strain map of the wall
-ncols = 26  # Cols in sim strain map of the wall
-data = pd.read_csv(low_f_data_name)
+    # Convert data
+    x1 = Y.astype(np.float32).reshape(nrows, ncols)
+    x2 = Z.astype(np.float32).reshape(nrows, ncols)
 
-Y, Z, E11, E22, E33, R11, R22, R33 = np.array(
-    [
-        data.values[:, x_col_index],
-        data.values[:, x_col_index + 1],
-        data.values[:, e_col_index],
-        data.values[:, e_col_index + 1],
-        data.values[:, e_col_index + 2],
-        data.values[:, e_col_index + 3],
-        data.values[:, e_col_index + 4],
-        data.values[:, e_col_index + 5],
-    ]
-)
+    # Sim_e11 = np.array(E11).astype(np.float32).reshape(nrows, ncols)
+    # Real_e11 = np.array(R11).astype(np.float32).reshape(nrows, ncols)
 
-x1 = Y.astype(np.float32).reshape(nrows, ncols)
-x2 = Z.astype(np.float32).reshape(nrows, ncols)
+    # Sim_e22 = np.array(E22).astype(np.float32).reshape(nrows, ncols)
+    # Real_e22 = np.array(R22).astype(np.float32).reshape(nrows, ncols)
 
-x1_norm = normalize_data(x1)
-x2_norm = normalize_data(x2)
+    # Sim_e33 = np.array(E33).astype(np.float32).reshape(nrows, ncols)
+    Real_e33 = np.array(R33).astype(np.float32).reshape(nrows, ncols)
 
-Sim_e11 = np.array(E11).astype(np.float32).reshape(nrows, ncols)
-Real_e11 = np.array(R11).astype(np.float32).reshape(nrows, ncols)
+    # Normalize data
+    x1_norm = normalize_data(x1)
+    x2_norm = normalize_data(x2)
 
-Sim_e22 = np.array(E22).astype(np.float32).reshape(nrows, ncols)
-Real_e22 = np.array(R22).astype(np.float32).reshape(nrows, ncols)
+    # Sim_e11_norm = normalize_data(Sim_e11)
+    # Real_e11_norm = normalize_data(Real_e11)
 
-Sim_e33 = np.array(E33).astype(np.float32).reshape(nrows, ncols)
-Real_e33 = np.array(R33).astype(np.float32).reshape(nrows, ncols)
+    # Sim_e22_norm = normalize_data(Sim_e22)
+    # Real_e22_norm = normalize_data(Real_e22)
 
-# Normalize data
-Sim_e11_norm = normalize_data(Sim_e11)
-Real_e11_norm = normalize_data(Real_e11)
+    # Sim_e33_norm = normalize_data(Sim_e33)
+    Real_e33_norm = normalize_data(Real_e33)
 
-Sim_e22_norm = normalize_data(Sim_e22)
-Real_e22_norm = normalize_data(Real_e22)
+    # read high fidelity data
+    # data = pd.read_csv(high_f_data_name)
 
-Sim_e33_norm = normalize_data(Sim_e33)
-Real_e33_norm = normalize_data(Real_e33)
+    # Y, Z, E11, E22, E33, R11, R22, R33 = np.array(
+    #    [
+    #        data.values[:, x_col_index],
+    #        data.values[:, x_col_index + 1],
+    #        data.values[:, e_col_index],
+    #        data.values[:, e_col_index + 1],
+    #        data.values[:, e_col_index + 2],
+    #        data.values[:, e_col_index + 3],
+    #        data.values[:, e_col_index + 4],
+    #        data.values[:, e_col_index + 5],
+    #    ]
+    # )
 
+    # Sim_hi_e11 = np.array(E11).astype(np.float32).reshape(nrows, ncols)
+    # Sim_hi_e22 = np.array(E22).astype(np.float32).reshape(nrows, ncols)
+    # Sim_hi_e33 = np.array(E33).astype(np.float32).reshape(nrows, ncols)
 
-data = pd.read_csv(high_f_data_name)
+    # Transpose the high-fidelity simulation data
+    # TODO, figure out why this is transposed
+    # Sim_hi_e11 = Sim_hi_e11.T
+    # Sim_hi_e22 = Sim_hi_e22.T
+    # Sim_hi_e33 = Sim_hi_e33.T
 
-Y, Z, E11, E22, E33, R11, R22, R33 = np.array(
-    [
-        data.values[:, x_col_index],
-        data.values[:, x_col_index + 1],
-        data.values[:, e_col_index],
-        data.values[:, e_col_index + 1],
-        data.values[:, e_col_index + 2],
-        data.values[:, e_col_index + 3],
-        data.values[:, e_col_index + 4],
-        data.values[:, e_col_index + 5],
-    ]
-)
+    # Normalize data
+    # Sim_hi_e11_norm = normalize_data(Sim_hi_e11)
+    # Sim_hi_e22_norm = normalize_data(Sim_hi_e22)
+    # Sim_hi_e33_norm = normalize_data(Sim_hi_e33)
 
-Sim_hi_e11 = np.array(E11).astype(np.float32).reshape(nrows, ncols)
-Sim_hi_e22 = np.array(E22).astype(np.float32).reshape(nrows, ncols)
-Sim_hi_e33 = np.array(E33).astype(np.float32).reshape(nrows, ncols)
+    # return data that will be used for the benchmark
+    return nrows, (x1_norm, x2_norm), Real_e33_norm
 
-# Transpose the high-fidelity simulation data
-# TODO, figure out why this is transposed
-Sim_hi_e11 = Sim_hi_e11.T
-Sim_hi_e22 = Sim_hi_e22.T
-Sim_hi_e33 = Sim_hi_e33.T
-
-# Normalize data
-Sim_hi_e11_norm = normalize_data(Sim_hi_e11)
-Sim_hi_e22_norm = normalize_data(Sim_hi_e22)
-Sim_hi_e33_norm = normalize_data(Sim_hi_e33)
 
 ###############
 
 ## Select data for ground truth
 
+ngrid, x_grids, y_grid = read_and_prepare_data()
+
 # default inputs
 INITIAL_BOUNDS = [[-1.0, 1.0], [-1.0, 1.0]]
 NUM_DIMS = len(INITIAL_BOUNDS)
 
-MESHGRID_SIZE = nrows
-INITIAL_MESHGRIDS = (x1_norm, x2_norm)
+MESHGRID_SIZE = ngrid
+INITIAL_MESHGRIDS = x_grids
 INITIAL_POINTS_TO_PREDICT = np.hstack([mg.reshape(-1, 1) for mg in INITIAL_MESHGRIDS])
 
-INITIAL_PREDICTIONS = Real_e33_norm.reshape(-1, 1)
+INITIAL_PREDICTIONS = y_grid.reshape(-1, 1)
 
 ###############
 
@@ -176,21 +186,23 @@ def strain_map(self, x) -> float:
 ###############
 
 # test parameters
-INITIAL_NUM_POINTS = 9
+INITIAL_NUM_POINTS = 25
 MAX_ITERATIONS = 150  # only allow a maximum of this many iterations in tests
 
 INITIAL_DATASET_X = np.random.uniform(-1.0, 1.0, size=(INITIAL_NUM_POINTS, NUM_DIMS)).tolist()
 
-TARGET_RMSE = 0.2
+TARGET_ERROR = 0.11
 
 
 def run_simulation(
     dataset_x: list[list[float]], dataset_y: list[float], strategy: str, strategy_args: object
 ) -> tuple[np.ndarray, np.ndarray]:
     # important "Hyper-parameters"
-    length_scale = 0.2
-    noise_level = NOISE_LEVEL
-    constant_value = 1.0
+    length_scale = 0.4
+    prior_std = 1.0
+
+    # Assume that there is some small noise in the measurements to stabilize the fit
+    statistics_y = Normal(loc='y', scale=NOISE_LEVEL)
 
     # modify a local copy of strategy_args
     strategy_args = strategy_args.copy()
@@ -198,51 +210,37 @@ def run_simulation(
 
     if backend == 'sable':
         # train model with new data
+        kernel = 'rbf'
         kernel_args = {
-            'x_range': [-1.0, 1.0],
-            'sigma_range': [1e-2, 1.0],
-            'gamma': 0.1,
+            'sigma_range': [5e-2, 1.0],
+            'gamma': 0.5,
         }
         backend_args = {
-            'n_features': 5000,
-            'alpha': constant_value * 0.05,
-            'p': 1.25,
-            'n_iter_irls': 100,
-            'noise_level': noise_level,
+            'prior_std': prior_std,
         }
 
-        client_state = DialWorkflowCreationParamsService(
-            dataset_x=dataset_x,
-            dataset_y=dataset_y,
-            bounds=INITIAL_BOUNDS,
-            kernel='rbf',
-            kernel_args=kernel_args,
-            y_is_good=False,
-            backend=backend,
-            backend_args=backend_args,
-            seed=-1,
-            dim_x=2,
-        )
     else:
         # train model with new data
-        client_state = DialWorkflowCreationParamsService(
-            dataset_x=dataset_x,
-            dataset_y=dataset_y,
-            bounds=INITIAL_BOUNDS,
-            kernel='rbf',
-            kernel_args={
-                'length_scale': length_scale,
-                'length_scale_bounds': 'fixed',
-                'noise_level': noise_level,
-                'noise_level_bounds': 'fixed',
-                'constant_value': constant_value,
-                'constant_value_bounds': 'fixed',
-            },
-            y_is_good=False,  # we wish to minimize y (the error)
-            backend=backend,  # "sklearn" or "gpax"
-            seed=-1,  # Use seed = -1 for random results
-            dim_x=2,
-        )
+        kernel = 'matern'
+        kernel_args = {
+            'length_scale': length_scale,
+            'constant_value': prior_std**2,
+        }
+        backend_args = {}
+
+    client_state = DialWorkflowCreationParamsService(
+        dataset_x=dataset_x,
+        dataset_y=dataset_y,
+        statistics_y=statistics_y,
+        bounds=INITIAL_BOUNDS,
+        kernel=kernel,
+        kernel_args=kernel_args,
+        y_is_good=False,
+        backend=backend,
+        backend_args=backend_args,
+        seed=-1,
+        dim_x=2,
+    )
 
     data = ServersideInputBase(client_state)
     model = dial_core.train_model(data)
@@ -268,10 +266,12 @@ def run_simulation(
         DialInputSingleOtherStrategy(
             workflow_id=MOCK_WORKFLOW_ID,
             bounds=INITIAL_BOUNDS,
-            y_is_good=False,  # we wish to minimize y (the error)
-            seed=-1,  # Use seed = -1 for random results
+            y_is_good=False,
+            seed=-1,
             strategy=strategy,
             strategy_args=strategy_args,
+            discrete_measurements=True,
+            discrete_measurement_grid_size=[MESHGRID_SIZE, MESHGRID_SIZE],
         ),
     )
     next_point = dial_core.get_next_point(data, model)
@@ -309,7 +309,7 @@ def graph(err_grid, dataset_x, strategy):
     )
     cbar = plt.colorbar()
     cbar.set_ticks(np.logspace(-2, 0, 7))
-    cbar.set_label('RMSE')
+    cbar.set_label('Overall error')
     plt.xlabel('Simulation Parameter #1')
     plt.ylabel('Simulation Parameter #2')
     # add black dots for data points and a red marker for the recommendation:
@@ -330,7 +330,6 @@ def accuracy_benchmark(
     """
 
     iterations = 0
-    total_rmse = float('inf')
 
     initial_dataset_x = np.random.uniform(-1.0, 1.0, size=(INITIAL_NUM_POINTS, NUM_DIMS)).tolist()
 
@@ -348,18 +347,25 @@ def accuracy_benchmark(
         # guess has no meaning here, take the last acquired datapoint
         guess = dataset_x[-1]
 
-        # compute the RMSE on the grid and the average
-        mse = err_grid**2 + std_grid**2
-        total_rmse = np.sqrt(np.mean(mse))
-        print(iterations, total_rmse)
+        # compute the RMSE and MAD on the grid and the average
+        # rmse = np.sqrt(err_grid**2 + std_grid**2)
+        mad = np.abs(err_grid)
+        # total_rmse = np.sqrt(np.mean(rmse**2))
+        total_mad = np.sqrt(np.mean(mad**2))
 
-        graph(np.sqrt(mse), dataset_x, f'{strategy}_{strategy_args.get("backend", "sklearn")}')
+        # pick the MAD over the RMSE as an error criterion, it only measures the mean deviation:
+        # it is less sensitive to wrong std and hyperparameter calibration, but does not measure how well
+        # the "error bars" (std_grid) quantify the actual ty
+        error = total_mad
+        print(iterations, error)
 
-        if total_rmse <= TARGET_RMSE:
+        graph(mad, dataset_x, f'{strategy}_{strategy_args.get("backend", "sklearn")}')
+
+        if error <= TARGET_ERROR:
             break
         iterations += 1
 
-    return iterations, total_rmse, guess
+    return iterations, error, guess
 
 
 TEST_PARAMS = (
@@ -398,7 +404,7 @@ def test_benchmark_strainmap_accuracy(
     )
     print(
         'Maximum early terminus value',
-        TARGET_RMSE,
+        TARGET_ERROR,
         ' with ',
         max_iterations,
         ' maximum iterations.',
@@ -458,7 +464,7 @@ def positive_int_type(arg):
     logger.info('Running benchmarks with %d iterations per strategy...', NUM_RUNS)
 
     results = {}
-    for strategy, strategy_args in strategies:
+    for strategy, strategy_args, _ in strategies:
         strategy_name = f'{strategy}' + (f' {json.dumps(strategy_args)}' if strategy_args else '')
         logger.info('\nBenchmarking: %s', strategy_name)
 
@@ -485,7 +491,7 @@ def positive_int_type(arg):
             'std_iterations': np.std(iterations_list),
             'avg_target': np.mean(targets_list),
             'std_target': np.std(targets_list),
-            'success_rate': sum(1 for t in targets_list if t <= TARGET_RMSE) / NUM_RUNS * 100,
+            'success_rate': sum(1 for t in targets_list if t <= TARGET_ERROR) / NUM_RUNS * 100,
         }
 
     # Generate plots
@@ -543,7 +549,9 @@ def positive_int_type(arg):
     ax2.set_xticklabels(
         [s.replace(' ', '\n') for s in strategy_names], rotation=0, ha='center', fontsize=8
     )
-    ax2.axhline(y=TARGET_RMSE, color='r', linestyle='--', label=f'Target Threshold ({TARGET_RMSE})')
+    ax2.axhline(
+        y=TARGET_ERROR, color='r', linestyle='--', label=f'Target Threshold ({TARGET_ERROR})'
+    )
     ax2.legend()
     ax2.grid(axis='y', alpha=0.3)
 
@@ -565,7 +573,7 @@ def positive_int_type(arg):
     bars3 = ax3.bar(range(len(strategy_names)), success_rates, alpha=0.7, color='green')
     ax3.set_xlabel('Strategy')
     ax3.set_ylabel('Success Rate (%)')
-    ax3.set_title(f'Success Rate (Target ≤ {TARGET_RMSE})')
+    ax3.set_title(f'Success Rate (Target ≤ {TARGET_ERROR})')
     ax3.set_xticks(range(len(strategy_names)))
     ax3.set_xticklabels(
         [s.replace(' ', '\n') for s in strategy_names], rotation=0, ha='center', fontsize=8
@@ -588,9 +596,7 @@ def positive_int_type(arg):
     # Plot 4: Box plot of iterations distribution
     ax4 = axes[1, 1]
     iterations_data = [results[s]['iterations'] for s in strategy_names]
-    bp = ax4.boxplot(
-        iterations_data, labels=[s.replace(' ', '\n') for s in strategy_names], patch_artist=True
-    )
+    bp = ax4.boxplot(iterations_data, patch_artist=True)
     for patch in bp['boxes']:
         patch.set_facecolor('lightblue')
     ax4.set_xlabel('Strategy')
@@ -695,7 +701,7 @@ def positive_int_type(arg):
         <p><strong>Generated:</strong> {datetime.datetime.now(tz=datetime.timezone.utc).astimezone().strftime('%Y-%m-%d %H:%M:%S')}</p>
         <p><strong>Number of Runs per Strategy:</strong> {NUM_RUNS}</p>
         <p><strong>Maximum Iterations:</strong> {MAX_ITERATIONS}</p>
-        <p><strong>Target Threshold:</strong> {TARGET_RMSE}</p>
+        <p><strong>Target Threshold:</strong> {TARGET_ERROR}</p>
         <p><strong>Initial Bounds:</strong> {INITIAL_BOUNDS}</p>
         <p><strong>Initial Points:</strong> {INITIAL_NUM_POINTS}</p>
     </div>
@@ -703,7 +709,7 @@ def positive_int_type(arg):
     <div class="summary">
         <h3>🎯 Test Objective</h3>
         <p>This benchmark evaluates different acquisition strategies for Bayesian optimization on a strain mapping experiment dataset.
-        The goal is to minimize the RMSE (error) within {MAX_ITERATIONS} iterations, starting from {INITIAL_NUM_POINTS} initial points.</p>
+        The goal is to minimize the total surrogate error (over all inputs) within {MAX_ITERATIONS} iterations, starting from {INITIAL_NUM_POINTS} initial points.</p>
     </div>
 
     <h2>📈 Benchmark Results</h2>
diff --git a/uv.lock b/uv.lock
index 0e595b0..9a8f535 100644
--- a/uv.lock
+++ b/uv.lock
@@ -645,6 +645,85 @@ toml = [
     { name = "tomli", marker = "python_full_version <= '3.11'" },
 ]
 
+[[package]]
+name = "cuda-bindings"
+version = "13.3.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "cuda-pathfinder", marker = "python_full_version < '3.15'" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a9/21/8464d133752951c154feafb3b65c297e7d80f301183d220bec4c830f1441/cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86", size = 6073403, upload-time = "2026-05-29T23:11:36.22Z" },
+    { url = "https://files.pythonhosted.org/packages/a8/1f/5ef51f5fbaa5d4d3201bb3d7555af028ec1aa4416275ccbf73c9e34e3d2d/cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9851b0caa8bfd3bc6fa054eaf57bea7c8e9c3a62db2d2621224677f49f3c53d0", size = 6675244, upload-time = "2026-05-29T23:11:38.664Z" },
+    { url = "https://files.pythonhosted.org/packages/51/6b/457ca12dad3ee9bfcc9a545cfd6b64b359ba49de40f776f6e028e678f262/cuda_bindings-13.3.1-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c5879712accf6e14bb01aa5e67440eb84998b8d104b509cc7a6dc0b8f656a474", size = 6053539, upload-time = "2026-05-29T23:11:43.19Z" },
+    { url = "https://files.pythonhosted.org/packages/95/7a/c5e3c34a409b148f5c0f5a4ea374158f95d488862c1dffedf9aa5c639df9/cuda_bindings-13.3.1-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:04436a9364059c84b8f9636f359eccda1cf814341f5b670c71d80d2f79dbc708", size = 6674166, upload-time = "2026-05-29T23:11:45.478Z" },
+    { url = "https://files.pythonhosted.org/packages/ce/67/5e7dba1ba576dd73da5dee894ca076ca5e959450dfff66d6d510a255d1f7/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49", size = 6025351, upload-time = "2026-05-29T23:11:49.685Z" },
+    { url = "https://files.pythonhosted.org/packages/39/2a/6d2e9047d1fb243dbaa364b01e0297534b9ed7fd27dba1c9f361519cf69b/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e32d08f71ebcdf00f0f41eab2eb37e8da94c8ed411cc9f7f7a019ce6b34abe3a", size = 6657965, upload-time = "2026-05-29T23:11:52.227Z" },
+    { url = "https://files.pythonhosted.org/packages/cc/6e/2394f8163360f8391f8f1b7e72d300a82724edb81a7b7084c799fbd4c91f/cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9efb21c1ee64981e184b9e0ba5eb3179e5ba3d4b51665a6cb52b8ef3d01a7cbf", size = 5920504, upload-time = "2026-05-29T23:11:56.883Z" },
+    { url = "https://files.pythonhosted.org/packages/34/c2/ef9b6a63f7dc432712a462c816662e662e00d38caa9b861c8c2588195d03/cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2732904099e0a4d4db774a5fc6d91ee95fae065b4d2ecabb4968c5fe2406c9d7", size = 6476660, upload-time = "2026-05-29T23:11:59.188Z" },
+    { url = "https://files.pythonhosted.org/packages/b1/81/bff68ce829999c1e4209c761bbf903b1c06ec570416ddb25020864ad5907/cuda_bindings-13.3.1-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ab2f74ed65bfef4163ba07a8db16f1085e0729291db12a2423aff84ee8278b8", size = 6013639, upload-time = "2026-05-29T23:12:03.509Z" },
+    { url = "https://files.pythonhosted.org/packages/d4/e0/c8a1f0c8f9ffdea4f5fe6dbab89b326cef4d85caf489dad39e209da89416/cuda_bindings-13.3.1-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:efd4c814d311ec08c981f6dded1dbe7d4b371067ee4f6c14cccec4bde9590f80", size = 6534419, upload-time = "2026-05-29T23:12:05.633Z" },
+    { url = "https://files.pythonhosted.org/packages/52/b8/83b1f563925b290f2d11a01a77a84013ba56052fe3653a5bef3ccfbb43d6/cuda_bindings-13.3.1-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c3c772dfff49681541d59630c90f858e173ac926b9c593a2b7123f2a1043cc76", size = 5809771, upload-time = "2026-05-29T23:12:10.422Z" },
+    { url = "https://files.pythonhosted.org/packages/12/20/e79b4bfe98f075195afb6343d41c498f9dbd2d161d7021d4d28bceb83581/cuda_bindings-13.3.1-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:36febb7c1079d68a981dbbd8d5a67235b399802b82075c9388624719607e52b9", size = 6358584, upload-time = "2026-05-29T23:12:12.767Z" },
+]
+
+[[package]]
+name = "cuda-pathfinder"
+version = "1.5.6"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d2/53/8fc9b0cdc5b7f62746e6a01b85b6461e5ae27f871010a5fcf8fa6950766d/cuda_pathfinder-1.5.6-py3-none-any.whl", hash = "sha256:7e4c07c117b78ba1fb35dac4c444d21f3677b1b1ff56175c53a8e3025c5b43c0", size = 52972, upload-time = "2026-06-30T00:58:04.34Z" },
+]
+
+[[package]]
+name = "cuda-toolkit"
+version = "13.0.3.0"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d1/c7/a79086a62c98befcdb8349656c6f114e2db3b8b2422f6e25c97a7f2a9a3c/cuda_toolkit-13.0.3.0-py2.py3-none-any.whl", hash = "sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f", size = 2512, upload-time = "2026-04-14T00:50:08.173Z" },
+]
+
+[package.optional-dependencies]
+cublas = [
+    { name = "nvidia-cublas", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+cudart = [
+    { name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+cufft = [
+    { name = "nvidia-cufft", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+cufile = [
+    { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+cupti = [
+    { name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+curand = [
+    { name = "nvidia-curand", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+cusolver = [
+    { name = "nvidia-cublas", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cusolver", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cusparse", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+cusparse = [
+    { name = "nvidia-cusparse", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+nvjitlink = [
+    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+nvrtc = [
+    { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+nvtx = [
+    { name = "nvidia-nvtx", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+]
+
 [[package]]
 name = "cycler"
 version = "0.12.1"
@@ -1909,6 +1988,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ad/3f/3d42e9a78fe5edf792a83c074b13b9b770092a4fbf3462872f4303135f09/ml_dtypes-0.5.4-cp314-cp314t-win_arm64.whl", hash = "sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d", size = 168825, upload-time = "2025-11-17T22:32:23.766Z" },
 ]
 
+[[package]]
+name = "mpmath"
+version = "1.3.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" },
+]
+
 [[package]]
 name = "msgpack"
 version = "1.1.2"
@@ -2055,6 +2143,36 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a0/c4/c2971a3ba4c6103a3d10c4b0f24f461ddc027f0f09763220cf35ca1401b3/nest_asyncio-1.6.0-py3-none-any.whl", hash = "sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c", size = 5195, upload-time = "2024-01-21T14:25:17.223Z" },
 ]
 
+[[package]]
+name = "networkx"
+version = "3.4.2"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version < '3.11'",
+]
+sdist = { url = "https://files.pythonhosted.org/packages/fd/1d/06475e1cd5264c0b870ea2cc6fdb3e37177c1e565c43f56ff17a10e3937f/networkx-3.4.2.tar.gz", hash = "sha256:307c3669428c5362aab27c8a1260aa8f47c4e91d3891f48be0141738d8d053e1", size = 2151368, upload-time = "2024-10-21T12:39:38.695Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl", hash = "sha256:df5d4365b724cf81b8c6a7312509d0c22386097011ad1abe274afd5e9d3bbc5f", size = 1723263, upload-time = "2024-10-21T12:39:36.247Z" },
+]
+
+[[package]]
+name = "networkx"
+version = "3.6.1"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "python_full_version >= '3.15' and sys_platform == 'win32'",
+    "python_full_version >= '3.15' and sys_platform == 'emscripten'",
+    "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'",
+    "python_full_version == '3.14.*'",
+    "python_full_version == '3.13.*'",
+    "python_full_version == '3.12.*'",
+    "python_full_version == '3.11.*'",
+]
+sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" },
+]
+
 [[package]]
 name = "nodeenv"
 version = "1.10.0"
@@ -2238,6 +2356,158 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/88/31/9b5da5995988437756bc3f1eead2e314d8916259875c6924cb41692f2b41/numpyro-0.19.0-py3-none-any.whl", hash = "sha256:1063a2c131a0785719e13c8e55f1b82e41850d814df149418097531f4dbdeda8", size = 370906, upload-time = "2025-08-05T10:26:31.35Z" },
 ]
 
+[[package]]
+name = "nvidia-cublas"
+version = "13.1.1.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" },
+    { url = "https://files.pythonhosted.org/packages/3b/cd/154ca20c38269e05eff77c1464e6c1da89f50a6390b565e9d82e06bc11e1/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:37936a16db8fe4ac1f065c2139360608a543a09275cb1a1af612e08cfa065436", size = 423138758, upload-time = "2026-04-08T18:46:58.655Z" },
+]
+
+[[package]]
+name = "nvidia-cuda-cupti"
+version = "13.0.85"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" },
+    { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" },
+]
+
+[[package]]
+name = "nvidia-cuda-nvrtc"
+version = "13.0.88"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" },
+    { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" },
+]
+
+[[package]]
+name = "nvidia-cuda-runtime"
+version = "13.0.96"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" },
+    { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" },
+]
+
+[[package]]
+name = "nvidia-cudnn-cu13"
+version = "9.20.0.48"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cublas", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" },
+    { url = "https://files.pythonhosted.org/packages/6e/5e/edb9c0ae051602c3ccaffe424256463636d639e27d7f302dde9975ef9e7a/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0c45dd8eeb50b603f07995b1b300c62ffe6a1980482b82b3bcf94a4ca9d49304", size = 366173588, upload-time = "2026-03-09T19:29:34.474Z" },
+]
+
+[[package]]
+name = "nvidia-cufft"
+version = "12.0.0.61"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" },
+    { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" },
+]
+
+[[package]]
+name = "nvidia-cufile"
+version = "1.15.1.6"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" },
+    { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" },
+]
+
+[[package]]
+name = "nvidia-curand"
+version = "10.4.0.35"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" },
+    { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" },
+]
+
+[[package]]
+name = "nvidia-cusolver"
+version = "12.0.4.66"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cublas", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-cusparse", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" },
+    { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" },
+]
+
+[[package]]
+name = "nvidia-cusparse"
+version = "12.6.3.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },
+    { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" },
+]
+
+[[package]]
+name = "nvidia-cusparselt-cu13"
+version = "0.8.1"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/46/e1/cdc1797eadf82d3a9a575a19b33fdc871a97edbec42c00b5b5e914f4aff4/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f", size = 221051344, upload-time = "2025-09-05T18:49:51.289Z" },
+    { url = "https://files.pythonhosted.org/packages/34/7d/2661f2fb3ac4302f3a246f5fc030213ac60c1fe0bce84f9783dbd831dbb7/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0", size = 170148586, upload-time = "2025-09-05T18:50:50.248Z" },
+]
+
+[[package]]
+name = "nvidia-nccl-cu13"
+version = "2.29.7"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/72/0d/daf50d44177ee0cbc7ff0a0c91eb5ff676c82be42f9a970bc7597f440c3a/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:674a12383e3c38a1bcccae7d4f3633b37852230b6047883cb2f4c2d1b36d9bf5", size = 206014712, upload-time = "2026-03-03T05:34:20.843Z" },
+    { url = "https://files.pythonhosted.org/packages/67/f4/58e4e91b6919367c7aafb8e36fce9aad1a3047e536bf7e2fd560927d3a4c/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d", size = 205976000, upload-time = "2026-03-03T05:36:24.472Z" },
+]
+
+[[package]]
+name = "nvidia-nvjitlink"
+version = "13.3.33"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f0/ee/580ca6f29dcab0221db8706badca1bbbb084f1975c4d4e83329c3a7e31f0/nvidia_nvjitlink-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:26a6de7fb4c8fdaa7703d3dad720d6d427ddfea5c48a528fd97c11733ad830e5", size = 40742423, upload-time = "2026-05-26T16:54:51.613Z" },
+    { url = "https://files.pythonhosted.org/packages/69/30/45414e35ff2eee7db3da037e5707037ccf9d2b5218ffbdb055ea4d5aa98a/nvidia_nvjitlink-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ce48b37dfeb3cb1eae4cf85adacb47d7a6539ea2272870c9a3628ce275c2037e", size = 39168635, upload-time = "2026-05-26T16:54:13.906Z" },
+]
+
+[[package]]
+name = "nvidia-nvshmem-cu13"
+version = "3.4.5"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" },
+    { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" },
+]
+
+[[package]]
+name = "nvidia-nvtx"
+version = "13.0.85"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" },
+    { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" },
+]
+
 [[package]]
 name = "opt-einsum"
 version = "3.4.0"
@@ -3208,11 +3478,12 @@ wheels = [
 
 [[package]]
 name = "sable"
-version = "0.1.0"
-source = { git = "https://code.ornl.gov/sable/sable.git#30350de7d14b730ab3de71d50c92035dc76cb9c6" }
+version = "0.1.1"
+source = { git = "https://code.ornl.gov/sable/sable.git#6bdbbf62f4ac26369ce202134273c10617fa6be3" }
 dependencies = [
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
+    { name = "torch" },
 ]
 
 [[package]]
@@ -3460,6 +3731,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/07/39/338d9219c4e87f3e708f18857ecd24d22a0c3094752393319553096b98af/scipy-1.17.1-cp314-cp314t-win_arm64.whl", hash = "sha256:200e1050faffacc162be6a486a984a0497866ec54149a01270adc8a59b7c7d21", size = 25489165, upload-time = "2026-02-23T00:22:29.563Z" },
 ]
 
+[[package]]
+name = "setuptools"
+version = "83.0.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" },
+]
+
 [[package]]
 name = "simplejson"
 version = "4.0.1"
@@ -3775,6 +4055,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/26/33/f1652d0c59fa51de18492ee2345b65372550501ad061daa38f950be390b6/statsmodels-0.14.6-cp314-cp314-win_amd64.whl", hash = "sha256:151b73e29f01fe619dbce7f66d61a356e9d1fe5e906529b78807df9189c37721", size = 9588010, upload-time = "2025-12-05T23:14:07.28Z" },
 ]
 
+[[package]]
+name = "sympy"
+version = "1.14.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "mpmath" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" },
+]
+
 [[package]]
 name = "tabulate"
 version = "0.10.0"
@@ -3922,6 +4214,54 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" },
 ]
 
+[[package]]
+name = "torch"
+version = "2.13.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "cuda-bindings", marker = "python_full_version < '3.15' and sys_platform == 'linux'" },
+    { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" },
+    { name = "filelock" },
+    { name = "fsspec" },
+    { name = "jinja2" },
+    { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
+    { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" },
+    { name = "setuptools" },
+    { name = "sympy" },
+    { name = "triton", marker = "python_full_version < '3.15' and sys_platform == 'linux'" },
+    { name = "typing-extensions" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7f/e7/19894fdb51c7dbaf94f5a79bb0871da0992e8e4241e579cb006da46d2e58/torch-2.13.0-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:94f0de129916f77b8dc2c7a8eff644cfeddfe59e39c9f55e9f6e17543410281d", size = 111178962, upload-time = "2026-07-08T16:05:49.855Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/5c/b1d5de470c54e339b30a92d96683a71bcebd78f5f2a7fc714cd6dc6bbd68/torch-2.13.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0ab4b69f3ee03a62a002cfbf77b1ca5e88aceb4ea64cb4388bb28f638ddbb045", size = 427198333, upload-time = "2026-07-08T16:05:36.847Z" },
+    { url = "https://files.pythonhosted.org/packages/50/c0/68a84105e1fcb8970144b388ff3d3e5dc15a3be28c1e247841f7d7247e41/torch-2.13.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:c78b7b4d04461855a764cf01bae9a462bb88bc93defcfa11235cbc8fdf3e12c4", size = 526555154, upload-time = "2026-07-08T16:05:06.507Z" },
+    { url = "https://files.pythonhosted.org/packages/2e/c9/0bb9d097b03cbaf96bb75b15e867347b8e41bfcdfe0539452d17d9e63993/torch-2.13.0-cp310-cp310-win_amd64.whl", hash = "sha256:2bd30b6b730d987fa386ce3898933762c5cb8cc82eb0535211d787cc3ce2dfeb", size = 122015602, upload-time = "2026-07-08T16:05:45.25Z" },
+    { url = "https://files.pythonhosted.org/packages/5b/fe/cba54dc58523434919b66f13a667e36e436deddd77ca519e96553617d4ec/torch-2.13.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:e76f9bcecc52b8ff711239a2f7547d5353df95878ab232f0773c1d95928b92f8", size = 111187938, upload-time = "2026-07-08T16:05:17.065Z" },
+    { url = "https://files.pythonhosted.org/packages/c2/59/1e3160e18e12aa3038390efab3ce02b36a9d4d6a527ecdd8520dca2e68d8/torch-2.13.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:092790c696a760c729fd5722835f50b9d81fd7c8f141571f3f3cf4081a8f664c", size = 427199369, upload-time = "2026-07-08T16:04:51.054Z" },
+    { url = "https://files.pythonhosted.org/packages/01/79/1f2d34ad7034ee1c7ffc1cf8bf0f8213af2a81df6ecdb3997ecec107c09d/torch-2.13.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:60fcdcb2f3876e21146cb4524ef06397d727ca9ad5f020818547e25075fe3cb7", size = 526574961, upload-time = "2026-07-08T16:04:07.075Z" },
+    { url = "https://files.pythonhosted.org/packages/6c/fd/0f2ce40f58aefbdb3392f9acce3c8171940943ae2d661f70558bfa73befb/torch-2.13.0-cp311-cp311-win_amd64.whl", hash = "sha256:a0d8b11f16a48d60e2015d8213aa0390744cbebb98e58b62b3514dddc656e330", size = 122015870, upload-time = "2026-07-08T16:05:27.59Z" },
+    { url = "https://files.pythonhosted.org/packages/c4/3a/ed0f4d4d1dcde03bced7aac9a28e800abcdc0cbd06b6775044c9fbd877b7/torch-2.13.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2fe228aba290d14b9f31b049be550dbd469c3fd3013d7a19705b30454da97027", size = 111213045, upload-time = "2026-07-08T16:05:22.997Z" },
+    { url = "https://files.pythonhosted.org/packages/df/a9/f6a2a4d763ff1df02e9a64c477029db614295bc9367f4131223791ccc243/torch-2.13.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:572df8be8ffb4599c88cbd6a0726f1f854f4da65d2e3c09f0e2c2283333cd6d4", size = 427210998, upload-time = "2026-07-08T16:04:37.708Z" },
+    { url = "https://files.pythonhosted.org/packages/f3/82/fea946351658e6534db52d2cc12bc53087cbf87f9440c5f180f367c1950b/torch-2.13.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:796633c4cdf0fe2cdced72d8f88f22e73dbcfce83132763162f6d4bff13b820b", size = 526605292, upload-time = "2026-07-08T16:04:22.81Z" },
+    { url = "https://files.pythonhosted.org/packages/21/d6/e8f3c6f7e01f626f77259de9860d2a78bc84c40539e28e79b7e98b0bb659/torch-2.13.0-cp312-cp312-win_amd64.whl", hash = "sha256:024c6cc0c1b085f2f91f20a3dc27b0471d021c31ce84b81be3afdc39f791fd9d", size = 122057313, upload-time = "2026-07-08T16:03:53.43Z" },
+    { url = "https://files.pythonhosted.org/packages/0d/fa/c1c10b7aff4a9a3e8956d4f0a5f468fa6db7abc3208805719076772b4833/torch-2.13.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:33449899ce5496c1b84b4853179d94fd102028ae1407314d9fb956bb79e70d09", size = 111213743, upload-time = "2026-07-08T16:03:28.579Z" },
+    { url = "https://files.pythonhosted.org/packages/11/18/9ecb37b56293a0be8d80f810bf672a72fe7e02f8b475d5ef1b9bf8a0d748/torch-2.13.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:1e09d6a722504957c694faceca843acde562786df1144ebcc5a74075ec7f6005", size = 427213008, upload-time = "2026-07-08T16:03:44.106Z" },
+    { url = "https://files.pythonhosted.org/packages/d4/5a/7c50ba1b7b713d71d34669c6d13dab0a11531a3eceb0307a5162dbfec0f7/torch-2.13.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:a3a9a21312872af8a26950b2c15680335a386a1f56ed03e780653d78b9607e9e", size = 526602329, upload-time = "2026-07-08T16:03:12.649Z" },
+    { url = "https://files.pythonhosted.org/packages/91/3d/e7adcc6aaf36961cd18f56cf8ad0f3058c3a5c84ccf391762176c94581b8/torch-2.13.0-cp313-cp313-win_amd64.whl", hash = "sha256:49b58f1e2c52440abb6f17c28f0335fe6c6d01ad1a7f55b0183b81e4b34d64e6", size = 122057920, upload-time = "2026-07-08T16:03:01.808Z" },
+    { url = "https://files.pythonhosted.org/packages/36/76/6dcc7f0c07052102dd36f83cbc5800842a909c8c3fbf1a7f8a5844954de9/torch-2.13.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d849b390e07d8d333ce8ecaf91b273c656c598379a19c9acf1318a883f6b391c", size = 111227066, upload-time = "2026-07-08T16:03:33.6Z" },
+    { url = "https://files.pythonhosted.org/packages/e9/09/2c10e8cd0e00fa5d23c052df6ce467eaa7182399f5e0f824f1e4ff42ccae/torch-2.13.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:a3893dc2da0a972a8ca5d698c85a9f967559ac5f8ee1797b77408aa8734d073c", size = 427226309, upload-time = "2026-07-08T16:02:53.127Z" },
+    { url = "https://files.pythonhosted.org/packages/76/c6/22c2102bbef14ca6a6cb4c20e42f088e49c5f812be4e160ae57502e325f9/torch-2.13.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:49f1ea385c754e54919408a9bb3b5a72b0b755bbe2c916c1d6f70afbec4908a2", size = 526614507, upload-time = "2026-07-08T16:02:16.441Z" },
+    { url = "https://files.pythonhosted.org/packages/2b/0c/7d1deb6bce5bc3e6042caf39100ac768eba3b9a098e1dddd16f75bd6489b/torch-2.13.0-cp314-cp314-win_amd64.whl", hash = "sha256:4f8573e3ce9ebcd53fe922f01077a6085ccdfbe5f12fd215883a9d87d7a744fd", size = 122051871, upload-time = "2026-07-08T16:03:23.521Z" },
+    { url = "https://files.pythonhosted.org/packages/f4/ce/aa8b7f9949d32e0f2f624f342bc3b48112c1b8a130288465938bc83bcbf9/torch-2.13.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:c28def70706c2f9ecc752574766e8ae4da9b810ab6676b611166761a78a9f1e1", size = 111537025, upload-time = "2026-07-08T16:02:44.28Z" },
+    { url = "https://files.pythonhosted.org/packages/69/d1/491e3a0389430946145888b0203f2b6a759ce2a61481b96a85c2da4f2ced/torch-2.13.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:31061ff56ed8fbf26c749806905aeb749ebeb819810fd5d52508aa5afd90dddc", size = 427219769, upload-time = "2026-07-08T16:02:31.18Z" },
+    { url = "https://files.pythonhosted.org/packages/9a/1d/38006e045bf0a1fc28ef01e757c554e59e59a8770c284bc4f47b14e60441/torch-2.13.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:cc26eead4cf51d0b544e31e364dcf000846549c273bd148936fe9d24d29acb92", size = 526571320, upload-time = "2026-07-08T16:01:59.348Z" },
+    { url = "https://files.pythonhosted.org/packages/56/94/655c91992a882bd5071aa0b5d22a07dbb130d801e872be97c0b627a7c693/torch-2.13.0-cp314-cp314t-win_amd64.whl", hash = "sha256:a7de8a313090dc5c7d7ba4bfe5c3be222528f9a4dba1acc83bddb1157360c4b8", size = 122306773, upload-time = "2026-07-08T16:02:39.832Z" },
+]
+
 [[package]]
 name = "tqdm"
 version = "4.67.3"
@@ -3947,6 +4287,25 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/43/2b/36e984399089c026a6499ac8f7401d38487cf0183839a4aa78140d373771/treescope-0.1.10-py3-none-any.whl", hash = "sha256:dde52f5314f4c29d22157a6fe4d3bd103f9cae02791c9e672eefa32c9aa1da51", size = 182255, upload-time = "2025-08-08T05:43:46.673Z" },
 ]
 
+[[package]]
+name = "triton"
+version = "3.7.1"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ec/ea/629cc37436ca5df93ce98956d09cd2ca1498bfee8ef4972d2fe48b9f958c/triton-3.7.1-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3daf64305d6cea88d3334c65ebc9bcd0c64c9564a977084366aa768d57cbcf64", size = 184551013, upload-time = "2026-06-17T20:03:37.551Z" },
+    { url = "https://files.pythonhosted.org/packages/15/76/c79c34311625227a288df3e483fc5cdf3d596624cbd4b4758c4cbdc14af3/triton-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ee89fbf782ec2ad50391dd1cf26cbea4f4467154c37f4773026da8fc31c0f58e", size = 197596267, upload-time = "2026-06-17T19:53:06.898Z" },
+    { url = "https://files.pythonhosted.org/packages/7b/f9/19d842d06a08559534fa1eaab6ca551b1bcf40f06620bddec1babaa2772d/triton-3.7.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d4a0e1cd4c4a76370ed74a8432a53cea28716827d19e40ffc732233e35ceb3f6", size = 184664887, upload-time = "2026-06-17T20:03:42.913Z" },
+    { url = "https://files.pythonhosted.org/packages/cd/5e/fce69606f7f240297f163e25539906732b199530d486ce67ae319877e821/triton-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6744957e9fd610a29680ec2346057d0c86948ed3812468670719f391e94b44a5", size = 197701306, upload-time = "2026-06-17T19:53:13.673Z" },
+    { url = "https://files.pythonhosted.org/packages/94/fa/f856e24deb462d5f18bd4b5a746957862ab9b6ee5834bda60605ec348366/triton-3.7.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9497f2e696ee368862a181a90b2dcc03ca978cc4f602abd67c7d81022a6988e1", size = 184692359, upload-time = "2026-06-17T20:03:48.288Z" },
+    { url = "https://files.pythonhosted.org/packages/c4/6f/fb96d15db6f36d6eae4cafb998c2e0353bf59d7c4ea1662d7497f269134a/triton-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e40869937a68206ec70d7f25bb7ec6433cb083f9135e1f36dbd318dc449a728", size = 197719725, upload-time = "2026-06-17T19:53:20.419Z" },
+    { url = "https://files.pythonhosted.org/packages/00/42/c5089d4d9327fcd1e862c599cc2927f39418f84dd11a84cb2ccff9d4787a/triton-3.7.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdbfc09d9ec58bc5e68321525653220de7515c199e7a8097a97c85e62b52cd0a", size = 184694629, upload-time = "2026-06-17T20:03:53.444Z" },
+    { url = "https://files.pythonhosted.org/packages/07/42/2c3ac59253ae8892b6f307875263dd23dc875cdf732d3aea40d6d41fb7cb/triton-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:58c0e131da05134a2a4788ccbcc0c1105cf0f54c8e98f19e34cd465396dc15eb", size = 197729241, upload-time = "2026-06-17T19:53:27.801Z" },
+    { url = "https://files.pythonhosted.org/packages/40/71/e01aa7ad573883ed9456f130226babdec70b005e098c4d6226a6238e761b/triton-3.7.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe4ea396a06171f1f1f58cbd39c70b09294398f7dd7c620939bab54ad6f934fa", size = 184705764, upload-time = "2026-06-17T20:03:59.064Z" },
+    { url = "https://files.pythonhosted.org/packages/a4/09/5683146fda6a2b569deb78ccfd8fbfea8bfe55f726b081c0a6bb18dd6f28/triton-3.7.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2020153b08280415ec0da6607834e79166442147e78e144df06b508c75b186d2", size = 197729537, upload-time = "2026-06-17T19:53:35.516Z" },
+    { url = "https://files.pythonhosted.org/packages/e9/f8/448220c3092019f9fdfab39ec47985968181d67da34b44f6a7f6280a5cbb/triton-3.7.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c58e4c61f0c73b5dba3b5d19b4a7093c32f90dc18b2a7f121a7c16ccd31107b7", size = 184814760, upload-time = "2026-06-17T20:04:04.984Z" },
+    { url = "https://files.pythonhosted.org/packages/f0/ac/229b7d4589d2e5937310e72c6d46e89599d16a4a12b479ffa1499fee8eb8/triton-3.7.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10ba85fa2cca4a2fbdeb36bf1cb082f2c252bda55bf9fccd74f65ec5bc647e68", size = 197824404, upload-time = "2026-06-17T19:53:42.772Z" },
+]
+
 [[package]]
 name = "typing-extensions"
 version = "4.15.0"

From f15bf241bce9e033b72a84c60480b96c7ed9ad56 Mon Sep 17 00:00:00 2001
From: Konstantin Pieper <pieperk@ornl.gov>
Date: Tue, 21 Jul 2026 18:37:36 -0400
Subject: [PATCH 3/8] add a heteroscedastic unit test

---
 tests/unit/test_internals.py | 49 ++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/tests/unit/test_internals.py b/tests/unit/test_internals.py
index a84d894..1b6f9c1 100644
--- a/tests/unit/test_internals.py
+++ b/tests/unit/test_internals.py
@@ -9,6 +9,7 @@
     DialInputMultiple,
     DialInputPredictions,
     DialInputSingleOtherStrategy,
+    Normal,
 )
 from dial_service import core
 from dial_service.serverside_data import (
@@ -258,6 +259,7 @@ def multiple_2D(backend, strategy, discrete_measurement_grid_size=None):
         dataset_x=[],
         dataset_y=[],
         dim_x=2,  # provide dim_x for empty dataset
+        dim_y=1,  # provide dim_y for empty dataset
         y_is_good=False,
         kernel='rbf',
         bounds=[[0, 100], [-1, 1]],
@@ -299,6 +301,30 @@ def prediction_1D(backend):
     return ServersideInputPrediction(workflow_state, params)
 
 
+def prediction_1D_heteroscedastic(backend):
+    workflow_state = DialWorkflowCreationParamsService(
+        dataset_x=[[1.0], [1.5], [2.0]],
+        dataset_y=[[-1.0, 1e-2], [0.0, 1e2], [1.0, 1e-5]],
+        labels_y=['y', 'yerr'],
+        dim_y=2,
+        statistics_y=Normal(loc='y', scale='yerr'),
+        bounds=[[1.0, 2.0]],
+        kernel='rbf',
+        kernel_args={
+            'length_scale': 0.5,
+        },
+        backend=backend,
+        preprocess_standardize=False,
+        y_is_good=True,
+        seed=42,
+    )
+    params = DialInputPredictions(
+        workflow_id=DUMMY_WORKFLOW_ID,
+        points_to_predict=[[1], [1.25], [1.5], [1.75], [2]],
+    )
+    return ServersideInputPrediction(workflow_state, params)
+
+
 ####### TESTS ###################
 
 
@@ -639,6 +665,29 @@ def test_surrogate(backend, expected_means, expected_stddevs):
     assert stddevs[1:4] == pytest.approx(expected_stddevs)
 
 
+@pytest.mark.parametrize(
+    ('backend', 'expected_means', 'expected_stddevs'),
+    [
+        (
+            'sklearn',
+            [-1.0, -0.65, 0.0, 0.65, 1.0],
+            [1e-2, 0.42, 0.59, 0.42, 1e-5],
+        ),
+        (
+            'sable',
+            [-1.0, -0.56, 0.0, 0.56, 1.0],
+            [1e-2, 0.29, 0.34, 0.29, 1e-5],
+        ),
+    ],
+)
+def test_surrogate_heteroscedastic(backend, expected_means, expected_stddevs):
+    data = prediction_1D_heteroscedastic(backend)
+    model = core.train_model(data)
+    means, stddevs, _ = core.get_surrogate_values(data, model)
+    assert means == pytest.approx(expected_means, rel=0.01, abs=1e-4)
+    assert stddevs == pytest.approx(expected_stddevs, rel=0.01, abs=1e-4)
+
+
 @pytest.mark.parametrize(
     ('backend'),
     [

From f6e5b3a8eb774b299528db01a883d85b8b9f2c8b Mon Sep 17 00:00:00 2001
From: Lance-Drane <Lance-Drane@users.noreply.github.com>
Date: Thu, 23 Jul 2026 21:08:47 -0400
Subject: [PATCH 4/8] make gpax and sable optional dependencies

Signed-off-by: Lance-Drane <Lance-Drane@users.noreply.github.com>
---
 pyproject.toml |  10 +-
 uv.lock        | 386 +++++++++++++++++++++++++------------------------
 2 files changed, 200 insertions(+), 196 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index b36ae29..eb50546 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,12 +20,10 @@ classifiers = ["Programming Language :: Python :: 3"]
 dependencies = [
   "intersect_sdk>=0.9.3,<0.10.0",
   "numpy",
-  "scikit-learn>=1.4.0,<2.0.0",                        # TODO consider making an optional dependency group
+  "scikit-learn>=1.4.0,<2.0.0",   # TODO consider making an optional dependency group
   "scipy>=1.12.0,<2.0.0",
-  "gpax>=0.1.8",                                       # TODO consider making an optional dependency group
-  "numpyro<0.20.1",                                    # depend on older numpyro for gpax (numpyro.contrib.module: random_haiku_module)
-  "pymongo>=4.12.1",                                   # TODO - this is only needed for dial_service, dial_dataclass can use a simple fixture for ObjectID representation which allows it to skip this dependency
-  "sable @ git+https://code.ornl.gov/sable/sable.git", # TODO this references the internal repo, transition to public version
+  "numpyro<0.20.1",               # depend on older numpyro for gpax (numpyro.contrib.module: random_haiku_module)
+  "pymongo>=4.12.1",              # TODO - this is only needed for dial_service, dial_dataclass can use a simple fixture for ObjectID representation which allows it to skip this dependency
 ]
 
 [project.urls]
@@ -36,6 +34,8 @@ Issues = "https://github.com/INTERSECT-DIAL/dial/issues"
 
 [project.optional-dependencies]
 docs = ["sphinx>=5.3.0", "furo>=2023.3.27"]
+gpax = ["gpax>=0.1.8"]
+sable = ["sable @ git+https://code.ornl.gov/sable/sable.git"]
 
 [dependency-groups]
 dev = [
diff --git a/uv.lock b/uv.lock
index 9a8f535..66d3f1b 100644
--- a/uv.lock
+++ b/uv.lock
@@ -374,7 +374,7 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/66/54/eb9bfc647b19f2009dd5c7f5ec51c4e6ca831725f1aea7a993034f483147/contourpy-1.3.2.tar.gz", hash = "sha256:b6945942715a034c671b7fc54f9588126b0b8bf23db2696e3ca8328f3ff0ab54", size = 13466130, upload-time = "2025-04-15T17:47:53.79Z" }
 wheels = [
@@ -450,7 +450,7 @@ resolution-markers = [
     "python_full_version == '3.11.*'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" }
 wheels = [
@@ -650,7 +650,7 @@ name = "cuda-bindings"
 version = "13.3.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "cuda-pathfinder", marker = "python_full_version < '3.15'" },
+    { name = "cuda-pathfinder" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/a9/21/8464d133752951c154feafb3b65c297e7d80f301183d220bec4c830f1441/cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86", size = 6073403, upload-time = "2026-05-29T23:11:36.22Z" },
@@ -685,43 +685,43 @@ wheels = [
 
 [package.optional-dependencies]
 cublas = [
-    { name = "nvidia-cublas", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
-    { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 cudart = [
-    { name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 cufft = [
-    { name = "nvidia-cufft", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
-    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 cufile = [
-    { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 cupti = [
-    { name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 curand = [
-    { name = "nvidia-curand", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 cusolver = [
-    { name = "nvidia-cublas", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
-    { name = "nvidia-cusolver", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
-    { name = "nvidia-cusparse", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
-    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 cusparse = [
-    { name = "nvidia-cusparse", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
-    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
+    { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 nvjitlink = [
-    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 nvrtc = [
-    { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 nvtx = [
-    { name = "nvidia-nvtx", marker = "(python_full_version < '3.15' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
+    { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
 
 [[package]]
@@ -738,14 +738,11 @@ name = "dial"
 version = "0.1.6"
 source = { editable = "." }
 dependencies = [
-    { name = "gpax", version = "0.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
-    { name = "gpax", version = "0.1.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" },
     { name = "intersect-sdk" },
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
     { name = "numpyro" },
     { name = "pymongo" },
-    { name = "sable" },
     { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "scikit-learn", version = "1.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
     { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
@@ -759,6 +756,13 @@ docs = [
     { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" },
     { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
 ]
+gpax = [
+    { name = "gpax", version = "0.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
+    { name = "gpax", version = "0.1.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" },
+]
+sable = [
+    { name = "sable" },
+]
 
 [package.dev-dependencies]
 dev = [
@@ -775,17 +779,17 @@ dev = [
 [package.metadata]
 requires-dist = [
     { name = "furo", marker = "extra == 'docs'", specifier = ">=2023.3.27" },
-    { name = "gpax", specifier = ">=0.1.8" },
+    { name = "gpax", marker = "extra == 'gpax'", specifier = ">=0.1.8" },
     { name = "intersect-sdk", specifier = ">=0.9.3,<0.10.0" },
     { name = "numpy" },
     { name = "numpyro", specifier = "<0.20.1" },
     { name = "pymongo", specifier = ">=4.12.1" },
-    { name = "sable", git = "https://code.ornl.gov/sable/sable.git" },
+    { name = "sable", marker = "extra == 'sable'", git = "https://code.ornl.gov/sable/sable.git" },
     { name = "scikit-learn", specifier = ">=1.4.0,<2.0.0" },
     { name = "scipy", specifier = ">=1.12.0,<2.0.0" },
     { name = "sphinx", marker = "extra == 'docs'", specifier = ">=5.3.0" },
 ]
-provides-extras = ["docs"]
+provides-extras = ["docs", "gpax", "sable"]
 
 [package.metadata.requires-dev]
 dev = [
@@ -877,13 +881,13 @@ wheels = [
 
 [package.optional-dependencies]
 epath = [
-    { name = "fsspec", marker = "python_full_version < '3.11'" },
-    { name = "importlib-resources", marker = "python_full_version < '3.11'" },
-    { name = "typing-extensions", marker = "python_full_version < '3.11'" },
-    { name = "zipp", marker = "python_full_version < '3.11'" },
+    { name = "fsspec" },
+    { name = "importlib-resources" },
+    { name = "typing-extensions" },
+    { name = "zipp" },
 ]
 epy = [
-    { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+    { name = "typing-extensions" },
 ]
 
 [[package]]
@@ -901,12 +905,12 @@ wheels = [
 
 [package.optional-dependencies]
 epath = [
-    { name = "fsspec", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "typing-extensions", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "zipp", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
+    { name = "fsspec" },
+    { name = "typing-extensions" },
+    { name = "zipp" },
 ]
 epy = [
-    { name = "typing-extensions", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
+    { name = "typing-extensions" },
 ]
 
 [[package]]
@@ -914,7 +918,7 @@ name = "exceptiongroup"
 version = "1.3.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+    { name = "typing-extensions" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" }
 wheels = [
@@ -938,15 +942,15 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "msgpack", marker = "python_full_version < '3.11'" },
-    { name = "optax", marker = "python_full_version < '3.11'" },
-    { name = "orbax-checkpoint", marker = "python_full_version < '3.11'" },
-    { name = "pyyaml", marker = "python_full_version < '3.11'" },
-    { name = "rich", marker = "python_full_version < '3.11'" },
-    { name = "tensorstore", version = "0.1.78", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "treescope", marker = "python_full_version < '3.11'" },
-    { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" } },
+    { name = "msgpack" },
+    { name = "optax" },
+    { name = "orbax-checkpoint" },
+    { name = "pyyaml" },
+    { name = "rich" },
+    { name = "tensorstore", version = "0.1.78", source = { registry = "https://pypi.org/simple" } },
+    { name = "treescope" },
+    { name = "typing-extensions" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/e6/76/4ea55a60a47e98fcff591238ee26ed4624cb4fdc4893aa3ebf78d0d021f4/flax-0.10.7.tar.gz", hash = "sha256:2930d6671e23076f6db3b96afacf45c5060898f5c189ecab6dda7e05d26c2085", size = 5136099, upload-time = "2025-07-02T06:10:07.819Z" }
 wheels = [
@@ -962,16 +966,16 @@ resolution-markers = [
     "python_full_version == '3.11.*'",
 ]
 dependencies = [
-    { name = "jax", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "msgpack", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "optax", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "orbax-checkpoint", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "pyyaml", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "rich", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "tensorstore", version = "0.1.82", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "treescope", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "typing-extensions", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
+    { name = "jax", version = "0.7.1", source = { registry = "https://pypi.org/simple" } },
+    { name = "msgpack" },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "optax" },
+    { name = "orbax-checkpoint" },
+    { name = "pyyaml" },
+    { name = "rich" },
+    { name = "tensorstore", version = "0.1.82", source = { registry = "https://pypi.org/simple" } },
+    { name = "treescope" },
+    { name = "typing-extensions" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/8b/02/2d4fbc31dcf5fb02e2b8ce733f8eecaffc7f24e905df2da39539809b195b/flax-0.12.0.tar.gz", hash = "sha256:cb3d4a2028666640c1a2e5b267dadfbd42ef14c6db41896fb4c765a7f05ebe74", size = 5038989, upload-time = "2025-09-25T23:59:00.919Z" }
 wheels = [
@@ -1083,10 +1087,10 @@ resolution-markers = [
     "python_full_version == '3.13.*'",
 ]
 dependencies = [
-    { name = "dm-haiku", marker = "python_full_version >= '3.13'" },
-    { name = "jax", version = "0.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
-    { name = "matplotlib", marker = "python_full_version >= '3.13'" },
-    { name = "numpyro", marker = "python_full_version >= '3.13'" },
+    { name = "dm-haiku" },
+    { name = "jax", version = "0.5.3", source = { registry = "https://pypi.org/simple" } },
+    { name = "matplotlib" },
+    { name = "numpyro" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/9e/25/74b8c1c77bad231a92674be105a5fd9152fe3a9eb77efa27d342d77fce20/gpax-0.1.8.tar.gz", hash = "sha256:7bc7b89bd58db2d31f9c6007515d3883810e064c69d5b8895543647422a368aa", size = 73263, upload-time = "2024-03-20T06:39:56.206Z" }
 wheels = [
@@ -1103,16 +1107,16 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "dm-haiku", marker = "python_full_version < '3.13'" },
-    { name = "flax", version = "0.10.7", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "dm-haiku" },
+    { name = "flax", version = "0.10.7", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "flax", version = "0.12.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "jax", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "jaxlib", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "jaxopt", marker = "python_full_version < '3.13'" },
-    { name = "matplotlib", marker = "python_full_version < '3.13'" },
-    { name = "numpyro", marker = "python_full_version < '3.13'" },
+    { name = "jaxopt" },
+    { name = "matplotlib" },
+    { name = "numpyro" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/9e/b8/b45ecc6fe6c6d8363f1c66fd07b3a14dc1c1d09c1d0b18d5202a74fb4e91/gpax-0.1.9.tar.gz", hash = "sha256:d86d2738ae3f83a5c5fdb65f6bc421a27d49c46198a9677464f1879a95fcfcca", size = 61256, upload-time = "2025-07-04T06:11:06.435Z" }
 wheels = [
@@ -1233,11 +1237,11 @@ resolution-markers = [
     "python_full_version == '3.13.*'",
 ]
 dependencies = [
-    { name = "jaxlib", version = "0.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
-    { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
-    { name = "opt-einsum", marker = "python_full_version >= '3.13'" },
-    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
+    { name = "jaxlib", version = "0.5.3", source = { registry = "https://pypi.org/simple" } },
+    { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "opt-einsum" },
+    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/13/e5/dabb73ab10330e9535aba14fc668b04a46fcd8e78f06567c4f4f1adce340/jax-0.5.3.tar.gz", hash = "sha256:f17fcb0fd61dc289394af6ce4de2dada2312f2689bb0d73642c6f026a95fbb2c", size = 2072748, upload-time = "2025-03-19T18:23:40.901Z" }
 wheels = [
@@ -1252,11 +1256,11 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "opt-einsum", marker = "python_full_version < '3.11'" },
-    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" } },
+    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } },
+    { name = "opt-einsum" },
+    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/cf/1e/267f59c8fb7f143c3f778c76cb7ef1389db3fd7e4540f04b9f42ca90764d/jax-0.6.2.tar.gz", hash = "sha256:a437d29038cbc8300334119692744704ca7941490867b9665406b7f90665cd96", size = 2334091, upload-time = "2025-06-17T23:10:27.186Z" }
 wheels = [
@@ -1272,11 +1276,11 @@ resolution-markers = [
     "python_full_version == '3.11.*'",
 ]
 dependencies = [
-    { name = "jaxlib", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "opt-einsum", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
+    { name = "jaxlib", version = "0.7.1", source = { registry = "https://pypi.org/simple" } },
+    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "opt-einsum" },
+    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/bc/e8/b393ee314d3b042bd66b986d38e52f4e6046590399d916381265c20467d3/jax-0.7.1.tar.gz", hash = "sha256:118f56338c503361d2791f069d24339d8d44a8db442ed851d2e591222fb7a56d", size = 2428411, upload-time = "2025-08-20T15:55:46.098Z" }
 wheels = [
@@ -1295,9 +1299,9 @@ resolution-markers = [
     "python_full_version == '3.13.*'",
 ]
 dependencies = [
-    { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
-    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
+    { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" } },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/2e/12/b1da8468ad843b30976b0e87c6b344ee621fb75ef8bbd39156a303f59059/jaxlib-0.5.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:48ff5c89fb8a0fe04d475e9ddc074b4879a91d7ab68a51cec5cd1e87f81e6c47", size = 63694868, upload-time = "2025-03-19T18:23:52.193Z" },
@@ -1327,9 +1331,9 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } },
+    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" } },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/15/c5/41598634c99cbebba46e6777286fb76abc449d33d50aeae5d36128ca8803/jaxlib-0.6.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:da4601b2b5dc8c23d6afb293eacfb9aec4e1d1871cb2f29c5a151d103e73b0f8", size = 54298019, upload-time = "2025-06-17T23:10:36.916Z" },
@@ -1361,9 +1365,9 @@ resolution-markers = [
     "python_full_version == '3.11.*'",
 ]
 dependencies = [
-    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
+    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" } },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/8b/af/5058d545e95f99a54289648f5430cc3c23263dd70a1391e7491f24ed328d/jaxlib-0.7.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3f32c3e4c167b7327c342e82d3df84079714ea0b43718be871d039999670b3c9", size = 57686934, upload-time = "2025-08-20T15:55:58.989Z" },
@@ -1395,13 +1399,13 @@ name = "jaxopt"
 version = "0.8.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "jax", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "jaxlib", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/3a/da/ff7d7fbd13b8ed5e8458e80308d075fc649062b9f8676d3fc56f2dc99a82/jaxopt-0.8.5.tar.gz", hash = "sha256:2790bd68ef132b216c083a8bc7a2704eceb35a92c0fc0a1e652e79dfb1e9e9ab", size = 121709, upload-time = "2025-04-14T17:59:01.618Z" }
@@ -1715,7 +1719,7 @@ name = "markdown-it-py"
 version = "4.0.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "mdurl", marker = "python_full_version < '3.13'" },
+    { name = "mdurl" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" }
 wheels = [
@@ -1919,7 +1923,7 @@ resolution-markers = [
     "python_full_version == '3.13.*'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/fd/15/76f86faa0902836cc133939732f7611ace68cf54148487a99c539c272dc8/ml_dtypes-0.4.1.tar.gz", hash = "sha256:fad5f2de464fd09127e49b7fd1252b9006fb43d2edc1ff112d390c324af5ca7a", size = 692594, upload-time = "2024-09-13T19:07:11.624Z" }
 wheels = [
@@ -1947,7 +1951,7 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" }
@@ -2361,7 +2365,7 @@ name = "nvidia-cublas"
 version = "13.1.1.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-cuda-nvrtc" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" },
@@ -2400,7 +2404,7 @@ name = "nvidia-cudnn-cu13"
 version = "9.20.0.48"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-cublas" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" },
@@ -2412,7 +2416,7 @@ name = "nvidia-cufft"
 version = "12.0.0.61"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-nvjitlink" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" },
@@ -2442,9 +2446,9 @@ name = "nvidia-cusolver"
 version = "12.0.4.66"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
-    { name = "nvidia-cusparse", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
-    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-cublas" },
+    { name = "nvidia-cusparse" },
+    { name = "nvidia-nvjitlink" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" },
@@ -2456,7 +2460,7 @@ name = "nvidia-cusparse"
 version = "12.6.3.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.15' and sys_platform == 'emscripten') or (python_full_version < '3.15' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
+    { name = "nvidia-nvjitlink" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },
@@ -2522,12 +2526,12 @@ name = "optax"
 version = "0.2.8"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "absl-py", marker = "python_full_version < '3.13'" },
-    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "absl-py" },
+    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "jax", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "jaxlib", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/8c/f9/e3d11ae6f298ee941a0690e353a323d158ba5dedc436e75621c310845c5c/optax-0.2.8.tar.gz", hash = "sha256:5b225b35066fc3eebaa4d798f1b4173b4d57d1a480610908981f8343b50af0b0", size = 301193, upload-time = "2026-03-20T23:30:05.465Z" }
@@ -2540,25 +2544,25 @@ name = "orbax-checkpoint"
 version = "0.11.36"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "absl-py", marker = "python_full_version < '3.13'" },
-    { name = "aiofiles", marker = "python_full_version < '3.13'" },
-    { name = "etils", version = "1.13.0", source = { registry = "https://pypi.org/simple" }, extra = ["epath", "epy"], marker = "python_full_version < '3.11'" },
+    { name = "absl-py" },
+    { name = "aiofiles" },
+    { name = "etils", version = "1.13.0", source = { registry = "https://pypi.org/simple" }, extra = ["epath", "epy"], marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "etils", version = "1.14.0", source = { registry = "https://pypi.org/simple" }, extra = ["epath", "epy"], marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "humanize", marker = "python_full_version < '3.13'" },
-    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "humanize" },
+    { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "jax", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "msgpack", marker = "python_full_version < '3.13'" },
-    { name = "nest-asyncio", marker = "python_full_version < '3.13' and sys_platform == 'win32'" },
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "msgpack" },
+    { name = "nest-asyncio", marker = "sys_platform == 'win32'" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "protobuf", marker = "python_full_version < '3.13'" },
-    { name = "psutil", marker = "python_full_version < '3.13'" },
-    { name = "pyyaml", marker = "python_full_version < '3.13'" },
-    { name = "simplejson", marker = "python_full_version < '3.13'" },
-    { name = "tensorstore", version = "0.1.78", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "protobuf" },
+    { name = "psutil" },
+    { name = "pyyaml" },
+    { name = "simplejson" },
+    { name = "tensorstore", version = "0.1.78", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "tensorstore", version = "0.1.82", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "typing-extensions", marker = "python_full_version < '3.13'" },
-    { name = "uvloop", marker = "python_full_version < '3.13' and sys_platform != 'win32'" },
+    { name = "typing-extensions" },
+    { name = "uvloop", marker = "sys_platform != 'win32'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/a9/95/2eb9afead8b1aaa5a5184ced09f57c7c0ee473ea283be8e36f22ea92680f/orbax_checkpoint-0.11.36.tar.gz", hash = "sha256:60ed7084a9b79385fb5b9e4b05d98c2db6f5892d05ee8d82df680cfac1622312", size = 585461, upload-time = "2026-04-14T17:03:47.475Z" }
 wheels = [
@@ -3312,8 +3316,8 @@ name = "rich"
 version = "15.0.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "markdown-it-py", marker = "python_full_version < '3.13'" },
-    { name = "pygments", marker = "python_full_version < '3.13'" },
+    { name = "markdown-it-py" },
+    { name = "pygments" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" }
 wheels = [
@@ -3494,10 +3498,10 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "joblib", marker = "python_full_version < '3.11'" },
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "threadpoolctl", marker = "python_full_version < '3.11'" },
+    { name = "joblib" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } },
+    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" } },
+    { name = "threadpoolctl" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/98/c2/a7855e41c9d285dfe86dc50b250978105dce513d6e459ea66a6aeb0e1e0c/scikit_learn-1.7.2.tar.gz", hash = "sha256:20e9e49ecd130598f1ca38a1d85090e1a600147b9c02fa6f15d69cb53d968fda", size = 7193136, upload-time = "2025-09-09T08:21:29.075Z" }
 wheels = [
@@ -3547,10 +3551,10 @@ resolution-markers = [
     "python_full_version == '3.11.*'",
 ]
 dependencies = [
-    { name = "joblib", marker = "python_full_version >= '3.11'" },
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
-    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
-    { name = "threadpoolctl", marker = "python_full_version >= '3.11'" },
+    { name = "joblib" },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" } },
+    { name = "threadpoolctl" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/0e/d4/40988bf3b8e34feec1d0e6a051446b1f66225f8529b9309becaeef62b6c4/scikit_learn-1.8.0.tar.gz", hash = "sha256:9bccbb3b40e3de10351f8f5068e105d0f4083b1a65fa07b6634fbc401a6287fd", size = 7335585, upload-time = "2025-12-10T07:08:53.618Z" }
 wheels = [
@@ -3600,7 +3604,7 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" }
 wheels = [
@@ -3665,7 +3669,7 @@ resolution-markers = [
     "python_full_version == '3.11.*'",
 ]
 dependencies = [
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" }
 wheels = [
@@ -3850,23 +3854,23 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "alabaster", marker = "python_full_version < '3.11'" },
-    { name = "babel", marker = "python_full_version < '3.11'" },
-    { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" },
-    { name = "docutils", version = "0.21.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "imagesize", marker = "python_full_version < '3.11'" },
-    { name = "jinja2", marker = "python_full_version < '3.11'" },
-    { name = "packaging", marker = "python_full_version < '3.11'" },
-    { name = "pygments", marker = "python_full_version < '3.11'" },
-    { name = "requests", marker = "python_full_version < '3.11'" },
-    { name = "snowballstemmer", marker = "python_full_version < '3.11'" },
-    { name = "sphinxcontrib-applehelp", marker = "python_full_version < '3.11'" },
-    { name = "sphinxcontrib-devhelp", marker = "python_full_version < '3.11'" },
-    { name = "sphinxcontrib-htmlhelp", marker = "python_full_version < '3.11'" },
-    { name = "sphinxcontrib-jsmath", marker = "python_full_version < '3.11'" },
-    { name = "sphinxcontrib-qthelp", marker = "python_full_version < '3.11'" },
-    { name = "sphinxcontrib-serializinghtml", marker = "python_full_version < '3.11'" },
-    { name = "tomli", marker = "python_full_version < '3.11'" },
+    { name = "alabaster" },
+    { name = "babel" },
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "docutils", version = "0.21.2", source = { registry = "https://pypi.org/simple" } },
+    { name = "imagesize" },
+    { name = "jinja2" },
+    { name = "packaging" },
+    { name = "pygments" },
+    { name = "requests" },
+    { name = "snowballstemmer" },
+    { name = "sphinxcontrib-applehelp" },
+    { name = "sphinxcontrib-devhelp" },
+    { name = "sphinxcontrib-htmlhelp" },
+    { name = "sphinxcontrib-jsmath" },
+    { name = "sphinxcontrib-qthelp" },
+    { name = "sphinxcontrib-serializinghtml" },
+    { name = "tomli" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/be0b61178fe2cdcb67e2a92fc9ebb488e3c51c4f74a36a7824c0adf23425/sphinx-8.1.3.tar.gz", hash = "sha256:43c1911eecb0d3e161ad78611bc905d1ad0e523e4ddc202a58a821773dc4c927", size = 8184611, upload-time = "2024-10-13T20:27:13.93Z" }
 wheels = [
@@ -3881,23 +3885,23 @@ resolution-markers = [
     "python_full_version == '3.11.*'",
 ]
 dependencies = [
-    { name = "alabaster", marker = "python_full_version == '3.11.*'" },
-    { name = "babel", marker = "python_full_version == '3.11.*'" },
-    { name = "colorama", marker = "python_full_version == '3.11.*' and sys_platform == 'win32'" },
-    { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" },
-    { name = "imagesize", marker = "python_full_version == '3.11.*'" },
-    { name = "jinja2", marker = "python_full_version == '3.11.*'" },
-    { name = "packaging", marker = "python_full_version == '3.11.*'" },
-    { name = "pygments", marker = "python_full_version == '3.11.*'" },
-    { name = "requests", marker = "python_full_version == '3.11.*'" },
-    { name = "roman-numerals", marker = "python_full_version == '3.11.*'" },
-    { name = "snowballstemmer", marker = "python_full_version == '3.11.*'" },
-    { name = "sphinxcontrib-applehelp", marker = "python_full_version == '3.11.*'" },
-    { name = "sphinxcontrib-devhelp", marker = "python_full_version == '3.11.*'" },
-    { name = "sphinxcontrib-htmlhelp", marker = "python_full_version == '3.11.*'" },
-    { name = "sphinxcontrib-jsmath", marker = "python_full_version == '3.11.*'" },
-    { name = "sphinxcontrib-qthelp", marker = "python_full_version == '3.11.*'" },
-    { name = "sphinxcontrib-serializinghtml", marker = "python_full_version == '3.11.*'" },
+    { name = "alabaster" },
+    { name = "babel" },
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "imagesize" },
+    { name = "jinja2" },
+    { name = "packaging" },
+    { name = "pygments" },
+    { name = "requests" },
+    { name = "roman-numerals" },
+    { name = "snowballstemmer" },
+    { name = "sphinxcontrib-applehelp" },
+    { name = "sphinxcontrib-devhelp" },
+    { name = "sphinxcontrib-htmlhelp" },
+    { name = "sphinxcontrib-jsmath" },
+    { name = "sphinxcontrib-qthelp" },
+    { name = "sphinxcontrib-serializinghtml" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/42/50/a8c6ccc36d5eacdfd7913ddccd15a9cee03ecafc5ee2bc40e1f168d85022/sphinx-9.0.4.tar.gz", hash = "sha256:594ef59d042972abbc581d8baa577404abe4e6c3b04ef61bd7fc2acbd51f3fa3", size = 8710502, upload-time = "2025-12-04T07:45:27.343Z" }
 wheels = [
@@ -3917,23 +3921,23 @@ resolution-markers = [
     "python_full_version == '3.12.*'",
 ]
 dependencies = [
-    { name = "alabaster", marker = "python_full_version >= '3.12'" },
-    { name = "babel", marker = "python_full_version >= '3.12'" },
-    { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" },
-    { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
-    { name = "imagesize", marker = "python_full_version >= '3.12'" },
-    { name = "jinja2", marker = "python_full_version >= '3.12'" },
-    { name = "packaging", marker = "python_full_version >= '3.12'" },
-    { name = "pygments", marker = "python_full_version >= '3.12'" },
-    { name = "requests", marker = "python_full_version >= '3.12'" },
-    { name = "roman-numerals", marker = "python_full_version >= '3.12'" },
-    { name = "snowballstemmer", marker = "python_full_version >= '3.12'" },
-    { name = "sphinxcontrib-applehelp", marker = "python_full_version >= '3.12'" },
-    { name = "sphinxcontrib-devhelp", marker = "python_full_version >= '3.12'" },
-    { name = "sphinxcontrib-htmlhelp", marker = "python_full_version >= '3.12'" },
-    { name = "sphinxcontrib-jsmath", marker = "python_full_version >= '3.12'" },
-    { name = "sphinxcontrib-qthelp", marker = "python_full_version >= '3.12'" },
-    { name = "sphinxcontrib-serializinghtml", marker = "python_full_version >= '3.12'" },
+    { name = "alabaster" },
+    { name = "babel" },
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "docutils", version = "0.22.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "imagesize" },
+    { name = "jinja2" },
+    { name = "packaging" },
+    { name = "pygments" },
+    { name = "requests" },
+    { name = "roman-numerals" },
+    { name = "snowballstemmer" },
+    { name = "sphinxcontrib-applehelp" },
+    { name = "sphinxcontrib-devhelp" },
+    { name = "sphinxcontrib-htmlhelp" },
+    { name = "sphinxcontrib-jsmath" },
+    { name = "sphinxcontrib-qthelp" },
+    { name = "sphinxcontrib-serializinghtml" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/cd/bd/f08eb0f4eed5c83f1ba2a3bd18f7745a2b1525fad70660a1c00224ec468a/sphinx-9.1.0.tar.gz", hash = "sha256:7741722357dd75f8190766926071fed3bdc211c74dd2d7d4df5404da95930ddb", size = 8718324, upload-time = "2025-12-31T15:09:27.646Z" }
 wheels = [
@@ -4084,8 +4088,8 @@ resolution-markers = [
     "python_full_version < '3.11'",
 ]
 dependencies = [
-    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/9f/ee/05eb424437f4db63331c90e4605025eedc0f71da3faff97161d5d7b405af/tensorstore-0.1.78.tar.gz", hash = "sha256:e26074ffe462394cf54197eb76d6569b500f347573cd74da3f4dd5f510a4ad7c", size = 6913502, upload-time = "2025-10-06T17:44:29.649Z" }
 wheels = [
@@ -4120,8 +4124,8 @@ resolution-markers = [
     "python_full_version == '3.11.*'",
 ]
 dependencies = [
-    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
-    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
+    { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } },
+    { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/cd/9b/43aedb544937f214dd7c665a7edf1b8b74f2f55d53ebd351c0ce69acf81a/tensorstore-0.1.82.tar.gz", hash = "sha256:ccfceffb7611fc61330f6da24b8b0abd9251d480ac8a5bac5a1729f9ed0c3a9f", size = 7160364, upload-time = "2026-03-13T00:22:16.888Z" }
 wheels = [
@@ -4279,7 +4283,7 @@ name = "treescope"
 version = "0.1.10"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
+    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/f0/2a/d13d3c38862632742d2fe2f7ae307c431db06538fd05ca03020d207b5dcc/treescope-0.1.10.tar.gz", hash = "sha256:20f74656f34ab2d8716715013e8163a0da79bdc2554c16d5023172c50d27ea95", size = 138870, upload-time = "2025-08-08T05:43:48.048Z" }

From f64c33d24e7a2ae079d9b25467aa8dc606f2822c Mon Sep 17 00:00:00 2001
From: Lance-Drane <Lance-Drane@users.noreply.github.com>
Date: Thu, 23 Jul 2026 21:31:03 -0400
Subject: [PATCH 5/8] numpyro version restriction only applies to gpax

Signed-off-by: Lance-Drane <Lance-Drane@users.noreply.github.com>
---
 pyproject.toml | 6 ++++--
 uv.lock        | 4 ++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index eb50546..42f3188 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,7 +22,6 @@ dependencies = [
   "numpy",
   "scikit-learn>=1.4.0,<2.0.0",   # TODO consider making an optional dependency group
   "scipy>=1.12.0,<2.0.0",
-  "numpyro<0.20.1",               # depend on older numpyro for gpax (numpyro.contrib.module: random_haiku_module)
   "pymongo>=4.12.1",              # TODO - this is only needed for dial_service, dial_dataclass can use a simple fixture for ObjectID representation which allows it to skip this dependency
 ]
 
@@ -34,7 +33,10 @@ Issues = "https://github.com/INTERSECT-DIAL/dial/issues"
 
 [project.optional-dependencies]
 docs = ["sphinx>=5.3.0", "furo>=2023.3.27"]
-gpax = ["gpax>=0.1.8"]
+gpax = [
+  "gpax>=0.1.8",
+  "numpyro<0.20.1", # depend on older numpyro for gpax (numpyro.contrib.module: random_haiku_module)
+]
 sable = ["sable @ git+https://code.ornl.gov/sable/sable.git"]
 
 [dependency-groups]
diff --git a/uv.lock b/uv.lock
index 66d3f1b..89d216d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -741,7 +741,6 @@ dependencies = [
     { name = "intersect-sdk" },
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
-    { name = "numpyro" },
     { name = "pymongo" },
     { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "scikit-learn", version = "1.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
@@ -759,6 +758,7 @@ docs = [
 gpax = [
     { name = "gpax", version = "0.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
     { name = "gpax", version = "0.1.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" },
+    { name = "numpyro" },
 ]
 sable = [
     { name = "sable" },
@@ -782,7 +782,7 @@ requires-dist = [
     { name = "gpax", marker = "extra == 'gpax'", specifier = ">=0.1.8" },
     { name = "intersect-sdk", specifier = ">=0.9.3,<0.10.0" },
     { name = "numpy" },
-    { name = "numpyro", specifier = "<0.20.1" },
+    { name = "numpyro", marker = "extra == 'gpax'", specifier = "<0.20.1" },
     { name = "pymongo", specifier = ">=4.12.1" },
     { name = "sable", marker = "extra == 'sable'", git = "https://code.ornl.gov/sable/sable.git" },
     { name = "scikit-learn", specifier = ">=1.4.0,<2.0.0" },

From be7b2f6012cf0270533a63f768552e1aadc03310 Mon Sep 17 00:00:00 2001
From: Lance-Drane <Lance-Drane@users.noreply.github.com>
Date: Fri, 24 Jul 2026 02:57:41 -0400
Subject: [PATCH 6/8] add guard to skip tests if backends are not installed

Signed-off-by: Lance-Drane <Lance-Drane@users.noreply.github.com>
---
 tests/__init__.py                  |    0
 tests/benchmarks/__init__.py       |    0
 tests/benchmarks/test_strainmap.py | 1736 ++++++++++++++--------------
 tests/helpers.py                   |   30 +
 tests/unit/__init__.py             |    0
 tests/unit/test_internals.py       |  121 +-
 6 files changed, 1025 insertions(+), 862 deletions(-)
 create mode 100644 tests/__init__.py
 create mode 100644 tests/benchmarks/__init__.py
 create mode 100644 tests/helpers.py
 create mode 100644 tests/unit/__init__.py

diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/benchmarks/__init__.py b/tests/benchmarks/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/benchmarks/test_strainmap.py b/tests/benchmarks/test_strainmap.py
index e91fa3e..3dd9725 100644
--- a/tests/benchmarks/test_strainmap.py
+++ b/tests/benchmarks/test_strainmap.py
@@ -1,846 +1,890 @@
-"""
-Benchmark which is meant to test core DIAL functionality (without the INTERSECT or MONGO pieces) for the Strain-Mapping benchmark problem.
-"""
-
-import logging
-import sys
-from dataclasses import dataclass
-from pathlib import Path
-
-import numpy as np
-import pandas as pd
-import pytest
-from scipy.interpolate import LinearNDInterpolator
-
-# from pytest_benchmark.fixture import BenchmarkFixture
-from dial_dataclass import (
-    DialInputPredictions,
-    DialInputSingleOtherStrategy,
-    Normal,
-)
-from dial_service import (
-    core as dial_core,
-)
-from dial_service.serverside_data import (
-    ServersideInputBase,
-    ServersideInputPrediction,
-    ServersideInputSingle,
-)
-from dial_service.service_specific_dataclasses import DialWorkflowCreationParamsService
-
-logger = logging.getLogger(__name__)
-
-MOCK_WORKFLOW_ID = '6984e6a6ef6e6290dabced91'
-"""fake ObjectID for testing purposes, we do not actually interact with a DB in these tests."""
-
-
-###############
-low_f_data_name = (
-    Path(__file__).parents[1] / 'fixtures' / 'adaptive_strain_manufacturing_low_fidelity.csv'
-)
-high_f_data_name = (
-    Path(__file__).parents[1] / 'fixtures' / 'adaptive_strain_manufacturing_high_fidelity.csv'
-)
-
-
-def read_and_prepare_data():
-    # normalization helper
-    def normalize_data(in_data):
-        m_data = (in_data.max() + in_data.min()) * 0.5
-        d_data = (in_data.max() - in_data.min()) * 0.5
-        return (in_data - m_data) / d_data
-
-    # define constants and read data
-    # wall_st_index = 0
-    # wall_st_end = 676
-    x_col_index = 0  # z follow next
-    e_col_index = 2  # e11, e22 and e33
-    nrows = 26  # Rows in sim strain map of the wall
-    ncols = 26  # Cols in sim strain map of the wall
-    data = pd.read_csv(low_f_data_name)
-
-    Y, Z, E11, E22, E33, R11, R22, R33 = np.array(
-        [
-            data.values[:, x_col_index],
-            data.values[:, x_col_index + 1],
-            data.values[:, e_col_index],
-            data.values[:, e_col_index + 1],
-            data.values[:, e_col_index + 2],
-            data.values[:, e_col_index + 3],
-            data.values[:, e_col_index + 4],
-            data.values[:, e_col_index + 5],
-        ]
-    )
-
-    # Convert data
-    x1 = Y.astype(np.float32).reshape(nrows, ncols)
-    x2 = Z.astype(np.float32).reshape(nrows, ncols)
-
-    # Sim_e11 = np.array(E11).astype(np.float32).reshape(nrows, ncols)
-    # Real_e11 = np.array(R11).astype(np.float32).reshape(nrows, ncols)
-
-    # Sim_e22 = np.array(E22).astype(np.float32).reshape(nrows, ncols)
-    # Real_e22 = np.array(R22).astype(np.float32).reshape(nrows, ncols)
-
-    # Sim_e33 = np.array(E33).astype(np.float32).reshape(nrows, ncols)
-    Real_e33 = np.array(R33).astype(np.float32).reshape(nrows, ncols)
-
-    # Normalize data
-    x1_norm = normalize_data(x1)
-    x2_norm = normalize_data(x2)
-
-    # Sim_e11_norm = normalize_data(Sim_e11)
-    # Real_e11_norm = normalize_data(Real_e11)
-
-    # Sim_e22_norm = normalize_data(Sim_e22)
-    # Real_e22_norm = normalize_data(Real_e22)
-
-    # Sim_e33_norm = normalize_data(Sim_e33)
-    Real_e33_norm = normalize_data(Real_e33)
-
-    # read high fidelity data
-    # data = pd.read_csv(high_f_data_name)
-
-    # Y, Z, E11, E22, E33, R11, R22, R33 = np.array(
-    #    [
-    #        data.values[:, x_col_index],
-    #        data.values[:, x_col_index + 1],
-    #        data.values[:, e_col_index],
-    #        data.values[:, e_col_index + 1],
-    #        data.values[:, e_col_index + 2],
-    #        data.values[:, e_col_index + 3],
-    #        data.values[:, e_col_index + 4],
-    #        data.values[:, e_col_index + 5],
-    #    ]
-    # )
-
-    # Sim_hi_e11 = np.array(E11).astype(np.float32).reshape(nrows, ncols)
-    # Sim_hi_e22 = np.array(E22).astype(np.float32).reshape(nrows, ncols)
-    # Sim_hi_e33 = np.array(E33).astype(np.float32).reshape(nrows, ncols)
-
-    # Transpose the high-fidelity simulation data
-    # TODO, figure out why this is transposed
-    # Sim_hi_e11 = Sim_hi_e11.T
-    # Sim_hi_e22 = Sim_hi_e22.T
-    # Sim_hi_e33 = Sim_hi_e33.T
-
-    # Normalize data
-    # Sim_hi_e11_norm = normalize_data(Sim_hi_e11)
-    # Sim_hi_e22_norm = normalize_data(Sim_hi_e22)
-    # Sim_hi_e33_norm = normalize_data(Sim_hi_e33)
-
-    # return data that will be used for the benchmark
-    return nrows, (x1_norm, x2_norm), Real_e33_norm
-
-
-###############
-
-## Select data for ground truth
-
-ngrid, x_grids, y_grid = read_and_prepare_data()
-
-# default inputs
-INITIAL_BOUNDS = [[-1.0, 1.0], [-1.0, 1.0]]
-NUM_DIMS = len(INITIAL_BOUNDS)
-
-MESHGRID_SIZE = ngrid
-INITIAL_MESHGRIDS = x_grids
-INITIAL_POINTS_TO_PREDICT = np.hstack([mg.reshape(-1, 1) for mg in INITIAL_MESHGRIDS])
-
-INITIAL_PREDICTIONS = y_grid.reshape(-1, 1)
-
-###############
-
-NOISE_LEVEL = 1.0e-2
-
-truth_interp = LinearNDInterpolator(INITIAL_POINTS_TO_PREDICT, INITIAL_PREDICTIONS)
-
-
-@dataclass
-class StrainMap:
-    """Strain Map function."""
-
-    noise_level: float
-
-    truth_interp: LinearNDInterpolator
-    """Strain interpolator that is used as 'ground_truth'"""
-
-    def strain_map(self, x) -> float:
-        """
-        Represents a measured strain as a function of two simulation parameters.
-        """
-        x = np.asarray(x).reshape((-1, 2))
-        x1, x2 = x[:, 0], x[:, 1]
-
-        y_true = truth_interp(np.asarray(x1), np.asarray(x2))
-
-        # print("Evaluated truth model", np.hstack((x, y_true)))
-
-        y_noise = y_true + self.noise_level * np.random.normal(size=y_true.shape)
-        return y_noise
-
-
-# build a strain map from the selected truth values
-truth_strain_map = StrainMap(truth_interp=truth_interp, noise_level=NOISE_LEVEL)
-
-###############
-
-# test parameters
-INITIAL_NUM_POINTS = 25
-MAX_ITERATIONS = 150  # only allow a maximum of this many iterations in tests
-
-INITIAL_DATASET_X = np.random.uniform(-1.0, 1.0, size=(INITIAL_NUM_POINTS, NUM_DIMS)).tolist()
-
-TARGET_ERROR = 0.11
-
-
-def run_simulation(
-    dataset_x: list[list[float]], dataset_y: list[float], strategy: str, strategy_args: object
-) -> tuple[np.ndarray, np.ndarray]:
-    # important "Hyper-parameters"
-    length_scale = 0.4
-    prior_std = 1.0
-
-    # Assume that there is some small noise in the measurements to stabilize the fit
-    statistics_y = Normal(loc='y', scale=NOISE_LEVEL)
-
-    # modify a local copy of strategy_args
-    strategy_args = strategy_args.copy()
-    backend = strategy_args.pop('backend', 'sklearn')
-
-    if backend == 'sable':
-        # train model with new data
-        kernel = 'rbf'
-        kernel_args = {
-            'sigma_range': [5e-2, 1.0],
-            'gamma': 0.5,
-        }
-        backend_args = {
-            'prior_std': prior_std,
-        }
-
-    else:
-        # train model with new data
-        kernel = 'matern'
-        kernel_args = {
-            'length_scale': length_scale,
-            'constant_value': prior_std**2,
-        }
-        backend_args = {}
-
-    client_state = DialWorkflowCreationParamsService(
-        dataset_x=dataset_x,
-        dataset_y=dataset_y,
-        statistics_y=statistics_y,
-        bounds=INITIAL_BOUNDS,
-        kernel=kernel,
-        kernel_args=kernel_args,
-        y_is_good=False,
-        backend=backend,
-        backend_args=backend_args,
-        seed=-1,
-        dim_x=2,
-    )
-
-    data = ServersideInputBase(client_state)
-    model = dial_core.train_model(data)
-
-    # use get_surrogate_values to predict mean and standard deviation on the inital points grid
-    data = ServersideInputPrediction(
-        client_state,
-        DialInputPredictions(
-            workflow_id=MOCK_WORKFLOW_ID,
-            points_to_predict=INITIAL_POINTS_TO_PREDICT,
-        ),
-    )
-    surrogate_mean, surrogate_std, _ = dial_core.get_surrogate_values(data, model)
-    mean_grid = np.array(surrogate_mean).reshape((-1, 1))
-
-    # subtract the true values and save mean absolute error and standard deviation
-    err_grid = np.abs(mean_grid - INITIAL_PREDICTIONS).reshape((MESHGRID_SIZE,) * NUM_DIMS)
-    std_grid = np.array(surrogate_std).reshape((MESHGRID_SIZE,) * NUM_DIMS)
-
-    # get_next_point
-    data = ServersideInputSingle(
-        client_state,
-        DialInputSingleOtherStrategy(
-            workflow_id=MOCK_WORKFLOW_ID,
-            bounds=INITIAL_BOUNDS,
-            y_is_good=False,
-            seed=-1,
-            strategy=strategy,
-            strategy_args=strategy_args,
-            discrete_measurements=True,
-            discrete_measurement_grid_size=[MESHGRID_SIZE, MESHGRID_SIZE],
-        ),
-    )
-    next_point = dial_core.get_next_point(data, model)
-
-    dataset_x.append(next_point)
-
-    # compute at next point
-    next_point_y = truth_strain_map.strain_map(next_point).reshape(-1).tolist()
-    dataset_y.append(next_point_y[0])
-
-    return err_grid, std_grid
-
-
-def graph(err_grid, dataset_x, strategy):
-    try:
-        import matplotlib as mpl
-        import matplotlib.pyplot as plt
-
-        mpl.use('Agg')  # Use non-interactive backend
-    except ImportError:
-        logger.error(  # noqa: TRY400
-            'Error: matplotlib is required for generating plots. Install it with: pip install matplotlib'
-        )
-        return
-
-    plt.clf()
-    data = np.array(err_grid)
-    plt.contourf(
-        INITIAL_MESHGRIDS[0],
-        INITIAL_MESHGRIDS[1],
-        data,
-        levels=np.logspace(-2, 0, 10),
-        norm='log',
-        extend='both',
-    )
-    cbar = plt.colorbar()
-    cbar.set_ticks(np.logspace(-2, 0, 7))
-    cbar.set_label('Overall error')
-    plt.xlabel('Simulation Parameter #1')
-    plt.ylabel('Simulation Parameter #2')
-    # add black dots for data points and a red marker for the recommendation:
-    X_train = np.array(dataset_x)
-    plt.scatter(X_train[:, 0], X_train[:, 1], color='black', marker='o')
-    plt.scatter(1.0, 1.0, s=300, color='None', edgecolors='black', marker='o')
-
-    plt.savefig(f'graph_{strategy}.png')
-
-
-def accuracy_benchmark(
-    strategy: str, strategy_args: object, max_iterations=MAX_ITERATIONS
-) -> tuple[int, float]:
-    """
-    returns:
-      - number of iterations taken to reach an acceptably accurate target
-      - target value achieved
-    """
-
-    iterations = 0
-
-    initial_dataset_x = np.random.uniform(-1.0, 1.0, size=(INITIAL_NUM_POINTS, NUM_DIMS)).tolist()
-
-    dataset_x = initial_dataset_x
-    dataset_y = truth_strain_map.strain_map(dataset_x).reshape(-1).tolist()
-
-    # run simulations until we reach an acceptable target range
-    while iterations < max_iterations:
-        try:
-            err_grid, std_grid = run_simulation(dataset_x, dataset_y, strategy, strategy_args)
-        except Exception as e:
-            logger.exception('Error during simulation')
-            raise AssertionError from e
-
-        # guess has no meaning here, take the last acquired datapoint
-        guess = dataset_x[-1]
-
-        # compute the RMSE and MAD on the grid and the average
-        # rmse = np.sqrt(err_grid**2 + std_grid**2)
-        mad = np.abs(err_grid)
-        # total_rmse = np.sqrt(np.mean(rmse**2))
-        total_mad = np.sqrt(np.mean(mad**2))
-
-        # pick the MAD over the RMSE as an error criterion, it only measures the mean deviation:
-        # it is less sensitive to wrong std and hyperparameter calibration, but does not measure how well
-        # the "error bars" (std_grid) quantify the actual ty
-        error = total_mad
-        print(iterations, error)
-
-        graph(mad, dataset_x, f'{strategy}_{strategy_args.get("backend", "sklearn")}')
-
-        if error <= TARGET_ERROR:
-            break
-        iterations += 1
-
-    return iterations, error, guess
-
-
-TEST_PARAMS = (
-    ('strategy', 'strategy_args', 'max_iterations'),
-    [
-        ('uncertainty', {}, 2),
-        ('uncertainty', {'backend': 'sable'}, 2),
-        ('random', {}, 2),
-        # use low number of iterations for unit test
-    ],
-)
-
-
-@pytest.mark.parametrize(*TEST_PARAMS)
-def test_benchmark_strainmap_accuracy(
-    # benchmark: BenchmarkFixture,
-    strategy: str,
-    strategy_args: dict,
-    max_iterations: int,
-) -> None:
-    # NUM_RUNS = 20
-    # for _ in range(NUM_RUNS):
-    # iterations, target = benchmark(
-    # partial(accuracy_benchmark, strategy, strategy_args)
-    # )
-    iterations, target, guess = accuracy_benchmark(strategy, strategy_args, max_iterations)
-    print(
-        'Iterations for',
-        strategy,
-        ': ',
-        iterations,
-        ' best guess:',
-        guess,
-        ' with target value:',
-        target,
-    )
-    print(
-        'Maximum early terminus value',
-        TARGET_ERROR,
-        ' with ',
-        max_iterations,
-        ' maximum iterations.',
-    )
-    print(
-        'Accuracy benchmark for strategy:',
-        strategy,
-        'reached' if iterations <= max_iterations else 'not reached',
-    )
-
-
-if __name__ == '__main__':
-    """Generate HTML benchmark report with plots comparing different strategies."""
-    import argparse
-    import datetime
-    import json
-    from pathlib import Path
-
-    logger = logging.getLogger(f'{__name__}_runner')
-
-    try:
-        import matplotlib as mpl
-        import matplotlib.pyplot as plt
-
-        mpl.use('Agg')  # Use non-interactive backend
-    except ImportError:
-        logger.error(  # noqa: TRY400
-            'Error: matplotlib is required for generating plots. Install it with: pip install matplotlib'
-        )
-        sys.exit(1)
-
-    def positive_int_type(arg):
-        try:
-            val = int(arg)
-        except ValueError as e:
-            msg = 'Must be an integer'
-            raise argparse.ArgumentTypeError(msg) from e
-        if val < 1:
-            msg = 'Argument must be a positive number'
-            raise argparse.ArgumentTypeError(msg)
-        return val
-
-    parser = argparse.ArgumentParser(description='Generate the Strainmap HTML benchmark pages.')
-    parser.add_argument(
-        '--num-runs',
-        '-n',
-        type=positive_int_type,
-        default=3,
-        help='Number of runs for each strategy.',
-    )
-    args = parser.parse_args()
-
-    strategies = TEST_PARAMS[1]
-
-    # Run multiple iterations for statistical analysis
-    NUM_RUNS = args.num_runs
-    logger.info('Running benchmarks with %d iterations per strategy...', NUM_RUNS)
-
-    results = {}
-    for strategy, strategy_args, _ in strategies:
-        strategy_name = f'{strategy}' + (f' {json.dumps(strategy_args)}' if strategy_args else '')
-        logger.info('\nBenchmarking: %s', strategy_name)
-
-        iterations_list = []
-        targets_list = []
-        guesses_list = []
-
-        for run in range(NUM_RUNS):
-            iterations, target, guess = accuracy_benchmark(strategy, strategy_args)
-            iterations_list.append(iterations)
-            targets_list.append(target)
-            guesses_list.append(guess)
-            logger.info(
-                '  Run %d/%d: iterations=%d, target=%.4f', run + 1, NUM_RUNS, iterations, target
-            )
-
-        results[strategy_name] = {
-            'strategy': strategy,
-            'strategy_args': strategy_args,
-            'iterations': iterations_list,
-            'targets': targets_list,
-            'guesses': guesses_list,
-            'avg_iterations': np.mean(iterations_list),
-            'std_iterations': np.std(iterations_list),
-            'avg_target': np.mean(targets_list),
-            'std_target': np.std(targets_list),
-            'success_rate': sum(1 for t in targets_list if t <= TARGET_ERROR) / NUM_RUNS * 100,
-        }
-
-    # Generate plots
-    output_dir = Path('reports/benchmarks')
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
-    fig.suptitle('Strainmap Optimization Benchmark Comparison', fontsize=16, fontweight='bold')
-
-    # Plot 1: Average iterations to convergence
-    ax1 = axes[0, 0]
-    strategy_names = list(results.keys())
-    avg_iterations = [results[s]['avg_iterations'] for s in strategy_names]
-    std_iterations = [results[s]['std_iterations'] for s in strategy_names]
-    bars1 = ax1.bar(
-        range(len(strategy_names)), avg_iterations, yerr=std_iterations, capsize=5, alpha=0.7
-    )
-    ax1.set_xlabel('Strategy')
-    ax1.set_ylabel('Average Iterations')
-    ax1.set_title('Iterations to Convergence (Lower is Better)')
-    ax1.set_xticks(range(len(strategy_names)))
-    ax1.set_xticklabels(
-        [s.replace(' ', '\n') for s in strategy_names], rotation=0, ha='center', fontsize=8
-    )
-    ax1.grid(axis='y', alpha=0.3)
-
-    # Add value labels on bars
-    for _i, (bar, val, std) in enumerate(zip(bars1, avg_iterations, std_iterations, strict=False)):
-        height = bar.get_height()
-        ax1.text(
-            bar.get_x() + bar.get_width() / 2.0,
-            height,
-            f'{val:.1f}±{std:.1f}',
-            ha='center',
-            va='bottom',
-            fontsize=9,
-        )
-
-    # Plot 2: Average target value achieved
-    ax2 = axes[0, 1]
-    avg_targets = [results[s]['avg_target'] for s in strategy_names]
-    std_targets = [results[s]['std_target'] for s in strategy_names]
-    bars2 = ax2.bar(
-        range(len(strategy_names)),
-        avg_targets,
-        yerr=std_targets,
-        capsize=5,
-        alpha=0.7,
-        color='orange',
-    )
-    ax2.set_xlabel('Strategy')
-    ax2.set_ylabel('Average Target Value')
-    ax2.set_title('Final Target Value (Lower is Better)')
-    ax2.set_xticks(range(len(strategy_names)))
-    ax2.set_xticklabels(
-        [s.replace(' ', '\n') for s in strategy_names], rotation=0, ha='center', fontsize=8
-    )
-    ax2.axhline(
-        y=TARGET_ERROR, color='r', linestyle='--', label=f'Target Threshold ({TARGET_ERROR})'
-    )
-    ax2.legend()
-    ax2.grid(axis='y', alpha=0.3)
-
-    # Add value labels on bars
-    for _i, (bar, val, std) in enumerate(zip(bars2, avg_targets, std_targets, strict=False)):
-        height = bar.get_height()
-        ax2.text(
-            bar.get_x() + bar.get_width() / 2.0,
-            height,
-            f'{val:.2f}±{std:.2f}',
-            ha='center',
-            va='bottom',
-            fontsize=9,
-        )
-
-    # Plot 3: Success rate
-    ax3 = axes[1, 0]
-    success_rates = [results[s]['success_rate'] for s in strategy_names]
-    bars3 = ax3.bar(range(len(strategy_names)), success_rates, alpha=0.7, color='green')
-    ax3.set_xlabel('Strategy')
-    ax3.set_ylabel('Success Rate (%)')
-    ax3.set_title(f'Success Rate (Target ≤ {TARGET_ERROR})')
-    ax3.set_xticks(range(len(strategy_names)))
-    ax3.set_xticklabels(
-        [s.replace(' ', '\n') for s in strategy_names], rotation=0, ha='center', fontsize=8
-    )
-    ax3.set_ylim([0, 105])
-    ax3.grid(axis='y', alpha=0.3)
-
-    # Add value labels on bars
-    for bar, val in zip(bars3, success_rates, strict=False):
-        height = bar.get_height()
-        ax3.text(
-            bar.get_x() + bar.get_width() / 2.0,
-            height,
-            f'{val:.1f}%',
-            ha='center',
-            va='bottom',
-            fontsize=9,
-        )
-
-    # Plot 4: Box plot of iterations distribution
-    ax4 = axes[1, 1]
-    iterations_data = [results[s]['iterations'] for s in strategy_names]
-    bp = ax4.boxplot(iterations_data, patch_artist=True)
-    for patch in bp['boxes']:
-        patch.set_facecolor('lightblue')
-    ax4.set_xlabel('Strategy')
-    ax4.set_ylabel('Iterations')
-    ax4.set_title('Iterations Distribution')
-    ax4.set_xticklabels([s.replace(' ', '\n') for s in strategy_names], fontsize=8)
-    ax4.grid(axis='y', alpha=0.3)
-
-    plt.tight_layout()
-    plot_path = output_dir / 'strainmap_benchmark.png'
-    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
-    logger.info('✓ Plot saved to %s', plot_path)
-    plt.close()
-
-    # Generate HTML report
-    html_content = f"""<!DOCTYPE html>
-<html lang="en">
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>Strainmap Optimization Benchmark Report</title>
-    <style>
-        body {{
-            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-            max-width: 1200px;
-            margin: 0 auto;
-            padding: 20px;
-            background-color: #f5f5f5;
-        }}
-        h1 {{
-            color: #2c3e50;
-            border-bottom: 3px solid #3498db;
-            padding-bottom: 10px;
-        }}
-        h2 {{
-            color: #34495e;
-            margin-top: 30px;
-        }}
-        .metadata {{
-            background-color: #ecf0f1;
-            padding: 15px;
-            border-radius: 5px;
-            margin-bottom: 20px;
-        }}
-        .metadata p {{
-            margin: 5px 0;
-        }}
-        table {{
-            width: 100%;
-            border-collapse: collapse;
-            margin: 20px 0;
-            background-color: white;
-            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-        }}
-        th, td {{
-            padding: 12px;
-            text-align: left;
-            border-bottom: 1px solid #ddd;
-        }}
-        th {{
-            background-color: #3498db;
-            color: white;
-            font-weight: bold;
-        }}
-        tr:hover {{
-            background-color: #f1f1f1;
-        }}
-        .best {{
-            background-color: #d4edda;
-            font-weight: bold;
-        }}
-        .plot {{
-            text-align: center;
-            margin: 30px 0;
-            background-color: white;
-            padding: 20px;
-            border-radius: 5px;
-            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-        }}
-        .plot img {{
-            max-width: 100%;
-            height: auto;
-        }}
-        .summary {{
-            background-color: #d1ecf1;
-            border-left: 4px solid #0c5460;
-            padding: 15px;
-            margin: 20px 0;
-        }}
-        code {{
-            background-color: #f4f4f4;
-            padding: 2px 6px;
-            border-radius: 3px;
-            font-family: 'Courier New', monospace;
-        }}
-    </style>
-</head>
-<body>
-    <h1>📊 Strainmap Optimization Benchmark Report</h1>
-
-    <div class="metadata">
-        <p><strong>Generated:</strong> {datetime.datetime.now(tz=datetime.timezone.utc).astimezone().strftime('%Y-%m-%d %H:%M:%S')}</p>
-        <p><strong>Number of Runs per Strategy:</strong> {NUM_RUNS}</p>
-        <p><strong>Maximum Iterations:</strong> {MAX_ITERATIONS}</p>
-        <p><strong>Target Threshold:</strong> {TARGET_ERROR}</p>
-        <p><strong>Initial Bounds:</strong> {INITIAL_BOUNDS}</p>
-        <p><strong>Initial Points:</strong> {INITIAL_NUM_POINTS}</p>
-    </div>
-
-    <div class="summary">
-        <h3>🎯 Test Objective</h3>
-        <p>This benchmark evaluates different acquisition strategies for Bayesian optimization on a strain mapping experiment dataset.
-        The goal is to minimize the total surrogate error (over all inputs) within {MAX_ITERATIONS} iterations, starting from {INITIAL_NUM_POINTS} initial points.</p>
-    </div>
-
-    <h2>📈 Benchmark Results</h2>
-
-    <div class="plot">
-        <h3>Performance Comparison</h3>
-        <img src="strainmap_benchmark.png" alt="Benchmark comparison plots">
-    </div>
-
-    <h2>📋 Detailed Statistics</h2>
-
-    <table>
-        <thead>
-            <tr>
-                <th>Strategy</th>
-                <th>Strategy Args</th>
-                <th>Avg Iterations</th>
-                <th>Std Iterations</th>
-                <th>Avg Target Value</th>
-                <th>Std Target Value</th>
-                <th>Success Rate</th>
-            </tr>
-        </thead>
-        <tbody>
-"""
-
-    # Find best performers
-    best_iterations_idx = min(range(len(strategy_names)), key=lambda i: avg_iterations[i])
-    best_target_idx = min(range(len(strategy_names)), key=lambda i: avg_targets[i])
-    best_success_idx = max(range(len(strategy_names)), key=lambda i: success_rates[i])
-
-    for i, strategy_name in enumerate(strategy_names):
-        result = results[strategy_name]
-        row_class = ''
-        if i in (best_iterations_idx, best_target_idx, best_success_idx):
-            row_class = ' class="best"'
-
-        args_str = json.dumps(result['strategy_args']) if result['strategy_args'] else 'None'
-
-        html_content += f"""            <tr{row_class}>
-                <td><code>{result['strategy']}</code></td>
-                <td><code>{args_str}</code></td>
-                <td>{result['avg_iterations']:.2f}</td>
-                <td>{result['std_iterations']:.2f}</td>
-                <td>{result['avg_target']:.4f}</td>
-                <td>{result['std_target']:.4f}</td>
-                <td>{result['success_rate']:.1f}%</td>
-            </tr>
-"""
-
-    html_content += """        </tbody>
-    </table>
-
-    <h2>🏆 Key Findings</h2>
-    <div class="summary">
-"""
-
-    # Add findings
-    best_strategy = strategy_names[best_iterations_idx]
-    html_content += f"""        <p><strong>Fastest Convergence:</strong> <code>{best_strategy}</code> with {avg_iterations[best_iterations_idx]:.2f} ± {std_iterations[best_iterations_idx]:.2f} iterations on average</p>
-"""
-
-    best_accuracy_strategy = strategy_names[best_target_idx]
-    html_content += f"""        <p><strong>Best Accuracy:</strong> <code>{best_accuracy_strategy}</code> with average target value {avg_targets[best_target_idx]:.4f} ± {std_targets[best_target_idx]:.4f}</p>
-"""
-
-    best_reliability_strategy = strategy_names[best_success_idx]
-    html_content += f"""        <p><strong>Most Reliable:</strong> <code>{best_reliability_strategy}</code> with {success_rates[best_success_idx]:.1f}% success rate</p>
-"""
-
-    html_content += """    </div>
-
-    <h2>📊 Raw Data</h2>
-    <details>
-        <summary>Click to expand raw results JSON</summary>
-        <pre style="background-color: #f4f4f4; padding: 15px; border-radius: 5px; overflow-x: auto;">
-"""
-
-    # Prepare JSON-serializable results
-    json_results = {}
-    for strategy_name, result in results.items():
-        json_results[strategy_name] = {
-            'strategy': result['strategy'],
-            'strategy_args': result['strategy_args'],
-            'iterations': result['iterations'],
-            'targets': result['targets'],
-            'guesses': [[float(x) for x in guess] for guess in result['guesses']],
-            'statistics': {
-                'avg_iterations': float(result['avg_iterations']),
-                'std_iterations': float(result['std_iterations']),
-                'avg_target': float(result['avg_target']),
-                'std_target': float(result['std_target']),
-                'success_rate': float(result['success_rate']),
-            },
-        }
-
-    html_content += json.dumps(json_results, indent=2)
-    html_content += """
-        </pre>
-    </details>
-
-    <footer style="margin-top: 50px; padding-top: 20px; border-top: 1px solid #ddd; color: #7f8c8d; text-align: center;">
-        <p>Generated by test_strainmap.py benchmark suite</p>
-    </footer>
-</body>
-</html>
-"""
-
-    # Save HTML report
-    html_path = output_dir / 'strainmap_benchmark.html'
-    html_path.write_text(html_content)
-    logger.info('✓ HTML report saved to %s', html_path)
-
-    # Save JSON data
-    json_path = output_dir / 'strainmap_benchmark.json'
-    json_path.write_text(json.dumps(json_results, indent=2))
-    logger.info('✓ JSON data saved to %s', json_path)
-
-    logger.info('\n%s', '=' * 60)
-    logger.info('📊 Benchmark Summary')
-    logger.info('%s', '=' * 60)
-    for strategy_name in strategy_names:
-        result = results[strategy_name]
-        logger.info('\n%s:', strategy_name)
-        logger.info(
-            '  Avg Iterations: %.2f ± %.2f', result['avg_iterations'], result['std_iterations']
-        )
-        logger.info('  Avg Target:     %.4f ± %.4f', result['avg_target'], result['std_target'])
-        logger.info('  Success Rate:   %.1f%%', result['success_rate'])
-
-    logger.info('\n%s', '=' * 60)
-    logger.info('✓ Open %s in your browser to view the full report', html_path)
-    logger.info('%s', '=' * 60)
-    # assert iterations <= MAX_ITERATIONS
+"""
+Benchmark which is meant to test core DIAL functionality (without the INTERSECT or MONGO pieces) for the Strain-Mapping benchmark problem.
+"""
+
+import logging
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pytest
+from scipy.interpolate import LinearNDInterpolator
+
+# from pytest_benchmark.fixture import BenchmarkFixture
+from dial_dataclass import (
+    DialInputPredictions,
+    DialInputSingleOtherStrategy,
+    Normal,
+)
+from dial_service import (
+    core as dial_core,
+)
+from dial_service.serverside_data import (
+    ServersideInputBase,
+    ServersideInputPrediction,
+    ServersideInputSingle,
+)
+from dial_service.service_specific_dataclasses import (
+    DialWorkflowCreationParamsService,
+)
+
+from ..helpers import generate_pytest_parameters
+
+logger = logging.getLogger(__name__)
+
+MOCK_WORKFLOW_ID = '6984e6a6ef6e6290dabced91'
+"""fake ObjectID for testing purposes, we do not actually interact with a DB in these tests."""
+
+
+###############
+low_f_data_name = (
+    Path(__file__).parents[1] / 'fixtures' / 'adaptive_strain_manufacturing_low_fidelity.csv'
+)
+high_f_data_name = (
+    Path(__file__).parents[1] / 'fixtures' / 'adaptive_strain_manufacturing_high_fidelity.csv'
+)
+
+
+def read_and_prepare_data():
+    # normalization helper
+    def normalize_data(in_data):
+        m_data = (in_data.max() + in_data.min()) * 0.5
+        d_data = (in_data.max() - in_data.min()) * 0.5
+        return (in_data - m_data) / d_data
+
+    # define constants and read data
+    # wall_st_index = 0
+    # wall_st_end = 676
+    x_col_index = 0  # z follow next
+    e_col_index = 2  # e11, e22 and e33
+    nrows = 26  # Rows in sim strain map of the wall
+    ncols = 26  # Cols in sim strain map of the wall
+    data = pd.read_csv(low_f_data_name)
+
+    Y, Z, E11, E22, E33, R11, R22, R33 = np.array(
+        [
+            data.values[:, x_col_index],
+            data.values[:, x_col_index + 1],
+            data.values[:, e_col_index],
+            data.values[:, e_col_index + 1],
+            data.values[:, e_col_index + 2],
+            data.values[:, e_col_index + 3],
+            data.values[:, e_col_index + 4],
+            data.values[:, e_col_index + 5],
+        ]
+    )
+
+    # Convert data
+    x1 = Y.astype(np.float32).reshape(nrows, ncols)
+    x2 = Z.astype(np.float32).reshape(nrows, ncols)
+
+    # Sim_e11 = np.array(E11).astype(np.float32).reshape(nrows, ncols)
+    # Real_e11 = np.array(R11).astype(np.float32).reshape(nrows, ncols)
+
+    # Sim_e22 = np.array(E22).astype(np.float32).reshape(nrows, ncols)
+    # Real_e22 = np.array(R22).astype(np.float32).reshape(nrows, ncols)
+
+    # Sim_e33 = np.array(E33).astype(np.float32).reshape(nrows, ncols)
+    Real_e33 = np.array(R33).astype(np.float32).reshape(nrows, ncols)
+
+    # Normalize data
+    x1_norm = normalize_data(x1)
+    x2_norm = normalize_data(x2)
+
+    # Sim_e11_norm = normalize_data(Sim_e11)
+    # Real_e11_norm = normalize_data(Real_e11)
+
+    # Sim_e22_norm = normalize_data(Sim_e22)
+    # Real_e22_norm = normalize_data(Real_e22)
+
+    # Sim_e33_norm = normalize_data(Sim_e33)
+    Real_e33_norm = normalize_data(Real_e33)
+
+    # read high fidelity data
+    # data = pd.read_csv(high_f_data_name)
+
+    # Y, Z, E11, E22, E33, R11, R22, R33 = np.array(
+    #    [
+    #        data.values[:, x_col_index],
+    #        data.values[:, x_col_index + 1],
+    #        data.values[:, e_col_index],
+    #        data.values[:, e_col_index + 1],
+    #        data.values[:, e_col_index + 2],
+    #        data.values[:, e_col_index + 3],
+    #        data.values[:, e_col_index + 4],
+    #        data.values[:, e_col_index + 5],
+    #    ]
+    # )
+
+    # Sim_hi_e11 = np.array(E11).astype(np.float32).reshape(nrows, ncols)
+    # Sim_hi_e22 = np.array(E22).astype(np.float32).reshape(nrows, ncols)
+    # Sim_hi_e33 = np.array(E33).astype(np.float32).reshape(nrows, ncols)
+
+    # Transpose the high-fidelity simulation data
+    # TODO, figure out why this is transposed
+    # Sim_hi_e11 = Sim_hi_e11.T
+    # Sim_hi_e22 = Sim_hi_e22.T
+    # Sim_hi_e33 = Sim_hi_e33.T
+
+    # Normalize data
+    # Sim_hi_e11_norm = normalize_data(Sim_hi_e11)
+    # Sim_hi_e22_norm = normalize_data(Sim_hi_e22)
+    # Sim_hi_e33_norm = normalize_data(Sim_hi_e33)
+
+    # return data that will be used for the benchmark
+    return nrows, (x1_norm, x2_norm), Real_e33_norm
+
+
+###############
+
+## Select data for ground truth
+
+ngrid, x_grids, y_grid = read_and_prepare_data()
+
+# default inputs
+INITIAL_BOUNDS = [[-1.0, 1.0], [-1.0, 1.0]]
+NUM_DIMS = len(INITIAL_BOUNDS)
+
+MESHGRID_SIZE = ngrid
+INITIAL_MESHGRIDS = x_grids
+INITIAL_POINTS_TO_PREDICT = np.hstack([mg.reshape(-1, 1) for mg in INITIAL_MESHGRIDS])
+
+INITIAL_PREDICTIONS = y_grid.reshape(-1, 1)
+
+###############
+
+NOISE_LEVEL = 1.0e-2
+
+truth_interp = LinearNDInterpolator(INITIAL_POINTS_TO_PREDICT, INITIAL_PREDICTIONS)
+
+
+@dataclass
+class StrainMap:
+    """Strain Map function."""
+
+    noise_level: float
+
+    truth_interp: LinearNDInterpolator
+    """Strain interpolator that is used as 'ground_truth'"""
+
+    def strain_map(self, x) -> float:
+        """
+        Represents a measured strain as a function of two simulation parameters.
+        """
+        x = np.asarray(x).reshape((-1, 2))
+        x1, x2 = x[:, 0], x[:, 1]
+
+        y_true = truth_interp(np.asarray(x1), np.asarray(x2))
+
+        # print("Evaluated truth model", np.hstack((x, y_true)))
+
+        y_noise = y_true + self.noise_level * np.random.normal(size=y_true.shape)
+        return y_noise
+
+
+# build a strain map from the selected truth values
+truth_strain_map = StrainMap(truth_interp=truth_interp, noise_level=NOISE_LEVEL)
+
+###############
+
+# test parameters
+INITIAL_NUM_POINTS = 25
+MAX_ITERATIONS = 150  # only allow a maximum of this many iterations in tests
+
+INITIAL_DATASET_X = np.random.uniform(-1.0, 1.0, size=(INITIAL_NUM_POINTS, NUM_DIMS)).tolist()
+
+TARGET_ERROR = 0.11
+
+
+def run_simulation(
+    dataset_x: list[list[float]],
+    dataset_y: list[float],
+    strategy: str,
+    strategy_args: object,
+) -> tuple[np.ndarray, np.ndarray]:
+    # important "Hyper-parameters"
+    length_scale = 0.4
+    prior_std = 1.0
+
+    # Assume that there is some small noise in the measurements to stabilize the fit
+    statistics_y = Normal(loc='y', scale=NOISE_LEVEL)
+
+    # modify a local copy of strategy_args
+    strategy_args = strategy_args.copy()
+    backend = strategy_args.pop('backend', 'sklearn')
+
+    if backend == 'sable':
+        # train model with new data
+        kernel = 'rbf'
+        kernel_args = {
+            'sigma_range': [5e-2, 1.0],
+            'gamma': 0.5,
+        }
+        backend_args = {
+            'prior_std': prior_std,
+        }
+
+    else:
+        # train model with new data
+        kernel = 'matern'
+        kernel_args = {
+            'length_scale': length_scale,
+            'constant_value': prior_std**2,
+        }
+        backend_args = {}
+
+    client_state = DialWorkflowCreationParamsService(
+        dataset_x=dataset_x,
+        dataset_y=dataset_y,
+        statistics_y=statistics_y,
+        bounds=INITIAL_BOUNDS,
+        kernel=kernel,
+        kernel_args=kernel_args,
+        y_is_good=False,
+        backend=backend,
+        backend_args=backend_args,
+        seed=-1,
+        dim_x=2,
+    )
+
+    data = ServersideInputBase(client_state)
+    model = dial_core.train_model(data)
+
+    # use get_surrogate_values to predict mean and standard deviation on the inital points grid
+    data = ServersideInputPrediction(
+        client_state,
+        DialInputPredictions(
+            workflow_id=MOCK_WORKFLOW_ID,
+            points_to_predict=INITIAL_POINTS_TO_PREDICT,
+        ),
+    )
+    surrogate_mean, surrogate_std, _ = dial_core.get_surrogate_values(data, model)
+    mean_grid = np.array(surrogate_mean).reshape((-1, 1))
+
+    # subtract the true values and save mean absolute error and standard deviation
+    err_grid = np.abs(mean_grid - INITIAL_PREDICTIONS).reshape((MESHGRID_SIZE,) * NUM_DIMS)
+    std_grid = np.array(surrogate_std).reshape((MESHGRID_SIZE,) * NUM_DIMS)
+
+    # get_next_point
+    data = ServersideInputSingle(
+        client_state,
+        DialInputSingleOtherStrategy(
+            workflow_id=MOCK_WORKFLOW_ID,
+            bounds=INITIAL_BOUNDS,
+            y_is_good=False,
+            seed=-1,
+            strategy=strategy,
+            strategy_args=strategy_args,
+            discrete_measurements=True,
+            discrete_measurement_grid_size=[MESHGRID_SIZE, MESHGRID_SIZE],
+        ),
+    )
+    next_point = dial_core.get_next_point(data, model)
+
+    dataset_x.append(next_point)
+
+    # compute at next point
+    next_point_y = truth_strain_map.strain_map(next_point).reshape(-1).tolist()
+    dataset_y.append(next_point_y[0])
+
+    return err_grid, std_grid
+
+
+def graph(err_grid, dataset_x, strategy):
+    try:
+        import matplotlib as mpl
+        import matplotlib.pyplot as plt
+
+        mpl.use('Agg')  # Use non-interactive backend
+    except ImportError:
+        logger.error(  # noqa: TRY400
+            'Error: matplotlib is required for generating plots. Install it with: pip install matplotlib'
+        )
+        return
+
+    plt.clf()
+    data = np.array(err_grid)
+    plt.contourf(
+        INITIAL_MESHGRIDS[0],
+        INITIAL_MESHGRIDS[1],
+        data,
+        levels=np.logspace(-2, 0, 10),
+        norm='log',
+        extend='both',
+    )
+    cbar = plt.colorbar()
+    cbar.set_ticks(np.logspace(-2, 0, 7))
+    cbar.set_label('Overall error')
+    plt.xlabel('Simulation Parameter #1')
+    plt.ylabel('Simulation Parameter #2')
+    # add black dots for data points and a red marker for the recommendation:
+    X_train = np.array(dataset_x)
+    plt.scatter(X_train[:, 0], X_train[:, 1], color='black', marker='o')
+    plt.scatter(1.0, 1.0, s=300, color='None', edgecolors='black', marker='o')
+
+    plt.savefig(f'graph_{strategy}.png')
+
+
+def accuracy_benchmark(
+    strategy: str, strategy_args: object, max_iterations=MAX_ITERATIONS
+) -> tuple[int, float, float]:
+    """
+    returns:
+      - number of iterations taken to reach an acceptably accurate target
+      - target value achieved
+    """
+
+    iterations = 0
+
+    initial_dataset_x = np.random.uniform(-1.0, 1.0, size=(INITIAL_NUM_POINTS, NUM_DIMS)).tolist()
+
+    dataset_x = initial_dataset_x
+    dataset_y = truth_strain_map.strain_map(dataset_x).reshape(-1).tolist()
+
+    # run simulations until we reach an acceptable target range
+    while iterations < max_iterations:
+        try:
+            err_grid, std_grid = run_simulation(dataset_x, dataset_y, strategy, strategy_args)
+        except Exception as e:
+            logger.exception('Error during simulation')
+            raise AssertionError from e
+
+        # guess has no meaning here, take the last acquired datapoint
+        guess = dataset_x[-1]
+
+        # compute the RMSE and MAD on the grid and the average
+        # rmse = np.sqrt(err_grid**2 + std_grid**2)
+        mad = np.abs(err_grid)
+        # total_rmse = np.sqrt(np.mean(rmse**2))
+        total_mad = np.sqrt(np.mean(mad**2))
+
+        # pick the MAD over the RMSE as an error criterion, it only measures the mean deviation:
+        # it is less sensitive to wrong std and hyperparameter calibration, but does not measure how well
+        # the "error bars" (std_grid) quantify the actual ty
+        error = total_mad
+        print(iterations, error)
+
+        graph(
+            mad,
+            dataset_x,
+            f'{strategy}_{strategy_args.get("backend", "sklearn")}',
+        )
+
+        if error <= TARGET_ERROR:
+            break
+        iterations += 1
+
+    return iterations, error, guess
+
+
+TEST_PARAMS = (
+    ('backend', 'strategy', 'strategy_args', 'max_iterations'),
+    [
+        ('sklearn', 'uncertainty', {}, 2),
+        ('sable', 'uncertainty', {}, 2),
+        ('sklearn', 'random', {}, 2),
+        # use low number of iterations for unit test
+    ],
+)
+
+
+@pytest.mark.parametrize(*generate_pytest_parameters(TEST_PARAMS, 0))
+def test_benchmark_strainmap_accuracy(
+    # benchmark: BenchmarkFixture,
+    backend: str,
+    strategy: str,
+    strategy_args: dict,
+    max_iterations: int,
+) -> None:
+    # NUM_RUNS = 20
+    # for _ in range(NUM_RUNS):
+    # iterations, target = benchmark(
+    # partial(accuracy_benchmark, strategy, strategy_args)
+    # )
+    iterations, target, guess = accuracy_benchmark(
+        strategy, {'backend': backend, **strategy_args}, max_iterations
+    )
+    print(
+        'Iterations for',
+        strategy,
+        ': ',
+        iterations,
+        ' best guess:',
+        guess,
+        ' with target value:',
+        target,
+    )
+    print(
+        'Maximum early terminus value',
+        TARGET_ERROR,
+        ' with ',
+        max_iterations,
+        ' maximum iterations.',
+    )
+    print(
+        'Accuracy benchmark for strategy:',
+        strategy,
+        'reached' if iterations <= max_iterations else 'not reached',
+    )
+
+
+if __name__ == '__main__':
+    """Generate HTML benchmark report with plots comparing different strategies."""
+    import argparse
+    import datetime
+    import json
+    from pathlib import Path
+
+    logger = logging.getLogger(f'{__name__}_runner')
+
+    try:
+        import matplotlib as mpl
+        import matplotlib.pyplot as plt
+
+        mpl.use('Agg')  # Use non-interactive backend
+    except ImportError:
+        logger.error(  # noqa: TRY400
+            'Error: matplotlib is required for generating plots. Install it with: pip install matplotlib'
+        )
+        sys.exit(1)
+
+    def positive_int_type(arg):
+        try:
+            val = int(arg)
+        except ValueError as e:
+            msg = 'Must be an integer'
+            raise argparse.ArgumentTypeError(msg) from e
+        if val < 1:
+            msg = 'Argument must be a positive number'
+            raise argparse.ArgumentTypeError(msg)
+        return val
+
+    parser = argparse.ArgumentParser(description='Generate the Strainmap HTML benchmark pages.')
+    parser.add_argument(
+        '--num-runs',
+        '-n',
+        type=positive_int_type,
+        default=3,
+        help='Number of runs for each strategy.',
+    )
+    args = parser.parse_args()
+
+    strategies = TEST_PARAMS[1]
+
+    # Run multiple iterations for statistical analysis
+    NUM_RUNS = args.num_runs
+    logger.info('Running benchmarks with %d iterations per strategy...', NUM_RUNS)
+
+    results = {}
+    for strategy, strategy_args, _ in strategies:
+        strategy_name = f'{strategy}' + (f' {json.dumps(strategy_args)}' if strategy_args else '')
+        logger.info('\nBenchmarking: %s', strategy_name)
+
+        iterations_list = []
+        targets_list = []
+        guesses_list = []
+
+        for run in range(NUM_RUNS):
+            iterations, target, guess = accuracy_benchmark(strategy, strategy_args)
+            iterations_list.append(iterations)
+            targets_list.append(target)
+            guesses_list.append(guess)
+            logger.info(
+                '  Run %d/%d: iterations=%d, target=%.4f',
+                run + 1,
+                NUM_RUNS,
+                iterations,
+                target,
+            )
+
+        results[strategy_name] = {
+            'strategy': strategy,
+            'strategy_args': strategy_args,
+            'iterations': iterations_list,
+            'targets': targets_list,
+            'guesses': guesses_list,
+            'avg_iterations': np.mean(iterations_list),
+            'std_iterations': np.std(iterations_list),
+            'avg_target': np.mean(targets_list),
+            'std_target': np.std(targets_list),
+            'success_rate': sum(1 for t in targets_list if t <= TARGET_ERROR) / NUM_RUNS * 100,
+        }
+
+    # Generate plots
+    output_dir = Path('reports/benchmarks')
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
+    fig.suptitle(
+        'Strainmap Optimization Benchmark Comparison',
+        fontsize=16,
+        fontweight='bold',
+    )
+
+    # Plot 1: Average iterations to convergence
+    ax1 = axes[0, 0]
+    strategy_names = list(results.keys())
+    avg_iterations = [results[s]['avg_iterations'] for s in strategy_names]
+    std_iterations = [results[s]['std_iterations'] for s in strategy_names]
+    bars1 = ax1.bar(
+        range(len(strategy_names)),
+        avg_iterations,
+        yerr=std_iterations,
+        capsize=5,
+        alpha=0.7,
+    )
+    ax1.set_xlabel('Strategy')
+    ax1.set_ylabel('Average Iterations')
+    ax1.set_title('Iterations to Convergence (Lower is Better)')
+    ax1.set_xticks(range(len(strategy_names)))
+    ax1.set_xticklabels(
+        [s.replace(' ', '\n') for s in strategy_names],
+        rotation=0,
+        ha='center',
+        fontsize=8,
+    )
+    ax1.grid(axis='y', alpha=0.3)
+
+    # Add value labels on bars
+    for _i, (bar, val, std) in enumerate(zip(bars1, avg_iterations, std_iterations, strict=False)):
+        height = bar.get_height()
+        ax1.text(
+            bar.get_x() + bar.get_width() / 2.0,
+            height,
+            f'{val:.1f}±{std:.1f}',
+            ha='center',
+            va='bottom',
+            fontsize=9,
+        )
+
+    # Plot 2: Average target value achieved
+    ax2 = axes[0, 1]
+    avg_targets = [results[s]['avg_target'] for s in strategy_names]
+    std_targets = [results[s]['std_target'] for s in strategy_names]
+    bars2 = ax2.bar(
+        range(len(strategy_names)),
+        avg_targets,
+        yerr=std_targets,
+        capsize=5,
+        alpha=0.7,
+        color='orange',
+    )
+    ax2.set_xlabel('Strategy')
+    ax2.set_ylabel('Average Target Value')
+    ax2.set_title('Final Target Value (Lower is Better)')
+    ax2.set_xticks(range(len(strategy_names)))
+    ax2.set_xticklabels(
+        [s.replace(' ', '\n') for s in strategy_names],
+        rotation=0,
+        ha='center',
+        fontsize=8,
+    )
+    ax2.axhline(
+        y=TARGET_ERROR,
+        color='r',
+        linestyle='--',
+        label=f'Target Threshold ({TARGET_ERROR})',
+    )
+    ax2.legend()
+    ax2.grid(axis='y', alpha=0.3)
+
+    # Add value labels on bars
+    for _i, (bar, val, std) in enumerate(zip(bars2, avg_targets, std_targets, strict=False)):
+        height = bar.get_height()
+        ax2.text(
+            bar.get_x() + bar.get_width() / 2.0,
+            height,
+            f'{val:.2f}±{std:.2f}',
+            ha='center',
+            va='bottom',
+            fontsize=9,
+        )
+
+    # Plot 3: Success rate
+    ax3 = axes[1, 0]
+    success_rates = [results[s]['success_rate'] for s in strategy_names]
+    bars3 = ax3.bar(range(len(strategy_names)), success_rates, alpha=0.7, color='green')
+    ax3.set_xlabel('Strategy')
+    ax3.set_ylabel('Success Rate (%)')
+    ax3.set_title(f'Success Rate (Target ≤ {TARGET_ERROR})')
+    ax3.set_xticks(range(len(strategy_names)))
+    ax3.set_xticklabels(
+        [s.replace(' ', '\n') for s in strategy_names],
+        rotation=0,
+        ha='center',
+        fontsize=8,
+    )
+    ax3.set_ylim([0, 105])
+    ax3.grid(axis='y', alpha=0.3)
+
+    # Add value labels on bars
+    for bar, val in zip(bars3, success_rates, strict=False):
+        height = bar.get_height()
+        ax3.text(
+            bar.get_x() + bar.get_width() / 2.0,
+            height,
+            f'{val:.1f}%',
+            ha='center',
+            va='bottom',
+            fontsize=9,
+        )
+
+    # Plot 4: Box plot of iterations distribution
+    ax4 = axes[1, 1]
+    iterations_data = [results[s]['iterations'] for s in strategy_names]
+    bp = ax4.boxplot(iterations_data, patch_artist=True)
+    for patch in bp['boxes']:
+        patch.set_facecolor('lightblue')
+    ax4.set_xlabel('Strategy')
+    ax4.set_ylabel('Iterations')
+    ax4.set_title('Iterations Distribution')
+    ax4.set_xticklabels([s.replace(' ', '\n') for s in strategy_names], fontsize=8)
+    ax4.grid(axis='y', alpha=0.3)
+
+    plt.tight_layout()
+    plot_path = output_dir / 'strainmap_benchmark.png'
+    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
+    logger.info('✓ Plot saved to %s', plot_path)
+    plt.close()
+
+    # Generate HTML report
+    html_content = f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Strainmap Optimization Benchmark Report</title>
+    <style>
+        body {{
+            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+            max-width: 1200px;
+            margin: 0 auto;
+            padding: 20px;
+            background-color: #f5f5f5;
+        }}
+        h1 {{
+            color: #2c3e50;
+            border-bottom: 3px solid #3498db;
+            padding-bottom: 10px;
+        }}
+        h2 {{
+            color: #34495e;
+            margin-top: 30px;
+        }}
+        .metadata {{
+            background-color: #ecf0f1;
+            padding: 15px;
+            border-radius: 5px;
+            margin-bottom: 20px;
+        }}
+        .metadata p {{
+            margin: 5px 0;
+        }}
+        table {{
+            width: 100%;
+            border-collapse: collapse;
+            margin: 20px 0;
+            background-color: white;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        }}
+        th, td {{
+            padding: 12px;
+            text-align: left;
+            border-bottom: 1px solid #ddd;
+        }}
+        th {{
+            background-color: #3498db;
+            color: white;
+            font-weight: bold;
+        }}
+        tr:hover {{
+            background-color: #f1f1f1;
+        }}
+        .best {{
+            background-color: #d4edda;
+            font-weight: bold;
+        }}
+        .plot {{
+            text-align: center;
+            margin: 30px 0;
+            background-color: white;
+            padding: 20px;
+            border-radius: 5px;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        }}
+        .plot img {{
+            max-width: 100%;
+            height: auto;
+        }}
+        .summary {{
+            background-color: #d1ecf1;
+            border-left: 4px solid #0c5460;
+            padding: 15px;
+            margin: 20px 0;
+        }}
+        code {{
+            background-color: #f4f4f4;
+            padding: 2px 6px;
+            border-radius: 3px;
+            font-family: 'Courier New', monospace;
+        }}
+    </style>
+</head>
+<body>
+    <h1>📊 Strainmap Optimization Benchmark Report</h1>
+
+    <div class="metadata">
+        <p><strong>Generated:</strong> {datetime.datetime.now(tz=datetime.timezone.utc).astimezone().strftime('%Y-%m-%d %H:%M:%S')}</p>
+        <p><strong>Number of Runs per Strategy:</strong> {NUM_RUNS}</p>
+        <p><strong>Maximum Iterations:</strong> {MAX_ITERATIONS}</p>
+        <p><strong>Target Threshold:</strong> {TARGET_ERROR}</p>
+        <p><strong>Initial Bounds:</strong> {INITIAL_BOUNDS}</p>
+        <p><strong>Initial Points:</strong> {INITIAL_NUM_POINTS}</p>
+    </div>
+
+    <div class="summary">
+        <h3>🎯 Test Objective</h3>
+        <p>This benchmark evaluates different acquisition strategies for Bayesian optimization on a strain mapping experiment dataset.
+        The goal is to minimize the total surrogate error (over all inputs) within {MAX_ITERATIONS} iterations, starting from {INITIAL_NUM_POINTS} initial points.</p>
+    </div>
+
+    <h2>📈 Benchmark Results</h2>
+
+    <div class="plot">
+        <h3>Performance Comparison</h3>
+        <img src="strainmap_benchmark.png" alt="Benchmark comparison plots">
+    </div>
+
+    <h2>📋 Detailed Statistics</h2>
+
+    <table>
+        <thead>
+            <tr>
+                <th>Strategy</th>
+                <th>Strategy Args</th>
+                <th>Avg Iterations</th>
+                <th>Std Iterations</th>
+                <th>Avg Target Value</th>
+                <th>Std Target Value</th>
+                <th>Success Rate</th>
+            </tr>
+        </thead>
+        <tbody>
+"""
+
+    # Find best performers
+    best_iterations_idx = min(range(len(strategy_names)), key=lambda i: avg_iterations[i])
+    best_target_idx = min(range(len(strategy_names)), key=lambda i: avg_targets[i])
+    best_success_idx = max(range(len(strategy_names)), key=lambda i: success_rates[i])
+
+    for i, strategy_name in enumerate(strategy_names):
+        result = results[strategy_name]
+        row_class = ''
+        if i in (best_iterations_idx, best_target_idx, best_success_idx):
+            row_class = ' class="best"'
+
+        args_str = json.dumps(result['strategy_args']) if result['strategy_args'] else 'None'
+
+        html_content += f"""            <tr{row_class}>
+                <td><code>{result['strategy']}</code></td>
+                <td><code>{args_str}</code></td>
+                <td>{result['avg_iterations']:.2f}</td>
+                <td>{result['std_iterations']:.2f}</td>
+                <td>{result['avg_target']:.4f}</td>
+                <td>{result['std_target']:.4f}</td>
+                <td>{result['success_rate']:.1f}%</td>
+            </tr>
+"""
+
+    html_content += """        </tbody>
+    </table>
+
+    <h2>🏆 Key Findings</h2>
+    <div class="summary">
+"""
+
+    # Add findings
+    best_strategy = strategy_names[best_iterations_idx]
+    html_content += f"""        <p><strong>Fastest Convergence:</strong> <code>{best_strategy}</code> with {avg_iterations[best_iterations_idx]:.2f} ± {std_iterations[best_iterations_idx]:.2f} iterations on average</p>
+"""
+
+    best_accuracy_strategy = strategy_names[best_target_idx]
+    html_content += f"""        <p><strong>Best Accuracy:</strong> <code>{best_accuracy_strategy}</code> with average target value {avg_targets[best_target_idx]:.4f} ± {std_targets[best_target_idx]:.4f}</p>
+"""
+
+    best_reliability_strategy = strategy_names[best_success_idx]
+    html_content += f"""        <p><strong>Most Reliable:</strong> <code>{best_reliability_strategy}</code> with {success_rates[best_success_idx]:.1f}% success rate</p>
+"""
+
+    html_content += """    </div>
+
+    <h2>📊 Raw Data</h2>
+    <details>
+        <summary>Click to expand raw results JSON</summary>
+        <pre style="background-color: #f4f4f4; padding: 15px; border-radius: 5px; overflow-x: auto;">
+"""
+
+    # Prepare JSON-serializable results
+    json_results = {}
+    for strategy_name, result in results.items():
+        json_results[strategy_name] = {
+            'strategy': result['strategy'],
+            'strategy_args': result['strategy_args'],
+            'iterations': result['iterations'],
+            'targets': result['targets'],
+            'guesses': [[float(x) for x in guess] for guess in result['guesses']],
+            'statistics': {
+                'avg_iterations': float(result['avg_iterations']),
+                'std_iterations': float(result['std_iterations']),
+                'avg_target': float(result['avg_target']),
+                'std_target': float(result['std_target']),
+                'success_rate': float(result['success_rate']),
+            },
+        }
+
+    html_content += json.dumps(json_results, indent=2)
+    html_content += """
+        </pre>
+    </details>
+
+    <footer style="margin-top: 50px; padding-top: 20px; border-top: 1px solid #ddd; color: #7f8c8d; text-align: center;">
+        <p>Generated by test_strainmap.py benchmark suite</p>
+    </footer>
+</body>
+</html>
+"""
+
+    # Save HTML report
+    html_path = output_dir / 'strainmap_benchmark.html'
+    html_path.write_text(html_content)
+    logger.info('✓ HTML report saved to %s', html_path)
+
+    # Save JSON data
+    json_path = output_dir / 'strainmap_benchmark.json'
+    json_path.write_text(json.dumps(json_results, indent=2))
+    logger.info('✓ JSON data saved to %s', json_path)
+
+    logger.info('\n%s', '=' * 60)
+    logger.info('📊 Benchmark Summary')
+    logger.info('%s', '=' * 60)
+    for strategy_name in strategy_names:
+        result = results[strategy_name]
+        logger.info('\n%s:', strategy_name)
+        logger.info(
+            '  Avg Iterations: %.2f ± %.2f',
+            result['avg_iterations'],
+            result['std_iterations'],
+        )
+        logger.info(
+            '  Avg Target:     %.4f ± %.4f',
+            result['avg_target'],
+            result['std_target'],
+        )
+        logger.info('  Success Rate:   %.1f%%', result['success_rate'])
+
+    logger.info('\n%s', '=' * 60)
+    logger.info('✓ Open %s in your browser to view the full report', html_path)
+    logger.info('%s', '=' * 60)
+    # assert iterations <= MAX_ITERATIONS
diff --git a/tests/helpers.py b/tests/helpers.py
new file mode 100644
index 0000000..930ea72
--- /dev/null
+++ b/tests/helpers.py
@@ -0,0 +1,30 @@
+from copy import deepcopy
+
+import pytest
+
+from dial_service.service_specific_dataclasses import AVAILABLE_DIAL_BACKENDS
+
+
+def generate_pytest_parameters(
+    old_params: tuple[tuple[str], tuple[list[object]]], backend_idx: int
+) -> tuple[tuple[str], tuple[list[object]]]:
+    """
+    Generates test safeguards from generic parameters.
+
+    This should only really be used if you simultaneously want to use generic parameters in contexts other than Pytest.
+
+    Params:
+      - old_params = your normal pytest.mark.parametrize parameters (don't use pytest.param)
+      - backend_idx = the index of the "backend" parameter in your args list"""
+    new_params = deepcopy(old_params)
+    for i, test in enumerate(old_params[1]):
+        backend_name = test[backend_idx]
+        if backend_name != 'sklearn':
+            new_params[1][i] = pytest.param(
+                *test,
+                marks=pytest.mark.skipif(
+                    backend_name not in AVAILABLE_DIAL_BACKENDS,
+                    reason=f'{backend_name} not installed',
+                ),
+            )
+    return new_params
diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/unit/test_internals.py b/tests/unit/test_internals.py
index 1b6f9c1..7d32e16 100644
--- a/tests/unit/test_internals.py
+++ b/tests/unit/test_internals.py
@@ -18,6 +18,7 @@
     ServersideInputSingle,
 )
 from dial_service.service_specific_dataclasses import (
+    AVAILABLE_DIAL_BACKENDS,
     DialWorkflowCreationParamsService,
 )
 
@@ -332,7 +333,13 @@ def prediction_1D_heteroscedastic(backend):
     ('backend', 'approx'),
     [
         ('sklearn', 1.842309),
-        # ('gpax', 2.0),
+        # pytest.param(
+        #     'gpax', 2.0,
+        #     marks=pytest.mark.skipif(
+        #         'gpax' not in AVAILABLE_DIAL_BACKENDS,
+        #         reason='gpax not installed',
+        #     ),
+        # ),
     ],
 )
 def test_EI_1D(backend, approx):
@@ -349,7 +356,13 @@ def test_EI_1D(backend, approx):
     ('backend', 'val'),
     [
         ('sklearn', 1.0),
-        # ('gpax', 2.0),
+        # pytest.param(
+        #     'gpax', 2.0,
+        #     marks=pytest.mark.skipif(
+        #         'gpax' not in AVAILABLE_DIAL_BACKENDS,
+        #         reason='gpax not installed',
+        #     ),
+        # ),
     ],
 )
 def test_EI_1D_discrete(backend, val):
@@ -370,7 +383,13 @@ def test_EI_1D_discrete(backend, val):
     ('backend', 'approx'),
     [
         ('sklearn', [1.705352, -1.682829]),
-        # ('gpax', [2.0, 2.0]),
+        # pytest.param(
+        #     'gpax', [2.0, 2,0],
+        #     marks=pytest.mark.skipif(
+        #         'gpax' not in AVAILABLE_DIAL_BACKENDS,
+        #         reason='gpax not installed',
+        #     ),
+        # ),
     ],
 )
 def test_EI_2D(backend, approx):
@@ -387,7 +406,12 @@ def test_EI_2D(backend, approx):
     ('backend', 'approx'),
     [
         ('sklearn', [2.000000, -1.143727, -1.859496]),
-        # ('gpax', [2.0, 2.0, -2.0],  # WAS: [2.0, 2.0, 2.0]
+        # pytest.param(
+        #     'gpax', [2.0, 2,0, 2.0,],  # WAS: [2.0,2.0,2,0]
+        #     marks=pytest.mark.skipif(
+        #         'gpax' not in AVAILABLE_DIAL_BACKENDS,
+        #         reason='gpax not installed',
+        #     ),
         # ),
     ],
 )
@@ -405,7 +429,13 @@ def test_EI_3D(backend, approx):
     ('backend', 'approx'),
     [
         ('sklearn', [1.5]),
-        # ('gpax',),
+        # pytest.param(
+        #     'gpax' [2.0],
+        #     marks=pytest.mark.skipif(
+        #         'gpax' not in AVAILABLE_DIAL_BACKENDS,
+        #         reason='gpax not installed',
+        #     ),
+        # ),
     ],
 )
 def test_uncertainty(backend, approx):
@@ -419,7 +449,13 @@ def test_uncertainty(backend, approx):
     ('backend', 'approx'),
     [
         ('sklearn', [1.790396262]),
-        # ('gpax', [2.0]),
+        # pytest.param(
+        #     'gpax' [2.0],
+        #     marks=pytest.mark.skipif(
+        #         'gpax' not in AVAILABLE_DIAL_BACKENDS,
+        #         reason='gpax not installed',
+        #     ),
+        # ),
     ],
 )
 def test_preprocessing_standardize(backend, approx):
@@ -435,7 +471,13 @@ def test_preprocessing_standardize(backend, approx):
     ('backend', 'val'),
     [
         ('sklearn', [1.0]),
-        # ('gpax', [2.0]),
+        # pytest.param(
+        #     'gpax' [2.0],
+        #     marks=pytest.mark.skipif(
+        #         'gpax' not in AVAILABLE_DIAL_BACKENDS,
+        #         reason='gpax not installed',
+        #     ),
+        # ),
     ],
 )
 def test_preprocessing_standardize_discrete(backend, val):
@@ -459,7 +501,13 @@ def test_preprocessing_standardize_discrete(backend, val):
     ('backend'),
     [
         ('sklearn'),
-        ('gpax'),
+        pytest.param(
+            'gpax',
+            marks=pytest.mark.skipif(
+                'gpax' not in AVAILABLE_DIAL_BACKENDS,
+                reason='gpax not installed',
+            ),
+        ),
     ],
 )
 def test_random(backend):
@@ -475,7 +523,13 @@ def test_random(backend):
     ('backend'),
     [
         ('sklearn'),
-        ('gpax'),
+        pytest.param(
+            'gpax',
+            marks=pytest.mark.skipif(
+                'gpax' not in AVAILABLE_DIAL_BACKENDS,
+                reason='gpax not installed',
+            ),
+        ),
     ],
 )
 def test_hypercube_single_point(backend):
@@ -505,7 +559,13 @@ def test_hypercube_single_point(backend):
     ('backend'),
     [
         ('sklearn'),
-        ('gpax'),
+        pytest.param(
+            'gpax',
+            marks=pytest.mark.skipif(
+                'gpax' not in AVAILABLE_DIAL_BACKENDS,
+                reason='gpax not installed',
+            ),
+        ),
     ],
 )
 def test_random_discrete(backend):
@@ -530,7 +590,13 @@ def test_random_discrete(backend):
     ('backend'),
     [
         ('sklearn'),
-        ('gpax'),
+        pytest.param(
+            'gpax',
+            marks=pytest.mark.skipif(
+                'gpax' not in AVAILABLE_DIAL_BACKENDS,
+                reason='gpax not installed',
+            ),
+        ),
     ],
 )
 def test_hypercube_single_point_discrete(backend):
@@ -566,7 +632,13 @@ def test_hypercube_single_point_discrete(backend):
     ('backend'),
     [
         ('sklearn'),
-        ('gpax'),
+        pytest.param(
+            'gpax',
+            marks=pytest.mark.skipif(
+                'gpax' not in AVAILABLE_DIAL_BACKENDS,
+                reason='gpax not installed',
+            ),
+        ),
     ],
 )
 def test_hypercube_single_point_discrete_2D(backend):
@@ -601,7 +673,13 @@ def test_hypercube_single_point_discrete_2D(backend):
     ('backend'),
     [
         ('sklearn'),
-        ('gpax'),
+        pytest.param(
+            'gpax',
+            marks=pytest.mark.skipif(
+                'gpax' not in AVAILABLE_DIAL_BACKENDS,
+                reason='gpax not installed',
+            ),
+        ),
     ],
 )
 def test_random_points(backend):
@@ -616,7 +694,13 @@ def test_random_points(backend):
     ('backend'),
     [
         ('sklearn'),
-        # ('gpax'),
+        # pytest.param(
+        #     'gpax',
+        #     marks=pytest.mark.skipif(
+        #         'gpax' not in AVAILABLE_DIAL_BACKENDS,
+        #         reason='gpax not installed',
+        #     ),
+        # ),
     ],
 )
 def test_hypercube_multiple_points(backend):
@@ -644,7 +728,7 @@ def test_hypercube_multiple_points(backend):
             ],
             [2.11126987e01, 2.96625069e01, 2.11126987e01],
         ),
-        # (
+        # pytest.param(
         # 'gpax',
         # [
         #     76.99987768175089,
@@ -654,6 +738,7 @@ def test_hypercube_multiple_points(backend):
         #     82.26569221517353,
         # ],
         # [3335.7290084812175, 3327.202331393974, 3335.7290084812175],
+        # marks=pytest.mark.skipif('gpax' not in AVAILABLE_DIAL_BACKENDS, reason='gpax not installed')
         # ),
     ],
 )
@@ -673,10 +758,14 @@ def test_surrogate(backend, expected_means, expected_stddevs):
             [-1.0, -0.65, 0.0, 0.65, 1.0],
             [1e-2, 0.42, 0.59, 0.42, 1e-5],
         ),
-        (
+        pytest.param(
             'sable',
             [-1.0, -0.56, 0.0, 0.56, 1.0],
             [1e-2, 0.29, 0.34, 0.29, 1e-5],
+            marks=pytest.mark.skipif(
+                'sable' not in AVAILABLE_DIAL_BACKENDS,
+                reason='sable not installed',
+            ),
         ),
     ],
 )

From bc8ec52d6c51923ee3a39881232f7ab88e419f7f Mon Sep 17 00:00:00 2001
From: Lance-Drane <Lance-Drane@users.noreply.github.com>
Date: Fri, 24 Jul 2026 11:10:49 -0400
Subject: [PATCH 7/8] add note about installing dependencies through pip in
 README

Signed-off-by: Lance-Drane <Lance-Drane@users.noreply.github.com>
---
 README.md | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index f17f8d4..a568e92 100644
--- a/README.md
+++ b/README.md
@@ -16,9 +16,12 @@ To install Dial from source, git clone this repository and run the following fro
 
 `pip install .`
 
-Alternatively, both intersect-sdk and Dial may be installed with the following:
+By default, **only scikit-learn is available as a backend**. We also support GPax and Sable as backends. To install these:
 
-`pip install -e .`
+- for GPax: `pip install ".[gpax]"`
+- for Sable: `pip install ".[sable]"`
+
+Use commas to add multiple backends: `pip install ".[gpax,sable]"`
 
 ## Installing (developers)
 

From 778342270117d38ba925074f4eddcd8586fc4c3c Mon Sep 17 00:00:00 2001
From: Konstantin Pieper <pieperk@ornl.gov>
Date: Fri, 24 Jul 2026 16:35:56 -0400
Subject: [PATCH 8/8] fix batch sampler and add a test for it

---
 src/dial_service/serverside_data.py      |  4 +--
 src/dial_service/utilities/strategies.py | 22 +++++++++++--
 tests/unit/test_internals.py             | 39 ++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/src/dial_service/serverside_data.py b/src/dial_service/serverside_data.py
index 0602da2..19ca396 100644
--- a/src/dial_service/serverside_data.py
+++ b/src/dial_service/serverside_data.py
@@ -18,8 +18,8 @@ def __init__(self, data: DialWorkflowCreationParamsService):
         self.dim_y = data.dim_y
         self.labels_x = data.labels_x
         self.labels_y = data.labels_y
-        self.dataset_x = np.array(data.dataset_x)
-        self.dataset_y = np.array(data.dataset_y).reshape((-1, self.dim_y))
+        self.dataset_x = np.array(data.dataset_x, float).reshape((-1, self.dim_x))
+        self.dataset_y = np.array(data.dataset_y, float).reshape((-1, self.dim_y))
         self.statistics_y = data.statistics_y
         # it seems like there should be a smarter way to do this, but stuff involving loops doesn't work with static autocompleters:
         self.bounds = data.bounds
diff --git a/src/dial_service/utilities/strategies.py b/src/dial_service/utilities/strategies.py
index 51d155a..592d796 100644
--- a/src/dial_service/utilities/strategies.py
+++ b/src/dial_service/utilities/strategies.py
@@ -190,13 +190,29 @@ def batch_sampling_acl(backend_module: AbstractBackend, model, data: ServersideI
       - ΔT(t) = max(current_batch_t_max, t) - current_batch_t_max (parallel reactor cost)
     """
 
+    if data.dim_x > 1:
+        msg = f'strategy batch_sampling_acl supports only one input dimension, but {data.dim_x=}'
+        raise ValueError(msg)
+
     x_grid = create_measurement_grid(data)
     x_grid = np.array(x_grid)
 
     data.set_x_predict(x_grid)
-    mean, sd_dev = backend_module.predict(model, data)
-    x_train = data.X_raw
-    _params = data.strategy_args
+    _, sd_dev = backend_module.predict(model, data)
+    x_train = data.dataset_x  # get raw x data without scaling
+
+    # set some default parameters
+    _params = {
+        'lambda_time': 0.0,
+        'lambda_near_train': 1.0,
+        'lambda_near_batch': 1.0,
+        'lambda_batchT': 0.0,
+        'radius_train_factor': 0.1,
+        'radius_batch_factor': 0.1,
+        'eps': 1.0e-3,
+    }
+    # update with provided values
+    _params.update(data.strategy_args or {})
 
     batch_size = data.points
     lambda_time = _params['lambda_time']  # penalty on large t
diff --git a/tests/unit/test_internals.py b/tests/unit/test_internals.py
index 7d32e16..c2f2329 100644
--- a/tests/unit/test_internals.py
+++ b/tests/unit/test_internals.py
@@ -255,6 +255,29 @@ def single_3D(backend, strategy, strategy_args, discrete_measurement_grid_size=N
     return ServersideInputSingle(workflow_state, params)
 
 
+def multiple_1D(backend, strategy, discrete_measurement_grid_size=None):
+    workflow_state = DialWorkflowCreationParamsService(
+        dataset_x=[],
+        dataset_y=[],
+        dim_x=1,  # provide dim_x for empty dataset
+        dim_y=1,  # provide dim_y for empty dataset
+        y_is_good=False,
+        kernel='rbf',
+        bounds=[[0, 100]],
+        backend=backend,
+        seed=42,
+    )
+    params = DialInputMultiple(
+        workflow_id=DUMMY_WORKFLOW_ID,
+        strategy=strategy,
+        bounds=[[0, 100]],
+        discrete_measurements=bool(discrete_measurement_grid_size),
+        discrete_measurement_grid_size=discrete_measurement_grid_size or [20, 20],
+        points=20,
+    )
+    return ServersideInputMultiple(workflow_state, params)
+
+
 def multiple_2D(backend, strategy, discrete_measurement_grid_size=None):
     workflow_state = DialWorkflowCreationParamsService(
         dataset_x=[],
@@ -669,6 +692,22 @@ def test_hypercube_single_point_discrete_2D(backend):
         data.point_index += 1
 
 
+@pytest.mark.parametrize(
+    ('backend', 'strategy'),
+    [
+        ('sklearn', 'polymer_acl_sampler'),
+    ],
+)
+def test_batch_points(backend, strategy):
+    data = multiple_1D(backend, strategy=strategy)
+    model = core.initialize_model(data)
+    batch_points = core.get_next_points(data, model)
+    assert len(batch_points) == data.points
+    for pt in batch_points:
+        print(pt)
+        assert 0 <= pt[0] <= 100
+
+
 @pytest.mark.parametrize(
     ('backend'),
     [