Merge pull request NSLS-II#52 from thomaswmorris/targeting

Targeting cabability for agent objectives

bloptools/bayesian/__init__.py

@@ -1,3 +1,3 @@
 from .agent import *  # noqa F401
-from .devices import *  # noqa F401
-from .objective import *  # noqa F401
+from .dofs import *  # noqa F401
+from .objectives import *  # noqa F401

bloptools/bayesian/acquisition/__init__.py

@@ -1,7 +1,6 @@
 import os
 
 import yaml
-from botorch.acquisition.objective import ScalarizedPosteriorTransform
 
 from . import analytic, monte_carlo
 from .analytic import *  # noqa F401
@@ -38,7 +37,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
     acq_func_name = parse_acq_func_identifier(identifier)
     acq_func_config = config["upper_confidence_bound"]
 
-    if config[acq_func_name]["multitask_only"] and (agent.num_tasks == 1):
+    if config[acq_func_name]["multitask_only"] and (len(agent.objectives) == 1):
         raise ValueError(f'Acquisition function "{acq_func_name}" is only for multi-task optimization problems!')
 
     # there is probably a better way to structure this
@@ -47,7 +46,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
             constraint=agent.constraint,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
-            posterior_transform=ScalarizedPosteriorTransform(weights=agent.objective_weights_torch, offset=0),
+            posterior_transform=agent.targeting_transform,
         )
         acq_func_meta = {"name": acq_func_name, "args": {}}
 
@@ -56,7 +55,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
             constraint=agent.constraint,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
-            posterior_transform=ScalarizedPosteriorTransform(weights=agent.objective_weights_torch, offset=0),
+            posterior_transform=agent.targeting_transform,
         )
         acq_func_meta = {"name": acq_func_name, "args": {}}
 
@@ -65,7 +64,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
             constraint=agent.constraint,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
-            posterior_transform=ScalarizedPosteriorTransform(weights=agent.objective_weights_torch, offset=0),
+            posterior_transform=agent.targeting_transform,
         )
         acq_func_meta = {"name": acq_func_name, "args": {}}
 
@@ -74,7 +73,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
             constraint=agent.constraint,
             model=agent.model,
             best_f=agent.max_scalarized_objective,
-            posterior_transform=ScalarizedPosteriorTransform(weights=agent.objective_weights_torch, offset=0),
+            posterior_transform=agent.targeting_transform,
         )
         acq_func_meta = {"name": acq_func_name, "args": {}}
 
@@ -103,7 +102,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
             constraint=agent.constraint,
             model=agent.model,
             beta=beta,
-            posterior_transform=ScalarizedPosteriorTransform(weights=agent.objective_weights_torch, offset=0),
+            posterior_transform=agent.targeting_transform,
         )
         acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}}
 
@@ -114,7 +113,7 @@ def get_acquisition_function(agent, identifier="qei", return_metadata=True, verb
             constraint=agent.constraint,
             model=agent.model,
             beta=beta,
-            posterior_transform=ScalarizedPosteriorTransform(weights=agent.objective_weights_torch, offset=0),
+            posterior_transform=agent.targeting_transform,
         )
         acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}}
 

bloptools/bayesian/acquisition/analytic.py

@@ -1,58 +1,10 @@
 import math
 
-import bluesky.plan_stubs as bps
-import bluesky.plans as bp
 import numpy as np
 import torch
 from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement, UpperConfidenceBound
 
 
-def list_scan_with_delay(*args, delay=0, **kwargs):
-    "Accepts all the normal 'scan' parameters, plus an optional delay."
-
-    def one_nd_step_with_delay(detectors, step, pos_cache):
-        "This is a copy of bluesky.plan_stubs.one_nd_step with a sleep added."
-        motors = step.keys()
-        yield from bps.move_per_step(step, pos_cache)
-        yield from bps.sleep(delay)
-        yield from bps.trigger_and_read(list(detectors) + list(motors))
-
-    kwargs.setdefault("per_step", one_nd_step_with_delay)
-    uid = yield from bp.list_scan(*args, **kwargs)
-    return uid
-
-
-def default_acquisition_plan(dofs, inputs, dets, **kwargs):
-    delay = kwargs.get("delay", 0)
-    args = []
-    for dof, points in zip(dofs, np.atleast_2d(inputs).T):
-        args.append(dof)
-        args.append(list(points))
-
-    uid = yield from list_scan_with_delay(dets, *args, delay=delay)
-    return uid
-
-
-# def sleepy_acquisition_plan(dofs, inputs, dets):
-
-#     args = []
-#     for dof, points in zip(dofs, np.atleast_2d(inputs).T):
-#         args.append(dof)
-#         args.append(list(points))
-
-#     for point in inputs:
-#         args = []
-#         for dof, value in zip(dofs, point):
-#             args.append(dof)
-#             args.append(value)
-
-#         yield from bps.mv(*args)
-#         yield from bps.count([*dets, *dofs])
-#         yield from bps.sleep(1)
-
-#     return uid
-
-
 class ConstrainedUpperConfidenceBound(UpperConfidenceBound):
     """Upper confidence bound, but scaled by some constraint.
     NOTE: Because the UCB can be negative, we constrain it by adjusting the Gaussian quantile.