From 3978d7bcc7a3cb783b6434d8ff14b7120dfd4741 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Mon, 17 Apr 2023 18:13:57 -0400 Subject: [PATCH 1/7] renamed + restructured the botorch interface --- bloptools/bo/__init__.py | 219 ++++++++++-------------- bloptools/bo/acquisition.py | 19 +- bloptools/bo/models.py | 161 ++++++++--------- bloptools/experiments/tests/__init__.py | 18 +- 4 files changed, 195 insertions(+), 222 deletions(-) diff --git a/bloptools/bo/__init__.py b/bloptools/bo/__init__.py index 5f1549d..c70211f 100644 --- a/bloptools/bo/__init__.py +++ b/bloptools/bo/__init__.py @@ -17,9 +17,9 @@ def load(filepath, **kwargs): with h5py.File(filepath, "r") as f: - params = f["params"][:] + X = f["X"][:] data = pd.read_hdf(filepath, key="data") - return BayesianOptimizationAgent(init_params=params, init_data=data, **kwargs) + return BayesianOptimizationAgent(init_X=X, init_data=data, **kwargs) # import bluesky_adaptive @@ -77,20 +77,20 @@ def __init__( self.test_grid = np.swapaxes(np.r_[np.meshgrid(*self.dim_mids, indexing="ij")], 0, -1) sampler = sp.stats.qmc.Halton(d=self.n_dof, scramble=True) - self.test_params = sampler.random(n=MAX_TEST_POINTS) * self.bounds.ptp(axis=1) + self.bounds.min(axis=1) + self.test_X = sampler.random(n=MAX_TEST_POINTS) * self.bounds.ptp(axis=1) + self.bounds.min(axis=1) - # convert params to x - self.params_trans_fun = lambda params: (params - self.bounds.min(axis=1)) / self.bounds.ptp(axis=1) + # convert X to x + self.X_trans_fun = lambda X: (X - self.bounds.min(axis=1)) / self.bounds.ptp(axis=1) - # convert x to params - self.inv_params_trans_fun = lambda x: x * self.bounds.ptp(axis=1) + self.bounds.min(axis=1) + # convert x to X + self.inv_X_trans_fun = lambda x: x * self.bounds.ptp(axis=1) + self.bounds.min(axis=1) # for actual prediction and optimization - self.params = np.empty((0, self.n_dof)) + self.X = np.empty((0, self.n_dof)) self.data = pd.DataFrame() - self.evaluator = models.GPR(MIN_SNR=experiment.MIN_SNR) - self.classifier = models.GPC() + self.evaluator = models.GPR(bounds=self.bounds, MIN_SNR=experiment.MIN_SNR) + self.classifier = models.GPC(bounds=self.bounds) def measurement_plan(self): yield from bp.count(detectors=self.dets) @@ -100,7 +100,7 @@ def unpack_run(self): def initialize( self, - init_params=None, + init_X=None, init_data=None, init_scheme=None, n_init=4, @@ -109,8 +109,8 @@ def initialize( yield from self.experiment.initialize() # now let's get bayesian - if (init_params is not None) and (init_data is not None): - self.tell(new_params=init_params, new_data=init_data, reuse_hypers=True, verbose=self.verbose) + if (init_X is not None) and (init_data is not None): + self.tell(new_X=init_X, new_data=init_data, reuse_hypers=True, verbose=self.verbose) elif init_scheme == "quasi-random": yield from self.learn( @@ -119,12 +119,12 @@ def initialize( else: raise Exception( - "Could not initialize model! Either pass initial params and data, or specify one of:" + "Could not initialize model! Either pass initial X and data, or specify one of:" "['quasi-random']." ) @property - def current_params(self): + def current_X(self): return np.array([dof.read()[dof.name]["value"] for dof in self.dofs]) @property @@ -135,13 +135,9 @@ def dof_names(self): def det_names(self): return self.experiment.DEPENDENT_COMPONENTS - @property - def current_X(self): - return self.inv_params_trans_fun(self.current_params) - @property def optimum(self): - return self.inv_params_trans_fun(self.evaluator.X[np.nanargmax(self.evaluator.y)]) + return self.evaluator.X[np.argmax(self.evaluator.Y)] def go_to_optimum(self): yield from bps.mv(*[_ for items in zip(self.dofs, np.atleast_1d(self.optimum).T) for _ in items]) @@ -169,18 +165,18 @@ def inspect_beam(self, index, border=None): def save(self, filepath="./model.gpo"): with h5py.File(filepath, "w") as f: - f.create_dataset("params", data=self.params) + f.create_dataset("X", data=self.X) self.data.to_hdf(filepath, key="data") - def tell(self, new_params, new_data, reuse_hypers=True, verbose=False): - self.params = np.r_[self.params, new_params] + def tell(self, new_X, new_data, reuse_hypers=True, verbose=False): + self.X = np.r_[self.X, new_X] self.data = pd.concat([self.data, new_data]) self.data.index = np.arange(len(self.data)) if hasattr(self.experiment, "IMAGE_NAME"): self.images = np.array([im for im in self.data[self.experiment.IMAGE_NAME].values]) - X = self.params_trans_fun(self.params) + X = self.X.copy() Y = self.data.fitness.values[:, None] c = (~np.isnan(Y).any(axis=-1)).astype(int) @@ -188,26 +184,27 @@ def tell(self, new_params, new_data, reuse_hypers=True, verbose=False): self.classifier.set_data(X, c) # self.timer.train(training_iter=self.training_iter, reuse_hypers=reuse_hypers, verbose=verbose) - self.evaluator.train(maxiter=self.training_iter) - self.classifier.train(maxiter=self.training_iter) - def acquire_with_bluesky(self, params, routing=True, verbose=False): + self.evaluator.train(step_limit=self.training_iter) + self.classifier.train(step_limit=self.training_iter) + + def acquire_with_bluesky(self, X, routing=True, verbose=False): if routing: - routing_index, _ = utils.get_routing(self.current_params, params) - ordered_params = params[routing_index] + routing_index, _ = utils.get_routing(self.current_X, X) + ordered_X = X[routing_index] else: - ordered_params = params + ordered_X = X table = pd.DataFrame(columns=[]) - # for _params in ordered_params: + # for _X in ordered_X: if verbose: - print(f"sampling {ordered_params}") + print(f"sampling {ordered_X}") try: uid = yield from bp.list_scan( - self.dets, *[_ for items in zip(self.dofs, np.atleast_2d(ordered_params).T) for _ in items] + self.dets, *[_ for items in zip(self.dofs, np.atleast_2d(ordered_X).T) for _ in items] ) _table = self.db[uid].table(fill=True) _table.loc[:, "uid"] = uid @@ -220,9 +217,9 @@ def acquire_with_bluesky(self, params, routing=True, verbose=False): table = pd.concat([table, _table]) - return ordered_params, table + return ordered_X, table - def recommend( + def ask( self, evaluator=None, classifier=None, @@ -250,11 +247,12 @@ def recommend( if strategy.lower() == "quasi-random": # if init and self.sample_center_on_init: # return np.r_[0.5 * np.ones((1, self.n_dof)), sampler.random(n=n-1)] - return sampler.random(n=n) + return sampler.random(n=n) * self.bounds.ptp(axis=1) + self.bounds.min(axis=1) if (not greedy) or (n == 1): # recommend some parameters that we might want to sample, with shape (., n, n_dof) - TEST_X = sampler.random(n=n * n_test).reshape(n_test, n, self.n_dof) + TEST_X = sampler.random(n=n * n_test) * self.bounds.ptp(axis=1) + self.bounds.min(axis=1) + TEST_X = TEST_X.reshape(n_test, n, self.n_dof) # how much will we have to change our parameters to sample these guys? DELTA_TEST_X = np.diff( @@ -274,7 +272,10 @@ def recommend( if strategy.lower() == "ei": # maximize the expected improvement objective = acquisition.expected_improvement(evaluator, classifier, TEST_X).sum(axis=1) - if strategy.lower() == "max_entropy": # maximize the total entropy + if strategy.lower() == "class_entropy": # maximize the class entropy + objective = classifier.entropy(TEST_X).sum(axis=1) + + if strategy.lower() == "total_entropy": # maximize the total entropy objective = (evaluator.normalized_entropy(TEST_X) + classifier.entropy(TEST_X)).sum(axis=1) if strategy.lower() == "egibbon": # maximize the expected GIBBON @@ -287,7 +288,7 @@ def recommend( X_to_sample = np.zeros((0, self.n_dof)) for i in range(n): - _X = self.recommend( + _X = self.ask( strategy=strategy, evaluator=dummy_evaluator, classifier=classifier, greedy=True, n=1 ) _y = self.evaluator.mean(_X) - disappointment * self.evaluator.sigma(_X) @@ -299,9 +300,6 @@ def recommend( return X_to_sample - def ask(self, **kwargs): - return self.inv_params_trans_fun(self.recommend(**kwargs)) - def learn( self, strategy, n_iter=1, n_per_iter=1, reuse_hypers=True, upsample=1, verbose=True, plots=[], **kwargs ): @@ -309,22 +307,22 @@ def learn( print(f'learning with strategy "{strategy}" ...') for i in range(n_iter): - params_to_sample = np.atleast_2d( + X_to_sample = np.atleast_2d( self.ask(n=n_per_iter, strategy=strategy, **kwargs) ) # get point(s) to sample from the strategizer - n_original = len(params_to_sample) + n_original = len(X_to_sample) n_upsample = upsample * n_original + 1 - upsampled_params_to_sample = sp.interpolate.interp1d( - np.arange(n_original + 1), np.r_[self.current_params[None], params_to_sample], axis=0 + upsampled_X_to_sample = sp.interpolate.interp1d( + np.arange(n_original + 1), np.r_[self.current_X[None], X_to_sample], axis=0 )(np.linspace(0, n_original, n_upsample)[1:]) - sampled_params, res_table = yield from self.acquire_with_bluesky( - upsampled_params_to_sample + sampled_X, res_table = yield from self.acquire_with_bluesky( + upsampled_X_to_sample ) # sample the point(s) - self.tell(new_params=sampled_params, new_data=res_table, reuse_hypers=reuse_hypers) + self.tell(new_X=sampled_X, new_data=res_table, reuse_hypers=reuse_hypers) if "state" in plots: self.plot_state(remake=False, gridded=self.gridded_plots) @@ -332,10 +330,10 @@ def learn( self.plot_fitness(remake=False) if verbose: - n_params = len(sampled_params) + n_X = len(sampled_X) df_to_print = pd.DataFrame( - np.c_[self.params, self.data.fitness.values], columns=[*self.dof_names, "fitness"] - ).iloc[-n_params:] + np.c_[self.X, self.data.fitness.values], columns=[*self.dof_names, "fitness"] + ).iloc[-n_X:] print(df_to_print) def plot_fitness(self, remake=True, **kwargs): @@ -410,24 +408,9 @@ def update_state_plots(self, gridded=False): ax.clear() ax.set_title("sampled fitness") - ax.scatter(*self.params.T[:2], s=s, edgecolor="k", facecolor="none") - - ref = ax.scatter(*self.params.T[:2], s=s, c=self.data.fitness.values, norm=fitness_norm) - - # max_improvement_params = self.inv_params_trans_fun( - # self.recommend(strategy="exploit", greedy=True, n=8, n_test=1024) - # ) - # max_information_params = self.inv_params_trans_fun( - # self.recommend(strategy="ev", greedy=True, n=8, n_test=1024) - # ) - - # max_improvement_routing_index, _ = utils.get_routing(self.current_params, max_improvement_params) - # ordered_max_improvement_params = max_improvement_params[max_improvement_routing_index] - - # max_information_routing_index, _ = utils.get_routing(self.current_params, max_information_params) - # ordered_max_information_params = max_information_params[max_information_routing_index] + ax.scatter(*self.X.T[:2], s=s, edgecolor="k", facecolor="none") - # ax.legend(fontsize=6) + ref = ax.scatter(*self.X.T[:2], s=s, c=self.data.fitness.values, norm=fitness_norm) clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) clb.set_label("fitness units") @@ -441,13 +424,11 @@ def update_state_plots(self, gridded=False): *self.dim_mids[:2], self.fitness_estimate(self.test_grid), norm=fitness_norm, shading="nearest" ) else: - ref = ax.scatter( - *self.test_params.T[:2], s=s, c=self.fitness_estimate(self.test_params), norm=fitness_norm - ) + ref = ax.scatter(*self.test_X.T[:2], s=s, c=self.fitness_estimate(self.test_X), norm=fitness_norm) clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) clb.set_label("fitness units") - ax.scatter(*self.params.T[:2], s=s, edgecolor="k", facecolor="none") + ax.scatter(*self.X.T[:2], s=s, edgecolor="k", facecolor="none") # plot the entropy rate of test points ax = self.state_axes[0, 2] @@ -457,15 +438,15 @@ def update_state_plots(self, gridded=False): ref = ax.pcolormesh(*self.dim_mids[:2], self.fitness_sigma(self.test_grid), shading="nearest") else: ref = ax.scatter( - *self.test_params.T[:2], + *self.test_X.T[:2], s=s, - c=self.fitness_sigma(self.test_params), + c=self.fitness_sigma(self.test_X), norm=mpl.colors.LogNorm(), ) clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) clb.set_label("fitness units") - ax.scatter(*self.params.T[:2], s=s, edgecolor="k", facecolor="none") + ax.scatter(*self.X.T[:2], s=s, edgecolor="k", facecolor="none") # plot the estimate of test points ax = self.state_axes[0, 3] @@ -474,38 +455,36 @@ def update_state_plots(self, gridded=False): if gridded: expected_improvement = acquisition.expected_improvement( - self.evaluator, self.classifier, self.params_trans_fun(self.test_grid) + self.evaluator, self.classifier, self.test_grid ) # expected_improvement[~(expected_improvement > 0)] = np.nan ref = ax.pcolormesh( *self.dim_mids[:2], expected_improvement, norm=mpl.colors.Normalize(), shading="nearest" ) else: - expected_improvement = acquisition.expected_improvement( - self.evaluator, self.classifier, self.params_trans_fun(self.test_params) - ) + expected_improvement = acquisition.expected_improvement(self.evaluator, self.classifier, self.test_X) # expected_improvement[~(expected_improvement > 0)] = np.nan ref = ax.scatter( - *self.test_params.T[:2], + *self.test_X.T[:2], s=s, c=expected_improvement, norm=mpl.colors.Normalize(), ) - # ax.plot(*ordered_max_improvement_params.T[:2], lw=1, color="k") - # ax.scatter(*ordered_max_improvement_params.T[:2], marker="*", color="k", s=s, label="max_improvement") + # ax.plot(*ordered_max_improvement_X.T[:2], lw=1, color="k") + # ax.scatter(*ordered_max_improvement_X.T[:2], marker="*", color="k", s=s, label="max_improvement") clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) clb.set_label("fitness units") - ax.scatter(*self.params.T[:2], s=s, edgecolor="k", facecolor="none") + ax.scatter(*self.X.T[:2], s=s, edgecolor="k", facecolor="none") # plot classification of data points ax = self.state_axes[1, 0] ax.clear() ax.set_title("sampled validity") - ax.scatter(*self.params.T[:2], s=s, edgecolor="k", facecolor="none") + ax.scatter(*self.X.T[:2], s=s, edgecolor="k", facecolor="none") - ref = ax.scatter(*self.params.T[:2], s=s, c=self.classifier.c, norm=mpl.colors.Normalize(vmin=0, vmax=1)) + ref = ax.scatter(*self.X.T[:2], s=s, c=self.classifier.c, norm=mpl.colors.Normalize(vmin=0, vmax=1)) clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) clb.set_ticks([0, 1]) clb.set_ticklabels(["invalid", "valid"]) @@ -516,22 +495,22 @@ def update_state_plots(self, gridded=False): if gridded: ref = ax.pcolormesh( *self.dim_mids[:2], - self.classifier.p(self.params_trans_fun(self.test_grid)), + self.classifier.p(self.test_grid), vmin=0, vmax=1, shading="nearest", ) else: ref = ax.scatter( - *self.test_params.T[:2], - c=self.classifier.p(self.params_trans_fun(self.test_params)), + *self.test_X.T[:2], + c=self.classifier.p(self.test_X), s=s, vmin=0, vmax=1, ) clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) clb.set_label("probability of validity") - ax.scatter(*self.params.T[:2], s=s, edgecolor="k", facecolor="none") + ax.scatter(*self.X.T[:2], s=s, edgecolor="k", facecolor="none") ax = self.state_axes[1, 2] ax.clear() @@ -539,42 +518,36 @@ def update_state_plots(self, gridded=False): if gridded: ref = ax.pcolormesh( *self.dim_mids[:2], - self.classifier.entropy(self.params_trans_fun(self.test_grid)), + self.classifier.entropy(self.test_grid), shading="nearest", ) else: - ref = ax.scatter( - *self.test_params.T[:2], c=self.classifier.entropy(self.params_trans_fun(self.test_params)), s=s - ) + ref = ax.scatter(*self.test_X.T[:2], c=self.classifier.entropy(self.test_X), s=s) clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) clb.set_label("nats") - ax.scatter(*self.params.T[:2], s=s, edgecolor="k", facecolor="none") + ax.scatter(*self.X.T[:2], s=s, edgecolor="k", facecolor="none") ax = self.state_axes[1, 3] ax.clear() ax.set_title("expected GIBBON") if gridded: - expected_improvement = acquisition.expected_gibbon( - self.evaluator, self.classifier, self.params_trans_fun(self.test_grid) - ) + expected_improvement = acquisition.expected_gibbon(self.evaluator, self.classifier, self.test_grid) # expected_improvement[~(expected_improvement > 0)] = np.nan ref = ax.pcolormesh( *self.dim_mids[:2], expected_improvement, norm=mpl.colors.Normalize(), shading="nearest" ) else: - expected_improvement = acquisition.expected_gibbon( - self.evaluator, self.classifier, self.params_trans_fun(self.test_params) - ) + expected_improvement = acquisition.expected_gibbon(self.evaluator, self.classifier, self.test_X) # expected_improvement[~(expected_improvement > 0)] = np.nan ref = ax.scatter( - *self.test_params.T[:2], + *self.test_X.T[:2], s=s, c=expected_improvement, norm=mpl.colors.Normalize(), ) clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) clb.set_label("units") - ax.scatter(*self.params.T[:2], s=s, edgecolor="k", facecolor="none") + ax.scatter(*self.X.T[:2], s=s, edgecolor="k", facecolor="none") for ax in self.state_axes.ravel(): ax.set_xlim(*self.bounds[0]) @@ -589,40 +562,36 @@ def show_state_plots(self, save_as=None): # talk to the model - def fitness_estimate(self, params): - return self.evaluator.mean(self.params_trans_fun(params).reshape(-1, self.n_dof)).reshape( - params.shape[:-1] - ) + def fitness_estimate(self, X): + return self.evaluator.mean(X.reshape(-1, self.n_dof)).reshape(X.shape[:-1]) - def fitness_sigma(self, params): - return self.evaluator.sigma(self.params_trans_fun(params).reshape(-1, self.n_dof)).reshape( - params.shape[:-1] - ) + def fitness_sigma(self, X): + return self.evaluator.sigma(X.reshape(-1, self.n_dof)).reshape(X.shape[:-1]) - def fitness_entropy(self, params): - return np.log(np.sqrt(2 * np.pi * np.e) * self.fitness_sigma(params) + 1e-12) + def fitness_entropy(self, X): + return np.log(np.sqrt(2 * np.pi * np.e) * self.fitness_sigma(X) + 1e-12) - def validate(self, params): - return self.classifier.p(self.params_trans_fun(params).reshape(-1, self.n_dof)).reshape(params.shape[:-1]) + def validate(self, X): + return self.classifier.p(X.reshape(-1, self.n_dof)).reshape(X.shape[:-1]) - def delay_estimate(self, params): - return self.timer.mean(self.params_trans_fun(params).reshape(-1, self.n_dof)).reshape(params.shape[:-1]) + def delay_estimate(self, X): + return self.timer.mean(X.reshape(-1, self.n_dof)).reshape(X.shape[:-1]) - def delay_sigma(self, params): - return self.timer.sigma(self.params_trans_fun(params).reshape(-1, self.n_dof)).reshape(params.shape[:-1]) + def delay_sigma(self, X): + return self.timer.sigma(X.reshape(-1, self.n_dof)).reshape(X.shape[:-1]) - def _negative_A_optimality(self, params): + def _negative_A_optimality(self, X): """ - The negative trace of the inverse Fisher information matrix contingent on sampling the passed params. + The negative trace of the inverse Fisher information matrix contingent on sampling the passed X. """ - test_X = self.params_trans_fun(params) + test_X = X invFIM = np.linalg.inv(self.evaluator._contingent_fisher_information_matrix(test_X, delta=1e-3)) return np.array(list(map(np.trace, invFIM))) - def _negative_D_optimality(self, params): + def _negative_D_optimality(self, X): """ - The negative determinant of the inverse Fisher information matrix contingent on sampling the passed params. + The negative determinant of the inverse Fisher information matrix contingent on sampling the passed X. """ - test_X = self.params_trans_fun(params) + test_X = X FIM_stack = self.evaluator._contingent_fisher_information_matrix(test_X, delta=1e-3) return -np.array(list(map(np.linalg.det, FIM_stack))) diff --git a/bloptools/bo/acquisition.py b/bloptools/bo/acquisition.py index 2b338f9..9bc8977 100644 --- a/bloptools/bo/acquisition.py +++ b/bloptools/bo/acquisition.py @@ -12,14 +12,17 @@ def expected_improvement(evaluator, classifier, X): Given a botorch fitness model "evaluator" and a botorch validation model "classifier", compute the expected improvement at parameters X. """ - torch_X = torch.as_tensor(X.reshape(-1, X.shape[-1])) + + model_inputs = evaluator.prepare_inputs(X.reshape(-1, X.shape[-1])) ei = np.exp( - LogExpectedImprovement(evaluator.model, best_f=evaluator.Y.max())(torch_X.unsqueeze(1)).detach().numpy() + LogExpectedImprovement(evaluator.model, best_f=evaluator.train_targets.max())(model_inputs.unsqueeze(1)) + .detach() + .numpy() ) - p_good = classifier.p(torch_X) + p_good = classifier.p(X) - return (ei * p_good).reshape(X.shape[:-1]) + return ei.reshape(X.shape[:-1]) * p_good.reshape(X.shape[:-1]) def expected_gibbon(evaluator, classifier, X, n_candidates=1024): @@ -27,15 +30,15 @@ def expected_gibbon(evaluator, classifier, X, n_candidates=1024): Given a botorch fitness model "evaluator" and a botorch validation model "classifier", compute the expected GIBBON at parameters X (https://www.jmlr.org/papers/volume22/21-0120/21-0120.pdf) """ - torch_X = torch.as_tensor(X.reshape(-1, X.shape[-1])) + model_inputs = evaluator.prepare_inputs(X.reshape(-1, X.shape[-1])) sampler = sp.stats.qmc.Halton(d=evaluator.X.shape[-1], scramble=True) candidate_set = torch.as_tensor(sampler.random(n=n_candidates)).double() - gibbon = qLowerBoundMaxValueEntropy(evaluator.model, candidate_set)(torch_X.unsqueeze(1)).detach().numpy() - p_good = classifier.p(torch_X) + gibbon = qLowerBoundMaxValueEntropy(evaluator.model, candidate_set)(model_inputs.unsqueeze(1)).detach().numpy() + p_good = classifier.p(X) - return (gibbon * p_good).reshape(X.shape[:-1]) + return gibbon.reshape(X.shape[:-1]) * p_good.reshape(X.shape[:-1]) # these return params that maximize the objective diff --git a/bloptools/bo/models.py b/bloptools/bo/models.py index 4cb2d89..59d75cc 100644 --- a/bloptools/bo/models.py +++ b/bloptools/bo/models.py @@ -42,16 +42,56 @@ def forward(self, x): return gpytorch.distributions.MultivariateNormal(mean_x, covar_x) -class GPR: +class BoTorchModelWrapper: + def __init__(self, bounds, MIN_SNR=1e-2): + self.model = None + self.bounds = bounds + self.MIN_SNR = MIN_SNR + + self.prepare_inputs = lambda inputs: torch.tensor( + (inputs - self.bounds.min(axis=1)) / self.bounds.ptp(axis=1) + ).double() + self.unprepare_inputs = lambda inputs: inputs.detach().numpy() * self.bounds.ptp(axis=1) + self.bounds.min( + axis=1 + ) + + def tell(self, X, Y): + self.model.set_train_data( + torch.cat([self.train_inputs, self.prepare_inputs(X)]), + torch.cat([self.train_targets, self.prepare_targets(Y)]), + strict=False, + ) + + def train(self, **kwargs): + botorch.optim.fit.fit_gpytorch_mll_torch(self.mll, optimizer=torch.optim.Adam, **kwargs) + + @property + def train_inputs(self): + return self.prepare_inputs(self.X) + + @property + def train_targets(self): + return self.prepare_targets(self.Y) + + @property + def n(self): + return self.X.shape[0] + + @property + def n_dof(self): + return self.X.shape[-1] + + @property + def n_tasks(self): + return self.Y.shape[-1] + + +class GPR(BoTorchModelWrapper): """ A Gaussian process regressor, with learning methods. """ - def __init__(self, MIN_SNR=1e-2): - self.model = None - self.MIN_SNR = MIN_SNR - def set_data(self, X, Y): """ Instantiate the GP with parameters and values. @@ -66,12 +106,24 @@ def set_data(self, X, Y): # prepare Gaussian process ingredients for the regressor and classifier # use only regressable points for the regressor + self.X, self.Y = X, Y + + self.target_means = Y.mean(axis=0) + self.target_scales = Y.std(axis=0) + + self.prepare_targets = lambda targets: torch.tensor( + (targets - self.target_means[None]) / self.target_scales[None] + ).double() + self.unprepare_targets = ( + lambda targets: targets.detach().numpy() * self.target_scales[None] + self.target_means[None] + ) + self.num_tasks = Y.shape[-1] - self.noise_upper_bound = np.square((1 if not Y.shape[0] > 1 else Y.std(axis=0)) / self.MIN_SNR) + self.noise_upper_bound = np.square(1 / self.MIN_SNR) likelihood_noise_constraint = gpytorch.constraints.Interval( - 0, torch.as_tensor(self.noise_upper_bound).double() + 0, torch.tensor(self.noise_upper_bound).double() ) likelihood = MultitaskGaussianLikelihood( @@ -80,52 +132,18 @@ def set_data(self, X, Y): ).double() self.model = BoTorchMultiTaskGP( - train_X=torch.as_tensor(X).double(), - train_Y=torch.as_tensor(Y).double(), + train_X=self.train_inputs, + train_Y=self.train_targets, likelihood=likelihood, ).double() self.mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.model.likelihood, self.model) - @property - def torch_inputs(self): - return self.model.train_inputs[0] - - @property - def torch_targets(self): - return self.model.train_targets - - @property - def X(self): - return self.torch_inputs.detach().numpy().astype(float) - - @property - def Y(self): - return self.torch_targets.detach().numpy().astype(float) - - @property - def n(self): - return self.Y.size - - @property - def n_dof(self): - return self.X.shape[-1] - - def tell(self, X, Y): - self.model.set_train_data( - torch.cat([self.torch_inputs, torch.as_tensor(np.atleast_2d(X))]).double(), - torch.cat([self.torch_targets, torch.as_tensor(np.atleast_2d(Y))]).double(), - strict=False, - ) - - def train(self, maxiter=100, lr=1e-2): - botorch.optim.fit.fit_gpytorch_mll_torch(self.mll, optimizer=torch.optim.Adam, step_limit=256) - def copy(self): if self.model is None: raise RuntimeError("You cannot copy a model with no data.") - dummy = GPR() + dummy = GPR(bounds=self.bounds, MIN_SNR=self.MIN_SNR) dummy.set_data(self.X, self.Y) dummy.model.load_state_dict(self.model.state_dict()) @@ -133,13 +151,13 @@ def copy(self): def mean(self, X): *input_shape, _ = X.shape - x = torch.as_tensor(X.reshape(-1, self.n_dof)).double() - return self.model.posterior(x).mean.detach().numpy().reshape(input_shape) + x = self.prepare_inputs(X).reshape(-1, self.n_dof) + return self.unprepare_targets(self.model.posterior(x).mean).reshape(input_shape) def sigma(self, X): *input_shape, _ = X.shape - x = torch.as_tensor(X.reshape(-1, self.n_dof)).double() - return self.model.posterior(x).stddev.detach().numpy().reshape(input_shape) + x = self.prepare_inputs(X).reshape(-1, self.n_dof) + return self.target_scales * self.model.posterior(x).stddev.detach().numpy().reshape(input_shape) @property def scale(self): @@ -202,14 +220,11 @@ def _contingent_fisher_information_matrix(self, test_X, delta=1e-3): return FIM_stack -class GPC: +class GPC(BoTorchModelWrapper): """ A Gaussian process classifier, with learning methods. """ - def __init__(self): - self.model = None - def set_data(self, X, c): """ Set the data with parameters and values. @@ -220,53 +235,31 @@ def set_data(self, X, c): Passed parameters must be between [-1, 1] in every dimension. Passed values must be integer labels. """ + self.X, self.c = X, c + + self.prepare_targets = lambda targets: torch.tensor(targets).int() + self.unprepare_targets = lambda targets: targets.detach().numpy().astype(int) + dirichlet_likelihood = gpytorch.likelihoods.DirichletClassificationLikelihood( torch.as_tensor(c).int(), learn_additional_noise=True ).double() self.model = BoTorchClassifier( - torch.as_tensor(X).double(), + self.train_inputs, dirichlet_likelihood.transformed_targets, dirichlet_likelihood, ).double() self.mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.model.likelihood, self.model) - @property - def torch_inputs(self): - return self.model.train_inputs[0] - - @property - def torch_targets(self): - return self.model.train_targets - - @property - def X(self): - return self.torch_inputs.detach().numpy().astype(float) - - @property - def c(self): - return self.torch_targets.detach().numpy().argmax(axis=0) - - @property - def n(self): - return self.c.size - - @property - def n_dof(self): - return self.X.shape[-1] - def tell(self, X, c): self.set_data(np.r_[self.X, np.atleast_2d(X)], np.r_[self.c, np.atleast_1d(c)]) - def train(self, maxiter=100, lr=1e-2): - botorch.optim.fit.fit_gpytorch_mll_torch(self.mll, optimizer=torch.optim.Adam, step_limit=256) - def copy(self): if self.model is None: raise RuntimeError("You cannot copy a model with no data.") - dummy = GPC() + dummy = GPC(bounds=self.bounds, MIN_SNR=self.MIN_SNR) dummy.set_data(self.X, self.c) dummy.model.load_state_dict(self.model.state_dict()) @@ -274,15 +267,11 @@ def copy(self): def p(self, X, n_samples=256): *input_shape, _ = X.shape - - x = torch.as_tensor(X.reshape(-1, self.n_dof)).double() - + x = self.prepare_inputs(X).reshape(-1, self.n_dof) samples = self.model.posterior(x).sample(torch.Size((n_samples,))).exp() - return (samples / samples.sum(-3, keepdim=True)).mean(0)[1].reshape(input_shape).detach().numpy() def entropy(self, X, n_samples=256): p = self.p(X, n_samples) q = 1 - p - return -p * np.log(p) - q * np.log(q) diff --git a/bloptools/experiments/tests/__init__.py b/bloptools/experiments/tests/__init__.py index 4e1b00d..eb2e856 100644 --- a/bloptools/experiments/tests/__init__.py +++ b/bloptools/experiments/tests/__init__.py @@ -8,7 +8,7 @@ def get_dofs(n=2): class BaseOptimizationTest: - MIN_SNR = 1e2 + MIN_SNR = 1e3 n_dof = 2 dofs = get_dofs(n=2) @@ -64,7 +64,7 @@ def fitness_func(X): class Rosenbrock(BaseOptimizationTest): """ - Ackley's function in $n$ dimensions (https://en.wikipedia.org/wiki/Ackley_function) + Rosenbrock function in $n$ dimensions (hhttps://en.wikipedia.org/wiki/Rosenbrock_function) """ def __init__(self, n_dof=2): @@ -78,9 +78,21 @@ def fitness_func(X): return -np.log((100 * (X[..., 1:] - X[..., :-1] ** 2) ** 2 + (1 - X[..., :-1]) ** 2).sum(axis=-1)) +class ConstrainedRosenbrock(Rosenbrock): + """ + Rosenbrock function in $n$ dimensions (hhttps://en.wikipedia.org/wiki/Rosenbrock_function) + """ + + @staticmethod + def fitness_func(X): + if (X**2).sum(axis=-1) > len(X): + return np.nan + return -np.log((100 * (X[..., 1:] - X[..., :-1] ** 2) ** 2 + (1 - X[..., :-1]) ** 2).sum(axis=-1)) + + class StyblinskiTang(BaseOptimizationTest): """ - Ackley's function in $n$ dimensions (https://en.wikipedia.org/wiki/Ackley_function) + Styblinski-Tang function in $n$ dimensions (https://en.wikipedia.org/wiki/Test_functions_for_optimization) """ def __init__(self, n_dof=2): From d646ac31825917f63e13e081364d8bee8c173210 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Tue, 18 Apr 2023 13:36:51 -0400 Subject: [PATCH 2/7] fixed glitch in .tell() method --- bloptools/bo/__init__.py | 4 ++-- bloptools/bo/models.py | 7 ++++++- bloptools/experiments/tests/__init__.py | 2 +- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/bloptools/bo/__init__.py b/bloptools/bo/__init__.py index c70211f..65a8767 100644 --- a/bloptools/bo/__init__.py +++ b/bloptools/bo/__init__.py @@ -69,7 +69,7 @@ def __init__( self.gridded_plots = True if self.n_dof == 2 else False - MAX_TEST_POINTS = 2**10 + MAX_TEST_POINTS = 2**11 n_bins_per_dim = int(np.power(MAX_TEST_POINTS, 1 / self.n_dof)) self.dim_bins = [np.linspace(*bounds, n_bins_per_dim + 1) for bounds in self.bounds] @@ -228,7 +228,7 @@ def ask( cost_model=None, n=1, n_test=1024, - disappointment=1, + disappointment=1.0, init=False, ): """ diff --git a/bloptools/bo/models.py b/bloptools/bo/models.py index 59d75cc..879dc1b 100644 --- a/bloptools/bo/models.py +++ b/bloptools/bo/models.py @@ -19,7 +19,9 @@ def __init__(self, train_X, train_Y, likelihood): super(BoTorchMultiTaskGP, self).__init__(train_X, train_Y, likelihood) self.mean_module = gpytorch.means.MultitaskMean(gpytorch.means.ConstantMean(), num_tasks=self._num_outputs) self.covar_module = gpytorch.kernels.MultitaskKernel( - kernels.LatentMaternKernel(n_dof=train_X.shape[-1], off_diag=True), num_tasks=self._num_outputs, rank=1 + kernels.LatentMaternKernel(n_dof=train_X.shape[-1], off_diag=False), + num_tasks=self._num_outputs, + rank=1, ) def forward(self, x): @@ -139,6 +141,9 @@ def set_data(self, X, Y): self.mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.model.likelihood, self.model) + def tell(self, X, Y): + self.set_data(np.r_[self.X, np.atleast_2d(X)], np.r_[self.Y, np.atleast_2d(Y)]) + def copy(self): if self.model is None: raise RuntimeError("You cannot copy a model with no data.") diff --git a/bloptools/experiments/tests/__init__.py b/bloptools/experiments/tests/__init__.py index eb2e856..7feda3d 100644 --- a/bloptools/experiments/tests/__init__.py +++ b/bloptools/experiments/tests/__init__.py @@ -8,7 +8,7 @@ def get_dofs(n=2): class BaseOptimizationTest: - MIN_SNR = 1e3 + MIN_SNR = 1e2 n_dof = 2 dofs = get_dofs(n=2) From 06616425695a28226de82ddec3026fb05bbe56fe Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Wed, 19 Apr 2023 12:25:53 -0400 Subject: [PATCH 3/7] added test functions to tests --- .flake8 | 2 +- bloptools/experiments/tests/__init__.py | 26 ++++++++++++++----- bloptools/tests/test_bayesian_himmelblau.py | 24 +++++++++++++++++ .../{test_boa.py => test_bayesian_shadow.py} | 2 ++ 4 files changed, 47 insertions(+), 7 deletions(-) create mode 100644 bloptools/tests/test_bayesian_himmelblau.py rename bloptools/tests/{test_boa.py => test_bayesian_shadow.py} (96%) diff --git a/.flake8 b/.flake8 index e601d73..41d1533 100644 --- a/.flake8 +++ b/.flake8 @@ -9,7 +9,7 @@ exclude = examples/benchmark.py, examples/prepare_bluesky.py, examples/prepare_tes_shadow.py, - bloptools/tests/test_boa.py, + bloptools/tests/test_bayesian_shadow.py, versioneer.py, max-line-length = 115 # Ignore some style 'errors' produced while formatting by 'black' diff --git a/bloptools/experiments/tests/__init__.py b/bloptools/experiments/tests/__init__.py index 7feda3d..30e7054 100644 --- a/bloptools/experiments/tests/__init__.py +++ b/bloptools/experiments/tests/__init__.py @@ -1,12 +1,26 @@ import bluesky.plan_stubs as bps import numpy as np -from ophyd import Signal +from ophyd import Component as Cpt +from ophyd import Device, Signal def get_dofs(n=2): return [Signal(name=f"x{i+1}", value=0) for i in range(n)] +def get_device(name="dofs", n=2): + components = {} + + for i in range(n): + components[f"{i+1}"] = Cpt(Signal, value=i + 1) + + cls = type("DOF", (Device,), components) + + device = cls(name=name) + + return [getattr(device, attr) for attr in device.read_attrs] + + class BaseOptimizationTest: MIN_SNR = 1e2 @@ -14,7 +28,7 @@ class BaseOptimizationTest: dofs = get_dofs(n=2) bounds = np.array([[-5.0, +5.0], [-5.0, +5.0]]) - DEPENDENT_COMPONENTS = [f"x{i+1}" for i in range(2)] + DEPENDENT_COMPONENTS = [dof.name for dof in dofs] def initialize(self): yield from bps.null() # do nothing @@ -34,7 +48,7 @@ def __init__(self, n_dof=2): self.n_dof = n_dof self.dofs = get_dofs(n=n_dof) self.bounds = np.array([[-5.0, +5.0] for i in range(self.n_dof)]) - self.DEPENDENT_COMPONENTS = [f"x{i+1}" for i in range(self.n_dof)] + self.DEPENDENT_COMPONENTS = [dof.name for dof in self.dofs] @staticmethod def fitness_func(X): @@ -55,7 +69,7 @@ class Himmelblau(BaseOptimizationTest): dofs = get_dofs(n=2) bounds = np.array([[-5.0, +5.0], [-5.0, +5.0]]) - DEPENDENT_COMPONENTS = [f"x{i+1}" for i in range(2)] + DEPENDENT_COMPONENTS = [dof.name for dof in dofs] @staticmethod def fitness_func(X): @@ -71,7 +85,7 @@ def __init__(self, n_dof=2): self.n_dof = n_dof self.dofs = get_dofs(n=n_dof) self.bounds = np.array([[-2.0, +2.0] for i in range(self.n_dof)]) - self.DEPENDENT_COMPONENTS = [f"x{i+1}" for i in range(self.n_dof)] + self.DEPENDENT_COMPONENTS = [dof.name for dof in self.dofs] @staticmethod def fitness_func(X): @@ -99,7 +113,7 @@ def __init__(self, n_dof=2): self.n_dof = n_dof self.dofs = get_dofs(n=n_dof) self.bounds = np.array([[-5.0, +5.0] for i in range(self.n_dof)]) - self.DEPENDENT_COMPONENTS = [f"x{i+1}" for i in range(self.n_dof)] + self.DEPENDENT_COMPONENTS = [dof.name for dof in self.dofs] @staticmethod def fitness_func(X): diff --git a/bloptools/tests/test_bayesian_himmelblau.py b/bloptools/tests/test_bayesian_himmelblau.py new file mode 100644 index 0000000..0391a4b --- /dev/null +++ b/bloptools/tests/test_bayesian_himmelblau.py @@ -0,0 +1,24 @@ +import pytest + +from bloptools.bo import BayesianOptimizationAgent +from bloptools.experiments.tests import Himmelblau + + +@pytest.mark.test_function +def test_himmelblau_boa(RE, db): + himmelblau = Himmelblau() + + boa = BayesianOptimizationAgent( + dofs=himmelblau.dofs, + dets=[], + bounds=himmelblau.bounds, + db=db, + experiment=himmelblau, + ) + + RE(boa.initialize(init_scheme="quasi-random", n_init=8)) + + RE(boa.learn(strategy="eI", n_iter=2, n_per_iter=3)) + RE(boa.learn(strategy="eGIBBON", n_iter=2, n_per_iter=3)) + + boa.plot_state(gridded=True) diff --git a/bloptools/tests/test_boa.py b/bloptools/tests/test_bayesian_shadow.py similarity index 96% rename from bloptools/tests/test_boa.py rename to bloptools/tests/test_bayesian_shadow.py index 3f0e330..2673afb 100644 --- a/bloptools/tests/test_boa.py +++ b/bloptools/tests/test_bayesian_shadow.py @@ -1,10 +1,12 @@ import numpy as np +import pytest from sirepo_bluesky.sirepo_ophyd import create_classes from bloptools.bo import BayesianOptimizationAgent from bloptools.experiments.shadow import tes +@pytest.mark.shadow def test_tes_shadow_boa(RE, db, shadow_tes_simulation): data, schema = shadow_tes_simulation.auth("shadow", "00000002") classes, objects = create_classes(shadow_tes_simulation.data, connection=shadow_tes_simulation) From 25523db0090abcb8b990e75c81ad341618a1c3d3 Mon Sep 17 00:00:00 2001 From: Max Rakitin Date: Sun, 23 Apr 2023 12:58:10 -0400 Subject: [PATCH 4/7] Fix the device generation funciton --- bloptools/experiments/tests/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bloptools/experiments/tests/__init__.py b/bloptools/experiments/tests/__init__.py index 30e7054..5eba20d 100644 --- a/bloptools/experiments/tests/__init__.py +++ b/bloptools/experiments/tests/__init__.py @@ -12,7 +12,7 @@ def get_device(name="dofs", n=2): components = {} for i in range(n): - components[f"{i+1}"] = Cpt(Signal, value=i + 1) + components[f"x{i+1}"] = Cpt(Signal, value=i + 1) cls = type("DOF", (Device,), components) From 21ef5c00bf421852f38af592122ac4aac4eb6973 Mon Sep 17 00:00:00 2001 From: Thomas Morris <41275226+thomaswmorris@users.noreply.github.com> Date: Sun, 23 Apr 2023 12:58:41 -0400 Subject: [PATCH 5/7] Update bloptools/experiments/tests/__init__.py Co-authored-by: Max Rakitin --- bloptools/experiments/tests/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bloptools/experiments/tests/__init__.py b/bloptools/experiments/tests/__init__.py index 5eba20d..4e5e581 100644 --- a/bloptools/experiments/tests/__init__.py +++ b/bloptools/experiments/tests/__init__.py @@ -78,7 +78,7 @@ def fitness_func(X): class Rosenbrock(BaseOptimizationTest): """ - Rosenbrock function in $n$ dimensions (hhttps://en.wikipedia.org/wiki/Rosenbrock_function) + Rosenbrock function in $n$ dimensions (https://en.wikipedia.org/wiki/Rosenbrock_function) """ def __init__(self, n_dof=2): From e2229644817dcbe452c5714ca0a6c342e0cb438c Mon Sep 17 00:00:00 2001 From: Thomas Morris <41275226+thomaswmorris@users.noreply.github.com> Date: Sun, 23 Apr 2023 13:42:30 -0400 Subject: [PATCH 6/7] Update bloptools/experiments/tests/__init__.py Co-authored-by: Max Rakitin --- bloptools/experiments/tests/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bloptools/experiments/tests/__init__.py b/bloptools/experiments/tests/__init__.py index 4e5e581..6a4c2d4 100644 --- a/bloptools/experiments/tests/__init__.py +++ b/bloptools/experiments/tests/__init__.py @@ -94,7 +94,7 @@ def fitness_func(X): class ConstrainedRosenbrock(Rosenbrock): """ - Rosenbrock function in $n$ dimensions (hhttps://en.wikipedia.org/wiki/Rosenbrock_function) + Rosenbrock function in $n$ dimensions (https://en.wikipedia.org/wiki/Rosenbrock_function) """ @staticmethod From f9e582f82dac53305e31452ee7ca5e694c088413 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Sun, 23 Apr 2023 13:45:59 -0400 Subject: [PATCH 7/7] notebook tweaks --- .flake8 | 4 +- bloptools/bo/__init__.py | 32 +++++------ bloptools/bo/models.py | 2 +- bloptools/experiments/shadow/chx.py | 66 +++++++++++++++++++++++ bloptools/experiments/tests/__init__.py | 16 ++++++ docs/source/notebooks/himmelblau-2d.ipynb | 34 +++++++++--- examples/prepare_chx_shadow.py | 22 ++++++++ 7 files changed, 151 insertions(+), 25 deletions(-) create mode 100644 bloptools/experiments/shadow/chx.py create mode 100644 examples/prepare_chx_shadow.py diff --git a/.flake8 b/.flake8 index 41d1533..a82760d 100644 --- a/.flake8 +++ b/.flake8 @@ -6,9 +6,7 @@ exclude = build, dist, docs/source/conf.py - examples/benchmark.py, - examples/prepare_bluesky.py, - examples/prepare_tes_shadow.py, + examples/*.py, bloptools/tests/test_bayesian_shadow.py, versioneer.py, max-line-length = 115 diff --git a/bloptools/bo/__init__.py b/bloptools/bo/__init__.py index 65a8767..f9eeb0b 100644 --- a/bloptools/bo/__init__.py +++ b/bloptools/bo/__init__.py @@ -74,7 +74,7 @@ def __init__( n_bins_per_dim = int(np.power(MAX_TEST_POINTS, 1 / self.n_dof)) self.dim_bins = [np.linspace(*bounds, n_bins_per_dim + 1) for bounds in self.bounds] self.dim_mids = [0.5 * (bins[1:] + bins[:-1]) for bins in self.dim_bins] - self.test_grid = np.swapaxes(np.r_[np.meshgrid(*self.dim_mids, indexing="ij")], 0, -1) + self.test_X_grid = np.swapaxes(np.r_[np.meshgrid(*self.dim_mids, indexing="ij")], 0, -1) sampler = sp.stats.qmc.Halton(d=self.n_dof, scramble=True) self.test_X = sampler.random(n=MAX_TEST_POINTS) * self.bounds.ptp(axis=1) + self.bounds.min(axis=1) @@ -421,7 +421,7 @@ def update_state_plots(self, gridded=False): ax.set_title("fitness estimate") if gridded: ref = ax.pcolormesh( - *self.dim_mids[:2], self.fitness_estimate(self.test_grid), norm=fitness_norm, shading="nearest" + *self.dim_mids[:2], self.fitness_estimate(self.test_X_grid), norm=fitness_norm, shading="nearest" ) else: ref = ax.scatter(*self.test_X.T[:2], s=s, c=self.fitness_estimate(self.test_X), norm=fitness_norm) @@ -435,13 +435,13 @@ def update_state_plots(self, gridded=False): ax.clear() ax.set_title("fitness uncertainty") if gridded: - ref = ax.pcolormesh(*self.dim_mids[:2], self.fitness_sigma(self.test_grid), shading="nearest") + ref = ax.pcolormesh(*self.dim_mids[:2], self.fitness_sigma(self.test_X_grid), shading="nearest") else: ref = ax.scatter( *self.test_X.T[:2], s=s, c=self.fitness_sigma(self.test_X), - norm=mpl.colors.LogNorm(), + norm=mpl.colors.Normalize(), ) clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) @@ -455,7 +455,7 @@ def update_state_plots(self, gridded=False): if gridded: expected_improvement = acquisition.expected_improvement( - self.evaluator, self.classifier, self.test_grid + self.evaluator, self.classifier, self.test_X_grid ) # expected_improvement[~(expected_improvement > 0)] = np.nan ref = ax.pcolormesh( @@ -495,7 +495,7 @@ def update_state_plots(self, gridded=False): if gridded: ref = ax.pcolormesh( *self.dim_mids[:2], - self.classifier.p(self.test_grid), + self.classifier.p(self.test_X_grid), vmin=0, vmax=1, shading="nearest", @@ -518,7 +518,7 @@ def update_state_plots(self, gridded=False): if gridded: ref = ax.pcolormesh( *self.dim_mids[:2], - self.classifier.entropy(self.test_grid), + self.classifier.entropy(self.test_X_grid), shading="nearest", ) else: @@ -529,24 +529,26 @@ def update_state_plots(self, gridded=False): ax = self.state_axes[1, 3] ax.clear() - ax.set_title("expected GIBBON") + ax.set_title("total entropy") if gridded: - expected_improvement = acquisition.expected_gibbon(self.evaluator, self.classifier, self.test_grid) - # expected_improvement[~(expected_improvement > 0)] = np.nan + expected_gibbon = acquisition.expected_gibbon(self.evaluator, self.classifier, self.test_X_grid) + # total_entropy = self.evaluator.normalized_entropy(self.test_X_grid) \ + # + self.classifier.entropy(self.test_X_grid) ref = ax.pcolormesh( - *self.dim_mids[:2], expected_improvement, norm=mpl.colors.Normalize(), shading="nearest" + *self.dim_mids[:2], expected_gibbon, norm=mpl.colors.Normalize(), shading="nearest" ) else: - expected_improvement = acquisition.expected_gibbon(self.evaluator, self.classifier, self.test_X) - # expected_improvement[~(expected_improvement > 0)] = np.nan + expected_gibbon = acquisition.expected_gibbon(self.evaluator, self.classifier, self.test_X) + # total_entropy = self.evaluator.normalized_entropy(self.test_X) \ + # + self.classifier.entropy(self.test_X) ref = ax.scatter( *self.test_X.T[:2], s=s, - c=expected_improvement, + c=expected_gibbon, norm=mpl.colors.Normalize(), ) clb = self.state_fig.colorbar(ref, ax=ax, location="bottom", aspect=32, shrink=0.8) - clb.set_label("units") + clb.set_label("nats") ax.scatter(*self.X.T[:2], s=s, edgecolor="k", facecolor="none") for ax in self.state_axes.ravel(): diff --git a/bloptools/bo/models.py b/bloptools/bo/models.py index 879dc1b..6d4da6b 100644 --- a/bloptools/bo/models.py +++ b/bloptools/bo/models.py @@ -19,7 +19,7 @@ def __init__(self, train_X, train_Y, likelihood): super(BoTorchMultiTaskGP, self).__init__(train_X, train_Y, likelihood) self.mean_module = gpytorch.means.MultitaskMean(gpytorch.means.ConstantMean(), num_tasks=self._num_outputs) self.covar_module = gpytorch.kernels.MultitaskKernel( - kernels.LatentMaternKernel(n_dof=train_X.shape[-1], off_diag=False), + kernels.LatentMaternKernel(n_dof=train_X.shape[-1], off_diag=True), num_tasks=self._num_outputs, rank=1, ) diff --git a/bloptools/experiments/shadow/chx.py b/bloptools/experiments/shadow/chx.py new file mode 100644 index 0000000..11b43f6 --- /dev/null +++ b/bloptools/experiments/shadow/chx.py @@ -0,0 +1,66 @@ +import bluesky.plan_stubs as bps +import numpy as np + +from ... import utils + +DEPENDENT_COMPONENTS = ["sample"] + +IMAGE_NAME = "sample_image" + +HORIZONTAL_EXTENT_NAME = "sample_horizontal_extent" +VERTICAL_EXTENT_NAME = "sample_vertical_extent" + +PCA_BEAM_PROP = 0.5 # how much of the first principle component we want to have in our bounding box +MIN_SEPARABILITY = 0.1 # the minimal variance proportion of the first SVD mode of the beam image + +MIN_SNR = 1e1 + + +def initialize(): + yield from bps.null() # do nothing + + +def parse_entry(entry): + # get the ingredient from our dependent variables + image = getattr(entry, IMAGE_NAME) + horizontal_extent = getattr(entry, HORIZONTAL_EXTENT_NAME) + vertical_extent = getattr(entry, VERTICAL_EXTENT_NAME) + + if not image.sum() > 0: + image = np.random.uniform(size=image.shape) + horizontal_extent = [np.nan, np.nan] + vertical_extent = [np.nan, np.nan] + + pix_x_min, pix_x_max, pix_y_min, pix_y_max, separability = utils.get_principal_component_bounds( + image, PCA_BEAM_PROP + ) + n_y, n_x = image.shape + x_min, x_max = np.interp([pix_x_min, pix_x_max], [0, n_x + 1], horizontal_extent) + y_min, y_max = np.interp([pix_y_min, pix_y_max], [0, n_y + 1], vertical_extent) + + width_x = x_max - x_min + width_y = y_max - y_min + + pos_x = (x_min + x_max) / 2 + pos_y = (y_min + y_max) / 2 + + flux = image.sum() + + bad = False + bad |= separability < MIN_SEPARABILITY + + if bad: + fitness = np.nan + + else: + fitness = np.log(flux * separability / (width_x**2 + width_y**2)) + + return ("fitness", "flux", "pos_x", "pos_y", "width_x", "width_y", "separability"), ( + fitness, + flux, + pos_x, + pos_y, + width_x, + width_y, + separability, + ) diff --git a/bloptools/experiments/tests/__init__.py b/bloptools/experiments/tests/__init__.py index 30e7054..8cf9659 100644 --- a/bloptools/experiments/tests/__init__.py +++ b/bloptools/experiments/tests/__init__.py @@ -76,6 +76,21 @@ def fitness_func(X): return -((X[0] ** 2 + X[1] - 11) ** 2 + (X[0] + X[1] ** 2 - 7) ** 2) +class ConstrainedHimmelblau(Himmelblau): + """ + Himmelblau's function (https://en.wikipedia.org/wiki/Himmelblau%27s_function) + Return NaN if ||X|| > dim(X) + """ + + bounds = np.array([[-10.0, +10.0], [-10.0, +10.0]]) + + @staticmethod + def fitness_func(X): + if np.sqrt(np.sum(np.square(X))) > 8: + return np.nan + return -((X[0] ** 2 + X[1] - 11) ** 2 + (X[0] + X[1] ** 2 - 7) ** 2) + + class Rosenbrock(BaseOptimizationTest): """ Rosenbrock function in $n$ dimensions (hhttps://en.wikipedia.org/wiki/Rosenbrock_function) @@ -95,6 +110,7 @@ def fitness_func(X): class ConstrainedRosenbrock(Rosenbrock): """ Rosenbrock function in $n$ dimensions (hhttps://en.wikipedia.org/wiki/Rosenbrock_function) + Return NaN if ||X|| > dim(X) """ @staticmethod diff --git a/docs/source/notebooks/himmelblau-2d.ipynb b/docs/source/notebooks/himmelblau-2d.ipynb index 444dbf0..3948bc8 100644 --- a/docs/source/notebooks/himmelblau-2d.ipynb +++ b/docs/source/notebooks/himmelblau-2d.ipynb @@ -67,6 +67,28 @@ "boa.plot_state(gridded=True)" ] }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "296d9fd2", + "metadata": {}, + "source": [ + "Now let's try the \"EI\" strategy to sample where we expect the largest improvement in the fitness:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f0e74651", + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "RE(boa.learn(strategy='eI', n_iter=1, n_per_iter=4))\n", + "boa.plot_state(gridded=True)" + ] + }, { "attachments": {}, "cell_type": "markdown", @@ -85,29 +107,29 @@ }, "outputs": [], "source": [ - "RE(boa.learn(strategy='eGIBBON', n_iter=2, n_per_iter=4))\n", + "RE(boa.learn(strategy='eGIBBON', n_iter=4, n_per_iter=4))\n", "boa.plot_state(gridded=True)" ] }, { "attachments": {}, "cell_type": "markdown", - "id": "296d9fd2", + "id": "e955233f", "metadata": {}, "source": [ - "Now let's try the \"EI\" strategy to sample where we expect the largest improvement in the fitness:" + "Eventually, we reach a point of saturation where no more improvement takes place:" ] }, { "cell_type": "code", "execution_count": null, - "id": "f0e74651", + "id": "d73e4fd5", "metadata": { "tags": [] }, "outputs": [], "source": [ - "RE(boa.learn(strategy='eI', n_iter=4, n_per_iter=4))\n", + "RE(boa.learn(strategy='eGIBBON', n_iter=8, n_per_iter=4))\n", "boa.plot_state(gridded=True)" ] } @@ -128,7 +150,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16 (main, Mar 8 2023, 14:00:05) \n[GCC 11.2.0]" + "version": "3.9.16" }, "vscode": { "interpreter": { diff --git a/examples/prepare_chx_shadow.py b/examples/prepare_chx_shadow.py new file mode 100644 index 0000000..252476f --- /dev/null +++ b/examples/prepare_chx_shadow.py @@ -0,0 +1,22 @@ +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) + +plt.ion() + +root_dir = "/tmp/sirepo-bluesky-data" +_ = make_dir_tree(datetime.datetime.now().year, base_path=root_dir) + +connection = SirepoBluesky("http://localhost:8000") + +data, schema = connection.auth("shadow", "I1Flcbdw") +classes, objects = create_classes(connection.data, connection=connection) +globals().update(**objects) + +# data["models"]["simulation"]["npoint"] = 100000 +# data["models"]["watchpointReport12"]["histogramBins"] = 32 +# w9.duration.kind = "hinted" + +bec.disable_baseline() +bec.disable_heading() +bec.disable_table()