From eea262e9dd328ddb9540315b016f89b90aa6264b Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 4 Sep 2023 22:21:25 -0400 Subject: [PATCH 01/43] added constrained acquisition functions --- bloptools/bayesian/__init__.py | 166 ++++++++++++------ bloptools/bayesian/acquisition.py | 36 ++++ bloptools/bayesian/models.py | 19 +- bloptools/functions.py | 4 + bloptools/test_functions.py | 38 +++- bloptools/tests/conftest.py | 24 ++- bloptools/tests/test_acq_funcs.py | 29 +++ bloptools/tests/test_agent.py | 3 +- bloptools/tests/test_bayesian_shadow.py | 4 +- bloptools/tests/test_passive_dofs.py | 2 +- bloptools/{utils.py => utils/__init__.py} | 0 bloptools/utils/prepare_re_env.py | 102 +++++++++++ docs/source/tutorials/hyperparameters.ipynb | 9 +- docs/source/tutorials/introduction.ipynb | 25 ++- .../tutorials/latent-toroid-dimensions.ipynb | 9 +- docs/source/tutorials/multi-task-sirepo.ipynb | 9 +- docs/source/tutorials/passive-dofs.ipynb | 9 +- requirements.txt | 2 +- 18 files changed, 385 insertions(+), 105 deletions(-) create mode 100644 bloptools/functions.py create mode 100644 bloptools/tests/test_acq_funcs.py rename bloptools/{utils.py => utils/__init__.py} (100%) create mode 100644 bloptools/utils/prepare_re_env.py diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index d1d643a..597a64d 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -14,14 +14,21 @@ import pandas as pd import scipy as sp import torch +import os from matplotlib import pyplot as plt from matplotlib.patches import Patch +from botorch.acquisition.objective import ScalarizedPosteriorTransform +from botorch.models.deterministic import GenericDeterministicModel +from botorch.models.model_list_gp_regression import ModelListGP + from .. import utils from . import models from .acquisition import default_acquisition_plan from .digestion import default_digestion_function +os.environ['KMP_DUPLICATE_LIB_OK']='True' + warnings.filterwarnings("ignore", category=botorch.exceptions.warnings.InputDataWarning) mpl.rc("image", cmap="coolwarm") @@ -29,44 +36,47 @@ DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen", "goldenrod"] DEFAULT_COLORMAP = "turbo" - MAX_TEST_INPUTS = 2**11 - TASK_CONFIG = {} +TASK_TRANSFORMS = {"log": lambda x: np.log(x)} + + ACQ_FUNC_CONFIG = { "quasi-random": { "identifiers": ["qr", "quasi-random"], "pretty_name": "Quasi-random", "description": "Sobol-sampled quasi-random points.", }, - "expected_mean": { - "identifiers": ["em", "expected_mean"], - "pretty_name": "Expected mean", - "description": "The expected value at each input.", - }, + # "expected_mean": { + # "identifiers": ["em", "expected_mean"], + # "pretty_name": "Expected mean", + # "description": "The expected value at each input.", + # }, "expected_improvement": { "identifiers": ["ei", "expected_improvement"], "pretty_name": "Expected improvement", "description": r"The expected value of max(f(x) - \nu, 0), where \nu is the current maximum.", }, + "expected_hypervolume_improvement": { + "identifiers": ["ehvi", "expected_hypervolume_improvement"], + "pretty_name": "Expected hypervolume improvement", + "description": r"It's like a big box. How big is the box?", + }, "probability_of_improvement": { "identifiers": ["pi", "probability_of_improvement"], "pretty_name": "Probability of improvement", "description": "The probability that this input improves on the current maximum.", }, - "upper_confidence_bound": { - "identifiers": ["ucb", "upper_confidence_bound"], - "default_args": {"z": 2}, - "pretty_name": "Upper confidence bound", - "description": r"The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).", - }, + # "upper_confidence_bound": { + # "identifiers": ["ucb", "upper_confidence_bound"], + # "default_args": {"z": 2}, + # "pretty_name": "Upper confidence bound", + # "description": r"The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).", + # }, } -TASK_TRANSFORMS = {"log": lambda x: np.log(x)} - - def _validate_and_prepare_dofs(dofs): for dof in dofs: if not isinstance(dof, Mapping): @@ -278,11 +288,11 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): outcome_transform=outcome_transform, ).double() - # this ensures that we have equal weight between task fitness and feasibility fitness - self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( - weights=torch.tensor([*torch.ones(self.n_tasks), self.fitness_variance.sum().sqrt()]).double(), - offset=0, - ) + # # this ensures that we have equal weight between task fitness and feasibility fitness + # self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( + # weights=torch.tensor([*torch.ones(self.n_tasks), self.fitness_variance.sum().sqrt()]).double(), + # offset=0, + # ) dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood( torch.tensor(feasibility).long(), learn_additional_noise=True @@ -309,11 +319,43 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): else: raise RuntimeError("Could not fit model on initialization!") - feasibility_fitness_model = botorch.models.deterministic.GenericDeterministicModel( - f=lambda X: -self.classifier.log_prob(X).square() - ) + self.constraint = GenericDeterministicModel(f=lambda x: self.classifier.probabilities(x)[..., -1].squeeze(-1)) - self.model_list = botorch.models.model.ModelList(*[task["model"] for task in self.tasks], feasibility_fitness_model) + # feasibility_fitness_model = botorch.models.deterministic.GenericDeterministicModel( + # f=lambda X: -self.classifier.log_prob(X).square() + # ) + + # self.model_list = botorch.models.model.ModelList(*[task["model"] for task in self.tasks], feasibility_fitness_model) + + @property + def model(self): + return ModelListGP(*[task["model"] for task in self.tasks]) + + @property + def target_names(self): + return [f'{task["key"]}_fitness' for task in self.tasks] + + @property + def train_inputs(self): + feasibility = ~self.task_fitnesses.isna().any(axis=1) + inputs = self.inputs.loc[feasibility, self._subset_dof_names(mode="on")].values + return torch.tensor(inputs).double() + + @property + def train_targets(self): + feasibility = ~self.task_fitnesses.isna().any(axis=1) + inputs = self.table.loc[feasibility, self.target_names].values + return torch.tensor(inputs).double() + + @property + def class_inputs(self): + inputs = self.inputs.loc[:, self._subset_dof_names(mode="on")].values + return torch.tensor(inputs).double() + + @property + def class_targets(self): + feasibility = ~self.task_fitnesses.isna().any(axis=1) + return torch.tensor(feasibility).double() @property def task_fitnesses(self): @@ -507,47 +549,59 @@ def acq_func_info(self): print("\n\n".join(entries)) def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=False, acq_func_args={}, **kwargs): + if not self._initialized: raise RuntimeError( f'Can\'t construct acquisition function "{acq_func_identifier}" (the agent is not initialized!)' ) if acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_improvement"]["identifiers"]: - acq_func = botorch.acquisition.analytic.LogExpectedImprovement( - self.model_list, - best_f=self.scalarized_fitness.max(), - posterior_transform=self.task_scalarization, + acq_func = acquisition.ConstrainedLogExpectedImprovement( + constraint=self.constraint, + model=self.model, + best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), + posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), **kwargs, ) acq_func_meta = {"name": "expected improvement", "args": {}} elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["probability_of_improvement"]["identifiers"]: - acq_func = botorch.acquisition.analytic.LogProbabilityOfImprovement( - self.model_list, - best_f=self.scalarized_fitness.max(), - posterior_transform=self.task_scalarization, + acq_func = acquisition.ConstrainedLogProbabilityOfImprovement( + constraint=self.constraint, + model=self.model, + best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), + posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), **kwargs, ) acq_func_meta = {"name": "probability of improvement", "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_mean"]["identifiers"]: - acq_func = botorch.acquisition.analytic.UpperConfidenceBound( - self.model_list, - beta=0, - posterior_transform=self.task_scalarization, - **kwargs, - ) - acq_func_meta = {"name": "expected mean"} - - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["upper_confidence_bound"]["identifiers"]: - beta = ACQ_FUNC_CONFIG["upper_confidence_bound"]["default_args"]["z"] ** 2 - acq_func = botorch.acquisition.analytic.UpperConfidenceBound( - self.model_list, - beta=beta, - posterior_transform=self.task_scalarization, - **kwargs, - ) - acq_func_meta = {"name": "upper confidence bound", "args": {"beta": beta}} + elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_hypervolume_improvement"]["identifiers"]: + acq_func = acquisition.qConstrainedNoisyExpectedHypervolumeImprovement( + constraint=self.constraint, + model=self.model, + ref_point=self.train_targets.min(dim=0).values, + X_baseline=self.train_inputs, + prune_baseline=True) + acq_func_meta = {"name": "expected hypervolume improvement", "args": {}} + + # elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_mean"]["identifiers"]: + # acq_func = botorch.acquisition.analytic.UpperConfidenceBound( + # self.model_list, + # beta=0, + # posterior_transform=self.task_scalarization, + # **kwargs, + # ) + # acq_func_meta = {"name": "expected mean"} + + # elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["upper_confidence_bound"]["identifiers"]: + # beta = ACQ_FUNC_CONFIG["upper_confidence_bound"]["default_args"]["z"] ** 2 + # acq_func = botorch.acquisition.analytic.UpperConfidenceBound( + # self.model_list, + # beta=beta, + # posterior_transform=self.task_scalarization, + # **kwargs, + # ) + # acq_func_meta = {"name": "upper confidence bound", "args": {"beta": beta}} else: raise ValueError(f'Unrecognized acquisition acq_func_identifier "{acq_func_identifier}".') @@ -952,13 +1006,13 @@ def _plot_feas_one_dof(self, size=16, lw=1e0): x = self.test_inputs_grid *input_shape, input_dim = x.shape - log_prob = self.classifier.log_prob(x.reshape(-1, 1, input_dim)).reshape(input_shape) + constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) self.feas_ax.scatter(self.inputs.values, ~self.task_fitnesses.isna().any(axis=1), s=size) on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - self.feas_ax.plot(x[..., on_dofs_are_active_mask].squeeze(), log_prob.exp().detach().numpy(), lw=lw) + self.feas_ax.plot(x[..., on_dofs_are_active_mask].squeeze(), constraint.detach().numpy(), lw=lw) self.feas_ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[0]["limits"]) @@ -979,13 +1033,13 @@ def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLO x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) *input_shape, input_dim = x.shape - log_prob = self.classifier.log_prob(x.reshape(-1, 1, input_dim)).reshape(input_shape) + constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) if gridded: self.feas_axes[1].pcolormesh( x[..., 0], x[..., 1], - log_prob.exp().detach().numpy(), + constraint.detach().numpy(), shading=shading, cmap=cmap, vmin=0, @@ -999,7 +1053,7 @@ def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLO self.feas_axes[1].scatter( x.detach().numpy()[..., axes[0]], x.detach().numpy()[..., axes[1]], - c=log_prob.exp().detach().numpy(), + c=constraint.detach().numpy(), ) self.feas_fig.colorbar(data_ax, ax=self.feas_axes[:2], location="bottom", aspect=32, shrink=0.8) diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index 3354d19..16eca12 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -1,3 +1,9 @@ +from botorch.acquisition.objective import ScalarizedPosteriorTransform +from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement +from botorch.acquisition.multi_objective.monte_carlo import ( + qNoisyExpectedHypervolumeImprovement, +) + import bluesky.plans as bp import numpy as np @@ -5,3 +11,33 @@ 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 + + +class ConstrainedLogExpectedImprovement(LogExpectedImprovement): + + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) + self.constraint(x).log() + +class ConstrainedLogProbabilityOfImprovement(LogProbabilityOfImprovement): + + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) + self.constraint(x).log() + +class qConstrainedNoisyExpectedHypervolumeImprovement(qNoisyExpectedHypervolumeImprovement): + + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) * self.constraint(x) + + diff --git a/bloptools/bayesian/models.py b/bloptools/bayesian/models.py index a160703..6f7b72a 100644 --- a/bloptools/bayesian/models.py +++ b/bloptools/bayesian/models.py @@ -21,21 +21,12 @@ def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs) ) -class LatentDirichletClassifier(botorch.models.gp_regression.SingleTaskGP): +class LatentDirichletClassifier(LatentGP): def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs): - super().__init__(train_inputs, train_targets, *args, **kwargs) + super().__init__(train_inputs, train_targets, skew_dims, *args, **kwargs) - self.mean_module = gpytorch.means.ConstantMean() - self.covar_module = kernels.LatentKernel( - num_inputs=train_inputs.shape[-1], - num_outputs=train_targets.shape[-1], - skew_dims=skew_dims, - diag_prior=True, - scale=True, - **kwargs - ) - - def log_prob(self, x, n_samples=256): + def probabilities(self, x, n_samples=256): *input_shape, n_dim = x.shape samples = self.posterior(x.reshape(-1, n_dim)).sample(torch.Size((n_samples,))).exp() - return torch.log((samples / samples.sum(-1, keepdim=True)).mean(0)[:, 1]).reshape(*input_shape, 1) + return (samples / samples.sum(-1, keepdim=True)).mean(0).reshape(*input_shape, -1) + diff --git a/bloptools/functions.py b/bloptools/functions.py new file mode 100644 index 0000000..2df3d49 --- /dev/null +++ b/bloptools/functions.py @@ -0,0 +1,4 @@ +import numpy as np + +def sigmoid(x): + return 1 / (1 + np.exp(-x)) \ No newline at end of file diff --git a/bloptools/test_functions.py b/bloptools/test_functions.py index d61b7f8..54be32d 100644 --- a/bloptools/test_functions.py +++ b/bloptools/test_functions.py @@ -77,19 +77,51 @@ def ackley(*x): def gaussian_beam_waist(x1, x2): """ - Simulating a misaligned Gaussian beam. The optimum is at (1, 1, 1, 1) + Simulating a misaligned Gaussian beam. The optimum is at (1, 1) """ return np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) -def himmelblau_digestion(db, uid): + + +# alpha = [1.0, 1.2, 3.0, 3.2]'; +# A = [10, 3, 17, 3.5, 1.7, 8; +# 0.05, 10, 17, 0.1, 8, 14; +# 3, 3.5, 1.7, 10, 17, 8; +# 17, 8, 0.05, 10, 0.1, 14]; +# P = 10^(-4) * [1312, 1696, 5569, 124, 8283, 5886; +# 2329, 4135, 8307, 3736, 1004, 9991; +# 2348, 1451, 3522, 2883, 3047, 6650; +# 4047, 8828, 8732, 5743, 1091, 381]; + + +def kb_tradeoff_2d(x1, x2): + + width = np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) + d = np.sqrt(x1**2 + x2**2) + flux = np.exp(-0.5 * np.where(d < 5, np.where(d > -5, 0, d + 5), d - 5)**2) + + return width, flux + + +def kb_tradeoff_4d(x1, x2, x3, x4): + + x_width = np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) + y_width = np.sqrt(1 + 0.25 * (x3 - x4) ** 2 + 16 * (x3 + x4 - 2) ** 2) + d = np.sqrt(x1**2 + x2**2 + x3**2 + x4**2) + flux = np.exp(-0.5 * np.where(d < 5, np.where(d > -5, 0, d + 5), d - 5)**2) + + return x_width, y_width, flux + + +def constrained_himmelblau_digestion(db, uid): """ Digests Himmelblau's function into the feedback. """ products = db[uid].table() for index, entry in products.iterrows(): - products.loc[index, "himmelblau"] = himmelblau(entry.x1, entry.x2) + products.loc[index, "himmelblau"] = constrained_himmelblau(entry.x1, entry.x2) return products diff --git a/bloptools/tests/conftest.py b/bloptools/tests/conftest.py index a02f9ca..bd27849 100644 --- a/bloptools/tests/conftest.py +++ b/bloptools/tests/conftest.py @@ -7,10 +7,8 @@ from bluesky.run_engine import RunEngine from databroker import Broker from ophyd.utils import make_dir_tree -from sirepo_bluesky.madx_handler import MADXFileHandler -from sirepo_bluesky.shadow_handler import ShadowFileHandler from sirepo_bluesky.sirepo_bluesky import SirepoBluesky -from sirepo_bluesky.srw_handler import SRWFileHandler + from bloptools.bayesian import Agent @@ -21,12 +19,28 @@ def db(): """Return a data broker""" # MongoDB backend: - db = Broker.named("local") # mongodb backend + db = Broker.named("temp") # mongodb backend + try: + databroker.assets.utils.install_sentinels(db.reg.config, version=1) + except Exception: + pass + + return db + +@pytest.fixture(scope="function") +def db_with_sirepo_bluesky(): + """Return a data broker""" + # MongoDB backend: + db = Broker.named("temp") # mongodb backend try: databroker.assets.utils.install_sentinels(db.reg.config, version=1) except Exception: pass + from sirepo_bluesky.madx_handler import MADXFileHandler + from sirepo_bluesky.shadow_handler import ShadowFileHandler + from sirepo_bluesky.srw_handler import SRWFileHandler + db.reg.register_handler("srw", SRWFileHandler, overwrite=True) db.reg.register_handler("shadow", ShadowFileHandler, overwrite=True) db.reg.register_handler("SIREPO_FLYER", SRWFileHandler, overwrite=True) @@ -71,7 +85,7 @@ def agent(db): agent = Agent( dofs=dofs, tasks=tasks, - digestion=test_functions.himmelblau_digestion, + digestion=test_functions.constrained_himmelblau_digestion, db=db, verbose=True, tolerate_acquisition_errors=False, diff --git a/bloptools/tests/test_acq_funcs.py b/bloptools/tests/test_acq_funcs.py new file mode 100644 index 0000000..9903ec8 --- /dev/null +++ b/bloptools/tests/test_acq_funcs.py @@ -0,0 +1,29 @@ +import pytest + +import bloptools +from bloptools import devices, test_functions +from bloptools.bayesian import Agent + + +@pytest.mark.test_func +def test_acq_funcs_single_task(RE, db): + + dofs = [ + {"device": devices.DOF(name="x1"), "limits": (-8, 8), "kind": "active"}, + {"device": devices.DOF(name="x2"), "limits": (-8, 8), "kind": "active"}, + ] + + tasks = [ + {"key": "himmelblau", "kind": "minimize"}, + ] + + agent = Agent( + dofs=dofs, + tasks=tasks, + digestion=test_functions.constrained_himmelblau_digestion, + db=db, + ) + + RE(agent.initialize("qr", n_init=64)) + RE(agent.learn("ei", n_iter=2)) + RE(agent.learn("pi", n_iter=2)) \ No newline at end of file diff --git a/bloptools/tests/test_agent.py b/bloptools/tests/test_agent.py index 66657ce..5e954df 100644 --- a/bloptools/tests/test_agent.py +++ b/bloptools/tests/test_agent.py @@ -1,5 +1,4 @@ -import os - +import os import pytest diff --git a/bloptools/tests/test_bayesian_shadow.py b/bloptools/tests/test_bayesian_shadow.py index 001f6e8..7f58b29 100644 --- a/bloptools/tests/test_bayesian_shadow.py +++ b/bloptools/tests/test_bayesian_shadow.py @@ -1,5 +1,4 @@ import pytest -from sirepo_bluesky.sirepo_ophyd import create_classes import bloptools from bloptools.experiments.sirepo.tes import w9_digestion @@ -7,6 +6,9 @@ @pytest.mark.shadow def test_bayesian_agent_tes_shadow(RE, db, shadow_tes_simulation): + + from sirepo_bluesky.sirepo_ophyd import create_classes + data, schema = shadow_tes_simulation.auth("shadow", "00000002") classes, objects = create_classes(connection=shadow_tes_simulation) globals().update(**objects) diff --git a/bloptools/tests/test_passive_dofs.py b/bloptools/tests/test_passive_dofs.py index f170a79..7760e38 100644 --- a/bloptools/tests/test_passive_dofs.py +++ b/bloptools/tests/test_passive_dofs.py @@ -20,7 +20,7 @@ def test_passive_dofs(RE, db): agent = Agent( dofs=dofs, tasks=tasks, - digestion=test_functions.himmelblau_digestion, + digestion=test_functions.constrained_himmelblau_digestion, db=db, verbose=True, tolerate_acquisition_errors=False, diff --git a/bloptools/utils.py b/bloptools/utils/__init__.py similarity index 100% rename from bloptools/utils.py rename to bloptools/utils/__init__.py diff --git a/bloptools/utils/prepare_re_env.py b/bloptools/utils/prepare_re_env.py new file mode 100644 index 0000000..86d1d1a --- /dev/null +++ b/bloptools/utils/prepare_re_env.py @@ -0,0 +1,102 @@ +import argparse +import datetime +import json # noqa F401 + +import bluesky.plan_stubs as bps # noqa F401 +import bluesky.plans as bp # noqa F401 +import databroker +import matplotlib.pyplot as plt +import numpy as np # noqa F401 +from bluesky.callbacks import best_effort +from bluesky.run_engine import RunEngine +from databroker import Broker +from ophyd.utils import make_dir_tree + +DEFAULT_DB_TYPE = "local" +DEFAULT_ROOT_DIR = "/tmp/sirepo-bluesky-data" +DEFAULT_ENV_TYPE = "stepper" +DEFAULT_USE_SIREPO = False + + +def re_env(db_type=DEFAULT_DB_TYPE, root_dir=DEFAULT_ROOT_DIR): + RE = RunEngine({}) + bec = best_effort.BestEffortCallback() + RE.subscribe(bec) + + db = Broker.named(db_type) + try: + databroker.assets.utils.install_sentinels(db.reg.config, version=1) + except Exception: + pass + RE.subscribe(db.insert) + + _ = make_dir_tree(datetime.datetime.now().year, base_path=root_dir) + + return dict(RE=RE, db=db, bec=bec) + + +def register_handlers(db, handlers): + for handler_spec, handler_class in handlers.items(): + db.reg.register_handler(handler_spec, handler_class, overwrite=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Prepare bluesky environment") + parser.add_argument( + "-d", + "--db-type", + dest="db_type", + default=DEFAULT_DB_TYPE, + help="Type of databroker ('local', 'temp', etc.)", + ) + parser.add_argument( + "-r", + "--root-dir", + dest="root_dir", + default=DEFAULT_ROOT_DIR, + help="The root dir to create YYYY/MM/DD dir structure.", + ) + + parser.add_argument( + "-s", + "--use-sirepo", + dest="use_sirepo", + default=DEFAULT_USE_SIREPO, + help="The root dir to create YYYY/MM/DD dir structure.", + ) + + env_choices = ["stepper", "flyer"] + parser.add_argument( + "-e", + "--env-type", + dest="env_type", + choices=env_choices, + default=DEFAULT_ENV_TYPE, + help="Type of RE environment.", + ) + + args = parser.parse_args() + kwargs_re = dict(db_type=args.db_type, root_dir=args.root_dir) + ret = re_env(**kwargs_re) + globals().update(**ret) + + if args.use_sirepo: + + from sirepo_bluesky.srw_handler import SRWFileHandler + + if args.env_type == "stepper": + from sirepo_bluesky.shadow_handler import ShadowFileHandler + + handlers = {"srw": SRWFileHandler, "SIREPO_FLYER": SRWFileHandler, "shadow": ShadowFileHandler} + plt.ion() + elif args.env_type == "flyer": + from sirepo_bluesky.madx_handler import MADXFileHandler + + handlers = {"srw": SRWFileHandler, "SIREPO_FLYER": SRWFileHandler, "madx": MADXFileHandler} + bec.disable_plots() # noqa: F821 + else: + raise RuntimeError( + f"Unknown environment type: {args.env_type}.\nAvailable environment types: {env_choices}" + ) + + register_handlers(db, handlers) # noqa: F821 diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb index d37b5c1..4c1900c 100644 --- a/docs/source/tutorials/hyperparameters.ipynb +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -68,7 +68,8 @@ }, "outputs": [], "source": [ - "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", + "from sirepo_bluesky import prepare_re_env\n", + "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", "from bloptools import devices\n", "from bloptools.bayesian import Agent\n", @@ -129,7 +130,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3.11.4 64-bit", "language": "python", "name": "python3" }, @@ -143,11 +144,11 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.11.5" }, "vscode": { "interpreter": { - "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e" } } }, diff --git a/docs/source/tutorials/introduction.ipynb b/docs/source/tutorials/introduction.ipynb index 1d4eef1..1ac8737 100644 --- a/docs/source/tutorials/introduction.ipynb +++ b/docs/source/tutorials/introduction.ipynb @@ -22,10 +22,22 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "22438de8", "metadata": {}, - "outputs": [], + "outputs": [ + { + "ename": "ModuleNotFoundError", + "evalue": "No module named 'numpy'", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mModuleNotFoundError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mmpl\u001b[39;00m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m pyplot \u001b[38;5;28;01mas\u001b[39;00m plt\n", + "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'numpy'" + ] + } + ], "source": [ "import numpy as np\n", "import matplotlib as mpl\n", @@ -118,7 +130,8 @@ }, "outputs": [], "source": [ - "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", + "from sirepo_bluesky import prepare_re_env\n", + "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", "from bloptools.bayesian import Agent\n", "\n", @@ -213,7 +226,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3.10.12 ('bluesky')", "language": "python", "name": "python3" }, @@ -227,11 +240,11 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.11.4" }, "vscode": { "interpreter": { - "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + "hash": "eee21ccc240bdddd7cf04478199e20f7257541e2f592ca1a4d34ebdc0225d742" } } }, diff --git a/docs/source/tutorials/latent-toroid-dimensions.ipynb b/docs/source/tutorials/latent-toroid-dimensions.ipynb index 7814d33..4409ac6 100644 --- a/docs/source/tutorials/latent-toroid-dimensions.ipynb +++ b/docs/source/tutorials/latent-toroid-dimensions.ipynb @@ -20,7 +20,8 @@ }, "outputs": [], "source": [ - "%run -i ../../../examples/prepare_bluesky.py\n", + "from bloptools.utils import prepare_re_env\n", + "%run -i $prepare_re_env.__file__ --db-type=temp\n", "%run -i ../../../examples/prepare_tes_shadow.py" ] }, @@ -90,7 +91,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3.11.5 64-bit", "language": "python", "name": "python3" }, @@ -104,11 +105,11 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.11.5" }, "vscode": { "interpreter": { - "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e" } } }, diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/source/tutorials/multi-task-sirepo.ipynb index b965327..a3ff943 100644 --- a/docs/source/tutorials/multi-task-sirepo.ipynb +++ b/docs/source/tutorials/multi-task-sirepo.ipynb @@ -22,7 +22,8 @@ }, "outputs": [], "source": [ - "%run -i ../../../examples/prepare_bluesky.py\n", + "from bloptools.utils import prepare_re_env\n", + "%run -i $prepare_re_env.__file__ --db-type=temp\n", "%run -i ../../../examples/prepare_tes_shadow.py" ] }, @@ -126,7 +127,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3.11.4 64-bit", "language": "python", "name": "python3" }, @@ -140,11 +141,11 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.11.5" }, "vscode": { "interpreter": { - "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e" } } }, diff --git a/docs/source/tutorials/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb index 6f5adea..7ea26b4 100644 --- a/docs/source/tutorials/passive-dofs.ipynb +++ b/docs/source/tutorials/passive-dofs.ipynb @@ -26,7 +26,8 @@ "metadata": {}, "outputs": [], "source": [ - "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", + "from bloptools.utils import prepare_re_env\n", + "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", "from bloptools import devices, test_functions\n", "from bloptools.bayesian import Agent\n", @@ -65,7 +66,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3.11.4 64-bit", "language": "python", "name": "python3" }, @@ -79,11 +80,11 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.11.5" }, "vscode": { "interpreter": { - "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e" } } }, diff --git a/requirements.txt b/requirements.txt index 0977478..433b4c5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,11 +2,11 @@ bluesky>=1.6.1 botorch databroker gpytorch +h5py matplotlib numpy ophyd ortools scipy -sirepo-bluesky>=0.7.0 tables torch From fac55462e02fa30735da864bf3036c2c8651997d Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 13:02:35 -0400 Subject: [PATCH 02/43] added more acquisition functions --- bloptools/bayesian/__init__.py | 113 +++++++++--------- bloptools/bayesian/acquisition.py | 52 +++++--- bloptools/bayesian/models.py | 3 + bloptools/functions.py | 3 +- bloptools/test_functions.py | 33 +++-- bloptools/tests/conftest.py | 2 +- bloptools/tests/test_acq_funcs.py | 4 +- bloptools/tests/test_agent.py | 3 +- bloptools/tests/test_bayesian_shadow.py | 1 - bloptools/utils/prepare_re_env.py | 5 +- docs/source/tutorials/hyperparameters.ipynb | 1 + docs/source/tutorials/introduction.ipynb | 17 +-- .../tutorials/latent-toroid-dimensions.ipynb | 1 + docs/source/tutorials/multi-task-sirepo.ipynb | 1 + docs/source/tutorials/passive-dofs.ipynb | 1 + 15 files changed, 133 insertions(+), 107 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 597a64d..06ff692 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -1,4 +1,5 @@ import logging +import os import time as ttime import warnings from collections import OrderedDict @@ -14,20 +15,18 @@ import pandas as pd import scipy as sp import torch -import os -from matplotlib import pyplot as plt -from matplotlib.patches import Patch - from botorch.acquisition.objective import ScalarizedPosteriorTransform from botorch.models.deterministic import GenericDeterministicModel from botorch.models.model_list_gp_regression import ModelListGP +from matplotlib import pyplot as plt +from matplotlib.patches import Patch from .. import utils -from . import models +from . import acquisition, models from .acquisition import default_acquisition_plan from .digestion import default_digestion_function -os.environ['KMP_DUPLICATE_LIB_OK']='True' +os.environ["KMP_DUPLICATE_LIB_OK"] = "True" warnings.filterwarnings("ignore", category=botorch.exceptions.warnings.InputDataWarning) @@ -49,34 +48,40 @@ "pretty_name": "Quasi-random", "description": "Sobol-sampled quasi-random points.", }, - # "expected_mean": { - # "identifiers": ["em", "expected_mean"], - # "pretty_name": "Expected mean", - # "description": "The expected value at each input.", - # }, + "expected_mean": { + "identifiers": ["em", "expected_mean"], + "pretty_name": "Expected mean", + "description": "The expected value at each input.", + }, "expected_improvement": { "identifiers": ["ei", "expected_improvement"], "pretty_name": "Expected improvement", "description": r"The expected value of max(f(x) - \nu, 0), where \nu is the current maximum.", }, - "expected_hypervolume_improvement": { - "identifiers": ["ehvi", "expected_hypervolume_improvement"], - "pretty_name": "Expected hypervolume improvement", + "noisy_expected_hypervolume_improvement": { + "identifiers": ["nehvi", "noisy_expected_hypervolume_improvement"], + "pretty_name": "Noisy expected hypervolume improvement", "description": r"It's like a big box. How big is the box?", }, + "lower_bound_max_value_entropy": { + "identifiers": ["lbmve", "lbmes", "gibbon", "lower_bound_max_value_entropy"], + "pretty_name": "Lower bound max value entropy", + "description": r"Max entropy search, basically", + }, "probability_of_improvement": { "identifiers": ["pi", "probability_of_improvement"], "pretty_name": "Probability of improvement", "description": "The probability that this input improves on the current maximum.", }, - # "upper_confidence_bound": { - # "identifiers": ["ucb", "upper_confidence_bound"], - # "default_args": {"z": 2}, - # "pretty_name": "Upper confidence bound", - # "description": r"The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).", - # }, + "upper_confidence_bound": { + "identifiers": ["ucb", "upper_confidence_bound"], + "default_args": {"beta": 4}, + "pretty_name": "Upper confidence bound", + "description": r"The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).", + }, } + def _validate_and_prepare_dofs(dofs): for dof in dofs: if not isinstance(dof, Mapping): @@ -321,15 +326,9 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): self.constraint = GenericDeterministicModel(f=lambda x: self.classifier.probabilities(x)[..., -1].squeeze(-1)) - # feasibility_fitness_model = botorch.models.deterministic.GenericDeterministicModel( - # f=lambda X: -self.classifier.log_prob(X).square() - # ) - - # self.model_list = botorch.models.model.ModelList(*[task["model"] for task in self.tasks], feasibility_fitness_model) - @property def model(self): - return ModelListGP(*[task["model"] for task in self.tasks]) + return ModelListGP(*[task["model"] for task in self.tasks]) if self.n_tasks > 1 else self.tasks[0]["model"] @property def target_names(self): @@ -548,8 +547,7 @@ def acq_func_info(self): print("\n\n".join(entries)) - def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=False, acq_func_args={}, **kwargs): - + def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=False, **acq_func_kwargs): if not self._initialized: raise RuntimeError( f'Can\'t construct acquisition function "{acq_func_identifier}" (the agent is not initialized!)' @@ -561,9 +559,8 @@ def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=Fal model=self.model, best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), - **kwargs, ) - acq_func_meta = {"name": "expected improvement", "args": {}} + acq_func_meta = {"name": "expected_improvement", "args": {}} elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["probability_of_improvement"]["identifiers"]: acq_func = acquisition.ConstrainedLogProbabilityOfImprovement( @@ -571,37 +568,43 @@ def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=Fal model=self.model, best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), - **kwargs, ) - acq_func_meta = {"name": "probability of improvement", "args": {}} + acq_func_meta = {"name": "probability_of_improvement", "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_hypervolume_improvement"]["identifiers"]: + elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["lower_bound_max_value_entropy"]["identifiers"]: + acq_func = acquisition.qConstrainedLowerBoundMaxValueEntropy( + constraint=self.constraint, + model=self.model, + candidate_set=self.test_inputs(n=1024).squeeze(1), + ) + acq_func_meta = {"name": "lower_bound_max_value_entropy", "args": {}} + + elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["noisy_expected_hypervolume_improvement"]["identifiers"]: acq_func = acquisition.qConstrainedNoisyExpectedHypervolumeImprovement( constraint=self.constraint, model=self.model, ref_point=self.train_targets.min(dim=0).values, X_baseline=self.train_inputs, - prune_baseline=True) - acq_func_meta = {"name": "expected hypervolume improvement", "args": {}} - - # elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_mean"]["identifiers"]: - # acq_func = botorch.acquisition.analytic.UpperConfidenceBound( - # self.model_list, - # beta=0, - # posterior_transform=self.task_scalarization, - # **kwargs, - # ) - # acq_func_meta = {"name": "expected mean"} - - # elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["upper_confidence_bound"]["identifiers"]: - # beta = ACQ_FUNC_CONFIG["upper_confidence_bound"]["default_args"]["z"] ** 2 - # acq_func = botorch.acquisition.analytic.UpperConfidenceBound( - # self.model_list, - # beta=beta, - # posterior_transform=self.task_scalarization, - # **kwargs, - # ) - # acq_func_meta = {"name": "upper confidence bound", "args": {"beta": beta}} + prune_baseline=True, + ) + acq_func_meta = {"name": "noisy_expected_hypervolume_improvement", "args": {}} + + elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["upper_confidence_bound"]["identifiers"]: + + config = ACQ_FUNC_CONFIG["upper_confidence_bound"] + beta = acq_func_kwargs.get("beta", config["default_args"]["beta"]) + + acq_func = acquisition.ConstrainedUpperConfidenceBound( + constraint=self.constraint, + model=self.model, + beta=beta, + posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), + ) + acq_func_meta = {"name": "upper_confidence_bound", "args": {"beta": beta}} + + elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_mean"]["identifiers"]: + acq_func = self.get_acquisition_function(acq_func_identifier="ucb", beta=0, return_metadata=False) + acq_func_meta = {"name": "expected_mean"} else: raise ValueError(f'Unrecognized acquisition acq_func_identifier "{acq_func_identifier}".') diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index 16eca12..253cce6 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -1,43 +1,67 @@ -from botorch.acquisition.objective import ScalarizedPosteriorTransform -from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement -from botorch.acquisition.multi_objective.monte_carlo import ( - qNoisyExpectedHypervolumeImprovement, -) - import bluesky.plans as bp import numpy as np +import math +import torch + +from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement, UpperConfidenceBound +from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement +from botorch.acquisition.max_value_entropy_search import qLowerBoundMaxValueEntropy 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 -class ConstrainedLogExpectedImprovement(LogExpectedImprovement): +class ConstrainedUpperConfidenceBound(UpperConfidenceBound): + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + + mean, sigma = self._mean_and_sigma(x) + + p_eff = 0.5 * (1 + torch.special.erf(self.beta.sqrt()/math.sqrt(2))) * torch.clamp(self.constraint(x), min=1e-6) + + return (mean if self.maximize else -mean) + sigma * np.sqrt(2) * torch.special.erfinv(2*p_eff-1) + + + +class ConstrainedLogExpectedImprovement(LogExpectedImprovement): def __init__(self, constraint, *args, **kwargs): super().__init__(*args, **kwargs) self.constraint = constraint - - def forward(self, x): + + def forward(self, x): return super().forward(x) + self.constraint(x).log() -class ConstrainedLogProbabilityOfImprovement(LogProbabilityOfImprovement): +class ConstrainedLogProbabilityOfImprovement(LogProbabilityOfImprovement): def __init__(self, constraint, *args, **kwargs): super().__init__(*args, **kwargs) self.constraint = constraint - - def forward(self, x): + + def forward(self, x): return super().forward(x) + self.constraint(x).log() + class qConstrainedNoisyExpectedHypervolumeImprovement(qNoisyExpectedHypervolumeImprovement): + def __init__(self, constraint, *args, **kwargs): + super().__init__(*args, **kwargs) + self.constraint = constraint + + def forward(self, x): + return super().forward(x) * self.constraint(x) + +class qConstrainedLowerBoundMaxValueEntropy(qLowerBoundMaxValueEntropy): def __init__(self, constraint, *args, **kwargs): super().__init__(*args, **kwargs) self.constraint = constraint - - def forward(self, x): + + def forward(self, x): return super().forward(x) * self.constraint(x) diff --git a/bloptools/bayesian/models.py b/bloptools/bayesian/models.py index 6f7b72a..2c11473 100644 --- a/bloptools/bayesian/models.py +++ b/bloptools/bayesian/models.py @@ -26,6 +26,9 @@ def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs) super().__init__(train_inputs, train_targets, skew_dims, *args, **kwargs) def probabilities(self, x, n_samples=256): + """ + Takes in a (..., m) dimension tensor and returns a (..., n_classes) tensor + """ *input_shape, n_dim = x.shape samples = self.posterior(x.reshape(-1, n_dim)).sample(torch.Size((n_samples,))).exp() return (samples / samples.sum(-1, keepdim=True)).mean(0).reshape(*input_shape, -1) diff --git a/bloptools/functions.py b/bloptools/functions.py index 2df3d49..2d7eff3 100644 --- a/bloptools/functions.py +++ b/bloptools/functions.py @@ -1,4 +1,5 @@ import numpy as np + def sigmoid(x): - return 1 / (1 + np.exp(-x)) \ No newline at end of file + return 1 / (1 + np.exp(-x)) diff --git a/bloptools/test_functions.py b/bloptools/test_functions.py index 54be32d..34d139d 100644 --- a/bloptools/test_functions.py +++ b/bloptools/test_functions.py @@ -84,32 +84,39 @@ def gaussian_beam_waist(x1, x2): -# alpha = [1.0, 1.2, 3.0, 3.2]'; -# A = [10, 3, 17, 3.5, 1.7, 8; -# 0.05, 10, 17, 0.1, 8, 14; -# 3, 3.5, 1.7, 10, 17, 8; -# 17, 8, 0.05, 10, 0.1, 14]; -# P = 10^(-4) * [1312, 1696, 5569, 124, 8283, 5886; -# 2329, 4135, 8307, 3736, 1004, 9991; -# 2348, 1451, 3522, 2883, 3047, 6650; -# 4047, 8828, 8732, 5743, 1091, 381]; -def kb_tradeoff_2d(x1, x2): +def hartmann6(*x): + + X = np.c_[x] + + alpha = np.array([1.0, 1.2, 3.0, 3.2]) + A = np.array([[10, 3, 17, 3.5, 1.7, 8], + [0.05, 10, 17, 0.1, 8, 14], + [3, 3.5, 1.7, 10, 17, 8], + [17, 8, 0.05, 10, 0.1, 14]]) + + P = 1e-4 * np.array([[1312, 1696, 5569, 124, 8283, 5886], + [2329, 4135, 8307, 3736, 1004, 9991], + [2348, 1451, 3522, 2883, 3047, 6650], + [4047, 8828, 8732, 5743, 1091, 381]]) + + return - (alpha * np.exp(-(A * np.square(X - P)).sum(axis=1))).sum() + +def kb_tradeoff_2d(x1, x2): width = np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) d = np.sqrt(x1**2 + x2**2) - flux = np.exp(-0.5 * np.where(d < 5, np.where(d > -5, 0, d + 5), d - 5)**2) + flux = np.exp(-0.5 * np.where(d < 5, np.where(d > -5, 0, d + 5), d - 5) ** 2) return width, flux def kb_tradeoff_4d(x1, x2, x3, x4): - x_width = np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) y_width = np.sqrt(1 + 0.25 * (x3 - x4) ** 2 + 16 * (x3 + x4 - 2) ** 2) d = np.sqrt(x1**2 + x2**2 + x3**2 + x4**2) - flux = np.exp(-0.5 * np.where(d < 5, np.where(d > -5, 0, d + 5), d - 5)**2) + flux = np.exp(-0.5 * np.where(d < 5, np.where(d > -5, 0, d + 5), d - 5) ** 2) return x_width, y_width, flux diff --git a/bloptools/tests/conftest.py b/bloptools/tests/conftest.py index bd27849..5e20d80 100644 --- a/bloptools/tests/conftest.py +++ b/bloptools/tests/conftest.py @@ -9,7 +9,6 @@ from ophyd.utils import make_dir_tree from sirepo_bluesky.sirepo_bluesky import SirepoBluesky - from bloptools.bayesian import Agent from .. import devices, test_functions @@ -27,6 +26,7 @@ def db(): return db + @pytest.fixture(scope="function") def db_with_sirepo_bluesky(): """Return a data broker""" diff --git a/bloptools/tests/test_acq_funcs.py b/bloptools/tests/test_acq_funcs.py index 9903ec8..0dcead6 100644 --- a/bloptools/tests/test_acq_funcs.py +++ b/bloptools/tests/test_acq_funcs.py @@ -1,13 +1,11 @@ import pytest -import bloptools from bloptools import devices, test_functions from bloptools.bayesian import Agent @pytest.mark.test_func def test_acq_funcs_single_task(RE, db): - dofs = [ {"device": devices.DOF(name="x1"), "limits": (-8, 8), "kind": "active"}, {"device": devices.DOF(name="x2"), "limits": (-8, 8), "kind": "active"}, @@ -26,4 +24,4 @@ def test_acq_funcs_single_task(RE, db): RE(agent.initialize("qr", n_init=64)) RE(agent.learn("ei", n_iter=2)) - RE(agent.learn("pi", n_iter=2)) \ No newline at end of file + RE(agent.learn("pi", n_iter=2)) diff --git a/bloptools/tests/test_agent.py b/bloptools/tests/test_agent.py index 5e954df..66657ce 100644 --- a/bloptools/tests/test_agent.py +++ b/bloptools/tests/test_agent.py @@ -1,4 +1,5 @@ -import os +import os + import pytest diff --git a/bloptools/tests/test_bayesian_shadow.py b/bloptools/tests/test_bayesian_shadow.py index 7f58b29..ccb66e2 100644 --- a/bloptools/tests/test_bayesian_shadow.py +++ b/bloptools/tests/test_bayesian_shadow.py @@ -6,7 +6,6 @@ @pytest.mark.shadow def test_bayesian_agent_tes_shadow(RE, db, shadow_tes_simulation): - from sirepo_bluesky.sirepo_ophyd import create_classes data, schema = shadow_tes_simulation.auth("shadow", "00000002") diff --git a/bloptools/utils/prepare_re_env.py b/bloptools/utils/prepare_re_env.py index 86d1d1a..1a7950b 100644 --- a/bloptools/utils/prepare_re_env.py +++ b/bloptools/utils/prepare_re_env.py @@ -81,7 +81,6 @@ def register_handlers(db, handlers): globals().update(**ret) if args.use_sirepo: - from sirepo_bluesky.srw_handler import SRWFileHandler if args.env_type == "stepper": @@ -95,8 +94,6 @@ def register_handlers(db, handlers): handlers = {"srw": SRWFileHandler, "SIREPO_FLYER": SRWFileHandler, "madx": MADXFileHandler} bec.disable_plots() # noqa: F821 else: - raise RuntimeError( - f"Unknown environment type: {args.env_type}.\nAvailable environment types: {env_choices}" - ) + raise RuntimeError(f"Unknown environment type: {args.env_type}.\nAvailable environment types: {env_choices}") register_handlers(db, handlers) # noqa: F821 diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb index 4c1900c..0421d8b 100644 --- a/docs/source/tutorials/hyperparameters.ipynb +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -69,6 +69,7 @@ "outputs": [], "source": [ "from sirepo_bluesky import prepare_re_env\n", + "\n", "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", "from bloptools import devices\n", diff --git a/docs/source/tutorials/introduction.ipynb b/docs/source/tutorials/introduction.ipynb index 1ac8737..a307fff 100644 --- a/docs/source/tutorials/introduction.ipynb +++ b/docs/source/tutorials/introduction.ipynb @@ -22,22 +22,10 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "22438de8", "metadata": {}, - "outputs": [ - { - "ename": "ModuleNotFoundError", - "evalue": "No module named 'numpy'", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mModuleNotFoundError\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[1], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[1;32m 2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mmpl\u001b[39;00m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m pyplot \u001b[38;5;28;01mas\u001b[39;00m plt\n", - "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'numpy'" - ] - } - ], + "outputs": [], "source": [ "import numpy as np\n", "import matplotlib as mpl\n", @@ -131,6 +119,7 @@ "outputs": [], "source": [ "from sirepo_bluesky import prepare_re_env\n", + "\n", "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", "from bloptools.bayesian import Agent\n", diff --git a/docs/source/tutorials/latent-toroid-dimensions.ipynb b/docs/source/tutorials/latent-toroid-dimensions.ipynb index 4409ac6..090bbc1 100644 --- a/docs/source/tutorials/latent-toroid-dimensions.ipynb +++ b/docs/source/tutorials/latent-toroid-dimensions.ipynb @@ -21,6 +21,7 @@ "outputs": [], "source": [ "from bloptools.utils import prepare_re_env\n", + "\n", "%run -i $prepare_re_env.__file__ --db-type=temp\n", "%run -i ../../../examples/prepare_tes_shadow.py" ] diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/source/tutorials/multi-task-sirepo.ipynb index a3ff943..5b61c5d 100644 --- a/docs/source/tutorials/multi-task-sirepo.ipynb +++ b/docs/source/tutorials/multi-task-sirepo.ipynb @@ -23,6 +23,7 @@ "outputs": [], "source": [ "from bloptools.utils import prepare_re_env\n", + "\n", "%run -i $prepare_re_env.__file__ --db-type=temp\n", "%run -i ../../../examples/prepare_tes_shadow.py" ] diff --git a/docs/source/tutorials/passive-dofs.ipynb b/docs/source/tutorials/passive-dofs.ipynb index 7ea26b4..2ce4245 100644 --- a/docs/source/tutorials/passive-dofs.ipynb +++ b/docs/source/tutorials/passive-dofs.ipynb @@ -27,6 +27,7 @@ "outputs": [], "source": [ "from bloptools.utils import prepare_re_env\n", + "\n", "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", "from bloptools import devices, test_functions\n", From 6ae35442a0eab88a204b4097f6d7092a14166adf Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 13:03:05 -0400 Subject: [PATCH 03/43] added sirepo-bluesky to deps --- requirements.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 433b4c5..eecfb9f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,5 +8,6 @@ numpy ophyd ortools scipy +sirepo-bluesky tables -torch +torch \ No newline at end of file From 57ccca64d588e0b739d13af5aad6db184b502ca6 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 14:17:51 -0400 Subject: [PATCH 04/43] fixed tests --- bloptools/bayesian/__init__.py | 3 +- bloptools/tests/conftest.py | 20 ++----------- bloptools/tests/test_bayesian_shadow.py | 9 +++--- docs/source/tutorials/multi-task-sirepo.ipynb | 29 ++++++++++++++++--- requirements.txt | 4 ++- 5 files changed, 36 insertions(+), 29 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 06ff692..f8d5359 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -184,7 +184,7 @@ def __init__( self.db = db self.verbose = kwargs.get("verbose", False) - self.allow_acquisition_errors = kwargs.get("allow_acquisition_errors", True) + self.allow_acquisition_errors = kwargs.get("allow_acquisition_errors", False) self.initialization = kwargs.get("initialization", None) self.acquisition_plan = kwargs.get("acquisition_plan", default_acquisition_plan) self.digestion = kwargs.get("digestion", default_digestion_function) @@ -590,7 +590,6 @@ def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=Fal acq_func_meta = {"name": "noisy_expected_hypervolume_improvement", "args": {}} elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["upper_confidence_bound"]["identifiers"]: - config = ACQ_FUNC_CONFIG["upper_confidence_bound"] beta = acq_func_kwargs.get("beta", config["default_args"]["beta"]) diff --git a/bloptools/tests/conftest.py b/bloptools/tests/conftest.py index 5e20d80..c1226ed 100644 --- a/bloptools/tests/conftest.py +++ b/bloptools/tests/conftest.py @@ -7,7 +7,10 @@ from bluesky.run_engine import RunEngine from databroker import Broker from ophyd.utils import make_dir_tree +from sirepo_bluesky.madx_handler import MADXFileHandler +from sirepo_bluesky.shadow_handler import ShadowFileHandler from sirepo_bluesky.sirepo_bluesky import SirepoBluesky +from sirepo_bluesky.srw_handler import SRWFileHandler from bloptools.bayesian import Agent @@ -24,23 +27,6 @@ def db(): except Exception: pass - return db - - -@pytest.fixture(scope="function") -def db_with_sirepo_bluesky(): - """Return a data broker""" - # MongoDB backend: - db = Broker.named("temp") # mongodb backend - try: - databroker.assets.utils.install_sentinels(db.reg.config, version=1) - except Exception: - pass - - from sirepo_bluesky.madx_handler import MADXFileHandler - from sirepo_bluesky.shadow_handler import ShadowFileHandler - from sirepo_bluesky.srw_handler import SRWFileHandler - db.reg.register_handler("srw", SRWFileHandler, overwrite=True) db.reg.register_handler("shadow", ShadowFileHandler, overwrite=True) db.reg.register_handler("SIREPO_FLYER", SRWFileHandler, overwrite=True) diff --git a/bloptools/tests/test_bayesian_shadow.py b/bloptools/tests/test_bayesian_shadow.py index ccb66e2..aa63e97 100644 --- a/bloptools/tests/test_bayesian_shadow.py +++ b/bloptools/tests/test_bayesian_shadow.py @@ -17,13 +17,13 @@ def test_bayesian_agent_tes_shadow(RE, db, shadow_tes_simulation): dofs = [ {"device": kbv.x_rot, "limits": (-0.1, 0.1), "kind": "active"}, - {"device": kbv.offz, "limits": (-0.5, 0.5), "kind": "active"}, + {"device": kbh.x_rot, "limits": (-0.1, 0.1), "kind": "active"}, ] tasks = [ - {"key": "flux", "kind": "maximize"}, - {"key": "w9_fwhm_x", "kind": "minimize"}, - {"key": "w9_fwhm_y", "kind": "minimize"}, + {"key": "flux", "kind": "maximize", "transform": "log"}, + {"key": "w9_fwhm_x", "kind": "minimize", "transform": "log"}, + {"key": "w9_fwhm_y", "kind": "minimize", "transform": "log"}, ] agent = bloptools.bayesian.Agent( @@ -35,7 +35,6 @@ def test_bayesian_agent_tes_shadow(RE, db, shadow_tes_simulation): ) RE(agent.initialize("qr", n_init=4)) - RE(agent.learn("ei", n_iter=2)) RE(agent.learn("pi", n_iter=2)) diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/source/tutorials/multi-task-sirepo.ipynb index 5b61c5d..3f33cf0 100644 --- a/docs/source/tutorials/multi-task-sirepo.ipynb +++ b/docs/source/tutorials/multi-task-sirepo.ipynb @@ -31,11 +31,22 @@ { "cell_type": "code", "execution_count": null, - "id": "8672bf9e", + "id": "fff28b50", "metadata": {}, "outputs": [], "source": [ - "toroid.offz" + "from bloptools.tests import conftest" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aa1d1473", + "metadata": {}, + "outputs": [], + "source": [ + "db = conftest.db\n", + "RE = conftest.RE()" ] }, { @@ -69,7 +80,17 @@ " db=db,\n", ")\n", "\n", - "RE(agent.initialize(\"qr\", n_init=16))" + "RE(agent.initialize(\"qr\", n_init=4))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2409e012", + "metadata": {}, + "outputs": [], + "source": [ + "RE(agent.learn(\"ei\"))" ] }, { @@ -142,7 +163,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.5" + "version": "3.10.12" }, "vscode": { "interpreter": { diff --git a/requirements.txt b/requirements.txt index eecfb9f..f5bf1cb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,8 @@ numpy ophyd ortools scipy +shadow3 sirepo-bluesky +srwpy tables -torch \ No newline at end of file +torch From 66c784ed7eb2d837d5c99eae78929ab94c4546c9 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 14:22:16 -0400 Subject: [PATCH 05/43] choosing acquisition functions --- bloptools/bayesian/__init__.py | 104 +++++++++++------------------- bloptools/bayesian/acquisition.py | 62 +++++++++++++++--- 2 files changed, 90 insertions(+), 76 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index f8d5359..7ef6abe 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -23,7 +23,7 @@ from .. import utils from . import acquisition, models -from .acquisition import default_acquisition_plan +from .acquisition import ACQ_FUNC_CONFIG, default_acquisition_plan from .digestion import default_digestion_function os.environ["KMP_DUPLICATE_LIB_OK"] = "True" @@ -42,46 +42,6 @@ TASK_TRANSFORMS = {"log": lambda x: np.log(x)} -ACQ_FUNC_CONFIG = { - "quasi-random": { - "identifiers": ["qr", "quasi-random"], - "pretty_name": "Quasi-random", - "description": "Sobol-sampled quasi-random points.", - }, - "expected_mean": { - "identifiers": ["em", "expected_mean"], - "pretty_name": "Expected mean", - "description": "The expected value at each input.", - }, - "expected_improvement": { - "identifiers": ["ei", "expected_improvement"], - "pretty_name": "Expected improvement", - "description": r"The expected value of max(f(x) - \nu, 0), where \nu is the current maximum.", - }, - "noisy_expected_hypervolume_improvement": { - "identifiers": ["nehvi", "noisy_expected_hypervolume_improvement"], - "pretty_name": "Noisy expected hypervolume improvement", - "description": r"It's like a big box. How big is the box?", - }, - "lower_bound_max_value_entropy": { - "identifiers": ["lbmve", "lbmes", "gibbon", "lower_bound_max_value_entropy"], - "pretty_name": "Lower bound max value entropy", - "description": r"Max entropy search, basically", - }, - "probability_of_improvement": { - "identifiers": ["pi", "probability_of_improvement"], - "pretty_name": "Probability of improvement", - "description": "The probability that this input improves on the current maximum.", - }, - "upper_confidence_bound": { - "identifiers": ["ucb", "upper_confidence_bound"], - "default_args": {"beta": 4}, - "pretty_name": "Upper confidence bound", - "description": r"The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).", - }, -} - - def _validate_and_prepare_dofs(dofs): for dof in dofs: if not isinstance(dof, Mapping): @@ -184,7 +144,7 @@ def __init__( self.db = db self.verbose = kwargs.get("verbose", False) - self.allow_acquisition_errors = kwargs.get("allow_acquisition_errors", False) + 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_function) @@ -295,7 +255,7 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): # # this ensures that we have equal weight between task fitness and feasibility fitness # self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( - # weights=torch.tensor([*torch.ones(self.n_tasks), self.fitness_variance.sum().sqrt()]).double(), + # weights=torch.tensor([*torch.ones(self.num_tasks), self.fitness_variance.sum().sqrt()]).double(), # offset=0, # ) @@ -328,7 +288,7 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): @property def model(self): - return ModelListGP(*[task["model"] for task in self.tasks]) if self.n_tasks > 1 else self.tasks[0]["model"] + return ModelListGP(*[task["model"] for task in self.tasks]) if self.num_tasks > 1 else self.tasks[0]["model"] @property def target_names(self): @@ -438,7 +398,7 @@ def _acq_func_bounds(self): ).T @property - def n_tasks(self): + def num_tasks(self): return len(self.tasks) @property @@ -548,38 +508,53 @@ def acq_func_info(self): print("\n\n".join(entries)) def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=False, **acq_func_kwargs): + """ + Generates an acquisition function from a supplied identifier. + """ + if not self._initialized: raise RuntimeError( f'Can\'t construct acquisition function "{acq_func_identifier}" (the agent is not initialized!)' ) - if acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_improvement"]["identifiers"]: + acq_func_name = None + for _acq_func_name in ACQ_FUNC_CONFIG.keys(): + if acq_func_identifier.lower() in ACQ_FUNC_CONFIG[_acq_func_name]["identifiers"]: + acq_func_name = _acq_func_name + + if acq_func_name is None: + raise ValueError(f'Unrecognized acquisition function "{acq_func_identifier}".') + + if ACQ_FUNC_CONFIG[acq_func_name]["multitask_only"] and (self.num_tasks == 1): + raise ValueError(f'Acquisition function "{acq_func_name}" is only for multi-task optimization problems!') + + if acq_func_name == "expected_improvement": acq_func = acquisition.ConstrainedLogExpectedImprovement( constraint=self.constraint, model=self.model, best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), ) - acq_func_meta = {"name": "expected_improvement", "args": {}} + acq_func_meta = {"name": acq_func_name, "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["probability_of_improvement"]["identifiers"]: + elif acq_func_name == "probability_of_improvement": acq_func = acquisition.ConstrainedLogProbabilityOfImprovement( constraint=self.constraint, model=self.model, best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), ) - acq_func_meta = {"name": "probability_of_improvement", "args": {}} + acq_func_meta = {"name": acq_func_name, "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["lower_bound_max_value_entropy"]["identifiers"]: + elif acq_func_name == "lower_bound_max_value_entropy": acq_func = acquisition.qConstrainedLowerBoundMaxValueEntropy( constraint=self.constraint, model=self.model, candidate_set=self.test_inputs(n=1024).squeeze(1), ) - acq_func_meta = {"name": "lower_bound_max_value_entropy", "args": {}} + acq_func_meta = {"name": acq_func_name, "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["noisy_expected_hypervolume_improvement"]["identifiers"]: + elif acq_func_name == "noisy_expected_hypervolume_improvement": acq_func = acquisition.qConstrainedNoisyExpectedHypervolumeImprovement( constraint=self.constraint, model=self.model, @@ -587,9 +562,9 @@ def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=Fal X_baseline=self.train_inputs, prune_baseline=True, ) - acq_func_meta = {"name": "noisy_expected_hypervolume_improvement", "args": {}} + acq_func_meta = {"name": acq_func_name, "args": {}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["upper_confidence_bound"]["identifiers"]: + elif acq_func_name == "upper_confidence_bound": config = ACQ_FUNC_CONFIG["upper_confidence_bound"] beta = acq_func_kwargs.get("beta", config["default_args"]["beta"]) @@ -599,14 +574,11 @@ def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=Fal beta=beta, posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), ) - acq_func_meta = {"name": "upper_confidence_bound", "args": {"beta": beta}} + acq_func_meta = {"name": acq_func_name, "args": {"beta": beta}} - elif acq_func_identifier.lower() in ACQ_FUNC_CONFIG["expected_mean"]["identifiers"]: + elif acq_func_name == "expected_mean": acq_func = self.get_acquisition_function(acq_func_identifier="ucb", beta=0, return_metadata=False) - acq_func_meta = {"name": "expected_mean"} - - else: - raise ValueError(f'Unrecognized acquisition acq_func_identifier "{acq_func_identifier}".') + acq_func_meta = {"name": acq_func_name, "args": {}} return (acq_func, acq_func_meta) if return_metadata else acq_func @@ -777,9 +749,9 @@ def plot_tasks(self, **kwargs): def _plot_tasks_one_dof(self, size=16, lw=1e0): self.task_fig, self.task_axes = plt.subplots( - self.n_tasks, + self.num_tasks, 1, - figsize=(6, 4 * self.n_tasks), + figsize=(6, 4 * self.num_tasks), sharex=True, constrained_layout=True, ) @@ -822,9 +794,9 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL gridded = self._len_subset_dofs(kind="active", mode="on") == 2 self.task_fig, self.task_axes = plt.subplots( - self.n_tasks, + self.num_tasks, 3, - figsize=(10, 4 * self.n_tasks), + figsize=(10, 4 * self.num_tasks), sharex=True, sharey=True, constrained_layout=True, @@ -1087,9 +1059,9 @@ def plot_history(self, x_key="index", show_all_tasks=False): num_task_plots = 1 if show_all_tasks: - num_task_plots = self.n_tasks + 1 + num_task_plots = self.num_tasks + 1 - self.n_tasks + 1 if self.n_tasks > 1 else 1 + self.num_tasks + 1 if self.num_tasks > 1 else 1 hist_fig, hist_axes = plt.subplots( num_task_plots, 1, figsize=(6, 4 * num_task_plots), sharex=True, constrained_layout=True, dpi=200 diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index 253cce6..73ae31e 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -1,18 +1,64 @@ +import math + import bluesky.plans as bp import numpy as np - -import math import torch - from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement, UpperConfidenceBound -from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement from botorch.acquisition.max_value_entropy_search import qLowerBoundMaxValueEntropy +from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement + 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 +ACQ_FUNC_CONFIG = { + "quasi-random": { + "identifiers": ["qr", "quasi-random"], + "pretty_name": "Quasi-random", + "description": "Sobol-sampled quasi-random points.", + "multitask_only": False, + }, + "expected_mean": { + "identifiers": ["em", "expected_mean"], + "pretty_name": "Expected mean", + "multitask_only": False, + "description": "The expected value at each input.", + }, + "expected_improvement": { + "identifiers": ["ei", "expected_improvement"], + "pretty_name": "Expected improvement", + "multitask_only": False, + "description": r"The expected value of max(f(x) - \nu, 0), where \nu is the current maximum.", + }, + "noisy_expected_hypervolume_improvement": { + "identifiers": ["nehvi", "noisy_expected_hypervolume_improvement"], + "pretty_name": "Noisy expected hypervolume improvement", + "multitask_only": True, + "description": r"It's like a big box. How big is the box?", + }, + "lower_bound_max_value_entropy": { + "identifiers": ["lbmve", "lbmes", "gibbon", "lower_bound_max_value_entropy"], + "pretty_name": "Lower bound max value entropy", + "multitask_only": False, + "description": r"Max entropy search, basically", + }, + "probability_of_improvement": { + "identifiers": ["pi", "probability_of_improvement"], + "pretty_name": "Probability of improvement", + "multitask_only": False, + "description": "The probability that this input improves on the current maximum.", + }, + "upper_confidence_bound": { + "identifiers": ["ucb", "upper_confidence_bound"], + "default_args": {"beta": 4}, + "pretty_name": "Upper confidence bound", + "multitask_only": False, + "description": r"The expected value, plus some multiple of the uncertainty (typically \mu + 2\sigma).", + }, +} + class ConstrainedUpperConfidenceBound(UpperConfidenceBound): def __init__(self, constraint, *args, **kwargs): @@ -20,13 +66,11 @@ def __init__(self, constraint, *args, **kwargs): self.constraint = constraint def forward(self, x): - mean, sigma = self._mean_and_sigma(x) - p_eff = 0.5 * (1 + torch.special.erf(self.beta.sqrt()/math.sqrt(2))) * torch.clamp(self.constraint(x), min=1e-6) - - return (mean if self.maximize else -mean) + sigma * np.sqrt(2) * torch.special.erfinv(2*p_eff-1) + p_eff = 0.5 * (1 + torch.special.erf(self.beta.sqrt() / math.sqrt(2))) * torch.clamp(self.constraint(x), min=1e-6) + return (mean if self.maximize else -mean) + sigma * np.sqrt(2) * torch.special.erfinv(2 * p_eff - 1) class ConstrainedLogExpectedImprovement(LogExpectedImprovement): @@ -63,5 +107,3 @@ def __init__(self, constraint, *args, **kwargs): def forward(self, x): return super().forward(x) * self.constraint(x) - - From cd49326d70039d2e23fd85b47d590e132e328a52 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 14:36:18 -0400 Subject: [PATCH 06/43] fix deps --- .github/workflows/testing.yml | 20 ++++++++++---------- requirements.txt | 2 +- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index f953f1e..98e075b 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -35,16 +35,16 @@ jobs: - name: Start MongoDB uses: supercharge/mongodb-github-action@1.6.0 - - name: Start Sirepo Docker container - uses: NSLS-II/start-sirepo-action@v2 - with: - docker-binary: docker + # - name: Start Sirepo Docker container + # uses: NSLS-II/start-sirepo-action@v2 + # with: + # docker-binary: docker - - name: Copy databroker config file - run: | - set -vxeuo pipefail - mkdir -v -p ~/.config/databroker/ - wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml + # - name: Copy databroker config file + # run: | + # set -vxeuo pipefail + # mkdir -v -p ~/.config/databroker/ + # wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml - name: Set up Python ${{ matrix.python-version }} with conda uses: conda-incubator/setup-miniconda@v2 @@ -72,5 +72,5 @@ jobs: - name: Test with pytest run: | set -vxeuo pipefail - coverage run -m pytest -vv -s + coverage run -m pytest -vv -s -m test_func coverage report -m diff --git a/requirements.txt b/requirements.txt index f5bf1cb..fc18783 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,6 +10,6 @@ ortools scipy shadow3 sirepo-bluesky -srwpy +srwpy==4.0.0b1 tables torch From 092768d377660ddf15240d46056bcf675c5f6db7 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 14:52:45 -0400 Subject: [PATCH 07/43] do not test in the docs --- .github/workflows/testing.yml | 4 +-- bloptools/bayesian/models.py | 1 - bloptools/test_functions.py | 31 +++++++++---------- .../tutorials/constrained-himmelblau.ipynb | 6 ++-- docs/source/tutorials/multi-task-sirepo.ipynb | 23 +------------- 5 files changed, 21 insertions(+), 44 deletions(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 98e075b..8b392aa 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -32,8 +32,8 @@ jobs: - name: Checkout the code uses: actions/checkout@v3 - - name: Start MongoDB - uses: supercharge/mongodb-github-action@1.6.0 + # - name: Start MongoDB + # uses: supercharge/mongodb-github-action@1.6.0 # - name: Start Sirepo Docker container # uses: NSLS-II/start-sirepo-action@v2 diff --git a/bloptools/bayesian/models.py b/bloptools/bayesian/models.py index 2c11473..8733ac8 100644 --- a/bloptools/bayesian/models.py +++ b/bloptools/bayesian/models.py @@ -32,4 +32,3 @@ def probabilities(self, x, n_samples=256): *input_shape, n_dim = x.shape samples = self.posterior(x.reshape(-1, n_dim)).sample(torch.Size((n_samples,))).exp() return (samples / samples.sum(-1, keepdim=True)).mean(0).reshape(*input_shape, -1) - diff --git a/bloptools/test_functions.py b/bloptools/test_functions.py index 34d139d..845e155 100644 --- a/bloptools/test_functions.py +++ b/bloptools/test_functions.py @@ -82,27 +82,26 @@ def gaussian_beam_waist(x1, x2): return np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) - - - - def hartmann6(*x): - X = np.c_[x] alpha = np.array([1.0, 1.2, 3.0, 3.2]) - A = np.array([[10, 3, 17, 3.5, 1.7, 8], - [0.05, 10, 17, 0.1, 8, 14], - [3, 3.5, 1.7, 10, 17, 8], - [17, 8, 0.05, 10, 0.1, 14]]) - - P = 1e-4 * np.array([[1312, 1696, 5569, 124, 8283, 5886], - [2329, 4135, 8307, 3736, 1004, 9991], - [2348, 1451, 3522, 2883, 3047, 6650], - [4047, 8828, 8732, 5743, 1091, 381]]) - - return - (alpha * np.exp(-(A * np.square(X - P)).sum(axis=1))).sum() + A = np.array( + [[10, 3, 17, 3.5, 1.7, 8], [0.05, 10, 17, 0.1, 8, 14], [3, 3.5, 1.7, 10, 17, 8], [17, 8, 0.05, 10, 0.1, 14]] + ) + + P = 1e-4 * np.array( + [ + [1312, 1696, 5569, 124, 8283, 5886], + [2329, 4135, 8307, 3736, 1004, 9991], + [2348, 1451, 3522, 2883, 3047, 6650], + [4047, 8828, 8732, 5743, 1091, 381], + ] + ) + + return -(alpha * np.exp(-(A * np.square(X - P)).sum(axis=1))).sum() + def kb_tradeoff_2d(x1, x2): width = np.sqrt(1 + 0.25 * (x1 - x2) ** 2 + 16 * (x1 + x2 - 2) ** 2) diff --git a/docs/source/tutorials/constrained-himmelblau.ipynb b/docs/source/tutorials/constrained-himmelblau.ipynb index 01c06d8..0333af8 100644 --- a/docs/source/tutorials/constrained-himmelblau.ipynb +++ b/docs/source/tutorials/constrained-himmelblau.ipynb @@ -210,7 +210,7 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": "Python 3.11.4 ('bluesky')", "language": "python", "name": "python3" }, @@ -224,11 +224,11 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.11.4" }, "vscode": { "interpreter": { - "hash": "9aced674e98d511b4f654e147532c84d38dc986fe042b1e92785fb9d8df41f75" + "hash": "eee21ccc240bdddd7cf04478199e20f7257541e2f592ca1a4d34ebdc0225d742" } } }, diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/source/tutorials/multi-task-sirepo.ipynb index 3f33cf0..906e457 100644 --- a/docs/source/tutorials/multi-task-sirepo.ipynb +++ b/docs/source/tutorials/multi-task-sirepo.ipynb @@ -28,27 +28,6 @@ "%run -i ../../../examples/prepare_tes_shadow.py" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "fff28b50", - "metadata": {}, - "outputs": [], - "source": [ - "from bloptools.tests import conftest" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "aa1d1473", - "metadata": {}, - "outputs": [], - "source": [ - "db = conftest.db\n", - "RE = conftest.RE()" - ] - }, { "cell_type": "code", "execution_count": null, @@ -163,7 +142,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.11.5" }, "vscode": { "interpreter": { From 52301fb1b56122058ad373394c05b1843a78c8fc Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 16:05:18 -0400 Subject: [PATCH 08/43] not all tasks need to be valid --- bloptools/bayesian/__init__.py | 150 +++++++++++++-------------------- bloptools/tests/test_plots.py | 3 +- 2 files changed, 62 insertions(+), 91 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 7ef6abe..c0f812e 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -212,28 +212,22 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): self.table = pd.concat([self.table, new_table]) if append else new_table self.table.index = np.arange(len(self.table)) - fitnesses = self.task_fitnesses # computes from self.table - - # update fitness estimates - self.table.loc[:, fitnesses.columns] = fitnesses.values - self.table.loc[:, "total_fitness"] = fitnesses.values.sum(axis=1) - skew_dims = [tuple(np.arange(self._len_subset_dofs(mode="on")))] if self._initialized: cached_hypers = self.hypers - feasibility = ~fitnesses.isna().any(axis=1) + inputs = self.inputs.loc[:, self._subset_dof_names(mode="on")].values - if not feasibility.sum() >= 2: - raise ValueError("There must be at least two feasible data points per task!") + for i, task in enumerate(self.tasks): + self.table.loc[:, f"{task['key']}_fitness"] = targets = self._get_task_fitness(i) + train_index = ~np.isnan(targets) - inputs = self.inputs.loc[feasibility, self._subset_dof_names(mode="on")].values - train_inputs = torch.tensor(inputs).double() # .unsqueeze(0) + if not len(train_index) >= 2: + raise ValueError("There must be at least two valid data points per task!") - for task in self.tasks: - targets = self.table.loc[feasibility, f'{task["key"]}_fitness'].values - train_targets = torch.tensor(targets).double().unsqueeze(-1) # .unsqueeze(0) + train_inputs = torch.tensor(inputs[train_index]).double() + train_targets = torch.tensor(targets[train_index]).double().unsqueeze(-1) # .unsqueeze(0) likelihood = gpytorch.likelihoods.GaussianLikelihood( noise_constraint=gpytorch.constraints.Interval( @@ -253,18 +247,12 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): outcome_transform=outcome_transform, ).double() - # # this ensures that we have equal weight between task fitness and feasibility fitness - # self.task_scalarization = botorch.acquisition.objective.ScalarizedPosteriorTransform( - # weights=torch.tensor([*torch.ones(self.num_tasks), self.fitness_variance.sum().sqrt()]).double(), - # offset=0, - # ) - dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood( - torch.tensor(feasibility).long(), learn_additional_noise=True + self.all_tasks_valid.long(), learn_additional_noise=True ).double() self.classifier = models.LatentDirichletClassifier( - train_inputs=torch.tensor(self.inputs.values).double(), + train_inputs=torch.tensor(inputs).double(), train_targets=dirichlet_likelihood.transformed_targets.transpose(-1, -2).double(), skew_dims=skew_dims, likelihood=dirichlet_likelihood, @@ -290,53 +278,43 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): def model(self): return ModelListGP(*[task["model"] for task in self.tasks]) if self.num_tasks > 1 else self.tasks[0]["model"] - @property - def target_names(self): - return [f'{task["key"]}_fitness' for task in self.tasks] + def _get_task_fitness(self, task_index): + task = self.tasks[task_index] - @property - def train_inputs(self): - feasibility = ~self.task_fitnesses.isna().any(axis=1) - inputs = self.inputs.loc[feasibility, self._subset_dof_names(mode="on")].values - return torch.tensor(inputs).double() + targets = self.table.loc[:, task["key"]].values.copy() - @property - def train_targets(self): - feasibility = ~self.task_fitnesses.isna().any(axis=1) - inputs = self.table.loc[feasibility, self.target_names].values - return torch.tensor(inputs).double() + if task["kind"] == "minimize": + targets *= -1 - @property - def class_inputs(self): - inputs = self.inputs.loc[:, self._subset_dof_names(mode="on")].values - return torch.tensor(inputs).double() + # check that task values are inside acceptable values + valid = (targets > task["limits"][0]) & (targets < task["limits"][1]) + targets = np.where(valid, targets, np.nan) - @property - def class_targets(self): - feasibility = ~self.task_fitnesses.isna().any(axis=1) - return torch.tensor(feasibility).double() + # transform if needed + if "transform" in task.keys(): + if task["transform"] == "log": + targets = np.where(targets > 0, np.log(targets), np.nan) - @property - def task_fitnesses(self): - df = pd.DataFrame(index=self.table.index) - for task in self.tasks: - name = f'{task["key"]}_fitness' + return targets - df.loc[:, name] = task["weight"] * self.table.loc[:, task["key"]] + @property + def fitnesses(self): + """ + Returns a (num_tasks x n_obs) array of fitnesses + """ + return torch.tensor(np.c_[*[self._get_task_fitness(i) for i in range(self.num_tasks)]]).double() - # check that task values are inside acceptable values - valid = (df.loc[:, name] > task["limits"][0]) & (df.loc[:, name] < task["limits"][1]) + @property + def scalarized_fitness(self): + return (self.fitnesses * self.task_weights).sum(axis=-1) - # transform if needed - if "transform" in task.keys(): - if task["transform"] == "log": - valid &= df.loc[:, name] > 0 - df.loc[valid, name] = np.log(df.loc[valid, name]) - df.loc[~valid, name] = np.nan + @property + def all_tasks_valid(self): + return ~torch.isnan(self.scalarized_fitness) - if task["kind"] == "minimize": - df.loc[valid, name] *= -1 - return df + @property + def target_names(self): + return [f'{task["key"]}_fitness' for task in self.tasks] def _dof_kind_mask(self, kind=None): return [dof["kind"] == kind if kind is not None else True for dof in self.dofs] @@ -721,17 +699,9 @@ def learn( def inputs(self): return self.table.loc[:, self._subset_dof_names(mode="on")].astype(float) - @property - def fitness_variance(self): - return torch.tensor(np.nanvar(self.task_fitnesses.values, axis=0)) - - @property - def scalarized_fitness(self): - return self.task_fitnesses.sum(axis=1) - # @property # def best_sum_of_tasks_inputs(self): - # return self.inputs[np.nanargmax(self.task_fitnesses.sum(axis=1))] + # return self.inputs[np.nanargmax(self.fitnesses.sum(axis=1))] @property def go_to(self, inputs): @@ -805,10 +775,8 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL self.task_axes = np.atleast_2d(self.task_axes) # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - feasible = ~self.task_fitnesses.isna().any(axis=1) - for itask, task in enumerate(self.tasks): - sampled_fitness = np.where(feasible, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) + sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) task_norm = mpl.colors.Normalize(task_vmin, task_vmax) @@ -968,37 +936,39 @@ def _plot_acq_many_dofs( ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) - def plot_feasibility(self, **kwargs): + def plot_validity(self, **kwargs): if self._len_subset_dofs(kind="active", mode="on") == 1: - self._plot_feas_one_dof(**kwargs) + self._plot_valid_one_dof(**kwargs) else: - self._plot_feas_many_dofs(**kwargs) + self._plot_valid_many_dofs(**kwargs) - def _plot_feas_one_dof(self, size=16, lw=1e0): - self.feas_fig, self.feas_ax = plt.subplots(1, 1, figsize=(4, 4), sharex=True, constrained_layout=True) + def _plot_valid_one_dof(self, size=16, lw=1e0): + self.valid_fig, self.valid_ax = plt.subplots(1, 1, figsize=(4, 4), sharex=True, constrained_layout=True) x = self.test_inputs_grid *input_shape, input_dim = x.shape constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) - self.feas_ax.scatter(self.inputs.values, ~self.task_fitnesses.isna().any(axis=1), s=size) + self.valid_ax.scatter(self.inputs.values, self.all_tasks_valid, s=size) on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - self.feas_ax.plot(x[..., on_dofs_are_active_mask].squeeze(), constraint.detach().numpy(), lw=lw) + self.valid_ax.plot(x[..., on_dofs_are_active_mask].squeeze(), constraint.detach().numpy(), lw=lw) - self.feas_ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[0]["limits"]) + self.valid_ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[0]["limits"]) - def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=16, gridded=None): - self.feas_fig, self.feas_axes = plt.subplots(1, 2, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True) + def _plot_valid_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=16, gridded=None): + self.valid_fig, self.valid_axes = plt.subplots( + 1, 2, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True + ) if gridded is None: gridded = self._len_subset_dofs(kind="active", mode="on") == 2 - data_ax = self.feas_axes[0].scatter( + data_ax = self.valid_axes[0].scatter( *self.inputs.values.T[:2], - c=~self.task_fitnesses.isna().any(axis=1), + c=self.all_tasks_valid, s=size, vmin=0, vmax=1, @@ -1010,7 +980,7 @@ def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLO constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) if gridded: - self.feas_axes[1].pcolormesh( + self.valid_axes[1].pcolormesh( x[..., 0], x[..., 1], constraint.detach().numpy(), @@ -1020,19 +990,19 @@ def _plot_feas_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLO vmax=1, ) - # self.acq_fig.colorbar(obj_ax, ax=self.feas_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) + # self.acq_fig.colorbar(obj_ax, ax=self.valid_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) else: - # self.feas_axes.set_title(acq_func_meta["name"]) - self.feas_axes[1].scatter( + # self.valid_axes.set_title(acq_func_meta["name"]) + self.valid_axes[1].scatter( x.detach().numpy()[..., axes[0]], x.detach().numpy()[..., axes[1]], c=constraint.detach().numpy(), ) - self.feas_fig.colorbar(data_ax, ax=self.feas_axes[:2], location="bottom", aspect=32, shrink=0.8) + self.valid_fig.colorbar(data_ax, ax=self.valid_axes[:2], location="bottom", aspect=32, shrink=0.8) - for ax in self.feas_axes.ravel(): + for ax in self.valid_axes.ravel(): ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) diff --git a/bloptools/tests/test_plots.py b/bloptools/tests/test_plots.py index 4675b01..aa1aa42 100644 --- a/bloptools/tests/test_plots.py +++ b/bloptools/tests/test_plots.py @@ -7,4 +7,5 @@ def test_plots(RE, agent): agent.plot_tasks() agent.plot_acquisition() - agent.plot_feasibility() + agent.plot_validity() + agent.plot_history() From fcd133323f3431db47bfb5bf3c4b1fdfed3141fd Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 16:14:42 -0400 Subject: [PATCH 09/43] not a syntax error but go off i guess --- bloptools/bayesian/__init__.py | 74 ++++++++++++++++++---------------- 1 file changed, 40 insertions(+), 34 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index c0f812e..98c6860 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -276,9 +276,15 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): @property def model(self): + """ + A model encompassing all the tasks. A single GP in the single-task case, or a model list. + """ return ModelListGP(*[task["model"] for task in self.tasks]) if self.num_tasks > 1 else self.tasks[0]["model"] def _get_task_fitness(self, task_index): + """ + Returns the fitness for a task given the task index. + """ task = self.tasks[task_index] targets = self.table.loc[:, task["key"]].values.copy() @@ -302,7 +308,7 @@ def fitnesses(self): """ Returns a (num_tasks x n_obs) array of fitnesses """ - return torch.tensor(np.c_[*[self._get_task_fitness(i) for i in range(self.num_tasks)]]).double() + return torch.cat([torch.tensor(self._get_task_fitness(i))[..., None] for i in range(self.num_tasks)], dim=1) @property def scalarized_fitness(self): @@ -316,39 +322,6 @@ def all_tasks_valid(self): def target_names(self): return [f'{task["key"]}_fitness' for task in self.tasks] - def _dof_kind_mask(self, kind=None): - return [dof["kind"] == kind if kind is not None else True for dof in self.dofs] - - def _dof_mode_mask(self, mode=None): - return [dof["mode"] == mode if mode is not None else True for dof in self.dofs] - - def _dof_mask(self, kind=None, mode=None): - return [(k and m) for k, m in zip(self._dof_kind_mask(kind), self._dof_mode_mask(mode))] - - def _subset_dofs(self, kind=None, mode=None): - return [dof for dof, m in zip(self.dofs, self._dof_mask(kind, mode)) if m] - - def _len_subset_dofs(self, kind=None, mode=None): - return len(self._subset_dofs(kind, mode)) - - def _subset_devices(self, kind=None, mode=None): - return [dof["device"] for dof in self._subset_dofs(kind, mode)] - - def _read_subset_devices(self, kind=None, mode=None): - return [device.read()[device.name]["value"] for device in self._subset_devices(kind, mode)] - - def _subset_dof_names(self, kind=None, mode=None): - return [device.name for device in self._subset_devices(kind, mode)] - - def _subset_dof_limits(self, kind=None, mode=None): - dofs_subset = self._subset_dofs(kind, mode) - if len(dofs_subset) > 0: - return torch.tensor([dof["limits"] for dof in dofs_subset], dtype=torch.float64).T - return torch.empty((2, 0)) - - def test_inputs(self, n=MAX_TEST_INPUTS): - return utils.sobol_sampler(self._acq_func_bounds, n=n) - @property def test_inputs_grid(self): n_side = int(MAX_TEST_INPUTS ** (1 / self._len_subset_dofs(kind="active", mode="on"))) @@ -399,6 +372,39 @@ def task_weights(self): def task_signs(self): return torch.tensor([(1 if task["kind"] == "maximize" else -1) for task in self.tasks], dtype=torch.long) + def _dof_kind_mask(self, kind=None): + return [dof["kind"] == kind if kind is not None else True for dof in self.dofs] + + def _dof_mode_mask(self, mode=None): + return [dof["mode"] == mode if mode is not None else True for dof in self.dofs] + + def _dof_mask(self, kind=None, mode=None): + return [(k and m) for k, m in zip(self._dof_kind_mask(kind), self._dof_mode_mask(mode))] + + def _subset_dofs(self, kind=None, mode=None): + return [dof for dof, m in zip(self.dofs, self._dof_mask(kind, mode)) if m] + + def _len_subset_dofs(self, kind=None, mode=None): + return len(self._subset_dofs(kind, mode)) + + def _subset_devices(self, kind=None, mode=None): + return [dof["device"] for dof in self._subset_dofs(kind, mode)] + + def _read_subset_devices(self, kind=None, mode=None): + return [device.read()[device.name]["value"] for device in self._subset_devices(kind, mode)] + + def _subset_dof_names(self, kind=None, mode=None): + return [device.name for device in self._subset_devices(kind, mode)] + + def _subset_dof_limits(self, kind=None, mode=None): + dofs_subset = self._subset_dofs(kind, mode) + if len(dofs_subset) > 0: + return torch.tensor([dof["limits"] for dof in dofs_subset], dtype=torch.float64).T + return torch.empty((2, 0)) + + def test_inputs(self, n=MAX_TEST_INPUTS): + return utils.sobol_sampler(self._acq_func_bounds, n=n) + def _subset_input_transform(self, kind=None, mode=None): limits = self._subset_dof_limits(kind, mode) offset = limits.min(dim=0).values From b9173842d62a58575e5b5edc9a465351639352b7 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 16:40:43 -0400 Subject: [PATCH 10/43] remove sirepo-bluesky dependency --- bloptools/bayesian/__init__.py | 11 +++- bloptools/tests/conftest.py | 65 ++++++++++++++----- bloptools/tests/test_passive_dofs.py | 2 +- .../tutorials/constrained-himmelblau.ipynb | 2 +- .../tutorials/latent-toroid-dimensions.ipynb | 2 +- requirements.txt | 3 - 6 files changed, 60 insertions(+), 25 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 98c6860..d9ac482 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -314,6 +314,11 @@ def fitnesses(self): def scalarized_fitness(self): return (self.fitnesses * self.task_weights).sum(axis=-1) + @property + def best_scalarized_fitness(self): + f = self.scalarized_fitness + return np.where(np.isnan(f), -np.inf, f).max() + @property def all_tasks_valid(self): return ~torch.isnan(self.scalarized_fitness) @@ -516,7 +521,7 @@ def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=Fal acq_func = acquisition.ConstrainedLogExpectedImprovement( constraint=self.constraint, model=self.model, - best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), + best_f=self.best_scalarized_fitness, posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), ) acq_func_meta = {"name": acq_func_name, "args": {}} @@ -525,7 +530,7 @@ def get_acquisition_function(self, acq_func_identifier="ei", return_metadata=Fal acq_func = acquisition.ConstrainedLogProbabilityOfImprovement( constraint=self.constraint, model=self.model, - best_f=(self.train_targets * self.task_weights).sum(axis=-1).max(), + best_f=self.best_scalarized_fitness, posterior_transform=ScalarizedPosteriorTransform(weights=self.task_weights, offset=0), ) acq_func_meta = {"name": acq_func_name, "args": {}} @@ -1057,7 +1062,7 @@ def plot_history(self, x_key="index", show_all_tasks=False): hist_axes[itask].plot(x, y, lw=5e-1, c="k") hist_axes[itask].set_ylabel(task["key"]) - y = self.table.total_fitness + y = self.scalarized_fitness cummax_y = np.array([np.nanmax(y[: i + 1]) for i in range(len(y))]) diff --git a/bloptools/tests/conftest.py b/bloptools/tests/conftest.py index c1226ed..07cc958 100644 --- a/bloptools/tests/conftest.py +++ b/bloptools/tests/conftest.py @@ -7,10 +7,6 @@ from bluesky.run_engine import RunEngine from databroker import Broker from ophyd.utils import make_dir_tree -from sirepo_bluesky.madx_handler import MADXFileHandler -from sirepo_bluesky.shadow_handler import ShadowFileHandler -from sirepo_bluesky.sirepo_bluesky import SirepoBluesky -from sirepo_bluesky.srw_handler import SRWFileHandler from bloptools.bayesian import Agent @@ -27,6 +23,41 @@ def db(): except Exception: pass + return db + + +@pytest.fixture(scope="function") +def RE(db): + loop = asyncio.new_event_loop() + loop.set_debug(True) + RE = RunEngine({}, loop=loop) + RE.subscribe(db.insert) + + bec = best_effort.BestEffortCallback() + RE.subscribe(bec) + + bec.disable_baseline() + bec.disable_heading() + bec.disable_table() + bec.disable_plots() + + return RE + + +@pytest.fixture(scope="function") +def db_with_bluesky(): + from sirepo_bluesky.madx_handler import MADXFileHandler + from sirepo_bluesky.shadow_handler import ShadowFileHandler + from sirepo_bluesky.srw_handler import SRWFileHandler + + """Return a data broker""" + # MongoDB backend: + db = Broker.named("local") # mongodb backend + try: + databroker.assets.utils.install_sentinels(db.reg.config, version=1) + except Exception: + pass + db.reg.register_handler("srw", SRWFileHandler, overwrite=True) db.reg.register_handler("shadow", ShadowFileHandler, overwrite=True) db.reg.register_handler("SIREPO_FLYER", SRWFileHandler, overwrite=True) @@ -36,11 +67,11 @@ def db(): @pytest.fixture(scope="function") -def RE(db): +def RE_with_bluesky(db_with_bluesky): loop = asyncio.new_event_loop() loop.set_debug(True) RE = RunEngine({}, loop=loop) - RE.subscribe(db.insert) + RE.subscribe(db_with_bluesky.insert) bec = best_effort.BestEffortCallback() RE.subscribe(bec) @@ -123,15 +154,17 @@ def make_dirs(): _ = make_dir_tree(datetime.datetime.now().year, base_path=root_dir) -@pytest.fixture(scope="function") -def srw_tes_simulation(make_dirs): - connection = SirepoBluesky("http://localhost:8000") - data, _ = connection.auth("srw", "00000002") - return connection +# @pytest.fixture(scope="function") +# def srw_tes_simulation(make_dirs): +# from sirepo_bluesky.sirepo_bluesky import SirepoBluesky +# connection = SirepoBluesky("http://localhost:8000") +# data, _ = connection.auth("srw", "00000002") +# return connection -@pytest.fixture(scope="function") -def shadow_tes_simulation(make_dirs): - connection = SirepoBluesky("http://localhost:8000") - data, _ = connection.auth("shadow", "00000002") - return connection +# @pytest.fixture(scope="function") +# def shadow_tes_simulation(make_dirs): +# from sirepo_bluesky.sirepo_bluesky import SirepoBluesky +# connection = SirepoBluesky("http://localhost:8000") +# data, _ = connection.auth("shadow", "00000002") +# return connection diff --git a/bloptools/tests/test_passive_dofs.py b/bloptools/tests/test_passive_dofs.py index 7760e38..d6aa6b7 100644 --- a/bloptools/tests/test_passive_dofs.py +++ b/bloptools/tests/test_passive_dofs.py @@ -30,4 +30,4 @@ def test_passive_dofs(RE, db): agent.plot_tasks() agent.plot_acquisition() - agent.plot_feasibility() + agent.plot_validity() diff --git a/docs/source/tutorials/constrained-himmelblau.ipynb b/docs/source/tutorials/constrained-himmelblau.ipynb index 0333af8..b16d59f 100644 --- a/docs/source/tutorials/constrained-himmelblau.ipynb +++ b/docs/source/tutorials/constrained-himmelblau.ipynb @@ -131,7 +131,7 @@ }, "outputs": [], "source": [ - "agent.plot_feasibility()" + "agent.plot_validity()" ] }, { diff --git a/docs/source/tutorials/latent-toroid-dimensions.ipynb b/docs/source/tutorials/latent-toroid-dimensions.ipynb index 090bbc1..dd93bd1 100644 --- a/docs/source/tutorials/latent-toroid-dimensions.ipynb +++ b/docs/source/tutorials/latent-toroid-dimensions.ipynb @@ -85,7 +85,7 @@ "outputs": [], "source": [ "agent.plot_tasks()\n", - "agent.plot_feasibility()\n", + "agent.plot_validity()\n", "agent.plot_acquisition(strategy=[\"ei\", \"pi\", \"ucb\"])" ] } diff --git a/requirements.txt b/requirements.txt index fc18783..433b4c5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,8 +8,5 @@ numpy ophyd ortools scipy -shadow3 -sirepo-bluesky -srwpy==4.0.0b1 tables torch From 5529fde77677458fd082fdd3a85a601d336cc94c Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 5 Sep 2023 18:00:50 -0400 Subject: [PATCH 11/43] fix deps --- docs/source/tutorials/hyperparameters.ipynb | 2 +- docs/source/tutorials/introduction.ipynb | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb index 0421d8b..e3908f5 100644 --- a/docs/source/tutorials/hyperparameters.ipynb +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -68,7 +68,7 @@ }, "outputs": [], "source": [ - "from sirepo_bluesky import prepare_re_env\n", + "from bloptools.utils import prepare_re_env\n", "\n", "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", diff --git a/docs/source/tutorials/introduction.ipynb b/docs/source/tutorials/introduction.ipynb index a307fff..9ccc3b0 100644 --- a/docs/source/tutorials/introduction.ipynb +++ b/docs/source/tutorials/introduction.ipynb @@ -118,7 +118,7 @@ }, "outputs": [], "source": [ - "from sirepo_bluesky import prepare_re_env\n", + "from bloptools.utils import prepare_re_env\n", "\n", "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", From 6326c606e1bc9a059dc150447de5b87c0ec742b5 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 6 Sep 2023 10:09:16 -0400 Subject: [PATCH 12/43] conda install srwpy to build docs --- .github/workflows/docs.yml | 3 +-- .github/workflows/testing.yml | 2 +- =4.0.0 | 0 =4.0.0b1 | 0 docs/source/tutorials/latent-toroid-dimensions.ipynb | 2 +- 5 files changed, 3 insertions(+), 4 deletions(-) create mode 100644 =4.0.0 create mode 100644 =4.0.0b1 diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 50cc42c..ae35ed6 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -63,8 +63,7 @@ jobs: set -vxeo pipefail conda env list - # mamba install -c conda-forge shadow3 srwpy pandoc - mamba install -c conda-forge pandoc + mamba install -c conda-forge shadow3 srwpy pandoc pip install --upgrade pip wheel pip install -v . pip install -r requirements-dev.txt diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 8b392aa..448bd72 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -62,7 +62,7 @@ jobs: set -vxeo pipefail conda env list - # mamba install -c conda-forge shadow3 srwpy + mamba install -c conda-forge shadow3 srwpy pip install --upgrade pip wheel pip install -v . pip install -r requirements-dev.txt diff --git a/=4.0.0 b/=4.0.0 new file mode 100644 index 0000000..e69de29 diff --git a/=4.0.0b1 b/=4.0.0b1 new file mode 100644 index 0000000..e69de29 diff --git a/docs/source/tutorials/latent-toroid-dimensions.ipynb b/docs/source/tutorials/latent-toroid-dimensions.ipynb index dd93bd1..8eb04b8 100644 --- a/docs/source/tutorials/latent-toroid-dimensions.ipynb +++ b/docs/source/tutorials/latent-toroid-dimensions.ipynb @@ -106,7 +106,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.5" + "version": "3.10.12" }, "vscode": { "interpreter": { From 939fa70e2dfce0d205602558c8820a6040850478 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 6 Sep 2023 10:09:53 -0400 Subject: [PATCH 13/43] conda install srwpy to build docs --- =4.0.0 | 0 =4.0.0b1 | 0 2 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 =4.0.0 delete mode 100644 =4.0.0b1 diff --git a/=4.0.0 b/=4.0.0 deleted file mode 100644 index e69de29..0000000 diff --git a/=4.0.0b1 b/=4.0.0b1 deleted file mode 100644 index e69de29..0000000 From 0ba210decbc88f1d6f707545c7b7535099e5da2a Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 6 Sep 2023 11:31:34 -0400 Subject: [PATCH 14/43] sirepo-bluesky in dev requirements --- .github/workflows/testing.yml | 22 +++++++++++----------- requirements-dev.txt | 3 +++ 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 448bd72..48b84be 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -32,19 +32,19 @@ jobs: - name: Checkout the code uses: actions/checkout@v3 - # - name: Start MongoDB - # uses: supercharge/mongodb-github-action@1.6.0 + - name: Start MongoDB + uses: supercharge/mongodb-github-action@1.6.0 - # - name: Start Sirepo Docker container - # uses: NSLS-II/start-sirepo-action@v2 - # with: - # docker-binary: docker + - name: Start Sirepo Docker container + uses: NSLS-II/start-sirepo-action@v2 + with: + docker-binary: docker - # - name: Copy databroker config file - # run: | - # set -vxeuo pipefail - # mkdir -v -p ~/.config/databroker/ - # wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml + - name: Copy databroker config file + run: | + set -vxeuo pipefail + mkdir -v -p ~/.config/databroker/ + wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml - name: Set up Python ${{ matrix.python-version }} with conda uses: conda-incubator/setup-miniconda@v2 diff --git a/requirements-dev.txt b/requirements-dev.txt index 6782e7c..fc2d41e 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -19,5 +19,8 @@ matplotlib nbsphinx numpydoc pandoc +shadow3 +sirepo-bluesky sphinx-copybutton sphinx_rtd_theme +srwpy From 61a39eec93f4b3cd907d72bd66ece05a0d963043 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 6 Sep 2023 11:59:21 -0400 Subject: [PATCH 15/43] tests --- .github/workflows/testing.yml | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 48b84be..5e166f9 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -16,7 +16,7 @@ jobs: strategy: matrix: host-os: ["ubuntu-latest"] - python-version: ["3.8", "3.9", "3.10"] + python-version: ["3.10"] fail-fast: false defaults: @@ -32,19 +32,19 @@ jobs: - name: Checkout the code uses: actions/checkout@v3 - - name: Start MongoDB - uses: supercharge/mongodb-github-action@1.6.0 + # - name: Start MongoDB + # uses: supercharge/mongodb-github-action@1.6.0 - - name: Start Sirepo Docker container - uses: NSLS-II/start-sirepo-action@v2 - with: - docker-binary: docker + # - name: Start Sirepo Docker container + # uses: NSLS-II/start-sirepo-action@v2 + # with: + # docker-binary: docker - - name: Copy databroker config file - run: | - set -vxeuo pipefail - mkdir -v -p ~/.config/databroker/ - wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml + # - name: Copy databroker config file + # run: | + # set -vxeuo pipefail + # mkdir -v -p ~/.config/databroker/ + # wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml - name: Set up Python ${{ matrix.python-version }} with conda uses: conda-incubator/setup-miniconda@v2 @@ -62,7 +62,7 @@ jobs: set -vxeo pipefail conda env list - mamba install -c conda-forge shadow3 srwpy + pip install --upgrade pip wheel pip install -v . pip install -r requirements-dev.txt From 864368fbd9775b1a383dbb9cfa9f853083542087 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 6 Sep 2023 12:06:07 -0400 Subject: [PATCH 16/43] deps --- requirements-dev.txt | 3 --- 1 file changed, 3 deletions(-) diff --git a/requirements-dev.txt b/requirements-dev.txt index fc2d41e..29110b8 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -18,9 +18,6 @@ jupyter matplotlib nbsphinx numpydoc -pandoc -shadow3 sirepo-bluesky sphinx-copybutton sphinx_rtd_theme -srwpy From a005f0d3604bea996cc2a98014f5dd2225c38033 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 6 Sep 2023 12:57:35 -0400 Subject: [PATCH 17/43] apply the minus sign AFTER the logarithm --- .github/workflows/testing.yml | 2 +- bloptools/bayesian/__init__.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 5e166f9..c574cb7 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -16,7 +16,7 @@ jobs: strategy: matrix: host-os: ["ubuntu-latest"] - python-version: ["3.10"] + python-version: ["3.8", "3.9", "3.10"] fail-fast: false defaults: diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index d9ac482..a65e33b 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -289,9 +289,6 @@ def _get_task_fitness(self, task_index): targets = self.table.loc[:, task["key"]].values.copy() - if task["kind"] == "minimize": - targets *= -1 - # check that task values are inside acceptable values valid = (targets > task["limits"][0]) & (targets < task["limits"][1]) targets = np.where(valid, targets, np.nan) @@ -301,6 +298,9 @@ def _get_task_fitness(self, task_index): if task["transform"] == "log": targets = np.where(targets > 0, np.log(targets), np.nan) + if task["kind"] == "minimize": + targets *= -1 + return targets @property From cd263656bc512d3535e9ac3c59d53e2b1b7ab263 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 6 Sep 2023 18:08:21 -0400 Subject: [PATCH 18/43] work at ATF --- bloptools/bayesian/__init__.py | 24 +++++++------ bloptools/bayesian/acquisition.py | 57 +++++++++++++++++++++++++++++-- bloptools/bayesian/digestion.py | 4 ++- 3 files changed, 71 insertions(+), 14 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index a65e33b..ad765dd 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -223,7 +223,7 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): self.table.loc[:, f"{task['key']}_fitness"] = targets = self._get_task_fitness(i) train_index = ~np.isnan(targets) - if not len(train_index) >= 2: + if not train_index.sum() >= 2: raise ValueError("There must be at least two valid data points per task!") train_inputs = torch.tensor(inputs[train_index]).double() @@ -232,7 +232,7 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): likelihood = gpytorch.likelihoods.GaussianLikelihood( noise_constraint=gpytorch.constraints.Interval( torch.tensor(1e-6).square(), - torch.tensor(1e-2).square(), + torch.tensor(1e0).square(), ), ).double() @@ -710,17 +710,19 @@ def learn( def inputs(self): return self.table.loc[:, self._subset_dof_names(mode="on")].astype(float) - # @property - # def best_sum_of_tasks_inputs(self): - # return self.inputs[np.nanargmax(self.fitnesses.sum(axis=1))] - @property - def go_to(self, inputs): - yield from bps.mv(*[_ for items in zip(self._subset_dofs(kind="active"), np.atleast_1d(inputs).T) for _ in items]) + def best_inputs(self): + return self.inputs.values[np.nanargmax(self.scalarized_fitness)] - # @property - # def go_to_best_sum_of_tasks(self): - # yield from self.go_to(self.best_sum_of_tasks_inputs) + def go_to(self, inputs): + args = [] + for device, value in zip(self._subset_devices(kind="active"), np.atleast_1d(inputs).T): + args.append(device) + args.append(value) + yield from bps.mv(*args) + + def go_to_best(self): + yield from self.go_to(self.best_inputs) def plot_tasks(self, **kwargs): if self._len_subset_dofs(kind="active", mode="on") == 1: diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index 73ae31e..a919cc9 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -1,18 +1,71 @@ import math - import bluesky.plans as bp +import bluesky.plan_stubs as bps import numpy as np import torch +import time from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement, UpperConfidenceBound from botorch.acquisition.max_value_entropy_search import qLowerBoundMaxValueEntropy from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement +# def default_acquisition_plan(dofs, inputs, dets): + +# args = [] +# for dof, points in zip(dofs, np.atleast_2d(inputs).T): +# args.append(dof) +# args.append(list(points)) + +# uid = yield from bp.list_scan(dets, *args) +# return uid + + +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): - uid = yield from bp.list_scan(dets, *[_ for items in zip(dofs, np.atleast_2d(inputs).T) for _ in items]) + + 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=1) 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 + + ACQ_FUNC_CONFIG = { "quasi-random": { "identifiers": ["qr", "quasi-random"], diff --git a/bloptools/bayesian/digestion.py b/bloptools/bayesian/digestion.py index 147d2d3..eed3199 100644 --- a/bloptools/bayesian/digestion.py +++ b/bloptools/bayesian/digestion.py @@ -1,2 +1,4 @@ def default_digestion_function(db, uid): - return db[uid].table(fill=True) + products = db[uid].table(fill=True) + print(products) + return products From fa95af0e7c14e67ec4bebf86055276c7cfc7973b Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 7 Sep 2023 09:59:23 -0400 Subject: [PATCH 19/43] use device name for dofs --- bloptools/bayesian/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index ad765dd..5aa9f9f 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -43,13 +43,14 @@ def _validate_and_prepare_dofs(dofs): - for dof in dofs: + for i_dof, dof in enumerate(dofs): if not isinstance(dof, Mapping): raise ValueError("Supplied dofs must be an iterable of mappings (e.g. a dict)!") if "device" not in dof.keys(): raise ValueError("Each DOF must have a device!") dof["device"].kind = "hinted" + dof["name"] = dof["device"].name if hasattr(dof["device"], "name") else f"x{i_dof+1}" if "limits" not in dof.keys(): dof["limits"] = (-np.inf, np.inf) From 6af10f20cfa53c117b6159aeeb3edc4b24a5d315 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 7 Sep 2023 13:33:05 -0400 Subject: [PATCH 20/43] optional trigger delay on default acquisition plan --- bloptools/bayesian/__init__.py | 4 +++- bloptools/bayesian/acquisition.py | 6 ++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 5aa9f9f..b1995d6 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -151,6 +151,8 @@ def __init__( self.digestion = kwargs.get("digestion", default_digestion_function) self.dets = list(np.atleast_1d(kwargs.get("dets", []))) + self.trigger_delay = kwargs.get("trigger_delay", 0) + self.acq_func_config = kwargs.get("acq_func_config", ACQ_FUNC_CONFIG) self.table = pd.DataFrame() @@ -657,7 +659,7 @@ def acquire(self, active_inputs): passive_devices = [*self._subset_devices(kind="passive"), *self._subset_devices(kind="active", mode="off")] uid = yield from self.acquisition_plan( - active_devices, active_inputs.astype(float), [*self.dets, *passive_devices] + active_devices, active_inputs.astype(float), [*self.dets, *passive_devices], delay=self.trigger_delay, ) products = self.digestion(self.db, uid) diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index a919cc9..f98be87 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -35,14 +35,16 @@ def one_nd_step_with_delay(detectors, step, pos_cache): return uid -def default_acquisition_plan(dofs, inputs, dets): +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=1) + uid = yield from list_scan_with_delay(dets, *args, delay=delay) return uid From 5335698a1c0aaa6e30a1658002866591d606147e Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 7 Sep 2023 13:56:40 -0400 Subject: [PATCH 21/43] tweaks --- bloptools/bayesian/__init__.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index b1995d6..3362157 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -604,8 +604,8 @@ def ask(self, acq_func_identifier="ei", n=1, route=True, return_metadata=False): active_X = np.concatenate(active_x_list, axis=0) self.forget(self.table.index[-(n - 1) :]) - if route: - active_X = active_X[utils.route(self._read_subset_devices(kind="active", mode="on"), active_X)] + if route: + active_X = active_X[utils.route(self._read_subset_devices(kind="active", mode="on"), active_X)] return (active_X, acq_func_meta) if return_metadata else active_X @@ -659,7 +659,10 @@ def acquire(self, active_inputs): passive_devices = [*self._subset_devices(kind="passive"), *self._subset_devices(kind="active", mode="off")] uid = yield from self.acquisition_plan( - active_devices, active_inputs.astype(float), [*self.dets, *passive_devices], delay=self.trigger_delay, + active_devices, + active_inputs.astype(float), + [*self.dets, *passive_devices], + delay=self.trigger_delay, ) products = self.digestion(self.db, uid) From 0dc720ff27f8170876b79333105084b6f273cad1 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 11 Sep 2023 11:19:47 -0400 Subject: [PATCH 22/43] tweak CI --- .github/workflows/docs.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index ae35ed6..169fea7 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -62,8 +62,10 @@ jobs: # For reference: https://www.gnu.org/software/bash/manual/html_node/The-Set-Builtin.html. set -vxeo pipefail + df -h conda env list - mamba install -c conda-forge shadow3 srwpy pandoc + pip install shadow3 srwpy + apt install pandoc pip install --upgrade pip wheel pip install -v . pip install -r requirements-dev.txt From 701e7b914e0ce777d8585d9c975330ba4966168c Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 11 Sep 2023 11:33:15 -0400 Subject: [PATCH 23/43] use native python in CI --- .github/workflows/docs.yml | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 169fea7..1ddd28f 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -47,25 +47,13 @@ jobs: mkdir -v -p ~/.config/databroker/ wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml - - name: Set up Python ${{ matrix.python-version }} with conda - uses: conda-incubator/setup-miniconda@v2 - with: - activate-environment: ${{ env.REPOSITORY_NAME }}-py${{ matrix.python-version }} - auto-update-conda: true - miniconda-version: "latest" - python-version: ${{ matrix.python-version }} - mamba-version: "*" - channels: conda-forge - - name: Install documentation-building requirements run: | # For reference: https://www.gnu.org/software/bash/manual/html_node/The-Set-Builtin.html. set -vxeo pipefail - df -h conda env list - pip install shadow3 srwpy - apt install pandoc + mamba install -c conda-forge shadow3 srwpy pandoc pip install --upgrade pip wheel pip install -v . pip install -r requirements-dev.txt From 96eaf65e614ae086b3fb0729e3af72906061fabb Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 11 Sep 2023 12:35:51 -0400 Subject: [PATCH 24/43] no sirepo in CI for now --- .github/workflows/docs.yml | 35 ++++++++++++------- docs/source/tutorials.rst | 2 -- .../tutorials/constrained-himmelblau.ipynb | 4 ++- .../latent-toroid-dimensions.ipynb | 0 .../tutorials => wip}/multi-task-sirepo.ipynb | 0 requirements-dev.txt | 3 +- 6 files changed, 28 insertions(+), 16 deletions(-) rename docs/{source/tutorials => wip}/latent-toroid-dimensions.ipynb (100%) rename docs/{source/tutorials => wip}/multi-task-sirepo.ipynb (100%) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 1ddd28f..68b9542 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -33,19 +33,29 @@ jobs: with: fetch-depth: 0 - - name: Start MongoDB - uses: supercharge/mongodb-github-action@1.6.0 + # - name: Start MongoDB + # uses: supercharge/mongodb-github-action@1.6.0 - - name: Start Sirepo Docker container - uses: NSLS-II/start-sirepo-action@v2 - with: - docker-binary: docker + # - name: Start Sirepo Docker container + # uses: NSLS-II/start-sirepo-action@v2 + # with: + # docker-binary: docker - - name: Copy databroker config file - run: | - set -vxeuo pipefail - mkdir -v -p ~/.config/databroker/ - wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml + # - name: Copy databroker config file + # run: | + # set -vxeuo pipefail + # mkdir -v -p ~/.config/databroker/ + # wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml + + - name: Set up Python ${{ matrix.python-version }} with conda + uses: conda-incubator/setup-miniconda@v2 + with: + activate-environment: ${{ env.REPOSITORY_NAME }}-py${{ matrix.python-version }} + auto-update-conda: true + miniconda-version: "latest" + python-version: ${{ matrix.python-version }} + mamba-version: "*" + channels: conda-forge - name: Install documentation-building requirements run: | @@ -53,7 +63,8 @@ jobs: set -vxeo pipefail conda env list - mamba install -c conda-forge shadow3 srwpy pandoc + # mamba install -c conda-forge shadow3 srwpy + mamba install -c conda-forge pandoc pip install --upgrade pip wheel pip install -v . pip install -r requirements-dev.txt diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst index 3a57df0..297b303 100644 --- a/docs/source/tutorials.rst +++ b/docs/source/tutorials.rst @@ -8,5 +8,3 @@ Tutorials tutorials/constrained-himmelblau.ipynb tutorials/hyperparameters.ipynb tutorials/passive-dofs.ipynb - tutorials/latent-toroid-dimensions.ipynb - tutorials/multi-task-sirepo.ipynb diff --git a/docs/source/tutorials/constrained-himmelblau.ipynb b/docs/source/tutorials/constrained-himmelblau.ipynb index b16d59f..95f863d 100644 --- a/docs/source/tutorials/constrained-himmelblau.ipynb +++ b/docs/source/tutorials/constrained-himmelblau.ipynb @@ -87,7 +87,9 @@ }, "outputs": [], "source": [ - "%run -i ../../../examples/prepare_bluesky.py # prepare the bluesky environment\n", + "from bloptools.utils import prepare_re_env\n", + "\n", + "%run -i $prepare_re_env.__file__ --db-type=temp\n", "\n", "from bloptools import devices\n", "from bloptools.bayesian import Agent\n", diff --git a/docs/source/tutorials/latent-toroid-dimensions.ipynb b/docs/wip/latent-toroid-dimensions.ipynb similarity index 100% rename from docs/source/tutorials/latent-toroid-dimensions.ipynb rename to docs/wip/latent-toroid-dimensions.ipynb diff --git a/docs/source/tutorials/multi-task-sirepo.ipynb b/docs/wip/multi-task-sirepo.ipynb similarity index 100% rename from docs/source/tutorials/multi-task-sirepo.ipynb rename to docs/wip/multi-task-sirepo.ipynb diff --git a/requirements-dev.txt b/requirements-dev.txt index 29110b8..2ec27e7 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -18,6 +18,7 @@ jupyter matplotlib nbsphinx numpydoc -sirepo-bluesky +pandoc +# sirepo-bluesky sphinx-copybutton sphinx_rtd_theme From 2bb18b29c2220a9c4fda1bf5756725366aa395f5 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 11 Sep 2023 12:36:21 -0400 Subject: [PATCH 25/43] no sirepo in CI for now --- bloptools/bayesian/acquisition.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index a919cc9..c0caa4a 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -1,14 +1,13 @@ import math -import bluesky.plans as bp + import bluesky.plan_stubs as bps +import bluesky.plans as bp import numpy as np import torch -import time from botorch.acquisition.analytic import LogExpectedImprovement, LogProbabilityOfImprovement, UpperConfidenceBound from botorch.acquisition.max_value_entropy_search import qLowerBoundMaxValueEntropy from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement - # def default_acquisition_plan(dofs, inputs, dets): # args = [] @@ -30,13 +29,12 @@ def one_nd_step_with_delay(detectors, 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) + 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): - args = [] for dof, points in zip(dofs, np.atleast_2d(inputs).T): args.append(dof) @@ -62,7 +60,7 @@ def default_acquisition_plan(dofs, inputs, dets): # yield from bps.mv(*args) # yield from bps.count([*dets, *dofs]) # yield from bps.sleep(1) - + # return uid From dd7ab75f7ee5b9bef097c91d248ed4ea0133eb5a Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 12 Sep 2023 09:17:20 -0400 Subject: [PATCH 26/43] some plots work --- bloptools/bayesian/__init__.py | 137 +++++++++++++----------- bloptools/bayesian/acquisition.py | 12 +-- bloptools/bayesian/plots.py | 55 ++++++++++ docs/wip/latent-toroid-dimensions.ipynb | 2 +- 4 files changed, 135 insertions(+), 71 deletions(-) create mode 100644 bloptools/bayesian/plots.py diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index ad765dd..2a5ef5b 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -22,7 +22,7 @@ from matplotlib.patches import Patch from .. import utils -from . import acquisition, models +from . import acquisition, models, plots from .acquisition import ACQ_FUNC_CONFIG, default_acquisition_plan from .digestion import default_digestion_function @@ -81,6 +81,8 @@ def _validate_and_prepare_tasks(tasks): raise ValueError("Supplied tasks must be an iterable of mappings (e.g. a dict)!") if task["kind"] not in ["minimize", "maximize"]: raise ValueError('"mode" must be specified as either "minimize" or "maximize"') + if "name" not in task.keys(): + task["name"] = task["key"] if "weight" not in task.keys(): task["weight"] = 1 if "limits" not in task.keys(): @@ -158,6 +160,8 @@ def __init__( self._train_models = True self.a_priori_hypers = None + self.plots = {} + def _subset_inputs_sampler(self, kind=None, mode=None, n=MAX_TEST_INPUTS): """ Returns $n$ quasi-randomly sampled inputs in the bounded parameter space @@ -656,7 +660,7 @@ def acquire(self, active_inputs): passive_devices = [*self._subset_devices(kind="passive"), *self._subset_devices(kind="active", mode="off")] uid = yield from self.acquisition_plan( - active_devices, active_inputs.astype(float), [*self.dets, *passive_devices] + active_devices, active_inputs.astype(float), [*self.dets, *passive_devices], ) products = self.digestion(self.db, uid) @@ -706,6 +710,10 @@ def learn( self.tell(new_table=new_table, reuse_hypers=reuse_hypers) + for plot in self.plots: + if plot.live: + plot.update() + @property def inputs(self): return self.table.loc[:, self._subset_dof_names(mode="on")].astype(float) @@ -724,11 +732,59 @@ def go_to(self, inputs): def go_to_best(self): yield from self.go_to(self.best_inputs) - def plot_tasks(self, **kwargs): + + + + def plot_tasks(self, live=True) + + def plot_tasks(self, live=False, **kwargs): + + self.plots["tasks"] = {} + if self._len_subset_dofs(kind="active", mode="on") == 1: - self._plot_tasks_one_dof(**kwargs) + self.tasks_plot = plots.TasksPlotOneDOF(self, live=live) else: - self._plot_tasks_many_dofs(**kwargs) + self.tasks_plot = plots.TasksPlotManyDOFs(self, live=live) + + + def _update_task_plots_many_dofs(self, axes=[0, 1]): + + gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + + for itask, task in enumerate(self.tasks): + + task_plots = self.plots["tasks"][task["name"]] + + sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) + task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) + task_norm = mpl.colors.Normalize(task_vmin, task_vmax) + task_plots["sampled"].set_norm(task_norm) + task_plots["sampled"].set_norm(task_norm) + + task_plots["sampled"].set_offsets(self.inputs.values[:, axes]) + task_plots["sampled"].set_array(sampled_fitness) + + x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + task_posterior = task["model"].posterior(x) + task_mean = task_posterior.mean + task_sigma = task_posterior.variance.sqrt() + + if gridded: + if not x.ndim == 3: + raise ValueError() + task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) + task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) + + else: + task_plots["pred_mean"].set_offsets(x.detach().numpy()[..., axes]) + task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) + task_plots["pred_sigma"].set_offsets(x.detach().numpy()[..., axes]) + task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) + + for ax in self.task_axes.ravel(): + ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) + def _plot_tasks_one_dof(self, size=16, lw=1e0): self.task_fig, self.task_axes = plt.subplots( @@ -742,6 +798,9 @@ def _plot_tasks_one_dof(self, size=16, lw=1e0): self.task_axes = np.atleast_1d(self.task_axes) for itask, task in enumerate(self.tasks): + + task_plots = self.plots["tasks"][task["name"]] = {} + color = DEFAULT_COLOR_LIST[itask] self.task_axes[itask].set_ylabel(task["key"]) @@ -751,12 +810,7 @@ def _plot_tasks_one_dof(self, size=16, lw=1e0): task_mean = task_posterior.mean.detach().numpy() task_sigma = task_posterior.variance.sqrt().detach().numpy() - self.task_axes[itask].scatter( - self.inputs.loc[:, self._subset_dof_names(kind="active", mode="on")], - self.table.loc[:, f'{task["key"]}_fitness'], - s=size, - color=color, - ) + task_plots["sampled"] = self.task_axes[itask].scatter([], [], s=size, color=color) on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] @@ -789,13 +843,12 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") for itask, task in enumerate(self.tasks): - sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) - task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) - task_norm = mpl.colors.Normalize(task_vmin, task_vmax) - # if task["transform"] == "log": - # task_norm = mpl.colors.LogNorm(task_vmin, task_vmax) - # else: + task_plots = self.plots["tasks"][task["name"]] = {} + + # sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) + # task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) + # task_norm = mpl.colors.Normalize() self.task_axes[itask, 0].set_ylabel(f'{task["key"]}_fitness') @@ -803,58 +856,22 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL self.task_axes[itask, 1].set_title("posterior mean") self.task_axes[itask, 2].set_title("posterior std. dev.") - data_ax = self.task_axes[itask, 0].scatter( - *self.inputs.values.T[axes], s=size, c=sampled_fitness, norm=task_norm, cmap=cmap - ) + task_plots["sampled"] = self.task_axes[itask, 0].scatter([], [], s=size, cmap=cmap) x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - task_posterior = task["model"].posterior(x) - task_mean = task_posterior.mean - task_sigma = task_posterior.variance.sqrt() - if gridded: if not x.ndim == 3: raise ValueError() - self.task_axes[itask, 1].pcolormesh( - x[..., 0], - x[..., 1], - task_mean[..., 0].detach().numpy(), - shading=shading, - cmap=cmap, - norm=task_norm, - ) - sigma_ax = self.task_axes[itask, 2].pcolormesh( - x[..., 0], - x[..., 1], - task_sigma[..., 0].detach().numpy(), - shading=shading, - cmap=cmap, - ) + task_plots["pred_mean"] = self.task_axes[itask, 1].imshow([[0]], cmap=cmap) + task_plots["pred_sigma"] = self.task_axes[itask, 2].imshow([[0]], cmap=cmap) else: - self.task_axes[itask, 1].scatter( - x.detach().numpy()[..., axes[0]], - x.detach().numpy()[..., axes[1]], - c=task_mean[..., 0].detach().numpy(), - s=size, - norm=task_norm, - cmap=cmap, - ) - sigma_ax = self.task_axes[itask, 2].scatter( - x.detach().numpy()[..., axes[0]], - x.detach().numpy()[..., axes[1]], - c=task_sigma[..., 0].detach().numpy(), - s=size, - cmap=cmap, - ) + task_plots["pred_mean"] = self.task_axes[itask, 1].scatter([], [], s=size, cmap=cmap) + task_plots["pred_sigma"] = self.task_axes[itask, 2].scatter([], [], s=size, cmap=cmap) - self.task_fig.colorbar(data_ax, ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) - self.task_fig.colorbar(sigma_ax, ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) - - for ax in self.task_axes.ravel(): - ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) - ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) + task_plots["colorbar_mean"] = self.task_fig.colorbar(task_plots["sampled"], ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) + task_plots["colorbar_sigma"] = self.task_fig.colorbar(task_plots["pred_sigma"], ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) def plot_acquisition(self, acq_funcs=["ei"], **kwargs): if self._len_subset_dofs(kind="active", mode="on") == 1: diff --git a/bloptools/bayesian/acquisition.py b/bloptools/bayesian/acquisition.py index c0caa4a..46e9f02 100644 --- a/bloptools/bayesian/acquisition.py +++ b/bloptools/bayesian/acquisition.py @@ -8,20 +8,12 @@ from botorch.acquisition.max_value_entropy_search import qLowerBoundMaxValueEntropy from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement -# def default_acquisition_plan(dofs, inputs, dets): - -# args = [] -# for dof, points in zip(dofs, np.atleast_2d(inputs).T): -# args.append(dof) -# args.append(list(points)) - -# uid = yield from bp.list_scan(dets, *args) -# return uid - def list_scan_with_delay(*args, delay=0, **kwargs): "Accepts all the normal 'scan' parameters, plus an optional delay." + delay = 0 + 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() diff --git a/bloptools/bayesian/plots.py b/bloptools/bayesian/plots.py new file mode 100644 index 0000000..2ce1244 --- /dev/null +++ b/bloptools/bayesian/plots.py @@ -0,0 +1,55 @@ +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np + + +class TasksPlotManyDOFs: + def __init__(self): + self.task_fig, self.task_axes = plt.subplots( + self.n_tasks, + 1, + figsize=(6, 4 * self.n_tasks), + sharex=True, + constrained_layout=True, + ) + + self.axes = [0, 1] + + self.update() + + def update(self, dofs, tasks): + gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + + for itask, task in enumerate(self.tasks): + task_plots = self.plots["tasks"][task["name"]] + + sampled_fitness = self.table.loc[:, f'{task["key"]}_fitness'].values + task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) + task_norm = mpl.colors.Normalize(task_vmin, task_vmax) + task_plots["sampled"].set_norm(task_norm) + task_plots["sampled"].set_norm(task_norm) + + task_plots["sampled"].set_offsets(self.inputs.values[:, self.axes]) + task_plots["sampled"].set_array(sampled_fitness) + + x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=1024) + + task_posterior = task["model"].posterior(x) + task_mean = task_posterior.mean + task_sigma = task_posterior.variance.sqrt() + + if gridded: + if not x.ndim == 3: + raise ValueError() + task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) + task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) + + else: + task_plots["pred_mean"].set_offsets(x.detach().numpy()[..., self.axes]) + task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) + task_plots["pred_sigma"].set_offsets(x.detach().numpy()[..., self.axes]) + task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) + + for ax in self.task_axes.ravel(): + ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[self.axes[0]]["limits"]) + ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[self.axes[1]]["limits"]) diff --git a/docs/wip/latent-toroid-dimensions.ipynb b/docs/wip/latent-toroid-dimensions.ipynb index 8eb04b8..dd93bd1 100644 --- a/docs/wip/latent-toroid-dimensions.ipynb +++ b/docs/wip/latent-toroid-dimensions.ipynb @@ -106,7 +106,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.12" + "version": "3.11.5" }, "vscode": { "interpreter": { From 7d1dcace8bf79e62783db1800835f9304fa969e0 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 13 Sep 2023 10:55:53 -0400 Subject: [PATCH 27/43] plot work --- bloptools/bayesian/__init__.py | 544 ++++++++++++++++----------------- bloptools/bayesian/plots.py | 84 +++-- 2 files changed, 327 insertions(+), 301 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 2a5ef5b..b61cf09 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -32,8 +32,6 @@ mpl.rc("image", cmap="coolwarm") -DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen", "goldenrod"] -DEFAULT_COLORMAP = "turbo" MAX_TEST_INPUTS = 2**11 @@ -735,371 +733,369 @@ def go_to_best(self): - def plot_tasks(self, live=True) - def plot_tasks(self, live=False, **kwargs): - self.plots["tasks"] = {} - if self._len_subset_dofs(kind="active", mode="on") == 1: self.tasks_plot = plots.TasksPlotOneDOF(self, live=live) else: self.tasks_plot = plots.TasksPlotManyDOFs(self, live=live) - def _update_task_plots_many_dofs(self, axes=[0, 1]): + # def _update_task_plots_many_dofs(self, axes=[0, 1]): - gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + # gridded = self._len_subset_dofs(kind="active", mode="on") == 2 - for itask, task in enumerate(self.tasks): + # for itask, task in enumerate(self.tasks): - task_plots = self.plots["tasks"][task["name"]] + # task_plots = self.plots["tasks"][task["name"]] - sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) - task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) - task_norm = mpl.colors.Normalize(task_vmin, task_vmax) - task_plots["sampled"].set_norm(task_norm) - task_plots["sampled"].set_norm(task_norm) + # sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) + # task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) + # task_norm = mpl.colors.Normalize(task_vmin, task_vmax) + # task_plots["sampled"].set_norm(task_norm) + # task_plots["sampled"].set_norm(task_norm) - task_plots["sampled"].set_offsets(self.inputs.values[:, axes]) - task_plots["sampled"].set_array(sampled_fitness) + # task_plots["sampled"].set_offsets(self.inputs.values[:, axes]) + # task_plots["sampled"].set_array(sampled_fitness) - x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - task_posterior = task["model"].posterior(x) - task_mean = task_posterior.mean - task_sigma = task_posterior.variance.sqrt() + # x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + # task_posterior = task["model"].posterior(x) + # task_mean = task_posterior.mean + # task_sigma = task_posterior.variance.sqrt() - if gridded: - if not x.ndim == 3: - raise ValueError() - task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) - task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) + # if gridded: + # if not x.ndim == 3: + # raise ValueError() + # task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) + # task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) - else: - task_plots["pred_mean"].set_offsets(x.detach().numpy()[..., axes]) - task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) - task_plots["pred_sigma"].set_offsets(x.detach().numpy()[..., axes]) - task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) - - for ax in self.task_axes.ravel(): - ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) - ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) - - - def _plot_tasks_one_dof(self, size=16, lw=1e0): - self.task_fig, self.task_axes = plt.subplots( - self.num_tasks, - 1, - figsize=(6, 4 * self.num_tasks), - sharex=True, - constrained_layout=True, - ) + # else: + # task_plots["pred_mean"].set_offsets(x.detach().numpy()[..., axes]) + # task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) + # task_plots["pred_sigma"].set_offsets(x.detach().numpy()[..., axes]) + # task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) - self.task_axes = np.atleast_1d(self.task_axes) + # for ax in self.task_axes.ravel(): + # ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + # ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) - for itask, task in enumerate(self.tasks): - task_plots = self.plots["tasks"][task["name"]] = {} + # def _plot_tasks_one_dof(self, size=16, lw=1e0): + # self.task_fig, self.task_axes = plt.subplots( + # self.num_tasks, + # 1, + # figsize=(6, 4 * self.num_tasks), + # sharex=True, + # constrained_layout=True, + # ) - color = DEFAULT_COLOR_LIST[itask] + # self.task_axes = np.atleast_1d(self.task_axes) - self.task_axes[itask].set_ylabel(task["key"]) + # for itask, task in enumerate(self.tasks): - x = self.test_inputs_grid - task_posterior = task["model"].posterior(x) - task_mean = task_posterior.mean.detach().numpy() - task_sigma = task_posterior.variance.sqrt().detach().numpy() + # task_plots = self.plots["tasks"][task["name"]] = {} - task_plots["sampled"] = self.task_axes[itask].scatter([], [], s=size, color=color) + # color = DEFAULT_COLOR_LIST[itask] - on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] + # self.task_axes[itask].set_ylabel(task["key"]) - for z in [0, 1, 2]: - self.task_axes[itask].fill_between( - x[..., on_dofs_are_active_mask].squeeze(), - (task_mean - z * task_sigma).squeeze(), - (task_mean + z * task_sigma).squeeze(), - lw=lw, - color=color, - alpha=0.5**z, - ) + # x = self.test_inputs_grid + # task_posterior = task["model"].posterior(x) + # task_mean = task_posterior.mean.detach().numpy() + # task_sigma = task_posterior.variance.sqrt().detach().numpy() - self.task_axes[itask].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) + # task_plots["sampled"] = self.task_axes[itask].scatter([], [], s=size, color=color) - def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16): - if gridded is None: - gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + # on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - self.task_fig, self.task_axes = plt.subplots( - self.num_tasks, - 3, - figsize=(10, 4 * self.num_tasks), - sharex=True, - sharey=True, - constrained_layout=True, - ) + # for z in [0, 1, 2]: + # self.task_axes[itask].fill_between( + # x[..., on_dofs_are_active_mask].squeeze(), + # (task_mean - z * task_sigma).squeeze(), + # (task_mean + z * task_sigma).squeeze(), + # lw=lw, + # color=color, + # alpha=0.5**z, + # ) - self.task_axes = np.atleast_2d(self.task_axes) - # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") + # self.task_axes[itask].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) - for itask, task in enumerate(self.tasks): + # def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16): + # if gridded is None: + # gridded = self._len_subset_dofs(kind="active", mode="on") == 2 - task_plots = self.plots["tasks"][task["name"]] = {} + # self.task_fig, self.task_axes = plt.subplots( + # self.num_tasks, + # 3, + # figsize=(10, 4 * self.num_tasks), + # sharex=True, + # sharey=True, + # constrained_layout=True, + # ) - # sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) - # task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) - # task_norm = mpl.colors.Normalize() + # self.task_axes = np.atleast_2d(self.task_axes) + # # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - self.task_axes[itask, 0].set_ylabel(f'{task["key"]}_fitness') + # for itask, task in enumerate(self.tasks): - self.task_axes[itask, 0].set_title("samples") - self.task_axes[itask, 1].set_title("posterior mean") - self.task_axes[itask, 2].set_title("posterior std. dev.") + # task_plots = self.plots["tasks"][task["name"]] = {} - task_plots["sampled"] = self.task_axes[itask, 0].scatter([], [], s=size, cmap=cmap) + # # sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) + # # task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) + # # task_norm = mpl.colors.Normalize() - x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + # self.task_axes[itask, 0].set_ylabel(f'{task["key"]}_fitness') - if gridded: - if not x.ndim == 3: - raise ValueError() - task_plots["pred_mean"] = self.task_axes[itask, 1].imshow([[0]], cmap=cmap) - task_plots["pred_sigma"] = self.task_axes[itask, 2].imshow([[0]], cmap=cmap) + # self.task_axes[itask, 0].set_title("samples") + # self.task_axes[itask, 1].set_title("posterior mean") + # self.task_axes[itask, 2].set_title("posterior std. dev.") - else: - task_plots["pred_mean"] = self.task_axes[itask, 1].scatter([], [], s=size, cmap=cmap) - task_plots["pred_sigma"] = self.task_axes[itask, 2].scatter([], [], s=size, cmap=cmap) + # task_plots["sampled"] = self.task_axes[itask, 0].scatter([], [], s=size, cmap=cmap) - task_plots["colorbar_mean"] = self.task_fig.colorbar(task_plots["sampled"], ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) - task_plots["colorbar_sigma"] = self.task_fig.colorbar(task_plots["pred_sigma"], ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) + # x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - def plot_acquisition(self, acq_funcs=["ei"], **kwargs): - if self._len_subset_dofs(kind="active", mode="on") == 1: - self._plot_acq_one_dof(acq_funcs=acq_funcs, **kwargs) + # if gridded: + # if not x.ndim == 3: + # raise ValueError() + # task_plots["pred_mean"] = self.task_axes[itask, 1].imshow([[0]], cmap=cmap) + # task_plots["pred_sigma"] = self.task_axes[itask, 2].imshow([[0]], cmap=cmap) - else: - self._plot_acq_many_dofs(acq_funcs=acq_funcs, **kwargs) - - def _plot_acq_one_dof(self, acq_funcs, lw=1e0, **kwargs): - self.acq_fig, self.acq_axes = plt.subplots( - 1, - len(acq_funcs), - figsize=(4 * len(acq_funcs), 4), - sharex=True, - constrained_layout=True, - ) + # else: + # task_plots["pred_mean"] = self.task_axes[itask, 1].scatter([], [], s=size, cmap=cmap) + # task_plots["pred_sigma"] = self.task_axes[itask, 2].scatter([], [], s=size, cmap=cmap) - self.acq_axes = np.atleast_1d(self.acq_axes) + # task_plots["colorbar_mean"] = self.task_fig.colorbar(task_plots["sampled"], ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) + # task_plots["colorbar_sigma"] = self.task_fig.colorbar(task_plots["pred_sigma"], ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) - for iacq_func, acq_func_identifier in enumerate(acq_funcs): - color = DEFAULT_COLOR_LIST[iacq_func] + # def plot_acquisition(self, acq_funcs=["ei"], **kwargs): + # if self._len_subset_dofs(kind="active", mode="on") == 1: + # self._plot_acq_one_dof(acq_funcs=acq_funcs, **kwargs) - acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) + # else: + # self._plot_acq_many_dofs(acq_funcs=acq_funcs, **kwargs) - x = self.test_inputs_grid - *input_shape, input_dim = x.shape - obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) + # def _plot_acq_one_dof(self, acq_funcs, lw=1e0, **kwargs): + # self.acq_fig, self.acq_axes = plt.subplots( + # 1, + # len(acq_funcs), + # figsize=(4 * len(acq_funcs), 4), + # sharex=True, + # constrained_layout=True, + # ) - if acq_func_identifier in ["ei", "pi"]: - obj = obj.exp() + # self.acq_axes = np.atleast_1d(self.acq_axes) - self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + # for iacq_func, acq_func_identifier in enumerate(acq_funcs): + # color = DEFAULT_COLOR_LIST[iacq_func] - on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - self.acq_axes[iacq_func].plot( - x[..., on_dofs_are_active_mask].squeeze(), obj.detach().numpy(), lw=lw, color=color - ) + # acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) - self.acq_axes[iacq_func].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) + # x = self.test_inputs_grid + # *input_shape, input_dim = x.shape + # obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) - def _plot_acq_many_dofs( - self, acq_funcs, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16, **kwargs - ): - self.acq_fig, self.acq_axes = plt.subplots( - 1, - len(acq_funcs), - figsize=(4 * len(acq_funcs), 4), - sharex=True, - sharey=True, - constrained_layout=True, - ) + # if acq_func_identifier in ["ei", "pi"]: + # obj = obj.exp() - if gridded is None: - gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + # self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) - self.acq_axes = np.atleast_1d(self.acq_axes) - # self.acq_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") + # on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] + # self.acq_axes[iacq_func].plot( + # x[..., on_dofs_are_active_mask].squeeze(), obj.detach().numpy(), lw=lw, color=color + # ) - x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - *input_shape, input_dim = x.shape + # self.acq_axes[iacq_func].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) - for iacq_func, acq_func_identifier in enumerate(acq_funcs): - acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) + # def _plot_acq_many_dofs( + # self, acq_funcs, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16, **kwargs + # ): + # self.acq_fig, self.acq_axes = plt.subplots( + # 1, + # len(acq_funcs), + # figsize=(4 * len(acq_funcs), 4), + # sharex=True, + # sharey=True, + # constrained_layout=True, + # ) - obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) - if acq_func_identifier in ["ei", "pi"]: - obj = obj.exp() + # if gridded is None: + # gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + + # self.acq_axes = np.atleast_1d(self.acq_axes) + # # self.acq_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - if gridded: - self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) - obj_ax = self.acq_axes[iacq_func].pcolormesh( - x[..., 0], - x[..., 1], - obj.detach().numpy(), - shading=shading, - cmap=cmap, - ) + # x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + # *input_shape, input_dim = x.shape - self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) + # for iacq_func, acq_func_identifier in enumerate(acq_funcs): + # acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) - else: - self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) - obj_ax = self.acq_axes[iacq_func].scatter( - x.detach().numpy()[..., axes[0]], - x.detach().numpy()[..., axes[1]], - c=obj.detach().numpy(), - ) + # obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) + # if acq_func_identifier in ["ei", "pi"]: + # obj = obj.exp() - self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) + # if gridded: + # self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + # obj_ax = self.acq_axes[iacq_func].pcolormesh( + # x[..., 0], + # x[..., 1], + # obj.detach().numpy(), + # shading=shading, + # cmap=cmap, + # ) - for ax in self.acq_axes.ravel(): - ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) - ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) + # self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) - def plot_validity(self, **kwargs): - if self._len_subset_dofs(kind="active", mode="on") == 1: - self._plot_valid_one_dof(**kwargs) + # else: + # self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + # obj_ax = self.acq_axes[iacq_func].scatter( + # x.detach().numpy()[..., axes[0]], + # x.detach().numpy()[..., axes[1]], + # c=obj.detach().numpy(), + # ) - else: - self._plot_valid_many_dofs(**kwargs) + # self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) + + # for ax in self.acq_axes.ravel(): + # ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + # ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) - def _plot_valid_one_dof(self, size=16, lw=1e0): - self.valid_fig, self.valid_ax = plt.subplots(1, 1, figsize=(4, 4), sharex=True, constrained_layout=True) + # def plot_validity(self, **kwargs): + # if self._len_subset_dofs(kind="active", mode="on") == 1: + # self._plot_valid_one_dof(**kwargs) - x = self.test_inputs_grid - *input_shape, input_dim = x.shape - constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) + # else: + # self._plot_valid_many_dofs(**kwargs) - self.valid_ax.scatter(self.inputs.values, self.all_tasks_valid, s=size) + # def _plot_valid_one_dof(self, size=16, lw=1e0): + # self.valid_fig, self.valid_ax = plt.subplots(1, 1, figsize=(4, 4), sharex=True, constrained_layout=True) - on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] + # x = self.test_inputs_grid + # *input_shape, input_dim = x.shape + # constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) - self.valid_ax.plot(x[..., on_dofs_are_active_mask].squeeze(), constraint.detach().numpy(), lw=lw) + # self.valid_ax.scatter(self.inputs.values, self.all_tasks_valid, s=size) + + # on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - self.valid_ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[0]["limits"]) + # self.valid_ax.plot(x[..., on_dofs_are_active_mask].squeeze(), constraint.detach().numpy(), lw=lw) - def _plot_valid_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=16, gridded=None): - self.valid_fig, self.valid_axes = plt.subplots( - 1, 2, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True - ) + # self.valid_ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[0]["limits"]) - if gridded is None: - gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + # def _plot_valid_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=16, gridded=None): + # self.valid_fig, self.valid_axes = plt.subplots( + # 1, 2, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True + # ) - data_ax = self.valid_axes[0].scatter( - *self.inputs.values.T[:2], - c=self.all_tasks_valid, - s=size, - vmin=0, - vmax=1, - cmap=cmap, - ) + # if gridded is None: + # gridded = self._len_subset_dofs(kind="active", mode="on") == 2 - x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - *input_shape, input_dim = x.shape - constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) - - if gridded: - self.valid_axes[1].pcolormesh( - x[..., 0], - x[..., 1], - constraint.detach().numpy(), - shading=shading, - cmap=cmap, - vmin=0, - vmax=1, - ) + # data_ax = self.valid_axes[0].scatter( + # *self.inputs.values.T[:2], + # c=self.all_tasks_valid, + # s=size, + # vmin=0, + # vmax=1, + # cmap=cmap, + # ) - # self.acq_fig.colorbar(obj_ax, ax=self.valid_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) + # x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + # *input_shape, input_dim = x.shape + # constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) - else: - # self.valid_axes.set_title(acq_func_meta["name"]) - self.valid_axes[1].scatter( - x.detach().numpy()[..., axes[0]], - x.detach().numpy()[..., axes[1]], - c=constraint.detach().numpy(), - ) + # if gridded: + # self.valid_axes[1].pcolormesh( + # x[..., 0], + # x[..., 1], + # constraint.detach().numpy(), + # shading=shading, + # cmap=cmap, + # vmin=0, + # vmax=1, + # ) - self.valid_fig.colorbar(data_ax, ax=self.valid_axes[:2], location="bottom", aspect=32, shrink=0.8) + # # self.acq_fig.colorbar(obj_ax, ax=self.valid_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) - for ax in self.valid_axes.ravel(): - ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) - ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) + # else: + # # self.valid_axes.set_title(acq_func_meta["name"]) + # self.valid_axes[1].scatter( + # x.detach().numpy()[..., axes[0]], + # x.detach().numpy()[..., axes[1]], + # c=constraint.detach().numpy(), + # ) - def inspect_beam(self, index, border=None): - im = self.images[index] + # self.plot.update(self) - x_min, x_max, y_min, y_max, width_x, width_y = self.table.loc[ - index, ["x_min", "x_max", "y_min", "y_max", "width_x", "width_y"] - ] + # self.valid_fig.colorbar(data_ax, ax=self.valid_axes[:2], location="bottom", aspect=32, shrink=0.8) - bbx = np.array([x_min, x_max])[[0, 0, 1, 1, 0]] - bby = np.array([y_min, y_max])[[0, 1, 1, 0, 0]] + # for ax in self.valid_axes.ravel(): + # ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + # ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) - plt.figure() - plt.imshow(im, cmap="gray_r") - plt.plot(bbx, bby, lw=4e0, c="r") + # def inspect_beam(self, index, border=None): + # im = self.images[index] - if border is not None: - plt.xlim(x_min - border * width_x, x_min + border * width_x) - plt.ylim(y_min - border * width_y, y_min + border * width_y) + # x_min, x_max, y_min, y_max, width_x, width_y = self.table.loc[ + # index, ["x_min", "x_max", "y_min", "y_max", "width_x", "width_y"] + # ] - def plot_history(self, x_key="index", show_all_tasks=False): - x = getattr(self.table, x_key).values + # bbx = np.array([x_min, x_max])[[0, 0, 1, 1, 0]] + # bby = np.array([y_min, y_max])[[0, 1, 1, 0, 0]] - num_task_plots = 1 - if show_all_tasks: - num_task_plots = self.num_tasks + 1 + # plt.figure() + # plt.imshow(im, cmap="gray_r") + # plt.plot(bbx, bby, lw=4e0, c="r") - self.num_tasks + 1 if self.num_tasks > 1 else 1 + # if border is not None: + # plt.xlim(x_min - border * width_x, x_min + border * width_x) + # plt.ylim(y_min - border * width_y, y_min + border * width_y) - hist_fig, hist_axes = plt.subplots( - num_task_plots, 1, figsize=(6, 4 * num_task_plots), sharex=True, constrained_layout=True, dpi=200 - ) - hist_axes = np.atleast_1d(hist_axes) + # def plot_history(self, x_key="index", show_all_tasks=False): + # x = getattr(self.table, x_key).values - unique_strategies, acq_func_index, acq_func_inverse = np.unique( - self.table.acq_func, return_index=True, return_inverse=True - ) + # num_task_plots = 1 + # if show_all_tasks: + # num_task_plots = self.num_tasks + 1 + + # self.num_tasks + 1 if self.num_tasks > 1 else 1 + + # hist_fig, hist_axes = plt.subplots( + # num_task_plots, 1, figsize=(6, 4 * num_task_plots), sharex=True, constrained_layout=True, dpi=200 + # ) + # hist_axes = np.atleast_1d(hist_axes) + + # unique_strategies, acq_func_index, acq_func_inverse = np.unique( + # self.table.acq_func, return_index=True, return_inverse=True + # ) - sample_colors = np.array(DEFAULT_COLOR_LIST)[acq_func_inverse] + # sample_colors = np.array(DEFAULT_COLOR_LIST)[acq_func_inverse] - if show_all_tasks: - for itask, task in enumerate(self.tasks): - y = self.table.loc[:, f'{task["key"]}_fitness'].values - hist_axes[itask].scatter(x, y, c=sample_colors) - hist_axes[itask].plot(x, y, lw=5e-1, c="k") - hist_axes[itask].set_ylabel(task["key"]) + # if show_all_tasks: + # for itask, task in enumerate(self.tasks): + # y = self.table.loc[:, f'{task["key"]}_fitness'].values + # hist_axes[itask].scatter(x, y, c=sample_colors) + # hist_axes[itask].plot(x, y, lw=5e-1, c="k") + # hist_axes[itask].set_ylabel(task["key"]) - y = self.scalarized_fitness + # y = self.scalarized_fitness - cummax_y = np.array([np.nanmax(y[: i + 1]) for i in range(len(y))]) + # cummax_y = np.array([np.nanmax(y[: i + 1]) for i in range(len(y))]) - hist_axes[-1].scatter(x, y, c=sample_colors) - hist_axes[-1].plot(x, y, lw=5e-1, c="k") + # hist_axes[-1].scatter(x, y, c=sample_colors) + # hist_axes[-1].plot(x, y, lw=5e-1, c="k") - hist_axes[-1].plot(x, cummax_y, lw=5e-1, c="k", ls=":") + # hist_axes[-1].plot(x, cummax_y, lw=5e-1, c="k", ls=":") - hist_axes[-1].set_ylabel("total_fitness") - hist_axes[-1].set_xlabel(x_key) + # hist_axes[-1].set_ylabel("total_fitness") + # hist_axes[-1].set_xlabel(x_key) - handles = [] - for i_acq_func, acq_func in enumerate(unique_strategies): - # i_acq_func = np.argsort(acq_func_index)[i_handle] - handles.append(Patch(color=DEFAULT_COLOR_LIST[i_acq_func], label=acq_func)) - legend = hist_axes[0].legend(handles=handles, fontsize=8) - legend.set_title("acquisition function") + # handles = [] + # for i_acq_func, acq_func in enumerate(unique_strategies): + # # i_acq_func = np.argsort(acq_func_index)[i_handle] + # handles.append(Patch(color=DEFAULT_COLOR_LIST[i_acq_func], label=acq_func)) + # legend = hist_axes[0].legend(handles=handles, fontsize=8) + # legend.set_title("acquisition function") - # plot_history(self, x_key="time") + # # plot_history(self, x_key="time") - # plt.savefig("bo-history.pdf", dpi=256) + # # plt.savefig("bo-history.pdf", dpi=256) diff --git a/bloptools/bayesian/plots.py b/bloptools/bayesian/plots.py index 2ce1244..da922a9 100644 --- a/bloptools/bayesian/plots.py +++ b/bloptools/bayesian/plots.py @@ -3,53 +3,83 @@ import numpy as np +DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen", "goldenrod"] +DEFAULT_COLORMAP = "turbo" +DEFAULT_SCATTER_SIZE = 16 + class TasksPlotManyDOFs: - def __init__(self): + + def __init__(self, agent, **kwargs): + + self.cmap = kwargs.get("cmap", DEFAULT_COLORMAP) + self.s = kwargs.get("s", DEFAULT_SCATTER_SIZE) + self.gridded = kwargs.get("gridded", False) + self.axes = kwargs.get("axes", [0, 1]) + + self.make_plots(agent) + + def make_plots(self, agent): + self.task_fig, self.task_axes = plt.subplots( - self.n_tasks, - 1, - figsize=(6, 4 * self.n_tasks), + len(agent.tasks), + 3, + figsize=(12, 4 * len(agent.tasks)), sharex=True, constrained_layout=True, ) - self.axes = [0, 1] + self.task_axes = np.atleast_2d(self.task_axes) - self.update() + for itask in range(len(agent.tasks)): - def update(self, dofs, tasks): - gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + self.task_axes[itask, 0].scatter([], [], s=self.s, cmap=self.cmap) - for itask, task in enumerate(self.tasks): - task_plots = self.plots["tasks"][task["name"]] + if self.gridded: + self.task_axes[itask, 1].imshow([[0]], cmap=self.cmap) + self.task_axes[itask, 2].imshow([[0]], cmap=self.cmap) - sampled_fitness = self.table.loc[:, f'{task["key"]}_fitness'].values + else: + self.task_axes[itask, 1].scatter([], [], s=self.s, cmap=self.cmap) + self.task_axes[itask, 2].scatter([], [], s=self.s, cmap=self.cmap) + + self.task_fig.colorbar(self.task_axes[itask, 0].collections[0], ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) + self.task_fig.colorbar(self.task_axes[itask, 2].collections[0], ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) + + + def update(self, agent): + + for itask, task in enumerate(agent.tasks): + + sampled_fitness = agent.table.loc[:, f'{task["key"]}_fitness'].values task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) task_norm = mpl.colors.Normalize(task_vmin, task_vmax) - task_plots["sampled"].set_norm(task_norm) - task_plots["sampled"].set_norm(task_norm) - task_plots["sampled"].set_offsets(self.inputs.values[:, self.axes]) - task_plots["sampled"].set_array(sampled_fitness) + data_scatter = self.task_axes[itask, 0].collections[0] + pred_mean = self.task_axes[itask, 1].collections[0] + pred_error = self.task_axes[itask, 2].collections[0] + + data_scatter.set_offsets(agent.inputs.values[:, self.axes]) + data_scatter.set_array(sampled_fitness) + data_scatter.set_norm(task_norm) - x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=1024) + x = agent.test_inputs_grid if self.gridded else agent.test_inputs(n=1024).squeeze() task_posterior = task["model"].posterior(x) - task_mean = task_posterior.mean - task_sigma = task_posterior.variance.sqrt() + task_pred_mean = task_posterior.mean.squeeze() + task_pred_sigma = task_posterior.variance.sqrt().squeeze() - if gridded: + if self.gridded: if not x.ndim == 3: raise ValueError() - task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) - task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) + pred_mean.set_array(task_pred_mean.detach()) + pred_error.set_array(task_pred_sigma.detach()) else: - task_plots["pred_mean"].set_offsets(x.detach().numpy()[..., self.axes]) - task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) - task_plots["pred_sigma"].set_offsets(x.detach().numpy()[..., self.axes]) - task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) + pred_mean.set_offsets(x.detach()[..., self.axes]) + pred_error.set_offsets(x.detach()[..., self.axes]) + pred_mean.set_array(task_pred_mean.detach()) + pred_error.set_array(task_pred_sigma.detach()) for ax in self.task_axes.ravel(): - ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[self.axes[0]]["limits"]) - ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[self.axes[1]]["limits"]) + ax.set_xlim(*agent._subset_dofs(kind="active", mode="on")[self.axes[0]]["limits"]) + ax.set_ylim(*agent._subset_dofs(kind="active", mode="on")[self.axes[1]]["limits"]) From 76a7010d73071d4ae5a9fdbf3cac53dc0d49348e Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 14 Sep 2023 14:44:22 -0400 Subject: [PATCH 28/43] able to specify which dimensions to treat as latent --- bloptools/bayesian/__init__.py | 570 ++++++++++++------------ bloptools/bayesian/plots.py | 85 ---- bloptools/test_functions.py | 28 ++ bloptools/tests/test_bayesian_shadow.py | 41 -- 4 files changed, 304 insertions(+), 420 deletions(-) delete mode 100644 bloptools/bayesian/plots.py delete mode 100644 bloptools/tests/test_bayesian_shadow.py diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 4168fc0..0974541 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -1,6 +1,7 @@ import logging import os import time as ttime +import uuid import warnings from collections import OrderedDict from collections.abc import Mapping @@ -22,7 +23,7 @@ from matplotlib.patches import Patch from .. import utils -from . import acquisition, models, plots +from . import acquisition, models from .acquisition import ACQ_FUNC_CONFIG, default_acquisition_plan from .digestion import default_digestion_function @@ -39,6 +40,10 @@ TASK_TRANSFORMS = {"log": lambda x: np.log(x)} +DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen", "goldenrod"] +DEFAULT_COLORMAP = "turbo" +DEFAULT_SCATTER_SIZE = 16 + def _validate_and_prepare_dofs(dofs): for i_dof, dof in enumerate(dofs): @@ -59,6 +64,7 @@ def _validate_and_prepare_dofs(dofs): if dof["kind"] not in ["active", "passive"]: raise ValueError('DOF kinds must be one of "active" or "passive"') + # active dofs are on by default, passive dofs are off by default dof["mode"] = dof.get("mode", "on" if dof["kind"] == "active" else "off") if dof["mode"] not in ["on", "off"]: raise ValueError('DOF modes must be one of "on" or "off"') @@ -161,14 +167,14 @@ def __init__( self._train_models = True self.a_priori_hypers = None - self.plots = {} + self.plots = {"tasks": {}} - def _subset_inputs_sampler(self, kind=None, mode=None, n=MAX_TEST_INPUTS): + def reset(self): """ - Returns $n$ quasi-randomly sampled inputs in the bounded parameter space + Reset the agent. """ - transform = self._subset_input_transform(kind=kind, mode=mode) - return transform.untransform(utils.normalized_sobol_sampler(n, d=self._len_subset_dofs(kind=kind, mode=mode))) + self.table = pd.DataFrame() + self._initialized = False def initialize( self, @@ -217,7 +223,7 @@ def tell(self, new_table=None, append=True, train=True, **kwargs): self.table = pd.concat([self.table, new_table]) if append else new_table self.table.index = np.arange(len(self.table)) - skew_dims = [tuple(np.arange(self._len_subset_dofs(mode="on")))] + skew_dims = self.latent_dim_tuples if self._initialized: cached_hypers = self.hypers @@ -412,6 +418,17 @@ def _subset_dof_limits(self, kind=None, mode=None): return torch.tensor([dof["limits"] for dof in dofs_subset], dtype=torch.float64).T return torch.empty((2, 0)) + @property + def latent_dim_tuples(self): + """ + Returns a list of tuples, where each tuple represent a group of dimension to find a latent representation of. + """ + + latent_dim_labels = [dof.get("latent_group", str(uuid.uuid4())) for dof in self._subset_dofs(mode="on")] + u, uinv = np.unique(latent_dim_labels, return_inverse=True) + + return [tuple(np.where(uinv == i)[0]) for i in range(len(u))] + def test_inputs(self, n=MAX_TEST_INPUTS): return utils.sobol_sampler(self._acq_func_bounds, n=n) @@ -423,6 +440,13 @@ def _subset_input_transform(self, kind=None, mode=None): d=limits.shape[-1], coefficient=coefficient, offset=offset ) + def _subset_inputs_sampler(self, kind=None, mode=None, n=MAX_TEST_INPUTS): + """ + Returns $n$ quasi-randomly sampled inputs in the bounded parameter space + """ + transform = self._subset_input_transform(kind=kind, mode=mode) + return transform.untransform(utils.normalized_sobol_sampler(n, d=self._len_subset_dofs(kind=kind, mode=mode))) + def save_data(self, filepath="./self_data.h5"): """ Save the sampled inputs and targets of the agent to a file, which can be used @@ -714,10 +738,6 @@ def learn( self.tell(new_table=new_table, reuse_hypers=reuse_hypers) - for plot in self.plots: - if plot.live: - plot.update() - @property def inputs(self): return self.table.loc[:, self._subset_dof_names(mode="on")].astype(float) @@ -736,366 +756,328 @@ def go_to(self, inputs): def go_to_best(self): yield from self.go_to(self.best_inputs) - def plot_tasks(self, live=False, **kwargs): + def plot_tasks(self, **kwargs): if self._len_subset_dofs(kind="active", mode="on") == 1: - self.tasks_plot = plots.TasksPlotOneDOF(self, live=live) + self._plot_tasks_one_dof(**kwargs) else: - self.tasks_plot = plots.TasksPlotManyDOFs(self, live=live) - - # def _update_task_plots_many_dofs(self, axes=[0, 1]): - - # gridded = self._len_subset_dofs(kind="active", mode="on") == 2 - - # for itask, task in enumerate(self.tasks): - - # task_plots = self.plots["tasks"][task["name"]] - - # sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) - # task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) - # task_norm = mpl.colors.Normalize(task_vmin, task_vmax) - # task_plots["sampled"].set_norm(task_norm) - # task_plots["sampled"].set_norm(task_norm) - - # task_plots["sampled"].set_offsets(self.inputs.values[:, axes]) - # task_plots["sampled"].set_array(sampled_fitness) - - # x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - # task_posterior = task["model"].posterior(x) - # task_mean = task_posterior.mean - # task_sigma = task_posterior.variance.sqrt() - - # if gridded: - # if not x.ndim == 3: - # raise ValueError() - # task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) - # task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) - - # else: - # task_plots["pred_mean"].set_offsets(x.detach().numpy()[..., axes]) - # task_plots["pred_mean"].set_array(task_mean[..., 0].detach().numpy()) - # task_plots["pred_sigma"].set_offsets(x.detach().numpy()[..., axes]) - # task_plots["pred_sigma"].set_array(task_sigma[..., 0].detach().numpy()) - - # for ax in self.task_axes.ravel(): - # ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) - # ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) - - # def _plot_tasks_one_dof(self, size=16, lw=1e0): - # self.task_fig, self.task_axes = plt.subplots( - # self.num_tasks, - # 1, - # figsize=(6, 4 * self.num_tasks), - # sharex=True, - # constrained_layout=True, - # ) - - # self.task_axes = np.atleast_1d(self.task_axes) + self._plot_tasks_many_dofs(**kwargs) + + def _plot_tasks_one_dof(self, size=16, lw=1e0): + self.task_fig, self.task_axes = plt.subplots( + self.num_tasks, + 1, + figsize=(6, 4 * self.num_tasks), + sharex=True, + constrained_layout=True, + ) - # for itask, task in enumerate(self.tasks): + self.task_axes = np.atleast_1d(self.task_axes) - # task_plots = self.plots["tasks"][task["name"]] = {} + for itask, task in enumerate(self.tasks): + task_plots = self.plots["tasks"][task["name"]] = {} - # color = DEFAULT_COLOR_LIST[itask] + color = DEFAULT_COLOR_LIST[itask] - # self.task_axes[itask].set_ylabel(task["key"]) + self.task_axes[itask].set_ylabel(task["key"]) - # x = self.test_inputs_grid - # task_posterior = task["model"].posterior(x) - # task_mean = task_posterior.mean.detach().numpy() - # task_sigma = task_posterior.variance.sqrt().detach().numpy() + x = self.test_inputs_grid + task_posterior = task["model"].posterior(x) + task_mean = task_posterior.mean.detach().numpy() + task_sigma = task_posterior.variance.sqrt().detach().numpy() - # task_plots["sampled"] = self.task_axes[itask].scatter([], [], s=size, color=color) + task_plots["sampled"] = self.task_axes[itask].scatter([], [], s=size, color=color) - # on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] + on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - # for z in [0, 1, 2]: - # self.task_axes[itask].fill_between( - # x[..., on_dofs_are_active_mask].squeeze(), - # (task_mean - z * task_sigma).squeeze(), - # (task_mean + z * task_sigma).squeeze(), - # lw=lw, - # color=color, - # alpha=0.5**z, - # ) + for z in [0, 1, 2]: + self.task_axes[itask].fill_between( + x[..., on_dofs_are_active_mask].squeeze(), + (task_mean - z * task_sigma).squeeze(), + (task_mean + z * task_sigma).squeeze(), + lw=lw, + color=color, + alpha=0.5**z, + ) - # self.task_axes[itask].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) + self.task_axes[itask].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) - # def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16): - # if gridded is None: - # gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16): + if gridded is None: + gridded = self._len_subset_dofs(kind="active", mode="on") == 2 - # self.task_fig, self.task_axes = plt.subplots( - # self.num_tasks, - # 3, - # figsize=(10, 4 * self.num_tasks), - # sharex=True, - # sharey=True, - # constrained_layout=True, - # ) + self.task_fig, self.task_axes = plt.subplots( + self.num_tasks, + 3, + figsize=(10, 4 * self.num_tasks), + sharex=True, + sharey=True, + constrained_layout=True, + ) - # self.task_axes = np.atleast_2d(self.task_axes) - # # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") + self.task_axes = np.atleast_2d(self.task_axes) + # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - # for itask, task in enumerate(self.tasks): + for itask, task in enumerate(self.tasks): + task_plots = self.plots["tasks"][task["name"]] = {} - # task_plots = self.plots["tasks"][task["name"]] = {} + # sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) + # task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) + # task_norm = mpl.colors.Normalize() - # # sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) - # # task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) - # # task_norm = mpl.colors.Normalize() + self.task_axes[itask, 0].set_ylabel(f'{task["key"]}_fitness') - # self.task_axes[itask, 0].set_ylabel(f'{task["key"]}_fitness') + self.task_axes[itask, 0].set_title("samples") + self.task_axes[itask, 1].set_title("posterior mean") + self.task_axes[itask, 2].set_title("posterior std. dev.") - # self.task_axes[itask, 0].set_title("samples") - # self.task_axes[itask, 1].set_title("posterior mean") - # self.task_axes[itask, 2].set_title("posterior std. dev.") + task_plots["sampled"] = self.task_axes[itask, 0].scatter([], [], s=size, cmap=cmap) - # task_plots["sampled"] = self.task_axes[itask, 0].scatter([], [], s=size, cmap=cmap) - - # x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - # if gridded: - # if not x.ndim == 3: - # raise ValueError() - # task_plots["pred_mean"] = self.task_axes[itask, 1].imshow([[0]], cmap=cmap) - # task_plots["pred_sigma"] = self.task_axes[itask, 2].imshow([[0]], cmap=cmap) + if gridded: + if not x.ndim == 3: + raise ValueError() + task_plots["pred_mean"] = self.task_axes[itask, 1].imshow([[0]], cmap=cmap) + task_plots["pred_sigma"] = self.task_axes[itask, 2].imshow([[0]], cmap=cmap) - # else: - # task_plots["pred_mean"] = self.task_axes[itask, 1].scatter([], [], s=size, cmap=cmap) - # task_plots["pred_sigma"] = self.task_axes[itask, 2].scatter([], [], s=size, cmap=cmap) + else: + task_plots["pred_mean"] = self.task_axes[itask, 1].scatter([], [], s=size, cmap=cmap) + task_plots["pred_sigma"] = self.task_axes[itask, 2].scatter([], [], s=size, cmap=cmap) - # task_plots["colorbar_mean"] = self.task_fig.colorbar(task_plots["sampled"], ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) - # task_plots["colorbar_sigma"] = self.task_fig.colorbar(task_plots["pred_sigma"], ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) + task_plots["colorbar_mean"] = self.task_fig.colorbar( + task_plots["sampled"], ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8 + ) + task_plots["colorbar_sigma"] = self.task_fig.colorbar( + task_plots["pred_sigma"], ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8 + ) - # def plot_acquisition(self, acq_funcs=["ei"], **kwargs): - # if self._len_subset_dofs(kind="active", mode="on") == 1: - # self._plot_acq_one_dof(acq_funcs=acq_funcs, **kwargs) + def plot_acquisition(self, acq_funcs=["ei"], **kwargs): + if self._len_subset_dofs(kind="active", mode="on") == 1: + self._plot_acq_one_dof(acq_funcs=acq_funcs, **kwargs) - # else: - # self._plot_acq_many_dofs(acq_funcs=acq_funcs, **kwargs) + else: + self._plot_acq_many_dofs(acq_funcs=acq_funcs, **kwargs) + + def _plot_acq_one_dof(self, acq_funcs, lw=1e0, **kwargs): + self.acq_fig, self.acq_axes = plt.subplots( + 1, + len(acq_funcs), + figsize=(4 * len(acq_funcs), 4), + sharex=True, + constrained_layout=True, + ) - # def _plot_acq_one_dof(self, acq_funcs, lw=1e0, **kwargs): - # self.acq_fig, self.acq_axes = plt.subplots( - # 1, - # len(acq_funcs), - # figsize=(4 * len(acq_funcs), 4), - # sharex=True, - # constrained_layout=True, - # ) + self.acq_axes = np.atleast_1d(self.acq_axes) - # self.acq_axes = np.atleast_1d(self.acq_axes) + for iacq_func, acq_func_identifier in enumerate(acq_funcs): + color = DEFAULT_COLOR_LIST[iacq_func] - # for iacq_func, acq_func_identifier in enumerate(acq_funcs): - # color = DEFAULT_COLOR_LIST[iacq_func] + acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) - # acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) + x = self.test_inputs_grid + *input_shape, input_dim = x.shape + obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) - # x = self.test_inputs_grid - # *input_shape, input_dim = x.shape - # obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) + if acq_func_identifier in ["ei", "pi"]: + obj = obj.exp() - # if acq_func_identifier in ["ei", "pi"]: - # obj = obj.exp() + self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) - # self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] + self.acq_axes[iacq_func].plot( + x[..., on_dofs_are_active_mask].squeeze(), obj.detach().numpy(), lw=lw, color=color + ) - # on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - # self.acq_axes[iacq_func].plot( - # x[..., on_dofs_are_active_mask].squeeze(), obj.detach().numpy(), lw=lw, color=color - # ) + self.acq_axes[iacq_func].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) - # self.acq_axes[iacq_func].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) + def _plot_acq_many_dofs( + self, acq_funcs, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16, **kwargs + ): + self.acq_fig, self.acq_axes = plt.subplots( + 1, + len(acq_funcs), + figsize=(4 * len(acq_funcs), 4), + sharex=True, + sharey=True, + constrained_layout=True, + ) - # def _plot_acq_many_dofs( - # self, acq_funcs, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16, **kwargs - # ): - # self.acq_fig, self.acq_axes = plt.subplots( - # 1, - # len(acq_funcs), - # figsize=(4 * len(acq_funcs), 4), - # sharex=True, - # sharey=True, - # constrained_layout=True, - # ) + if gridded is None: + gridded = self._len_subset_dofs(kind="active", mode="on") == 2 - # if gridded is None: - # gridded = self._len_subset_dofs(kind="active", mode="on") == 2 - - # self.acq_axes = np.atleast_1d(self.acq_axes) - # # self.acq_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") + self.acq_axes = np.atleast_1d(self.acq_axes) + # self.acq_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - # x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - # *input_shape, input_dim = x.shape + x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + *input_shape, input_dim = x.shape - # for iacq_func, acq_func_identifier in enumerate(acq_funcs): - # acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) + for iacq_func, acq_func_identifier in enumerate(acq_funcs): + acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) - # obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) - # if acq_func_identifier in ["ei", "pi"]: - # obj = obj.exp() + obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) + if acq_func_identifier in ["ei", "pi"]: + obj = obj.exp() - # if gridded: - # self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) - # obj_ax = self.acq_axes[iacq_func].pcolormesh( - # x[..., 0], - # x[..., 1], - # obj.detach().numpy(), - # shading=shading, - # cmap=cmap, - # ) + if gridded: + self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + obj_ax = self.acq_axes[iacq_func].pcolormesh( + x[..., 0], + x[..., 1], + obj.detach().numpy(), + shading=shading, + cmap=cmap, + ) - # self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) + self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) - # else: - # self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) - # obj_ax = self.acq_axes[iacq_func].scatter( - # x.detach().numpy()[..., axes[0]], - # x.detach().numpy()[..., axes[1]], - # c=obj.detach().numpy(), - # ) + else: + self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + obj_ax = self.acq_axes[iacq_func].scatter( + x.detach().numpy()[..., axes[0]], + x.detach().numpy()[..., axes[1]], + c=obj.detach().numpy(), + ) - # self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) - - # for ax in self.acq_axes.ravel(): - # ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) - # ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) + self.acq_fig.colorbar(obj_ax, ax=self.acq_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) - # def plot_validity(self, **kwargs): - # if self._len_subset_dofs(kind="active", mode="on") == 1: - # self._plot_valid_one_dof(**kwargs) + for ax in self.acq_axes.ravel(): + ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) - # else: - # self._plot_valid_many_dofs(**kwargs) + def plot_validity(self, **kwargs): + if self._len_subset_dofs(kind="active", mode="on") == 1: + self._plot_valid_one_dof(**kwargs) - # def _plot_valid_one_dof(self, size=16, lw=1e0): - # self.valid_fig, self.valid_ax = plt.subplots(1, 1, figsize=(4, 4), sharex=True, constrained_layout=True) + else: + self._plot_valid_many_dofs(**kwargs) - # x = self.test_inputs_grid - # *input_shape, input_dim = x.shape - # constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) + def _plot_valid_one_dof(self, size=16, lw=1e0): + self.valid_fig, self.valid_ax = plt.subplots(1, 1, figsize=(4, 4), sharex=True, constrained_layout=True) - # self.valid_ax.scatter(self.inputs.values, self.all_tasks_valid, s=size) - - # on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] + x = self.test_inputs_grid + *input_shape, input_dim = x.shape + constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) - # self.valid_ax.plot(x[..., on_dofs_are_active_mask].squeeze(), constraint.detach().numpy(), lw=lw) + self.valid_ax.scatter(self.inputs.values, self.all_tasks_valid, s=size) - # self.valid_ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[0]["limits"]) + on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - # def _plot_valid_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=16, gridded=None): - # self.valid_fig, self.valid_axes = plt.subplots( - # 1, 2, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True - # ) + self.valid_ax.plot(x[..., on_dofs_are_active_mask].squeeze(), constraint.detach().numpy(), lw=lw) - # if gridded is None: - # gridded = self._len_subset_dofs(kind="active", mode="on") == 2 + self.valid_ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[0]["limits"]) - # data_ax = self.valid_axes[0].scatter( - # *self.inputs.values.T[:2], - # c=self.all_tasks_valid, - # s=size, - # vmin=0, - # vmax=1, - # cmap=cmap, - # ) + def _plot_valid_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, size=16, gridded=None): + self.valid_fig, self.valid_axes = plt.subplots( + 1, 2, figsize=(8, 4), sharex=True, sharey=True, constrained_layout=True + ) - # x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) - # *input_shape, input_dim = x.shape - # constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) + if gridded is None: + gridded = self._len_subset_dofs(kind="active", mode="on") == 2 - # if gridded: - # self.valid_axes[1].pcolormesh( - # x[..., 0], - # x[..., 1], - # constraint.detach().numpy(), - # shading=shading, - # cmap=cmap, - # vmin=0, - # vmax=1, - # ) + data_ax = self.valid_axes[0].scatter( + *self.inputs.values.T[:2], + c=self.all_tasks_valid, + s=size, + vmin=0, + vmax=1, + cmap=cmap, + ) - # # self.acq_fig.colorbar(obj_ax, ax=self.valid_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) + x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + *input_shape, input_dim = x.shape + constraint = self.classifier.probabilities(x.reshape(-1, 1, input_dim))[..., -1].reshape(input_shape) + + if gridded: + self.valid_axes[1].pcolormesh( + x[..., 0], + x[..., 1], + constraint.detach().numpy(), + shading=shading, + cmap=cmap, + vmin=0, + vmax=1, + ) - # else: - # # self.valid_axes.set_title(acq_func_meta["name"]) - # self.valid_axes[1].scatter( - # x.detach().numpy()[..., axes[0]], - # x.detach().numpy()[..., axes[1]], - # c=constraint.detach().numpy(), - # ) + # self.acq_fig.colorbar(obj_ax, ax=self.valid_axes[iacq_func], location="bottom", aspect=32, shrink=0.8) - # self.plot.update(self) + else: + # self.valid_axes.set_title(acq_func_meta["name"]) + self.valid_axes[1].scatter( + x.detach().numpy()[..., axes[0]], + x.detach().numpy()[..., axes[1]], + c=constraint.detach().numpy(), + ) - # self.valid_fig.colorbar(data_ax, ax=self.valid_axes[:2], location="bottom", aspect=32, shrink=0.8) + self.valid_fig.colorbar(data_ax, ax=self.valid_axes[:2], location="bottom", aspect=32, shrink=0.8) - # for ax in self.valid_axes.ravel(): - # ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) - # ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) + for ax in self.valid_axes.ravel(): + ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) - # def inspect_beam(self, index, border=None): - # im = self.images[index] + def inspect_beam(self, index, border=None): + im = self.images[index] - # x_min, x_max, y_min, y_max, width_x, width_y = self.table.loc[ - # index, ["x_min", "x_max", "y_min", "y_max", "width_x", "width_y"] - # ] + x_min, x_max, y_min, y_max, width_x, width_y = self.table.loc[ + index, ["x_min", "x_max", "y_min", "y_max", "width_x", "width_y"] + ] - # bbx = np.array([x_min, x_max])[[0, 0, 1, 1, 0]] - # bby = np.array([y_min, y_max])[[0, 1, 1, 0, 0]] + bbx = np.array([x_min, x_max])[[0, 0, 1, 1, 0]] + bby = np.array([y_min, y_max])[[0, 1, 1, 0, 0]] - # plt.figure() - # plt.imshow(im, cmap="gray_r") - # plt.plot(bbx, bby, lw=4e0, c="r") + plt.figure() + plt.imshow(im, cmap="gray_r") + plt.plot(bbx, bby, lw=4e0, c="r") - # if border is not None: - # plt.xlim(x_min - border * width_x, x_min + border * width_x) - # plt.ylim(y_min - border * width_y, y_min + border * width_y) + if border is not None: + plt.xlim(x_min - border * width_x, x_min + border * width_x) + plt.ylim(y_min - border * width_y, y_min + border * width_y) - # def plot_history(self, x_key="index", show_all_tasks=False): - # x = getattr(self.table, x_key).values + def plot_history(self, x_key="index", show_all_tasks=False): + x = getattr(self.table, x_key).values - # num_task_plots = 1 - # if show_all_tasks: - # num_task_plots = self.num_tasks + 1 + num_task_plots = 1 + if show_all_tasks: + num_task_plots = self.num_tasks + 1 - # self.num_tasks + 1 if self.num_tasks > 1 else 1 + self.num_tasks + 1 if self.num_tasks > 1 else 1 - # hist_fig, hist_axes = plt.subplots( - # num_task_plots, 1, figsize=(6, 4 * num_task_plots), sharex=True, constrained_layout=True, dpi=200 - # ) - # hist_axes = np.atleast_1d(hist_axes) + hist_fig, hist_axes = plt.subplots( + num_task_plots, 1, figsize=(6, 4 * num_task_plots), sharex=True, constrained_layout=True, dpi=200 + ) + hist_axes = np.atleast_1d(hist_axes) - # unique_strategies, acq_func_index, acq_func_inverse = np.unique( - # self.table.acq_func, return_index=True, return_inverse=True - # ) + unique_strategies, acq_func_index, acq_func_inverse = np.unique( + self.table.acq_func, return_index=True, return_inverse=True + ) - # sample_colors = np.array(DEFAULT_COLOR_LIST)[acq_func_inverse] + sample_colors = np.array(DEFAULT_COLOR_LIST)[acq_func_inverse] - # if show_all_tasks: - # for itask, task in enumerate(self.tasks): - # y = self.table.loc[:, f'{task["key"]}_fitness'].values - # hist_axes[itask].scatter(x, y, c=sample_colors) - # hist_axes[itask].plot(x, y, lw=5e-1, c="k") - # hist_axes[itask].set_ylabel(task["key"]) + if show_all_tasks: + for itask, task in enumerate(self.tasks): + y = self.table.loc[:, f'{task["key"]}_fitness'].values + hist_axes[itask].scatter(x, y, c=sample_colors) + hist_axes[itask].plot(x, y, lw=5e-1, c="k") + hist_axes[itask].set_ylabel(task["key"]) - # y = self.scalarized_fitness + y = self.scalarized_fitness - # cummax_y = np.array([np.nanmax(y[: i + 1]) for i in range(len(y))]) + cummax_y = np.array([np.nanmax(y[: i + 1]) for i in range(len(y))]) - # hist_axes[-1].scatter(x, y, c=sample_colors) - # hist_axes[-1].plot(x, y, lw=5e-1, c="k") + hist_axes[-1].scatter(x, y, c=sample_colors) + hist_axes[-1].plot(x, y, lw=5e-1, c="k") - # hist_axes[-1].plot(x, cummax_y, lw=5e-1, c="k", ls=":") + hist_axes[-1].plot(x, cummax_y, lw=5e-1, c="k", ls=":") - # hist_axes[-1].set_ylabel("total_fitness") - # hist_axes[-1].set_xlabel(x_key) + hist_axes[-1].set_ylabel("total_fitness") + hist_axes[-1].set_xlabel(x_key) - # handles = [] - # for i_acq_func, acq_func in enumerate(unique_strategies): - # # i_acq_func = np.argsort(acq_func_index)[i_handle] - # handles.append(Patch(color=DEFAULT_COLOR_LIST[i_acq_func], label=acq_func)) - # legend = hist_axes[0].legend(handles=handles, fontsize=8) - # legend.set_title("acquisition function") + handles = [] + for i_acq_func, acq_func in enumerate(unique_strategies): + # i_acq_func = np.argsort(acq_func_index)[i_handle] + handles.append(Patch(color=DEFAULT_COLOR_LIST[i_acq_func], label=acq_func)) + legend = hist_axes[0].legend(handles=handles, fontsize=8) + legend.set_title("acquisition function") - # # plot_history(self, x_key="time") + # plot_history(self, x_key="time") - # # plt.savefig("bo-history.pdf", dpi=256) + # plt.savefig("bo-history.pdf", dpi=256) diff --git a/bloptools/bayesian/plots.py b/bloptools/bayesian/plots.py deleted file mode 100644 index da922a9..0000000 --- a/bloptools/bayesian/plots.py +++ /dev/null @@ -1,85 +0,0 @@ -import matplotlib as mpl -import matplotlib.pyplot as plt -import numpy as np - - -DEFAULT_COLOR_LIST = ["dodgerblue", "tomato", "mediumseagreen", "goldenrod"] -DEFAULT_COLORMAP = "turbo" -DEFAULT_SCATTER_SIZE = 16 - -class TasksPlotManyDOFs: - - def __init__(self, agent, **kwargs): - - self.cmap = kwargs.get("cmap", DEFAULT_COLORMAP) - self.s = kwargs.get("s", DEFAULT_SCATTER_SIZE) - self.gridded = kwargs.get("gridded", False) - self.axes = kwargs.get("axes", [0, 1]) - - self.make_plots(agent) - - def make_plots(self, agent): - - self.task_fig, self.task_axes = plt.subplots( - len(agent.tasks), - 3, - figsize=(12, 4 * len(agent.tasks)), - sharex=True, - constrained_layout=True, - ) - - self.task_axes = np.atleast_2d(self.task_axes) - - for itask in range(len(agent.tasks)): - - self.task_axes[itask, 0].scatter([], [], s=self.s, cmap=self.cmap) - - if self.gridded: - self.task_axes[itask, 1].imshow([[0]], cmap=self.cmap) - self.task_axes[itask, 2].imshow([[0]], cmap=self.cmap) - - else: - self.task_axes[itask, 1].scatter([], [], s=self.s, cmap=self.cmap) - self.task_axes[itask, 2].scatter([], [], s=self.s, cmap=self.cmap) - - self.task_fig.colorbar(self.task_axes[itask, 0].collections[0], ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) - self.task_fig.colorbar(self.task_axes[itask, 2].collections[0], ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) - - - def update(self, agent): - - for itask, task in enumerate(agent.tasks): - - sampled_fitness = agent.table.loc[:, f'{task["key"]}_fitness'].values - task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) - task_norm = mpl.colors.Normalize(task_vmin, task_vmax) - - data_scatter = self.task_axes[itask, 0].collections[0] - pred_mean = self.task_axes[itask, 1].collections[0] - pred_error = self.task_axes[itask, 2].collections[0] - - data_scatter.set_offsets(agent.inputs.values[:, self.axes]) - data_scatter.set_array(sampled_fitness) - data_scatter.set_norm(task_norm) - - x = agent.test_inputs_grid if self.gridded else agent.test_inputs(n=1024).squeeze() - - task_posterior = task["model"].posterior(x) - task_pred_mean = task_posterior.mean.squeeze() - task_pred_sigma = task_posterior.variance.sqrt().squeeze() - - if self.gridded: - if not x.ndim == 3: - raise ValueError() - pred_mean.set_array(task_pred_mean.detach()) - pred_error.set_array(task_pred_sigma.detach()) - - else: - pred_mean.set_offsets(x.detach()[..., self.axes]) - pred_error.set_offsets(x.detach()[..., self.axes]) - pred_mean.set_array(task_pred_mean.detach()) - pred_error.set_array(task_pred_sigma.detach()) - - for ax in self.task_axes.ravel(): - ax.set_xlim(*agent._subset_dofs(kind="active", mode="on")[self.axes[0]]["limits"]) - ax.set_ylim(*agent._subset_dofs(kind="active", mode="on")[self.axes[1]]["limits"]) diff --git a/bloptools/test_functions.py b/bloptools/test_functions.py index 845e155..377cf16 100644 --- a/bloptools/test_functions.py +++ b/bloptools/test_functions.py @@ -29,6 +29,34 @@ def constrained_himmelblau(x1, x2): return np.where(x1**2 + x2**2 < 50, himmelblau(x1, x2), np.nan) +def binh_korn(x1, x2): + """ + Binh and Korn function + """ + f1 = 4 * x1**2 + 4 * x2**2 + f2 = (x1 - 5) ** 2 + (x2 - 5) ** 2 + g1 = (x1 - 5) ** 2 + x2**2 <= 25 + g2 = (x1 - 8) ** 2 + (x2 + 3) ** 2 >= 7.7 + + c = g1 & g2 + + return np.where(c, f1, np.nan), np.where(c, f2, np.nan) + + +def binh_korn_digestion(db, uid): + """ + Digests Himmelblau's function into the feedback. + """ + products = db[uid].table() + + for index, entry in products.iterrows(): + f1, f2 = binh_korn(entry.x1, entry.x2) + products.loc[index, "f1"] = f1 + products.loc[index, "f2"] = f2 + + return products + + def skewed_himmelblau(x1, x2): """ Himmelblau's function, with skewed coordinates diff --git a/bloptools/tests/test_bayesian_shadow.py b/bloptools/tests/test_bayesian_shadow.py deleted file mode 100644 index aa63e97..0000000 --- a/bloptools/tests/test_bayesian_shadow.py +++ /dev/null @@ -1,41 +0,0 @@ -import pytest - -import bloptools -from bloptools.experiments.sirepo.tes import w9_digestion - - -@pytest.mark.shadow -def test_bayesian_agent_tes_shadow(RE, db, shadow_tes_simulation): - from sirepo_bluesky.sirepo_ophyd import create_classes - - data, schema = shadow_tes_simulation.auth("shadow", "00000002") - classes, objects = create_classes(connection=shadow_tes_simulation) - globals().update(**objects) - - data["models"]["simulation"]["npoint"] = 100000 - data["models"]["watchpointReport12"]["histogramBins"] = 32 - - dofs = [ - {"device": kbv.x_rot, "limits": (-0.1, 0.1), "kind": "active"}, - {"device": kbh.x_rot, "limits": (-0.1, 0.1), "kind": "active"}, - ] - - tasks = [ - {"key": "flux", "kind": "maximize", "transform": "log"}, - {"key": "w9_fwhm_x", "kind": "minimize", "transform": "log"}, - {"key": "w9_fwhm_y", "kind": "minimize", "transform": "log"}, - ] - - agent = bloptools.bayesian.Agent( - dofs=dofs, - tasks=tasks, - dets=[w9], - digestion=w9_digestion, - db=db, - ) - - RE(agent.initialize("qr", n_init=4)) - RE(agent.learn("ei", n_iter=2)) - RE(agent.learn("pi", n_iter=2)) - - agent.plot_tasks() From 208b139401dd43ba0d47c8dc4e9c795915a662f4 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 14 Sep 2023 16:12:27 -0400 Subject: [PATCH 29/43] old plots --- bloptools/bayesian/__init__.py | 107 ++++++++++++++++----------------- 1 file changed, 52 insertions(+), 55 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 0974541..a4e2715 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -806,9 +806,9 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL gridded = self._len_subset_dofs(kind="active", mode="on") == 2 self.task_fig, self.task_axes = plt.subplots( - self.num_tasks, + self.n_tasks, 3, - figsize=(10, 4 * self.num_tasks), + figsize=(10, 4 * self.n_tasks), sharex=True, sharey=True, constrained_layout=True, @@ -817,12 +817,16 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL self.task_axes = np.atleast_2d(self.task_axes) # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") + feasible = ~self.task_fitnesses.isna().any(axis=1) + for itask, task in enumerate(self.tasks): - task_plots = self.plots["tasks"][task["name"]] = {} + sampled_fitness = np.where(feasible, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) + task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) + task_norm = mpl.colors.Normalize(task_vmin, task_vmax) - # sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) - # task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) - # task_norm = mpl.colors.Normalize() + # if task["transform"] == "log": + # task_norm = mpl.colors.LogNorm(task_vmin, task_vmax) + # else: self.task_axes[itask, 0].set_ylabel(f'{task["key"]}_fitness') @@ -830,65 +834,58 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL self.task_axes[itask, 1].set_title("posterior mean") self.task_axes[itask, 2].set_title("posterior std. dev.") - task_plots["sampled"] = self.task_axes[itask, 0].scatter([], [], s=size, cmap=cmap) + data_ax = self.task_axes[itask, 0].scatter( + *self.inputs.values.T[axes], s=size, c=sampled_fitness, norm=task_norm, cmap=cmap + ) x = self.test_inputs_grid.squeeze() if gridded else self.test_inputs(n=MAX_TEST_INPUTS) + task_posterior = task["model"].posterior(x) + task_mean = task_posterior.mean + task_sigma = task_posterior.variance.sqrt() + if gridded: if not x.ndim == 3: raise ValueError() - task_plots["pred_mean"] = self.task_axes[itask, 1].imshow([[0]], cmap=cmap) - task_plots["pred_sigma"] = self.task_axes[itask, 2].imshow([[0]], cmap=cmap) + self.task_axes[itask, 1].pcolormesh( + x[..., 0], + x[..., 1], + task_mean[..., 0].detach().numpy(), + shading=shading, + cmap=cmap, + norm=task_norm, + ) + sigma_ax = self.task_axes[itask, 2].pcolormesh( + x[..., 0], + x[..., 1], + task_sigma[..., 0].detach().numpy(), + shading=shading, + cmap=cmap, + ) else: - task_plots["pred_mean"] = self.task_axes[itask, 1].scatter([], [], s=size, cmap=cmap) - task_plots["pred_sigma"] = self.task_axes[itask, 2].scatter([], [], s=size, cmap=cmap) - - task_plots["colorbar_mean"] = self.task_fig.colorbar( - task_plots["sampled"], ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8 - ) - task_plots["colorbar_sigma"] = self.task_fig.colorbar( - task_plots["pred_sigma"], ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8 - ) - - def plot_acquisition(self, acq_funcs=["ei"], **kwargs): - if self._len_subset_dofs(kind="active", mode="on") == 1: - self._plot_acq_one_dof(acq_funcs=acq_funcs, **kwargs) - - else: - self._plot_acq_many_dofs(acq_funcs=acq_funcs, **kwargs) - - def _plot_acq_one_dof(self, acq_funcs, lw=1e0, **kwargs): - self.acq_fig, self.acq_axes = plt.subplots( - 1, - len(acq_funcs), - figsize=(4 * len(acq_funcs), 4), - sharex=True, - constrained_layout=True, - ) - - self.acq_axes = np.atleast_1d(self.acq_axes) - - for iacq_func, acq_func_identifier in enumerate(acq_funcs): - color = DEFAULT_COLOR_LIST[iacq_func] - - acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) - - x = self.test_inputs_grid - *input_shape, input_dim = x.shape - obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) - - if acq_func_identifier in ["ei", "pi"]: - obj = obj.exp() - - self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + self.task_axes[itask, 1].scatter( + x.detach().numpy()[..., axes[0]], + x.detach().numpy()[..., axes[1]], + c=task_mean[..., 0].detach().numpy(), + s=size, + norm=task_norm, + cmap=cmap, + ) + sigma_ax = self.task_axes[itask, 2].scatter( + x.detach().numpy()[..., axes[0]], + x.detach().numpy()[..., axes[1]], + c=task_sigma[..., 0].detach().numpy(), + s=size, + cmap=cmap, + ) - on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] - self.acq_axes[iacq_func].plot( - x[..., on_dofs_are_active_mask].squeeze(), obj.detach().numpy(), lw=lw, color=color - ) + self.task_fig.colorbar(data_ax, ax=self.task_axes[itask, :2], location="bottom", aspect=32, shrink=0.8) + self.task_fig.colorbar(sigma_ax, ax=self.task_axes[itask, 2], location="bottom", aspect=32, shrink=0.8) - self.acq_axes[iacq_func].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) + for ax in self.task_axes.ravel(): + ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) + ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) def _plot_acq_many_dofs( self, acq_funcs, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16, **kwargs From d60475202c62e6b54118b1821ef1ecc8602e5636 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 21 Sep 2023 10:03:20 -0400 Subject: [PATCH 30/43] iss tweaks --- bloptools/bayesian/__init__.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index a4e2715..ddf78dd 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -717,6 +717,7 @@ def learn( n_iter=1, n_per_iter=1, reuse_hypers=True, + train=True, upsample=1, verbose=True, plots=[], @@ -736,7 +737,7 @@ def learn( new_table.loc[:, "acq_func"] = acq_func_meta["name"] - self.tell(new_table=new_table, reuse_hypers=reuse_hypers) + self.tell(new_table=new_table, train=train) @property def inputs(self): @@ -806,9 +807,9 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL gridded = self._len_subset_dofs(kind="active", mode="on") == 2 self.task_fig, self.task_axes = plt.subplots( - self.n_tasks, + len(self.tasks), 3, - figsize=(10, 4 * self.n_tasks), + figsize=(10, 4 * len(self.tasks)), sharex=True, sharey=True, constrained_layout=True, From 008a08967383977fce80d4dfe9aef1ee4a178815 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 21 Sep 2023 10:41:54 -0400 Subject: [PATCH 31/43] fix routing bug --- bloptools/bayesian/__init__.py | 4 +--- bloptools/utils/__init__.py | 3 ++- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index ddf78dd..6f468d0 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -818,10 +818,8 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL self.task_axes = np.atleast_2d(self.task_axes) # self.task_fig.suptitle(f"(x,y)=({self.dofs[axes[0]].name},{self.dofs[axes[1]].name})") - feasible = ~self.task_fitnesses.isna().any(axis=1) - for itask, task in enumerate(self.tasks): - sampled_fitness = np.where(feasible, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) + sampled_fitness = np.where(self.all_tasks_valid, self.table.loc[:, f'{task["key"]}_fitness'].values, np.nan) task_vmin, task_vmax = np.nanpercentile(sampled_fitness, q=[1, 99]) task_norm = mpl.colors.Normalize(task_vmin, task_vmax) diff --git a/bloptools/utils/__init__.py b/bloptools/utils/__init__.py index 67e1ec1..4b98dba 100644 --- a/bloptools/utils/__init__.py +++ b/bloptools/utils/__init__.py @@ -50,7 +50,8 @@ def route(start_point, points): """ total_points = np.r_[np.atleast_2d(start_point), points] - normalized_points = (total_points - total_points.min(axis=0)) / total_points.ptp(axis=0) + points_scale = total_points.ptp(axis=0) + normalized_points = (total_points - total_points.min(axis=0)) / np.where(np.isnan(points_scale), 0, points_scale) delay_matrix = np.sqrt(np.square(normalized_points[:, None, :] - normalized_points[None, :, :]).sum(axis=-1)) delay_matrix = (1e4 * delay_matrix).astype(int) # it likes integers idk From 17c141d768a134966614519693952c97bac42bd0 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 21 Sep 2023 11:04:35 -0400 Subject: [PATCH 32/43] actually fix routing bug --- bloptools/bayesian/__init__.py | 2 +- bloptools/utils/__init__.py | 8 +++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 6f468d0..6ecd1f8 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -651,7 +651,7 @@ def ask_single( ) BATCH_SIZE = 1 - NUM_RESTARTS = 8 + NUM_RESTARTS = 4 RAW_SAMPLES = 256 candidates, _ = botorch.optim.optimize_acqf( diff --git a/bloptools/utils/__init__.py b/bloptools/utils/__init__.py index 4b98dba..0804f6d 100644 --- a/bloptools/utils/__init__.py +++ b/bloptools/utils/__init__.py @@ -51,7 +51,13 @@ def route(start_point, points): total_points = np.r_[np.atleast_2d(start_point), points] points_scale = total_points.ptp(axis=0) - normalized_points = (total_points - total_points.min(axis=0)) / np.where(np.isnan(points_scale), 0, points_scale) + dim_mask = points_scale > 0 + + if dim_mask.sum() == 0: + return np.arange(len(points)) + + normalized_points = (total_points - total_points.min(axis=0))[dim_mask] / points_scale[dim_mask] + delay_matrix = np.sqrt(np.square(normalized_points[:, None, :] - normalized_points[None, :, :]).sum(axis=-1)) delay_matrix = (1e4 * delay_matrix).astype(int) # it likes integers idk From 33c4e207975bd611326d35ac616dd0da152bfb4f Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 21 Sep 2023 11:08:22 -0400 Subject: [PATCH 33/43] wrong dimensiongit add . --- bloptools/utils/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bloptools/utils/__init__.py b/bloptools/utils/__init__.py index 0804f6d..b59f122 100644 --- a/bloptools/utils/__init__.py +++ b/bloptools/utils/__init__.py @@ -56,7 +56,7 @@ def route(start_point, points): if dim_mask.sum() == 0: return np.arange(len(points)) - normalized_points = (total_points - total_points.min(axis=0))[dim_mask] / points_scale[dim_mask] + normalized_points = (total_points - total_points.min(axis=0))[:, dim_mask] / points_scale[dim_mask] delay_matrix = np.sqrt(np.square(normalized_points[:, None, :] - normalized_points[None, :, :]).sum(axis=-1)) delay_matrix = (1e4 * delay_matrix).astype(int) # it likes integers idk From e75cf50ccb59fd511730e57c8fb7d90357b68723 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Fri, 22 Sep 2023 13:34:32 -0400 Subject: [PATCH 34/43] why did we get rid of acquisition plots??? --- bloptools/bayesian/__init__.py | 39 ++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 6ecd1f8..df634ab 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -886,6 +886,45 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL ax.set_xlim(*self._subset_dofs(kind="active", mode="on")[axes[0]]["limits"]) ax.set_ylim(*self._subset_dofs(kind="active", mode="on")[axes[1]]["limits"]) + def plot_acquisition(self, acq_funcs=["ei"], **kwargs): + if self._len_subset_dofs(kind="active", mode="on") == 1: + self._plot_acq_one_dof(acq_funcs=acq_funcs, **kwargs) + + else: + self._plot_acq_many_dofs(acq_funcs=acq_funcs, **kwargs) + + def _plot_acq_one_dof(self, acq_funcs, lw=1e0, **kwargs): + self.acq_fig, self.acq_axes = plt.subplots( + 1, + len(acq_funcs), + figsize=(4 * len(acq_funcs), 4), + sharex=True, + constrained_layout=True, + ) + + self.acq_axes = np.atleast_1d(self.acq_axes) + + for iacq_func, acq_func_identifier in enumerate(acq_funcs): + color = DEFAULT_COLOR_LIST[iacq_func] + + acq_func, acq_func_meta = self.get_acquisition_function(acq_func_identifier, return_metadata=True) + + x = self.test_inputs_grid + *input_shape, input_dim = x.shape + obj = acq_func.forward(x.reshape(-1, 1, input_dim)).reshape(input_shape) + + if acq_func_identifier in ["ei", "pi"]: + obj = obj.exp() + + self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) + + on_dofs_are_active_mask = [dof["kind"] == "active" for dof in self._subset_dofs(mode="on")] + self.acq_axes[iacq_func].plot( + x[..., on_dofs_are_active_mask].squeeze(), obj.detach().numpy(), lw=lw, color=color + ) + + self.acq_axes[iacq_func].set_xlim(self._subset_dofs(kind="active", mode="on")[0]["limits"]) + def _plot_acq_many_dofs( self, acq_funcs, axes=[0, 1], shading="nearest", cmap=DEFAULT_COLORMAP, gridded=None, size=16, **kwargs ): From c17a02848b5211f679374770cfc8d56f900e8816 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 25 Sep 2023 09:22:23 -0400 Subject: [PATCH 35/43] tweaks --- bloptools/bayesian/__init__.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index df634ab..def8b24 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -728,7 +728,7 @@ def learn( It should be passed to a Bluesky RunEngine. """ - for iteration in range(n_iter): + for i in range(n_iter): x, acq_func_meta = self.ask( n=n_per_iter, acq_func_identifier=acq_func_identifier, return_metadata=True, **kwargs ) @@ -749,9 +749,10 @@ def best_inputs(self): def go_to(self, inputs): args = [] - for device, value in zip(self._subset_devices(kind="active"), np.atleast_1d(inputs).T): - args.append(device) - args.append(value) + for dof, value in zip(self._subset_dofs(mode="on"), np.atleast_1d(inputs).T): + if dof["kind"] == "active": + args.append(dof["device"]) + args.append(value) yield from bps.mv(*args) def go_to_best(self): From 9b60b3b559357dbafc0bc1f1203d4b86bf72434d Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 25 Sep 2023 09:49:16 -0400 Subject: [PATCH 36/43] all dofs if no tags --- bloptools/bayesian/__init__.py | 56 +++++++++++++++++++++------------- 1 file changed, 35 insertions(+), 21 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index def8b24..b25ca97 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -59,6 +59,9 @@ def _validate_and_prepare_dofs(dofs): dof["limits"] = (-np.inf, np.inf) dof["limits"] = tuple(np.array(dof["limits"], dtype=float)) + if "tags" not in dof.keys(): + dof["tags"] = [] + # dofs are passive by default dof["kind"] = dof.get("kind", "passive") if dof["kind"] not in ["active", "passive"]: @@ -161,6 +164,8 @@ def __init__( self.acq_func_config = kwargs.get("acq_func_config", ACQ_FUNC_CONFIG) + self.sample_center_on_init = kwargs.get("sample_center_on_init", False) + self.table = pd.DataFrame() self._initialized = False @@ -196,6 +201,11 @@ def initialize( if hypers is not None: self.a_priori_hypers = self.load_hypers(hypers) + if data is not None: + new_table = yield from self.acquire(self._acq_func_bounds.mean(axis=0)) + new_table.loc[:, "acq_func"] = "sample_center_on_init" + self.tell(new_table=new_table, train=False) + if data is not None: if type(data) == str: self.tell(new_table=pd.read_hdf(data, key="table")) @@ -394,26 +404,32 @@ def _dof_kind_mask(self, kind=None): def _dof_mode_mask(self, mode=None): return [dof["mode"] == mode if mode is not None else True for dof in self.dofs] - def _dof_mask(self, kind=None, mode=None): - return [(k and m) for k, m in zip(self._dof_kind_mask(kind), self._dof_mode_mask(mode))] + def _dof_tags_mask(self, tags=[]): + return [np.isin(dof["tags"], tags).any() if tags else True for dof in self.dofs] - def _subset_dofs(self, kind=None, mode=None): - return [dof for dof, m in zip(self.dofs, self._dof_mask(kind, mode)) if m] + def _dof_mask(self, kind=None, mode=None, tags=[]): + return [ + (k and m and t) + for k, m, t in zip(self._dof_kind_mask(kind), self._dof_mode_mask(mode), self._dof_tags_mask(tags)) + ] + + def _subset_dofs(self, kind=None, mode=None, tags=[]): + return [dof for dof, m in zip(self.dofs, self._dof_mask(kind, mode, tags)) if m] - def _len_subset_dofs(self, kind=None, mode=None): - return len(self._subset_dofs(kind, mode)) + def _len_subset_dofs(self, kind=None, mode=None, tags=[]): + return len(self._subset_dofs(kind, mode, tags)) - def _subset_devices(self, kind=None, mode=None): - return [dof["device"] for dof in self._subset_dofs(kind, mode)] + def _subset_devices(self, kind=None, mode=None, tags=[]): + return [dof["device"] for dof in self._subset_dofs(kind, mode, tags)] - def _read_subset_devices(self, kind=None, mode=None): - return [device.read()[device.name]["value"] for device in self._subset_devices(kind, mode)] + def _read_subset_devices(self, kind=None, mode=None, tags=[]): + return [device.read()[device.name]["value"] for device in self._subset_devices(kind, mode, tags)] - def _subset_dof_names(self, kind=None, mode=None): - return [device.name for device in self._subset_devices(kind, mode)] + def _subset_dof_names(self, kind=None, mode=None, tags=[]): + return [device.name for device in self._subset_devices(kind, mode, tags)] - def _subset_dof_limits(self, kind=None, mode=None): - dofs_subset = self._subset_dofs(kind, mode) + def _subset_dof_limits(self, kind=None, mode=None, tags=[]): + dofs_subset = self._subset_dofs(kind, mode, tags) if len(dofs_subset) > 0: return torch.tensor([dof["limits"] for dof in dofs_subset], dtype=torch.float64).T return torch.empty((2, 0)) @@ -432,20 +448,20 @@ def latent_dim_tuples(self): def test_inputs(self, n=MAX_TEST_INPUTS): return utils.sobol_sampler(self._acq_func_bounds, n=n) - def _subset_input_transform(self, kind=None, mode=None): - limits = self._subset_dof_limits(kind, mode) + def _subset_input_transform(self, kind=None, mode=None, tags=[]): + limits = self._subset_dof_limits(kind, mode, tags) offset = limits.min(dim=0).values coefficient = limits.max(dim=0).values - offset return botorch.models.transforms.input.AffineInputTransform( d=limits.shape[-1], coefficient=coefficient, offset=offset ) - def _subset_inputs_sampler(self, kind=None, mode=None, n=MAX_TEST_INPUTS): + def _subset_inputs_sampler(self, kind=None, mode=None, tags=[], n=MAX_TEST_INPUTS): """ Returns $n$ quasi-randomly sampled inputs in the bounded parameter space """ - transform = self._subset_input_transform(kind=kind, mode=mode) - return transform.untransform(utils.normalized_sobol_sampler(n, d=self._len_subset_dofs(kind=kind, mode=mode))) + transform = self._subset_input_transform(kind, mode, tags) + return transform.untransform(utils.normalized_sobol_sampler(n, d=self._len_subset_dofs(kind, mode, tags))) def save_data(self, filepath="./self_data.h5"): """ @@ -734,9 +750,7 @@ def learn( ) new_table = yield from self.acquire(x) - new_table.loc[:, "acq_func"] = acq_func_meta["name"] - self.tell(new_table=new_table, train=train) @property From 05449e66bbee4ce00980b4f2bc1dfdaeff9390ed Mon Sep 17 00:00:00 2001 From: Max Rakitin Date: Mon, 25 Sep 2023 09:58:17 -0400 Subject: [PATCH 37/43] CI: install packages via built-in system Python (not conda) --- .github/workflows/docs.yml | 20 +++++++++++++++----- .github/workflows/testing.yml | 22 ++++++++++++---------- 2 files changed, 27 insertions(+), 15 deletions(-) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 68b9542..636daa0 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -57,19 +57,29 @@ jobs: mamba-version: "*" channels: conda-forge - - name: Install documentation-building requirements + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + + - name: Install documentation-building requirements with apt/dpkg + run: | + set -vxeuo pipefail + wget --progress=dot:giga "https://github.com/jgm/pandoc/releases/download/3.1.6.1/pandoc-3.1.6.1-1-amd64.deb" -O /tmp/pandoc.deb + sudo dpkg -i /tmp/pandoc.deb + # conda install -c conda-forge -y pandoc + which pandoc + pandoc --version + + - name: Install documentation-building requirements with pip run: | # For reference: https://www.gnu.org/software/bash/manual/html_node/The-Set-Builtin.html. set -vxeo pipefail - conda env list - # mamba install -c conda-forge shadow3 srwpy - mamba install -c conda-forge pandoc pip install --upgrade pip wheel pip install -v . pip install -r requirements-dev.txt pip list - conda list - name: Build Docs run: make -C docs/ html diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index c574cb7..2304149 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -46,28 +46,30 @@ jobs: # mkdir -v -p ~/.config/databroker/ # wget https://raw.githubusercontent.com/NSLS-II/sirepo-bluesky/main/examples/local.yml -O ~/.config/databroker/local.yml - - name: Set up Python ${{ matrix.python-version }} with conda - uses: conda-incubator/setup-miniconda@v2 + # - name: Set up Python ${{ matrix.python-version }} with conda + # uses: conda-incubator/setup-miniconda@v2 + # with: + # activate-environment: ${{ env.REPOSITORY_NAME }}-py${{ matrix.python-version }} + # auto-update-conda: true + # miniconda-version: "latest" + # python-version: ${{ matrix.python-version }} + # mamba-version: "*" + # channels: conda-forge + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 with: - activate-environment: ${{ env.REPOSITORY_NAME }}-py${{ matrix.python-version }} - auto-update-conda: true - miniconda-version: "latest" python-version: ${{ matrix.python-version }} - mamba-version: "*" - channels: conda-forge - name: Install the package and its dependencies run: | # For reference: https://www.gnu.org/software/bash/manual/html_node/The-Set-Builtin.html. set -vxeo pipefail - conda env list - pip install --upgrade pip wheel pip install -v . pip install -r requirements-dev.txt pip list - conda list - name: Test with pytest run: | From ceaf6d88b0cd84db1faa83004c4b60522a245c10 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 25 Sep 2023 10:42:58 -0400 Subject: [PATCH 38/43] beam image parser --- bloptools/bayesian/__init__.py | 8 +++++++ bloptools/utils/__init__.py | 41 ++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index b25ca97..af4afb6 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -413,6 +413,14 @@ def _dof_mask(self, kind=None, mode=None, tags=[]): for k, m, t in zip(self._dof_kind_mask(kind), self._dof_mode_mask(mode), self._dof_tags_mask(tags)) ] + def activate_dofs(self, kind=None, mode=None, tags=[]): + for dof in self._subset_dofs(kind, mode, tags): + dof["mode"] = "on" + + def deactivate_dofs(self, kind=None, mode=None, tags=[]): + for dof in self._subset_dofs(kind, mode, tags): + dof["mode"] = "off" + def _subset_dofs(self, kind=None, mode=None, tags=[]): return [dof for dof, m in zip(self.dofs, self._dof_mask(kind, mode, tags)) if m] diff --git a/bloptools/utils/__init__.py b/bloptools/utils/__init__.py index b59f122..89021e9 100644 --- a/bloptools/utils/__init__.py +++ b/bloptools/utils/__init__.py @@ -144,3 +144,44 @@ def get_principal_component_bounds(image, beam_prop=0.5): y_max, separability, ) + + +def get_beam_bounding_box(image, thresh=0.5): + """ + Returns the bounding box in pixel units of an image, along with a goodness of fit parameter. + This should go off without a hitch as long as beam_prop is less than 1. + """ + + if image.sum() == 0: + return np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan + + # filter the image + zim = sp.ndimage.median_filter(image.astype(float), size=3) + zim -= np.median(zim, axis=0) + zim -= np.median(zim, axis=1)[:, None] + + x_sum = zim.sum(axis=0) + y_sum = zim.sum(axis=1) + + x_sum_min_val = thresh * x_sum.max() + y_sum_min_val = thresh * y_sum.max() + + gtt_x = x_sum > x_sum_min_val + gtt_y = y_sum > y_sum_min_val + + i_x_min_start = np.where(~gtt_x[:-1] & gtt_x[1:])[0][0] + i_x_max_start = np.where(gtt_x[:-1] & ~gtt_x[1:])[0][-1] + i_y_min_start = np.where(~gtt_y[:-1] & gtt_y[1:])[0][0] + i_y_max_start = np.where(gtt_y[:-1] & ~gtt_y[1:])[0][-1] + + x_min = np.interp(x_sum_min_val, x_sum[[i_x_min_start, i_x_min_start + 1]], [i_x_min_start, i_x_min_start + 1]) + y_min = np.interp(y_sum_min_val, y_sum[[i_y_min_start, i_y_min_start + 1]], [i_y_min_start, i_y_min_start + 1]) + x_max = np.interp(x_sum_min_val, x_sum[[i_x_max_start + 1, i_x_max_start]], [i_x_max_start + 1, i_x_max_start]) + y_max = np.interp(y_sum_min_val, y_sum[[i_y_max_start + 1, i_y_max_start]], [i_y_max_start + 1, i_y_max_start]) + + return ( + x_min, + x_max, + y_min, + y_max, + ) From 26546b6278eb92850cf380bf710eec6cd068f553 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 25 Sep 2023 12:36:34 -0400 Subject: [PATCH 39/43] beam image parser bug --- bloptools/utils/__init__.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/bloptools/utils/__init__.py b/bloptools/utils/__init__.py index 89021e9..326d58e 100644 --- a/bloptools/utils/__init__.py +++ b/bloptools/utils/__init__.py @@ -152,6 +152,8 @@ def get_beam_bounding_box(image, thresh=0.5): This should go off without a hitch as long as beam_prop is less than 1. """ + n_y, n_x = image.shape + if image.sum() == 0: return np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan @@ -174,10 +176,26 @@ def get_beam_bounding_box(image, thresh=0.5): i_y_min_start = np.where(~gtt_y[:-1] & gtt_y[1:])[0][0] i_y_max_start = np.where(gtt_y[:-1] & ~gtt_y[1:])[0][-1] - x_min = np.interp(x_sum_min_val, x_sum[[i_x_min_start, i_x_min_start + 1]], [i_x_min_start, i_x_min_start + 1]) - y_min = np.interp(y_sum_min_val, y_sum[[i_y_min_start, i_y_min_start + 1]], [i_y_min_start, i_y_min_start + 1]) - x_max = np.interp(x_sum_min_val, x_sum[[i_x_max_start + 1, i_x_max_start]], [i_x_max_start + 1, i_x_max_start]) - y_max = np.interp(y_sum_min_val, y_sum[[i_y_max_start + 1, i_y_max_start]], [i_y_max_start + 1, i_y_max_start]) + x_min = ( + 0 + if gtt_x[0] + else np.interp(x_sum_min_val, x_sum[[i_x_min_start, i_x_min_start + 1]], [i_x_min_start, i_x_min_start + 1]) + ) + y_min = ( + 0 + if gtt_y[0] + else np.interp(y_sum_min_val, y_sum[[i_y_min_start, i_y_min_start + 1]], [i_y_min_start, i_y_min_start + 1]) + ) + x_max = ( + n_x - 2 + if gtt_x[-1] + else np.interp(x_sum_min_val, x_sum[[i_x_max_start + 1, i_x_max_start]], [i_x_max_start + 1, i_x_max_start]) + ) + y_max = ( + n_y - 2 + if gtt_y[-1] + else np.interp(y_sum_min_val, y_sum[[i_y_max_start + 1, i_y_max_start]], [i_y_max_start + 1, i_y_max_start]) + ) return ( x_min, From 090f303fce3e7655784df94ed17494f4fe0a7285 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 25 Sep 2023 13:16:27 -0400 Subject: [PATCH 40/43] better parser --- bloptools/utils/__init__.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/bloptools/utils/__init__.py b/bloptools/utils/__init__.py index 326d58e..6f3bd85 100644 --- a/bloptools/utils/__init__.py +++ b/bloptools/utils/__init__.py @@ -203,3 +203,25 @@ def get_beam_bounding_box(image, thresh=0.5): y_min, y_max, ) + + +def best_image_feedback(image): + n_y, n_x = image.shape + + fim = sp.ndimage.median_filter(image, size=3) + + masked_image = fim * (fim - fim.mean() > 0.5 * fim.ptp()) + + x_weight = masked_image.sum(axis=0) + y_weight = masked_image.sum(axis=1) + + x = np.arange(n_x) + y = np.arange(n_y) + + x0 = np.sum(x_weight * x) / np.sum(x_weight) + y0 = np.sum(y_weight * y) / np.sum(y_weight) + + xw = 2 * np.sqrt((np.sum(x_weight * x**2) / np.sum(x_weight) - x0**2)) + yw = 2 * np.sqrt((np.sum(y_weight * y**2) / np.sum(y_weight) - y0**2)) + + return x0, xw, y0, yw From 4af51f60c63147b128364630ee31ad56e1dce931 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 26 Sep 2023 12:43:27 -0400 Subject: [PATCH 41/43] detach tensors before plotting --- bloptools/bayesian/__init__.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index af4afb6..35ba934 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -870,16 +870,16 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL if not x.ndim == 3: raise ValueError() self.task_axes[itask, 1].pcolormesh( - x[..., 0], - x[..., 1], + x[..., 0].detach(), + x[..., 1].detach(), task_mean[..., 0].detach().numpy(), shading=shading, cmap=cmap, norm=task_norm, ) sigma_ax = self.task_axes[itask, 2].pcolormesh( - x[..., 0], - x[..., 1], + x[..., 0].detach(), + x[..., 1].detach(), task_sigma[..., 0].detach().numpy(), shading=shading, cmap=cmap, @@ -979,8 +979,8 @@ def _plot_acq_many_dofs( if gridded: self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) obj_ax = self.acq_axes[iacq_func].pcolormesh( - x[..., 0], - x[..., 1], + x[..., 0].detach(), + x[..., 1].detach(), obj.detach().numpy(), shading=shading, cmap=cmap, @@ -1047,8 +1047,8 @@ def _plot_valid_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL if gridded: self.valid_axes[1].pcolormesh( - x[..., 0], - x[..., 1], + x[..., 0].detach(), + x[..., 1].detach(), constraint.detach().numpy(), shading=shading, cmap=cmap, From 26e61995784a0a87cb57bbf25e61b0581cb3ca47 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 26 Sep 2023 12:50:30 -0400 Subject: [PATCH 42/43] add .numpy() --- bloptools/bayesian/__init__.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/bloptools/bayesian/__init__.py b/bloptools/bayesian/__init__.py index 35ba934..b951ce0 100644 --- a/bloptools/bayesian/__init__.py +++ b/bloptools/bayesian/__init__.py @@ -870,16 +870,16 @@ def _plot_tasks_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL if not x.ndim == 3: raise ValueError() self.task_axes[itask, 1].pcolormesh( - x[..., 0].detach(), - x[..., 1].detach(), + x[..., 0].detach().numpy(), + x[..., 1].detach().numpy(), task_mean[..., 0].detach().numpy(), shading=shading, cmap=cmap, norm=task_norm, ) sigma_ax = self.task_axes[itask, 2].pcolormesh( - x[..., 0].detach(), - x[..., 1].detach(), + x[..., 0].detach().numpy(), + x[..., 1].detach().numpy(), task_sigma[..., 0].detach().numpy(), shading=shading, cmap=cmap, @@ -979,8 +979,8 @@ def _plot_acq_many_dofs( if gridded: self.acq_axes[iacq_func].set_title(acq_func_meta["name"]) obj_ax = self.acq_axes[iacq_func].pcolormesh( - x[..., 0].detach(), - x[..., 1].detach(), + x[..., 0].detach().numpy(), + x[..., 1].detach().numpy(), obj.detach().numpy(), shading=shading, cmap=cmap, @@ -1047,8 +1047,8 @@ def _plot_valid_many_dofs(self, axes=[0, 1], shading="nearest", cmap=DEFAULT_COL if gridded: self.valid_axes[1].pcolormesh( - x[..., 0].detach(), - x[..., 1].detach(), + x[..., 0].detach().numpy(), + x[..., 1].detach().numpy(), constraint.detach().numpy(), shading=shading, cmap=cmap, From 397baed9ce0af6ad70026edf4c6a7f63530d4839 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 28 Sep 2023 17:31:51 -0400 Subject: [PATCH 43/43] tweak the requirements --- requirements.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 433b4c5..e48f1be 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -bluesky>=1.6.1 +bluesky botorch databroker gpytorch @@ -10,3 +10,4 @@ ortools scipy tables torch +zict<3.0.0