From 2b05f34fbbe260738a3062aef9dc3f1c37361612 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Sat, 6 Jul 2024 21:55:09 -0400 Subject: [PATCH] add simulated KB mirrors --- pyproject.toml | 1 + src/blop/agent.py | 2 +- src/blop/sim/__init__.py | 1 + src/blop/sim/beamline.py | 199 ++++++++++++++++++++++++++++++++++++ src/blop/tests/test_sims.py | 35 +++++++ src/blop/utils/__init__.py | 163 +++++++---------------------- 6 files changed, 273 insertions(+), 128 deletions(-) create mode 100644 src/blop/sim/__init__.py create mode 100644 src/blop/sim/beamline.py create mode 100644 src/blop/tests/test_sims.py diff --git a/pyproject.toml b/pyproject.toml index 8b532d9..6b4bcb1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ maintainers = [ requires-python = ">=3.9" dependencies = [ + "area-detector-handlers", "bluesky", "botorch", "databroker", diff --git a/src/blop/agent.py b/src/blop/agent.py index 80ef55f..a7de1aa 100644 --- a/src/blop/agent.py +++ b/src/blop/agent.py @@ -842,7 +842,7 @@ 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) + # x = self.dofs(active=True).transform(x) p = torch.ones(x.shape[:-1]) for obj in self.objectives(active=True): diff --git a/src/blop/sim/__init__.py b/src/blop/sim/__init__.py new file mode 100644 index 0000000..70c7e45 --- /dev/null +++ b/src/blop/sim/__init__.py @@ -0,0 +1 @@ +from .beamline import Beamline, Detector # noqa diff --git a/src/blop/sim/beamline.py b/src/blop/sim/beamline.py new file mode 100644 index 0000000..a70b966 --- /dev/null +++ b/src/blop/sim/beamline.py @@ -0,0 +1,199 @@ +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 area_detector_handlers.handlers import HandlerBase +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 + + +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 + + +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() + for item in items: + yield item + + 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.5 + 5e-1 * (self.parent.kbh_ush.get() + self.parent.kbh_dsh.get() - 1) ** 2) + y_width = np.sqrt(0.25 + 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 = 5e-3 * np.random.standard_normal(size=X.shape) + pink_noise = 5e-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/tests/test_sims.py b/src/blop/tests/test_sims.py new file mode 100644 index 0000000..25a7e3b --- /dev/null +++ b/src/blop/tests/test_sims.py @@ -0,0 +1,35 @@ +import numpy as np + +from blop import DOF, Agent, Objective +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=(1, np.inf)), + Objective(name="bl_det_wid_x", target="min", transform="log"), # , latent_groups=[("x1", "x2")]), + Objective(name="bl_det_wid_y", target="min", transform="log"), # , latent_groups=[("x1", "x2")])] + ] + + agent = Agent( + dofs=dofs, + objectives=objectives, + dets=[beamline.det], + verbose=True, + db=db, + tolerate_acquisition_errors=False, + train_every=3, + ) + + RE(agent.learn("qr", n=32)) + RE(agent.learn("qei", n=4, iterations=4)) diff --git a/src/blop/utils/__init__.py b/src/blop/utils/__init__.py index a25bf4f..7ff1c8b 100644 --- a/src/blop/utils/__init__.py +++ b/src/blop/utils/__init__.py @@ -5,6 +5,42 @@ from ortools.constraint_solver import pywrapcp, routing_enums_pb2 +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, + "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)))] @@ -102,130 +138,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