diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 55516cd..4c964f4 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -4,8 +4,8 @@ on: push: pull_request: workflow_dispatch: - # schedule: - # - cron: '00 4 * * *' # daily at 4AM + schedule: + - cron: '00 6 * * *' # daily at 6AM UTC jobs: run_tests: diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst index d146376..64bd580 100644 --- a/docs/source/tutorials.rst +++ b/docs/source/tutorials.rst @@ -8,3 +8,4 @@ Tutorials tutorials/hyperparameters.ipynb tutorials/pareto-fronts.ipynb tutorials/passive-dofs.ipynb + tutorials/kb-mirrors.ipynb diff --git a/docs/source/tutorials/hyperparameters.ipynb b/docs/source/tutorials/hyperparameters.ipynb index e2c03b5..f58f10c 100644 --- a/docs/source/tutorials/hyperparameters.ipynb +++ b/docs/source/tutorials/hyperparameters.ipynb @@ -153,7 +153,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.0" + "version": "3.11.9" }, "vscode": { "interpreter": { diff --git a/docs/source/tutorials/kb-mirrors.ipynb b/docs/source/tutorials/kb-mirrors.ipynb new file mode 100644 index 0000000..397ea78 --- /dev/null +++ b/docs/source/tutorials/kb-mirrors.ipynb @@ -0,0 +1,142 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "c1ac0e4b-d065-41af-97ec-b73bbc7dad7d", + "metadata": {}, + "source": [ + "# Hyperparameters\n", + "\n", + "This example simulates the alignment of a KB mirror endstation (with four degrees of freedom)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2a1798fa-b251-409e-9d2d-097240372b03", + "metadata": {}, + "outputs": [], + "source": [ + "from blop.utils import prepare_re_env\n", + "\n", + "%run -i $prepare_re_env.__file__ --db-type=temp\n", + "bec.disable_plots()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0ca926ab-5859-49f3-a96a-73c401cc18e6", + "metadata": {}, + "outputs": [], + "source": [ + "from blop.sim import Beamline\n", + "\n", + "beamline = Beamline(name=\"bl\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "28ee6cfc-428a-4472-b503-bb072f939866", + "metadata": {}, + "outputs": [], + "source": [ + "from blop import DOF, Objective, Agent\n", + "from blop.digestion import beam_stats_digestion\n", + "\n", + "dofs = [\n", + " DOF(description=\"KBV downstream\", device=beamline.kbv_dsv, search_domain=(-5.0, 5.0)),\n", + " DOF(description=\"KBV upstream\", device=beamline.kbv_usv, search_domain=(-5.0, 5.0)),\n", + " DOF(description=\"KBH downstream\", device=beamline.kbh_dsh, search_domain=(-5.0, 5.0)),\n", + " DOF(description=\"KBH upstream\", device=beamline.kbh_ush, search_domain=(-5.0, 5.0)),\n", + "]\n", + "\n", + "objectives = [\n", + " Objective(name=\"bl_det_sum\", target=\"max\", transform=\"log\", trust_domain=(200, np.inf)),\n", + " Objective(name=\"bl_det_wid_x\", target=\"min\", transform=\"log\", latent_groups=[(\"bl_kbh_dsh\", \"bl_kbh_ush\")]),\n", + " Objective(name=\"bl_det_wid_y\", target=\"min\", transform=\"log\", latent_groups=[(\"bl_kbv_dsv\", \"bl_kbv_usv\")]),\n", + "]\n", + "\n", + "agent = Agent(\n", + " dofs=dofs,\n", + " objectives=objectives,\n", + " detectors=[beamline.det],\n", + " digestion=beam_stats_digestion,\n", + " digestion_kwargs={\"image_key\": \"bl_det_image\"},\n", + " verbose=True,\n", + " db=db,\n", + " tolerate_acquisition_errors=False,\n", + " enforce_all_objectives_valid=True,\n", + " train_every=3,\n", + ")\n", + "\n", + "(uid,) = RE(agent.learn(\"qr\", n=32))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "491636be-1e04-45a2-a622-dc21e192d208", + "metadata": {}, + "outputs": [], + "source": [ + "RE(agent.learn(\"qei\", n=4, iterations=4))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d6434fc9-682a-45fa-b100-465dfda3aff1", + "metadata": {}, + "outputs": [], + "source": [ + "plt.imshow(agent.best.bl_det_image)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ef89c0cd-65be-4d1f-ab6d-c8a4ce251d7b", + "metadata": {}, + "outputs": [], + "source": [ + "agent.plot_objectives(axes=(2, 3))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb56db05-0fe5-46a5-8435-383f1e34f55d", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3.12.2 64-bit", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.2" + }, + "vscode": { + "interpreter": { + "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e" + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index 8b532d9..0106599 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,13 +15,14 @@ maintainers = [ requires-python = ">=3.9" dependencies = [ + "area-detector-handlers", "bluesky", "botorch", "databroker", "gpytorch", "h5py", "matplotlib", - "numpy", + "numpy<2", "ophyd", "ortools", "scipy", diff --git a/src/blop/agent.py b/src/blop/agent.py index 80ef55f..93f8bfd 100644 --- a/src/blop/agent.py +++ b/src/blop/agent.py @@ -29,6 +29,7 @@ from . import plotting, utils from .bayesian import acquisition, models from .bayesian.acquisition import _construct_acqf, parse_acqf_identifier +from .bayesian.models import construct_single_task_model, train_model # from .bayesian.transforms import TargetingPosteriorTransform from .digestion import default_digestion_function @@ -66,13 +67,15 @@ def __init__( dofs: Sequence[DOF], objectives: Sequence[Objective], db: Broker = None, - dets: Sequence[Signal] = [], + detectors: Sequence[Signal] = None, acquistion_plan=default_acquisition_plan, digestion: Callable = default_digestion_function, digestion_kwargs: dict = {}, verbose: bool = False, - tolerate_acquisition_errors=False, - sample_center_on_init=False, + enforce_all_objectives_valid: bool = True, + model_inactive_objectives: bool = False, + tolerate_acquisition_errors: bool = False, + sample_center_on_init: bool = False, trigger_delay: float = 0, train_every: int = 1, ): @@ -85,7 +88,7 @@ def __init__( The degrees of freedom that the agent can control, which determine the output of the model. objectives : iterable of Objective objects The objectives which the agent will try to optimize. - dets : iterable of ophyd objects + detectors : iterable of ophyd objects Detectors to trigger during acquisition. acquisition_plan : optional A plan that samples the beamline for some given inputs. @@ -125,31 +128,36 @@ def __init__( self.dofs = DOFList(list(np.atleast_1d(dofs))) self.objectives = ObjectiveList(list(np.atleast_1d(objectives))) + self.detectors = list(np.atleast_1d(detectors or [])) _validate_dofs_and_objs(self.dofs, self.objectives) self.db = db - - self.dets = dets self.acquisition_plan = acquistion_plan self.digestion = digestion self.digestion_kwargs = digestion_kwargs self.verbose = verbose + self.model_inactive_objectives = model_inactive_objectives self.tolerate_acquisition_errors = tolerate_acquisition_errors + self.enforce_all_objectives_valid = enforce_all_objectives_valid self.train_every = train_every self.trigger_delay = trigger_delay self.sample_center_on_init = sample_center_on_init - self.table = pd.DataFrame() + self._table = pd.DataFrame() self.initialized = False self.a_priori_hypers = None self.n_last_trained = 0 + @property + def table(self): + return self._table + def unpack_run(self): return @@ -172,7 +180,7 @@ def refresh(self): self._train_all_models() def redigest(self): - self.table = self.digestion(self.table, **self.digestion_kwargs) + self._table = self.digestion(self._table, **self.digestion_kwargs) def sample(self, n: int = DEFAULT_MAX_SAMPLES, method: str = "quasi-random") -> torch.Tensor: """ @@ -249,7 +257,7 @@ def ask(self, acqf="qei", n=1, route=True, sequential=True, upsample=1, **acqf_k for obj in active_objs: if obj.model_dofs != set(active_dofs.names): self._construct_model(obj) - self._train_model(obj.model) + train_model(obj.model) if acqf_config["type"] == "analytic" and n > 1: raise ValueError("Can't generate multiple design points for analytic acquisition functions.") @@ -351,11 +359,12 @@ def tell( raise ValueError("All supplies values must be the same length!") new_table = pd.DataFrame(data) - self.table = pd.concat([self.table, new_table]) if append else new_table - self.table.index = np.arange(len(self.table)) + self._table = pd.concat([self._table, new_table]) if append else new_table + self._table.index = np.arange(len(self._table)) if update_models: - for obj in self.objectives(active=True): + objectives_to_model = self.objectives if self.model_inactive_objectives else self.objectives(active=True) + for obj in objectives_to_model: t0 = ttime.monotonic() cached_hypers = obj.model.state_dict() if hasattr(obj, "model") else None @@ -369,12 +378,12 @@ def tell( if len(obj.model.train_targets) >= 4: if train: t0 = ttime.monotonic() - self._train_model(obj.model) + train_model(obj.model) if self.verbose: print(f"trained model '{obj.name}' in {1e3*(ttime.monotonic() - t0):.00f} ms") else: - self._train_model(obj.model, hypers=cached_hypers) + train_model(obj.model, hypers=cached_hypers) def learn( self, @@ -498,10 +507,10 @@ def acquire(self, points): uid = yield from self.acquisition_plan( acquisition_dofs, points, - [*self.dets, *self.dofs.devices], + [*self.detectors, *self.dofs.devices], delay=self.trigger_delay, ) - products = self.digestion(self.db[uid].table(), **self.digestion_kwargs) + products = self.digestion(self.db[uid].table(fill=True), **self.digestion_kwargs) except KeyboardInterrupt as interrupt: raise interrupt @@ -521,12 +530,12 @@ def acquire(self, points): def load_data(self, data_file, append=True): new_table = pd.read_hdf(data_file, key="table") - self.table = pd.concat([self.table, new_table]) if append else new_table + self._table = pd.concat([self._table, new_table]) if append else new_table self.refresh() def reset(self): """Reset the agent.""" - self.table = pd.DataFrame() + self._table = pd.DataFrame() for obj in self.objectives(active=True): if hasattr(obj, "model"): @@ -584,10 +593,11 @@ def fitness_model(self): @property def evaluated_constraints(self): constraint_objectives = self.objectives(kind="constraint") + raw_targets_dict = self.raw_targets() if len(constraint_objectives): - return torch.cat([obj.constrain(self.raw_targets(obj.name)) for obj in constraint_objectives], dim=-1) + return torch.cat([obj.constrain(raw_targets_dict[obj.name]) for obj in constraint_objectives], dim=-1) else: - return torch.ones(size=(len(self.table), 0), dtype=torch.bool) + return torch.ones(size=(len(self._table), 0), dtype=torch.bool) def fitness_scalarization(self, weights="default"): fitness_objectives = self.objectives(active=True, kind="fitness") @@ -610,10 +620,12 @@ def scalarized_fitnesses(self, weights="default", constrained=True): """ fitness_objs = self.objectives(kind="fitness") if len(fitness_objs) >= 1: - f = self.fitness_scalarization(weights=weights).evaluate(self.train_targets(active=True, kind="fitness")) + f = self.fitness_scalarization(weights=weights).evaluate( + self.train_targets(active=True, kind="fitness", concatenate=True) + ) f = torch.where(f.isnan(), -np.inf, f) # remove all nans else: - f = torch.zeros(len(self.table), dtype=torch.double) # if there are no fitnesses, use a constant dummy fitness + f = torch.zeros(len(self._table), dtype=torch.double) # if there are no fitnesses, use a constant dummy fitness if constrained: # how many constraints are satisfied? c = self.evaluated_constraints.sum(axis=-1) @@ -632,7 +644,7 @@ def pareto_mask(self): Returns a mask of all points that satisfy all constraints and are Pareto efficient. A point is Pareto efficient if it is there is no other point that is better at every objective. """ - Y = self.train_targets(active=True, kind="fitness") + Y = self.train_targets(active=True, kind="fitness", concatenate=True) # nuke the bad points Y[~self.evaluated_constraints.all(axis=-1)] = -np.inf @@ -646,30 +658,22 @@ def pareto_front(self): """ A subset of the data table containing only points on the Pareto front. """ - return self.table.loc[self.pareto_mask.numpy()] + return self._table.loc[self.pareto_mask.numpy()] @property def min_ref_point(self): - y = self.train_targets()[:, self.objectives.type == "fitness"] + y = self.train_targets(concatenate=True)[:, self.objectives.type == "fitness"] return y[y.argmax(axis=0)].min(axis=0).values @property def random_ref_point(self): - return self.train_targets(active=True, kind="fitness")[self.argmax_best_f(weights="random")] + return self.train_targets(active=True, kind="fitness", concatenate=True)[self.argmax_best_f(weights="random")] @property def all_objectives_valid(self): """A mask of whether all objectives are valid for each data point.""" return ~torch.isnan(self.scalarized_fitnesses()) - def _train_model(self, model, hypers=None, **kwargs): - """Fit all of the agent's models. All kwargs are passed to `botorch.fit.fit_gpytorch_mll`.""" - if hypers is not None: - model.load_state_dict(hypers) - else: - botorch.fit.fit_gpytorch_mll(gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model), **kwargs) - model.trained = True - def _construct_model(self, obj, skew_dims=None): """ Construct an untrained model for an objective. @@ -677,30 +681,20 @@ def _construct_model(self, obj, skew_dims=None): skew_dims = skew_dims if skew_dims is not None else self._latent_dim_tuples(obj.name) - likelihood = gpytorch.likelihoods.GaussianLikelihood( - noise_constraint=gpytorch.constraints.Interval( - torch.tensor(obj.min_noise), - torch.tensor(obj.max_noise), - ), - ) - - outcome_transform = botorch.models.transforms.outcome.Standardize(m=1) # , batch_shape=torch.Size((1,))) - train_inputs = self.train_inputs(active=True) - train_targets = self.train_targets(obj.name) + train_targets = self.train_targets()[obj.name].unsqueeze(-1) inputs_are_trusted = ~torch.isnan(train_inputs).any(axis=1) targets_are_trusted = ~torch.isnan(train_targets).any(axis=1) trusted = inputs_are_trusted & targets_are_trusted - obj.model = models.LatentGP( - train_inputs=train_inputs[trusted], - train_targets=train_targets[trusted], - likelihood=likelihood, - skew_dims=skew_dims, - input_transform=self.input_normalization, - outcome_transform=outcome_transform, + obj.model = construct_single_task_model( + X=train_inputs[trusted], + y=train_targets[trusted], + min_noise=obj.min_noise, + max_noise=obj.max_noise, + skew_dims=self._latent_dim_tuples()[obj.name], ) obj.model_dofs = set(self.dofs(active=True).names) # if these change, retrain the model on self.ask() @@ -728,21 +722,23 @@ def _construct_model(self, obj, skew_dims=None): def _construct_all_models(self): """Construct a model for each objective.""" - for obj in self.objectives(active=True): + objectives_to_construct = self.objectives if self.model_inactive_objectives else self.objectives(active=True) + for obj in objectives_to_construct: self._construct_model(obj) def _train_all_models(self, **kwargs): """Fit all of the agent's models. All kwargs are passed to `botorch.fit.fit_gpytorch_mll`.""" t0 = ttime.monotonic() - for obj in self.objectives(active=True): - self._train_model(obj.model) + objectives_to_train = self.objectives if self.model_inactive_objectives else self.objectives(active=True) + for obj in objectives_to_train: + train_model(obj.model) if obj.validity_conjugate_model is not None: - self._train_model(obj.validity_conjugate_model) + train_model(obj.validity_conjugate_model) if self.verbose: print(f"trained models in {ttime.monotonic() - t0:.01f} seconds") - self.n_last_trained = len(self.table) + self.n_last_trained = len(self._table) def _get_acquisition_function(self, identifier, return_metadata=False): """Returns a BoTorch acquisition function for a given identifier. Acquisition functions can be @@ -808,7 +804,7 @@ def save_data(self, path="./data.h5"): save_dir, _ = os.path.split(path) pathlib.Path(save_dir).mkdir(parents=True, exist_ok=True) - self.table.to_hdf(path, key="table") + self._table.to_hdf(path, key="table") def forget(self, last=None, index=None, train=True): """ @@ -823,12 +819,12 @@ def forget(self, last=None, index=None, train=True): """ if last is not None: - if last > len(self.table): - raise ValueError(f"Cannot forget last {last} data points (only {len(self.table)} samples have been taken).") - self.forget(index=self.table.index.values[-last:], train=train) + if last > len(self._table): + raise ValueError(f"Cannot forget last {last} data points (only {len(self._table)} samples have been taken).") + self.forget(index=self._table.index.values[-last:], train=train) elif index is not None: - self.table.drop(index=index, inplace=True) + self._table.drop(index=index, inplace=True) self._construct_all_models() if train: self._train_all_models() @@ -842,8 +838,6 @@ def _set_hypers(self, hypers): self.validity_constraint.load_state_dict(hypers["validity_constraint"]) def constraint(self, x): - x = self.dofs(active=True).transform(x) - p = torch.ones(x.shape[:-1]) for obj in self.objectives(active=True): # if the targeting constraint is non-trivial @@ -914,7 +908,7 @@ def raw_inputs(self, index=None, **subset_kwargs): """ if index is None: return torch.cat([self.raw_inputs(dof.name) for dof in self.dofs(**subset_kwargs)], dim=-1) - return torch.tensor(self.table.loc[:, self.dofs[index].name].values, dtype=torch.double).unsqueeze(-1) + return torch.tensor(self._table.loc[:, self.dofs[index].name].values, dtype=torch.double).unsqueeze(-1) def train_inputs(self, index=None, **subset_kwargs): """A two-dimensional tensor of all DOF values.""" @@ -935,34 +929,52 @@ def raw_targets(self, index=None, **subset_kwargs): """ Get the raw, untransformed inputs for an objective (or for a subset). """ - if index is None: - return torch.cat([self.raw_targets(index=obj.name) for obj in self.objectives(**subset_kwargs)], dim=-1) - return torch.tensor(self.table.loc[:, self.objectives[index].name].values, dtype=torch.double).unsqueeze(-1) + values = {} - def train_targets(self, index=None, **subset_kwargs): + for obj in self.objectives(**subset_kwargs): + # return torch.cat([self.raw_targets(index=obj.name) for obj in self.objectives(**subset_kwargs)], dim=-1) + values[obj.name] = torch.tensor(self._table.loc[:, obj.name].values, dtype=torch.double) + + return values + + def train_targets(self, concatenate=False, **subset_kwargs): """Returns the values associated with an objective name.""" - if index is None: - return torch.cat([self.train_targets(obj.name) for obj in self.objectives(**subset_kwargs)], dim=-1) + targets_dict = {} + raw_targets_dict = self.raw_targets(**subset_kwargs) + + for obj in self.objectives(**subset_kwargs): + y = raw_targets_dict[obj.name] + + # check that targets values are inside acceptable values + valid = (y >= obj._trust_domain[0]) & (y <= obj._trust_domain[1]) + y = torch.where(valid, y, np.nan) + + targets_dict[obj.name] = obj._transform(y) + + if self.enforce_all_objectives_valid: + all_valid_mask = True + + for name, values in targets_dict.items(): + all_valid_mask &= ~values.isnan() - obj = self.objectives[index] - raw_targets = self.raw_targets(index=index, **subset_kwargs) + for name in targets_dict.keys(): + targets_dict[name] = targets_dict[name].where(all_valid_mask, np.nan) - # check that targets values are inside acceptable values - valid = (raw_targets >= obj._trust_domain[0]) & (raw_targets <= obj._trust_domain[1]) - raw_targets = torch.where(valid, raw_targets, np.nan) + if concatenate: + return torch.cat([values.unsqueeze(-1) for values in targets_dict.values()], axis=-1) - return obj._transform(raw_targets) + return targets_dict @property def best(self): """Returns all data for the best point.""" - return self.table.loc[self.argmax_best_f()] + return self._table.loc[self.argmax_best_f()] @property def best_inputs(self): """Returns the value of each DOF at the best point.""" - return self.table.loc[self.argmax_best_f(), self.dofs.names].to_dict() + return self._table.loc[self.argmax_best_f(), self.dofs.names].to_dict() def go_to(self, **positions): """Set all settable DOFs to a given position. DOF/value pairs should be supplied as kwargs, e.g. as diff --git a/src/blop/bayesian/models.py b/src/blop/bayesian/models.py index 59ca6e9..7ee3faf 100644 --- a/src/blop/bayesian/models.py +++ b/src/blop/bayesian/models.py @@ -1,3 +1,4 @@ +import botorch import gpytorch import torch from botorch.models.gp_regression import SingleTaskGP @@ -5,6 +6,45 @@ from . import kernels +def train_model(model, hypers=None, **kwargs): + """Fit all of the agent's models. All kwargs are passed to `botorch.fit.fit_gpytorch_mll`.""" + if hypers is not None: + model.load_state_dict(hypers) + else: + botorch.fit.fit_gpytorch_mll(gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model), **kwargs) + model.trained = True + + +def construct_single_task_model(X, y, skew_dims=None, min_noise=1e-6, max_noise=1e0): + """ + Construct an untrained model for an objective. + """ + + likelihood = gpytorch.likelihoods.GaussianLikelihood( + noise_constraint=gpytorch.constraints.Interval( + torch.tensor(min_noise), + torch.tensor(max_noise), + ), + ) + + input_transform = botorch.models.transforms.input.Normalize(d=X.shape[-1]) + outcome_transform = botorch.models.transforms.outcome.Standardize(m=1) # , batch_shape=torch.Size((1,))) + + if not X.isfinite().all(): + raise ValueError("'X' must not contain points that are inf or NaN.") + if not y.isfinite().all(): + raise ValueError("'y' must not contain points that are inf or NaN.") + + return LatentGP( + train_inputs=X, + train_targets=y, + likelihood=likelihood, + skew_dims=skew_dims, + input_transform=input_transform, + outcome_transform=outcome_transform, + ) + + class LatentGP(SingleTaskGP): def __init__(self, train_inputs, train_targets, skew_dims=True, *args, **kwargs): super().__init__(train_inputs, train_targets, *args, **kwargs) diff --git a/src/blop/digestion/__init__.py b/src/blop/digestion/__init__.py index bdfd1b3..0728b11 100644 --- a/src/blop/digestion/__init__.py +++ b/src/blop/digestion/__init__.py @@ -1,5 +1,18 @@ import pandas as pd +from ..utils import get_beam_stats + def default_digestion_function(df: pd.DataFrame) -> pd.DataFrame: return df + + +def beam_stats_digestion(df: pd.DataFrame, image_key, **kwargs) -> pd.DataFrame: + for index, entry in df.iterrows(): + stats = get_beam_stats(entry.loc[image_key], **kwargs) + + for attr, value in stats.items(): + if attr not in ["bbox"]: + df.loc[index, attr] = value + + return df diff --git a/src/blop/plotting.py b/src/blop/plotting.py index 5492a65..91b7ca7 100644 --- a/src/blop/plotting.py +++ b/src/blop/plotting.py @@ -33,7 +33,7 @@ def _plot_fitness_objs_one_dof(agent, size=16, lw=1e0): test_model_inputs = agent.dofs(active=True).transform(test_inputs) for obj_index, obj in enumerate(fitness_objs): - obj_values = agent.train_targets(obj.name).squeeze(-1).numpy() + obj_values = agent.train_targets()(obj.name).numpy() color = DEFAULT_COLOR_LIST[obj_index] @@ -84,7 +84,7 @@ def _plot_constraint_objs_one_dof(agent, size=16, lw=1e0): val_ax = agent.obj_axes[obj_index, 0] con_ax = agent.obj_axes[obj_index, 1] - obj_values = agent.train_targets(obj.name).squeeze(-1).numpy() + obj_values = agent.train_targets()[obj.name].numpy() color = DEFAULT_COLOR_LIST[obj_index] @@ -158,7 +158,7 @@ def _plot_objs_many_dofs(agent, axes=(0, 1), shading="nearest", cmap=DEFAULT_COL model_inputs = agent.dofs(active=True).transform(test_inputs) for obj_index, obj in enumerate(agent.objectives): - targets = agent.train_targets(obj.name).squeeze(-1).numpy() + targets = agent.train_targets()[obj.name].numpy() values = obj._untransform(targets) @@ -547,7 +547,7 @@ def _plot_pareto_front(agent, obj_indices=(0, 1)): fig, ax = plt.subplots(1, 1, figsize=(6, 6)) - y = agent.train_targets(kind="fitness") + y = agent.train_targets(kind="fitness", concatenate=True) pareto_mask = agent.pareto_mask constraint = agent.evaluated_constraints.all(axis=-1) diff --git a/src/blop/sim/__init__.py b/src/blop/sim/__init__.py new file mode 100644 index 0000000..2a71bc9 --- /dev/null +++ b/src/blop/sim/__init__.py @@ -0,0 +1,2 @@ +from .beamline import Beamline, Detector # noqa +from .handlers import HDF5Handler # noqa diff --git a/src/blop/sim/beamline.py b/src/blop/sim/beamline.py new file mode 100644 index 0000000..b770473 --- /dev/null +++ b/src/blop/sim/beamline.py @@ -0,0 +1,170 @@ +import itertools +from collections import deque +from datetime import datetime +from pathlib import Path + +import h5py +import numpy as np +import scipy as sp +from event_model import compose_resource +from ophyd import Component as Cpt +from ophyd import Device, Signal +from ophyd.sim import NullStatus, new_uid +from ophyd.utils import make_dir_tree + +from ..utils import get_beam_stats +from .handlers import ExternalFileReference + + +class Detector(Device): + sum = Cpt(Signal, kind="hinted") + max = Cpt(Signal, kind="normal") + area = Cpt(Signal, kind="normal") + cen_x = Cpt(Signal, kind="hinted") + cen_y = Cpt(Signal, kind="hinted") + wid_x = Cpt(Signal, kind="hinted") + wid_y = Cpt(Signal, kind="hinted") + image = Cpt(ExternalFileReference, kind="normal") + image_shape = Cpt(Signal, value=(300, 400), kind="normal") + noise = Cpt(Signal, kind="normal") + + def __init__(self, root_dir: str = "/tmp/blop/sim", verbose: bool = True, noise: bool = True, *args, **kwargs): + super().__init__(*args, **kwargs) + + _ = make_dir_tree(datetime.now().year, base_path=root_dir) + + self._root_dir = root_dir + self._verbose = verbose + + # Used for the emulated cameras only. + self._img_dir = None + + # Resource/datum docs related variables. + self._asset_docs_cache = deque() + self._resource_document = None + self._datum_factory = None + + self.noise.put(noise) + + def trigger(self): + super().trigger() + raw_image = self.generate_beam(noise=self.noise.get()) + + current_frame = next(self._counter) + + self._dataset.resize((current_frame + 1, *self.image_shape.get())) + + self._dataset[current_frame, :, :] = raw_image + + datum_document = self._datum_factory(datum_kwargs={"frame": current_frame}) + self._asset_docs_cache.append(("datum", datum_document)) + + stats = get_beam_stats(raw_image) + self.image.put(datum_document["datum_id"]) + + for attr in ["max", "sum", "cen_x", "cen_y", "wid_x", "wid_y"]: + getattr(self, attr).put(stats[attr]) + + super().trigger() + + return NullStatus() + + def stage(self): + super().stage() + date = datetime.now() + self._assets_dir = date.strftime("%Y/%m/%d") + data_file = f"{new_uid()}.h5" + + self._resource_document, self._datum_factory, _ = compose_resource( + start={"uid": "needed for compose_resource() but will be discarded"}, + spec="HDF5", + root=self._root_dir, + resource_path=str(Path(self._assets_dir) / Path(data_file)), + resource_kwargs={}, + ) + + self._data_file = str(Path(self._resource_document["root"]) / Path(self._resource_document["resource_path"])) + + # now discard the start uid, a real one will be added later + self._resource_document.pop("run_start") + self._asset_docs_cache.append(("resource", self._resource_document)) + + self._h5file_desc = h5py.File(self._data_file, "x") + group = self._h5file_desc.create_group("/entry") + self._dataset = group.create_dataset( + "image", + data=np.full(fill_value=np.nan, shape=(1, *self.image_shape.get())), + maxshape=(None, *self.image_shape.get()), + chunks=(1, *self.image_shape.get()), + dtype="float64", + compression="lzf", + ) + self._counter = itertools.count() + + def unstage(self): + super().unstage() + del self._dataset + self._h5file_desc.close() + self._resource_document = None + self._datum_factory = None + + def collect_asset_docs(self): + items = list(self._asset_docs_cache) + self._asset_docs_cache.clear() + yield from items + + def generate_beam(self, noise: bool = True): + nx, ny = self.image_shape.get() + + x = np.linspace(-10, 10, ny) + y = np.linspace(-10, 10, nx) + X, Y = np.meshgrid(x, y) + + x0 = self.parent.kbh_ush.get() - self.parent.kbh_dsh.get() + y0 = self.parent.kbv_usv.get() - self.parent.kbv_dsv.get() + x_width = np.sqrt(0.2 + 5e-1 * (self.parent.kbh_ush.get() + self.parent.kbh_dsh.get() - 1) ** 2) + y_width = np.sqrt(0.1 + 5e-1 * (self.parent.kbv_usv.get() + self.parent.kbv_dsv.get() - 2) ** 2) + + beam = np.exp(-0.5 * (((X - x0) / x_width) ** 4 + ((Y - y0) / y_width) ** 4)) / ( + np.sqrt(2 * np.pi) * x_width * y_width + ) + + mask = X > self.parent.ssa_inboard.get() + mask &= X < self.parent.ssa_outboard.get() + mask &= Y > self.parent.ssa_lower.get() + mask &= Y < self.parent.ssa_upper.get() + mask = sp.ndimage.gaussian_filter(mask.astype(float), sigma=1) + + image = beam * mask + + if noise: + kx = np.fft.fftfreq(n=len(x), d=0.1) + ky = np.fft.fftfreq(n=len(y), d=0.1) + KX, KY = np.meshgrid(kx, ky) + + power_spectrum = 1 / (1e-2 + KX**2 + KY**2) + + white_noise = 1e-3 * np.random.standard_normal(size=X.shape) + pink_noise = 1e-3 * np.real(np.fft.ifft2(power_spectrum * np.fft.fft2(np.random.standard_normal(size=X.shape)))) + # background = 5e-3 * (X - Y) / X.max() + + image += white_noise + pink_noise + + return image + + +class Beamline(Device): + det = Cpt(Detector) + + kbh_ush = Cpt(Signal, kind="hinted") + kbh_dsh = Cpt(Signal, kind="hinted") + kbv_usv = Cpt(Signal, kind="hinted") + kbv_dsv = Cpt(Signal, kind="hinted") + + ssa_inboard = Cpt(Signal, value=-5.0, kind="hinted") + ssa_outboard = Cpt(Signal, value=5.0, kind="hinted") + ssa_lower = Cpt(Signal, value=-5.0, kind="hinted") + ssa_upper = Cpt(Signal, value=5.0, kind="hinted") + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) diff --git a/src/blop/sim/handlers.py b/src/blop/sim/handlers.py new file mode 100644 index 0000000..b95f310 --- /dev/null +++ b/src/blop/sim/handlers.py @@ -0,0 +1,31 @@ +import h5py +from area_detector_handlers.handlers import HandlerBase +from ophyd import Signal + + +class HDF5Handler(HandlerBase): + specs = {"HDF5"} + + def __init__(self, filename): + self._name = filename + + def __call__(self, frame): + with h5py.File(self._name, "r") as f: + entry = f["/entry/image"] + return entry[frame] + + +class ExternalFileReference(Signal): + """ + A pure software Signal that describe()s an image in an external file. + """ + + def describe(self): + resource_document_data = super().describe() + resource_document_data[self.name].update( + dict( + external="FILESTORE:", + dtype="array", + ) + ) + return resource_document_data diff --git a/src/blop/tests/conftest.py b/src/blop/tests/conftest.py index 9de6513..84acb41 100644 --- a/src/blop/tests/conftest.py +++ b/src/blop/tests/conftest.py @@ -11,6 +11,7 @@ from blop import DOF, Agent, Objective from blop.digestion.tests import chankong_and_haimes_digestion, sketchy_himmelblau_digestion from blop.dofs import BrownianMotion +from blop.sim import HDF5Handler @pytest.fixture(scope="function") @@ -23,6 +24,8 @@ def db(): except Exception: pass + db.reg.register_handler("HDF5", HDF5Handler, overwrite=True) + return db diff --git a/src/blop/tests/test_acqfs.py b/src/blop/tests/test_acqfs.py index 3548d18..bc03155 100644 --- a/src/blop/tests/test_acqfs.py +++ b/src/blop/tests/test_acqfs.py @@ -7,7 +7,7 @@ @pytest.mark.parametrize("agent", all_agents, indirect=True) def test_analytic_acqfs(agent, RE, db, acqf): agent.db = db - RE(agent.learn("qr", n=16)) + RE(agent.learn("qr", n=4)) RE(agent.learn(acqf, n=1)) getattr(agent, acqf) diff --git a/src/blop/tests/test_agents.py b/src/blop/tests/test_agents.py index 84f434c..598117d 100644 --- a/src/blop/tests/test_agents.py +++ b/src/blop/tests/test_agents.py @@ -10,7 +10,7 @@ def test_agent(agent, RE, db): """ agent.db = db - RE(agent.learn("qr", n=16)) + RE(agent.learn("qr", n=32)) best = agent.best assert [dof.name in best for dof in agent.dofs] diff --git a/src/blop/tests/test_sims.py b/src/blop/tests/test_sims.py new file mode 100644 index 0000000..5906ac1 --- /dev/null +++ b/src/blop/tests/test_sims.py @@ -0,0 +1,39 @@ +import numpy as np + +from blop import DOF, Agent, Objective +from blop.digestion import beam_stats_digestion +from blop.sim import Beamline + + +def test_kb_simulation(RE, db): + beamline = Beamline(name="bl") + beamline.det.noise.put(False) + + dofs = [ + DOF(description="KBV downstream", device=beamline.kbv_dsv, search_domain=(-5.0, 5.0)), + DOF(description="KBV upstream", device=beamline.kbv_usv, search_domain=(-5.0, 5.0)), + DOF(description="KBH downstream", device=beamline.kbh_dsh, search_domain=(-5.0, 5.0)), + DOF(description="KBH upstream", device=beamline.kbh_ush, search_domain=(-5.0, 5.0)), + ] + + objectives = [ + Objective(name="bl_det_sum", target="max", transform="log", trust_domain=(200, np.inf)), + Objective(name="bl_det_wid_x", target="min", transform="log", latent_groups=[("bl_kbh_dsh", "bl_kbh_ush")]), + Objective(name="bl_det_wid_y", target="min", transform="log", latent_groups=[("bl_kbv_dsv", "bl_kbv_usv")]), + ] + + agent = Agent( + dofs=dofs, + objectives=objectives, + detectors=[beamline.det], + digestion=beam_stats_digestion, + digestion_kwargs={"image_key": "bl_det_image"}, + verbose=True, + db=db, + tolerate_acquisition_errors=False, + enforce_all_objectives_valid=True, + train_every=3, + ) + + RE(agent.learn("qr", n=16)) + RE(agent.learn("qei", n=4, iterations=4)) diff --git a/src/blop/utils/__init__.py b/src/blop/utils/__init__.py index a25bf4f..92fabfb 100644 --- a/src/blop/utils/__init__.py +++ b/src/blop/utils/__init__.py @@ -4,6 +4,48 @@ import torch from ortools.constraint_solver import pywrapcp, routing_enums_pb2 +from . import functions # noqa + + +def get_beam_stats(image, threshold=0.5): + ny, nx = image.shape + + fim = image.copy() + fim -= np.median(fim, axis=0) + fim -= np.median(fim, axis=1)[:, None] + + fim = sp.ndimage.median_filter(fim, size=3) + fim = sp.ndimage.gaussian_filter(fim, sigma=1) + + m = fim > threshold * fim.max() + area = m.sum() + + cs_x = np.cumsum(m.sum(axis=0)) / area + cs_y = np.cumsum(m.sum(axis=1)) / area + + q_min, q_max = [0.15865, 0.84135] # one sigma + q_min, q_max = [0.05, 0.95] # 90% + + x_min, x_max = np.interp([q_min, q_max], cs_x, np.arange(nx)) + y_min, y_max = np.interp([q_min, q_max], cs_y, np.arange(ny)) + + stats = { + "max": fim.max(), + "sum": fim.sum(), + "area": area, + "cen_x": (x_min + x_max) / 2, + "cen_y": (y_min + y_max) / 2, + "wid_x": x_max - x_min, + "wid_y": y_max - y_min, + "x_min": x_min, + "x_max": x_max, + "y_min": y_min, + "y_max": y_max, + "bbox": [[x_min, x_max, x_max, x_min, x_min], [y_min, y_min, y_max, y_max, y_min]], + } + + return stats + def cummax(x): return [np.nanmax(x[: i + 1]) for i in range(len(np.atleast_1d(x)))] @@ -54,7 +96,7 @@ def route(start_point, points, dim_weights=1): """ total_points = np.r_[np.atleast_2d(start_point), points] - points_scale = total_points.ptp(axis=0) + points_scale = np.ptp(total_points, axis=0) dim_mask = points_scale > 0 if dim_mask.sum() == 0: @@ -102,130 +144,3 @@ def get_movement_time(x, v_max, a): return 2 * np.sqrt(np.abs(x) / a) * (np.abs(x) < v_max**2 / a) + (np.abs(x) / v_max + v_max / a) * ( np.abs(x) > v_max**2 / a ) - - -def get_principal_component_bounds(image, beam_prop=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 - - # print(f'{extent = }') - - u, s, v = np.linalg.svd(image - image.mean()) - separability = np.square(s[0]) / np.square(s).sum() - - # q refers to "quantile" - q_min, q_max = 0.5 * (1 - beam_prop), 0.5 * (1 + beam_prop) - - # these represent the cumulative proportion of the beam, as captured by the SVD. - cs_beam_x = np.cumsum(v[0]) / np.sum(v[0]) - cs_beam_y = np.cumsum(u[:, 0]) / np.sum(u[:, 0]) - cs_beam_x[0], cs_beam_y[0] = 0, 0 # change this lol - - # the first coordinate where the cumulative beam is greater than the minimum - i_q_min_x = np.where((cs_beam_x[1:] > q_min) & (cs_beam_x[:-1] < q_min))[0][0] - i_q_min_y = np.where((cs_beam_y[1:] > q_min) & (cs_beam_y[:-1] < q_min))[0][0] - - # the last coordinate where the cumulative beam is less than the maximum - i_q_max_x = np.where((cs_beam_x[1:] > q_max) & (cs_beam_x[:-1] < q_max))[0][-1] - i_q_max_y = np.where((cs_beam_y[1:] > q_max) & (cs_beam_y[:-1] < q_max))[0][-1] - - # interpolate, so that we can go finer than one pixel. this quartet is the "bounding box", from 0 to 1. - # (let's make this more efficient later) - x_min = np.interp(q_min, cs_beam_x[[i_q_min_x, i_q_min_x + 1]], [i_q_min_x, i_q_min_x + 1]) - x_max = np.interp(q_max, cs_beam_x[[i_q_max_x, i_q_max_x + 1]], [i_q_max_x, i_q_max_x + 1]) - y_min = np.interp(q_min, cs_beam_y[[i_q_min_y, i_q_min_y + 1]], [i_q_min_y, i_q_min_y + 1]) - y_max = np.interp(q_max, cs_beam_y[[i_q_max_y, i_q_max_y + 1]], [i_q_max_y, i_q_max_y + 1]) - - return ( - x_min, - x_max, - y_min, - 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. - """ - - 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 - - # 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 = ( - 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, - x_max, - 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 diff --git a/src/blop/utils/misc.py b/src/blop/utils/misc.py deleted file mode 100644 index 6f3bd85..0000000 --- a/src/blop/utils/misc.py +++ /dev/null @@ -1,227 +0,0 @@ -import botorch -import numpy as np -import scipy as sp -import torch -from ortools.constraint_solver import pywrapcp, routing_enums_pb2 - - -def sobol_sampler(bounds, n, q=1): - """ - Returns $n$ quasi-randomly sampled points within the bounds (a 2 by d tensor) - and batch size $q$ - """ - return botorch.utils.sampling.draw_sobol_samples(bounds, n=n, q=q) - - -def normalized_sobol_sampler(n, d): - """ - Returns $n$ quasi-randomly sampled points in the [0,1]^d hypercube - """ - normalized_bounds = torch.outer(torch.tensor([0, 1]), torch.ones(d)) - return sobol_sampler(normalized_bounds, n=n, q=1) - - -def estimate_root_indices(x): - # or, indices_before_sign_changes - i_whole = np.where(np.sign(x[1:]) != np.sign(x[:-1]))[0] - i_part = 1 - x[i_whole + 1] / (x[i_whole + 1] - x[i_whole]) - return i_whole + i_part - - -def _fast_psd_inverse(M): - """ - About twice as fast as np.linalg.inv for large, PSD matrices. - """ - cholesky, dpotrf_info = sp.linalg.lapack.dpotrf(M) - invM, dpotri_info = sp.linalg.lapack.dpotri(cholesky) - return np.where(invM, invM, invM.T) - - -def mprod(*M): - res = M[0] - for m in M[1:]: - res = np.matmul(res, m) - return res - - -def route(start_point, points): - """ - Returns the indices of the most efficient way to visit `points`, starting from `start_point`. - """ - - total_points = np.r_[np.atleast_2d(start_point), points] - points_scale = total_points.ptp(axis=0) - 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 - - manager = pywrapcp.RoutingIndexManager(len(total_points), 1, 0) # number of depots, number of salesmen, starting index - routing = pywrapcp.RoutingModel(manager) - - def delay_callback(from_index, to_index): - to_node = manager.IndexToNode(to_index) - if to_node == 0: - return 0 # it is free to return to the depot from anywhere; we just won't do it - from_node = manager.IndexToNode(from_index) - return delay_matrix[from_node][to_node] - - transit_callback_index = routing.RegisterTransitCallback(delay_callback) - routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index) - search_parameters = pywrapcp.DefaultRoutingSearchParameters() - search_parameters.first_solution_strategy = routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC - - solution = routing.SolveWithParameters(search_parameters) - - index = routing.Start(0) - route_indices, route_delays = [0], [] - while not routing.IsEnd(index): - previous_index = index - index = solution.Value(routing.NextVar(index)) - route_delays.append(routing.GetArcCostForVehicle(previous_index, index, 0)) - route_indices.append(index) - - # omit the first and last indices, which correspond to the start - return np.array(route_indices)[1:-1] - 1 - - -def get_movement_time(x, v_max, a): - """ - How long does it take an object to go distance $x$ with acceleration $a$ and maximum velocity $v_max$? - That's what this function answers. - """ - return 2 * np.sqrt(np.abs(x) / a) * (np.abs(x) < v_max**2 / a) + (np.abs(x) / v_max + v_max / a) * ( - np.abs(x) > v_max**2 / a - ) - - -def get_principal_component_bounds(image, beam_prop=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 - - # print(f'{extent = }') - - u, s, v = np.linalg.svd(image - image.mean()) - separability = np.square(s[0]) / np.square(s).sum() - - # q refers to "quantile" - q_min, q_max = 0.5 * (1 - beam_prop), 0.5 * (1 + beam_prop) - - # these represent the cumulative proportion of the beam, as captured by the SVD. - cs_beam_x = np.cumsum(v[0]) / np.sum(v[0]) - cs_beam_y = np.cumsum(u[:, 0]) / np.sum(u[:, 0]) - cs_beam_x[0], cs_beam_y[0] = 0, 0 # change this lol - - # the first coordinate where the cumulative beam is greater than the minimum - i_q_min_x = np.where((cs_beam_x[1:] > q_min) & (cs_beam_x[:-1] < q_min))[0][0] - i_q_min_y = np.where((cs_beam_y[1:] > q_min) & (cs_beam_y[:-1] < q_min))[0][0] - - # the last coordinate where the cumulative beam is less than the maximum - i_q_max_x = np.where((cs_beam_x[1:] > q_max) & (cs_beam_x[:-1] < q_max))[0][-1] - i_q_max_y = np.where((cs_beam_y[1:] > q_max) & (cs_beam_y[:-1] < q_max))[0][-1] - - # interpolate, so that we can go finer than one pixel. this quartet is the "bounding box", from 0 to 1. - # (let's make this more efficient later) - x_min = np.interp(q_min, cs_beam_x[[i_q_min_x, i_q_min_x + 1]], [i_q_min_x, i_q_min_x + 1]) - x_max = np.interp(q_max, cs_beam_x[[i_q_max_x, i_q_max_x + 1]], [i_q_max_x, i_q_max_x + 1]) - y_min = np.interp(q_min, cs_beam_y[[i_q_min_y, i_q_min_y + 1]], [i_q_min_y, i_q_min_y + 1]) - y_max = np.interp(q_max, cs_beam_y[[i_q_max_y, i_q_max_y + 1]], [i_q_max_y, i_q_max_y + 1]) - - return ( - x_min, - x_max, - y_min, - 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. - """ - - 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 - - # 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 = ( - 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, - x_max, - 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 diff --git a/src/blop/utils/prepare_re_env.py b/src/blop/utils/prepare_re_env.py index fdf715e..01c7484 100644 --- a/src/blop/utils/prepare_re_env.py +++ b/src/blop/utils/prepare_re_env.py @@ -12,6 +12,8 @@ from databroker import Broker from ophyd.utils import make_dir_tree +from blop.sim import HDF5Handler + DEFAULT_DB_TYPE = "local" DEFAULT_ROOT_DIR = "/tmp/sirepo-bluesky-data" DEFAULT_ENV_TYPE = "stepper" @@ -24,6 +26,7 @@ def re_env(db_type=DEFAULT_DB_TYPE, root_dir=DEFAULT_ROOT_DIR): RE.subscribe(bec) db = Broker.named(db_type) + db.reg.register_handler("HDF5", HDF5Handler, overwrite=True) try: databroker.assets.utils.install_sentinels(db.reg.config, version=1) except Exception: