search

bluesky / bluesky-adaptive

Reference implementation for tightly integrated adaptive scans.
https://blueskyproject.io/bluesky-adaptive
BSD 3-Clause "New" or "Revised" License
4 stars 10 forks source link

API: Fix ask many #16

Closed maffettone closed 1 year ago

maffettone commented 1 year ago

Change the behavior of ask to return a sequence of documents and sequence suggestions/queries instead of a single document.

Description

This assigns many ask documents a common uid associated with a given ask call, and logs the batch index and the suggestion by default. Changes were propagated to the MonarchSubject agent base, and Botorch and simple agents. The Sklearn agents are passive (no ask) and remain unaffected. This creates an API burden to the user to enforce single docs be encapsulated in a list (as single points already were).

Core internal change with default keys is here, with other changes made to docs and for agent consistency. 

docs, next_points = self.ask(batch_size)
uid = str(uuid.uuid4())
for batch_idx, (doc, next_point) in enumerate(zip(docs, next_points)):
    doc["suggestion"] = next_point
    doc["batch_idx"] = batch_idx
    doc["batch_size"] = len(next_points)
    self._write_event("ask", doc, uid=uid)

Motivation and Context

This ensures that variable batch sizes can be stacked in a common data array. Initially it was though sufficient to just update a single ask doc as many docs with common metadata and different suggestion fields. However, there is often metadata associated with each individual suggestion that should be propagated through (e.g. in BoTorch we want to carry the suggestions and their associated acquisition value).

Summary of Changes for Release Notes

Fixed

Added

Changed

API change for ask -> Tuple[Sequence[Dict[str, ArrayLike]], Sequence[ArrayLike]]

Removed

How Has This Been Tested?

Running all unit tests sequentially using docker container for qserver, httpserver, mongo, and kafka.