better docs

bloptools/bayesian/__init__.py

@@ -13,6 +13,8 @@
 import pandas as pd
 import scipy as sp
 import torch
+from acquisition import default_acquisition_plan
+from digestion import default_digestion_function
 from matplotlib import pyplot as plt
 from matplotlib.patches import Patch
 
@@ -27,15 +29,6 @@
 DEFAULT_COLORMAP = "turbo"
 
 
-def default_acquisition_plan(dofs, inputs, dets):
-    uid = yield from bp.list_scan(dets, *[_ for items in zip(dofs, np.atleast_2d(inputs).T) for _ in items])
-    return uid
-
-
-def default_digestion_plan(db, uid):
-    return db[uid].table(fill=True)
-
-
 MAX_TEST_INPUTS = 2**11
 
 
@@ -172,7 +165,7 @@ def __init__(
         self.allow_acquisition_errors = kwargs.get("allow_acquisition_errors", True)
         self.initialization = kwargs.get("initialization", None)
         self.acquisition_plan = kwargs.get("acquisition_plan", default_acquisition_plan)
-        self.digestion = kwargs.get("digestion", default_digestion_plan)
+        self.digestion = kwargs.get("digestion", default_digestion_function)
         self.dets = list(np.atleast_1d(kwargs.get("dets", [])))
 
         self.acqf_config = kwargs.get("acqf_config", ACQF_CONFIG)

bloptools/bayesian/acquisition.py

@@ -1,37 +1,7 @@
-import torch
-from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement
+import bluesky.plans as bp
+import numpy as np
 
 
-class Acquisition:
-    def __init__(self, *args, **kwargs):
-        ...
-
-    @staticmethod
-    def log_expected_improvement(candidates, agent):
-        *input_shape, n_dim = candidates.shape
-
-        x = torch.as_tensor(candidates.reshape(-1, 1, n_dim)).double()
-
-        LEI = LogExpectedImprovement(
-            model=agent.task_model, best_f=agent.best_sum_of_tasks, posterior_transform=agent.scalarization
-        )
-
-        lei = LEI.forward(x)
-        feas_log_prob = agent.feas_model(x)
-
-        return (lei.reshape(input_shape) + feas_log_prob.reshape(input_shape)).detach().numpy()
-
-    @staticmethod
-    def log_probability_of_improvement(candidates, agent):
-        *input_shape, n_dim = candidates.shape
-
-        x = torch.as_tensor(candidates.reshape(-1, 1, n_dim)).double()
-
-        LPI = LogProbabilityOfImprovement(
-            model=agent.task_model, best_f=agent.best_sum_of_tasks, posterior_transform=agent.scalarization
-        )
-
-        lpi = LPI.forward(x)
-        feas_log_prob = agent.feas_model(x)
-
-        return (lpi.reshape(input_shape) + feas_log_prob.reshape(input_shape)).detach().numpy()
+def default_acquisition_plan(dofs, inputs, dets):
+    uid = yield from bp.list_scan(dets, *[_ for items in zip(dofs, np.atleast_2d(inputs).T) for _ in items])
+    return uid

bloptools/bayesian/digestion.py

@@ -0,0 +1,2 @@
+def default_digestion_function(db, uid):
+    return db[uid].table(fill=True)

bloptools/test_functions.py

@@ -2,43 +2,70 @@
 
 
 def booth(x1, x2):
+    """
+    The Booth function (https://en.wikipedia.org/wiki/Test_functions_for_optimization)
+    """
     return (x1 + 2 * x2 - 7) ** 2 + (2 * x1 + x2 - 5) ** 2
 
 
 def matyas(x1, x2):
+    """
+    The Matyas function (https://en.wikipedia.org/wiki/Test_functions_for_optimization)
+    """
     return 0.26 * (x1**2 + x2**2) - 0.48 * x1 * x2
 
 
 def himmelblau(x1, x2):
+    """
+    Himmelblau's function (https://en.wikipedia.org/wiki/Himmelblau%27s_function)
+    """
     return (x1**2 + x2 - 11) ** 2 + (x1 + x2**2 - 7) ** 2
 
 
 def constrained_himmelblau(x1, x2):
+    """
+    Himmelblau's function, returns NaN outside the constraint
+    """
     return np.where(x1**2 + x2**2 < 50, himmelblau(x1, x2), np.nan)
 
 
 def skewed_himmelblau(x1, x2):
+    """
+    Himmelblau's function, with skewed coordinates
+    """
     _x1 = 2 * x1 + x2
     _x2 = 0.5 * (x1 - 2 * x2)
 
     return constrained_himmelblau(_x1, _x2)
 
 
 def bukin(x1, x2):
+    """
+    Bukin function N.6 (https://en.wikipedia.org/wiki/Test_functions_for_optimization)
+    """
     return 100 * np.sqrt(np.abs(x2 - 1e-2 * x1**2)) + 0.01 * np.abs(x1)
 
 
 def rastrigin(*x):
+    """
+    The Rastrigin function in arbitrary dimensions (https://en.wikipedia.org/wiki/Rastrigin_function)
+    """
     X = np.c_[x]
     return 10 * X.shape[-1] + (X**2 - 10 * np.cos(2 * np.pi * X)).sum(axis=1)
 
 
 def styblinski_tang(*x):
+    """
+    Styblinski-Tang function in arbitrary dimensions (https://en.wikipedia.org/wiki/Test_functions_for_optimization)
+    """
     X = np.c_[x]
     return 0.5 * (X**4 - 16 * X**2 + 5 * X).sum(axis=1)
 
 
 def ackley(*x):
+    """
+    The Ackley function in arbitrary dimensions (https://en.wikipedia.org/wiki/Ackley_function)
+    """
     X = np.c_[x]
     return (
         -20 * np.exp(-0.2 * np.sqrt(0.5 * (X**2).sum(axis=1)))
@@ -49,10 +76,16 @@ def ackley(*x):
 
 
 def gaussian_beam_waist(x1, x2):
+    """
+    Simulating a misaligned Gaussian beam. The optimum is at (1, 1, 1, 1)
+    """
     return np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2)
 
 
 def himmelblau_digestion(db, uid):
+    """
+    Digests Himmelblau's function into the feedback.
+    """
     products = db[uid].table()
 
     for index, entry in products.iterrows():
@@ -63,7 +96,7 @@ def himmelblau_digestion(db, uid):
 
 def mock_kbs_digestion(db, uid):
     """
-    Simulating a misaligned Gaussian beam. The optimum is at (1, 1, 1, 1)
+    Digests a beam waist and height into the feedback.
     """
 
     products = db[uid].table()

docs/source/usage.rst

@@ -4,30 +4,98 @@ Usage
 
 Working in the Bluesky environment, we need to pass four ingredients to the Bayesian agent:
 
-* ``dofs``: A list of degrees of freedom for the agent.
-* ``dets`` (Optional): A list of detectors to be triggered during acquisition.
+* ``dofs``: A list of degrees of freedom for the agent to optimize over.
+* ``dets``: A list of detectors to be triggered during acquisition.
 * ``tasks``: A list of tasks for the agent to maximize.
 * ``digestion``: A function that processes the output of the acquisition into the task values.
 
-.. code-block:: python
 
-    import bloptools
 
-    dofs = [
-        {"device": some_motor, "limits": (-0.5, 0.5), "kind": "active"},
-        {"device": another_motor, "limits": (-0.5, 0.5), "kind": "active"},
+
+
+Degrees of freedom
+++++++++++++++++++
+
+Degrees of freedom (DOFs) are passed as an iterable of dicts, each containing at least the device and set of limits.
+
+.. code-block:: python
+
+    my_dofs = [
+        {"device": some_motor, "limits": (lower_limit, upper_limit)},
+        {"device": another_motor, "limits": (lower_limit, upper_limit)},
     ]
 
-    tasks = [
-        {"key": "flux", "kind": "maximize", "transform": "log"}
+Here ``some_motor`` and ``another_motor`` are ``ophyd`` objects.
+
+
+Detectors
++++++++++
+
+Detectors are triggered for each input.
+
+.. code-block:: python
+
+    my_dets = [some_detector, some_other_detector]
+
+
+Tasks
++++++
+
+Degrees of freedom (DOFs) are passed as an iterable of dicts, each containing at least the device and set of limits.
+
+.. code-block:: python
+
+    my_tasks = [
+        {"key": "value_to_maximize"}
         ]
 
-    agent = bloptools.bayesian.Agent(
-        dofs=dofs,
-        tasks=tasks,
-        dets=[some_detector, another_detector],
-        digestion=your_digestion_function,
-        db=db,
+
+
+Digestion
++++++++++
+
+The digestion function is how we go from what is spit out by the acquisition to the actual values of the tasks.
+
+.. code-block:: python
+
+    def my_digestion_function(db, uid):
+
+        products = db[uid].table(fill=True) # a pandas DataFrame
+
+        # for each entry, do some
+        for index, entry in products.iterrows():
+
+            raw_output_1 = entry.raw_output_1
+            raw_output_2 = entry.raw_output_2
+
+            entry.loc[index, "value_to_maximize"] = some_fitness_function(raw_output_1, raw_output_2)
+
+        return products
+
+
+
+Building the agent
+++++++++++++++++++
+
+Combining these with a databroker instance will construct an agent.
+
+.. code-block:: python
+
+    import bloptools
+
+    my_agent = bloptools.bayesian.Agent(
+        dofs=my_dofs,
+        dets=my_dets,
+        tasks=my_tasks,
+        digestion=my_digestion_function,
+        db=db, # a databroker instance
     )
 
     RE(agent.initialize("qr", n_init=24))
+
+
+In the example below, the agent will loop over the following steps in each iteration of learning.
+
+#. Find the most interesting point (or points) to sample, and move the degrees of freedom there.
+#. For each point, run an acquisition plan (e.g., trigger and read the detectors).
+#. Digest the results of the acquisition to find the value of the task.