Add Tsuchinoko agent

.github/workflows/tests.yml

@@ -25,7 +25,7 @@ jobs:
       - name: Install Docker Compose
         run: |
           sudo apt-get update
-          sudo apt-get install -y docker-compose
+          sudo apt-get install -y docker-compose freeglut3-dev libopenblas-dev
 
       - name: Download and build bluesky-pods
         run: |
@@ -48,7 +48,7 @@ jobs:
         run: |
           set -vxeuo pipefail
           pip install --upgrade pip wheel
-          pip install .[dev,agents]
+          pip install .[dev,agents,tsuchinoko]
           mkdir -p /home/runner/.config/tiled/profiles
           cp ./bluesky_adaptive/tests/podman/tiled_client_config.yml /home/runner/.config/tiled/profiles/tiled_client_config.yml
           pip list
@@ -90,4 +90,4 @@ jobs:
         run: |
           set -vxeuo pipefail
           coverage run -m pytest -v
-          coverage report
+          coverage report -i

bluesky_adaptive/agents/tsuchinoko.py

@@ -0,0 +1,175 @@
+import pickle
+import time
+import warnings
+from abc import ABC
+from collections.abc import Sequence
+from logging import getLogger
+
+import numpy as np
+import zmq
+from numpy._typing import ArrayLike
+
+from .base import Agent
+
+logger = getLogger("bluesky_adaptive.agents")
+
+SLEEP_FOR_AGENT_TIME = 0.1
+SLEEP_FOR_TSUCHINOKO_TIME = 0.1
+FORCE_KICKSTART_TIME = 5
+
+
+class TsuchinokoBase:
+    def __init__(self, *args, host: str = "127.0.0.1", port: int = 5557, **kwargs):
+        """
+
+        Parameters
+        ----------
+        args
+            args passed through to `bluesky_adaptive.agents.base.Agent.__init__()`
+        host
+            A host address target for the zmq socket.
+        port
+            The port used for the zmq socket.
+        kwargs
+            kwargs passed through to `bluesky_adaptive.agents.base.Agent.__init__()`
+        """
+
+        super().__init__(*args, **kwargs)
+        self.host = host
+        self.port = port
+        self.outbound_measurements = []
+        self.context = None
+        self.socket = None
+        self.setup_socket()
+        self.last_targets_received = time.time()
+        self.kickstart()
+
+    def kickstart(self):
+        self.send_payload({"send_targets": True})  # kickstart to recover from shutdowns
+        self.last_targets_received = time.time()  # forgive lack of response until now
+
+    def setup_socket(self):
+        self.context = zmq.Context()
+        self.socket = self.context.socket(zmq.PAIR)
+
+        # Attempt to connect, retry every second if fails
+        while True:
+            try:
+                self.socket.connect(f"tcp://{self.host}:{self.port}")
+            except zmq.ZMQError:
+                logger.info(f"Unable to connect to tcp://{self.host}:{self.port}. Retrying in 1 second...")
+                time.sleep(1)
+            else:
+                logger.info(f"Connected to tcp://{self.host}:{self.port}.")
+                break
+
+    def ingest(self, x, yv):
+        """
+        Send measurement to BlueskyAdaptiveEngine
+        """
+        payload = {"target_measured": (x, yv)}
+        self.send_payload(payload)
+
+    def suggest(self, batch_size: int = 1) -> Sequence[ArrayLike]:
+        """
+        Wait until at least one target is received, also exhaust the queue of
+        received targets, overwriting old ones
+        """
+        payload = None
+        while True:
+            try:
+                payload = self.recv_payload(flags=zmq.NOBLOCK)
+            except zmq.ZMQError:
+                if payload is not None:
+                    break
+                else:
+                    time.sleep(SLEEP_FOR_TSUCHINOKO_TIME)
+                    if time.time() > self.last_targets_received + FORCE_KICKSTART_TIME:
+                        self.kickstart()
+        assert "candidate" in payload
+        self.last_targets_received = time.time()
+        return payload
+
+    def send_payload(self, payload: dict):
+        logger.info(f"message: {payload}")
+        self.socket.send(pickle.dumps(payload))
+
+    def recv_payload(self, flags=0) -> dict:
+        payload_response = pickle.loads(self.socket.recv(flags=flags))
+        logger.info(f"response: {payload_response}")
+        return payload_response
+
+
+class TsuchinokoAgent(TsuchinokoBase, Agent, ABC):
+    """
+    A Bluesky-Adaptive 'Agent'. This Agent communicates with Tsuchinoko over zmq
+    to request new targets and report back measurements. This is an abstract
+    class that must be subclassed.
+
+    A `tsuchinoko.execution.bluesky_adaptive.BlueskyAdaptiveEngine` is required
+    for the Tsuchinoko server to complement one of these `TsuchinokoAgent`.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._targets_shape = None
+
+    def ingest(self, x, yv) -> dict[str, ArrayLike]:
+        super().ingest(x, yv)
+        return self.get_ingest_document(x, yv)
+
+    def suggest(self, batch_size: int = 1) -> tuple[Sequence[dict[str, ArrayLike]], Sequence[ArrayLike]]:
+        targets = super().suggest(batch_size)
+        optimizer_state = targets.pop("optimizer")
+        return self.get_suggest_documents(targets, optimizer_state), targets
+
+    def get_ingest_document(self, x, yv) -> dict[str, ArrayLike]:
+        """
+        Return any single document corresponding to 'tell'-ing Tsuchinoko about the newly measured `x`, `y` data
+
+        Parameters
+        ----------
+        x :
+            Independent variable for data observed
+        yv :
+            Dependent variable for data observed, concatenated with variance
+        Returns
+        -------
+        dict
+            Dictionary to be unpacked or added to a document
+
+        """
+        y, v = yv
+        return {"independent": np.asarray(x), "observable": np.asarray(y), "variance": np.asarray(v)}
+
+    def get_suggest_documents(
+        self, targets: Sequence[ArrayLike], optimizer_state: dict
+    ) -> Sequence[dict[str, ArrayLike]]:
+        """
+        Ask the agent for a new batch of points to measure.
+
+        Parameters
+        ----------
+        targets : List[Tuple]
+            The new target positions to be measured received during this `ask`.
+        optimizer_state: Dict
+            The serialized state of a GPOptimizer instance
+
+        Returns
+        -------
+        docs : Sequence[dict]
+            Documents of key metadata from the ask approach for each point in next_points.
+            Must be length of batch size.
+
+        """
+
+        # check if targets length changes
+        if not self._targets_shape:
+            self._targets_shape = len(targets)
+        if self._targets_shape != len(targets):
+            warnings.warn(
+                "The length of the target queue has changed. A new databroker run will be generated", stacklevel=2
+            )
+            self.close_and_restart()
+
+        return [targets | optimizer_state]

bluesky_adaptive/tests/test_tsuchinoko_agents.py

@@ -0,0 +1,118 @@
+import sys
+from threading import Thread
+from typing import Union
+
+import numpy as np
+import pytest
+from numpy.typing import ArrayLike
+from pytest import fixture
+from tsuchinoko.adaptive.gpCAM_in_process import GPCAMInProcessEngine
+from tsuchinoko.core import CoreState, ZMQCore
+from tsuchinoko.execution.bluesky_adaptive import BlueskyAdaptiveEngine
+from xarray import Dataset
+
+from bluesky_adaptive.agents.tsuchinoko import TsuchinokoAgent
+from bluesky_adaptive.typing import BlueskyRunLike
+from bluesky_adaptive.utils.offline import OfflineAgent
+
+from .conftest import catalog  # noqa: F401
+
+
+@fixture
+def gpcam_engine():
+    # Define a gpCAM adaptive engine with initial parameters
+    adaptive = GPCAMInProcessEngine(
+        dimensionality=2,
+        parameter_bounds=[(0, 1), (0, 1)],
+        hyperparameters=[255, 100, 100],
+        hyperparameter_bounds=[(0, 1e5), (0, 1e5), (0, 1e5)],
+    )
+    return adaptive
+
+
+@fixture
+def execution_engine(gpcam_engine):
+    execution = BlueskyAdaptiveEngine(gpcam_engine)
+    yield execution
+
+
+@fixture
+def core(gpcam_engine, execution_engine):
+    print("starting setup")
+    assert gpcam_engine is execution_engine.adaptive_engine
+    core = ZMQCore()
+    core.set_adaptive_engine(gpcam_engine)
+    core.set_execution_engine(execution_engine)
+
+    class ErrorThread(Thread):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.error = None
+
+        def run(self):
+            try:
+                super().run()
+            except Exception as ex:
+                self.error = ex
+                print(ex)
+                print("The above error occurred in the tsuchinoko core thread.")
+
+    server_thread = ErrorThread(target=core.main)
+    server_thread.start()
+    core.state = CoreState.Starting
+    print("setup complete")
+
+    yield core
+
+    core.exit()
+    server_thread.join()
+    print("teardown complete")
+    if server_thread.error:
+        raise server_thread.error  # most likely a test would have already failed by this point
+
+
+class GPTestAgent(TsuchinokoAgent, OfflineAgent):
+    def measurement_plan(self, point: ArrayLike) -> tuple[str, list, dict]:
+        return self.measurement_plan_name, [1.5], {}
+
+    def unpack_run(self, run: BlueskyRunLike) -> tuple[Union[float, ArrayLike], Union[float, ArrayLike]]:
+        self.counter += 1
+        y = np.random.rand(10)
+        v = 0.1
+        return self.counter, (y, v)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="GPCam Requires Python 3.11 or higher")
+def test_gp_agent(catalog, core):  # noqa: F811
+    # Test ingest, suggest, and report; uses Tiled functionality
+    agent = GPTestAgent(tiled_data_node=catalog, tiled_agent_node=catalog)
+    agent.start()
+    agent_uid = agent._compose_run_bundle.start_doc["uid"]
+    for i in range(5):
+        uid = f"uid{i}"
+        x = np.random.rand(2)
+        y = 1 - np.sum(np.sin(x)) + np.random.rand() * 0.01
+        v = 0.1
+        yv = y, v
+        doc = agent.ingest(x, yv)
+        doc["exp_uid"] = uid
+        agent._write_event("ingest", doc)
+        agent.known_uid_cache.append(uid)
+    # agent.generate_report()
+    doc, query = agent.suggest()
+    agent._write_event("suggest", doc[0])
+    doc, query = agent.suggest()
+    agent._write_event("suggest", doc[0])
+    agent.stop()
+    assert len(query) == 1
+    run = catalog[agent_uid]
+    xr = run.suggest.read()
+    assert xr["candidate"].shape == (2, 1, 2)
+    # Acquisition value isn't available here; this is outside of separation of responsibilities for tsuchinoko's
+    # execution engines.
+    # assert xr["acquisition_value"].shape == (2,)
+    # TODO: getting AttributeError on 'report' here.
+    # xr = run.report.read()
+    # assert xr["STATEDICT-beta"].shape == (1,)
+    xr = run.ingest.read()
+    assert isinstance(xr, Dataset)

docs/source/reference/tsuchinoko-agent.md

@@ -0,0 +1,26 @@
+# Tsuchinoko Agent
+
+The Tsuchinoko agent enables using both the [gpCAM](https://gpcam.readthedocs.io/en/stable/) suggestion engine and the 
+[Tsuchinoko](https://tsuchinoko.readthedocs.io/en/latest/) graphical user interface.
+
+## Key Features of Tsuchinoko Agent
+
+- **Bayesian Optimization**: Utilizes a GP-based surrogate model for decision-making.
+- **Highly Customizable**: Modular flexibility of the suggestion algorithm with acquisition, kernel, noise, prior mean, and cost functions.
+- **High Performance**: Fast training and prediction with options to support distributed processing on HPC.
+- **Feedback and Control**: Visualization and live control of the agent from the Tsuchinoko desktop application keeps the user _in the loop_.
+
+```{eval-rst}
+.. autoclass:: bluesky_adaptive.agents.tsuchinoko.TsuchinokoAgent
+```
+
+To utilize the Tsuchinoko agent, as with other agent classes, the `measurement_plan` and `unpack_run` abstract methods
+must be defined in a subclass.
+
+To run the Tsuchinoko agent, you would need:
+- An installation of `bluesky_adaptive` with the optional `tsuchinoko` dependencies installed.
+- A running `TsuchinokoAgent`
+- A running `tsuchinoko` instance with a `BlueskyAdaptiveEngine` as its execution engine and a `GPCAMInProcessEngine` as
+  its adaptive engine.
+
+See `tests/test_tsuchinoko_agents.py` for refenece.

pyproject.toml

@@ -64,6 +64,7 @@ dev = [
   "databroker[all]>=2.0.0b1,<3.0.0",
   "bluesky-tiled-plugins",
 ]
+tsuchinoko = ["tsuchinoko>=1.1.23"]
 
 [project.urls]
 GitHub = "https://github.com/bluesky/bluesky-adaptive"