Expanded output for get_surrogate_values()

.github/workflows/docker-release.yml

@@ -5,8 +5,10 @@ on:
     branches:
       - develop
       - main
+      # TODO remove this once this branch is merged into develop
+      - expanded-output
     tags:
-      - 'v*.*.*'
+      - "v*.*.*"
   workflow_dispatch:
 
 permissions:

scripts/1d_sable_client.py

@@ -215,17 +215,16 @@ def callback_message(self, operation: str, **kwargs) -> IntersectClientCallback:
         )
 
     def handle_surrogate_values(self, payload):
-        response_data = payload['data']
+        means = payload['values']
+        transformed_stddevs = payload['transformed_stddevs']
         if self.at_grids:
-            self.stddev_grid = np.array(response_data[1]).reshape(
-                (self.meshgrid_size,) * self.num_dims
-            )
-            self.mean_grid = np.array(response_data[0]).reshape(
+            self.stddev_grid = np.array(transformed_stddevs).reshape(
                 (self.meshgrid_size,) * self.num_dims
             )
+            self.mean_grid = np.array(means).reshape((self.meshgrid_size,) * self.num_dims)
         else:
-            self.stddev_test = np.array(response_data[1])
-            self.mean_test = np.array(response_data[0])
+            self.stddev_test = np.array(transformed_stddevs)
+            self.mean_test = np.array(means)
             print(
                 f'Values at testing points {self.x_test.reshape(-1)}: Mean: {self.mean_test}, Stddev: {self.stddev_test}'
             )

scripts/1d_sinusoidal_growth_client.py

@@ -187,17 +187,16 @@ def callback_message(self, operation: str, **kwargs) -> IntersectClientCallback:
         )
 
     def handle_surrogate_values(self, payload):
-        response_data = payload['data']
+        means = payload['values']
+        transformed_stddevs = payload['transformed_stddevs']
         if self.at_grids:
-            self.variance_grid = np.array(response_data[1]).reshape(
-                (self.meshgrid_size,) * self.num_dims
-            )
-            self.mean_grid = np.array(response_data[0]).reshape(
+            self.variance_grid = np.array(transformed_stddevs).reshape(
                 (self.meshgrid_size,) * self.num_dims
             )
+            self.mean_grid = np.array(means).reshape((self.meshgrid_size,) * self.num_dims)
         else:
-            self.variance_test = np.array(response_data[1])
-            self.mean_test = np.array(response_data[0])
+            self.variance_test = np.array(transformed_stddevs)
+            self.mean_test = np.array(means)
             print(f'Test Mean: {self.mean_test}, Variance: {self.variance_test}')
 
         # end of active learning loop after max_iter

scripts/2d_rosenbrock_client.py

@@ -208,8 +208,8 @@ def __call__(
         if (
             operation == 'dial.get_surrogate_values'
         ):  # if we receive a grid of surrogate values, record it for graphing, then ask for the next recommended point
-            data = payload['data']
-            self.mean_grid = np.array(data[0]).reshape((MESHGRID_SIZE,) * NUM_DIMS)
+            means = payload['values']
+            self.mean_grid = np.array(means).reshape((MESHGRID_SIZE,) * NUM_DIMS)
             return self.assemble_message('get_next_point')
 
         if operation == 'dial.get_next_point':

scripts/manual_client.py

@@ -187,7 +187,7 @@ def __call__(
         if (
             operation == 'dial.get_surrogate_values'
         ):  # if we receive a grid of surrogate values, record it for graphing, then ask for the next recommended point
-            self.mean_grid = np.array(payload[0]).reshape(XX.shape)
+            self.mean_grid = np.array(payload['values']).reshape(XX.shape)
             return self.assemble_message('get_next_point')
         if operation == 'dial.get_next_point':
             # if we receive an EI recommendation, record it, show the user the current graph, and ask the user for the results of their experiment:

src/dial_dataclass/__init__.py

@@ -12,5 +12,7 @@
 from .dial_dataclass_responses import (
     DialDataResponse1D,
     DialDataResponse2D,
+    DialSurrogateValuesResponse,
+    DialWorkflowFullState,
 )
 from .pydantic_helpers import ValidatedObjectId