diff --git a/scripts/2d_rosenbrock_client_with_strategies.py b/scripts/2d_rosenbrock_client_with_strategies.py new file mode 100644 index 0000000..f0149ed --- /dev/null +++ b/scripts/2d_rosenbrock_client_with_strategies.py @@ -0,0 +1,624 @@ +import argparse +import bisect +import json +import logging +import math +import os +import sys +from pathlib import Path +from typing import Any + +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np +from intersect_sdk import ( + INTERSECT_RESPONSE_VALUE, + HierarchyConfig, + IntersectClient, + IntersectClientCallback, + IntersectClientConfig, + IntersectDirectMessageParams, + default_intersect_lifecycle_loop, +) + +# from scipy.stats import qmc +from dial_dataclass import ( + DialInputMultipleOtherStrategy, + DialInputPredictions, + DialInputSingleOtherStrategy, + DialWorkflowCreationParamsClient, + DialWorkflowDatasetUpdate, + DialWorkflowDatasetUpdates, +) + +mpl.use('agg') +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +def rosenbrock(x, y) -> np.ndarray: + x = np.asarray(x) + y = np.asarray(y) + return (1 - x) ** 2 + 100 * (y - x**2) ** 2 + + +# default inputs +BOUNDS = [[-2.0, 2.0], [-2.0, 2.0]] +NUM_DIMS = len(BOUNDS) + +MESHGRID_SIZE = 201 +SURROGATE_MESHGRID = np.meshgrid( + *[np.linspace(dim_bounds[0], dim_bounds[1], MESHGRID_SIZE) for dim_bounds in BOUNDS], + indexing='ij', +) +POINTS_TO_PREDICT = np.hstack([mg.reshape(-1, 1) for mg in SURROGATE_MESHGRID]).tolist() + +GROUND_TRUTH_X = np.hstack([mg.reshape(-1, 1) for mg in SURROGATE_MESHGRID]) +GROUND_TRUTH_Y = rosenbrock(SURROGATE_MESHGRID[0], SURROGATE_MESHGRID[1]).tolist() + + +NUM_ITERATIONS = 200 + + +SAMPLING_GRID_SIZE = [5, 5] +STRATEGY_SCHEDULE = [ + {'strategy': 'center', 'num_samples': 1}, + {'strategy': 'corners', 'num_samples': 2**NUM_DIMS}, + # # {'strategy': 'grid', 'num_samples': np.prod(SAMPLING_GRID_SIZE), 'args': {'grid_size': SAMPLING_GRID_SIZE}}, + # {'strategy': 'chebyshev', 'num_samples': np.prod(SAMPLING_GRID_SIZE), 'args': {'grid_size': SAMPLING_GRID_SIZE}}, + # { + # 'strategy': 'latin_hypercube', + # 'num_samples': np.prod(SAMPLING_GRID_SIZE), + # 'args': {'grid_size': SAMPLING_GRID_SIZE}, + # }, + # { + # 'batch_strategy': 'liar', + # 'strategy': 'corners', + # 'num_samples': 4, + # 'args': {'liar_value': 0}, + # }, + { + 'batch_strategy': 'liar', + 'strategy': 'grid', + 'num_samples': 25, + 'args': {'liar_value': 'median', 'grid_size': [4, 4]}, + }, + { + 'batch_strategy': 'liar', + 'strategy': 'chebyshev', + 'num_samples': 25, + 'args': {'liar_value': 100, 'grid_size': [5, 5]}, + }, + { + 'batch_strategy': 'believer', + 'strategy': 'latin_hypercube', + 'num_samples': 25, + 'args': {'grid_size': [5, 5]}, + }, + { + 'batch_strategy': 'liar', + 'strategy': 'expected_improvement', + 'num_samples': 10, + 'args': {'liar_value': 0}, + }, + { + 'batch_strategy': 'liar', + 'strategy': 'expected_improvement', + 'num_samples': 10, + 'args': {'liar_value': 'mean'}, + }, + { + 'batch_strategy': 'believer', + 'strategy': 'expected_improvement', + 'num_samples': 10, + 'args': {'liar_value': 'mean'}, + }, + {'strategy': 'expected_improvement', 'num_samples': 10}, + { + 'batch_strategy': 'believer', + 'strategy': 'expected_improvement', + 'num_samples': 30, + 'args': {'liar_value': 'mean'}, + }, + { + 'batch_strategy': 'liar', + 'strategy': 'expected_improvement', + 'num_samples': 10, + 'args': {'liar_value': 'min'}, + }, + {'strategy': 'expected_improvement'}, +] + +# KERNEL HYPERPARAMETERS +LENGTH_SCALE = 0.2 +NOISE_LEVEL = 10e-6 +CONSTANT_VALUE = 1.0 + + +class Scheduler: + def __init__(self, strategy_schedule: list[dict[str, Any]]): + self.strategy_schedule = strategy_schedule + self.strategy_break_points = np.cumsum( + [s.get('num_samples', 1e6) for s in strategy_schedule], dtype=int + ) + + def __call__(self, sample_index): + index = bisect.bisect_right(self.strategy_break_points, sample_index) + return self.strategy_schedule[min(index, len(self.strategy_schedule) - 1)] + + def get_strategy_index(self, sample_index): + index = bisect.bisect_right(self.strategy_break_points, sample_index) + return min(index, len(self.strategy_schedule) - 1) + + +class Plotter: + """Class to handle plotting of the surrogate model and optimization progress.""" + + def __init__(self, scheduler: Scheduler, max_cols: int = 4): + """Initialize with a grid of subplots based on the number of strategies""" + self.save_path = 'graph.png' + self.scheduler = scheduler + + num_plots = len(scheduler.strategy_schedule) + 1 + num_cols = min(max_cols, num_plots) + num_rows = math.ceil(num_plots / num_cols) + + self._graph_fig, graph_axes = plt.subplots( + num_rows, num_cols, figsize=(6 * num_cols, 6 * num_rows), squeeze=False + ) + self._graph_axes = graph_axes.ravel() + + # Hide empty subplot positions + for unused_ax in self._graph_axes[num_plots:]: + unused_ax.set_visible(False) + + self._graph_colorbars = {} + + break_points = [0, *list(scheduler.strategy_break_points)] + for i, ax in enumerate(self._graph_axes[1:num_plots]): + current_strategy = scheduler(break_points[i]) + batch_strategy = current_strategy.get('batch_strategy', None) + if batch_strategy and batch_strategy.lower() == 'liar': + batch_strategy = f'{batch_strategy} ({current_strategy["args"]["liar_value"]})' + strategy_name = (f'{batch_strategy} + ' if batch_strategy else '') + current_strategy[ + 'strategy' + ] + strategy_name += ( + f', {current_strategy["num_samples"]} samples' + if 'num_samples' in current_strategy + else '' + ) + + # ax.set_xlabel('Simulation Parameter #1') + # ax.set_ylabel('Simulation Parameter #2') + ax.set_title(f'{strategy_name}') + + colorbar = self._graph_fig.colorbar(None, ax=ax) + colorbar.set_ticks(np.logspace(-2, 4, 7)) + # colorbar.set_label('Simulation Result') + self._graph_colorbars[ax] = colorbar + + self.plot_ground_truth() + + def plot_ground_truth(self): + """Plot the ground truth Rosenbrock function on the first subplot""" + contourf = self._graph_axes[0].contourf( + SURROGATE_MESHGRID[0], + SURROGATE_MESHGRID[1], + GROUND_TRUTH_Y, + levels=np.logspace(-2, 4, 101), + norm='log', + extend='both', + ) + self._graph_axes[0].contour( + SURROGATE_MESHGRID[0], + SURROGATE_MESHGRID[1], + GROUND_TRUTH_Y, + levels=np.logspace(-2, 4, 10), + norm='log', + extend='both', + colors='k', + ) + self._graph_axes[0].scatter( + 1, + 1, + color='red', + marker='*', + s=200, + zorder=10, + ) + self._graph_axes[0].set_xlabel('Simulation Parameter #1') + self._graph_axes[0].set_ylabel('Simulation Parameter #2') + self._graph_axes[0].set_title('Ground Truth Rosenbrock Function') + + colorbar = self._graph_fig.colorbar(contourf, ax=self._graph_axes[0]) + colorbar.set_ticks(np.logspace(-2, 4, 7)) + # colorbar.set_label('Simulation Result') + self._graph_colorbars[self._graph_axes[0]] = colorbar + + self._graph_fig.tight_layout() + self._graph_fig.savefig(self.save_path, bbox_inches='tight') + + def __call__( + self, + new_x: list[float] | list[list[float]], + train_x, + train_y, + surrogate_y=None, + final: bool = False, + ): + strategy_index = self.scheduler.get_strategy_index(len(train_x)) + ax = self._graph_axes[strategy_index + 1] + + # Save attributes before clearing + t = ax.get_title() + xl = ax.get_xlabel() + yl = ax.get_ylabel() + # Clear the axis content only + ax.cla() + # Restore the attributes + ax.set_title(t) + ax.set_xlabel(xl) + ax.set_ylabel(yl) + + if NUM_DIMS == 2: + new_x = np.asarray(new_x, dtype=float).reshape(-1, NUM_DIMS) + + if surrogate_y is not None: + data = np.maximum(np.asarray(surrogate_y), 0.11) + else: + data = np.zeros((MESHGRID_SIZE, MESHGRID_SIZE)) + + contourf = ax.contourf( + SURROGATE_MESHGRID[0], + SURROGATE_MESHGRID[1], + data, + levels=np.logspace(-2, 4, 101), + norm='log', + extend='both', + ) + ax.scatter( + 1.0, + 1.0, + s=300, + facecolors='none', + edgecolors='black', + marker='o', + label='True Minimum', + zorder=10, + ) + + self._graph_colorbars[ax].update_normal(contourf) + + optimal_coords = None + minpos = None + if len(train_x) > 0: + train_x = np.asarray(train_x) + + prev_strategies_points = self.scheduler.strategy_break_points[ + max(0, strategy_index - 1) + ] + + ax.scatter( + train_x[:prev_strategies_points, 0], + train_x[:prev_strategies_points, 1], + color='black', + marker='s', + label='Previous strategies', + s=30, + alpha=0.3, + zorder=10, + ) + ax.scatter( + train_x[prev_strategies_points:, 0], + train_x[prev_strategies_points:, 1], + color='black', + marker='o', + label='Current strategy', + s=50, + zorder=10, + ) + + minpos = int(np.argmin(train_y)) + optimal_coords = np.asarray(train_x[minpos]) + + ax.scatter( + optimal_coords[0], + optimal_coords[1], + color='red', + marker='*', + s=200, + label='Best Point Estimate', + zorder=10, + ) + + if final and optimal_coords is not None: + ax.set_title('final surrogate') + final_x = ', '.join(f'{coord:.2f}' for coord in optimal_coords) + + self._graph_fig.suptitle( + f'Best point estimate so far is x=({final_x}), ' + f'y={train_y[minpos]:.3f}; ' + 'true minimum is at x=(1.00, 1.00), y=0.00', + x=0.5, + y=1.0, + ) + else: + ax.scatter( + new_x[:, 0], + new_x[:, 1], + color='red', + marker='o', + label='Recommended Points', + s=50, + zorder=10, + ) + + # Extract and deduplicate handles and labels from ALL subplots + handles, labels = [], [] + for ax in self._graph_fig.axes: + h, lab = ax.get_legend_handles_labels() + handles.extend(h) + labels.extend(lab) + # Use a dictionary to keep only the first occurrence of each unique label + by_label = dict(zip(labels, handles, strict=False)) + # Create the clean figure-level legend + self._graph_fig.legend( + by_label.values(), + by_label.keys(), + loc='outside lower center', + ncol=3, + bbox_to_anchor=(0.5, -0.05), + ) + + self._graph_fig.tight_layout() + self._graph_fig.savefig(self.save_path, bbox_inches='tight') + + +class ActiveLearningOrchestrator: + def __init__(self, service_destination: str, rosenbrock_destination: str): + self.service_destination = service_destination + self.rosenbrock_destination = rosenbrock_destination + + # This value gets populated from the return value of initializing the workflow + self.workflow_id = '' + + # The full dataset object state only needs to exist for the purposes of generating the graph and determining a stop-workflow order + # if we don't care about "step by step" data, we technically do NOT need to save these as stateful, as we can get the data at the end by calling "dial.get_workflow_data" + self.dataset_x = [] + self.dataset_y: list[float] = [] + + self.next_x = None + self.surrogate_y = None + + self.scheduler = Scheduler(STRATEGY_SCHEDULE) + self.plotter = Plotter(scheduler=self.scheduler) + + @property + def current_strategy(self): + return self.scheduler(len(self.dataset_y)) + + # create a message to send to the server + def assemble_message(self, operation: str, **kwargs: Any) -> IntersectClientCallback: + print(f'assembling dial message: {operation}') + if operation == 'initialize_workflow': + payload = DialWorkflowCreationParamsClient( + dataset_x=self.dataset_x, + dataset_y=self.dataset_y, + bounds=BOUNDS, + dim_x=NUM_DIMS, # Explicitly set the dimension based on the 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', + }, + length_per_dimension=False, # allow the matern to use separate length scales for the two parameters + y_is_good=False, # we wish to minimize y (the error) + backend='sklearn', # "sklearn" or "gpax" + seed=-1, # Use seed = -1 for random results + ) + elif operation == 'update_workflow_with_data': + payload = DialWorkflowDatasetUpdate( + workflow_id=self.workflow_id, + **kwargs, + ) + elif operation == 'update_workflow_with_batch_data': + payload = DialWorkflowDatasetUpdates( + workflow_id=self.workflow_id, + **kwargs, + ) + elif operation == 'get_next_point': + payload = DialInputSingleOtherStrategy( + workflow_id=self.workflow_id, + strategy=self.current_strategy['strategy'], + strategy_args=self.current_strategy.get('args', None), + bounds=BOUNDS, + ) + elif operation == 'get_next_points': + if self.current_strategy.get('batch_strategy', None) is None: + return self.assemble_message('get_next_point') + payload = DialInputMultipleOtherStrategy( + workflow_id=self.workflow_id, + batch_strategy=self.current_strategy.get('batch_strategy'), + strategy=self.current_strategy['strategy'], + strategy_args=self.current_strategy.get('args', None), + points=self.current_strategy.get('num_samples', 1), + bounds=BOUNDS, + ) + elif operation == 'get_surrogate_values': + payload = DialInputPredictions( + workflow_id=self.workflow_id, + points_to_predict=POINTS_TO_PREDICT, + ) + else: + err_msg = f'Invalid operation {operation}' + raise Exception(err_msg) # noqa: TRY002 + return IntersectClientCallback( + messages_to_send=[ + IntersectDirectMessageParams( + destination=self.service_destination, + operation=f'dial.{operation}', + payload=payload, + ) + ] + ) + + def assemble_rosenbrock_message(self, operation: str) -> IntersectClientCallback: + print(f'assembling rosenbrock message: {operation}') + if operation == 'rosenbrock': + payload = { + 'x': self.next_x[0], + 'y': self.next_x[1], + } + elif operation == 'rosenbrock_bulk': + payload = [{'x': x[0], 'y': x[1]} for x in self.next_x] + else: + err_msg = f'Invalid operation {operation}' + raise Exception(err_msg) # noqa: TRY002 + return IntersectClientCallback( + messages_to_send=[ + IntersectDirectMessageParams( + destination=self.rosenbrock_destination, + operation=f'Rosenbrock.{operation}', + payload=payload, + ) + ] + ) + + # The callback function. This is called whenever the server responds to our message. + # This could instead be implemented by defining a callback method (and passing it later), but here we chose to directly make the object callable. + def __call__( + self, + _source: str, + operation: str, + has_error: bool, + payload: INTERSECT_RESPONSE_VALUE, + ) -> IntersectClientCallback: + if has_error: + print('============ERROR==============', file=sys.stderr) + print(operation, file=sys.stderr) + print(payload, file=sys.stderr) + print(file=sys.stderr) + msg = f'Error in operation {operation}: payload = {payload}' + raise Exception(msg) # noqa: TRY002 (break INTERSECT loop) + if operation == 'Rosenbrock.rosenbrock': + coord_str = ', '.join([f'{coord:.2f}' for coord in self.next_x]) + print(f'Running simulation at ({coord_str}): {payload:.3f}') + self.dataset_x.append(self.next_x) + self.dataset_y.append(payload) + return self.assemble_message( + 'update_workflow_with_data', next_x=self.next_x, next_y=payload + ) + if operation == 'Rosenbrock.rosenbrock_bulk': + self.dataset_x += self.next_x + self.dataset_y += payload + coord_str = '\n '.join( + [ + '(' + ', '.join([f'{coord:.2f}' for coord in next_x]) + f') -> {next_y:.3f}' + for next_x, next_y in zip(self.next_x, payload, strict=False) + ] + ) + print(f'Running simulation at\n {coord_str}') + return self.assemble_message( + 'update_workflow_with_batch_data', next_x_list=self.next_x, next_y_list=payload + ) + if operation == 'dial.initialize_workflow': + self.workflow_id: str = payload + return self.assemble_message('get_next_points') + if operation == 'dial.update_workflow_with_data': + return self.assemble_message('get_surrogate_values') + if operation == 'dial.update_workflow_with_batch_data': + return self.assemble_message('get_surrogate_values') + if operation == 'dial.get_surrogate_values': + means = payload['values'] + self.surrogate_y = np.array(means).reshape((MESHGRID_SIZE,) * NUM_DIMS) + if len(self.dataset_x) >= NUM_ITERATIONS: + minpos = np.argmin(self.dataset_y) + x_opt = self.dataset_x[minpos] + y_opt = self.dataset_y[minpos] + # self.graph(x_opt, True) + self.plotter( + new_x=x_opt, + train_x=self.dataset_x, + train_y=self.dataset_y, + surrogate_y=self.surrogate_y, + final=True, + ) + coord_str = ', '.join([f'{coord:.2f}' for coord in x_opt]) + print( + f'Optimal simulated datapoint at ({coord_str}), y={y_opt:.3f}', + end='\n', + flush=True, + ) + msg = 'Client simulation completed successfully.' + raise Exception(msg) # noqa: TRY002 (INTERSECT interaction mechanism, do not need custom exception) + return self.assemble_message('get_next_points') + if operation == 'dial.get_next_point': + # if we receive an EI recommendation, record it, show the user the current graph, and run the "simulation": + self.next_x = payload['data'] + self.plotter( + new_x=self.next_x, + train_x=self.dataset_x, + train_y=self.dataset_y, + surrogate_y=self.surrogate_y, + final=False, + ) + return self.assemble_rosenbrock_message('rosenbrock') + if operation == 'dial.get_next_points': + # if we receive an EI recommendation, record it, show the user the current graph, and run the "simulation": + self.next_x = payload['data'] + self.plotter( + new_x=self.next_x, + train_x=self.dataset_x, + train_y=self.dataset_y, + surrogate_y=self.surrogate_y, + final=False, + ) + return self.assemble_rosenbrock_message('rosenbrock_bulk') + + err_msg = f'Unknown operation received: {operation}' + raise Exception(err_msg) # noqa: TRY002 (INTERSECT interaction mechanism) + + +if __name__ == '__main__': + # In production, everything in this dictionary should come from a configuration file, command line arguments, or environment variables. + parser = argparse.ArgumentParser(description='Automated client') + parser.add_argument( + '--config', + type=Path, + default=os.environ.get('DIAL_CONFIG_FILE', Path(__file__).parents[1] / 'local-conf.json'), + ) + args = parser.parse_args() + try: + with Path(args.config).open('rb') as f: + from_config_file = json.load(f) + except (json.decoder.JSONDecodeError, OSError) as e: + logger.critical('unable to load config file: %s', str(e)) + sys.exit(1) + + active_learning = ActiveLearningOrchestrator( + service_destination=HierarchyConfig( + **from_config_file['intersect-hierarchy'] + ).hierarchy_string('.'), + rosenbrock_destination=HierarchyConfig( + **from_config_file['rosenbrock-hierarchy'] + ).hierarchy_string('.'), + ) + config = IntersectClientConfig( + initial_message_event_config=active_learning.assemble_message('initialize_workflow'), + **from_config_file['intersect'], + ) + # use the orchestator to create the client + client = IntersectClient( + config=config, + # the callback (here we use a callable object, as discussed above) + user_callback=active_learning, + ) + # This will run the send message -> wait for response -> callback -> repeat cycle until we have 25 points (and then raise an Exception) + default_intersect_lifecycle_loop( + client, + ) diff --git a/src/dial_dataclass/dial_dataclass.py b/src/dial_dataclass/dial_dataclass.py index 8f45ebc..d5f5289 100644 --- a/src/dial_dataclass/dial_dataclass.py +++ b/src/dial_dataclass/dial_dataclass.py @@ -399,20 +399,36 @@ class DialInputSingleConfidenceBound(BaseModel): ] -class DialInputSingleOtherStrategy(BaseModel): - """This class is used to request a single next point using a given strategy.""" +SingleStrategyType = Literal[ + ############################# + # surrogate free strategies + 'center', + 'corners', + 'grid', + 'chebyshev', + 'latin_hypercube', + 'random', + ############################# + # surrogate based strategies + 'uncertainty', + 'expected_improvement', + 'upper_confidence_bound', + 'upper_confidence_bound_nomad', + 'polymer_acl_sampler', +] + +MultipleStrategyType = Literal[ + ############################# + # batch strategies + 'liar', 'believer' +] + + +class DialInputSingleOtherStrategy(BaseModel): workflow_id: ValidatedObjectId - strategy: Literal[ - 'random', - 'hypercube', - 'uncertainty', - 'expected_improvement', - 'upper_confidence_bound', - 'upper_confidence_bound_nomad', - 'polymer_acl_sampler', - ] - strategy_args: dict[str, float | int | bool] | None = Field(default=None) + strategy: SingleStrategyType + strategy_args: dict[str, float | int | bool | list[int | float]] | None = Field(default=None) y_is_good: Annotated[ bool | None, Field( @@ -473,16 +489,15 @@ class DialInputMultipleOtherStrategy(BaseModel): workflow_id: ValidatedObjectId points: PositiveIntType - strategy: Literal[ - 'random', - 'uncertainty', - 'expected_improvement', - 'upper_confidence_bound', - 'upper_confidence_bound_nomad', - 'polymer_acl_sampler', - 'hypercube', - ] - strategy_args: dict[str, float | int | bool] | None = Field(default=None) + strategy: SingleStrategyType + batch_strategy: MultipleStrategyType | None = Field( + default=None, + description='If a batch strategy is specified, it will be used to generate multiple points from the single point strategy. ' + 'If None, the single point strategy will be used to generate multiple points.', + ) + strategy_args: dict[str, str | float | int | bool | list[int | float]] | None = Field( + default=None + ) y_is_good: Annotated[ bool | None, Field( diff --git a/src/dial_service/core.py b/src/dial_service/core.py index d7b616d..31ced9e 100644 --- a/src/dial_service/core.py +++ b/src/dial_service/core.py @@ -16,8 +16,11 @@ ServersideInputSingle, ) from .utilities.strategies import ( + INDEXED_STRATEGIES, + batch_sampling, create_measurement_grid, hypercube, + indexed_selection, random_in_bounds, ) @@ -34,6 +37,10 @@ def get_next_point(data: ServersideInputSingle, model: Any) -> list[float]: Returns: list[float]: The selected point for the next iteration. """ + + if data.strategy in INDEXED_STRATEGIES: + return indexed_selection(data) + # If it's random point, we don't need to train a model or anything else if data.strategy == 'random': if data.discrete_measurements: @@ -78,6 +85,13 @@ def get_next_points(data: ServersideInputMultiple, model: Any) -> list[list[floa """ # model = self._train_model(data) #this will be needed when we add qEI/constant liars output_points = None + + backend = data.backend.lower() + module = get_backend_module(backend) + + if data.batch_strategy is not None: + return batch_sampling(module, model, data) + match data.strategy: case 'random': output_points = [ diff --git a/src/dial_service/dial_service.py b/src/dial_service/dial_service.py index bc1a5be..153c007 100644 --- a/src/dial_service/dial_service.py +++ b/src/dial_service/dial_service.py @@ -55,7 +55,7 @@ def initialize_workflow(self, client_data: DialWorkflowCreationParamsService) -> """ try: server_data = ServersideInputBase(client_data) - if client_data.dataset_x: + if client_data.dataset_x and len(client_data.dataset_y) > 0: # the user provided some initial data, so train a model model = pickle.dumps(core.train_model(server_data), protocol=5) else: diff --git a/src/dial_service/serverside_data.py b/src/dial_service/serverside_data.py index 0602da2..d99fbb7 100644 --- a/src/dial_service/serverside_data.py +++ b/src/dial_service/serverside_data.py @@ -62,14 +62,11 @@ def scale_X(self, X: np.ndarray) -> np.ndarray: return (X - lows) / span - def _extract_y_train_from_dataset(self): + @cached_property + def y_train_raw(self) -> np.ndarray: """ - Find output y and error values yerr in dataset_y, and save. + Return the raw training target values extracted from the dataset. """ - if hasattr(self, 'y_train_raw') and hasattr(self, 'yerr_train_raw'): - # only compute this on first invocation - return - y_label = self.statistics_y.loc if not isinstance(y_label, str): msg = 'statistics_y.loc must be a Label (str).' @@ -78,21 +75,27 @@ def _extract_y_train_from_dataset(self): # Use the label from self.statistics_y.loc to find the data column with the mean y data # this may trigger a ValueError, if the label does not exist, but should be handled by dataclass validation pos_y = self.labels_y.index(y_label) - self.y_train_raw = self.dataset_y[:, pos_y] + return self.dataset_y[:, pos_y] + @cached_property + def yerr_train_raw(self) -> any: + """ + Return the raw training error values extracted from the dataset. + """ yerr_label = self.statistics_y.scale if isinstance(yerr_label, float): - self.yerr_train_raw = yerr_label + _yerr_train_raw = yerr_label else: # yerr_label is str # this may trigger a ValueError, but should be handled by dataclass validation pos_yerr = self.labels_y.index(yerr_label) - self.yerr_train_raw = self.dataset_y[:, pos_yerr] + _yerr_train_raw = self.dataset_y[:, pos_yerr] - if np.any(self.yerr_train_raw < 0): - idxs = np.where(np.yerr_train_raw < 0) - msg = f'yerr values in statistics_y.scale must be non-negative, found {np.yerr_train_raw[idxs[0]]} at {idxs[0]}.' + if np.any(_yerr_train_raw < 0): + idxs = np.where(_yerr_train_raw < 0) + msg = f'yerr values in statistics_y.scale must be non-negative, found {_yerr_train_raw[idxs[0]]} at {idxs[0]}.' raise ValueError(msg) + return _yerr_train_raw @cached_property def Y_train(self) -> np.ndarray: @@ -100,9 +103,6 @@ def Y_train(self) -> np.ndarray: Find output y and error values yerr in dataset_y, and apply transformation. Return transformed y value. """ - # ensure that self.y_train_raw, self.yerr_train_raw are populated - self._extract_y_train_from_dataset() - y, _ = self.transform_Y(self.y_train_raw, self.yerr_train_raw) # return only y, to conform to interface @@ -114,9 +114,6 @@ def Yerr_train(self) -> any: Find output y and error values in dataset y, and apply transformation. Return transformed yerr value. """ - # ensure that self.y_train_raw, self.yerr_train_raw are populated - self._extract_y_train_from_dataset() - # recompute transformation, at some overhead (probably not worth to optimize) _, yerr = self.transform_Y(self.y_train_raw, self.yerr_train_raw) @@ -127,9 +124,6 @@ def _transform_Y_params(self) -> tuple[float, float]: """ Return the appropriate mean and scaling of the raw y data for normalization """ - # ensure that self.y_train_raw, self.yerr_train_raw are populated - self._extract_y_train_from_dataset() - # find y_std from y_train_raw y_train = self.y_train_raw if len(y_train) > 0 and self.preprocess_standardize: @@ -177,6 +171,29 @@ def inverse_transform_Y(self, y: np.ndarray, yerr: any) -> tuple[np.ndarray, any def Y_best(self) -> float: return self.Y_train.max() if self.y_is_good else self.Y_train.min() + def clear_cached_properties(self) -> None: + # Track attribute names that have already been encountered. + # Classes are inspected in Method Resolution Order (MRO) order, + # starting with the most-derived class + resolved_names = set() + + # Walk through the class hierarchy: + # DerivedClass -> BaseClass -> ... -> object + for cls in type(self).__mro__: + # Inspect only attributes defined directly on the current class + for name, attr in cls.__dict__.items(): + # Skip names already defined by a more-derived class + if name in resolved_names: + continue + + resolved_names.add(name) + + # A cached_property descriptor is stored on the class, + # while its computed value is stored in the instance __dict__ + if isinstance(attr, cached_property): + # Remove the cached value if it exists + self.__dict__.pop(name, None) + class ServersideInputSingle(ServersideInputBase): def __init__( @@ -232,6 +249,7 @@ def __init__( # set new inputs super().__init__(workflow_state) self.strategy = params.strategy + self.batch_strategy = params.batch_strategy self.points = params.points self.strategy = params.strategy self.strategy_args = params.strategy_args diff --git a/src/dial_service/utilities/strategies.py b/src/dial_service/utilities/strategies.py index 51d155a..6f53949 100644 --- a/src/dial_service/utilities/strategies.py +++ b/src/dial_service/utilities/strategies.py @@ -1,5 +1,6 @@ import itertools import logging +from numbers import Real import numpy as np from scipy.optimize import minimize @@ -92,6 +93,113 @@ def confidence_bound(mean, stddev, data): } +############################################################################### +# Surrogate-free indexed strategies + + +def domain_center(data, indices=None): # noqa: ARG001 + return [[0.5 * (data.bounds[i][1] + data.bounds[i][0]) for i in range(data.dim_x)]] + + +def domain_corners(data, indices: list[int]): + points = [] + for index in indices: + # convert flat index into coordinate indices for each dimension + coo_indices = np.unravel_index(index, [2] * data.dim_x) # 2 corners per dimension + points.append([data.bounds[i][coo_indices[i]] for i in range(data.dim_x)]) + return points + + +def uniform_grid(data, indices: list[int]): + grid_size = data.strategy_args['grid_size'] + + # grid spacing in each dimension + steps = [ + (data.bounds[i][1] - data.bounds[i][0]) / max(1, grid_size[i] - 1) + for i in range(data.dim_x) + ] + + points = [] + for index in indices: + # convert flat index into coordinate indices for each dimension + coo_indices = np.unravel_index(index, grid_size) + points.append( + [ + data.bounds[i][0] + coo_indices[i] * steps[i] + if grid_size[i] > 1 + else 0.5 * (data.bounds[i][1] + data.bounds[i][0]) + for i in range(data.dim_x) + ] + ) + return points + + +def chebyshev_grid(data, indices: list[int]): + grid_size = data.strategy_args['grid_size'] + + points = [] + for index in indices: + # convert flat index into coordinate indices for each dimension + coo_indices = np.unravel_index(index, grid_size) + # Chebyshev nodes in each dimension, or the midpoint if only one point is specified + x = [ + np.cos(coo_indices[i] * np.pi / (grid_size[i] - 1)) + if grid_size[i] > 1 + else 0.5 * (data.bounds[i][1] + data.bounds[i][0]) + for i in range(data.dim_x) + ] + points.append( + [ + 0.5 * (data.bounds[i][1] - data.bounds[i][0]) * (x[i] + 1) + data.bounds[i][0] + for i in range(data.dim_x) + ] + ) + return points + + +def latin_hypercube(data, indices: list[int]): + """For latin hypercube grid_size is the number of intervals per dimension""" + grid_size = data.strategy_args['grid_size'] + + points = [] + for index in indices: + # convert flat index into interval indices for each dimension + interval_indices = np.unravel_index(index, grid_size) + points.append( + [ + data.numpy_rng.uniform( + data.bounds[i][0] + + interval_indices[i] * (data.bounds[i][1] - data.bounds[i][0]) / grid_size[i], + data.bounds[i][0] + + (interval_indices[i] + 1) + * (data.bounds[i][1] - data.bounds[i][0]) + / grid_size[i], + ) + for i in range(data.dim_x) + ] + ) + return points + + +INDEXED_STRATEGIES = { + 'center': domain_center, + 'corners': domain_corners, + 'grid': uniform_grid, + 'chebyshev': chebyshev_grid, + 'latin_hypercube': latin_hypercube, +} + +MAX_INDEXED_POINTS = { + 'center': lambda data: 1, # noqa: ARG005 + 'corners': lambda data: 2**data.dim_x, + 'grid': lambda data: np.prod(data.strategy_args['grid_size']), + 'chebyshev': lambda data: np.prod(data.strategy_args['grid_size']), + 'latin_hypercube': lambda data: np.prod(data.strategy_args['grid_size']), +} + +############################################################################### + + def hypercube( bounds: list[list[float]], num_points: int, rng: np.random.RandomState ) -> list[list[float]]: @@ -173,6 +281,112 @@ def to_minimize(_x: np.ndarray): return selected_point.tolist() +def indexed_selection(data: ServersideInputSingle): + try: + strategy_ = INDEXED_STRATEGIES[data.strategy] + except KeyError as exc: + msg = f'Invalid strategy: {data.strategy}' + raise ValueError(msg) from exc + + start_index = data.strategy_args.get('start_index', 0) if data.strategy_args is not None else 0 + index = (len(data.dataset_y) - start_index) % MAX_INDEXED_POINTS[data.strategy](data) + selected_point = strategy_(data, [index])[0] + + return selected_point + + +def batch_sampling(backend_module: AbstractBackend, model, data: ServersideInputMultiple): + """ + Greedy batch selection using the liar or believer strategies. + """ + + if data.points <= 0: + return [] + + selected_points: list[list[float]] = [] + initial_x = np.asarray(data.dataset_x, dtype=float) + initial_y = np.asarray(data.dataset_y, dtype=float) + + if data.batch_strategy is not None: + if data.batch_strategy == 'liar': + liar_setting = ( + data.strategy_args.get('liar_value', 'mean') + if data.strategy_args is not None + else 'mean' + ) + + if isinstance(liar_setting, Real): + liar_value = float(liar_setting) + elif isinstance(liar_setting, str): + match liar_setting: + case 'mean': + liar_value = np.mean(data.dataset_y) + case 'max': + liar_value = np.max(data.dataset_y) + case 'min': + liar_value = np.min(data.dataset_y) + case 'random': + liar_value = data.numpy_rng.uniform( + np.min(data.dataset_y), np.max(data.dataset_y) + ) + case 'median': + liar_value = np.median(data.dataset_y) + case _: + liar_value = np.mean(data.dataset_y) + elif callable(liar_setting): + liar_value = liar_setting(data.dataset_y) + + def predictor(point): # noqa: ARG001 + return liar_value + + elif data.batch_strategy == 'believer': + believer_setting = ( + data.strategy_args.get('believer_type', 'kriging') + if data.strategy_args is not None + else 'kriging' + ) + + match believer_setting: + case 'kriging': + + def predictor(point): + data.set_x_predict(point) + return backend_module.predict(model, data)[0][0] + + current_model = model + try: + for _ in range(data.points): + if data.strategy in INDEXED_STRATEGIES: + point = indexed_selection(data) + else: + point = greedy_sampling(backend_module, current_model, data) + selected_points.append([float(v) for v in point]) + + x_arr = np.asarray(point, dtype=float).reshape(1, -1) + y_arr = np.asarray([[predictor(selected_points[-1])]], dtype=float) + + if data.dataset_x.size == 0: + data.dataset_x = x_arr + else: + data.dataset_x = np.vstack([np.asarray(data.dataset_x, dtype=float), x_arr]) + data.dataset_y = np.concatenate([np.asarray(data.dataset_y, dtype=float), y_arr]) + + # Why do we need to strip cached properties here? + # Because we are modifying dataset_x and dataset_y, which are used in cached properties like stddev, Y_best, etc. + # If we don't clear these caches, they might return outdated values based on the old dataset + # By stripping the cached properties, we ensure that the next time these properties are accessed, they will be recalculated based on the updated dataset + data.clear_cached_properties() + + current_model = backend_module.train_model(data) + finally: + # Restore original state so pseudo-observations never leak outside this method. + data.dataset_x = initial_x + data.dataset_y = initial_y + data.clear_cached_properties() + + return selected_points + + def batch_sampling_acl(backend_module: AbstractBackend, model, data: ServersideInputMultiple): """ Greedy batch selection using GP std and multiple penalties: