allow dataset_y updates with multiple values

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)
 

pyproject.toml

@@ -20,12 +20,9 @@ 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
+  "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 +33,11 @@ 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",
+  "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]
 dev = [

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()
 

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

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

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:

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

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

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 '')